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 __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) # The dbfs:/ path ultimately used for artifact operations should not contain the # Databricks profile info, so strip it before setting ``artifact_uri``. super(DbfsRestArtifactRepository, self).__init__( remove_databricks_profile_info_from_artifact_uri(artifact_uri)) databricks_profile_uri = get_databricks_profile_uri_from_artifact_uri(artifact_uri) if databricks_profile_uri: hostcreds_from_uri = get_databricks_host_creds(databricks_profile_uri) self.get_host_creds = lambda: hostcreds_from_uri else: self.get_host_creds = _get_host_creds_from_default_store()
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)
def test_remove_databricks_profile_info_from_artifact_uri(uri, result): assert remove_databricks_profile_info_from_artifact_uri(uri) == result