Claude's cleanup code

Author: forsyth2

zstash/globus.py

@@ -16,6 +16,7 @@
     set_up_TransferData,
     submit_transfer_with_checks,
 )
+from .hpss import cleanup_completed_transfers
 from .settings import logger
 from .utils import ts_utc
 
@@ -331,3 +332,24 @@ def globus_finalize(non_blocking: bool = False):
             globus_wait(task_id)
         if last_task_id:
             globus_wait(last_task_id)
+
+    # Final cleanup of any remaining pending transfers
+    if non_blocking:
+        # Wait a bit longer for any final transfers to complete
+        logger.info(f"{ts_utc()}: Final cleanup check for non-blocking transfers")
+        cleanup_completed_transfers()
+
+        # If there are still pending transfers, we could either:
+        # 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(
+                f"{ts_utc()}: {len(pending_cleanup)} transfers still pending - files not deleted"
+            )
+            for task_id, (file_path, _) in pending_cleanup.items():
+                logger.warning(
+                    f"{ts_utc()}: File {file_path} (task {task_id}) cleanup deferred"
+                )

zstash/hpss.py

@@ -2,7 +2,7 @@
 
 import os.path
 import subprocess
-from typing import List
+from typing import Dict, List, Tuple
 
 from six.moves.urllib.parse import urlparse
 
@@ -12,9 +12,80 @@
 
 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 hpss_transfer(
+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)
+def hpss_transfer(  # noqa: C901
     hpss: str,
     file_path: str,
     transfer_type: str,
@@ -25,6 +96,11 @@ def hpss_transfer(
 ):
     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":
+        cleanup_completed_transfers()
 
     logger.info(
         f"{ts_utc()}: in hpss_transfer, prev_transfers is starting as {prev_transfers}"
@@ -118,6 +194,17 @@ def hpss_transfer(
             # or perhaps transfer is hanging. We should decide whether to ignore it, or cancel it, but
             # we'd need the task_id to issue a cancellation.  Perhaps we should have globus_transfer
             # return a tuple (task_id, status).
+
+            # 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
+
+                if task_id:
+                    pending_cleanup[task_id] = (file_path, ts_utc())
+                    logger.debug(
+                        f"{ts_utc()}: Added {file_path} to pending cleanup for task {task_id}"
+                    )
         else:
             # Transfer file using `hsi`
             command: str = 'hsi -q "cd {}; {} {}"'.format(hpss, transfer_command, name)
@@ -130,19 +217,26 @@ def hpss_transfer(
 
         if transfer_type == "put":
             if not keep:
-                if (scheme != "globus") or (globus_status == "SUCCEEDED"):
-                    # Note: This is intended to fulfill the default removal of successfully-transfered
-                    # tar files when keep=False, irrespective of non-blocking status
+                if scheme != "globus":
+                    # HSI transfers - delete immediately as before
                     logger.debug(
-                        f"{ts_utc()}: deleting transfered files {prev_transfers}"
+                        f"{ts_utc()}: deleting transferred files {prev_transfers}"
                     )
                     for src_path in prev_transfers:
                         os.remove(src_path)
                     prev_transfers = curr_transfers
                     curr_transfers = list()
-                    logger.info(
-                        f"{ts_utc()}: prev_transfers has been set to {prev_transfers}"
-                    )
+                elif not non_blocking:
+                    # Blocking Globus transfers - delete immediately if succeeded
+                    if globus_status == "SUCCEEDED":
+                        logger.debug(
+                            f"{ts_utc()}: deleting transferred files {prev_transfers}"
+                        )
+                        for src_path in prev_transfers:
+                            os.remove(src_path)
+                        prev_transfers = curr_transfers
+                        curr_transfers = list()
+                # else: non-blocking Globus - cleanup handled by pending_cleanup mechanism
 
 
 def hpss_put(