X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/8f5755f5ec65da08ffe9dfc6856848ec747a31a8..436d545670d3637dcd96b429e538b995ee98a49b:/sdk/python/arvados/commands/put.py diff --git a/sdk/python/arvados/commands/put.py b/sdk/python/arvados/commands/put.py index 86670266ae..01bae2fead 100644 --- a/sdk/python/arvados/commands/put.py +++ b/sdk/python/arvados/commands/put.py @@ -11,9 +11,12 @@ import fcntl import hashlib import json import os +import signal import sys import tempfile +CAUGHT_SIGNALS = [signal.SIGINT, signal.SIGQUIT, signal.SIGTERM] + def parse_arguments(arguments): parser = argparse.ArgumentParser( description='Copy data from the local filesystem to Keep.') @@ -169,7 +172,7 @@ class ResumeCache(object): md5 = hashlib.md5() md5.update(arvados.config.get('ARVADOS_API_HOST', '!nohost')) realpaths = sorted(os.path.realpath(path) for path in args.paths) - md5.update(''.join(realpaths)) + md5.update('\0'.join(realpaths)) if any(os.path.isdir(path) for path in realpaths): md5.update(str(max(args.max_manifest_depth, -1))) elif args.filename: @@ -227,10 +230,7 @@ class ArvPutCollectionWriter(arvados.ResumableCollectionWriter): self.bytes_written = 0 self._seen_inputs = [] self.cache = cache - if reporter is None: - self.report_progress = lambda bytes_w, bytes_e: None - else: - self.report_progress = reporter + self.reporter = reporter self.bytes_expected = bytes_expected super(ArvPutCollectionWriter, self).__init__() @@ -246,11 +246,9 @@ class ArvPutCollectionWriter(arvados.ResumableCollectionWriter): else: return writer - def preresume_hook(self): - print >>sys.stderr, "arv-put: Resuming previous upload. Bypass with the --no-resume option." - self.report_progress(self.bytes_written, self.bytes_expected) - def cache_state(self): + if self.cache is None: + return state = self.dump_state() # Transform attributes for serialization. for attr, value in state.items(): @@ -260,11 +258,19 @@ class ArvPutCollectionWriter(arvados.ResumableCollectionWriter): state[attr] = list(value) self.cache.save(state) + def report_progress(self): + if self.reporter is not None: + self.reporter(self.bytes_written, self.bytes_expected) + def flush_data(self): - bytes_buffered = self._data_buffer_len + start_buffer_len = self._data_buffer_len + start_block_count = self.bytes_written / self.KEEP_BLOCK_SIZE super(ArvPutCollectionWriter, self).flush_data() - self.bytes_written += (bytes_buffered - self._data_buffer_len) - self.report_progress(self.bytes_written, self.bytes_expected) + if self._data_buffer_len < start_buffer_len: # We actually PUT data. + self.bytes_written += (start_buffer_len - self._data_buffer_len) + self.report_progress() + if (self.bytes_written / self.KEEP_BLOCK_SIZE) > start_block_count: + self.cache_state() def _record_new_input(self, input_type, source_name, dest_name): # The key needs to be a list because that's what we'll get back @@ -319,8 +325,10 @@ def progress_writer(progress_func, outfile=sys.stderr): outfile.write(progress_func(bytes_written, bytes_expected)) return write_progress +def exit_signal_handler(sigcode, frame): + sys.exit(-sigcode) + def main(arguments=None): - ResumeCache.setup_user_cache() args = parse_arguments(arguments) if args.progress: @@ -329,39 +337,55 @@ def main(arguments=None): reporter = progress_writer(machine_progress) else: reporter = None + bytes_expected = expected_bytes_for(args.paths) try: + ResumeCache.setup_user_cache() resume_cache = ResumeCache(ResumeCache.make_path(args)) - if not args.resume: - resume_cache.restart() + except (IOError, OSError): + # Couldn't open cache directory/file. Continue without it. + resume_cache = None + writer = ArvPutCollectionWriter(resume_cache, reporter, bytes_expected) except ResumeCacheConflict: print "arv-put: Another process is already uploading this data." sys.exit(1) - - try: + else: + if not args.resume: + resume_cache.restart() writer = ArvPutCollectionWriter.from_cache( - resume_cache, reporter, expected_bytes_for(args.paths)) - - # Copy file data to Keep. - for path in args.paths: - if os.path.isdir(path): - writer.write_directory_tree( - path, max_manifest_depth=args.max_manifest_depth) - else: - writer.start_new_stream() - writer.write_file(path, args.filename or os.path.basename(path)) - except (Exception, KeyboardInterrupt): - writer.cache_state() - raise + resume_cache, reporter, bytes_expected) + + # 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 writer.bytes_written > 0: # We're resuming a previous upload. + print >>sys.stderr, "\n".join([ + "arv-put: Resuming previous upload from last checkpoint.", + " Use the --no-resume option to start over."]) + writer.report_progress() + + writer.do_queued_work() # Do work resumed from cache. + for path in args.paths: # Copy file data to Keep. + if os.path.isdir(path): + writer.write_directory_tree( + path, max_manifest_depth=args.max_manifest_depth) + else: + writer.start_new_stream() + writer.write_file(path, args.filename or os.path.basename(path)) + writer.finish_current_stream() + + if args.progress: # Print newline to split stderr from stdout for humans. + print >>sys.stderr if args.stream: print writer.manifest_text(), elif args.raw: - writer.finish_current_stream() print ','.join(writer.data_locators()) else: # Register the resulting collection in Arvados. - arvados.api().collections().create( + collection = arvados.api().collections().create( body={ 'uuid': writer.finish(), 'manifest_text': writer.manifest_text(), @@ -369,8 +393,13 @@ def main(arguments=None): ).execute() # Print the locator (uuid) of the new collection. - print writer.finish() - resume_cache.destroy() + print collection['uuid'] + + for sigcode, orig_handler in orig_signal_handlers.items(): + signal.signal(sigcode, orig_handler) + + if resume_cache is not None: + resume_cache.destroy() if __name__ == '__main__': main()