X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/a3222e35cda68c8e48a17921c33ac37ecb5c3bac..0eb72b526bf8bbb011551ecf019f604e17a534f1:/crunch_scripts/crunchutil/robust_put.py

diff --git a/crunch_scripts/crunchutil/robust_put.py b/crunch_scripts/crunchutil/robust_put.py
index 158ceb13e6..27b0bf3456 100644
--- 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