Add support for 'mlflow-artifacts' source in import

mlflow_export_import/copy/copy_utils.py

@@ -6,8 +6,8 @@ def get_model_name(artifact_path):
     """
     Return 'my-model' from '/foo/artifacts/my-model'
     """
-    idx = artifact_path.find("artifacts")
-    idx += len("artifacts") + 1
+    idx = artifact_path.find("artifacts/")
+    idx += len("artifacts/")
     return artifact_path[idx:]
 
 

mlflow_export_import/model/import_model.py

@@ -326,13 +326,13 @@ def _extract_model_path(source, run_id):
     if idx == -1:
         raise MlflowExportImportException(f"Cannot find run ID '{run_id}' in registered model version source field '{source}'", http_status_code=404)
     model_path = source[1+idx+len(run_id):]
-    pattern = "artifacts"
+    pattern = "artifacts/"
 
     idx = source.find(pattern)
     if idx == -1: # Bizarre - sometimes there is no 'artifacts' after run_id
         model_path = ""
     else:
-        model_path = source[1+idx+len(pattern):]
+        model_path = source[idx+len(pattern):]
     return model_path
 
 

mlflow_export_import/model_version/import_model_version.py

@@ -111,7 +111,7 @@ def _import_model_version(
     ):
     start_time = time.time()
     dst_source = dst_source.replace("file://","") # OSS MLflow
-    if not (dst_source.startswith("dbfs:") or dst_source.startswith("s3:")) and not os.path.exists(dst_source):
+    if not (dst_source.startswith("dbfs:") or dst_source.startswith("s3:") or dst_source.startswith("mlflow-artifacts:")) and not os.path.exists(dst_source):
         raise MlflowExportImportException(f"'source' argument for MLflowClient.create_model_version does not exist: {dst_source}", http_status_code=404)
 
     tags = src_vr["tags"]
@@ -164,11 +164,11 @@ def _extract_model_path(source):
     :param source: 'source' field of registered model version
     :return: relative path to the model artifact
     """
-    pattern = "artifacts"
+    pattern = "artifacts/"
     idx = source.find(pattern)
     if idx == -1:
         return None
-    return source[1+idx+len(pattern):]
+    return source[idx+len(pattern):]
 
 
 def _set_source_tags_for_field(dct, tags):

mlflow_export_import/run/run_utils.py

@@ -5,8 +5,8 @@
 from mlflow_export_import.common.find_artifacts import find_run_model_names
 
 def get_model_name(artifact_path):
-    idx = artifact_path.find("artifacts")
-    idx += len("artifacts") + 1
+    idx = artifact_path.find("artifacts/")
+    idx += len("artifacts/")
     return artifact_path[idx:]
 
 

tests/open_source/test_models.py

@@ -176,6 +176,7 @@ def _create_run(client):
 _exp_id = "1812"
 _run_id = "48cf29167ddb4e098da780f0959fb4cf"
 _local_path_base = os.path.join("dbfs:/databricks/mlflow-tracking", _exp_id, _run_id)
+_mlflow_artifacts_path_base = os.path.join("mlflow-artifacts:", _exp_id, _run_id)
 
 
 def test_extract_no_artifacts():
@@ -216,6 +217,15 @@ def test_extract_no_run_id():
     except MlflowExportImportException:
         pass
 
+def test_mlflow_artifacts_path_no_artifacts():
+    source = os.path.join(_mlflow_artifacts_path_base)
+    model_path = _extract_model_path(source, _run_id)
+    assert model_path == ""
+
+def test_mlflow_artifacts_path_extract_model():
+    source = os.path.join(_mlflow_artifacts_path_base, "model")
+    model_path = _extract_model_path(source, _run_id)
+    assert model_path == "model"
 
 # == Test that both model names are either UC model names or WS model names