diff --git a/modal/_output.py b/modal/_output.py index dd901b417..a7c6e7c89 100644 --- a/modal/_output.py +++ b/modal/_output.py @@ -462,18 +462,21 @@ def progress( reset: Optional[bool] = False, complete: Optional[bool] = False, ) -> Optional[TaskID]: - if task_id is not None: - if reset: - return self._reset_sub_task(task_id) + try: + if task_id is not None: + if reset: + return self._reset_sub_task(task_id) + elif complete: + return self._complete_sub_task(task_id) + elif advance is not None: + return self._advance_sub_task(task_id, advance) + elif name is not None and size is not None: + return self._add_sub_task(name, size) elif complete: - return self._complete_sub_task(task_id) - elif advance is not None: - return self._advance_sub_task(task_id, advance) - elif name is not None and size is not None: - return self._add_sub_task(name, size) - elif complete: - return self._complete_progress() - + return self._complete_progress() + except Exception as exc: + # Liberal exception handling to avoid crashing downloads and uploads. + logger.error(f"failed progress update: {exc}") raise NotImplementedError( "Unknown action to take with args: " + f"name={name} " diff --git a/modal/_utils/blob_utils.py b/modal/_utils/blob_utils.py index 0c02f46f6..ad6bd6874 100644 --- a/modal/_utils/blob_utils.py +++ b/modal/_utils/blob_utils.py @@ -114,8 +114,6 @@ async def safe_read(): await writer.write(chunk) self.progress_report_cb(advance=len(chunk)) - self.progress_report_cb(complete=True) - def remaining_bytes(self): return self.segment_length - self.num_bytes_read @@ -282,6 +280,9 @@ async def _blob_upload( content_md5_b64=upload_hashes.md5_base64, ) + if progress_report_cb: + progress_report_cb(complete=True) + return blob_id