def test_init_validation_and_cleaning(self):
repo = DbfsArtifactRepository('dbfs:/test/',
lambda: MlflowHostCreds('http://host'))
assert repo.artifact_uri == 'dbfs:/test'
with pytest.raises(MlflowException):
DbfsArtifactRepository('s3://test',
lambda: MlflowHostCreds('http://host'))
def test_init_validation_and_cleaning(self):
with mock.patch('mlflow.store.dbfs_artifact_repo._get_host_creds_from_default_store') \
as get_creds_mock:
get_creds_mock.return_value = lambda: MlflowHostCreds('http://host'
)
repo = DbfsArtifactRepository('dbfs:/test/')
assert repo.artifact_uri == 'dbfs:/test'
with pytest.raises(MlflowException):
DbfsArtifactRepository('s3://test')
def from_artifact_uri(artifact_uri, store):
"""
Given an artifact URI for an Experiment Run (e.g., /local/file/path or s3://my/bucket),
returns an ArtifactReposistory instance capable of logging and downloading artifacts
on behalf of this URI.
:param store: An instance of AbstractStore which the artifacts are registered in.
"""
if artifact_uri.startswith("s3:/"):
# Import these locally to avoid creating a circular import loop
from mlflow.store.s3_artifact_repo import S3ArtifactRepository
return S3ArtifactRepository(artifact_uri)
elif artifact_uri.startswith("gs:/"):
from mlflow.store.gcs_artifact_repo import GCSArtifactRepository
return GCSArtifactRepository(artifact_uri)
elif artifact_uri.startswith("wasbs:/"):
from mlflow.store.azure_blob_artifact_repo import AzureBlobArtifactRepository
return AzureBlobArtifactRepository(artifact_uri)
elif artifact_uri.startswith("sftp:/"):
from mlflow.store.sftp_artifact_repo import SFTPArtifactRepository
return SFTPArtifactRepository(artifact_uri)
elif artifact_uri.startswith("dbfs:/"):
from mlflow.store.dbfs_artifact_repo import DbfsArtifactRepository
if not isinstance(store, DatabricksStore):
raise MlflowException(
'`store` must be an instance of DatabricksStore.')
return DbfsArtifactRepository(artifact_uri,
store.http_request_kwargs)
else:
from mlflow.store.local_artifact_repo import LocalArtifactRepository
return LocalArtifactRepository(artifact_uri)
def dbfs_artifact_repo():
return DbfsArtifactRepository('dbfs:/test/',
lambda: MlflowHostCreds('http://host'))
def test_init_validation_and_cleaning(self):
repo = DbfsArtifactRepository('dbfs:/test/', {})
assert repo.artifact_uri == 'dbfs:/test'
with pytest.raises(MlflowException):
DbfsArtifactRepository('s3://test', {})
def dbfs_artifact_repo():
return DbfsArtifactRepository('dbfs:/test/', {})
def dbfs_artifact_repo():
with mock.patch('mlflow.store.dbfs_artifact_repo._get_host_creds_from_default_store') \
as get_creds_mock:
get_creds_mock.return_value = lambda: MlflowHostCreds('http://host')
return DbfsArtifactRepository('dbfs:/test/')