From 94e2722fe5d0a8dd709e7e2480fe1a9cf7dd9cb6 Mon Sep 17 00:00:00 2001 From: Nancy Hung Date: Wed, 30 Oct 2024 17:12:40 -0700 Subject: [PATCH 1/3] wrap upload in error handling too --- composer/utils/object_store/uc_object_store.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/composer/utils/object_store/uc_object_store.py b/composer/utils/object_store/uc_object_store.py index b7801d9757..5516c4123f 100644 --- a/composer/utils/object_store/uc_object_store.py +++ b/composer/utils/object_store/uc_object_store.py @@ -157,8 +157,11 @@ def upload_object( """ # remove unused variable del callback - with open(filename, 'rb') as f: - self.client.files.upload(self._get_object_path(object_name), f) + try: + with open(filename, 'rb') as f: + self.client.files.upload(self._get_object_path(object_name), f) + except Exception as e: + _wrap_errors(self.get_uri(object_name), e) def download_object( self, @@ -192,7 +195,6 @@ def download_object( tmp_path = str(filename) + f'{uuid.uuid4()}.tmp' try: - try: contents = self.client.files.download(self._get_object_path(object_name)).contents assert contents is not None From b39ccfb304bd2d1847e56276f0043849640ecb78 Mon Sep 17 00:00:00 2001 From: Nancy Hung Date: Wed, 30 Oct 2024 17:17:39 -0700 Subject: [PATCH 2/3] wrap another function in mlflow --- composer/utils/object_store/mlflow_object_store.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/composer/utils/object_store/mlflow_object_store.py b/composer/utils/object_store/mlflow_object_store.py index 6a6655fba8..705dc97b87 100644 --- a/composer/utils/object_store/mlflow_object_store.py +++ b/composer/utils/object_store/mlflow_object_store.py @@ -517,13 +517,15 @@ def _get_artifact_info(self, object_name): Returns: Optional[FileInfo]: The :class:`~mlflow.entities.FileInfo` for the object, or None if it does not exist. """ + from mlflow.exceptions import MlflowException # MLflow doesn't support info for a singleton artifact, so we need to list all artifacts in the # parent path and find the one with the matching name. artifact_path = self.get_artifact_path(object_name) artifact_dir = os.path.dirname(artifact_path) - artifacts = self._mlflow_client.list_artifacts(self.run_id, artifact_dir) - for artifact in artifacts: - if not artifact.is_dir and artifact.path == artifact_path: - return artifact - - return None + try: + artifacts = self._mlflow_client.list_artifacts(self.run_id, artifact_dir) + for artifact in artifacts: + if not artifact.is_dir and artifact.path == artifact_path: + return artifact + except MlflowException as e: + _wrap_mlflow_exceptions(object_name, e) From 8af758ce1889dff4dbd0b3e64ca64b8710b69549 Mon Sep 17 00:00:00 2001 From: Nancy Hung Date: Thu, 14 Nov 2024 11:53:30 -0800 Subject: [PATCH 3/3] lint? --- composer/utils/object_store/mlflow_object_store.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/composer/utils/object_store/mlflow_object_store.py b/composer/utils/object_store/mlflow_object_store.py index 705dc97b87..e1d3df8d61 100644 --- a/composer/utils/object_store/mlflow_object_store.py +++ b/composer/utils/object_store/mlflow_object_store.py @@ -518,10 +518,12 @@ def _get_artifact_info(self, object_name): Optional[FileInfo]: The :class:`~mlflow.entities.FileInfo` for the object, or None if it does not exist. """ from mlflow.exceptions import MlflowException + # MLflow doesn't support info for a singleton artifact, so we need to list all artifacts in the # parent path and find the one with the matching name. artifact_path = self.get_artifact_path(object_name) artifact_dir = os.path.dirname(artifact_path) + try: artifacts = self._mlflow_client.list_artifacts(self.run_id, artifact_dir) for artifact in artifacts: