def test_run_relative_artifact_repo_root_path(self, artifact_uri, expected_relative_path):
with mock.patch(
DATABRICKS_ARTIFACT_REPOSITORY + "._get_run_artifact_root"
) as get_run_artifact_root_mock:
get_run_artifact_root_mock.return_value = MOCK_RUN_ROOT_URI
# Basic artifact uri
repo = get_artifact_repository(artifact_uri)
assert repo.run_id == MOCK_RUN_ID
assert repo.run_relative_artifact_repo_root_path == expected_relative_path
def get_model_version_artifact_handler():
from querystring_parser import parser
query_string = request.query_string.decode("utf-8")
request_dict = parser.parse(query_string, normalized=True)
name = request_dict.get("name")
version = request_dict.get("version")
artifact_uri = _get_model_registry_store().get_model_version_download_uri(name, version)
return _send_artifact(get_artifact_repository(artifact_uri), request_dict["path"])
def test_init_validation_and_cleaning(self):
with mock.patch(DBFS_ARTIFACT_REPOSITORY_PACKAGE + '._get_host_creds_from_default_store') \
as get_creds_mock:
get_creds_mock.return_value = lambda: MlflowHostCreds('http://host'
)
repo = get_artifact_repository('dbfs:/test/')
assert repo.artifact_uri == 'dbfs:/test'
with pytest.raises(MlflowException):
DbfsRestArtifactRepository('s3://test')
def get_model_version_artifact_handler():
query_string = request.query_string.decode('utf-8')
request_dict = parser.parse(query_string, normalized=True)
name = request_dict.get('name')
version = request_dict.get('version')
artifact_uri = _get_model_registry_store().get_model_version_download_uri(
name, version)
return _send_artifact(get_artifact_repository(artifact_uri),
request_dict['path'])
def test_extract_run_id(self):
with mock.patch(DATABRICKS_ARTIFACT_REPOSITORY + '._get_run_artifact_root') \
as get_run_artifact_root_mock:
get_run_artifact_root_mock.return_value = MOCK_RUN_ROOT_URI
expected_run_id = "RUN_ID"
repo = get_artifact_repository(
'dbfs:/databricks/mlflow-tracking/EXP/RUN_ID/artifact')
assert repo.run_id == expected_run_id
repo = get_artifact_repository(
'dbfs:/databricks/mlflow-tracking/EXP_ID/RUN_ID/artifacts')
assert repo.run_id == expected_run_id
repo = get_artifact_repository(
'dbfs:/databricks///mlflow-tracking///EXP_ID///RUN_ID///artifacts/'
)
assert repo.run_id == expected_run_id
repo = get_artifact_repository(
'dbfs:/databricks///mlflow-tracking//EXP_ID//RUN_ID///artifacts//'
)
assert repo.run_id == expected_run_id
def log_artifacts(self, run_id, local_dir, artifact_path=None):
"""
Write a directory of files to the remote ``artifact_uri``.
:param local_dir: Path to the directory of files to write.
:param artifact_path: If provided, the directory in ``artifact_uri`` to write to.
"""
run = self.get_run(run_id)
artifact_repo = get_artifact_repository(run.info.artifact_uri)
artifact_repo.log_artifacts(local_dir, artifact_path)
def __init__(self, artifact_uri):
from mlflow.store.artifact.artifact_repository_registry import get_artifact_repository
super().__init__(artifact_uri)
if is_using_databricks_registry(artifact_uri):
# Use the DatabricksModelsArtifactRepository if a databricks profile is being used.
self.repo = DatabricksModelsArtifactRepository(artifact_uri)
else:
uri = ModelsArtifactRepository.get_underlying_uri(artifact_uri)
self.repo = get_artifact_repository(uri)
def list_artifacts(run_id, artifact_path):
"""
Return all the artifacts directly under run's root artifact directory,
or a sub-directory. The output is a JSON-formatted list.
"""
artifact_path = artifact_path if artifact_path is not None else ""
store = _get_store()
artifact_uri = store.get_run(run_id).info.artifact_uri
artifact_repo = get_artifact_repository(artifact_uri)
file_infos = artifact_repo.list_artifacts(artifact_path)
click.echo(_file_infos_to_json(file_infos))
def test_plugin_registration():
artifact_repository_registry = ArtifactRepositoryRegistry()
mock_plugin = mock.Mock()
artifact_repository_registry.register("mock-scheme", mock_plugin)
assert "mock-scheme" in artifact_repository_registry._registry
repository_instance = artifact_repository_registry.get_artifact_repository(
artifact_uri="mock-scheme://fake-host/fake-path")
assert repository_instance == mock_plugin.return_value
mock_plugin.assert_called_once_with("mock-scheme://fake-host/fake-path")
def _get_artifact_repo_mlflow_artifacts():
"""
Get an artifact repository specified by ``--artifacts-destination`` option for ``mlflow server``
command.
"""
from mlflow.server import ARTIFACTS_DESTINATION_ENV_VAR
global _artifact_repo
if _artifact_repo is None:
_artifact_repo = get_artifact_repository(
os.environ[ARTIFACTS_DESTINATION_ENV_VAR])
return _artifact_repo
def test_file_artifact_is_logged_and_downloaded_successfully(s3_artifact_root, tmpdir):
file_name = "test.txt"
file_path = os.path.join(str(tmpdir), file_name)
file_text = "Hello world!"
with open(file_path, "w") as f:
f.write(file_text)
repo = get_artifact_repository(posixpath.join(s3_artifact_root, "some/path"))
repo.log_artifact(file_path)
downloaded_text = open(repo.download_artifacts(file_name)).read()
assert downloaded_text == file_text
def log_artifacts(self, run_id, local_dir, artifact_path=None):
"""
Write a directory of files to the remote ``artifact_uri``.
:param local_dir: Path to the directory of files to write.
:param artifact_path: If provided, the directory in ``artifact_uri`` to write to.
"""
run = self.get_run(run_id)
# TODO: use add_databricks_profile_info_to_artifact_uri(artifact_root, self.tracking_uri)
# for DBFS artifact repositories to use the correct tracking URI.
artifact_repo = get_artifact_repository(run.info.artifact_uri)
artifact_repo.log_artifacts(local_dir, artifact_path)
def log_artifacts(local_dir, run_id, artifact_path):
"""
Log the files within a local directory as an artifact of a run, optionally
within a run-specific artifact path. Run artifacts can be organized into
directories, so you can place the artifact in a directory this way.
"""
store = _get_store()
artifact_uri = store.get_run(run_id).info.artifact_uri
artifact_repo = get_artifact_repository(artifact_uri)
artifact_repo.log_artifacts(local_dir, artifact_path)
_logger.info("Logged artifact from local dir %s to artifact_path=%s",
local_dir, artifact_path)
def _download_artifact_from_uri(artifact_uri, output_path=None):
"""
:param artifact_uri: The *absolute* URI of the artifact to download.
:param output_path: The local filesystem path to which to download the artifact. If unspecified,
a local output path will be created.
"""
if os.path.exists(artifact_uri):
if os.name != "nt":
# If we're dealing with local files, just reference the direct pathing.
# non-nt-based file systems can directly reference path information, while nt-based
# systems need to url-encode special characters in directory listings to be able to
# resolve them (i.e., spaces converted to %20 within a file name or path listing)
root_uri = os.path.dirname(artifact_uri)
artifact_path = os.path.basename(artifact_uri)
return get_artifact_repository(
artifact_uri=root_uri).download_artifacts(
artifact_path=artifact_path, dst_path=output_path)
else: # if we're dealing with nt-based systems, we need to utilize pathname2url to encode.
artifact_uri = path_to_local_file_uri(artifact_uri)
parsed_uri = urllib.parse.urlparse(str(artifact_uri))
prefix = ""
if parsed_uri.scheme and not parsed_uri.path.startswith("/"):
# relative path is a special case, urllib does not reconstruct it properly
prefix = parsed_uri.scheme + ":"
parsed_uri = parsed_uri._replace(scheme="")
# For models:/ URIs, it doesn't make sense to initialize a ModelsArtifactRepository with only
# the model name portion of the URI, then call download_artifacts with the version info.
if ModelsArtifactRepository.is_models_uri(artifact_uri):
root_uri = artifact_uri
artifact_path = ""
else:
artifact_path = posixpath.basename(parsed_uri.path)
parsed_uri = parsed_uri._replace(
path=posixpath.dirname(parsed_uri.path))
root_uri = prefix + urllib.parse.urlunparse(parsed_uri)
return get_artifact_repository(artifact_uri=root_uri).download_artifacts(
artifact_path=artifact_path, dst_path=output_path)
def test_list_artifacts_with_relative_path(self):
list_artifact_file_proto_mock = [
FileInfo(path=posixpath.join(MOCK_SUBDIR, "a.txt"),
is_dir=False,
file_size=0)
]
list_artifacts_dir_proto_mock = [
FileInfo(path=posixpath.join(MOCK_SUBDIR, "test/a.txt"),
is_dir=False,
file_size=100),
FileInfo(path=posixpath.join(MOCK_SUBDIR, "test/dir"),
is_dir=True,
file_size=0),
]
with mock.patch(DATABRICKS_ARTIFACT_REPOSITORY +
"._get_run_artifact_root"
) as get_run_artifact_root_mock, mock.patch(
DATABRICKS_ARTIFACT_REPOSITORY_PACKAGE +
".message_to_json") as message_mock, mock.patch(
DATABRICKS_ARTIFACT_REPOSITORY +
"._call_endpoint") as call_endpoint_mock:
get_run_artifact_root_mock.return_value = MOCK_RUN_ROOT_URI
list_artifact_response_proto = ListArtifacts.Response(
root_uri="",
files=list_artifacts_dir_proto_mock,
next_page_token=None)
call_endpoint_mock.return_value = list_artifact_response_proto
message_mock.return_value = None
databricks_artifact_repo = get_artifact_repository(
MOCK_SUBDIR_ROOT_URI)
artifacts = databricks_artifact_repo.list_artifacts("test")
assert isinstance(artifacts, list)
assert isinstance(artifacts[0], FileInfoEntity)
assert len(artifacts) == 2
assert artifacts[0].path == "test/a.txt"
assert artifacts[0].is_dir is False
assert artifacts[0].file_size == 100
assert artifacts[1].path == "test/dir"
assert artifacts[1].is_dir is True
assert artifacts[1].file_size is None
message_mock.assert_called_with(
ListArtifacts(run_id=MOCK_RUN_ID,
path=posixpath.join(MOCK_SUBDIR, "test")))
# Calling list_artifacts() on a relative path that's a file should return an empty list
list_artifact_response_proto = ListArtifacts.Response(
root_uri="",
files=list_artifact_file_proto_mock,
next_page_token=None)
call_endpoint_mock.return_value = list_artifact_response_proto
artifacts = databricks_artifact_repo.list_artifacts("a.txt")
assert len(artifacts) == 0
def list_artifacts(self, run_id, path=None):
"""
List the artifacts for a run.
:param run_id: The run to list artifacts from.
:param path: The run's relative artifact path to list from. By default it is set to None
or the root artifact path.
:return: List of :py:class:`mlflow.entities.FileInfo`
"""
run = self.get_run(run_id)
artifact_root = run.info.artifact_uri
artifact_repo = get_artifact_repository(artifact_root)
return artifact_repo.list_artifacts(path)
def test_init_validation_and_cleaning(self):
with mock.patch(
DBFS_ARTIFACT_REPOSITORY_PACKAGE +
"._get_host_creds_from_default_store") as get_creds_mock:
get_creds_mock.return_value = lambda: MlflowHostCreds("http://host"
)
repo = get_artifact_repository("dbfs:/test/")
assert repo.artifact_uri == "dbfs:/test"
with pytest.raises(MlflowException):
DbfsRestArtifactRepository("s3://test")
with pytest.raises(MlflowException):
DbfsRestArtifactRepository(
"dbfs://profile@notdatabricks/test/")
def test_download_file_artifact_succeeds_when_artifact_root_is_s3_bucket_root(
s3_artifact_root, tmpdir):
file_a_name = "a.txt"
file_a_text = "A"
file_a_path = os.path.join(str(tmpdir), file_a_name)
with open(file_a_path, "w") as f:
f.write(file_a_text)
repo = get_artifact_repository(s3_artifact_root)
repo.log_artifact(file_a_path)
downloaded_file_path = repo.download_artifacts(file_a_name)
with open(downloaded_file_path, "r") as f:
assert f.read() == file_a_text
def test_dbfs_artifact_repo_delegates_to_correct_repo(is_dbfs_fuse_available,
host_creds_mock): # pylint: disable=unused-argument
# fuse available
is_dbfs_fuse_available.return_value = True
artifact_uri = "dbfs:/databricks/my/absolute/dbfs/path"
repo = get_artifact_repository(artifact_uri)
assert isinstance(repo, LocalArtifactRepository)
assert repo.artifact_dir == os.path.join(os.path.sep, "dbfs", "databricks",
"my", "absolute", "dbfs", "path")
# fuse available but a model repository DBFS location
repo = get_artifact_repository(
"dbfs:/databricks/mlflow-registry/version12345/models")
assert isinstance(repo, DbfsRestArtifactRepository)
# fuse not available
with mock.patch.dict(os.environ,
{'MLFLOW_ENABLE_DBFS_FUSE_ARTIFACT_REPO': 'false'}):
fuse_disabled_repo = get_artifact_repository(artifact_uri)
assert isinstance(fuse_disabled_repo, DbfsRestArtifactRepository)
assert fuse_disabled_repo.artifact_uri == artifact_uri
is_dbfs_fuse_available.return_value = False
rest_repo = get_artifact_repository(artifact_uri)
assert isinstance(rest_repo, DbfsRestArtifactRepository)
assert rest_repo.artifact_uri == artifact_uri
def test_get_s3_client_verify_param_set_correctly(s3_artifact_root,
ignore_tls_env, verify):
from unittest.mock import ANY
with mock.patch.dict("os.environ",
{"MLFLOW_S3_IGNORE_TLS": ignore_tls_env},
clear=True):
with mock.patch("boto3.client") as mock_get_s3_client:
repo = get_artifact_repository(
posixpath.join(s3_artifact_root, "some/path"))
repo._get_s3_client()
mock_get_s3_client.assert_called_with("s3",
config=ANY,
endpoint_url=ANY,
verify=verify)
def test_plugin_registration_via_installed_package():
"""This test requires the package in tests/resources/mlflow-test-plugin to be installed"""
reload(artifact_repository_registry)
assert "file-plugin" in artifact_repository_registry._artifact_repository_registry._registry
from mlflow_test_plugin.local_artifact import PluginLocalArtifactRepository
test_uri = "file-plugin:test-path"
plugin_repo = artifact_repository_registry.get_artifact_repository(test_uri)
assert isinstance(plugin_repo, PluginLocalArtifactRepository)
assert plugin_repo.is_plugin
def list_artifacts(self, run_id, path=None):
"""
List the artifacts for a run.
:param run_id: The run to list artifacts from.
:param path: The run's relative artifact path to list from. By default it is set to None
or the root artifact path.
:return: List of :py:class:`mlflow.entities.FileInfo`
"""
run = self.get_run(run_id)
artifact_root = run.info.artifact_uri
# TODO: use add_databricks_profile_info_to_artifact_uri(artifact_root, self.tracking_uri)
# for DBFS artifact repositories to use the correct tracking URI.
artifact_repo = get_artifact_repository(artifact_root)
return artifact_repo.list_artifacts(path)
def log_artifact(self, run_id, local_path, artifact_path=None):
"""
Write a local file or directory to the remote ``artifact_uri``.
:param local_path: Path to the file or directory to write.
:param artifact_path: If provided, the directory in ``artifact_uri`` to write to.
"""
run = self.get_run(run_id)
artifact_repo = get_artifact_repository(run.info.artifact_uri)
if os.path.isdir(local_path):
dir_name = os.path.basename(os.path.normpath(local_path))
path_name = os.path.join(artifact_path, dir_name) \
if artifact_path is not None else dir_name
artifact_repo.log_artifacts(local_path, path_name)
else:
artifact_repo.log_artifact(local_path, artifact_path)
def test_file_artifact_is_logged_with_content_metadata(s3_artifact_root, tmpdir):
file_name = "test.txt"
file_path = os.path.join(str(tmpdir), file_name)
file_text = "Hello world!"
with open(file_path, "w") as f:
f.write(file_text)
repo = get_artifact_repository(posixpath.join(s3_artifact_root, "some/path"))
repo.log_artifact(file_path)
bucket, _ = repo.parse_s3_uri(s3_artifact_root)
s3_client = repo._get_s3_client()
response = s3_client.head_object(Bucket=bucket, Key="some/path/test.txt")
assert response.get("ContentType") == "text/plain"
assert response.get("ContentEncoding") is None
def _get_artifact_repo(self, run_id):
# Attempt to fetch the artifact repo from a local cache
cached_repo = TrackingServiceClient._artifact_repos_cache.get(run_id)
if cached_repo is not None:
return cached_repo
else:
run = self.get_run(run_id)
artifact_uri = add_databricks_profile_info_to_artifact_uri(
run.info.artifact_uri, self.tracking_uri)
artifact_repo = get_artifact_repository(artifact_uri)
# Cache the artifact repo to avoid a future network call, removing the oldest
# entry in the cache if there are too many elements
if len(TrackingServiceClient._artifact_repos_cache) > 1024:
TrackingServiceClient._artifact_repos_cache.popitem(last=False)
TrackingServiceClient._artifact_repos_cache[run_id] = artifact_repo
return artifact_repo
def _download_artifact_from_uri(artifact_uri, output_path=None):
"""
:param artifact_uri: The *absolute* URI of the artifact to download.
:param output_path: The local filesystem path to which to download the artifact. If unspecified,
a local output path will be created.
"""
parsed_uri = urllib.parse.urlparse(artifact_uri)
prefix = ""
if parsed_uri.scheme and not parsed_uri.path.startswith("/"):
# relative path is a special case, urllib does not reconstruct it properly
prefix = parsed_uri.scheme + ":"
parsed_uri = parsed_uri._replace(scheme="")
artifact_path = posixpath.basename(parsed_uri.path)
parsed_uri = parsed_uri._replace(path=posixpath.dirname(parsed_uri.path))
root_uri = prefix + urllib.parse.urlunparse(parsed_uri)
return get_artifact_repository(artifact_uri=root_uri).download_artifacts(
artifact_path=artifact_path, dst_path=output_path)
def test_download_directory_artifact_succeeds_when_artifact_root_is_s3_bucket_root(
s3_artifact_root, tmpdir):
file_a_name = "a.txt"
file_a_text = "A"
subdir_path = str(tmpdir.mkdir("subdir"))
nested_path = os.path.join(subdir_path, "nested")
os.makedirs(nested_path)
with open(os.path.join(nested_path, file_a_name), "w") as f:
f.write(file_a_text)
repo = get_artifact_repository(s3_artifact_root)
repo.log_artifacts(subdir_path)
downloaded_dir_path = repo.download_artifacts("nested")
assert file_a_name in os.listdir(downloaded_dir_path)
with open(os.path.join(downloaded_dir_path, file_a_name), "r") as f:
assert f.read() == file_a_text
def log_artifact(self, run_id, local_path, artifact_path=None):
"""
Write a local file or directory to the remote ``artifact_uri``.
:param local_path: Path to the file or directory to write.
:param artifact_path: If provided, the directory in ``artifact_uri`` to write to.
"""
run = self.get_run(run_id)
# TODO: use add_databricks_profile_info_to_artifact_uri(artifact_root, self.tracking_uri)
# for DBFS artifact repositories to use the correct tracking URI.
artifact_repo = get_artifact_repository(run.info.artifact_uri)
if os.path.isdir(local_path):
dir_name = os.path.basename(os.path.normpath(local_path))
path_name = os.path.join(artifact_path, dir_name) \
if artifact_path is not None else dir_name
artifact_repo.log_artifacts(local_path, path_name)
else:
artifact_repo.log_artifact(local_path, artifact_path)
def test_init_validation_and_cleaning(self):
with mock.patch(DATABRICKS_ARTIFACT_REPOSITORY + '._get_run_artifact_root') \
as get_run_artifact_root_mock:
get_run_artifact_root_mock.return_value = MOCK_RUN_ROOT_URI
# Basic artifact uri
repo = get_artifact_repository(
'dbfs:/databricks/mlflow-tracking/MOCK-EXP/MOCK-RUN-ID/artifacts'
)
assert repo.artifact_uri == 'dbfs:/databricks/mlflow-tracking/' \
'MOCK-EXP/MOCK-RUN-ID/artifacts'
assert repo.run_id == MOCK_RUN_ID
assert repo.run_relative_artifact_repo_root_path == ""
with pytest.raises(MlflowException):
DatabricksArtifactRepository('s3://test')
with pytest.raises(MlflowException):
DatabricksArtifactRepository(
'dbfs:/databricks/mlflow/EXP/RUN/artifact')
def download_artifacts(self, run_id, path, dst_path=None):
"""
Download an artifact file or directory from a run to a local directory if applicable,
and return a local path for it.
:param run_id: The run to download artifacts from.
:param path: Relative source path to the desired artifact.
:param dst_path: Absolute path of the local filesystem destination directory to which to
download the specified artifacts. This directory must already exist.
If unspecified, the artifacts will either be downloaded to a new
uniquely-named directory on the local filesystem or will be returned
directly in the case of the LocalArtifactRepository.
:return: Local path of desired artifact.
"""
run = self.get_run(run_id)
artifact_root = run.info.artifact_uri
artifact_repo = get_artifact_repository(artifact_root)
return artifact_repo.download_artifacts(path, dst_path)
