Migrate logic for default artifact root

Signed-off-by: Ben Wilson <benjamin.wilson@databricks.com>
mlflow · Nov 24, 2021 · 09b84eb · 09b84eb
1 parent 5c1e158
commit 09b84eb
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 11 deletions.
diff --git a/mlflow/cli.py b/mlflow/cli.py
@@ -20,7 +20,7 @@
 from mlflow.utils.annotations import experimental
 from mlflow.utils.logging_utils import eprint
 from mlflow.utils.process import ShellCommandException
-from mlflow.utils.uri import is_local_uri, resolve_default_artifact_root
+from mlflow.utils.uri import resolve_default_artifact_root
 from mlflow.entities.lifecycle_stage import LifecycleStage
 from mlflow.exceptions import MlflowException
 
@@ -271,7 +271,7 @@ def ui(
         backend_store_uri = DEFAULT_LOCAL_FILE_AND_ARTIFACT_PATH
 
     default_artifact_root = resolve_default_artifact_root(
-        serve_artifacts, default_artifact_root, backend_store_uri
+        serve_artifacts, default_artifact_root, backend_store_uri, True
     )
 
     try:
@@ -403,13 +403,6 @@ def server(
     if not backend_store_uri:
         backend_store_uri = DEFAULT_LOCAL_FILE_AND_ARTIFACT_PATH
 
-    if not serve_artifacts and not default_artifact_root and not is_local_uri(backend_store_uri):
-        eprint(
-            "Option 'default-artifact-root' is required, when backend store is not "
-            "local file based."
-        )
-        sys.exit(1)
-
     default_artifact_root = resolve_default_artifact_root(
         serve_artifacts, default_artifact_root, backend_store_uri
     )

diff --git a/mlflow/utils/uri.py b/mlflow/utils/uri.py
@@ -1,3 +1,4 @@
+import sys
 import posixpath
 import urllib.parse
 
@@ -6,6 +7,7 @@
 from mlflow.store.db.db_types import DATABASE_ENGINES
 from mlflow.store.tracking import DEFAULT_LOCAL_FILE_AND_ARTIFACT_PATH, DEFAULT_ARTIFACTS_URI
 from mlflow.utils.validation import _validate_db_type_string
+from mlflow.utils.logging_utils import eprint
 
 _INVALID_DB_URI_MSG = (
     "Please refer to https://mlflow.org/docs/latest/tracking.html#storage for "
@@ -297,13 +299,21 @@ def dbfs_hdfs_uri_to_fuse_path(dbfs_uri):
     return _DBFS_FUSE_PREFIX + dbfs_uri[len(_DBFS_HDFS_URI_PREFIX) :]
 
 
-def resolve_default_artifact_root(serve_artifacts, default_artifact_root, backend_store_uri):
+def resolve_default_artifact_root(
+    serve_artifacts, default_artifact_root, backend_store_uri, resolve_to_local=False
+):
 
     if serve_artifacts and not default_artifact_root:
         default_artifact_root = DEFAULT_ARTIFACTS_URI
     elif not serve_artifacts and not default_artifact_root:
         if is_local_uri(backend_store_uri):
             default_artifact_root = backend_store_uri
-        else:
+        elif resolve_to_local:
             default_artifact_root = DEFAULT_LOCAL_FILE_AND_ARTIFACT_PATH
+        else:
+            eprint(
+                "Option 'default-artifact-root' is required, when backend store is not "
+                "local file based."
+            )
+            sys.exit(1)
     return default_artifact_root