wrap another function in mlflow

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)