Merge branch '8784-dir-listings'

[arvados.git] / crunch_scripts / crunchutil / robust_put.py

diff --git a/crunch_scripts/crunchutil/robust_put.py b/crunch_scripts/crunchutil/robust_put.py
index 158ceb13e68f9f8742947cd440054b391ded631d..27b0bf34567b0489da7095be63f836344bb7b47c 100644 (file)
--- a/crunch_scripts/crunchutil/robust_put.py
+++ b/crunch_scripts/crunchutil/robust_put.py
@@ -1,7 +1,12 @@
+# Copyright (C) The Arvados Authors. All rights reserved.
+#
+# SPDX-License-Identifier: Apache-2.0
+
 import arvados
 import arvados.commands.put as put
 import os
 import logging
+import time
 
 def machine_progress(bytes_written, bytes_expected):
     return "upload wrote {} total {}\n".format(
@@ -15,7 +20,10 @@ class Args(object):
 
 # Upload to Keep with error recovery.
 # Return a uuid or raise an exception if there are too many failures.
-def upload(source_dir):
+def upload(source_dir, logger=None):
+    if logger is None:
+        logger = logging.getLogger("arvados")
+
     source_dir = os.path.abspath(source_dir)
     done = False
     if 'TASK_WORK' in os.environ:
@@ -34,15 +42,15 @@ def upload(source_dir):
             outuuid = out.finish()
             done = True
         except KeyboardInterrupt as e:
-            logging.critical("caught interrupt signal 2")
+            logger.critical("caught interrupt signal 2")
             raise e
         except Exception as e:
-            logging.exception("caught exception:")
+            logger.exception("caught exception:")
             backoff *= 2
             if backoff > 256:
-                logging.critical("Too many upload failures, giving up")
+                logger.critical("Too many upload failures, giving up")
                 raise e
             else:
-                logging.warning("Sleeping for %s seconds before trying again" % backoff)
+                logger.warning("Sleeping for %s seconds before trying again" % backoff)
                 time.sleep(backoff)
     return outuuid