crunch_scripts/crunchutil/robust_put.py

   1 import arvados
   2 import arvados.commands.put as put
   3 import os
   4 import logging
   5 import time
   6
   7 def machine_progress(bytes_written, bytes_expected):
   8     return "upload wrote {} total {}\n".format(
   9         bytes_written, -1 if (bytes_expected is None) else bytes_expected)
  10
  11 class Args(object):
  12     def __init__(self, fn):
  13         self.filename = None
  14         self.paths = [fn]
  15         self.max_manifest_depth = 0
  16
  17 # Upload to Keep with error recovery.
  18 # Return a uuid or raise an exception if there are too many failures.
  19 def upload(source_dir):
  20     source_dir = os.path.abspath(source_dir)
  21     done = False
  22     if 'TASK_WORK' in os.environ:
  23         resume_cache = put.ResumeCache(os.path.join(arvados.current_task().tmpdir, "upload-output-checkpoint"))
  24     else:
  25         resume_cache = put.ResumeCache(put.ResumeCache.make_path(Args(source_dir)))
  26     reporter = put.progress_writer(machine_progress)
  27     bytes_expected = put.expected_bytes_for([source_dir])
  28     backoff = 1
  29     outuuid = None
  30     while not done:
  31         try:
  32             out = put.ArvPutCollectionWriter.from_cache(resume_cache, reporter, bytes_expected)
  33             out.do_queued_work()
  34             out.write_directory_tree(source_dir, max_manifest_depth=0)
  35             outuuid = out.finish()
  36             done = True
  37         except KeyboardInterrupt as e:
  38             logging.critical("caught interrupt signal 2")
  39             raise e
  40         except Exception as e:
  41             logging.exception("caught exception:")
  42             backoff *= 2
  43             if backoff > 256:
  44                 logging.critical("Too many upload failures, giving up")
  45                 raise e
  46             else:
  47                 logging.warning("Sleeping for %s seconds before trying again" % backoff)
  48                 time.sleep(backoff)
  49     return outuuid