X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/8303a85b6c3fed92710a590fa5653557b73439cf..ebb2559b3a09636ff687316bbe512e0e8a86b168:/sdk/python/arvados/commands/put.py diff --git a/sdk/python/arvados/commands/put.py b/sdk/python/arvados/commands/put.py index e5c85a92a7..714281cc95 100644 --- a/sdk/python/arvados/commands/put.py +++ b/sdk/python/arvados/commands/put.py @@ -24,6 +24,7 @@ import tempfile import threading import time from apiclient import errors as apiclient_errors +from arvados._version import __version__ import arvados.commands._util as arv_cmd @@ -32,6 +33,9 @@ api_client = None upload_opts = argparse.ArgumentParser(add_help=False) +upload_opts.add_argument('--version', action='version', + version="%s %s" % (sys.argv[0], __version__), + help='Print version and exit.') upload_opts.add_argument('paths', metavar='path', type=str, nargs='*', help=""" Local file or directory. Default: read from standard input. @@ -48,6 +52,12 @@ Normalize the manifest by re-ordering files and streams after writing data. """) +_group.add_argument('--dry-run', action='store_true', default=False, + help=""" +Don't actually upload files, but only check if any file should be +uploaded. Exit with code=2 when files are pending for upload. +""") + _group = upload_opts.add_mutually_exclusive_group() _group.add_argument('--as-stream', action='store_true', dest='stream', @@ -223,6 +233,29 @@ class ResumeCacheConflict(Exception): pass +class ArvPutArgumentConflict(Exception): + pass + + +class ArvPutUploadIsPending(Exception): + pass + + +class ArvPutUploadNotPending(Exception): + pass + + +class FileUploadList(list): + def __init__(self, dry_run=False): + list.__init__(self) + self.dry_run = dry_run + + def append(self, other): + if self.dry_run: + raise ArvPutUploadIsPending() + super(FileUploadList, self).append(other) + + class ResumeCache(object): CACHE_DIR = '.cache/arvados/arv-put' @@ -315,7 +348,8 @@ class ArvPutUploadJob(object): def __init__(self, paths, resume=True, use_cache=True, reporter=None, bytes_expected=None, name=None, owner_uuid=None, ensure_unique_name=False, num_retries=None, replication_desired=None, - filename=None, update_time=20.0, update_collection=None): + filename=None, update_time=20.0, update_collection=None, + logger=logging.getLogger('arvados.arv_put'), dry_run=False): self.paths = paths self.resume = resume self.use_cache = use_cache @@ -340,11 +374,18 @@ class ArvPutUploadJob(object): self._file_paths = [] # Files to be updated in remote collection self._stop_checkpointer = threading.Event() self._checkpointer = threading.Thread(target=self._update_task) + self._checkpointer.daemon = True self._update_task_time = update_time # How many seconds wait between update runs - self.logger = logging.getLogger('arvados.arv_put') + self._files_to_upload = FileUploadList(dry_run=dry_run) + self.logger = logger + self.dry_run = dry_run if not self.use_cache and self.resume: - raise ArgumentError('resume cannot be True when use_cache is False') + raise ArvPutArgumentConflict('resume cannot be True when use_cache is False') + + # Check for obvious dry-run responses + if self.dry_run and (not self.use_cache or not self.resume): + raise ArvPutUploadIsPending() # Load cached data if any and if needed self._setup_state(update_collection) @@ -353,12 +394,14 @@ class ArvPutUploadJob(object): """ Start supporting thread & file uploading """ - self._checkpointer.daemon = True - self._checkpointer.start() + if not self.dry_run: + self._checkpointer.start() try: for path in self.paths: # Test for stdin first, in case some file named '-' exist if path == '-': + if self.dry_run: + raise ArvPutUploadIsPending() self._write_stdin(self.filename or 'stdin') elif os.path.isdir(path): # Use absolute paths on cache index so CWD doesn't interfere @@ -371,24 +414,38 @@ class ArvPutUploadJob(object): dirs.sort() files.sort() for f in files: - self._write_file(os.path.join(root, f), + self._check_file(os.path.join(root, f), os.path.join(root[len(prefixdir):], f)) else: - self._write_file(os.path.abspath(path), + self._check_file(os.path.abspath(path), self.filename or os.path.basename(path)) + # If dry-mode is on, and got up to this point, then we should notify that + # there aren't any file to upload. + if self.dry_run: + raise ArvPutUploadNotPending() + # Remove local_collection's files that don't exist locally anymore, so the + # bytes_written count is correct. + for f in self.collection_file_paths(self._local_collection, + path_prefix=""): + if f != 'stdin' and f != self.filename and not f in self._file_paths: + self._local_collection.remove(f) + # Update bytes_written from current local collection and + # report initial progress. + self._update() + # Actual file upload + self._upload_files() finally: - # Stop the thread before doing anything else - self._stop_checkpointer.set() - self._checkpointer.join() - # Commit all pending blocks & one last _update() - self._local_collection.manifest_text() - self._update(final=True) + if not self.dry_run: + # Stop the thread before doing anything else + self._stop_checkpointer.set() + self._checkpointer.join() + # Commit all pending blocks & one last _update() + self._local_collection.manifest_text() + self._update(final=True) + if save_collection: + self.save_collection() if self.use_cache: self._cache_file.close() - if save_collection: - self.save_collection() - # Correct the final written bytes count - self.bytes_written -= self.bytes_skipped def save_collection(self): if self.update: @@ -412,7 +469,7 @@ class ArvPutUploadJob(object): num_retries=self.num_retries) def destroy_cache(self): - if self.resume: + if self.use_cache: try: os.unlink(self._cache_filename) except OSError as error: @@ -468,11 +525,11 @@ class ArvPutUploadJob(object): self._write(sys.stdin, output) output.close() - def _write_file(self, source, filename): + def _check_file(self, source, filename): + """Check if this file needs to be uploaded""" resume_offset = 0 should_upload = False new_file_in_cache = False - # Record file path for updating the remote collection before exiting self._file_paths.append(filename) @@ -501,22 +558,35 @@ class ArvPutUploadJob(object): if not file_in_local_collection: # File not uploaded yet, upload it completely should_upload = True + elif file_in_local_collection.permission_expired(): + # Permission token expired, re-upload file. This will change whenever + # we have a API for refreshing tokens. + should_upload = True + self._local_collection.remove(filename) elif cached_file_data['size'] == file_in_local_collection.size(): # File already there, skip it. self.bytes_skipped += cached_file_data['size'] elif cached_file_data['size'] > file_in_local_collection.size(): # File partially uploaded, resume! resume_offset = file_in_local_collection.size() + self.bytes_skipped += resume_offset should_upload = True else: # Inconsistent cache, re-upload the file should_upload = True + self._local_collection.remove(filename) self.logger.warning("Uploaded version of file '{}' is bigger than local version, will re-upload it from scratch.".format(source)) # Local file differs from cached data, re-upload it. else: + if file_in_local_collection: + self._local_collection.remove(filename) should_upload = True if should_upload: + self._files_to_upload.append((source, resume_offset, filename)) + + def _upload_files(self): + for source, resume_offset, filename in self._files_to_upload: with open(source, 'r') as source_fd: with self._state_lock: self._state['files'][source]['mtime'] = os.path.getmtime(source) @@ -525,7 +595,6 @@ class ArvPutUploadJob(object): # Start upload where we left off output = self._local_collection.open(filename, 'a') source_fd.seek(resume_offset) - self.bytes_skipped += resume_offset else: # Start from scratch output = self._local_collection.open(filename, 'w') @@ -568,9 +637,14 @@ class ArvPutUploadJob(object): if self.filename: md5.update(self.filename) cache_filename = md5.hexdigest() - self._cache_file = open(os.path.join( + cache_filepath = os.path.join( arv_cmd.make_home_conf_dir(self.CACHE_DIR, 0o700, 'raise'), - cache_filename), 'a+') + cache_filename) + if self.resume: + self._cache_file = open(cache_filepath, 'a+') + else: + # --no-resume means start with a empty cache file. + self._cache_file = open(cache_filepath, 'w+') self._cache_filename = self._cache_file.name self._lock_file(self._cache_file) self._cache_file.seek(0) @@ -590,9 +664,17 @@ class ArvPutUploadJob(object): self._state = copy.deepcopy(self.EMPTY_STATE) # Load the previous manifest so we can check if files were modified remotely. self._local_collection = arvados.collection.Collection(self._state['manifest'], replication_desired=self.replication_desired) - # Load how many bytes were uploaded on previous run - with self._collection_lock: - self.bytes_written = self._collection_size(self._local_collection) + + def collection_file_paths(self, col, path_prefix='.'): + """Return a list of file paths by recursively go through the entire collection `col`""" + file_paths = [] + for name, item in col.items(): + if isinstance(item, arvados.arvfile.ArvadosFile): + file_paths.append(os.path.join(path_prefix, name)) + elif isinstance(item, arvados.collection.Subcollection): + new_prefix = os.path.join(path_prefix, name) + file_paths += self.collection_file_paths(item, path_prefix=new_prefix) + return file_paths def _lock_file(self, fileobj): try: @@ -718,6 +800,7 @@ def desired_project_uuid(api_client, project_uuid, num_retries): def main(arguments=None, stdout=sys.stdout, stderr=sys.stderr): global api_client + logger = logging.getLogger('arvados.arv_put') args = parse_arguments(arguments) status = 0 if api_client is None: @@ -726,7 +809,10 @@ def main(arguments=None, stdout=sys.stdout, stderr=sys.stderr): # Determine the name to use if args.name: if args.stream or args.raw: - print >>stderr, "Cannot use --name with --stream or --raw" + logger.error("Cannot use --name with --stream or --raw") + sys.exit(1) + elif args.update_collection: + logger.error("Cannot use --name with --update-collection") sys.exit(1) collection_name = args.name else: @@ -736,7 +822,7 @@ def main(arguments=None, stdout=sys.stdout, stderr=sys.stderr): socket.gethostname()) if args.project_uuid and (args.stream or args.raw): - print >>stderr, "Cannot use --project-uuid with --stream or --raw" + logger.error("Cannot use --project-uuid with --stream or --raw") sys.exit(1) # Determine the parent project @@ -744,7 +830,7 @@ def main(arguments=None, stdout=sys.stdout, stderr=sys.stderr): project_uuid = desired_project_uuid(api_client, args.project_uuid, args.retries) except (apiclient_errors.Error, ValueError) as error: - print >>stderr, error + logger.error(error) sys.exit(1) if args.progress: @@ -768,38 +854,53 @@ def main(arguments=None, stdout=sys.stdout, stderr=sys.stderr): name = collection_name, owner_uuid = project_uuid, ensure_unique_name = True, - update_collection = args.update_collection) + update_collection = args.update_collection, + logger=logger, + dry_run=args.dry_run) except ResumeCacheConflict: - print >>stderr, "\n".join([ + logger.error("\n".join([ "arv-put: Another process is already uploading this data.", - " Use --no-cache if this is really what you want."]) + " Use --no-cache if this is really what you want."])) sys.exit(1) except CollectionUpdateError as error: - print >>stderr, "\n".join([ - "arv-put: %s" % str(error)]) + logger.error("\n".join([ + "arv-put: %s" % str(error)])) sys.exit(1) + except ArvPutUploadIsPending: + # Dry run check successful, return proper exit code. + sys.exit(2) + except ArvPutUploadNotPending: + # No files pending for upload + sys.exit(0) # Install our signal handler for each code in CAUGHT_SIGNALS, and save # the originals. orig_signal_handlers = {sigcode: signal.signal(sigcode, exit_signal_handler) for sigcode in CAUGHT_SIGNALS} - if not args.update_collection and args.resume and writer.bytes_written > 0: - print >>stderr, "\n".join([ - "arv-put: Resuming previous upload from last checkpoint.", - " Use the --no-resume option to start over."]) + if not args.dry_run and not args.update_collection and args.resume and writer.bytes_written > 0: + logger.warning("\n".join([ + "arv-put: Resuming previous upload from last checkpoint.", + " Use the --no-resume option to start over."])) - writer.report_progress() + if not args.dry_run: + writer.report_progress() output = None try: writer.start(save_collection=not(args.stream or args.raw)) except arvados.errors.ApiError as error: - print >>stderr, "\n".join([ - "arv-put: %s" % str(error)]) + logger.error("\n".join([ + "arv-put: %s" % str(error)])) sys.exit(1) + except ArvPutUploadIsPending: + # Dry run check successful, return proper exit code. + sys.exit(2) + except ArvPutUploadNotPending: + # No files pending for upload + sys.exit(0) if args.progress: # Print newline to split stderr from stdout for humans. - print >>stderr + logger.info("\n") if args.stream: if args.normalize: @@ -811,15 +912,15 @@ def main(arguments=None, stdout=sys.stdout, stderr=sys.stderr): else: try: if args.update_collection: - print >>stderr, "Collection updated: '{}'".format(writer.collection_name()) + logger.info("Collection updated: '{}'".format(writer.collection_name())) else: - print >>stderr, "Collection saved as '{}'".format(writer.collection_name()) + logger.info("Collection saved as '{}'".format(writer.collection_name())) if args.portable_data_hash: output = writer.portable_data_hash() else: output = writer.manifest_locator() except apiclient_errors.Error as error: - print >>stderr, ( + logger.error( "arv-put: Error creating Collection on project: {}.".format( error)) status = 1