Fix circular import but tars still get deleted

Author: forsyth2

zstash/globus.py

@@ -1,7 +1,8 @@
 from __future__ import absolute_import, print_function
 
+import os
 import sys
-from typing import List, Optional
+from typing import Dict, List, Optional, Tuple
 
 from globus_sdk import TransferAPIError, TransferClient, TransferData
 from globus_sdk.response import GlobusHTTPResponse
@@ -16,7 +17,6 @@
     set_up_TransferData,
     submit_transfer_with_checks,
 )
-from .hpss import cleanup_completed_transfers
 from .settings import logger
 from .utils import ts_utc
 
@@ -26,6 +26,76 @@
 transfer_data: TransferData = None
 task_id = None
 archive_directory_listing: IterableTransferResponse = None
+# Global tracking for non-blocking transfers
+pending_cleanup: Dict[str, Tuple[str, str]] = {}  # task_id -> (file_path, timestamp)
+
+
+def cleanup_completed_transfers():
+    """
+    Check status of previously submitted non-blocking transfers and delete
+    files for any that have completed successfully.
+    """
+    from .globus import transfer_client
+
+    global pending_cleanup
+
+    if not pending_cleanup:
+        return
+
+    logger.debug(
+        f"{ts_utc()}: Checking {len(pending_cleanup)} pending transfers for cleanup"
+    )
+
+    completed_tasks = []
+
+    for task_id, (file_path, submit_time) in pending_cleanup.items():
+        try:
+            # Check transfer status
+            task = transfer_client.get_task(task_id)
+            task_status = task["status"]
+
+            if task_status == "SUCCEEDED":
+                # Transfer succeeded - delete the file
+                if os.path.exists(file_path):
+                    logger.info(
+                        f"{ts_utc()}: Deleting successfully transferred file {file_path}"
+                    )
+                    os.remove(file_path)
+                else:
+                    logger.warning(f"{ts_utc()}: File {file_path} already missing")
+                completed_tasks.append(task_id)
+
+            elif task_status in ["FAILED", "CANCELED"]:
+                # Transfer failed - keep the file but stop tracking
+                logger.warning(
+                    f"{ts_utc()}: Transfer {task_id} failed/canceled, keeping file {file_path}"
+                )
+                completed_tasks.append(task_id)
+
+            elif task_status in ["ACTIVE", "PENDING"]:
+                # Still in progress - keep tracking
+                logger.debug(f"{ts_utc()}: Transfer {task_id} still {task_status}")
+
+            else:
+                # Unknown status - log and remove from tracking
+                logger.warning(
+                    f"{ts_utc()}: Unknown transfer status {task_status} for {task_id}"
+                )
+                completed_tasks.append(task_id)
+
+        except Exception as e:
+            # Error checking status - remove from tracking to avoid infinite retries
+            logger.error(f"Error checking transfer status for {task_id}: {e}")
+            completed_tasks.append(task_id)
+
+    # Remove completed/failed transfers from tracking
+    for task_id in completed_tasks:
+        del pending_cleanup[task_id]
+
+    if completed_tasks:
+        logger.debug(
+            f"{ts_utc()}: Cleaned up {len(completed_tasks)} completed transfers"
+        )
 
 
 def globus_activate(hpss: str):
@@ -293,6 +363,7 @@ def globus_finalize(non_blocking: bool = False):
     global transfer_data
     global task_id
     global global_variable_tarfiles_pushed
+    global pending_cleanup
 
     last_task_id = None
 
@@ -343,7 +414,6 @@ def globus_finalize(non_blocking: bool = False):
         # 1. Wait for them (defeating the non-blocking purpose)
         # 2. Leave them for manual cleanup
         # 3. Log a warning
-        from .hpss import pending_cleanup
 
         if pending_cleanup:
             logger.warning(

zstash/hpss.py

@@ -2,86 +2,16 @@
 
 import os.path
 import subprocess
-from typing import Dict, List, Tuple
+from typing import List
 
 from six.moves.urllib.parse import urlparse
 
-from .globus import globus_transfer
+from .globus import cleanup_completed_transfers, globus_transfer
 from .settings import get_db_filename, logger
 from .utils import run_command, ts_utc
 
 prev_transfers: List[str] = list()
 curr_transfers: List[str] = list()
-# Global tracking for non-blocking transfers
-pending_cleanup: Dict[str, Tuple[str, str]] = {}  # task_id -> (file_path, timestamp)
-
-
-def cleanup_completed_transfers():
-    """
-    Check status of previously submitted non-blocking transfers and delete
-    files for any that have completed successfully.
-    """
-    from .globus import transfer_client
-
-    global pending_cleanup
-
-    if not pending_cleanup:
-        return
-
-    logger.debug(
-        f"{ts_utc()}: Checking {len(pending_cleanup)} pending transfers for cleanup"
-    )
-
-    completed_tasks = []
-
-    for task_id, (file_path, submit_time) in pending_cleanup.items():
-        try:
-            # Check transfer status
-            task = transfer_client.get_task(task_id)
-            task_status = task["status"]
-
-            if task_status == "SUCCEEDED":
-                # Transfer succeeded - delete the file
-                if os.path.exists(file_path):
-                    logger.info(
-                        f"{ts_utc()}: Deleting successfully transferred file {file_path}"
-                    )
-                    os.remove(file_path)
-                else:
-                    logger.warning(f"{ts_utc()}: File {file_path} already missing")
-                completed_tasks.append(task_id)
-
-            elif task_status in ["FAILED", "CANCELED"]:
-                # Transfer failed - keep the file but stop tracking
-                logger.warning(
-                    f"{ts_utc()}: Transfer {task_id} failed/canceled, keeping file {file_path}"
-                )
-                completed_tasks.append(task_id)
-
-            elif task_status in ["ACTIVE", "PENDING"]:
-                # Still in progress - keep tracking
-                logger.debug(f"{ts_utc()}: Transfer {task_id} still {task_status}")
-
-            else:
-                # Unknown status - log and remove from tracking
-                logger.warning(
-                    f"{ts_utc()}: Unknown transfer status {task_status} for {task_id}"
-                )
-                completed_tasks.append(task_id)
-
-        except Exception as e:
-            # Error checking status - remove from tracking to avoid infinite retries
-            logger.error(f"Error checking transfer status for {task_id}: {e}")
-            completed_tasks.append(task_id)
-
-    # Remove completed/failed transfers from tracking
-    for task_id in completed_tasks:
-        del pending_cleanup[task_id]
-
-    if completed_tasks:
-        logger.debug(
-            f"{ts_utc()}: Cleaned up {len(completed_tasks)} completed transfers"
-        )
 
 
 # C901 'hpss_transfer' is too complex (19)
@@ -96,7 +26,6 @@ def hpss_transfer(  # noqa: C901
 ):
     global prev_transfers
     global curr_transfers
-    global pending_cleanup
 
     # NEW: Clean up any completed transfers from previous calls
     if non_blocking and transfer_type == "put":
@@ -198,7 +127,7 @@ def hpss_transfer(  # noqa: C901
             # NEW: Track non-blocking transfers for later cleanup
             if non_blocking and transfer_type == "put" and not keep:
                 # Get the current task_id from globus module
-                from .globus import task_id
+                from .globus import pending_cleanup, task_id
 
                 if task_id:
                     pending_cleanup[task_id] = (file_path, ts_utc())