def __init__(self, artifact_uri):
super(DatabricksArtifactRepository, self).__init__(artifact_uri)
if not artifact_uri.startswith('dbfs:/'):
raise MlflowException(
message=
'DatabricksArtifactRepository URI must start with dbfs:/',
error_code=INVALID_PARAMETER_VALUE)
if not is_databricks_acled_artifacts_uri(artifact_uri):
raise MlflowException(
message=('Artifact URI incorrect. Expected path prefix to be'
' databricks/mlflow-tracking/path/to/artifact/..'),
error_code=INVALID_PARAMETER_VALUE)
self.run_id = self._extract_run_id(self.artifact_uri)
# Fetch the artifact root for the MLflow Run associated with `artifact_uri` and compute
# the path of `artifact_uri` relative to the MLflow Run's artifact root
# (the `run_relative_artifact_repo_root_path`). All operations performed on this artifact
# repository will be performed relative to this computed location
artifact_repo_root_path = extract_and_normalize_path(artifact_uri)
run_artifact_root_uri = self._get_run_artifact_root(self.run_id)
run_artifact_root_path = extract_and_normalize_path(
run_artifact_root_uri)
run_relative_root_path = posixpath.relpath(
path=artifact_repo_root_path, start=run_artifact_root_path)
# If the paths are equal, then use empty string over "./" for ListArtifact compatibility.
self.run_relative_artifact_repo_root_path = \
"" if run_artifact_root_path == artifact_repo_root_path else run_relative_root_path
def __init__(self, artifact_uri):
if not is_valid_dbfs_uri(artifact_uri):
raise MlflowException(message="DBFS URI must be of the form dbfs:/<path> or " +
"dbfs://profile@databricks/<path>",
error_code=INVALID_PARAMETER_VALUE)
if not is_databricks_acled_artifacts_uri(artifact_uri):
raise MlflowException(message=('Artifact URI incorrect. Expected path prefix to be'
' databricks/mlflow-tracking/path/to/artifact/..'),
error_code=INVALID_PARAMETER_VALUE)
# The dbfs:/ path ultimately used for artifact operations should not contain the
# Databricks profile info, so strip it before setting ``artifact_uri``.
super(DatabricksArtifactRepository, self).__init__(
remove_databricks_profile_info_from_artifact_uri(artifact_uri))
self.databricks_profile_uri = get_databricks_profile_uri_from_artifact_uri(artifact_uri) \
or mlflow.tracking.get_tracking_uri()
self.run_id = self._extract_run_id(self.artifact_uri)
# Fetch the artifact root for the MLflow Run associated with `artifact_uri` and compute
# the path of `artifact_uri` relative to the MLflow Run's artifact root
# (the `run_relative_artifact_repo_root_path`). All operations performed on this artifact
# repository will be performed relative to this computed location
artifact_repo_root_path = extract_and_normalize_path(artifact_uri)
run_artifact_root_uri = self._get_run_artifact_root(self.run_id)
run_artifact_root_path = extract_and_normalize_path(run_artifact_root_uri)
run_relative_root_path = posixpath.relpath(
path=artifact_repo_root_path, start=run_artifact_root_path
)
# If the paths are equal, then use empty string over "./" for ListArtifact compatibility.
self.run_relative_artifact_repo_root_path = \
"" if run_artifact_root_path == artifact_repo_root_path else run_relative_root_path
def dbfs_artifact_repo_factory(artifact_uri):
"""
Returns an ArtifactRepository subclass for storing artifacts on DBFS.
This factory method is used with URIs of the form ``dbfs:/<path>``. DBFS-backed artifact
storage can only be used together with the RestStore.
In the special case where the URI is of the form
`dbfs:/databricks/mlflow-tracking/<Exp-ID>/<Run-ID>/<path>',
a DatabricksArtifactRepository is returned. This is capable of storing access controlled
artifacts.
:param artifact_uri: DBFS root artifact URI (string).
:return: Subclass of ArtifactRepository capable of storing artifacts on DBFS.
"""
cleaned_artifact_uri = artifact_uri.rstrip('/')
uri_scheme = get_uri_scheme(artifact_uri)
if uri_scheme != 'dbfs':
raise MlflowException(
"DBFS URI must be of the form "
"dbfs:/<path>, but received {uri}".format(uri=artifact_uri))
if is_databricks_acled_artifacts_uri(artifact_uri):
return DatabricksArtifactRepository(cleaned_artifact_uri)
elif mlflow.utils.databricks_utils.is_dbfs_fuse_available() \
and os.environ.get(USE_FUSE_ENV_VAR, "").lower() != "false" \
and not artifact_uri.startswith("dbfs:/databricks/mlflow-registry"):
# If the DBFS FUSE mount is available, write artifacts directly to /dbfs/... using
# local filesystem APIs
file_uri = "file:///dbfs/{}".format(
strip_prefix(cleaned_artifact_uri, "dbfs:/"))
return LocalArtifactRepository(file_uri)
return DbfsRestArtifactRepository(cleaned_artifact_uri)
def test_is_databricks_acled_artifacts_uri():
assert is_databricks_acled_artifacts_uri(
"dbfs:databricks/mlflow-tracking/EXP_ID/RUN_ID/artifacts")
assert is_databricks_acled_artifacts_uri(
"dbfs:/databricks/mlflow-tracking/EXP_ID/RUN_ID/artifacts")
assert is_databricks_acled_artifacts_uri(
"dbfs:///databricks/mlflow-tracking/EXP_ID/RUN_ID/artifacts")
assert is_databricks_acled_artifacts_uri(
"dbfs:/databricks///mlflow-tracking///EXP_ID///RUN_ID///artifacts/")
assert is_databricks_acled_artifacts_uri(
"dbfs:///databricks///mlflow-tracking//EXP_ID//RUN_ID///artifacts//")
assert is_databricks_acled_artifacts_uri(
"dbfs:databricks///mlflow-tracking//EXP_ID//RUN_ID///artifacts//")
assert not is_databricks_acled_artifacts_uri(
"dbfs:/databricks/mlflow//EXP_ID//RUN_ID///artifacts//")
def __init__(self, artifact_uri):
if not is_valid_dbfs_uri(artifact_uri):
raise MlflowException(
message="DBFS URI must be of the form dbfs:/<path> or "
+ "dbfs://profile@databricks/<path>",
error_code=INVALID_PARAMETER_VALUE,
)
if not is_databricks_acled_artifacts_uri(artifact_uri):
raise MlflowException(
message=(
"Artifact URI incorrect. Expected path prefix to be"
" databricks/mlflow-tracking/path/to/artifact/.."
),
error_code=INVALID_PARAMETER_VALUE,
)
# The dbfs:/ path ultimately used for artifact operations should not contain the
# Databricks profile info, so strip it before setting ``artifact_uri``.
super().__init__(remove_databricks_profile_info_from_artifact_uri(artifact_uri))
self.databricks_profile_uri = (
get_databricks_profile_uri_from_artifact_uri(artifact_uri)
or mlflow.tracking.get_tracking_uri()
)
self.run_id = self._extract_run_id(self.artifact_uri)
# Fetch the artifact root for the MLflow Run associated with `artifact_uri` and compute
# the path of `artifact_uri` relative to the MLflow Run's artifact root
# (the `run_relative_artifact_repo_root_path`). All operations performed on this artifact
# repository will be performed relative to this computed location
artifact_repo_root_path = extract_and_normalize_path(artifact_uri)
run_artifact_root_uri = self._get_run_artifact_root(self.run_id)
run_artifact_root_path = extract_and_normalize_path(run_artifact_root_uri)
run_relative_root_path = posixpath.relpath(
path=artifact_repo_root_path, start=run_artifact_root_path
)
# If the paths are equal, then use empty string over "./" for ListArtifact compatibility.
self.run_relative_artifact_repo_root_path = (
"" if run_artifact_root_path == artifact_repo_root_path else run_relative_root_path
)
# Limit the number of threads used for artifact uploads, using at most 8 threads or
# 2 * the number of CPU cores available on the system (whichever is smaller)
num_cpus = os.cpu_count() or 4
num_artifact_workers = min(num_cpus * 2, 8)
self.thread_pool = ThreadPoolExecutor(max_workers=num_artifact_workers)
def dbfs_artifact_repo_factory(artifact_uri):
"""
Returns an ArtifactRepository subclass for storing artifacts on DBFS.
This factory method is used with URIs of the form ``dbfs:/<path>``. DBFS-backed artifact
storage can only be used together with the RestStore.
In the special case where the URI is of the form
`dbfs:/databricks/mlflow-tracking/<Exp-ID>/<Run-ID>/<path>',
a DatabricksArtifactRepository is returned. This is capable of storing access controlled
artifacts.
:param artifact_uri: DBFS root artifact URI (string).
:return: Subclass of ArtifactRepository capable of storing artifacts on DBFS.
"""
if not is_valid_dbfs_uri(artifact_uri):
raise MlflowException(
"DBFS URI must be of the form dbfs:/<path> or "
+ "dbfs://profile@databricks/<path>, but received "
+ artifact_uri
)
cleaned_artifact_uri = artifact_uri.rstrip("/")
db_profile_uri = get_databricks_profile_uri_from_artifact_uri(cleaned_artifact_uri)
if is_databricks_acled_artifacts_uri(artifact_uri):
return DatabricksArtifactRepository(cleaned_artifact_uri)
elif (
mlflow.utils.databricks_utils.is_dbfs_fuse_available()
and os.environ.get(USE_FUSE_ENV_VAR, "").lower() != "false"
and not is_databricks_model_registry_artifacts_uri(artifact_uri)
and (db_profile_uri is None or db_profile_uri == "databricks")
):
# If the DBFS FUSE mount is available, write artifacts directly to
# /dbfs/... using local filesystem APIs.
# Note: it is possible for a named Databricks profile to point to the current workspace,
# but we're going to avoid doing a complex check and assume users will use `databricks`
# to mean the current workspace. Using `DbfsRestArtifactRepository` to access the current
# workspace's DBFS should still work; it just may be slower.
final_artifact_uri = remove_databricks_profile_info_from_artifact_uri(cleaned_artifact_uri)
file_uri = "file:///dbfs/{}".format(strip_prefix(final_artifact_uri, "dbfs:/"))
return LocalArtifactRepository(file_uri)
return DbfsRestArtifactRepository(cleaned_artifact_uri)
