diff --git a/composer/utils/object_store/mlflow_object_store.py b/composer/utils/object_store/mlflow_object_store.py index 6a6655fba8..e1d3df8d61 100644 --- a/composer/utils/object_store/mlflow_object_store.py +++ b/composer/utils/object_store/mlflow_object_store.py @@ -517,13 +517,17 @@ 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) 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