diff --git a/zstash/globus.py b/zstash/globus.py index b99a6392..90eff420 100644 --- a/zstash/globus.py +++ b/zstash/globus.py @@ -81,165 +81,186 @@ def update_cumulative_tarfiles_pushed( ) -# C901 'globus_transfer' is too complex (20) -def globus_transfer( # noqa: C901 - transfer_manager: TransferManager, - remote_ep: str, - remote_path: str, - name: str, - transfer_type: str, - non_blocking: bool, -) -> TaskStatus: +# --------------------------------------------------------------------------- +# globus_transfer helpers +# --------------------------------------------------------------------------- - logger.info(f"{ts_utc()}: Entered globus_transfer() for name = {name}") - logger.debug(f"{ts_utc()}: non_blocking = {non_blocking}") - if (not transfer_manager.globus_config) or ( - not transfer_manager.globus_config.transfer_client + +def _ensure_globus_config(transfer_manager: TransferManager, remote_ep: str) -> None: + """Initialise transfer_manager.globus_config if it is not already set.""" + if ( + not transfer_manager.globus_config + or not transfer_manager.globus_config.transfer_client ): transfer_manager.globus_config = globus_activate("globus://" + remote_ep) - if (not transfer_manager.globus_config) or ( - not transfer_manager.globus_config.transfer_client + if ( + not transfer_manager.globus_config + or not transfer_manager.globus_config.transfer_client ): sys.exit(1) - if transfer_type == "get": - if not transfer_manager.globus_config.archive_directory_listing: - transfer_manager.globus_config.archive_directory_listing = ( - transfer_manager.globus_config.transfer_client.operation_ls( - transfer_manager.globus_config.remote_endpoint, remote_path - ) - ) - if not file_exists( - transfer_manager.globus_config.archive_directory_listing, name - ): - logger.error( - "Remote file globus://{}{}/{} does not exist".format( - remote_ep, remote_path, name - ) + +def _get_globus_config(transfer_manager: TransferManager) -> GlobusConfig: + """ + Return transfer_manager.globus_config, asserting it is not None. + Call this inside helpers that run after _ensure_globus_config has already + been called (i.e. anywhere inside globus_transfer and its callees). + """ + gc = transfer_manager.globus_config + assert gc is not None, "globus_config must be set before calling Globus helpers" + assert ( + gc.transfer_client is not None + ), "transfer_client must be set before calling Globus helpers" + return gc + + +def _verify_remote_file_exists( + transfer_manager: TransferManager, + remote_ep: str, + remote_path: str, + name: str, +) -> None: + """ + For 'get' transfers: populate the cached directory listing if needed, then + exit if the requested file is not present on the remote endpoint. + """ + gc = _get_globus_config(transfer_manager) + assert gc.transfer_client is not None + if not gc.archive_directory_listing: + gc.archive_directory_listing = gc.transfer_client.operation_ls( + gc.remote_endpoint, remote_path + ) + if not file_exists(gc.archive_directory_listing, name): + logger.error( + "Remote file globus://{}{}/{} does not exist".format( + remote_ep, remote_path, name ) - sys.exit(1) + ) + sys.exit(1) + - mrb: Optional[TransferBatch] = transfer_manager.get_most_recent_batch() - if not mrb: +def _add_file_to_current_batch( + transfer_manager: TransferManager, + remote_ep: str, + remote_path: str, + name: str, + transfer_type: str, +) -> None: + """ + Add *name* to the TransferData on the most recent batch, creating a new + TransferData object if the batch does not yet have one. + """ + gc = _get_globus_config(transfer_manager) + assert ( + gc.local_endpoint is not None + ), "local_endpoint must be set after globus_activate" + assert ( + gc.remote_endpoint is not None + ), "remote_endpoint must be set after globus_activate" + local_endpoint: str = gc.local_endpoint + remote_endpoint: str = gc.remote_endpoint + + mrb = transfer_manager.get_most_recent_batch() + if mrb is None: raise RuntimeError( - "The transfer manager should always have at least one batch by the time globus_transfer is called, however, the batch list is empty." + "No batch exists; hpss_transfer() must create one before calling globus_transfer()." ) - if transfer_manager.globus_config.local_endpoint: - local_endpoint: str = transfer_manager.globus_config.local_endpoint - else: - raise ValueError("Local endpoint ID is not set.") - if transfer_manager.globus_config.remote_endpoint: - remote_endpoint: str = transfer_manager.globus_config.remote_endpoint - else: - raise ValueError("Remote endpoint ID is not set.") - label: str = get_label(remote_path, name) - transfer_data: TransferData - if mrb.transfer_data: - # We already have a TransferData for this batch. - transfer_data = mrb.transfer_data - else: - # We need to create a new TransferData for this batch. - transfer_data = create_TransferData( + label = get_label(remote_path, name) + if mrb.transfer_data is None: + mrb.transfer_data = create_TransferData( transfer_type, local_endpoint, remote_endpoint, - transfer_manager.globus_config.transfer_client, + gc.transfer_client, label, ) + add_file_to_TransferData( transfer_type, local_endpoint, remote_endpoint, remote_path, name, - transfer_data, + mrb.transfer_data, label, ) - task: GlobusHTTPResponse - try: - if mrb.task_id: - # This the current transfer task associated with the most recent batch. - task = transfer_manager.globus_config.transfer_client.get_task(mrb.task_id) - # Update the most recent batch's task_status based on the current status from Globus API. - mrb.task_status = TaskStatus.convert_from_status_from_globus_sdk(task) - if mrb.task_status == TaskStatus.ACTIVE: - # The most recent transfer (mrb) is still active. - logger.info( - f"{ts_utc()}: Previous task_id {mrb.task_id} Still Active. Returning ACTIVE." - ) - if non_blocking: - # Globus allows up to 3 simulataneous transfers, - # but zstash is currently configured to only ever allow 1. - # If we're in this block, then we're already at 1 active transfer. - # We will therefore wait to submit a new transfer until it's done. - # So, we'll simply return and the next run of globus_transfer - # (i.e., on the next tar) will evaluate if the active transfer has finished. - return TaskStatus.ACTIVE - else: - # If we're in this block, then the blocking wait - # for the previous transfer to finish was unsuccessful. - # This is an unexpected state and so we raise an error. - error_str: str = ( - "task_status='ACTIVE', but in blocking mode, the previous transfer should have waited through globus_block_wait" - ) - logger.error(error_str) - raise RuntimeError(error_str) - elif mrb.task_status == TaskStatus.SUCCEEDED: - logger.info( - f"{ts_utc()}: Previous task_id {mrb.task_id} status = SUCCEEDED." - ) - src_ep = task["source_endpoint_id"] - dst_ep = task["destination_endpoint_id"] - label = task["label"] - ts = ts_utc() - logger.info( - f"{ts}:Globus transfer {mrb.task_id}, from {src_ep} to {dst_ep}: {label} succeeded" - ) - # The previous transfer succeeded. - # That means we can transfer the current batch now. - else: - # The previous transfer is in an unexpected state (i.e., "INACTIVE", "FAILED"). - # Either way, the previous transfer is effectively terminated, - # so we will proceed with the current transfer attempt. - # (I.e., we will not return yet). - # Note: any status we manually set - # (I.e., "UNKNOWN", "SUBMITTED", "EXHAUSTED_TIMEOUT_RETRIES") is NOT possible here, - # because we're using `task["status"]` from the globus_sdk TransferClient. - logger.warning( - f"{ts_utc()}: Previous task_id {mrb.task_id} status = {mrb.task_status}." - ) - update_cumulative_tarfiles_pushed(transfer_manager, transfer_data) +def _should_defer_submission( + transfer_manager: TransferManager, + non_blocking: bool, +) -> bool: + """ + Check the status of the previously submitted Globus task (if any). - logger.info(f"{ts_utc()}: DIVING: Submit Transfer for {transfer_data['label']}") - # Submit the current transfer_data - # ALWAYS submit. If we've gotten to this point, we're ready to submit. - task = submit_transfer_with_checks( - transfer_manager.globus_config.transfer_client, transfer_data - ) - task_id = task.get("task_id") - logger.info( - f"{ts_utc()}: SURFACE Submit Transfer returned new task_id = {task_id}, with last tarfile having label: {transfer_data['label']}" - ) + Returns True if submission should be deferred because the previous task is + still ACTIVE and we are in non-blocking mode. - # Update the current batch with the task info - # The batch was already created in hpss_transfer with files added to it - # We just need to mark it as submitted - if transfer_manager.batches: - # Update these two fields of the most recent batch - # (which is still available in this function as `mrb`). - transfer_manager.batches[-1].task_id = task_id - transfer_manager.batches[-1].task_status = TaskStatus.SUBMITTED + Raises RuntimeError if we are in blocking mode but the previous task is + somehow still ACTIVE (that would indicate a bug in the blocking-wait logic). + """ + mrb = transfer_manager.get_most_recent_batch() + if mrb is None or not mrb.task_id: + # No previous submission to worry about. + return False + + gc = _get_globus_config(transfer_manager) + assert gc.transfer_client is not None + task = gc.transfer_client.get_task(mrb.task_id) + mrb.task_status = TaskStatus.convert_from_status_from_globus_sdk(task) + + if mrb.task_status == TaskStatus.ACTIVE: + if non_blocking: + # The previous batch is still transferring; accumulate this file into + # the pending TransferData and come back to it later. + logger.info( + f"{ts_utc()}: Previous task_id {mrb.task_id} still ACTIVE; " + f"deferring submission (non-blocking mode)." + ) + return True else: - # This block should be impossible to reach. - # By now, we've ensured that `get_most_recent_batch()` returns a batch, - # and we haven't removed any batches since then, - # so there should always be at least one batch in `batches`. - error_str = "transfer_manager has no batches" + error_str = ( + "task_status='ACTIVE' in blocking mode — the previous transfer " + "should have completed via globus_block_wait before reaching here." + ) logger.error(error_str) raise RuntimeError(error_str) + + if mrb.task_status == TaskStatus.SUCCEEDED: + src_ep = task["source_endpoint_id"] + dst_ep = task["destination_endpoint_id"] + label = task["label"] + logger.info( + f"{ts_utc()}: Previous task_id {mrb.task_id} SUCCEEDED " + f"(from {src_ep} to {dst_ep}: {label}). Proceeding with next submission." + ) + else: + # INACTIVE, FAILED, or an unexpected status. The previous transfer is + # effectively terminal; log a warning and proceed with the new submission. + logger.warning( + f"{ts_utc()}: Previous task_id {mrb.task_id} has unexpected " + f"status={mrb.task_status}; proceeding anyway." + ) + + return False + + +def _submit_current_batch(transfer_manager: TransferManager) -> str: + """ + Submit the TransferData that has been accumulated on the most recent batch. + Returns the new Globus task_id. + """ + mrb = transfer_manager.get_most_recent_batch() + if mrb is None or mrb.transfer_data is None: + raise RuntimeError("No pending TransferData to submit.") + + update_cumulative_tarfiles_pushed(transfer_manager, mrb.transfer_data) + + logger.info(f"{ts_utc()}: Submitting Globus transfer: {mrb.transfer_data['label']}") + try: + gc = _get_globus_config(transfer_manager) + task = submit_transfer_with_checks(gc.transfer_client, mrb.transfer_data) except TransferAPIError as e: if e.code == "NoCredException": logger.error( @@ -254,34 +275,81 @@ def globus_transfer( # noqa: C901 logger.error("Exception: {}".format(e)) sys.exit(1) - task_status: TaskStatus = TaskStatus.UNKNOWN - if mrb.task_id: - if not non_blocking: - # If blocking, wait for the task to complete and get the final status, - # before we proceed with any more transfers. - mrb.task_status = globus_block_wait( - transfer_manager.globus_config.transfer_client, - task_id=mrb.task_id, - ) - task_status = mrb.task_status - else: - logger.info( - f"{ts_utc()}: NO BLOCKING (task_wait) for task_id {mrb.task_id}" - ) - else: - # This block should be impossible to reach. - # By now, we've set `transfer_manager.batches[-1].task_id = task_id` or else raised an error, so `mrb.task_id` should always be set. - error_str = "No task_id found for most recent batch after submission" - logger.error(f"{ts_utc()}: {error_str}") - raise RuntimeError(error_str) + task_id: str = task.get("task_id") + logger.info( + f"{ts_utc()}: Submitted transfer, new task_id={task_id} " + f"(label: {mrb.transfer_data['label']})" + ) - if transfer_type == "put": - return task_status + if not transfer_manager.batches: + raise RuntimeError("transfer_manager has no batches after submission.") + transfer_manager.batches[-1].task_id = task_id + transfer_manager.batches[-1].task_status = TaskStatus.SUBMITTED - if transfer_type == "get" and task_id: - globus_wait(transfer_manager.globus_config.transfer_client, task_id) + return task_id - return task_status + +# --------------------------------------------------------------------------- +# Public entry point +# --------------------------------------------------------------------------- + + +def globus_transfer( + transfer_manager: TransferManager, + remote_ep: str, + remote_path: str, + name: str, + transfer_type: str, + non_blocking: bool, +) -> TaskStatus: + """ + Transfer a single file to or from a Globus endpoint. + + For 'put' (non-blocking): the file is added to the pending TransferData. + Submission is deferred while the previous Globus task is still ACTIVE. + When the previous task finishes (or there is none), the accumulated batch + is submitted as one Globus task. + + For 'put' (blocking): the batch is submitted immediately after each file and + we wait for it to complete before returning. + + For 'get': the file is submitted immediately and we always wait for completion. + """ + logger.info(f"{ts_utc()}: globus_transfer() called for name={name!r}") + logger.debug(f"{ts_utc()}: non_blocking={non_blocking}") + + _ensure_globus_config(transfer_manager, remote_ep) + + if transfer_type == "get": + _verify_remote_file_exists(transfer_manager, remote_ep, remote_path, name) + + _add_file_to_current_batch( + transfer_manager, remote_ep, remote_path, name, transfer_type + ) + + if _should_defer_submission(transfer_manager, non_blocking): + return TaskStatus.ACTIVE + + task_id = _submit_current_batch(transfer_manager) + gc = _get_globus_config(transfer_manager) + + if transfer_type == "get": + # 'get' transfers always block until complete. + globus_wait(gc.transfer_client, task_id) + return TaskStatus.SUCCEEDED + + if not non_blocking: + # Blocking 'put': wait for this task before processing the next tar. + status = globus_block_wait(gc.transfer_client, task_id=task_id) + transfer_manager.batches[-1].task_status = status + return status + + return TaskStatus.SUBMITTED + + +# --------------------------------------------------------------------------- +# Wait helpers +# --------------------------------------------------------------------------- def globus_block_wait( @@ -290,77 +358,73 @@ def globus_block_wait( wait_timeout: int = 7200, # 7200/3600 = 2 hours max_retries: int = 5, ) -> TaskStatus: - # Poll every "polling_interval" seconds to speed up small transfers. - # Report every "wait_timeout" seconds, and stop waiting after "max_retries" reports. - # By default: report every 2 hours, stop waiting after 5*2 = 10 hours - logger.info( - f"{ts_utc()}: BLOCKING START: invoking task_wait for task_id = {task_id}" - ) + """ + Block until the given Globus task reaches a terminal state, or until + max_retries * wait_timeout seconds have elapsed. + + Polls every 10 seconds; reports progress every wait_timeout seconds. + Default limits: report every 2 hours, give up after 5 × 2 = 10 hours. + """ + logger.info(f"{ts_utc()}: Blocking wait started for task_id={task_id}") task_status: TaskStatus = TaskStatus.UNKNOWN retry_count: int = 0 while retry_count < max_retries: try: logger.info( - f"{ts_utc()}: on task_wait try {retry_count + 1} out of {max_retries}" + f"{ts_utc()}: task_wait attempt {retry_count + 1} of {max_retries}" ) - # Wait for the task to complete. This is what makes this function BLOCKING. - # From https://globus-sdk-python.readthedocs.io/en/stable/services/transfer.html#globus_sdk.TransferClient.task_wait: Wait until a Task is complete or fails, with a time limit. If the task is “ACTIVE” after time runs out, returns False. Otherwise returns True. - task_is_not_active: bool = transfer_client.task_wait( + # task_wait returns True when the task has reached a terminal state, + # False if it is still ACTIVE after `timeout` seconds. + task_is_terminal: bool = transfer_client.task_wait( task_id, timeout=wait_timeout, polling_interval=10 ) - if task_is_not_active: + if task_is_terminal: curr_task: GlobusHTTPResponse = transfer_client.get_task(task_id) task_status = TaskStatus.convert_from_status_from_globus_sdk(curr_task) if task_status == TaskStatus.SUCCEEDED: - break # Break out of the while-loop. The transfer already succeeded, so no need to retry. + break elif task_status == TaskStatus.FAILED: - error_str = f"{ts_utc()}: task_wait returned True, but task_status={task_status} for task_id {task_id}. No reason to keep retrying now." - logger.warning(error_str) - # We still need to break, because no matter how long we wait now, nothing will change with the transfer status. + logger.warning( + f"{ts_utc()}: task_id={task_id} FAILED; no point retrying." + ) break else: - error_str = f"{ts_utc()}: task_wait returned True, but task_status={task_status} for task_id {task_id}. Will retry waiting until max_retries is reached." - logger.warning(error_str) - # Don't break -- continue retries - logger.info(f"{ts_utc()}: done with wait") + logger.warning( + f"{ts_utc()}: task_id={task_id} reached unexpected terminal " + f"status={task_status}; will retry up to max_retries." + ) + logger.info(f"{ts_utc()}: task_wait returned (not yet terminal)") except Exception as e: logger.error(f"Unexpected Exception: {e}") finally: retry_count += 1 logger.info( - f"{ts_utc()}: BLOCKING retry_count = {retry_count} of {max_retries} of timeout {wait_timeout} seconds" + f"{ts_utc()}: blocking wait retry_count={retry_count}/{max_retries}, " + f"timeout={wait_timeout}s" ) if retry_count == max_retries: logger.info( - f"{ts_utc()}: BLOCKING EXHAUSTED {max_retries} of timeout {wait_timeout} seconds" + f"{ts_utc()}: Exhausted {max_retries} wait attempts of {wait_timeout}s each" ) task_status = TaskStatus.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 wait ended for task_id={task_id}, status={task_status}" ) - return task_status def globus_wait(transfer_client: TransferClient, task_id: str): + """ + Poll until the given Globus task reaches a terminal state, then log the outcome. + Exits the process on API errors. + """ try: - """ - A Globus transfer job (task) can be in one of the four states: - {ACTIVE, SUCCEEDED, FAILED, INACTIVE} - according to https://docs.globus.org/api/transfer/task/#task_fields. - The script every 20 seconds polls a - status of the transfer job (task) from the Globus Transfer service, - with 20 second timeout limit. If the task is ACTIVE after time runs - out 'task_wait' returns False, and True otherwise. - """ + # Poll every 20 seconds; re-poll indefinitely until terminal. while not 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: GlobusHTTPResponse = transfer_client.get_task(task_id) if TaskStatus.convert_from_status_from_globus_sdk(task) == TaskStatus.SUCCEEDED: src_ep = task["source_endpoint_id"] @@ -388,13 +452,18 @@ def globus_wait(transfer_client: TransferClient, task_id: str): sys.exit(1) +# --------------------------------------------------------------------------- +# globus_finalize helpers +# --------------------------------------------------------------------------- + + def _submit_pending_transfer_data( transfer_client: TransferClient, transfer_manager: TransferManager, ) -> Optional[str]: """ - If the most recent batch has unsubmitted TransferData, submit it and return task_id. - Otherwise return None. + If the most recent batch has unsubmitted TransferData, submit it and return + the new task_id. Returns None if there is nothing pending. """ transfer: Optional[TransferBatch] = transfer_manager.get_most_recent_batch() if not transfer or not transfer.transfer_data: @@ -403,13 +472,12 @@ def _submit_pending_transfer_data( update_cumulative_tarfiles_pushed(transfer_manager, transfer.transfer_data) logger.info( - f"{ts_utc()}: DIVING: Submit Transfer for {transfer.transfer_data['label']}" + f"{ts_utc()}: Submitting final pending transfer: {transfer.transfer_data['label']}" ) try: last_task = submit_transfer_with_checks(transfer_client, transfer.transfer_data) task_id = last_task.get("task_id") - # Best-effort: if this batch represents the submission, store the task_id. if task_id and transfer.is_globus and not transfer.task_id: transfer.task_id = task_id @@ -434,7 +502,9 @@ def _collect_globus_task_ids( transfer_manager: TransferManager, extra_task_id: Optional[str], keep: bool ) -> Tuple[List[str], Dict[str, TransferBatch]]: """ - Return (ordered unique task_ids, task_id->batch mapping for first occurrence). + Return (ordered unique task_ids, task_id -> first-seen batch mapping). + Skips batches that have already had their files deleted (local_paths_to_delete + is empty), unless keep=True in which case deletion is never tracked. """ task_ids: List[str] = [] seen: Set[str] = set() @@ -444,15 +514,13 @@ def _collect_globus_task_ids( if not keep: # NOTE: This is always true if `keep` is set, # since we never track files for deletion if `keep` is set. - already_deleted: bool = not batch.file_paths + already_deleted: bool = not batch.local_paths_to_delete if already_deleted: - # This batch has already been processed and files deleted, so we can skip it. continue if (not batch.is_globus) or (not batch.task_id): continue - # By this point, we know batch.task_id is not None tid: str = batch.task_id if tid in seen: continue @@ -461,8 +529,8 @@ def _collect_globus_task_ids( task_ids.append(tid) task_to_batch[tid] = batch - # Always include extra_task_id (e.g., just-submitted transfer), - # even if not yet reflected in batches. + # Always include extra_task_id (e.g., a just-submitted transfer that may not + # yet be reflected in the batches list). if extra_task_id and (extra_task_id not in seen): task_ids.append(extra_task_id) @@ -475,8 +543,8 @@ def _refresh_batch_status( task_to_batch: Dict[str, TransferBatch], ) -> Optional[TaskStatus]: """ - Fetch Globus task status and update corresponding batch.task_status if present. - Returns status, or None if fetch fails. + Fetch the current Globus status for task_id and update the corresponding + batch. Returns the status, or None if the fetch fails. """ try: task: GlobusHTTPResponse = transfer_client.get_task(task_id) @@ -487,7 +555,8 @@ def _refresh_batch_status( return status except Exception as e: logger.warning( - f"{ts_utc()}: Could not fetch status for task_id={task_id}; will wait anyway. ({e})" + f"{ts_utc()}: Could not fetch status for task_id={task_id}; " + f"will wait anyway. ({e})" ) return None @@ -498,8 +567,9 @@ def _wait_for_all_tasks( task_to_batch: Dict[str, TransferBatch], ) -> None: """ - For each task_id, refresh status; if not SUCCEEDED, block via globus_wait; - then refresh status again for deletion logic. + For each outstanding Globus task: refresh its status; if it has not already + succeeded, block until it reaches a terminal state; then refresh once more + so the batch status is accurate for the subsequent deletion step. """ for tid in task_ids: status = _refresh_batch_status(transfer_client, tid, task_to_batch) @@ -507,30 +577,36 @@ def _wait_for_all_tasks( logger.info(f"{ts_utc()}: task_id={tid} already SUCCEEDED; skipping wait") continue - logger.info( - f"{ts_utc()}: Waiting for transfer task_id={tid} to complete (status={status})" - ) + logger.info(f"{ts_utc()}: Waiting for transfer task_id={tid} (status={status})") globus_wait(transfer_client, tid) - # After wait returns, task is terminal; refresh once more. + # Refresh once more so deletion logic sees the final status. _refresh_batch_status(transfer_client, tid, task_to_batch) def _prune_empty_batches(transfer_manager: TransferManager) -> None: - """ - Remove batches which have no remaining files to manage. - - Note: we only prune batches whose file_paths is empty, regardless of Globus/HPSS. - That matches current semantics where file_paths=[] means "processed". - """ + """Remove batches that have no remaining files to manage.""" before = len(transfer_manager.batches) - transfer_manager.batches = [b for b in transfer_manager.batches if b.file_paths] + transfer_manager.batches = [ + b for b in transfer_manager.batches if b.local_paths_to_delete + ] after = len(transfer_manager.batches) if after != before: logger.debug(f"{ts_utc()}: Pruned {before - after} empty transfer batches") def globus_finalize(transfer_manager: TransferManager, keep: bool) -> None: + """ + Called once at the end of a create/update run to flush any remaining + pending transfers and wait for all outstanding Globus tasks to complete. + + Steps: + 1. Submit any TransferData that was accumulated but not yet sent. + 2. Collect the task_ids of all batches that still have files to delete. + 3. Wait for every outstanding task to reach a terminal state. + 4. Delete the local tar files for successfully completed transfers. + 5. Prune batches that no longer have any files to manage. + """ if transfer_manager.globus_config is None: logger.debug("No GlobusConfig object provided for finalization") return @@ -538,21 +614,25 @@ def globus_finalize(transfer_manager: TransferManager, keep: bool) -> None: logger.debug("GlobusConfig provided but transfer_client is None") return - # By this point, we know transfer_client is not None transfer_client: TransferClient = transfer_manager.globus_config.transfer_client + # 1. Submit any pending (unsubmitted) TransferData. last_task_id: Optional[str] = _submit_pending_transfer_data( transfer_client, transfer_manager ) + # 2. Collect all task_ids that still have associated local files. task_ids: List[str] task_to_batch: Dict[str, TransferBatch] task_ids, task_to_batch = _collect_globus_task_ids( transfer_manager, last_task_id, keep ) + # 3. Wait for every outstanding task. _wait_for_all_tasks(transfer_client, task_ids, task_to_batch) + # 4. Delete local tar files from succeeded transfers. transfer_manager.delete_successfully_transferred_files() + # 5. Remove empty (fully-processed) batches. _prune_empty_batches(transfer_manager) diff --git a/zstash/hpss.py b/zstash/hpss.py index 2c59145a..6bb0cbad 100644 --- a/zstash/hpss.py +++ b/zstash/hpss.py @@ -7,10 +7,144 @@ from six.moves.urllib.parse import urlparse from .globus import globus_transfer -from .settings import get_db_filename, logger -from .transfer_tracking import GlobusConfig, TaskStatus, TransferBatch, TransferManager +from .settings import logger +from .transfer_tracking import GlobusConfig, TransferBatch, TransferManager from .utils import run_command, ts_utc +# --------------------------------------------------------------------------- +# Internal helpers for each transfer variant +# --------------------------------------------------------------------------- + + +def _local_put( + file_path: str, + cache: str, + is_index: bool, +) -> None: + """ + Handle hpss='none' for a put: do nothing for the index DB or for tar files + when keep=True. For tar files (keep=False implied by caller), remove write + permissions so the local archive is read-only. + """ + if is_index: + # Nothing to do; the DB is always kept locally. + return + + # Remove write permissions from the tar file so the local-only archive + # behaves like an immutable store. + logger.info("put (local): removing write permissions from {}".format(file_path)) + + display_cmd: List[str] = "stat --format '%a' {}".format(file_path).split() + original_mode: bytes = subprocess.check_output(display_cmd).strip() + logger.info("{!r} original mode={!r}".format(file_path, original_mode)) + + subprocess.check_output("chmod ugo-w {}".format(file_path).split()) + + new_mode: bytes = subprocess.check_output(display_cmd).strip() + logger.info("{!r} new mode={!r}".format(file_path, new_mode)) + + +def _hsi_transfer(hpss: str, file_path: str, transfer_type: str) -> None: + """Transfer a single file to or from HPSS using the hsi command-line tool.""" + if transfer_type == "put": + transfer_word, transfer_command = "to", "put" + else: + transfer_word, transfer_command = "from", "get" + + _, name = os.path.split(file_path) + logger.info("Transferring file {} HPSS: {}".format(transfer_word, file_path)) + + command = 'hsi -q "cd {}; {} {}"'.format(hpss, transfer_command, name) + error_str = "Transferring file {} HPSS: {}".format(transfer_word, name) + run_command(command, error_str) + + +def _globus_put_or_get( + hpss: str, + file_path: str, + transfer_type: str, + keep: bool, + non_blocking: bool, + is_index: bool, + transfer_manager: TransferManager, +) -> None: + """ + Transfer a file using the Globus Transfer Service, then delete local tar + files for any batches whose transfer has succeeded. + """ + url = urlparse(hpss) + endpoint: str = url.netloc + url_path: str = url.path + path, name = os.path.split(file_path) + + if not keep and not is_index: + # Track this tar for deletion once its Globus transfer succeeds. + transfer_manager.batches[-1].local_paths_to_delete.append(file_path) + logger.debug( + f"{ts_utc()}: Tracking {file_path} for deletion after transfer; " + f"batch now has {len(transfer_manager.batches[-1].local_paths_to_delete)} file(s)" + ) + + # hsi requires us to be in the directory containing the file. + cwd = os.getcwd() + if path: + if transfer_type == "get" and not os.path.isdir(path): + os.makedirs(path) + os.chdir(path) + + logger.info(f"{ts_utc()}: globus_transfer() -> name={name!r}") + globus_transfer( + transfer_manager, endpoint, url_path, name, transfer_type, non_blocking + ) + logger.info(f"{ts_utc()}: globus_transfer() returned for name={name!r}") + + if path: + os.chdir(cwd) + + if transfer_type == "put" and not keep: + transfer_manager.delete_successfully_transferred_files() + + +def _hsi_put_or_get( + hpss: str, + file_path: str, + transfer_type: str, + keep: bool, + is_index: bool, + transfer_manager: TransferManager, +) -> None: + """ + Transfer a file using the hsi command-line tool, then delete local tar + files for any batches whose transfer has succeeded. + """ + path, _ = os.path.split(file_path) + + if not keep and not is_index: + transfer_manager.batches[-1].local_paths_to_delete.append(file_path) + logger.debug( + f"{ts_utc()}: Tracking {file_path} for deletion after transfer; " + f"batch now has {len(transfer_manager.batches[-1].local_paths_to_delete)} file(s)" + ) + + cwd = os.getcwd() + if path: + if transfer_type == "get" and not os.path.isdir(path): + os.makedirs(path) + os.chdir(path) + + _hsi_transfer(hpss, file_path, transfer_type) + + if path: + os.chdir(cwd) + + if transfer_type == "put" and not keep: + transfer_manager.delete_successfully_transferred_files() + + +# --------------------------------------------------------------------------- +# Public API +# --------------------------------------------------------------------------- + def hpss_transfer( hpss: str, @@ -28,122 +162,41 @@ def hpss_transfer( url = urlparse(hpss) scheme = url.scheme - # Create a new batch if needed (before we start adding files) + # Ensure there is an open batch to add files to. A new batch is needed + # when none exists yet, or when the last batch has already been submitted + # (task_id is set), meaning we are starting a new submission window. if not transfer_manager.batches or transfer_manager.batches[-1].task_id: - # Either no batches exist, or the last batch was already submitted new_batch = TransferBatch() new_batch.is_globus = scheme == "globus" transfer_manager.batches.append(new_batch) logger.debug( - f"{ts_utc()}: Created new TransferBatch, total batches: {len(transfer_manager.batches)}" + f"{ts_utc()}: Created new TransferBatch " + f"(total batches: {len(transfer_manager.batches)})" ) if hpss == "none": logger.info("{}: HPSS is unavailable".format(transfer_type)) - if transfer_type == "put" and file_path != get_db_filename(cache): - # We are adding a file (that is not the cache) to the local non-HPSS archive - logger.info( - "{}: Keeping tar files locally and removing write permissions".format( - transfer_type - ) - ) - # https://unix.stackexchange.com/questions/46915/get-the-chmod-numerical-value-for-a-file - display_mode_command: List[str] = "stat --format '%a' {}".format( - file_path - ).split() - display_mode_output: bytes = subprocess.check_output( - display_mode_command - ).strip() - logger.info( - "{!r} original mode={!r}".format(file_path, display_mode_output) - ) - # https://www.washington.edu/doit/technology-tips-chmod-overview - # Remove write-permission from user, group, and others, - # without changing read or execute permissions for any. - change_mode_command: List[str] = "chmod ugo-w {}".format(file_path).split() - # An error will be raised if this line fails. - subprocess.check_output(change_mode_command) - new_display_mode_output: bytes = subprocess.check_output( - display_mode_command - ).strip() - logger.info("{!r} new mode={!r}".format(file_path, new_display_mode_output)) - # else: no action needed - else: - transfer_word: str - transfer_command: str - if transfer_type == "put": - transfer_word = "to" - transfer_command = "put" - elif transfer_type == "get": - transfer_word = "from" - transfer_command = "get" - else: - raise ValueError("Invalid transfer_type={}".format(transfer_type)) - logger.info("Transferring file {} HPSS: {}".format(transfer_word, file_path)) - - endpoint: str = url.netloc - url_path = url.path - path: str - name: str - path, name = os.path.split(file_path) - - # Never track index.db for deletion, only the tar files - if (not keep) and (not is_index): - # Add this tar file to the current batch - transfer_manager.batches[-1].file_paths.append(file_path) - logger.debug( - f"{ts_utc()}: Added {file_path} to current batch, batch now has {len(transfer_manager.batches[-1].file_paths)} files" - ) - - # Need to be in local directory for `hsi` to work - cwd = os.getcwd() - if path != "": - if (transfer_type == "get") and (not os.path.isdir(path)): - # We are getting a file from HPSS. - # The directory the file is in doesn't exist locally. - # So, make the path locally - os.makedirs(path) - # Enter the path (directory) - # For `put`, this directory contains the file we want to transfer to HPSS. - # For `get`, this directory is where the file we get from HPSS will go. - os.chdir(path) - - globus_status: TaskStatus = TaskStatus.UNKNOWN - if scheme == "globus": - if not transfer_manager.globus_config: - transfer_manager.globus_config = GlobusConfig() - # Transfer file using the Globus Transfer Service - logger.info(f"{ts_utc()}: DIVING: hpss calls globus_transfer(name={name})") - task_status: TaskStatus = globus_transfer( - transfer_manager, endpoint, url_path, name, transfer_type, non_blocking - ) - logger.info( - f"{ts_utc()}: SURFACE: hpss globus_transfer(name={name}) returned task_status={task_status}" - ) - mrb: Optional[TransferBatch] = transfer_manager.get_most_recent_batch() - if mrb and mrb.task_status: - globus_status = mrb.task_status - logger.info( - f"{ts_utc()}: Most recent globus_transfer returned task_status={globus_status}" - ) - # NOTE: Here, the status could be "EXHAUSTED_TIMEOUT_RETRIES", meaning a very long 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). - else: - # Transfer file using `hsi` - command: str = 'hsi -q "cd {}; {} {}"'.format(hpss, transfer_command, name) - error_str: str = "Transferring file {} HPSS: {}".format(transfer_word, name) - run_command(command, error_str) - - # Return to original working directory - if path != "": - os.chdir(cwd) - if transfer_type == "put": - if not keep: - # We never delete if `--keep` is set. - transfer_manager.delete_successfully_transferred_files() + _local_put(file_path, cache, is_index) + # get with hpss='none' means the file is already local; nothing to do. + return + + if scheme == "globus": + if not transfer_manager.globus_config: + transfer_manager.globus_config = GlobusConfig() + _globus_put_or_get( + hpss, + file_path, + transfer_type, + keep, + non_blocking, + is_index, + transfer_manager, + ) + else: + _hsi_put_or_get( + hpss, file_path, transfer_type, keep, is_index, transfer_manager + ) def hpss_put( @@ -153,11 +206,9 @@ def hpss_put( transfer_manager: TransferManager, keep: bool = True, non_blocking: bool = False, - is_index=False, + is_index: bool = False, ): - """ - Put a file to the HPSS archive. - """ + """Put a file to the HPSS archive.""" hpss_transfer( hpss, file_path, @@ -176,13 +227,11 @@ def hpss_get( cache: str, transfer_manager: Optional[TransferManager] = None, ): - """ - Get a file from the HPSS archive. - """ + """Get a file from the HPSS archive.""" url = urlparse(hpss) if not transfer_manager: transfer_manager = TransferManager() - if (url.scheme == "globus") and not (transfer_manager.globus_config): + if url.scheme == "globus" and not transfer_manager.globus_config: transfer_manager.globus_config = GlobusConfig() hpss_transfer( hpss, file_path, "get", cache, False, transfer_manager=transfer_manager @@ -190,17 +239,11 @@ def hpss_get( def hpss_chgrp(hpss: str, group: str, recurse: bool = False): - """ - Change the group of the HPSS archive. - """ + """Change the group of the HPSS archive.""" if hpss == "none": logger.info("chgrp: HPSS is unavailable") else: - recurse_str: str - if recurse: - recurse_str = "-R " - else: - recurse_str = "" - command: str = "hsi chgrp {}{} {}".format(recurse_str, group, hpss) - error_str: str = "Changing group of HPSS archive {} to {}".format(hpss, group) + recurse_str = "-R " if recurse else "" + command = "hsi chgrp {}{} {}".format(recurse_str, group, hpss) + error_str = "Changing group of HPSS archive {} to {}".format(hpss, group) run_command(command, error_str) diff --git a/zstash/hpss_utils.py b/zstash/hpss_utils.py index 0956cd89..34a10694 100644 --- a/zstash/hpss_utils.py +++ b/zstash/hpss_utils.py @@ -16,10 +16,17 @@ from .transfer_tracking import TransferManager from .utils import create_tars_table, tars_table_exists, ts_utc +# --------------------------------------------------------------------------- +# DevOptions +# --------------------------------------------------------------------------- + -# This class holds parameters for developer options. -# I.e., these parameters should only ever be activated by developers during debugging and/or testing. class DevOptions(object): + """ + Parameters that activate deliberate misbehaviour for testing/debugging. + None of these should ever be set in production. + """ + def __init__( self, error_on_duplicate_tar: bool, @@ -55,18 +62,211 @@ def simulate_row_existing( ) +# --------------------------------------------------------------------------- +# TarWrapper +# --------------------------------------------------------------------------- + + class TarWrapper(object): + """ + Wraps a single tar archive being built during a create or update run. + + Lifecycle (driven by construct_tars): + 1. __init__: open a new tar file in the cache directory. + 2. process_file (× N): add each source file to the tar. + 3. finalize: close the tar, upload it to HPSS/Globus, and record it in + the database. Broken into three named steps below so that each + concern is easy to find. + """ + def __init__(self, tar_num: int, cache: str, do_hash: bool, follow_symlinks: bool): - # Create a hex value at least 6 digits long + # Derive the tar filename from the sequential tar number (hex, min 6 digits). tname: str = "{0:0{1}x}".format(tar_num, 6) - # Create the tar file name by adding ".tar" self.tfname: str = f"{tname}.tar" logger.info(f"{ts_utc()}: Creating new tar archive {self.tfname}") - # Open that tar file in the cache self.tarFileObject = HashIO(os.path.join(cache, self.tfname), "wb", do_hash) # FIXME: error: Argument "fileobj" to "open" has incompatible type "HashIO"; expected "Optional[IO[bytes]]" self.tar = tarfile.open(mode="w", fileobj=self.tarFileObject, dereference=follow_symlinks) # type: ignore + # ------------------------------------------------------------------ + # Step 1 of finalize: close the open tar file and capture its hash + # ------------------------------------------------------------------ + + def _close_tar(self) -> Tuple[int, Optional[str]]: + """ + Close the tar archive and return (tar_size_bytes, tar_md5). + Must be called before _upload_tar or _record_tar_in_database. + """ + logger.debug(f"{ts_utc()}: Closing tar archive {self.tfname}") + self.tar.close() + tar_size: int = self.tarFileObject.tell() + tar_md5: Optional[str] = self.tarFileObject.md5() + self.tarFileObject.close() + logger.info(f"{ts_utc()}: Closed archive {self.tfname} ({tar_size} bytes)") + return tar_size, tar_md5 + + # ------------------------------------------------------------------ + # Step 2 of finalize: upload the tar to HPSS / Globus + # ------------------------------------------------------------------ + + def _upload_tar( + self, + cache: str, + keep: bool, + non_blocking: bool, + transfer_manager: TransferManager, + ) -> None: + """Submit the closed tar file to the transfer system.""" + if config.hpss is None: + raise TypeError("Invalid config.hpss={}".format(config.hpss)) + hpss: str = config.hpss + + logger.debug(f"Cache contents before upload: {os.listdir(cache)}") + logger.info( + f"{ts_utc()}: Uploading {self.tfname} " + f"[keep={keep}, non_blocking={non_blocking}]" + ) + hpss_put( + hpss, + os.path.join(cache, self.tfname), + cache, + transfer_manager, + keep, + non_blocking, + is_index=False, + ) + logger.info(f"{ts_utc()}: Upload dispatched for {self.tfname}") + + # ------------------------------------------------------------------ + # Step 3 of finalize: record the tar and its files in the database + # ------------------------------------------------------------------ + + def _record_tar_in_database( + self, + tar_size: int, + tar_md5: Optional[str], + skip_tars_table: bool, + cur: sqlite3.Cursor, + con: sqlite3.Connection, + dev_options: DevOptions, + archived: List[TupleFilesRowNoId], + ) -> None: + """ + Insert the tar itself into the 'tars' table (unless skip_tars_table is + set) and insert every file it contains into the 'files' table. + """ + if not skip_tars_table: + self._insert_tar_row(tar_size, tar_md5, cur, con, dev_options) + + # Record each individual file that was archived into this tar. + cur.executemany("insert into files values (NULL,?,?,?,?,?,?)", archived) + con.commit() + + def _insert_tar_row( + self, + tar_size: int, + tar_md5: Optional[str], + cur: sqlite3.Cursor, + con: sqlite3.Connection, + dev_options: DevOptions, + ) -> None: + """ + Insert (or update, depending on DevOptions) a row in the 'tars' table + for this tar archive. Handles duplicate-tar detection and the various + developer-only corruption-simulation modes. + """ + tar_tuple: TupleTarsRowNoId = (self.tfname, tar_size, tar_md5) + logger.info("tar name={}, tar size={}, tar md5={}".format(*tar_tuple)) + + if not tars_table_exists(cur): + create_tars_table(cur, con) + + # Developer-only: optionally insert a duplicate row before the main logic. + dev_options.simulate_row_existing( + self.tfname, cur, tar_tuple, tar_size, tar_md5 + ) + + cur.execute("SELECT COUNT(*) FROM tars WHERE name = ?", (self.tfname,)) + tar_count: int = cur.fetchone()[0] + + if tar_count != 0: + self._handle_duplicate_tar( + tar_size, tar_md5, tar_tuple, cur, con, dev_options + ) + elif dev_options.force_database_corruption == "simulate_no_correct_size": + # Tested by database_corruption.bash Case 6 + logger.info( + f"TESTING/DEBUGGING ONLY: Simulating no correct size for {self.tfname}." + ) + cur.execute( + "INSERT INTO tars VALUES (NULL,?,?,?)", + (self.tfname, tar_size + 1000, tar_md5), + ) + cur.execute( + "INSERT INTO tars VALUES (NULL,?,?,?)", + (self.tfname, tar_size + 2000, tar_md5), + ) + elif ( + dev_options.force_database_corruption == "simulate_bad_size_for_most_recent" + ): + # Tested by database_corruption.bash Case 8 + logger.info( + f"TESTING/DEBUGGING ONLY: Simulating bad size for most recent " + f"entry for {self.tfname}." + ) + cur.execute( + "INSERT INTO tars VALUES (NULL,?,?,?)", + (self.tfname, tar_size, tar_md5), + ) + cur.execute( + "INSERT INTO tars VALUES (NULL,?,?,?)", + (self.tfname, tar_size + 2000, tar_md5), + ) + else: + # Tested by database_corruption.bash Cases 1, 2 — normal path. + logger.info(f"Adding {self.tfname} to the database.") + cur.execute("INSERT INTO tars VALUES (NULL,?,?,?)", tar_tuple) + + con.commit() + + def _handle_duplicate_tar( + self, + tar_size: int, + tar_md5: Optional[str], + tar_tuple: TupleTarsRowNoId, + cur: sqlite3.Cursor, + con: sqlite3.Connection, + dev_options: DevOptions, + ) -> None: + """ + React to finding that this tar's name is already present in the database. + Behaviour is controlled by DevOptions flags. + """ + error_str = ( + f"Database corruption detected! {self.tfname} is already in the database." + ) + if dev_options.error_on_duplicate_tar: + # Tested by database_corruption.bash Case 3 + logger.error(error_str) + raise RuntimeError(error_str) + elif dev_options.overwrite_duplicate_tars: + # Tested by database_corruption.bash Case 4 + logger.warning(error_str) + logger.warning(f"Overwriting existing tar entry for {self.tfname}.") + cur.execute( + "UPDATE tars SET size = ?, md5 = ? WHERE name = ?", + (tar_size, tar_md5, self.tfname), + ) + else: + # Tested by database_corruption.bash Cases 5, 7 + logger.warning(error_str) + logger.warning(f"Adding a new entry for {self.tfname}.") + cur.execute("INSERT INTO tars VALUES (NULL,?,?,?)", tar_tuple) + + # ------------------------------------------------------------------ + # Public entry point: add a single source file to the open tar + # ------------------------------------------------------------------ + def process_file( self, current_file: str, @@ -74,6 +274,12 @@ def process_file( archived: List[TupleFilesRowNoId], failures: List[str], ) -> int: + """ + Add *current_file* to the tar archive. + + Appends a row to *archived* on success, or to *failures* on error. + Returns the current cumulative tar size (0 on failure). + """ logger.info(f"Archiving {current_file}") tar_size: int = 0 try: @@ -84,26 +290,19 @@ def process_file( offset, size, mtime, md5 = add_file_to_tar_archive( self.tar, current_file, tar_info ) - t: TupleFilesRowNoId = ( - current_file, - size, - mtime, - md5, - self.tfname, - offset, - ) - archived.append(t) - # Increase tar_size by the size of the current file. - # Use `tell()` to also include the tar's metadata in the size. + archived.append((current_file, size, mtime, md5, self.tfname, offset)) tar_size = self.tarFileObject.tell() except Exception: - # Catch all exceptions here. traceback.print_exc() logger.error(f"Archiving {current_file}") failures.append(current_file) return tar_size - def process_tar( + # ------------------------------------------------------------------ + # Public entry point: close, upload, and record this tar + # ------------------------------------------------------------------ + + def finalize( self, cache: str, keep: bool, @@ -114,134 +313,51 @@ def process_tar( con: sqlite3.Connection, dev_options: DevOptions, archived: List[TupleFilesRowNoId], - ): - # 1. Close the tar #################################################### - logger.debug(f"{ts_utc()}: Closing tar archive {self.tfname}") - self.tar.close() - - tar_size = self.tarFileObject.tell() - tar_md5: Optional[str] = self.tarFileObject.md5() - self.tarFileObject.close() - logger.info(f"{ts_utc()}: (process_tar): Completed archive file {self.tfname}") - - # 2. Submit the tar to the transfer manager's batch transfer system ### - if config.hpss is not None: - hpss: str = config.hpss - else: - raise TypeError("Invalid config.hpss={}".format(config.hpss)) - - logger.debug(f"Contents of the cache prior to `hpss_put`: {os.listdir(cache)}") - - logger.info( - f"{ts_utc()}: DIVING: (process_tar): Calling hpss_put to dispatch archive file {self.tfname} [keep, non_blocking] = [{keep}, {non_blocking}]" + ) -> None: + """ + Complete processing of this tar archive: + 1. Close the tar file and capture its size and MD5. + 2. Upload it to HPSS / Globus. + 3. Record the tar and its constituent files in the database. + """ + tar_size, tar_md5 = self._close_tar() + self._upload_tar(cache, keep, non_blocking, transfer_manager) + self._record_tar_in_database( + tar_size, tar_md5, skip_tars_table, cur, con, dev_options, archived ) - # Actually submit the tar file - hpss_put( - hpss, - os.path.join(cache, self.tfname), + + # Keep the old name as an alias so any external callers are not broken. + # Prefer finalize() in new code. + def process_tar( + self, + cache: str, + keep: bool, + non_blocking: bool, + transfer_manager: TransferManager, + skip_tars_table: bool, + cur: sqlite3.Cursor, + con: sqlite3.Connection, + dev_options: DevOptions, + archived: List[TupleFilesRowNoId], + ) -> None: + self.finalize( cache, - transfer_manager, keep, non_blocking, - is_index=False, - ) - logger.info( - f"{ts_utc()}: SURFACE (process_tar): Called hpss_put to dispatch archive file {self.tfname}" + transfer_manager, + skip_tars_table, + cur, + con, + dev_options, + archived, ) - # 3. Add the tar itself to the tars table ############################# - if not skip_tars_table: - tar_tuple: TupleTarsRowNoId = (self.tfname, tar_size, tar_md5) - logger.info("tar name={}, tar size={}, tar md5={}".format(*tar_tuple)) - if not tars_table_exists(cur): - # Need to create tars table - create_tars_table(cur, con) - - # For developers only! For debugging/testing purposes only! - dev_options.simulate_row_existing( - self.tfname, cur, tar_tuple, tar_size, tar_md5 - ) - # We're done adding files to the tar. - # And we've transferred it to HPSS. - # Now we can insert the tar into the database. - cur.execute("SELECT COUNT(*) FROM tars WHERE name = ?", (self.tfname,)) - tar_count: int = cur.fetchone()[0] - if tar_count != 0: - error_str: str = ( - f"Database corruption detected! {self.tfname} is already in the database." - ) - if dev_options.error_on_duplicate_tar: - # Tested by database_corruption.bash Case 3 - # Exists - error out - logger.error(error_str) - raise RuntimeError(error_str) - elif dev_options.overwrite_duplicate_tars: - # Tested by database_corruption.bash Case 4 - # Exists - update with new size and md5 - logger.warning(error_str) - logger.warning(f"Updating existing tar {self.tfname} to proceed.") - cur.execute( - "UPDATE tars SET size = ?, md5 = ? WHERE name = ?", - (tar_size, tar_md5, self.tfname), - ) - else: - # Tested by database_corruption.bash Cases 5,7 - # Proceed as if we're in the typical case -- insert new - logger.warning(error_str) - logger.warning(f"Adding a new entry for {self.tfname}.") - cur.execute("INSERT INTO tars VALUES (NULL,?,?,?)", tar_tuple) - elif dev_options.force_database_corruption == "simulate_no_correct_size": - # Tested by database_corruption.bash Case 6 - # For developers only! For debugging purposes only! - # Add this tar twice, with different sizes. - logger.info( - f"TESTING/DEBUGGING ONLY: Simulating no correct size for {self.tfname}." - ) - cur.execute( - "INSERT INTO tars VALUES (NULL,?,?,?)", - (self.tfname, tar_size + 1000, tar_md5), - ) - cur.execute( - "INSERT INTO tars VALUES (NULL,?,?,?)", - (self.tfname, tar_size + 2000, tar_md5), - ) - elif ( - dev_options.force_database_corruption - == "simulate_bad_size_for_most_recent" - ): - # Tested by database_corruption.bash Case 8 - # For developers only! For debugging purposes only! - # Add this tar twice, second time with bad size. - logger.info( - f"TESTING/DEBUGGING ONLY: Simulating bad size for most recent entry for {self.tfname}." - ) - cur.execute( - "INSERT INTO tars VALUES (NULL,?,?,?)", - (self.tfname, tar_size, tar_md5), - ) - cur.execute( - "INSERT INTO tars VALUES (NULL,?,?,?)", - (self.tfname, tar_size + 2000, tar_md5), - ) - else: - # Tested by database_corruption.bash Cases 1,2 - # Typical case - # Doesn't exist - insert new - logger.info(f"Adding {self.tfname} to the database.") - cur.execute("INSERT INTO tars VALUES (NULL,?,?,?)", tar_tuple) - - con.commit() - - # 4. Add the files included in this tar to the files table ############ - # Update database with the individual files that have been archived - # Add a row to the "files" table, - # the last 6 columns matching the values of `archived` - cur.executemany("insert into files values (NULL,?,?,?,?,?,?)", archived) - con.commit() +# --------------------------------------------------------------------------- +# HashIO — minimal file-like object that tracks position and MD5 as we write +# --------------------------------------------------------------------------- -# Minimum output file object class HashIO(object): def __init__(self, name: str, mode: str, do_hash: bool): self.f = open(name, mode) @@ -258,16 +374,8 @@ def tell(self) -> int: def write(self, s): """ - This is called implicitly. - In TarWrapper.__init__: - - ``` - self.tarFileObject = HashIO(os.path.join(cache, self.tfname), "wb", do_hash) - self.tar = tarfile.open(mode="w", fileobj=self.tarFileObject, dereference=follow_symlinks) - ``` - - tarfile.open requires that the fileobj argument has a write() method. - It calls that method to write data to the tar file. + Called implicitly by tarfile as it streams data into the archive. + (tarfile.open requires the fileobj to have a write() method.) """ self.f.write(s) if self.hash: @@ -275,59 +383,61 @@ def write(self, s): self.position += len(s) def md5(self) -> Optional[str]: - md5: Optional[str] if self.hash: - md5 = self.hash.hexdigest() - else: - md5 = None - return md5 + return self.hash.hexdigest() + return None def close(self): if self.closed: return - self.f.close() self.closed = True +# --------------------------------------------------------------------------- +# Tar-size estimation +# --------------------------------------------------------------------------- + + def estimate_tar_entry_size(file_size: int) -> int: """ - Estimate how much space a file of a given size would take in the tar archive, - including metadata and padding. + Estimate how much space a file of the given size will occupy in the tar + archive, including the per-file header and block-alignment padding. """ TAR_BLOCK_SIZE = 512 - TAR_HEADER_SIZE = 512 # per file header - # This formula computes: ceil(file_size / TAR_BLOCK_SIZE) - # But faster and avoiding floats. + TAR_HEADER_SIZE = 512 data_blocks = (file_size + TAR_BLOCK_SIZE - 1) // TAR_BLOCK_SIZE return TAR_HEADER_SIZE + (data_blocks * TAR_BLOCK_SIZE) -# Add file to tar archive while computing its hash -# Return file offset (in tar archive), size and md5 hash +# --------------------------------------------------------------------------- +# Adding a single file to a tar archive +# --------------------------------------------------------------------------- + + def add_file_to_tar_archive( tar: tarfile.TarFile, file_name: str, tar_info: tarfile.TarInfo ) -> Tuple[int, int, datetime, Optional[str]]: - offset = tar.offset + """ + Add *file_name* to *tar* while computing its MD5 hash. + Returns (offset_in_tar, file_size, mtime, md5). + md5 is None for directories and symlinks. + """ + offset = tar.offset md5: Optional[str] = None - # For files/hardlinks if tar_info.isfile() or tar_info.islnk(): if tar_info.size > 0: - # Non-empty files: stream with hash computation hash_md5 = hashlib.md5() with open(file_name, "rb") as f: - wrapper = HashingFileWrapper(f, hash_md5) - tar.addfile(tar_info, wrapper) + tar.addfile(tar_info, HashingFileWrapper(f, hash_md5)) md5 = hash_md5.hexdigest() else: - # Empty files: just add to tar, compute hash of empty data tar.addfile(tar_info) - md5 = hashlib.md5(b"").hexdigest() # MD5 of empty bytes + md5 = hashlib.md5(b"").hexdigest() else: - # Directories, symlinks, etc. - # md5 will be None in these cases. + # Directories, symlinks, etc. — no file data to hash. tar.addfile(tar_info) size = tar_info.size @@ -335,6 +445,11 @@ def add_file_to_tar_archive( return offset, size, mtime, md5 +# --------------------------------------------------------------------------- +# Main loop: pack files into tars and dispatch each tar for upload +# --------------------------------------------------------------------------- + + def construct_tars( cur: sqlite3.Cursor, con: sqlite3.Connection, @@ -348,63 +463,54 @@ def construct_tars( skip_tars_table: bool = False, non_blocking: bool = False, ) -> List[str]: + """ + Pack *file_stats* into a sequence of tar archives (each no larger than + config.maxsize), upload each archive, and record everything in the database. + + *itar* is the index of the last existing tar (-1 for a fresh create, or the + highest existing tar number for an update). The first new tar will be + itar+1. + Returns a list of file paths that could not be archived. + """ failures: List[str] = [] files: List[str] = list(file_stats.keys()) nfiles: int = len(files) - if config.maxsize is not None: - max_size: int = config.maxsize - else: + if config.maxsize is None: raise TypeError(f"Invalid config.maxsize={config.maxsize}") + max_size: int = config.maxsize - operation: str - if itar == -1: - operation = "creation" - else: - operation = "update" + operation = "creation" if itar == -1 else "update" i_file: int = 0 while i_file < nfiles: - # Each iteration of this loop constructs one tar - - # `create` passes in itar=-1, so the first tar will be 000000.tar - # `update` passes in itar=max existing tar number, so the first tar will be max+1 + # Each iteration of this outer loop produces exactly one tar archive. itar += 1 cumulative_tar_size: int = 0 archived: List[TupleFilesRowNoId] = [] - # Open a new tar - # Note: if we're not skipping the tars table, then we DO want to calculate the hash of the tars. - # That is, we DO want to add the tar to the tars table in the database. - # That means we need to calculate the hash of the tar file as well. - # - # We ALWAYS want to calculate the hashes of the individual files, regardless of skip_tars_table, - # because we need to add those to the files table. tar_wrapper = TarWrapper( tar_num=itar, cache=cache, + # We need the tar's hash iff we are writing it to the tars table. do_hash=not skip_tars_table, follow_symlinks=follow_symlinks, ) - # Add files to the tar until we reach the max size + # Add files until this tar would exceed max_size. while i_file < nfiles: current_file: str = files[i_file] - current_file_size: int current_file_size, _ = file_stats[current_file] estimated_entry_size: int = estimate_tar_entry_size(current_file_size) - if (cumulative_tar_size != 0) and ( + + if cumulative_tar_size != 0 and ( cumulative_tar_size + estimated_entry_size > max_size ): - # Over the size limit: time to close and transfer this tar archive. - # Done adding files to this particular tar. - # Break out of the inner while-loop + # This file would push us over the limit; start a new tar. break - # If we make it this far, - # we know we can add the current file without going over the max size. - # (Either that, or the tar is currently empty, - # in which case we add the file even if it's over the max size.) + + # Attempt to get the tarinfo for the current file. try: tar_info = tar_wrapper.tar.gettarinfo(current_file) if tar_info.islnk(): @@ -417,17 +523,16 @@ def construct_tars( ) else: raise - new_cumulative_tar_size = tar_wrapper.process_file( + + new_size = tar_wrapper.process_file( current_file, tar_info, archived, failures ) - if new_cumulative_tar_size != 0: - # Update the cumulative tar size with the new tar size returned by process_file. - cumulative_tar_size = new_cumulative_tar_size - # Else: process_file failed, so we should keep the original cumulative_tar_size + if new_size != 0: + cumulative_tar_size = new_size i_file += 1 - # Close the tar, submit it to the batch transfer system, and update the database with the archived files (and optionally the tar as well, depending on skip_tars_table) - tar_wrapper.process_tar( + # Close this tar, upload it, and record it in the database. + tar_wrapper.finalize( cache, keep, non_blocking, @@ -442,7 +547,11 @@ def construct_tars( return failures -# Create a wrapper that computes hash while data passes through +# --------------------------------------------------------------------------- +# HashingFileWrapper — streams data through a hasher as tarfile reads it +# --------------------------------------------------------------------------- + + class HashingFileWrapper: def __init__(self, fileobj, hasher): self.fileobj = fileobj diff --git a/zstash/transfer_tracking.py b/zstash/transfer_tracking.py index 1cac18d5..fc6f4175 100644 --- a/zstash/transfer_tracking.py +++ b/zstash/transfer_tracking.py @@ -58,50 +58,65 @@ def __str__(self) -> str: class TransferBatch: - """Represents one batch of files being transferred""" + """ + Represents one batch of files submitted (or to be submitted) as a single + Globus transfer task. + + Lifecycle: + 1. Created in hpss_transfer() when a new batch is needed. + 2. Files are added to local_paths_to_delete (and to transfer_data) as + hpss_put() is called for each tar. + 3. The batch is submitted to Globus (task_id is set). + 4. Once the Globus task succeeds, local_paths_to_delete are removed from disk. + """ def __init__(self): - self.file_paths: List[str] = [] + # Local tar files in this batch; deleted once the Globus transfer succeeds. + self.local_paths_to_delete: List[str] = [] self.task_id: Optional[str] = None self.task_status: Optional[TaskStatus] = None self.is_globus: bool = False self.transfer_data: Optional[TransferData] = None # Only for Globus - def delete_files(self): - for src_path in self.file_paths: + def delete_local_files(self): + """Delete all local tar files tracked by this batch.""" + for path in self.local_paths_to_delete: try: - os.remove(src_path) + os.remove(path) except FileNotFoundError: - logger.warning(f"File already deleted: {src_path}") + logger.warning(f"File already deleted: {path}") class TransferManager: def __init__(self): - # All transfer batches (Globus or HPSS) + # All transfer batches (Globus or HPSS), in submission order. self.batches: List[TransferBatch] = [] self.cumulative_tarfiles_pushed: int = 0 - # Connection state (Globus-specific, None if not using Globus) + # Globus connection state; None when not using Globus. self.globus_config: Optional[GlobusConfig] = None def get_most_recent_batch(self) -> Optional[TransferBatch]: - """Get the last batch added to the manager, or None if no batches exist""" + """Return the last batch, or None if no batches exist.""" return self.batches[-1] if self.batches else None def delete_successfully_transferred_files(self): - """Check transfer status and delete files from successful transfers""" + """ + Delete local tar files for every batch whose Globus transfer has + succeeded (or for every non-Globus batch, which transfers synchronously). + Batches whose files have already been deleted are skipped. + """ logger.info( f"{ts_utc()}: Checking for successfully transferred files to delete" ) - # Clean up empty batches first - self.batches = [batch for batch in self.batches if batch.file_paths] - # Now delete files for successful transfers for batch in self.batches: - if (not batch.is_globus) or (batch.task_status == TaskStatus.SUCCEEDED): - # The files were transferred successfully, so delete them - logger.info( - f"{ts_utc()}: Deleting {len(batch.file_paths)} files from successful transfer" - ) - batch.delete_files() - logger.debug("Deletion completed") - batch.file_paths = [] # Mark as processed + if not batch.local_paths_to_delete: + continue # Already processed + if batch.is_globus and batch.task_status != TaskStatus.SUCCEEDED: + continue # Globus transfer not yet confirmed successful + logger.info( + f"{ts_utc()}: Deleting {len(batch.local_paths_to_delete)} files " + f"from successful transfer" + ) + batch.delete_local_files() + batch.local_paths_to_delete = [] # Mark as processed