zstash create refactored

Author: forsyth2

zstash/create.py

@@ -80,17 +80,17 @@ def create():
 
     # Transfer to HPSS. Always keep a local copy of the database.
     logger.debug(f"{ts_utc()}: calling hpss_put() for {command_info.get_db_name()}")
-    # TODO: (A) Continue refactoring from here
     hpss_put(command_info, command_info.get_db_name(), is_index=True)
 
-    logger.debug(f"{ts_utc()}: calling globus_finalize()")
-    globus_finalize(non_blocking=args.non_blocking)
+    if command_info.globus_info:
+        logger.debug(f"{ts_utc()}: calling globus_finalize()")
+        globus_finalize(command_info.globus_info, non_blocking=args.non_blocking)
 
     if len(failures) > 0:
         # List the failures
         logger.warning("Some files could not be archived")
         for file_path in failures:
-            logger.error("Failed to archive {}".format(file_path))
+            logger.error(f"Failed to archive {file_path}")
 
 
 def setup_create(ci: CommandInfo) -> argparse.Namespace:

zstash/globus.py

@@ -31,13 +31,6 @@
     r"perlmutter.*\.nersc\.gov": "6bdc7956-fc0f-4ad2-989c-7aa5ee643a79",
 }
 
-remote_endpoint = None
-local_endpoint = None
-transfer_client: TransferClient = None
-transfer_data: TransferData = None
-task_id = None
-archive_directory_listing: IterableTransferResponse = None
-
 
 def check_endpoint_version_5(ep_id):
     output = transfer_client.get_endpoint(ep_id)
@@ -49,13 +42,13 @@ def check_endpoint_version_5(ep_id):
     return False
 
 
-def submit_transfer_with_checks(transfer_data):
+def submit_transfer_with_checks(globus_info: GlobusInfo):
     try:
-        task = transfer_client.submit_transfer(transfer_data)
+        task = globus_info.transfer_client.submit_transfer(globus_info.transfer_data)
     except TransferAPIError as err:
         if err.info.consent_required:
             scopes = "urn:globus:auth:scope:transfer.api.globus.org:all["
-            for ep_id in [remote_endpoint, local_endpoint]:
+            for ep_id in [globus_info.remote_endpoint, globus_info.local_endpoint]:
                 if check_endpoint_version_5(ep_id):
                     scopes += f" *https://auth.globus.org/scopes/{ep_id}/data_access"
             scopes += " ]"
@@ -147,72 +140,61 @@ def file_exists(name: str) -> bool:
     return False
 
 
-global_variable_tarfiles_pushed = 0
-
-
 # C901 'globus_transfer' is too complex (20)
 def globus_transfer(  # noqa: C901
-    remote_ep: str, remote_path: str, name: str, transfer_type: str, non_blocking: bool
+    globus_info: GlobusInfo, remote_ep: str, remote_path: str, name: str, transfer_type: str, non_blocking: bool
 ):
-    global transfer_client
-    global local_endpoint
-    global remote_endpoint
-    global transfer_data
-    global task_id
-    global archive_directory_listing
-    global global_variable_tarfiles_pushed
 
     logger.info(f"{ts_utc()}: Entered globus_transfer() for name = {name}")
     logger.debug(f"{ts_utc()}: non_blocking = {non_blocking}")
-    if not transfer_client:
+    if not globus_info.transfer_client:
         globus_activate("globus://" + remote_ep)
-    if not transfer_client:
+    if not globus_info.transfer_client:
+        logger.info(f"{ts_utc()}: Could not instantiate transfer client.")
         sys.exit(1)
 
     if transfer_type == "get":
-        if not archive_directory_listing:
-            archive_directory_listing = transfer_client.operation_ls(
-                remote_endpoint, remote_path
+        if not globus_info.archive_directory_listing:
+            globus_info.archive_directory_listing = globus_info.transfer_client.operation_ls(
+                globus_info.remote_endpoint, remote_path
             )
         if not file_exists(name):
             logger.error(
-                "Remote file globus://{}{}/{} does not exist".format(
-                    remote_ep, remote_path, name
-                )
+                f"Remote file globus://{remote_ep}{remote_path}/{name} does not exist"
             )
             sys.exit(1)
 
     if transfer_type == "get":
-        src_ep = remote_endpoint
+        src_ep = globus_info.remote_endpoint
         src_path = os.path.join(remote_path, name)
-        dst_ep = local_endpoint
+        dst_ep = globus_info.local_endpoint
         dst_path = os.path.join(os.getcwd(), name)
     else:
-        src_ep = local_endpoint
+        src_ep = globus_info.local_endpoint
         src_path = os.path.join(os.getcwd(), name)
-        dst_ep = remote_endpoint
+        dst_ep = globus_info.remote_endpoint
         dst_path = os.path.join(remote_path, name)
 
     subdir = os.path.basename(os.path.normpath(remote_path))
     subdir_label = re.sub("[^A-Za-z0-9_ -]", "", subdir)
     filename = name.split(".")[0]
     label = subdir_label + " " + filename
 
-    if not transfer_data:
-        transfer_data = TransferData(
-            transfer_client,
+    if not globus_info.transfer_data:
+        globus_info.transfer_data = TransferData(
+            globus_info.transfer_client,
             src_ep,
             dst_ep,
             label=label,
             verify_checksum=True,
             preserve_timestamp=True,
             fail_on_quota_errors=True,
         )
-    transfer_data.add_item(src_path, dst_path)
-    transfer_data["label"] = label
+    globus_info.transfer_data.add_item(src_path, dst_path)
+    globus_info.transfer_data["label"] = label
     try:
-        if task_id:
-            task = transfer_client.get_task(task_id)
+        if globus_info.task_id:
+            task = globus_info.transfer_client.get_task(globus_info.task_id)
             prev_task_status = task["status"]
             # one  of {ACTIVE, SUCCEEDED, FAILED, CANCELED, PENDING, INACTIVE}
             # NOTE: How we behave here depends upon whether we want to support mutliple active transfers.
@@ -231,78 +213,73 @@ def globus_transfer(  # noqa: C901
                 label = task["label"]
                 ts = ts_utc()
                 logger.info(
-                    "{}:Globus transfer {}, from {} to {}: {} succeeded".format(
-                        ts, task_id, src_ep, dst_ep, label
-                    )
+                    f"{ts}:Globus transfer {globus_info.task_id}, from {src_ep} to {dst_ep}: {label} succeeded"
                 )
             else:
                 logger.error(
-                    f"{ts_utc()}: Previous task_id {task_id} status = {prev_task_status}."
+                    f"{ts_utc()}: Previous task_id {globus_info.task_id} status = {prev_task_status}."
                 )
 
         # DEBUG: review accumulated items in TransferData
         logger.info(f"{ts_utc()}: TransferData: accumulated items:")
-        attribs = transfer_data.__dict__
+        attribs = globus_info.transfer_data.__dict__
         for item in attribs["data"]["DATA"]:
             if item["DATA_TYPE"] == "transfer_item":
-                global_variable_tarfiles_pushed += 1
+                globus_info.tarfiles_pushed += 1
                 print(
-                    f"   (routine)  PUSHING (#{global_variable_tarfiles_pushed}) STORED source item: {item['source_path']}",
+                    f"   (routine)  PUSHING (#{globus_info.tarfiles_pushed}) STORED source item: {item['source_path']}",
                     flush=True,
                 )
 
         # SUBMIT new transfer here
         logger.info(f"{ts_utc()}: DIVING: Submit Transfer for {transfer_data['label']}")
-        task = submit_transfer_with_checks(transfer_data)
-        task_id = task.get("task_id")
+        task = submit_transfer_with_checks(globus_info)
+        globus_info.task_id = task.get("task_id")
         # NOTE: This log message is misleading. If we have accumulated multiple tar files for transfer,
         # the "lable" given here refers only to the LAST tarfile in the TransferData list.
         logger.info(
-            f"{ts_utc()}: SURFACE Submit Transfer returned new task_id = {task_id} for label {transfer_data['label']}"
+            f"{ts_utc()}: SURFACE Submit Transfer returned new task_id = {globus_info.task_id} for label {transfer_data['label']}"
         )
 
         # Nullify the submitted transfer data structure so that a new one will be created on next call.
-        transfer_data = None
+        globus_info.transfer_data = None
     except TransferAPIError as e:
         if e.code == "NoCredException":
             logger.error(
-                "{}. Please go to https://app.globus.org/endpoints and activate the endpoint.".format(
-                    e.message
-                )
+                f"{e.message}. Please go to https://app.globus.org/endpoints and activate the endpoint."
             )
         else:
             logger.error(e)
         sys.exit(1)
     except Exception as e:
-        logger.error("Exception: {}".format(e))
+        logger.error(f"Exception: {e}")
         sys.exit(1)
 
     # test for blocking on new task_id
     task_status = "UNKNOWN"
     if not non_blocking:
         task_status = globus_block_wait(
-            task_id=task_id, wait_timeout=7200, polling_interval=10, max_retries=5
+            globus_info, wait_timeout=7200, polling_interval=10, max_retries=5
         )
     else:
-        logger.info(f"{ts_utc()}: NO BLOCKING (task_wait) for task_id {task_id}")
+        logger.info(f"{ts_utc()}: NO BLOCKING (task_wait) for task_id {globus_info.task_id}")
 
     if transfer_type == "put":
         return task_status
 
     if transfer_type == "get" and task_id:
-        globus_wait(task_id)
+        globus_wait(globus_info)
 
     return task_status
 
 
 def globus_block_wait(
-    task_id: str, wait_timeout: int, polling_interval: int, max_retries: int
+    globus_info: GlobusInfo, wait_timeout: int, polling_interval: int, max_retries: int
 ):
-    global transfer_client
 
     # poll every "polling_interval" seconds to speed up small transfers.  Report every 2 hours, stop waiting aftert 5*2 = 10 hours
     logger.info(
-        f"{ts_utc()}: BLOCKING START: invoking task_wait for task_id = {task_id}"
+        f"{ts_utc()}: BLOCKING START: invoking task_wait for task_id = {globus_info.task_id}"
     )
     task_status = "UNKNOWN"
     retry_count = 0
@@ -312,14 +289,14 @@ def globus_block_wait(
             logger.info(
                 f"{ts_utc()}: on task_wait try {retry_count+1} out of {max_retries}"
             )
-            transfer_client.task_wait(
-                task_id, timeout=wait_timeout, polling_interval=10
+            globus_info.transfer_client.task_wait(
+                globus_info.task_id, timeout=wait_timeout, polling_interval=10
             )
             logger.info(f"{ts_utc()}: done with wait")
         except Exception as e:
             logger.error(f"Unexpected Exception: {e}")
         else:
-            curr_task = transfer_client.get_task(task_id)
+            curr_task = globus_info.transfer_client.get_task(task_id)
             task_status = curr_task["status"]
             if task_status == "SUCCEEDED":
                 break
@@ -336,15 +313,17 @@ def globus_block_wait(
         task_status = "EXHAUSTED_TIMEOUT_RETRIES"
 
     logger.info(
-        f"{ts_utc()}: BLOCKING ENDS: task_id {task_id} returned from task_wait with status {task_status}"
+        f"{ts_utc()}: BLOCKING ENDS: task_id {globus_info.task_id} returned from task_wait with status {task_status}"
     )
 
     return task_status
 
 
-def globus_wait(task_id: str):
-    global transfer_client
-
+def globus_wait(globus_info: GlobusInfo, alternative_task_id=None):
+    if alternative_task_id:
+        task_id = alternative_task_id
+    else:
+        task_id = globus_info.task_id
     try:
         """
         A Globus transfer job (task) can be in one of the three states:
@@ -353,80 +332,69 @@ def globus_wait(task_id: str):
         with 20 second timeout limit. If the task is ACTIVE after time runs
         out 'task_wait' returns False, and True otherwise.
         """
-        while not transfer_client.task_wait(task_id, timeout=300, polling_interval=20):
+        while not globus_info.transfer_client.task_wait(task_id, timeout=300, polling_interval=20):
             pass
         """
         The Globus transfer job (task) has been finished (SUCCEEDED or FAILED).
         Check if the transfer SUCCEEDED or FAILED.
         """
-        task = transfer_client.get_task(task_id)
+        task = globus_info.transfer_client.get_task(task_id)
         if task["status"] == "SUCCEEDED":
             src_ep = task["source_endpoint_id"]
             dst_ep = task["destination_endpoint_id"]
             label = task["label"]
             logger.info(
-                "Globus transfer {}, from {} to {}: {} succeeded".format(
-                    task_id, src_ep, dst_ep, label
-                )
+                f"Globus transfer {task_id}, from {src_ep} to {dst_ep}: {label} succeeded"
             )
         else:
             logger.error("Transfer FAILED")
     except TransferAPIError as e:
         if e.code == "NoCredException":
             logger.error(
-                "{}. Please go to https://app.globus.org/endpoints and activate the endpoint.".format(
-                    e.message
-                )
+                f"{e.message}. Please go to https://app.globus.org/endpoints and activate the endpoint."
             )
         else:
             logger.error(e)
         sys.exit(1)
     except Exception as e:
-        logger.error("Exception: {}".format(e))
+        logger.error(f"Exception: {e}")
         sys.exit(1)
 
 
-def globus_finalize(non_blocking: bool = False):
-    global transfer_client
-    global transfer_data
-    global task_id
-    global global_variable_tarfiles_pushed
-
+def globus_finalize(globus_info: GlobusInfo, non_blocking: bool = False):
     last_task_id = None
 
-    if transfer_data:
+    if globus_info.transfer_data:
         # DEBUG: review accumulated items in TransferData
         logger.info(f"{ts_utc()}: FINAL TransferData: accumulated items:")
-        attribs = transfer_data.__dict__
+        attribs = globus_info.transfer_data.__dict__
         for item in attribs["data"]["DATA"]:
             if item["DATA_TYPE"] == "transfer_item":
-                global_variable_tarfiles_pushed += 1
+                globus_info.tarfiles_pushed += 1
                 print(
-                    f"    (finalize) PUSHING ({global_variable_tarfiles_pushed}) source item: {item['source_path']}",
+                    f"    (finalize) PUSHING ({globus_info.tarfiles_pushed}) source item: {item['source_path']}",
                     flush=True,
                 )
 
         # SUBMIT new transfer here
-        logger.info(f"{ts_utc()}: DIVING: Submit Transfer for {transfer_data['label']}")
+        logger.info(f"{ts_utc()}: DIVING: Submit Transfer for {globus_info.transfer_data['label']}")
         try:
-            last_task = submit_transfer_with_checks(transfer_data)
+            last_task = submit_transfer_with_checks(globus_info.transfer_data)
             last_task_id = last_task.get("task_id")
         except TransferAPIError as e:
             if e.code == "NoCredException":
                 logger.error(
-                    "{}. Please go to https://app.globus.org/endpoints and activate the endpoint.".format(
-                        e.message
-                    )
+                    f"{e.message}. Please go to https://app.globus.org/endpoints and activate the endpoint."
                 )
             else:
                 logger.error(e)
             sys.exit(1)
         except Exception as e:
-            logger.error("Exception: {}".format(e))
+            logger.error(f"Exception: {e}")
             sys.exit(1)
 
     if not non_blocking:
-        if task_id:
-            globus_wait(task_id)
+        if globus_info.task_id:
+            globus_wait(globus_info)
         if last_task_id:
-            globus_wait(last_task_id)
+            globus_wait(globus_info, last_task_id)