star

venv/lib/python3.9/site-packages/mlflow/store/artifact/artifact_repo.py

@@ -0,0 +1,466 @@
+import json
+import logging
+import os
+import posixpath
+import tempfile
+import traceback
+from abc import ABC, ABCMeta, abstractmethod
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Any, Optional
+
+from mlflow.entities.file_info import FileInfo
+from mlflow.entities.multipart_upload import (
+    CreateMultipartUploadResponse,
+    MultipartUploadPart,
+)
+from mlflow.exceptions import (
+    MlflowException,
+    MlflowTraceDataCorrupted,
+    MlflowTraceDataNotFound,
+)
+from mlflow.protos.databricks_pb2 import (
+    INVALID_PARAMETER_VALUE,
+    RESOURCE_DOES_NOT_EXIST,
+)
+from mlflow.tracing.artifact_utils import TRACE_DATA_FILE_NAME
+from mlflow.utils.annotations import developer_stable
+from mlflow.utils.async_logging.async_artifacts_logging_queue import (
+    AsyncArtifactsLoggingQueue,
+)
+from mlflow.utils.file_utils import ArtifactProgressBar, create_tmp_dir
+from mlflow.utils.validation import bad_path_message, path_not_unique
+
+# Constants used to determine max level of parallelism to use while uploading/downloading artifacts.
+# Max threads to use for parallelism.
+_NUM_MAX_THREADS = 20
+# Max threads per CPU
+_NUM_MAX_THREADS_PER_CPU = 2
+assert _NUM_MAX_THREADS >= _NUM_MAX_THREADS_PER_CPU
+assert _NUM_MAX_THREADS_PER_CPU > 0
+# Default number of CPUs to assume on the machine if unavailable to fetch it using os.cpu_count()
+_NUM_DEFAULT_CPUS = _NUM_MAX_THREADS // _NUM_MAX_THREADS_PER_CPU
+_logger = logging.getLogger(__name__)
+
+
+def _truncate_error(err: str, max_length: int = 10_000) -> str:
+    if len(err) <= max_length:
+        return err
+    half = max_length // 2
+    return err[:half] + "\n\n*** Error message is too long, truncated ***\n\n" + err[-half:]
+
+
+def _retry_with_new_creds(try_func, creds_func, orig_creds=None):
+    """
+    Attempt the try_func with the original credentials (og_creds) if provided, or by generating the
+    credentials using creds_func. If the try_func throws, then try again with new credentials
+    provided by creds_func.
+    """
+    try:
+        first_creds = creds_func() if orig_creds is None else orig_creds
+        return try_func(first_creds)
+    except Exception as e:
+        _logger.info(
+            f"Failed to complete request, possibly due to credential expiration (Error: {e})."
+            " Refreshing credentials and trying again..."
+        )
+        new_creds = creds_func()
+        return try_func(new_creds)
+
+
+@developer_stable
+class ArtifactRepository:
+    """
+    Abstract artifact repo that defines how to upload (log) and download potentially large
+    artifacts from different storage backends.
+    """
+
+    __metaclass__ = ABCMeta
+
+    def __init__(self, artifact_uri):
+        self.artifact_uri = artifact_uri
+        # Limit the number of threads used for artifact uploads/downloads. Use at most
+        # constants._NUM_MAX_THREADS threads or 2 * the number of CPU cores available on the
+        # system (whichever is smaller)
+        self.thread_pool = self._create_thread_pool()
+
+        def log_artifact_handler(filename, artifact_path=None, artifact=None):
+            with tempfile.TemporaryDirectory() as tmp_dir:
+                tmp_path = os.path.join(tmp_dir, filename)
+                if artifact is not None:
+                    # User should already have installed PIL to log a PIL image
+                    from PIL import Image
+
+                    if isinstance(artifact, Image.Image):
+                        artifact.save(tmp_path)
+                self.log_artifact(tmp_path, artifact_path)
+
+        self._async_logging_queue = AsyncArtifactsLoggingQueue(log_artifact_handler)
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(artifact_uri={self.artifact_uri!r})"
+
+    def _create_thread_pool(self):
+        return ThreadPoolExecutor(
+            max_workers=self.max_workers, thread_name_prefix=f"Mlflow{self.__class__.__name__}"
+        )
+
+    def flush_async_logging(self):
+        """
+        Flushes the async logging queue, ensuring that all pending logging operations have
+        completed.
+        """
+        if self._async_logging_queue._is_activated:
+            self._async_logging_queue.flush()
+
+    @abstractmethod
+    def log_artifact(self, local_file, artifact_path=None):
+        """
+        Log a local file as an artifact, optionally taking an ``artifact_path`` to place it in
+        within the run's artifacts. Run artifacts can be organized into directories, so you can
+        place the artifact in a directory this way.
+
+        Args:
+            local_file: Path to artifact to log.
+            artifact_path: Directory within the run's artifact directory in which to log the
+                artifact.
+        """
+
+    def _log_artifact_async(self, filename, artifact_path=None, artifact=None):
+        """
+        Asynchronously log a local file as an artifact, optionally taking an ``artifact_path`` to
+        place it within the run's artifacts. Run artifacts can be organized into directory, so you
+        can place the artifact in the directory this way. Cleanup tells the function whether to
+        cleanup the local_file after running log_artifact, since it could be a Temporary
+        Directory.
+
+        Args:
+            filename: Filename of the artifact to be logged.
+            artifact_path: Directory within the run's artifact directory in which to log the
+                artifact.
+            artifact: The artifact to be logged.
+
+        Returns:
+            An :py:class:`mlflow.utils.async_logging.run_operations.RunOperations` instance
+            that represents future for logging operation.
+        """
+
+        if not self._async_logging_queue.is_active():
+            self._async_logging_queue.activate()
+
+        return self._async_logging_queue.log_artifacts_async(
+            filename=filename, artifact_path=artifact_path, artifact=artifact
+        )
+
+    @abstractmethod
+    def log_artifacts(self, local_dir, artifact_path=None):
+        """
+        Log the files in the specified local directory as artifacts, optionally taking
+        an ``artifact_path`` to place them in within the run's artifacts.
+
+        Args:
+            local_dir: Directory of local artifacts to log.
+            artifact_path: Directory within the run's artifact directory in which to log the
+                artifacts.
+        """
+
+    @abstractmethod
+    def list_artifacts(self, path: Optional[str] = None) -> list:
+        """
+        Return all the artifacts for this run_id directly under path. If path is a file, returns
+        an empty list. Will error if path is neither a file nor directory.
+
+        Args:
+            path: Relative source path that contains desired artifacts.
+
+        Returns:
+            List of artifacts as FileInfo listed directly under path.
+        """
+
+    def _is_directory(self, artifact_path):
+        listing = self.list_artifacts(artifact_path)
+        return len(listing) > 0
+
+    def _create_download_destination(self, src_artifact_path, dst_local_dir_path=None):
+        """
+        Creates a local filesystem location to be used as a destination for downloading the artifact
+        specified by `src_artifact_path`. The destination location is a subdirectory of the
+        specified `dst_local_dir_path`, which is determined according to the structure of
+        `src_artifact_path`. For example, if `src_artifact_path` is `dir1/file1.txt`, then the
+        resulting destination path is `<dst_local_dir_path>/dir1/file1.txt`. Local directories are
+        created for the resulting destination location if they do not exist.
+
+        Args:
+            src_artifact_path: A relative, POSIX-style path referring to an artifact stored
+                within the repository's artifact root location. `src_artifact_path` should be
+                specified relative to the repository's artifact root location.
+            dst_local_dir_path: The absolute path to a local filesystem directory in which the
+                local destination path will be contained. The local destination path may be
+                contained in a subdirectory of `dst_root_dir` if `src_artifact_path` contains
+                subdirectories.
+
+        Returns:
+            The absolute path to a local filesystem location to be used as a destination
+            for downloading the artifact specified by `src_artifact_path`.
+        """
+        src_artifact_path = src_artifact_path.rstrip("/")  # Ensure correct dirname for trailing '/'
+        dirpath = posixpath.dirname(src_artifact_path)
+        local_dir_path = os.path.join(dst_local_dir_path, dirpath)
+        local_file_path = os.path.join(dst_local_dir_path, src_artifact_path)
+        if not os.path.exists(local_dir_path):
+            os.makedirs(local_dir_path, exist_ok=True)
+        return local_file_path
+
+    def _iter_artifacts_recursive(self, path):
+        dir_content = [
+            file_info
+            for file_info in self.list_artifacts(path)
+            # prevent infinite loop, sometimes the dir is recursively included
+            if file_info.path not in [".", path]
+        ]
+        # Empty directory
+        if not dir_content:
+            yield FileInfo(path=path, is_dir=True, file_size=None)
+            return
+
+        for file_info in dir_content:
+            if file_info.is_dir:
+                yield from self._iter_artifacts_recursive(file_info.path)
+            else:
+                yield file_info
+
+    def download_artifacts(self, artifact_path, dst_path=None):
+        """
+        Download an artifact file or directory to a local directory if applicable, and return a
+        local path for it.
+        The caller is responsible for managing the lifecycle of the downloaded artifacts.
+
+        Args:
+            artifact_path: Relative source path to the desired artifacts.
+            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.
+
+        Returns:
+            Absolute path of the local filesystem location containing the desired artifacts.
+        """
+        if dst_path:
+            dst_path = os.path.abspath(dst_path)
+            if not os.path.exists(dst_path):
+                raise MlflowException(
+                    message=(
+                        "The destination path for downloaded artifacts does not"
+                        f" exist! Destination path: {dst_path}"
+                    ),
+                    error_code=RESOURCE_DOES_NOT_EXIST,
+                )
+            elif not os.path.isdir(dst_path):
+                raise MlflowException(
+                    message=(
+                        "The destination path for downloaded artifacts must be a directory!"
+                        f" Destination path: {dst_path}"
+                    ),
+                    error_code=INVALID_PARAMETER_VALUE,
+                )
+        else:
+            dst_path = create_tmp_dir()
+
+        def _download_file(src_artifact_path, dst_local_dir_path):
+            dst_local_file_path = self._create_download_destination(
+                src_artifact_path=src_artifact_path, dst_local_dir_path=dst_local_dir_path
+            )
+            return self.thread_pool.submit(
+                self._download_file,
+                remote_file_path=src_artifact_path,
+                local_path=dst_local_file_path,
+            )
+
+        # Submit download tasks
+        futures = {}
+        if self._is_directory(artifact_path):
+            for file_info in self._iter_artifacts_recursive(artifact_path):
+                if file_info.is_dir:  # Empty directory
+                    os.makedirs(os.path.join(dst_path, file_info.path), exist_ok=True)
+                else:
+                    fut = _download_file(file_info.path, dst_path)
+                    futures[fut] = file_info.path
+        else:
+            fut = _download_file(artifact_path, dst_path)
+            futures[fut] = artifact_path
+
+        # Wait for downloads to complete and collect failures
+        failed_downloads = {}
+        tracebacks = {}
+        with ArtifactProgressBar.files(desc="Downloading artifacts", total=len(futures)) as pbar:
+            for f in as_completed(futures):
+                try:
+                    f.result()
+                    pbar.update()
+                except Exception as e:
+                    path = futures[f]
+                    failed_downloads[path] = e
+                    tracebacks[path] = traceback.format_exc()
+
+        if failed_downloads:
+            if _logger.isEnabledFor(logging.DEBUG):
+                template = "##### File {path} #####\n{error}\nTraceback:\n{traceback}\n"
+            else:
+                template = "##### File {path} #####\n{error}"
+
+            failures = "\n".join(
+                template.format(path=path, error=error, traceback=tracebacks[path])
+                for path, error in failed_downloads.items()
+            )
+            raise MlflowException(
+                message=(
+                    "The following failures occurred while downloading one or more"
+                    f" artifacts from {self.artifact_uri}:\n{_truncate_error(failures)}"
+                )
+            )
+
+        return os.path.join(dst_path, artifact_path)
+
+    @abstractmethod
+    def _download_file(self, remote_file_path, local_path):
+        """
+        Download the file at the specified relative remote path and saves
+        it at the specified local path.
+
+        Args:
+            remote_file_path: Source path to the remote file, relative to the root
+                directory of the artifact repository.
+            local_path: The path to which to save the downloaded file.
+        """
+
+    def delete_artifacts(self, artifact_path=None):
+        """
+        Delete the artifacts at the specified location.
+        Supports the deletion of a single file or of a directory. Deletion of a directory
+        is recursive.
+
+        Args:
+            artifact_path: Path of the artifact to delete.
+        """
+
+    @property
+    def max_workers(self) -> int:
+        """Compute the number of workers to use for multi-threading."""
+        num_cpus = os.cpu_count() or _NUM_DEFAULT_CPUS
+        return min(num_cpus * _NUM_MAX_THREADS_PER_CPU, _NUM_MAX_THREADS)
+
+    def download_trace_data(self) -> dict[str, Any]:
+        """
+        Download the trace data.
+
+        Returns:
+            The trace data as a dictionary.
+
+        Raises:
+            - `MlflowTraceDataNotFound`: The trace data is not found.
+            - `MlflowTraceDataCorrupted`: The trace data is corrupted.
+        """
+        with tempfile.TemporaryDirectory() as temp_dir:
+            temp_file = Path(temp_dir, TRACE_DATA_FILE_NAME)
+            try:
+                self._download_file(TRACE_DATA_FILE_NAME, temp_file)
+            except Exception as e:
+                # `MlflowTraceDataNotFound` is caught in `TrackingServiceClient.search_traces` and
+                # is used to filter out traces with failed trace data download.
+                raise MlflowTraceDataNotFound(artifact_path=TRACE_DATA_FILE_NAME) from e
+            return try_read_trace_data(temp_file)
+
+    def upload_trace_data(self, trace_data: str) -> None:
+        """
+        Upload the trace data.
+
+        Args:
+            trace_data: The json-serialized trace data to upload.
+        """
+        with write_local_temp_trace_data_file(trace_data) as temp_file:
+            self.log_artifact(temp_file)
+
+
+@contextmanager
+def write_local_temp_trace_data_file(trace_data: str):
+    with tempfile.TemporaryDirectory() as temp_dir:
+        temp_file = Path(temp_dir, TRACE_DATA_FILE_NAME)
+        temp_file.write_text(trace_data, encoding="utf-8")
+        yield temp_file
+
+
+def try_read_trace_data(trace_data_path):
+    if not os.path.exists(trace_data_path):
+        raise MlflowTraceDataNotFound(artifact_path=trace_data_path)
+    with open(trace_data_path, encoding="utf-8") as f:
+        data = f.read()
+    if not data:
+        raise MlflowTraceDataNotFound(artifact_path=trace_data_path)
+    try:
+        return json.loads(data)
+    except json.decoder.JSONDecodeError as e:
+        raise MlflowTraceDataCorrupted(artifact_path=trace_data_path) from e
+
+
+class MultipartUploadMixin(ABC):
+    @abstractmethod
+    def create_multipart_upload(
+        self, local_file: str, num_parts: int, artifact_path: Optional[str] = None
+    ) -> CreateMultipartUploadResponse:
+        """
+        Initiate a multipart upload and retrieve the pre-signed upload URLS and upload id.
+
+        Args:
+            local_file: Path of artifact to upload.
+            num_parts: Number of parts to upload. Only required by S3 and GCS.
+            artifact_path: Directory within the run's artifact directory in which to upload the
+                artifact.
+
+        """
+
+    @abstractmethod
+    def complete_multipart_upload(
+        self,
+        local_file: str,
+        upload_id: str,
+        parts: list[MultipartUploadPart],
+        artifact_path: Optional[str] = None,
+    ) -> None:
+        """
+        Complete a multipart upload.
+
+        Args:
+            local_file: Path of artifact to upload.
+            upload_id: The upload ID. Only required by S3 and GCS.
+            parts: A list containing the metadata of each part that has been uploaded.
+            artifact_path: Directory within the run's artifact directory in which to upload the
+                artifact.
+
+        """
+
+    @abstractmethod
+    def abort_multipart_upload(
+        self,
+        local_file: str,
+        upload_id: str,
+        artifact_path: Optional[str] = None,
+    ) -> None:
+        """
+        Abort a multipart upload.
+
+        Args:
+            local_file: Path of artifact to upload.
+            upload_id: The upload ID. Only required by S3 and GCS.
+            artifact_path: Directory within the run's artifact directory in which to upload the
+                artifact.
+
+        """
+
+
+def verify_artifact_path(artifact_path):
+    if artifact_path and path_not_unique(artifact_path):
+        raise MlflowException(
+            f"Invalid artifact path: '{artifact_path}'. {bad_path_message(artifact_path)}"
+        )

venv/lib/python3.9/site-packages/mlflow/store/artifact/artifact_repository_registry.py

@@ -0,0 +1,143 @@
+import warnings
+
+from mlflow.exceptions import MlflowException
+from mlflow.store.artifact.artifact_repo import ArtifactRepository
+from mlflow.store.artifact.azure_blob_artifact_repo import AzureBlobArtifactRepository
+from mlflow.store.artifact.azure_data_lake_artifact_repo import AzureDataLakeArtifactRepository
+from mlflow.store.artifact.dbfs_artifact_repo import dbfs_artifact_repo_factory
+from mlflow.store.artifact.ftp_artifact_repo import FTPArtifactRepository
+from mlflow.store.artifact.gcs_artifact_repo import GCSArtifactRepository
+from mlflow.store.artifact.hdfs_artifact_repo import HdfsArtifactRepository
+from mlflow.store.artifact.http_artifact_repo import HttpArtifactRepository
+from mlflow.store.artifact.local_artifact_repo import LocalArtifactRepository
+from mlflow.store.artifact.mlflow_artifacts_repo import MlflowArtifactsRepository
+from mlflow.store.artifact.models_artifact_repo import ModelsArtifactRepository
+from mlflow.store.artifact.r2_artifact_repo import R2ArtifactRepository
+from mlflow.store.artifact.runs_artifact_repo import RunsArtifactRepository
+from mlflow.store.artifact.s3_artifact_repo import S3ArtifactRepository
+from mlflow.store.artifact.sftp_artifact_repo import SFTPArtifactRepository
+from mlflow.store.artifact.uc_volume_artifact_repo import uc_volume_artifact_repo_factory
+from mlflow.utils.plugins import get_entry_points
+from mlflow.utils.uri import get_uri_scheme, is_uc_volumes_uri
+
+
+class ArtifactRepositoryRegistry:
+    """Scheme-based registry for artifact repository implementations
+
+    This class allows the registration of a function or class to provide an implementation for a
+    given scheme of `artifact_uri` through the `register` method. Implementations declared though
+    the entrypoints `mlflow.artifact_repository` group can be automatically registered through the
+    `register_entrypoints` method.
+
+    When instantiating an artifact repository through the `get_artifact_repository` method, the
+    scheme of the artifact URI provided will be used to select which implementation to instantiate,
+    which will be called with same arguments passed to the `get_artifact_repository` method.
+    """
+
+    def __init__(self):
+        self._registry = {}
+
+    def register(self, scheme, repository):
+        """Register artifact repositories provided by other packages"""
+        self._registry[scheme] = repository
+
+    def register_entrypoints(self):
+        # Register artifact repositories provided by other packages
+        for entrypoint in get_entry_points("mlflow.artifact_repository"):
+            try:
+                self.register(entrypoint.name, entrypoint.load())
+            except (AttributeError, ImportError) as exc:
+                warnings.warn(
+                    'Failure attempting to register artifact repository for scheme "{}": {}'.format(
+                        entrypoint.name, str(exc)
+                    ),
+                    stacklevel=2,
+                )
+
+    def get_artifact_repository(self, artifact_uri):
+        """
+        Get an artifact repository from the registry based on the scheme of artifact_uri
+
+        Args:
+            artifact_uri: The artifact store URI. This URI is used to select which artifact
+                repository implementation to instantiate and is passed to the constructor of the
+                implementation.
+
+        Returns:
+            An instance of `mlflow.store.ArtifactRepository` that fulfills the artifact URI
+            requirements.
+        """
+        scheme = get_uri_scheme(artifact_uri)
+        repository = self._registry.get(scheme)
+        if repository is None:
+            raise MlflowException(
+                f"Could not find a registered artifact repository for: {artifact_uri}. "
+                f"Currently registered schemes are: {list(self._registry.keys())}"
+            )
+        return repository(artifact_uri)
+
+    def get_registered_artifact_repositories(self):
+        """
+        Get all registered artifact repositories.
+
+        Returns:
+            A dictionary mapping string artifact URI schemes to artifact repositories.
+        """
+        return self._registry
+
+
+def _dbfs_artifact_repo_factory(artifact_uri: str) -> ArtifactRepository:
+    return (
+        uc_volume_artifact_repo_factory(artifact_uri)
+        if is_uc_volumes_uri(artifact_uri)
+        else dbfs_artifact_repo_factory(artifact_uri)
+    )
+
+
+_artifact_repository_registry = ArtifactRepositoryRegistry()
+
+_artifact_repository_registry.register("", LocalArtifactRepository)
+_artifact_repository_registry.register("file", LocalArtifactRepository)
+_artifact_repository_registry.register("s3", S3ArtifactRepository)
+_artifact_repository_registry.register("r2", R2ArtifactRepository)
+_artifact_repository_registry.register("gs", GCSArtifactRepository)
+_artifact_repository_registry.register("wasbs", AzureBlobArtifactRepository)
+_artifact_repository_registry.register("ftp", FTPArtifactRepository)
+_artifact_repository_registry.register("sftp", SFTPArtifactRepository)
+_artifact_repository_registry.register("dbfs", _dbfs_artifact_repo_factory)
+_artifact_repository_registry.register("hdfs", HdfsArtifactRepository)
+_artifact_repository_registry.register("viewfs", HdfsArtifactRepository)
+_artifact_repository_registry.register("runs", RunsArtifactRepository)
+_artifact_repository_registry.register("models", ModelsArtifactRepository)
+for scheme in ["http", "https"]:
+    _artifact_repository_registry.register(scheme, HttpArtifactRepository)
+_artifact_repository_registry.register("mlflow-artifacts", MlflowArtifactsRepository)
+_artifact_repository_registry.register("abfss", AzureDataLakeArtifactRepository)
+
+_artifact_repository_registry.register_entrypoints()
+
+
+def get_artifact_repository(artifact_uri: str) -> ArtifactRepository:
+    """
+    Get an artifact repository from the registry based on the scheme of artifact_uri
+
+    Args:
+        artifact_uri: The artifact store URI. This URI is used to select which artifact
+            repository implementation to instantiate and is passed to the constructor of the
+            implementation.
+
+    Returns:
+        An instance of `mlflow.store.ArtifactRepository` that fulfills the artifact URI
+        requirements.
+    """
+    return _artifact_repository_registry.get_artifact_repository(artifact_uri)
+
+
+def get_registered_artifact_repositories() -> dict[str, ArtifactRepository]:
+    """
+    Get all registered artifact repositories.
+
+    Returns:
+        A dictionary mapping string artifact URI schemes to artifact repositories.
+    """
+    return _artifact_repository_registry.get_registered_artifact_repositories()

venv/lib/python3.9/site-packages/mlflow/store/artifact/azure_blob_artifact_repo.py

@@ -0,0 +1,275 @@
+import base64
+import datetime
+import os
+import posixpath
+import re
+import urllib.parse
+from typing import Union
+
+from mlflow.entities import FileInfo
+from mlflow.entities.multipart_upload import (
+    CreateMultipartUploadResponse,
+    MultipartUploadCredential,
+)
+from mlflow.environment_variables import MLFLOW_ARTIFACT_UPLOAD_DOWNLOAD_TIMEOUT
+from mlflow.exceptions import MlflowException
+from mlflow.store.artifact.artifact_repo import ArtifactRepository, MultipartUploadMixin
+from mlflow.utils.credentials import get_default_host_creds
+
+
+def encode_base64(data: Union[str, bytes]) -> str:
+    if isinstance(data, str):
+        data = data.encode("utf-8")
+    encoded = base64.b64encode(data)
+    return encoded.decode("utf-8")
+
+
+def decode_base64(encoded: str) -> str:
+    decoded_bytes = base64.b64decode(encoded)
+    return decoded_bytes.decode("utf-8")
+
+
+class AzureBlobArtifactRepository(ArtifactRepository, MultipartUploadMixin):
+    """
+    Stores artifacts on Azure Blob Storage.
+
+    This repository is used with URIs of the form
+    ``wasbs://<container-name>@<ystorage-account-name>.blob.core.windows.net/<path>``,
+    following the same URI scheme as Hadoop on Azure blob storage. It requires either that:
+    - Azure storage connection string is in the env var ``AZURE_STORAGE_CONNECTION_STRING``
+    - Azure storage access key is in the env var ``AZURE_STORAGE_ACCESS_KEY``
+    - DefaultAzureCredential is configured
+    """
+
+    def __init__(self, artifact_uri, client=None):
+        super().__init__(artifact_uri)
+
+        _DEFAULT_TIMEOUT = 600  # 10 minutes
+        self.write_timeout = MLFLOW_ARTIFACT_UPLOAD_DOWNLOAD_TIMEOUT.get() or _DEFAULT_TIMEOUT
+
+        # Allow override for testing
+        if client:
+            self.client = client
+            return
+
+        from azure.storage.blob import BlobServiceClient
+
+        (_, account, _, api_uri_suffix) = AzureBlobArtifactRepository.parse_wasbs_uri(artifact_uri)
+        if "AZURE_STORAGE_CONNECTION_STRING" in os.environ:
+            self.client = BlobServiceClient.from_connection_string(
+                conn_str=os.environ.get("AZURE_STORAGE_CONNECTION_STRING"),
+                connection_verify=get_default_host_creds(artifact_uri).verify,
+            )
+        elif "AZURE_STORAGE_ACCESS_KEY" in os.environ:
+            account_url = f"https://{account}.{api_uri_suffix}"
+            self.client = BlobServiceClient(
+                account_url=account_url,
+                credential=os.environ.get("AZURE_STORAGE_ACCESS_KEY"),
+                connection_verify=get_default_host_creds(artifact_uri).verify,
+            )
+        else:
+            try:
+                from azure.identity import DefaultAzureCredential
+            except ImportError as exc:
+                raise ImportError(
+                    "Using DefaultAzureCredential requires the azure-identity package. "
+                    "Please install it via: pip install azure-identity"
+                ) from exc
+
+            account_url = f"https://{account}.{api_uri_suffix}"
+            self.client = BlobServiceClient(
+                account_url=account_url,
+                credential=DefaultAzureCredential(),
+                connection_verify=get_default_host_creds(artifact_uri).verify,
+            )
+
+    @staticmethod
+    def parse_wasbs_uri(uri):
+        """Parse a wasbs:// URI, returning (container, storage_account, path, api_uri_suffix)."""
+        parsed = urllib.parse.urlparse(uri)
+        if parsed.scheme != "wasbs":
+            raise Exception(f"Not a WASBS URI: {uri}")
+
+        match = re.match(
+            r"([^@]+)@([^.]+)\.(blob\.core\.(windows\.net|chinacloudapi\.cn))", parsed.netloc
+        )
+
+        if match is None:
+            raise Exception(
+                "WASBS URI must be of the form "
+                "<container>@<account>.blob.core.windows.net"
+                " or <container>@<account>.blob.core.chinacloudapi.cn"
+            )
+        container = match.group(1)
+        storage_account = match.group(2)
+        api_uri_suffix = match.group(3)
+        path = parsed.path
+        if path.startswith("/"):
+            path = path[1:]
+        return container, storage_account, path, api_uri_suffix
+
+    def log_artifact(self, local_file, artifact_path=None):
+        (container, _, dest_path, _) = self.parse_wasbs_uri(self.artifact_uri)
+        container_client = self.client.get_container_client(container)
+        if artifact_path:
+            dest_path = posixpath.join(dest_path, artifact_path)
+        dest_path = posixpath.join(dest_path, os.path.basename(local_file))
+        with open(local_file, "rb") as file:
+            container_client.upload_blob(
+                dest_path, file, overwrite=True, timeout=self.write_timeout
+            )
+
+    def log_artifacts(self, local_dir, artifact_path=None):
+        (container, _, dest_path, _) = self.parse_wasbs_uri(self.artifact_uri)
+        container_client = self.client.get_container_client(container)
+        if artifact_path:
+            dest_path = posixpath.join(dest_path, artifact_path)
+        local_dir = os.path.abspath(local_dir)
+        for root, _, filenames in os.walk(local_dir):
+            upload_path = dest_path
+            if root != local_dir:
+                rel_path = os.path.relpath(root, local_dir)
+                upload_path = posixpath.join(dest_path, rel_path)
+            for f in filenames:
+                remote_file_path = posixpath.join(upload_path, f)
+                local_file_path = os.path.join(root, f)
+                with open(local_file_path, "rb") as file:
+                    container_client.upload_blob(
+                        remote_file_path, file, overwrite=True, timeout=self.write_timeout
+                    )
+
+    def list_artifacts(self, path=None):
+        # Newer versions of `azure-storage-blob` (>= 12.4.0) provide a public
+        # `azure.storage.blob.BlobPrefix` object to signify that a blob is a directory,
+        # while older versions only expose this API internally as
+        # `azure.storage.blob._models.BlobPrefix`
+        try:
+            from azure.storage.blob import BlobPrefix
+        except ImportError:
+            from azure.storage.blob._models import BlobPrefix
+
+        def is_dir(result):
+            return isinstance(result, BlobPrefix)
+
+        (container, _, artifact_path, _) = self.parse_wasbs_uri(self.artifact_uri)
+        container_client = self.client.get_container_client(container)
+        dest_path = artifact_path
+        if path:
+            dest_path = posixpath.join(dest_path, path)
+        infos = []
+        prefix = dest_path if dest_path.endswith("/") else dest_path + "/"
+        results = container_client.walk_blobs(name_starts_with=prefix)
+
+        for result in results:
+            if (
+                dest_path == result.name
+            ):  # result isn't actually a child of the path we're interested in, so skip it
+                continue
+
+            if not result.name.startswith(artifact_path):
+                raise MlflowException(
+                    "The name of the listed Azure blob does not begin with the specified"
+                    f" artifact path. Artifact path: {artifact_path}. Blob name: {result.name}"
+                )
+
+            if is_dir(result):
+                subdir = posixpath.relpath(path=result.name, start=artifact_path)
+                if subdir.endswith("/"):
+                    subdir = subdir[:-1]
+                infos.append(FileInfo(subdir, is_dir=True, file_size=None))
+            else:  # Just a plain old blob
+                file_name = posixpath.relpath(path=result.name, start=artifact_path)
+                infos.append(FileInfo(file_name, is_dir=False, file_size=result.size))
+
+        # The list_artifacts API expects us to return an empty list if the
+        # the path references a single file.
+        rel_path = dest_path[len(artifact_path) + 1 :]
+        if (len(infos) == 1) and not infos[0].is_dir and (infos[0].path == rel_path):
+            return []
+        return sorted(infos, key=lambda f: f.path)
+
+    def _download_file(self, remote_file_path, local_path):
+        (container, _, remote_root_path, _) = self.parse_wasbs_uri(self.artifact_uri)
+        container_client = self.client.get_container_client(container)
+        remote_full_path = posixpath.join(remote_root_path, remote_file_path)
+        blob = container_client.download_blob(remote_full_path)
+        with open(local_path, "wb") as file:
+            blob.readinto(file)
+
+    def delete_artifacts(self, artifact_path=None):
+        from azure.core.exceptions import ResourceNotFoundError
+
+        (container, _, dest_path, _) = self.parse_wasbs_uri(self.artifact_uri)
+        container_client = self.client.get_container_client(container)
+        if artifact_path:
+            dest_path = posixpath.join(dest_path, artifact_path)
+
+        try:
+            blobs = container_client.list_blobs(name_starts_with=dest_path)
+            blob_list = list(blobs)
+            if not blob_list:
+                raise MlflowException(f"No such file or directory: '{dest_path}'")
+
+            for blob in blob_list:
+                container_client.delete_blob(blob.name)
+        except ResourceNotFoundError:
+            raise MlflowException(f"No such file or directory: '{dest_path}'")
+
+    def create_multipart_upload(self, local_file, num_parts=1, artifact_path=None):
+        from azure.storage.blob import BlobSasPermissions, generate_blob_sas
+
+        (container, _, dest_path, _) = self.parse_wasbs_uri(self.artifact_uri)
+        if artifact_path:
+            dest_path = posixpath.join(dest_path, artifact_path)
+        dest_path = posixpath.join(dest_path, os.path.basename(local_file))
+
+        # Put Block: https://learn.microsoft.com/en-us/rest/api/storageservices/put-block?tabs=microsoft-entra-id
+        # SDK: https://learn.microsoft.com/en-us/python/api/azure-storage-blob/azure.storage.blob.blobclient?view=azure-python#azure-storage-blob-blobclient-stage-block
+        blob_url = posixpath.join(self.client.url, container, dest_path)
+        sas_token = generate_blob_sas(
+            account_name=self.client.account_name,
+            container_name=container,
+            blob_name=dest_path,
+            account_key=self.client.credential.account_key,
+            permission=BlobSasPermissions(read=True, write=True),
+            expiry=datetime.datetime.utcnow() + datetime.timedelta(hours=1),
+        )
+        credentials = []
+        for i in range(1, num_parts + 1):
+            block_id = f"mlflow_block_{i}"
+            # see https://github.com/Azure/azure-sdk-for-python/blob/18a66ef98c6f2153491489d3d7d2fe4a5849e4ac/sdk/storage/azure-storage-blob/azure/storage/blob/_blob_client.py#L2468
+            safe_block_id = urllib.parse.quote(encode_base64(block_id), safe="")
+            url = f"{blob_url}?comp=block&blockid={safe_block_id}&{sas_token}"
+            credentials.append(
+                MultipartUploadCredential(
+                    url=url,
+                    part_number=i,
+                    headers={},
+                )
+            )
+        return CreateMultipartUploadResponse(
+            credentials=credentials,
+            upload_id=None,
+        )
+
+    def complete_multipart_upload(self, local_file, upload_id, parts=None, artifact_path=None):
+        (container, _, dest_path, _) = self.parse_wasbs_uri(self.artifact_uri)
+        if artifact_path:
+            dest_path = posixpath.join(dest_path, artifact_path)
+        dest_path = posixpath.join(dest_path, os.path.basename(local_file))
+
+        block_ids = []
+        for part in parts:
+            qs = urllib.parse.urlparse(part.url).query
+            block_id = urllib.parse.parse_qs(qs)["blockid"][0]
+            block_id = decode_base64(urllib.parse.unquote(block_id))
+            block_ids.append(block_id)
+        blob_client = self.client.get_blob_client(container, dest_path)
+        blob_client.commit_block_list(block_ids)
+
+    def abort_multipart_upload(self, local_file, upload_id, artifact_path=None):
+        # There is no way to delete uncommitted blocks in Azure Blob Storage.
+        # Instead, they are garbage collected within 7 days.
+        # See https://docs.microsoft.com/en-us/rest/api/storageservices/put-block-list#remarks
+        # The blob may already exist so we cannot delete it either.
+        pass

venv/lib/python3.9/site-packages/mlflow/store/artifact/azure_data_lake_artifact_repo.py

@@ -0,0 +1,293 @@
+import os
+import posixpath
+import re
+import urllib.parse
+
+import requests
+
+from mlflow.azure.client import patch_adls_file_upload, patch_adls_flush, put_adls_file_creation
+from mlflow.entities import FileInfo
+from mlflow.environment_variables import (
+    MLFLOW_ARTIFACT_UPLOAD_DOWNLOAD_TIMEOUT,
+    MLFLOW_ENABLE_MULTIPART_UPLOAD,
+    MLFLOW_MULTIPART_UPLOAD_CHUNK_SIZE,
+)
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_artifacts_pb2 import ArtifactCredentialInfo
+from mlflow.store.artifact.artifact_repo import _retry_with_new_creds
+from mlflow.store.artifact.cloud_artifact_repo import (
+    CloudArtifactRepository,
+    _complete_futures,
+    _compute_num_chunks,
+)
+
+
+def _parse_abfss_uri(uri):
+    """
+    Parse an ABFSS URI in the format
+    "abfss://<file_system>@<account_name>.<domain_suffix>/<path>",
+    returning a tuple consisting of the filesystem, account name, domain suffix, and path
+
+    See more details about ABFSS URIs at
+    https://learn.microsoft.com/en-us/azure/storage/blobs/data-lake-storage-abfs-driver#uri-scheme-to-reference-data.
+    Also, see different domain suffixes for:
+    * Azure China: https://learn.microsoft.com/en-us/azure/china/resources-developer-guide
+    * Azure Government: https://learn.microsoft.com/en-us/azure/azure-government/compare-azure-government-global-azure#guidance-for-developers
+    * Azure Private Link: https://learn.microsoft.com/en-us/azure/private-link/private-endpoint-dns#government
+    Args:
+        uri: ABFSS URI to parse
+
+    Returns:
+        A tuple containing the name of the filesystem, account name, domain suffix,
+        and path
+    """
+    parsed = urllib.parse.urlparse(uri)
+    if parsed.scheme != "abfss":
+        raise MlflowException(f"Not an ABFSS URI: {uri}")
+
+    match = re.match(r"([^@]+)@([^.]+)\.(.*)", parsed.netloc)
+
+    if match is None:
+        raise MlflowException(
+            "ABFSS URI must be of the form abfss://<filesystem>@<account>.<domain_suffix>"
+        )
+    filesystem = match.group(1)
+    account_name = match.group(2)
+    domain_suffix = match.group(3)
+    path = parsed.path
+    if path.startswith("/"):
+        path = path[1:]
+    return filesystem, account_name, domain_suffix, path
+
+
+def _get_data_lake_client(account_url, credential):
+    from azure.storage.filedatalake import DataLakeServiceClient
+
+    return DataLakeServiceClient(account_url, credential)
+
+
+class AzureDataLakeArtifactRepository(CloudArtifactRepository):
+    """
+    Stores artifacts on Azure Data Lake Storage Gen2.
+
+    This repository is used with URIs of the form
+    ``abfs[s]://file_system@account_name.dfs.core.windows.net/<path>/<path>``.
+
+    Args
+        credential: Azure credential (see options in https://learn.microsoft.com/en-us/python/api/azure-core/azure.core.credentials?view=azure-python)
+            to use to authenticate to storage
+    """
+
+    def __init__(
+        self,
+        artifact_uri,
+        credential=None,
+        credential_refresh_def=None,
+    ):
+        super().__init__(artifact_uri)
+        _DEFAULT_TIMEOUT = 600  # 10 minutes
+        self.write_timeout = MLFLOW_ARTIFACT_UPLOAD_DOWNLOAD_TIMEOUT.get() or _DEFAULT_TIMEOUT
+        self._parse_credentials(credential)
+        self._credential_refresh_def = credential_refresh_def
+
+    def _parse_credentials(self, credential):
+        (filesystem, account_name, domain_suffix, path) = _parse_abfss_uri(self.artifact_uri)
+        account_url = f"https://{account_name}.{domain_suffix}"
+        self.sas_token = ""
+        if credential is None:
+            if sas_token := os.environ.get("AZURE_STORAGE_SAS_TOKEN"):
+                self.sas_token = f"?{sas_token}"
+                account_url += self.sas_token
+            else:
+                from azure.identity import DefaultAzureCredential
+
+                credential = DefaultAzureCredential()
+        self.credential = credential
+        data_lake_client = _get_data_lake_client(
+            account_url=account_url, credential=self.credential
+        )
+        self.fs_client = data_lake_client.get_file_system_client(filesystem)
+        self.domain_suffix = domain_suffix
+        self.base_data_lake_directory = path
+        self.account_name = account_name
+        self.container = filesystem
+
+    def _refresh_credentials(self):
+        if not self._credential_refresh_def:
+            return self.fs_client
+        new_creds = self._credential_refresh_def()
+        self._parse_credentials(new_creds["credential"])
+        return self.fs_client
+
+    def log_artifact(self, local_file, artifact_path=None):
+        dest_path = self.base_data_lake_directory
+        if artifact_path:
+            dest_path = posixpath.join(dest_path, artifact_path)
+        local_file_path = os.path.abspath(local_file)
+        file_name = os.path.basename(local_file_path)
+
+        def try_func(creds):
+            dir_client = creds.get_directory_client(dest_path)
+            file_client = dir_client.get_file_client(file_name)
+            if os.path.getsize(local_file_path) == 0:
+                file_client.create_file()
+            else:
+                with open(local_file_path, "rb") as file:
+                    file_client.upload_data(data=file, overwrite=True)
+
+        _retry_with_new_creds(
+            try_func=try_func, creds_func=self._refresh_credentials, orig_creds=self.fs_client
+        )
+
+    def list_artifacts(self, path=None):
+        directory_to_list = self.base_data_lake_directory
+        if path:
+            directory_to_list = posixpath.join(directory_to_list, path)
+        infos = []
+        for result in self.fs_client.get_paths(path=directory_to_list, recursive=False):
+            if (
+                directory_to_list == result.name
+            ):  # result isn't actually a child of the path we're interested in, so skip it
+                continue
+            if result.is_directory:
+                subdir = posixpath.relpath(path=result.name, start=self.base_data_lake_directory)
+                if subdir.endswith("/"):
+                    subdir = subdir[:-1]
+                infos.append(FileInfo(subdir, is_dir=True, file_size=None))
+            else:
+                file_name = posixpath.relpath(path=result.name, start=self.base_data_lake_directory)
+                infos.append(FileInfo(file_name, is_dir=False, file_size=result.content_length))
+
+        # The list_artifacts API expects us to return an empty list if the
+        # the path references a single file.
+        rel_path = directory_to_list[len(self.base_data_lake_directory) + 1 :]
+        if (len(infos) == 1) and not infos[0].is_dir and (infos[0].path == rel_path):
+            return []
+        return sorted(infos, key=lambda f: f.path)
+
+    def _download_from_cloud(self, remote_file_path, local_path):
+        remote_full_path = posixpath.join(self.base_data_lake_directory, remote_file_path)
+        base_dir = posixpath.dirname(remote_full_path)
+
+        def try_func(creds):
+            dir_client = creds.get_directory_client(base_dir)
+            filename = posixpath.basename(remote_full_path)
+            file_client = dir_client.get_file_client(filename)
+            with open(local_path, "wb") as file:
+                file_client.download_file().readinto(file)
+
+        _retry_with_new_creds(
+            try_func=try_func, creds_func=self._refresh_credentials, orig_creds=self.fs_client
+        )
+
+    def delete_artifacts(self, artifact_path=None):
+        raise NotImplementedError("This artifact repository does not support deleting artifacts")
+
+    def _upload_to_cloud(self, cloud_credential_info, src_file_path, artifact_file_path):
+        if (
+            MLFLOW_ENABLE_MULTIPART_UPLOAD.get()
+            and os.path.getsize(src_file_path) > MLFLOW_MULTIPART_UPLOAD_CHUNK_SIZE.get()
+        ):
+            self._multipart_upload(cloud_credential_info, src_file_path, artifact_file_path)
+        else:
+            artifact_subdir = posixpath.dirname(artifact_file_path)
+            self.log_artifact(src_file_path, artifact_subdir)
+
+    def _retryable_adls_function(self, func, artifact_file_path, **kwargs):
+        # Attempt to call the passed function.  Retry if the credentials have expired
+        try:
+            func(**kwargs)
+        except requests.HTTPError as e:
+            if e.response.status_code in [403]:
+                new_credentials = self._get_write_credential_infos([artifact_file_path])[0]
+                kwargs["sas_url"] = new_credentials.signed_uri
+                func(**kwargs)
+            else:
+                raise e
+
+    def _multipart_upload(self, credentials, src_file_path, artifact_file_path):
+        """
+        Uploads a file to a given Azure storage location using the ADLS gen2 API.
+        """
+        try:
+            headers = self._extract_headers_from_credentials(credentials.headers)
+            # try to create the file
+            self._retryable_adls_function(
+                func=put_adls_file_creation,
+                artifact_file_path=artifact_file_path,
+                sas_url=credentials.signed_uri,
+                headers=headers,
+            )
+            # next try to append the file
+            futures = {}
+            file_size = os.path.getsize(src_file_path)
+            num_chunks = _compute_num_chunks(
+                src_file_path, MLFLOW_MULTIPART_UPLOAD_CHUNK_SIZE.get()
+            )
+            use_single_part_upload = num_chunks == 1
+            for index in range(num_chunks):
+                start_byte = index * MLFLOW_MULTIPART_UPLOAD_CHUNK_SIZE.get()
+                future = self.chunk_thread_pool.submit(
+                    self._retryable_adls_function,
+                    func=patch_adls_file_upload,
+                    artifact_file_path=artifact_file_path,
+                    sas_url=credentials.signed_uri,
+                    local_file=src_file_path,
+                    start_byte=start_byte,
+                    size=MLFLOW_MULTIPART_UPLOAD_CHUNK_SIZE.get(),
+                    position=start_byte,
+                    headers=headers,
+                    is_single=use_single_part_upload,
+                )
+                futures[future] = index
+
+            _, errors = _complete_futures(futures, src_file_path)
+            if errors:
+                raise MlflowException(
+                    f"Failed to upload at least one part of {artifact_file_path}. Errors: {errors}"
+                )
+
+            # finally try to flush the file
+            if not use_single_part_upload:
+                self._retryable_adls_function(
+                    func=patch_adls_flush,
+                    artifact_file_path=artifact_file_path,
+                    sas_url=credentials.signed_uri,
+                    position=file_size,
+                    headers=headers,
+                )
+        except Exception as err:
+            raise MlflowException(err)
+
+    def _get_presigned_uri(self, artifact_file_path):
+        """
+        Gets the presigned URL required to upload a file to or download a file from a given Azure
+        storage location.
+
+        Args:
+            artifact_file_path: Path of the file relative to the artifact repository root.
+
+        Returns:
+            a string presigned URL.
+        """
+        sas_token = (
+            f"?{self.credential.signature}"
+            if hasattr(self.credential, "signature")
+            else self.sas_token
+        )
+        return (
+            f"https://{self.account_name}.{self.domain_suffix}/{self.container}/"
+            f"{self.base_data_lake_directory}/{artifact_file_path}{sas_token}"
+        )
+
+    def _get_write_credential_infos(self, remote_file_paths) -> list[ArtifactCredentialInfo]:
+        return [
+            ArtifactCredentialInfo(signed_uri=self._get_presigned_uri(path))
+            for path in remote_file_paths
+        ]
+
+    def _get_read_credential_infos(self, remote_file_paths) -> list[ArtifactCredentialInfo]:
+        return [
+            ArtifactCredentialInfo(signed_uri=self._get_presigned_uri(path))
+            for path in remote_file_paths
+        ]

venv/lib/python3.9/site-packages/mlflow/store/artifact/cli.py

@@ -0,0 +1,141 @@
+import logging
+
+import click
+
+from mlflow.artifacts import download_artifacts as _download_artifacts
+from mlflow.store.artifact.artifact_repository_registry import get_artifact_repository
+from mlflow.tracking import _get_store
+from mlflow.utils.proto_json_utils import message_to_json
+
+_logger = logging.getLogger(__name__)
+
+
+@click.group("artifacts")
+def commands():
+    """
+    Upload, list, and download artifacts from an MLflow artifact repository.
+
+    To manage artifacts for a run associated with a tracking server, set the MLFLOW_TRACKING_URI
+    environment variable to the URL of the desired server.
+    """
+
+
+@commands.command("log-artifact")
+@click.option("--local-file", "-l", required=True, help="Local path to artifact to log")
+@click.option("--run-id", "-r", required=True, help="Run ID into which we should log the artifact.")
+@click.option(
+    "--artifact-path",
+    "-a",
+    help="If specified, we will log the artifact into this subdirectory of the "
+    + "run's artifact directory.",
+)
+def log_artifact(local_file, run_id, artifact_path):
+    """
+    Log a local file 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_artifact(local_file, artifact_path)
+    _logger.info(
+        "Logged artifact from local file %s to artifact_path=%s", local_file, artifact_path
+    )
+
+
+@commands.command("log-artifacts")
+@click.option("--local-dir", "-l", required=True, help="Directory of local artifacts to log")
+@click.option("--run-id", "-r", required=True, help="Run ID into which we should log the artifact.")
+@click.option(
+    "--artifact-path",
+    "-a",
+    help="If specified, we will log the artifact into this subdirectory of the "
+    + "run's artifact directory.",
+)
+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)
+
+
+@commands.command("list")
+@click.option("--run-id", "-r", required=True, help="Run ID to be listed")
+@click.option(
+    "--artifact-path",
+    "-a",
+    help="If specified, a path relative to the run's root directory to list.",
+)
+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 _file_infos_to_json(file_infos):
+    json_list = [message_to_json(file_info.to_proto()) for file_info in file_infos]
+    return "[" + ", ".join(json_list) + "]"
+
+
+@commands.command("download")
+@click.option("--run-id", "-r", help="Run ID from which to download")
+@click.option(
+    "--artifact-path",
+    "-a",
+    help="For use with Run ID: if specified, a path relative to the run's root "
+    "directory to download",
+)
+@click.option(
+    "--artifact-uri",
+    "-u",
+    help="URI pointing to the artifact file or artifacts directory; use as an "
+    "alternative to specifying --run_id and --artifact-path",
+)
+@click.option(
+    "--dst-path",
+    "-d",
+    help=(
+        "Path of the local filesystem destination directory to which to download the"
+        " specified artifacts. If the directory does not exist, it is created. If unspecified"
+        " the artifacts are downloaded to a new uniquely-named directory on the local filesystem,"
+        " unless the artifacts already exist on the local filesystem, in which case their local"
+        " path is returned directly"
+    ),
+)
+def download_artifacts(run_id, artifact_path, artifact_uri, dst_path):
+    """
+    Download an artifact file or directory to a local directory.
+    The output is the name of the file or directory on the local filesystem.
+
+    Either ``--artifact-uri`` or ``--run-id`` must be provided.
+    """
+    # Preserve preexisting behavior in MLflow <= 1.24.0 where specifying `artifact_uri` and
+    # `artifact_path` together did not throw an exception (unlike
+    # `mlflow.artifacts.download_artifacts()`) and instead used `artifact_uri` while ignoring
+    # `run_id` and `artifact_path`
+    if artifact_uri is not None:
+        run_id = None
+        artifact_path = None
+
+    downloaded_local_artifact_location = _download_artifacts(
+        artifact_uri=artifact_uri, run_id=run_id, artifact_path=artifact_path, dst_path=dst_path
+    )
+    click.echo(f"\n{downloaded_local_artifact_location}")
+
+
+if __name__ == "__main__":
+    commands()

venv/lib/python3.9/site-packages/mlflow/store/artifact/cloud_artifact_repo.py

@@ -0,0 +1,331 @@
+import logging
+import math
+import os
+import posixpath
+import time
+from abc import abstractmethod
+from collections import namedtuple
+from concurrent.futures import as_completed
+
+from mlflow.environment_variables import (
+    _MLFLOW_MPD_NUM_RETRIES,
+    _MLFLOW_MPD_RETRY_INTERVAL_SECONDS,
+    MLFLOW_ENABLE_MULTIPART_DOWNLOAD,
+    MLFLOW_MULTIPART_DOWNLOAD_CHUNK_SIZE,
+    MLFLOW_MULTIPART_DOWNLOAD_MINIMUM_FILE_SIZE,
+    MLFLOW_MULTIPART_UPLOAD_CHUNK_SIZE,
+)
+from mlflow.exceptions import MlflowException
+from mlflow.store.artifact.artifact_repo import ArtifactRepository
+from mlflow.utils import chunk_list
+from mlflow.utils.file_utils import (
+    ArtifactProgressBar,
+    parallelized_download_file_using_http_uri,
+    relative_path_to_artifact_path,
+    remove_on_error,
+)
+from mlflow.utils.request_utils import download_chunk
+from mlflow.utils.uri import is_fuse_or_uc_volumes_uri
+
+_logger = logging.getLogger(__name__)
+_ARTIFACT_UPLOAD_BATCH_SIZE = (
+    50  # Max number of artifacts for which to fetch write credentials at once.
+)
+_AWS_MIN_CHUNK_SIZE = 5 * 1024**2  # 5 MB is the minimum chunk size for S3 multipart uploads
+_AWS_MAX_CHUNK_SIZE = 5 * 1024**3  # 5 GB is the maximum chunk size for S3 multipart uploads
+
+
+def _readable_size(size: int) -> str:
+    return f"{size / 1024**2:.2f} MB"
+
+
+def _validate_chunk_size_aws(chunk_size: int) -> None:
+    """
+    Validates the specified chunk size in bytes is in valid range for AWS multipart uploads.
+    """
+    if chunk_size < _AWS_MIN_CHUNK_SIZE or chunk_size > _AWS_MAX_CHUNK_SIZE:
+        raise MlflowException(
+            message=(
+                f"Multipart chunk size {_readable_size(chunk_size)} must be in range: "
+                f"{_readable_size(_AWS_MIN_CHUNK_SIZE)} to {_readable_size(_AWS_MAX_CHUNK_SIZE)}."
+            )
+        )
+
+
+def _compute_num_chunks(local_file: os.PathLike, chunk_size: int) -> int:
+    """
+    Computes the number of chunks to use for a multipart upload of the specified file.
+    """
+    return math.ceil(os.path.getsize(local_file) / chunk_size)
+
+
+def _complete_futures(futures_dict, file):
+    """
+    Waits for the completion of all the futures in the given dictionary and returns
+    a tuple of two dictionaries. The first dictionary contains the results of the
+    futures (unordered) and the second contains the errors (unordered) that occurred
+    during the execution of the futures.
+    """
+    results = {}
+    errors = {}
+
+    with ArtifactProgressBar.chunks(
+        os.path.getsize(file),
+        f"Uploading {file}",
+        MLFLOW_MULTIPART_UPLOAD_CHUNK_SIZE.get(),
+    ) as pbar:
+        for future in as_completed(futures_dict):
+            key = futures_dict[future]
+            try:
+                results[key] = future.result()
+                pbar.update()
+            except Exception as e:
+                errors[key] = repr(e)
+
+    return results, errors
+
+
+StagedArtifactUpload = namedtuple(
+    "StagedArtifactUpload",
+    [
+        # Local filesystem path of the source file to upload
+        "src_file_path",
+        # Base artifact URI-relative path specifying the upload destination
+        "artifact_file_path",
+    ],
+)
+
+
+class CloudArtifactRepository(ArtifactRepository):
+    def __init__(self, artifact_uri):
+        super().__init__(artifact_uri)
+        # Use an isolated thread pool executor for chunk uploads/downloads to avoid a deadlock
+        # caused by waiting for a chunk-upload/download task within a file-upload/download task.
+        # See https://superfastpython.com/threadpoolexecutor-deadlock/#Deadlock_1_Submit_and_Wait_for_a_Task_Within_a_Task
+        # for more details
+        self.chunk_thread_pool = self._create_thread_pool()
+
+    # Write APIs
+
+    def log_artifacts(self, local_dir, artifact_path=None):
+        """
+        Parallelized implementation of `log_artifacts`.
+        """
+
+        artifact_path = artifact_path or ""
+
+        staged_uploads = []
+        for dirpath, _, filenames in os.walk(local_dir):
+            artifact_subdir = artifact_path
+            if dirpath != local_dir:
+                rel_path = os.path.relpath(dirpath, local_dir)
+                rel_path = relative_path_to_artifact_path(rel_path)
+                artifact_subdir = posixpath.join(artifact_path, rel_path)
+            for name in filenames:
+                src_file_path = os.path.join(dirpath, name)
+                src_file_name = os.path.basename(src_file_path)
+                staged_uploads.append(
+                    StagedArtifactUpload(
+                        src_file_path=src_file_path,
+                        artifact_file_path=posixpath.join(artifact_subdir, src_file_name),
+                    )
+                )
+
+        # Join futures to ensure that all artifacts have been uploaded prior to returning
+        failed_uploads = {}
+
+        # For each batch of files, upload them in parallel and wait for completion
+        # TODO: change to class method
+        def upload_artifacts_iter():
+            for staged_upload_chunk in chunk_list(staged_uploads, _ARTIFACT_UPLOAD_BATCH_SIZE):
+                write_credential_infos = self._get_write_credential_infos(
+                    remote_file_paths=[
+                        staged_upload.artifact_file_path for staged_upload in staged_upload_chunk
+                    ],
+                )
+
+                inflight_uploads = {}
+                for staged_upload, write_credential_info in zip(
+                    staged_upload_chunk, write_credential_infos
+                ):
+                    upload_future = self.thread_pool.submit(
+                        self._upload_to_cloud,
+                        cloud_credential_info=write_credential_info,
+                        src_file_path=staged_upload.src_file_path,
+                        artifact_file_path=staged_upload.artifact_file_path,
+                    )
+                    inflight_uploads[staged_upload.src_file_path] = upload_future
+
+                yield from inflight_uploads.items()
+
+        with ArtifactProgressBar.files(
+            desc="Uploading artifacts", total=len(staged_uploads)
+        ) as pbar:
+            for src_file_path, upload_future in upload_artifacts_iter():
+                try:
+                    upload_future.result()
+                    pbar.update()
+                except Exception as e:
+                    failed_uploads[src_file_path] = repr(e)
+
+        if len(failed_uploads) > 0:
+            raise MlflowException(
+                message=(
+                    "The following failures occurred while uploading one or more artifacts"
+                    f" to {self.artifact_uri}: {failed_uploads}"
+                )
+            )
+
+    @abstractmethod
+    def _get_write_credential_infos(self, remote_file_paths):
+        """
+        Retrieve write credentials for a batch of remote file paths, including presigned URLs.
+
+        Args:
+            remote_file_paths: List of file paths in the remote artifact repository.
+
+        Returns:
+            List of ArtifactCredentialInfo objects corresponding to each file path.
+        """
+
+    @abstractmethod
+    def _upload_to_cloud(self, cloud_credential_info, src_file_path, artifact_file_path):
+        """
+        Upload a single file to the cloud.
+
+        Args:
+            cloud_credential_info: ArtifactCredentialInfo object with presigned URL for the file.
+            src_file_path: Local source file path for the upload.
+            artifact_file_path: Path in the artifact repository where the artifact will be logged.
+
+        """
+
+    # Read APIs
+
+    def _extract_headers_from_credentials(self, headers):
+        """
+        Returns:
+            A python dictionary of http headers converted from the protobuf credentials.
+        """
+        return {header.name: header.value for header in headers}
+
+    def _parallelized_download_from_cloud(self, file_size, remote_file_path, local_path):
+        read_credentials = self._get_read_credential_infos([remote_file_path])
+        # Read credentials for only one file were requested. So we expected only one value in
+        # the response.
+        assert len(read_credentials) == 1
+        cloud_credential_info = read_credentials[0]
+
+        with remove_on_error(local_path):
+            parallel_download_subproc_env = os.environ.copy()
+            failed_downloads = parallelized_download_file_using_http_uri(
+                thread_pool_executor=self.chunk_thread_pool,
+                http_uri=cloud_credential_info.signed_uri,
+                download_path=local_path,
+                remote_file_path=remote_file_path,
+                file_size=file_size,
+                uri_type=cloud_credential_info.type,
+                chunk_size=MLFLOW_MULTIPART_DOWNLOAD_CHUNK_SIZE.get(),
+                env=parallel_download_subproc_env,
+                headers=self._extract_headers_from_credentials(cloud_credential_info.headers),
+            )
+            num_retries = _MLFLOW_MPD_NUM_RETRIES.get()
+            interval = _MLFLOW_MPD_RETRY_INTERVAL_SECONDS.get()
+            failed_downloads = list(failed_downloads)
+            while failed_downloads and num_retries > 0:
+                self._refresh_credentials()
+                new_cloud_creds = self._get_read_credential_infos([remote_file_path])[0]
+                new_signed_uri = new_cloud_creds.signed_uri
+                new_headers = self._extract_headers_from_credentials(new_cloud_creds.headers)
+
+                futures = {
+                    self.chunk_thread_pool.submit(
+                        download_chunk,
+                        range_start=chunk.start,
+                        range_end=chunk.end,
+                        headers=new_headers,
+                        download_path=local_path,
+                        http_uri=new_signed_uri,
+                    ): chunk
+                    for chunk in failed_downloads
+                }
+
+                new_failed_downloads = []
+
+                for future in as_completed(futures):
+                    chunk = futures[future]
+                    try:
+                        future.result()
+                    except Exception as e:
+                        _logger.info(
+                            f"Failed to download chunk {chunk.index} for {chunk.path}: {e}. "
+                            f"The download of this chunk will be retried later."
+                        )
+                        new_failed_downloads.append(chunk)
+
+                failed_downloads = new_failed_downloads
+                num_retries -= 1
+                time.sleep(interval)
+
+            if failed_downloads:
+                raise MlflowException(
+                    message=("All retries have been exhausted. Download has failed.")
+                )
+
+    def _download_file(self, remote_file_path, local_path):
+        # list_artifacts API only returns a list of FileInfos at the specified path
+        # if it's a directory. To get file size, we need to iterate over FileInfos
+        # contained by the parent directory. A bad path could result in there being
+        # no matching FileInfos (by path), so fall back to None size to prevent
+        # parallelized download.
+        parent_dir = posixpath.dirname(remote_file_path)
+        file_infos = self.list_artifacts(parent_dir)
+        file_info = [info for info in file_infos if info.path == remote_file_path]
+        file_size = file_info[0].file_size if len(file_info) == 1 else None
+
+        # NB: FUSE mounts do not support file write from a non-0th index seek position.
+        # Due to this limitation (writes must start at the beginning of a file),
+        # offset writes are disabled if FUSE is the local_path destination.
+        if (
+            not MLFLOW_ENABLE_MULTIPART_DOWNLOAD.get()
+            or not file_size
+            or file_size < MLFLOW_MULTIPART_DOWNLOAD_MINIMUM_FILE_SIZE.get()
+            or is_fuse_or_uc_volumes_uri(local_path)
+            # DatabricksSDKModelsArtifactRepository can only download file via databricks sdk
+            # rather than presigned uri used in _parallelized_download_from_cloud.
+            or type(self).__name__ == "DatabricksSDKModelsArtifactRepository"
+        ):
+            self._download_from_cloud(remote_file_path, local_path)
+        else:
+            self._parallelized_download_from_cloud(file_size, remote_file_path, local_path)
+
+    @abstractmethod
+    def _get_read_credential_infos(self, remote_file_paths):
+        """
+        Retrieve read credentials for a batch of remote file paths, including presigned URLs.
+
+        Args:
+            remote_file_paths: List of file paths in the remote artifact repository.
+
+        Returns:
+            List of ArtifactCredentialInfo objects corresponding to each file path.
+        """
+
+    @abstractmethod
+    def _download_from_cloud(self, remote_file_path, local_path):
+        """
+        Download a file from the input `remote_file_path` and save it to `local_path`.
+
+        Args:
+            remote_file_path: Path to file in the remote artifact repository.
+            local_path: Local path to download file to.
+
+        """
+
+    def _refresh_credentials(self):
+        """
+        Refresh credentials for user in the case of credential expiration
+
+        Args:
+            None
+        """

venv/lib/python3.9/site-packages/mlflow/store/artifact/databricks_artifact_repo.py

@@ -0,0 +1,761 @@
+import base64
+import json
+import logging
+import os
+import posixpath
+import uuid
+from typing import Any, Optional
+
+import requests
+
+import mlflow.tracking
+from mlflow.azure.client import (
+    patch_adls_file_upload,
+    patch_adls_flush,
+    put_adls_file_creation,
+    put_block,
+    put_block_list,
+)
+from mlflow.entities import FileInfo
+from mlflow.environment_variables import (
+    MLFLOW_MULTIPART_DOWNLOAD_CHUNK_SIZE,
+    MLFLOW_MULTIPART_UPLOAD_CHUNK_SIZE,
+    MLFLOW_MULTIPART_UPLOAD_MINIMUM_FILE_SIZE,
+)
+from mlflow.exceptions import (
+    MlflowException,
+    MlflowTraceDataCorrupted,
+    MlflowTraceDataNotFound,
+)
+from mlflow.protos.databricks_artifacts_pb2 import (
+    ArtifactCredentialType,
+    CompleteMultipartUpload,
+    CreateMultipartUpload,
+    DatabricksMlflowArtifactsService,
+    GetCredentialsForRead,
+    GetCredentialsForTraceDataDownload,
+    GetCredentialsForTraceDataUpload,
+    GetCredentialsForWrite,
+    GetPresignedUploadPartUrl,
+    PartEtag,
+)
+from mlflow.protos.databricks_pb2 import (
+    INTERNAL_ERROR,
+    INVALID_PARAMETER_VALUE,
+)
+from mlflow.protos.service_pb2 import GetRun, ListArtifacts, MlflowService
+from mlflow.store.artifact.artifact_repo import write_local_temp_trace_data_file
+from mlflow.store.artifact.cloud_artifact_repo import (
+    CloudArtifactRepository,
+    _complete_futures,
+    _compute_num_chunks,
+    _validate_chunk_size_aws,
+)
+from mlflow.utils import chunk_list
+from mlflow.utils.databricks_utils import get_databricks_host_creds
+from mlflow.utils.file_utils import (
+    download_file_using_http_uri,
+    read_chunk,
+)
+from mlflow.utils.proto_json_utils import message_to_json
+from mlflow.utils.request_utils import cloud_storage_http_request
+from mlflow.utils.rest_utils import (
+    _REST_API_PATH_PREFIX,
+    augmented_raise_for_status,
+    call_endpoint,
+    extract_api_info_for_service,
+)
+from mlflow.utils.uri import (
+    extract_and_normalize_path,
+    get_databricks_profile_uri_from_artifact_uri,
+    is_databricks_acled_artifacts_uri,
+    is_valid_dbfs_uri,
+    remove_databricks_profile_info_from_artifact_uri,
+)
+
+_logger = logging.getLogger(__name__)
+_MAX_CREDENTIALS_REQUEST_SIZE = 2000  # Max number of artifact paths in a single credentials request
+_SERVICE_AND_METHOD_TO_INFO = {
+    service: extract_api_info_for_service(service, _REST_API_PATH_PREFIX)
+    for service in [MlflowService, DatabricksMlflowArtifactsService]
+}
+
+
+class DatabricksArtifactRepository(CloudArtifactRepository):
+    """
+    Performs storage operations on artifacts in the access-controlled
+    `dbfs:/databricks/mlflow-tracking` location.
+
+    Signed access URIs for S3 / Azure Blob Storage are fetched from the MLflow service and used to
+    read and write files from/to this location.
+
+    The artifact_uri is expected to be of the form
+    dbfs:/databricks/mlflow-tracking/<EXP_ID>/<RUN_ID>/
+    """
+
+    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)
+        self._run_relative_artifact_repo_root_path = None
+
+    @property
+    def run_relative_artifact_repo_root_path(self):
+        """
+        Lazily computes the run-relative artifact repository root path to skip the run existence
+        check when downloading/uploading trace data.
+        """
+        if self._run_relative_artifact_repo_root_path is None:
+            # 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(self.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
+            )
+
+        return self._run_relative_artifact_repo_root_path
+
+    @staticmethod
+    def _extract_run_id(artifact_uri):
+        """
+        The artifact_uri is expected to be
+        dbfs:/databricks/mlflow-tracking/<EXP_ID>/<RUN_ID>/artifacts/<path>
+        Once the path from the input uri is extracted and normalized, it is
+        expected to be of the form
+        databricks/mlflow-tracking/<EXP_ID>/<RUN_ID>/artifacts/<path>
+
+        Hence the run_id is the 4th element of the normalized path.
+
+        Returns:
+            run_id extracted from the artifact_uri.
+        """
+        artifact_path = extract_and_normalize_path(artifact_uri)
+        return artifact_path.split("/")[3]
+
+    def _call_endpoint(self, service, api, json_body=None, path_params=None):
+        """
+        Calls the specified REST endpoint with the specified JSON body and path parameters.
+
+        Args:
+            service: The service to call.
+            api: The API to call.
+            json_body: The JSON body of the request.
+            path_params: The path parameters to substitute into the endpoint URI.
+
+        Returns:
+            The response from the REST endpoint.
+        """
+        db_creds = get_databricks_host_creds(self.databricks_profile_uri)
+        endpoint, method = _SERVICE_AND_METHOD_TO_INFO[service][api]
+        if path_params:
+            endpoint = endpoint.format(**path_params)
+        response_proto = api.Response()
+        return call_endpoint(db_creds, endpoint, method, json_body, response_proto)
+
+    def _get_run_artifact_root(self, run_id):
+        json_body = message_to_json(GetRun(run_id=run_id))
+        run_response = self._call_endpoint(MlflowService, GetRun, json_body)
+        return run_response.run.info.artifact_uri
+
+    def _get_credential_infos(self, request_message_class, run_id, paths):
+        """
+        Issue one or more requests for artifact credentials, providing read or write
+        access to the specified run-relative artifact `paths` within the MLflow Run specified
+        by `run_id`. The type of access credentials, read or write, is specified by
+        `request_message_class`.
+
+        Args:
+            request_message_class: Specifies the type of access credentials, read or write.
+            run_id: The specified MLflow Run.
+            paths: The specified run-relative artifact paths within the MLflow Run.
+
+        Returns:
+            A list of `ArtifactCredentialInfo` objects providing read access to the specified
+            run-relative artifact `paths` within the MLflow Run specified by `run_id`.
+        """
+        credential_infos = []
+
+        for paths_chunk in chunk_list(paths, _MAX_CREDENTIALS_REQUEST_SIZE):
+            page_token = None
+            while True:
+                json_body = message_to_json(
+                    request_message_class(run_id=run_id, path=paths_chunk, page_token=page_token)
+                )
+                response = self._call_endpoint(
+                    DatabricksMlflowArtifactsService, request_message_class, json_body
+                )
+                credential_infos += response.credential_infos
+                page_token = response.next_page_token
+                if not page_token or len(response.credential_infos) == 0:
+                    break
+
+        return credential_infos
+
+    def _get_write_credential_infos(self, remote_file_paths):
+        """
+        A list of `ArtifactCredentialInfo` objects providing write access to the specified
+        run-relative artifact `paths` within the MLflow Run specified by `run_id`.
+        """
+        run_relative_remote_paths = [
+            posixpath.join(self.run_relative_artifact_repo_root_path, p or "")
+            for p in remote_file_paths
+        ]
+        return self._get_credential_infos(
+            GetCredentialsForWrite, self.run_id, run_relative_remote_paths
+        )
+
+    def download_trace_data(self) -> dict[str, Any]:
+        cred = self._call_endpoint(
+            DatabricksMlflowArtifactsService,
+            GetCredentialsForTraceDataDownload,
+            path_params={"request_id": self.run_id},
+        )
+        signed_uri = cred.credential_info.signed_uri
+        headers = self._extract_headers_from_credentials(cred.credential_info.headers)
+        with cloud_storage_http_request("get", signed_uri, headers=headers) as resp:
+            try:
+                augmented_raise_for_status(resp)
+            except requests.HTTPError as e:
+                if e.response.status_code == 404:
+                    raise MlflowTraceDataNotFound(request_id=self.run_id) from e
+                raise
+
+            try:
+                return json.loads(resp.content)
+            except json.JSONDecodeError as e:
+                raise MlflowTraceDataCorrupted(request_id=self.run_id) from e
+
+    def _get_upload_trace_data_cred_info(self):
+        res = self._call_endpoint(
+            DatabricksMlflowArtifactsService,
+            GetCredentialsForTraceDataUpload,
+            path_params={"request_id": self.run_id},
+        )
+        return res.credential_info
+
+    def upload_trace_data(self, trace_data: str) -> None:
+        cred = self._get_upload_trace_data_cred_info()
+        with write_local_temp_trace_data_file(trace_data) as temp_file:
+            if cred.type == ArtifactCredentialType.AZURE_ADLS_GEN2_SAS_URI:
+                self._azure_adls_gen2_upload_file(
+                    credentials=cred,
+                    local_file=temp_file,
+                    artifact_file_path=None,
+                    get_credentials=lambda artifact_paths: [
+                        self._get_upload_trace_data_cred_info()
+                    ],
+                )
+            elif cred.type == ArtifactCredentialType.AZURE_SAS_URI:
+                self._azure_upload_file(
+                    credentials=cred,
+                    local_file=temp_file,
+                    artifact_file_path=None,
+                    get_credentials=lambda artifact_paths: [
+                        self._get_upload_trace_data_cred_info()
+                    ],
+                )
+            elif (
+                cred.type == ArtifactCredentialType.AWS_PRESIGNED_URL
+                or cred.type == ArtifactCredentialType.GCP_SIGNED_URL
+            ):
+                self._signed_url_upload_file(cred, temp_file)
+
+    def _get_read_credential_infos(self, remote_file_paths):
+        """
+        Returns:
+            A list of `ArtifactCredentialInfo` objects providing read access to the specified
+            run-relative artifact `paths` within the MLflow Run specified by `run_id`.
+        """
+        if type(remote_file_paths) == str:
+            remote_file_paths = [remote_file_paths]
+        if type(remote_file_paths) != list:
+            raise MlflowException(
+                f"Expected `paths` to be a list of strings. Got {type(remote_file_paths)}"
+            )
+        run_relative_remote_paths = [
+            posixpath.join(self.run_relative_artifact_repo_root_path, p) for p in remote_file_paths
+        ]
+        return self._get_credential_infos(
+            GetCredentialsForRead, self.run_id, run_relative_remote_paths
+        )
+
+    def _extract_headers_from_credentials(self, headers):
+        """
+        Returns:
+            A python dictionary of http headers converted from the protobuf credentials.
+        """
+        return {header.name: header.value for header in headers}
+
+    def _azure_upload_chunk(
+        self,
+        credentials,
+        headers,
+        local_file,
+        artifact_file_path,
+        start_byte,
+        size,
+        get_credentials,
+    ):
+        """
+        Uploads a chunk of a file to a given Azure storage location.
+
+        Args:
+            credentials: The credentials for the upload.
+            headers: The headers for the upload.
+            local_file: The local file to upload.
+            artifact_file_path: The path to the artifact file.
+            start_byte: The starting byte of the chunk.
+            size: The size of the chunk.
+            get_credentials: The function to call to get new credentials.
+        """
+        # Base64-encode a UUID, producing a UTF8-encoded bytestring. Then, decode
+        # the bytestring for compliance with Azure Blob Storage API requests
+        block_id = base64.b64encode(uuid.uuid4().hex.encode()).decode("utf-8")
+        chunk = read_chunk(local_file, size, start_byte)
+        try:
+            put_block(credentials.signed_uri, block_id, chunk, headers=headers)
+        except requests.HTTPError as e:
+            if e.response.status_code in [401, 403]:
+                _logger.info(
+                    "Failed to authorize request, possibly due to credential expiration."
+                    " Refreshing credentials and trying again..."
+                )
+                credential_info = get_credentials([artifact_file_path])[0]
+                put_block(credential_info.signed_uri, block_id, chunk, headers=headers)
+            else:
+                raise e
+        return block_id
+
+    def _azure_upload_file(self, credentials, local_file, artifact_file_path, get_credentials):
+        """
+        Uploads a file to a given Azure storage location.
+        The function uses a file chunking generator with 100 MB being the size limit for each chunk.
+        This limit is imposed by the stage_block API in azure-storage-blob.
+        In the case the file size is large and the upload takes longer than the validity of the
+        given credentials, a new set of credentials are generated and the operation continues. This
+        is the reason for the first nested try-except block
+        Finally, since the prevailing credentials could expire in the time between the last
+        stage_block and the commit, a second try-except block refreshes credentials if needed.
+
+        Args:
+            credentials: The credentials for the upload.
+            local_file: The local file to upload.
+            artifact_file_path: The path to the artifact file.
+            get_credentials: The function to call to get new credentials.
+        """
+        try:
+            headers = self._extract_headers_from_credentials(credentials.headers)
+            futures = {}
+            num_chunks = _compute_num_chunks(local_file, MLFLOW_MULTIPART_UPLOAD_CHUNK_SIZE.get())
+            for index in range(num_chunks):
+                start_byte = index * MLFLOW_MULTIPART_UPLOAD_CHUNK_SIZE.get()
+                future = self.chunk_thread_pool.submit(
+                    self._azure_upload_chunk,
+                    credentials=credentials,
+                    headers=headers,
+                    local_file=local_file,
+                    artifact_file_path=artifact_file_path,
+                    start_byte=start_byte,
+                    size=MLFLOW_MULTIPART_UPLOAD_CHUNK_SIZE.get(),
+                    get_credentials=get_credentials,
+                )
+                futures[future] = index
+
+            results, errors = _complete_futures(futures, local_file)
+            if errors:
+                raise MlflowException(
+                    f"Failed to upload at least one part of {local_file}. Errors: {errors}"
+                )
+            # Sort results by the chunk index
+            uploading_block_list = [results[index] for index in sorted(results)]
+
+            try:
+                put_block_list(credentials.signed_uri, uploading_block_list, headers=headers)
+            except requests.HTTPError as e:
+                if e.response.status_code in [401, 403]:
+                    _logger.info(
+                        "Failed to authorize request, possibly due to credential expiration."
+                        " Refreshing credentials and trying again..."
+                    )
+                    credential_info = get_credentials([artifact_file_path])[0]
+                    put_block_list(
+                        credential_info.signed_uri, uploading_block_list, headers=headers
+                    )
+                else:
+                    raise e
+        except Exception as err:
+            raise MlflowException(err)
+
+    def _retryable_adls_function(self, func, artifact_file_path, get_credentials, **kwargs):
+        """
+        Calls the passed function, retrying if the credentials have expired.
+
+        Args:
+            func: The function to call.
+            artifact_file_path: The artifact file path.
+            get_credentials: The function to call to get new credentials.
+            **kwargs: The keyword arguments to pass to the function.
+        """
+        # Attempt to call the passed function.  Retry if the credentials have expired
+        try:
+            func(**kwargs)
+        except requests.HTTPError as e:
+            if e.response.status_code in [403]:
+                _logger.info(
+                    "Failed to authorize ADLS operation, possibly due "
+                    "to credential expiration. Refreshing credentials and trying again..."
+                )
+                new_credentials = get_credentials([artifact_file_path])[0]
+                kwargs["sas_url"] = new_credentials.signed_uri
+                func(**kwargs)
+            else:
+                raise e
+
+    def _azure_adls_gen2_upload_file(
+        self, credentials, local_file, artifact_file_path, get_credentials
+    ):
+        """
+        Uploads a file to a given Azure storage location using the ADLS gen2 API.
+
+        Args:
+            credentials: The credentials for the upload.
+            local_file: The local file to upload.
+            artifact_file_path: The path to the artifact file.
+            get_credentials: The function to call to get new credentials.
+        """
+        try:
+            headers = self._extract_headers_from_credentials(credentials.headers)
+
+            # try to create the file
+            self._retryable_adls_function(
+                func=put_adls_file_creation,
+                artifact_file_path=artifact_file_path,
+                get_credentials=get_credentials,
+                sas_url=credentials.signed_uri,
+                headers=headers,
+            )
+
+            # next try to append the file
+            futures = {}
+            file_size = os.path.getsize(local_file)
+            num_chunks = _compute_num_chunks(local_file, MLFLOW_MULTIPART_UPLOAD_CHUNK_SIZE.get())
+            use_single_part_upload = num_chunks == 1
+            for index in range(num_chunks):
+                start_byte = index * MLFLOW_MULTIPART_UPLOAD_CHUNK_SIZE.get()
+                future = self.chunk_thread_pool.submit(
+                    self._retryable_adls_function,
+                    func=patch_adls_file_upload,
+                    artifact_file_path=artifact_file_path,
+                    get_credentials=get_credentials,
+                    sas_url=credentials.signed_uri,
+                    local_file=local_file,
+                    start_byte=start_byte,
+                    size=MLFLOW_MULTIPART_UPLOAD_CHUNK_SIZE.get(),
+                    position=start_byte,
+                    headers=headers,
+                    is_single=use_single_part_upload,
+                )
+                futures[future] = index
+
+            _, errors = _complete_futures(futures, local_file)
+            if errors:
+                raise MlflowException(
+                    f"Failed to upload at least one part of {artifact_file_path}. Errors: {errors}"
+                )
+
+            # finally try to flush the file
+            if not use_single_part_upload:
+                self._retryable_adls_function(
+                    func=patch_adls_flush,
+                    artifact_file_path=artifact_file_path,
+                    get_credentials=get_credentials,
+                    sas_url=credentials.signed_uri,
+                    position=file_size,
+                    headers=headers,
+                )
+        except Exception as err:
+            raise MlflowException(err)
+
+    def _signed_url_upload_file(self, credentials, local_file):
+        try:
+            headers = self._extract_headers_from_credentials(credentials.headers)
+            signed_write_uri = credentials.signed_uri
+            # Putting an empty file in a request by reading file bytes gives 501 error.
+            if os.stat(local_file).st_size == 0:
+                with cloud_storage_http_request(
+                    "put", signed_write_uri, data="", headers=headers
+                ) as response:
+                    augmented_raise_for_status(response)
+            else:
+                with open(local_file, "rb") as file:
+                    with cloud_storage_http_request(
+                        "put", signed_write_uri, data=file, headers=headers
+                    ) as response:
+                        augmented_raise_for_status(response)
+        except Exception as err:
+            raise MlflowException(err)
+
+    def _upload_to_cloud(self, cloud_credential_info, src_file_path, artifact_file_path):
+        """
+        Upload a local file to the cloud. Note that in this artifact repository, files are uploaded
+        to run-relative artifact file paths in the artifact repository.
+
+        Args:
+            cloud_credential_info: ArtifactCredentialInfo object with presigned URL for the file.
+            src_file_path: Local source file path for the upload.
+            artifact_file_path: Path in the artifact repository, relative to the run root path,
+                where the artifact will be logged.
+
+        """
+        if cloud_credential_info.type == ArtifactCredentialType.AZURE_SAS_URI:
+            self._azure_upload_file(
+                cloud_credential_info,
+                src_file_path,
+                artifact_file_path,
+                get_credentials=self._get_write_credential_infos,
+            )
+        elif cloud_credential_info.type == ArtifactCredentialType.AZURE_ADLS_GEN2_SAS_URI:
+            self._azure_adls_gen2_upload_file(
+                cloud_credential_info,
+                src_file_path,
+                artifact_file_path,
+                self._get_write_credential_infos,
+            )
+        elif cloud_credential_info.type == ArtifactCredentialType.AWS_PRESIGNED_URL:
+            if os.path.getsize(src_file_path) > MLFLOW_MULTIPART_UPLOAD_MINIMUM_FILE_SIZE.get():
+                _validate_chunk_size_aws(MLFLOW_MULTIPART_UPLOAD_CHUNK_SIZE.get())
+                self._multipart_upload(src_file_path, artifact_file_path)
+            else:
+                self._signed_url_upload_file(cloud_credential_info, src_file_path)
+        elif cloud_credential_info.type == ArtifactCredentialType.GCP_SIGNED_URL:
+            self._signed_url_upload_file(cloud_credential_info, src_file_path)
+        else:
+            raise MlflowException(
+                message="Cloud provider not supported.", error_code=INTERNAL_ERROR
+            )
+
+    def _download_from_cloud(self, remote_file_path, local_path):
+        """
+        Download a file from the input `remote_file_path` and save it to `local_path`.
+
+        Args:
+            remote_file_path: Path relative to the run root path to file in remote artifact
+                repository.
+            local_path: Local path to download file to.
+
+        """
+        read_credentials = self._get_read_credential_infos(remote_file_path)
+        # Read credentials for only one file were requested. So we expected only one value in
+        # the response.
+        assert len(read_credentials) == 1
+        cloud_credential_info = read_credentials[0]
+
+        if cloud_credential_info.type not in [
+            ArtifactCredentialType.AZURE_SAS_URI,
+            ArtifactCredentialType.AZURE_ADLS_GEN2_SAS_URI,
+            ArtifactCredentialType.AWS_PRESIGNED_URL,
+            ArtifactCredentialType.GCP_SIGNED_URL,
+        ]:
+            raise MlflowException(
+                message="Cloud provider not supported.", error_code=INTERNAL_ERROR
+            )
+        try:
+            download_file_using_http_uri(
+                cloud_credential_info.signed_uri,
+                local_path,
+                MLFLOW_MULTIPART_DOWNLOAD_CHUNK_SIZE.get(),
+                self._extract_headers_from_credentials(cloud_credential_info.headers),
+            )
+        except Exception as err:
+            raise MlflowException(err)
+
+    def _create_multipart_upload(self, run_id, path, num_parts):
+        return self._call_endpoint(
+            DatabricksMlflowArtifactsService,
+            CreateMultipartUpload,
+            message_to_json(CreateMultipartUpload(run_id=run_id, path=path, num_parts=num_parts)),
+        )
+
+    def _get_presigned_upload_part_url(self, run_id, path, upload_id, part_number):
+        return self._call_endpoint(
+            DatabricksMlflowArtifactsService,
+            GetPresignedUploadPartUrl,
+            message_to_json(
+                GetPresignedUploadPartUrl(
+                    run_id=run_id, path=path, upload_id=upload_id, part_number=part_number
+                )
+            ),
+        )
+
+    def _upload_part(self, cred_info, data):
+        headers = self._extract_headers_from_credentials(cred_info.headers)
+        with cloud_storage_http_request(
+            "put",
+            cred_info.signed_uri,
+            data=data,
+            headers=headers,
+        ) as response:
+            augmented_raise_for_status(response)
+            return response.headers["ETag"]
+
+    def _upload_part_retry(self, cred_info, upload_id, part_number, local_file, start_byte, size):
+        data = read_chunk(local_file, size, start_byte)
+        try:
+            return self._upload_part(cred_info, data)
+        except requests.HTTPError as e:
+            if e.response.status_code not in (401, 403):
+                raise e
+            _logger.info(
+                "Failed to authorize request, possibly due to credential expiration."
+                " Refreshing credentials and trying again..."
+            )
+            resp = self._get_presigned_upload_part_url(
+                cred_info.run_id, cred_info.path, upload_id, part_number
+            )
+            return self._upload_part(resp.upload_credential_info, data)
+
+    def _upload_parts(self, local_file, create_mpu_resp):
+        futures = {}
+        for index, cred_info in enumerate(create_mpu_resp.upload_credential_infos):
+            part_number = index + 1
+            start_byte = index * MLFLOW_MULTIPART_UPLOAD_CHUNK_SIZE.get()
+            future = self.chunk_thread_pool.submit(
+                self._upload_part_retry,
+                cred_info=cred_info,
+                upload_id=create_mpu_resp.upload_id,
+                part_number=part_number,
+                local_file=local_file,
+                start_byte=start_byte,
+                size=MLFLOW_MULTIPART_UPLOAD_CHUNK_SIZE.get(),
+            )
+            futures[future] = part_number
+
+        results, errors = _complete_futures(futures, local_file)
+        if errors:
+            raise MlflowException(
+                f"Failed to upload at least one part of {local_file}. Errors: {errors}"
+            )
+
+        return [
+            PartEtag(part_number=part_number, etag=results[part_number])
+            for part_number in sorted(results)
+        ]
+
+    def _complete_multipart_upload(self, run_id, path, upload_id, part_etags):
+        return self._call_endpoint(
+            DatabricksMlflowArtifactsService,
+            CompleteMultipartUpload,
+            message_to_json(
+                CompleteMultipartUpload(
+                    run_id=run_id,
+                    path=path,
+                    upload_id=upload_id,
+                    part_etags=part_etags,
+                )
+            ),
+        )
+
+    def _abort_multipart_upload(self, cred_info):
+        headers = self._extract_headers_from_credentials(cred_info.headers)
+        with cloud_storage_http_request(
+            "delete", cred_info.signed_uri, headers=headers
+        ) as response:
+            augmented_raise_for_status(response)
+            return response
+
+    def _multipart_upload(self, local_file, artifact_file_path):
+        run_relative_artifact_path = posixpath.join(
+            self.run_relative_artifact_repo_root_path, artifact_file_path or ""
+        )
+        num_parts = _compute_num_chunks(local_file, MLFLOW_MULTIPART_UPLOAD_CHUNK_SIZE.get())
+        create_mpu_resp = self._create_multipart_upload(
+            self.run_id, run_relative_artifact_path, num_parts
+        )
+        try:
+            part_etags = self._upload_parts(local_file, create_mpu_resp)
+            self._complete_multipart_upload(
+                self.run_id,
+                run_relative_artifact_path,
+                create_mpu_resp.upload_id,
+                part_etags,
+            )
+        except Exception as e:
+            _logger.warning(
+                "Encountered an unexpected error during multipart upload: %s, aborting", e
+            )
+            self._abort_multipart_upload(create_mpu_resp.abort_credential_info)
+            raise e
+
+    def log_artifact(self, local_file, artifact_path=None):
+        src_file_name = os.path.basename(local_file)
+        artifact_file_path = posixpath.join(artifact_path or "", src_file_name)
+        write_credential_info = self._get_write_credential_infos([artifact_file_path])[0]
+        self._upload_to_cloud(
+            cloud_credential_info=write_credential_info,
+            src_file_path=local_file,
+            artifact_file_path=artifact_file_path,
+        )
+
+    def list_artifacts(self, path: Optional[str] = None) -> list:
+        if path:
+            run_relative_path = posixpath.join(self.run_relative_artifact_repo_root_path, path)
+        else:
+            run_relative_path = self.run_relative_artifact_repo_root_path
+        infos = []
+        page_token = None
+        while True:
+            json_body = message_to_json(
+                ListArtifacts(run_id=self.run_id, path=run_relative_path, page_token=page_token)
+            )
+            response = self._call_endpoint(MlflowService, ListArtifacts, json_body)
+            artifact_list = response.files
+            # If `path` is a file, ListArtifacts returns a single list element with the
+            # same name as `path`. The list_artifacts API expects us to return an empty list in this
+            # case, so we do so here.
+            if (
+                len(artifact_list) == 1
+                and artifact_list[0].path == run_relative_path
+                and not artifact_list[0].is_dir
+            ):
+                return []
+            for output_file in artifact_list:
+                file_rel_path = posixpath.relpath(
+                    path=output_file.path, start=self.run_relative_artifact_repo_root_path
+                )
+                artifact_size = None if output_file.is_dir else output_file.file_size
+                infos.append(FileInfo(file_rel_path, output_file.is_dir, artifact_size))
+            if len(artifact_list) == 0 or not response.next_page_token:
+                break
+            page_token = response.next_page_token
+        return infos
+
+    def delete_artifacts(self, artifact_path=None):
+        raise MlflowException("Not implemented yet")

venv/lib/python3.9/site-packages/mlflow/store/artifact/databricks_models_artifact_repo.py

@@ -0,0 +1,216 @@
+import json
+import logging
+import os
+import posixpath
+from typing import Optional
+
+import mlflow.tracking
+from mlflow.entities import FileInfo
+from mlflow.environment_variables import (
+    MLFLOW_ENABLE_MULTIPART_DOWNLOAD,
+    MLFLOW_MULTIPART_DOWNLOAD_CHUNK_SIZE,
+)
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
+from mlflow.store.artifact.artifact_repo import ArtifactRepository
+from mlflow.store.artifact.utils.models import (
+    get_model_name_and_version,
+    is_using_databricks_registry,
+)
+from mlflow.utils.databricks_utils import (
+    get_databricks_host_creds,
+    warn_on_deprecated_cross_workspace_registry_uri,
+)
+from mlflow.utils.file_utils import (
+    download_chunk_retries,
+    download_file_using_http_uri,
+    parallelized_download_file_using_http_uri,
+    remove_on_error,
+)
+from mlflow.utils.rest_utils import http_request
+from mlflow.utils.uri import get_databricks_profile_uri_from_artifact_uri
+
+_logger = logging.getLogger(__name__)
+# The constant REGISTRY_LIST_ARTIFACT_ENDPOINT is defined as @developer_stable
+REGISTRY_LIST_ARTIFACTS_ENDPOINT = "/api/2.0/mlflow/model-versions/list-artifacts"
+# The constant REGISTRY_ARTIFACT_PRESIGNED_URI_ENDPOINT is defined as @developer_stable
+REGISTRY_ARTIFACT_PRESIGNED_URI_ENDPOINT = "/api/2.0/mlflow/model-versions/get-signed-download-uri"
+
+
+class DatabricksModelsArtifactRepository(ArtifactRepository):
+    """
+    Performs storage operations on artifacts controlled by a Databricks-hosted model registry.
+
+    Signed access URIs for the appropriate cloud storage locations are fetched from the
+    MLflow service and used to download model artifacts.
+
+    The artifact_uri is expected to be of the form
+    - `models:/<model_name>/<model_version>`
+    - `models:/<model_name>/<stage>`  (refers to the latest model version in the given stage)
+    - `models:/<model_name>/latest`  (refers to the latest of all model versions)
+    - `models://<profile>/<model_name>/<model_version or stage or 'latest'>`
+
+    Note : This artifact repository is meant is to be instantiated by the ModelsArtifactRepository
+    when the client is pointing to a Databricks-hosted model registry.
+    """
+
+    def __init__(self, artifact_uri):
+        if not is_using_databricks_registry(artifact_uri):
+            raise MlflowException(
+                message="A valid databricks profile is required to instantiate this repository",
+                error_code=INVALID_PARAMETER_VALUE,
+            )
+        super().__init__(artifact_uri)
+        from mlflow.tracking.client import MlflowClient
+
+        self.databricks_profile_uri = (
+            get_databricks_profile_uri_from_artifact_uri(artifact_uri) or mlflow.get_registry_uri()
+        )
+        warn_on_deprecated_cross_workspace_registry_uri(self.databricks_profile_uri)
+        client = MlflowClient(registry_uri=self.databricks_profile_uri)
+        self.model_name, self.model_version = get_model_name_and_version(client, artifact_uri)
+        # Use an isolated thread pool executor for chunk uploads/downloads to avoid a deadlock
+        # caused by waiting for a chunk-upload/download task within a file-upload/download task.
+        # See https://superfastpython.com/threadpoolexecutor-deadlock/#Deadlock_1_Submit_and_Wait_for_a_Task_Within_a_Task
+        # for more details
+        self.chunk_thread_pool = self._create_thread_pool()
+
+    def _call_endpoint(self, json, endpoint):
+        db_creds = get_databricks_host_creds(self.databricks_profile_uri)
+        return http_request(host_creds=db_creds, endpoint=endpoint, method="GET", params=json)
+
+    def _make_json_body(self, path, page_token=None):
+        body = {"name": self.model_name, "version": self.model_version, "path": path}
+        if page_token:
+            body["page_token"] = page_token
+        return body
+
+    def list_artifacts(self, path: Optional[str] = None) -> list[FileInfo]:
+        infos = []
+        page_token = None
+        if not path:
+            path = ""
+        while True:
+            json_body = self._make_json_body(path, page_token)
+            response = self._call_endpoint(json_body, REGISTRY_LIST_ARTIFACTS_ENDPOINT)
+            try:
+                response.raise_for_status()
+                json_response = json.loads(response.text)
+            except Exception:
+                raise MlflowException(
+                    f"API request to list files under path `{path}` failed with status code "
+                    f"{response.status_code}. Response body: {response.text}"
+                )
+            artifact_list = json_response.get("files", [])
+            next_page_token = json_response.get("next_page_token", None)
+            # If `path` is a file, ListArtifacts returns a single list element with the
+            # same name as `path`. The list_artifacts API expects us to return an empty list in this
+            # case, so we do so here.
+            if (
+                len(artifact_list) == 1
+                and artifact_list[0]["path"] == path
+                and not artifact_list[0]["is_dir"]
+            ):
+                return []
+            for output_file in artifact_list:
+                artifact_size = None if output_file["is_dir"] else output_file["file_size"]
+                infos.append(FileInfo(output_file["path"], output_file["is_dir"], artifact_size))
+            if len(artifact_list) == 0 or not next_page_token:
+                break
+            page_token = next_page_token
+        return infos
+
+    # TODO: Change the implementation of this to match how databricks_artifact_repo.py handles this
+    def _get_signed_download_uri(self, path=None):
+        if not path:
+            path = ""
+        json_body = self._make_json_body(path)
+        response = self._call_endpoint(json_body, REGISTRY_ARTIFACT_PRESIGNED_URI_ENDPOINT)
+        try:
+            json_response = json.loads(response.text)
+        except ValueError:
+            raise MlflowException(
+                f"API request to get presigned uri to for file under path `{path}` failed with"
+                f" status code {response.status_code}. Response body: {response.text}"
+            )
+        return json_response.get("signed_uri", None), json_response.get("headers", None)
+
+    def _extract_headers_from_signed_url(self, headers):
+        if headers is None:
+            return {}
+        filtered_headers = filter(lambda h: "name" in h and "value" in h, headers)
+        return {header.get("name"): header.get("value") for header in filtered_headers}
+
+    def _parallelized_download_from_cloud(
+        self, signed_uri, headers, file_size, dst_local_file_path, dst_run_relative_artifact_path
+    ):
+        from mlflow.utils.databricks_utils import get_databricks_env_vars
+
+        with remove_on_error(dst_local_file_path):
+            parallel_download_subproc_env = os.environ.copy()
+            parallel_download_subproc_env.update(
+                get_databricks_env_vars(self.databricks_profile_uri)
+            )
+            failed_downloads = parallelized_download_file_using_http_uri(
+                thread_pool_executor=self.chunk_thread_pool,
+                http_uri=signed_uri,
+                download_path=dst_local_file_path,
+                remote_file_path=dst_run_relative_artifact_path,
+                file_size=file_size,
+                # URI type is not known in this context
+                uri_type=None,
+                chunk_size=MLFLOW_MULTIPART_DOWNLOAD_CHUNK_SIZE.get(),
+                env=parallel_download_subproc_env,
+                headers=headers,
+            )
+            if failed_downloads:
+                new_signed_uri, new_headers = self._get_signed_download_uri(
+                    dst_run_relative_artifact_path
+                )
+                new_headers = self._extract_headers_from_signed_url(new_headers)
+                download_chunk_retries(
+                    chunks=list(failed_downloads),
+                    http_uri=new_signed_uri,
+                    headers=new_headers,
+                    download_path=dst_local_file_path,
+                )
+
+    def _download_file(self, remote_file_path, local_path):
+        try:
+            parent_dir, _ = posixpath.split(remote_file_path)
+            file_infos = self.list_artifacts(parent_dir)
+            file_info = [info for info in file_infos if info.path == remote_file_path]
+            file_size = file_info[0].file_size if len(file_info) == 1 else None
+            signed_uri, raw_headers = self._get_signed_download_uri(remote_file_path)
+            headers = {}
+            if raw_headers is not None:
+                # Don't send None to _extract_headers_from_signed_url
+                headers = self._extract_headers_from_signed_url(raw_headers)
+            if (
+                not file_size
+                or file_size <= MLFLOW_MULTIPART_DOWNLOAD_CHUNK_SIZE.get()
+                or not MLFLOW_ENABLE_MULTIPART_DOWNLOAD.get()
+            ):
+                download_file_using_http_uri(
+                    signed_uri, local_path, MLFLOW_MULTIPART_DOWNLOAD_CHUNK_SIZE.get(), headers
+                )
+            else:
+                self._parallelized_download_from_cloud(
+                    signed_uri,
+                    headers,
+                    file_size,
+                    local_path,
+                    remote_file_path,
+                )
+
+        except Exception as err:
+            raise MlflowException(err)
+
+    def log_artifact(self, local_file, artifact_path=None):
+        raise MlflowException("This repository does not support logging artifacts.")
+
+    def log_artifacts(self, local_dir, artifact_path=None):
+        raise MlflowException("This repository does not support logging artifacts.")
+
+    def delete_artifacts(self, artifact_path=None):
+        raise NotImplementedError("This artifact repository does not support deleting artifacts")

venv/lib/python3.9/site-packages/mlflow/store/artifact/databricks_sdk_models_artifact_repo.py

@@ -0,0 +1,95 @@
+import posixpath
+from typing import Optional
+
+from databricks.sdk.errors.platform import NotFound
+
+from mlflow.entities import FileInfo
+from mlflow.store.artifact.cloud_artifact_repo import CloudArtifactRepository
+
+DOWNLOAD_CHUNK_SIZE = 1024 * 1024 * 1024
+
+
+def _get_databricks_workspace_client():
+    from databricks.sdk import WorkspaceClient
+
+    return WorkspaceClient()
+
+
+class DatabricksSDKModelsArtifactRepository(CloudArtifactRepository):
+    """
+    Stores and retrieves model artifacts via Databricks SDK, agnostic to the underlying cloud
+    that stores the model artifacts.
+    """
+
+    def __init__(self, model_name, model_version):
+        self.model_name = model_name
+        self.model_version = model_version
+        self.model_base_path = f"/Models/{model_name.replace('.', '/')}/{model_version}"
+        self.client = _get_databricks_workspace_client()
+        super().__init__(self.model_base_path)
+
+    def list_artifacts(self, path: Optional[str] = None) -> list[FileInfo]:
+        dest_path = self.model_base_path
+        if path:
+            dest_path = posixpath.join(dest_path, path)
+
+        file_infos = []
+
+        # check if dest_path is file, if so return empty dir
+        if not self._is_dir(dest_path):
+            return file_infos
+
+        resp = self.client.files.list_directory_contents(dest_path)
+        for directory_entry in resp:
+            relative_path = posixpath.relpath(directory_entry.path, self.model_base_path)
+            file_infos.append(
+                FileInfo(
+                    path=relative_path,
+                    is_dir=directory_entry.is_directory,
+                    file_size=directory_entry.file_size,
+                )
+            )
+
+        return sorted(file_infos, key=lambda f: f.path)
+
+    def _is_dir(self, artifact_path):
+        try:
+            self.client.files.get_directory_metadata(artifact_path)
+        except NotFound:
+            return False
+        return True
+
+    def _upload_to_cloud(self, cloud_credential_info, src_file_path, artifact_file_path=None):
+        dest_path = self.model_base_path
+        if artifact_file_path:
+            dest_path = posixpath.join(dest_path, artifact_file_path)
+
+        with open(src_file_path, "rb") as f:
+            self.client.files.upload(dest_path, f, overwrite=True)
+
+    def log_artifact(self, local_file, artifact_path=None):
+        self._upload_to_cloud(
+            cloud_credential_info=None,
+            src_file_path=local_file,
+            artifact_file_path=artifact_path,
+        )
+
+    def _download_from_cloud(self, remote_file_path, local_path):
+        dest_path = self.model_base_path
+        if remote_file_path:
+            dest_path = posixpath.join(dest_path, remote_file_path)
+
+        resp = self.client.files.download(dest_path)
+        contents = resp.contents
+
+        with open(local_path, "wb") as f:
+            while chunk := contents.read(DOWNLOAD_CHUNK_SIZE):
+                f.write(chunk)
+
+    def _get_write_credential_infos(self, remote_file_paths):
+        # Databricks sdk based model download/upload don't need any extra credentials
+        return [None] * len(remote_file_paths)
+
+    def _get_read_credential_infos(self, remote_file_paths):
+        # Databricks sdk based model download/upload don't need any extra credentials
+        return [None] * len(remote_file_paths)

venv/lib/python3.9/site-packages/mlflow/store/artifact/dbfs_artifact_repo.py

@@ -0,0 +1,232 @@
+import json
+import os
+import posixpath
+from typing import Optional
+
+import mlflow.utils.databricks_utils
+from mlflow.entities import FileInfo
+from mlflow.environment_variables import MLFLOW_ENABLE_DBFS_FUSE_ARTIFACT_REPO
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
+from mlflow.store.artifact.artifact_repo import ArtifactRepository
+from mlflow.store.artifact.databricks_artifact_repo import DatabricksArtifactRepository
+from mlflow.store.artifact.local_artifact_repo import LocalArtifactRepository
+from mlflow.store.tracking.rest_store import RestStore
+from mlflow.tracking._tracking_service import utils
+from mlflow.utils.databricks_utils import get_databricks_host_creds
+from mlflow.utils.file_utils import relative_path_to_artifact_path
+from mlflow.utils.rest_utils import (
+    RESOURCE_NON_EXISTENT,
+    http_request,
+    http_request_safe,
+)
+from mlflow.utils.string_utils import strip_prefix
+from mlflow.utils.uri import (
+    get_databricks_profile_uri_from_artifact_uri,
+    is_databricks_acled_artifacts_uri,
+    is_databricks_model_registry_artifacts_uri,
+    is_valid_dbfs_uri,
+    remove_databricks_profile_info_from_artifact_uri,
+    strip_scheme,
+)
+
+# The following constants are defined as @developer_stable
+LIST_API_ENDPOINT = "/api/2.0/dbfs/list"
+GET_STATUS_ENDPOINT = "/api/2.0/dbfs/get-status"
+DOWNLOAD_CHUNK_SIZE = 1024
+
+
+class DbfsRestArtifactRepository(ArtifactRepository):
+    """
+    Stores artifacts on DBFS using the DBFS REST API.
+
+    This repository is used with URIs of the form ``dbfs:/<path>``. The repository can only be used
+    together with the RestStore.
+    """
+
+    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().__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 _databricks_api_request(self, endpoint, method, **kwargs):
+        host_creds = self.get_host_creds()
+        return http_request_safe(host_creds=host_creds, endpoint=endpoint, method=method, **kwargs)
+
+    def _dbfs_list_api(self, json):
+        host_creds = self.get_host_creds()
+        return http_request(
+            host_creds=host_creds, endpoint=LIST_API_ENDPOINT, method="GET", params=json
+        )
+
+    def _dbfs_download(self, output_path, endpoint):
+        with open(output_path, "wb") as f:
+            response = self._databricks_api_request(endpoint=endpoint, method="GET", stream=True)
+            try:
+                for content in response.iter_content(chunk_size=DOWNLOAD_CHUNK_SIZE):
+                    f.write(content)
+            finally:
+                response.close()
+
+    def _is_directory(self, artifact_path):
+        dbfs_path = self._get_dbfs_path(artifact_path) if artifact_path else self._get_dbfs_path("")
+        return self._dbfs_is_dir(dbfs_path)
+
+    def _dbfs_is_dir(self, dbfs_path):
+        response = self._databricks_api_request(
+            endpoint=GET_STATUS_ENDPOINT, method="GET", params={"path": dbfs_path}
+        )
+        json_response = json.loads(response.text)
+        try:
+            return json_response["is_dir"]
+        except KeyError:
+            raise MlflowException(f"DBFS path {dbfs_path} does not exist")
+
+    def _get_dbfs_path(self, artifact_path):
+        return "/{}/{}".format(
+            strip_scheme(self.artifact_uri).lstrip("/"),
+            artifact_path.lstrip("/"),
+        )
+
+    def _get_dbfs_endpoint(self, artifact_path):
+        return f"/dbfs{self._get_dbfs_path(artifact_path)}"
+
+    def log_artifact(self, local_file, artifact_path=None):
+        basename = os.path.basename(local_file)
+        if artifact_path:
+            http_endpoint = self._get_dbfs_endpoint(posixpath.join(artifact_path, basename))
+        else:
+            http_endpoint = self._get_dbfs_endpoint(basename)
+        if os.stat(local_file).st_size == 0:
+            # The API frontend doesn't like it when we post empty files to it using
+            # `requests.request`, potentially due to the bug described in
+            # https://github.com/requests/requests/issues/4215
+            self._databricks_api_request(
+                endpoint=http_endpoint, method="POST", data="", allow_redirects=False
+            )
+        else:
+            with open(local_file, "rb") as f:
+                self._databricks_api_request(
+                    endpoint=http_endpoint, method="POST", data=f, allow_redirects=False
+                )
+
+    def log_artifacts(self, local_dir, artifact_path=None):
+        artifact_path = artifact_path or ""
+        for dirpath, _, filenames in os.walk(local_dir):
+            artifact_subdir = artifact_path
+            if dirpath != local_dir:
+                rel_path = os.path.relpath(dirpath, local_dir)
+                rel_path = relative_path_to_artifact_path(rel_path)
+                artifact_subdir = posixpath.join(artifact_path, rel_path)
+            for name in filenames:
+                file_path = os.path.join(dirpath, name)
+                self.log_artifact(file_path, artifact_subdir)
+
+    def list_artifacts(self, path: Optional[str] = None) -> list:
+        dbfs_path = self._get_dbfs_path(path) if path else self._get_dbfs_path("")
+        dbfs_list_json = {"path": dbfs_path}
+        response = self._dbfs_list_api(dbfs_list_json)
+        try:
+            json_response = json.loads(response.text)
+        except ValueError:
+            raise MlflowException(
+                f"API request to list files under DBFS path {dbfs_path} failed with "
+                f"status code {response.status_code}. Response body: {response.text}"
+            )
+        # /api/2.0/dbfs/list will not have the 'files' key in the response for empty directories.
+        infos = []
+        artifact_prefix = strip_prefix(self.artifact_uri, "dbfs:")
+        if json_response.get("error_code", None) == RESOURCE_NON_EXISTENT:
+            return []
+        dbfs_files = json_response.get("files", [])
+        for dbfs_file in dbfs_files:
+            stripped_path = strip_prefix(dbfs_file["path"], artifact_prefix + "/")
+            # If `path` is a file, the DBFS list API returns a single list element with the
+            # same name as `path`. The list_artifacts API expects us to return an empty list in this
+            # case, so we do so here.
+            if stripped_path == path:
+                return []
+            is_dir = dbfs_file["is_dir"]
+            artifact_size = None if is_dir else dbfs_file["file_size"]
+            infos.append(FileInfo(stripped_path, is_dir, artifact_size))
+        return sorted(infos, key=lambda f: f.path)
+
+    def _download_file(self, remote_file_path, local_path):
+        self._dbfs_download(
+            output_path=local_path, endpoint=self._get_dbfs_endpoint(remote_file_path)
+        )
+
+    def delete_artifacts(self, artifact_path=None):
+        raise MlflowException("Not implemented yet")
+
+
+def _get_host_creds_from_default_store():
+    store = utils._get_store()
+    if not isinstance(store, RestStore):
+        raise MlflowException(
+            "Failed to get credentials for DBFS; they are read from the "
+            + "Databricks CLI credentials or MLFLOW_TRACKING* environment "
+            + "variables."
+        )
+    return store.get_host_creds
+
+
+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.
+
+    Args:
+        artifact_uri: DBFS root artifact URI.
+
+    Returns:
+        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 MLFLOW_ENABLE_DBFS_FUSE_ARTIFACT_REPO.get()
+        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)

venv/lib/python3.9/site-packages/mlflow/store/artifact/ftp_artifact_repo.py

@@ -0,0 +1,133 @@
+import ftplib
+import os
+import posixpath
+import urllib.parse
+from contextlib import contextmanager
+from ftplib import FTP
+from urllib.parse import unquote
+
+from mlflow.entities.file_info import FileInfo
+from mlflow.exceptions import MlflowException
+from mlflow.store.artifact.artifact_repo import ArtifactRepository
+from mlflow.utils.file_utils import relative_path_to_artifact_path
+
+
+class FTPArtifactRepository(ArtifactRepository):
+    """Stores artifacts as files in a remote directory, via ftp."""
+
+    def __init__(self, artifact_uri):
+        self.uri = artifact_uri
+        parsed = urllib.parse.urlparse(artifact_uri)
+        self.config = {
+            "host": parsed.hostname,
+            "port": 21 if parsed.port is None else parsed.port,
+            "username": parsed.username,
+            "password": parsed.password,
+        }
+        self.path = parsed.path or "/"
+
+        if self.config["host"] is None:
+            self.config["host"] = "localhost"
+        if self.config["password"] is None:
+            self.config["password"] = ""
+        else:
+            self.config["password"] = unquote(parsed.password)
+
+        super().__init__(artifact_uri)
+
+    @contextmanager
+    def get_ftp_client(self):
+        ftp = FTP()
+        ftp.connect(self.config["host"], self.config["port"])
+        ftp.login(self.config["username"], self.config["password"])
+        yield ftp
+        ftp.close()
+
+    @staticmethod
+    def _is_dir(ftp, full_file_path):
+        try:
+            ftp.cwd(full_file_path)
+            return True
+        except ftplib.error_perm:
+            return False
+
+    @staticmethod
+    def _mkdir(ftp, artifact_dir):
+        try:
+            if not FTPArtifactRepository._is_dir(ftp, artifact_dir):
+                ftp.mkd(artifact_dir)
+        except ftplib.error_perm:
+            head, _ = posixpath.split(artifact_dir)
+            FTPArtifactRepository._mkdir(ftp, head)
+            FTPArtifactRepository._mkdir(ftp, artifact_dir)
+
+    @staticmethod
+    def _size(ftp, full_file_path):
+        ftp.voidcmd("TYPE I")
+        size = ftp.size(full_file_path)
+        ftp.voidcmd("TYPE A")
+        return size
+
+    def log_artifact(self, local_file, artifact_path=None):
+        with self.get_ftp_client() as ftp:
+            artifact_dir = posixpath.join(self.path, artifact_path) if artifact_path else self.path
+            self._mkdir(ftp, artifact_dir)
+            with open(local_file, "rb") as f:
+                ftp.cwd(artifact_dir)
+                ftp.storbinary("STOR " + os.path.basename(local_file), f)
+
+    def log_artifacts(self, local_dir, artifact_path=None):
+        dest_path = posixpath.join(self.path, artifact_path) if artifact_path else self.path
+
+        local_dir = os.path.abspath(local_dir)
+        for root, _, filenames in os.walk(local_dir):
+            upload_path = dest_path
+            if root != local_dir:
+                rel_path = os.path.relpath(root, local_dir)
+                rel_upload_path = relative_path_to_artifact_path(rel_path)
+                upload_path = posixpath.join(dest_path, rel_upload_path)
+            if not filenames:
+                with self.get_ftp_client() as ftp:
+                    self._mkdir(ftp, upload_path)
+            for f in filenames:
+                if os.path.isfile(os.path.join(root, f)):
+                    self.log_artifact(os.path.join(root, f), upload_path)
+
+    def _is_directory(self, artifact_path):
+        artifact_dir = self.path
+        list_dir = posixpath.join(artifact_dir, artifact_path) if artifact_path else artifact_dir
+        with self.get_ftp_client() as ftp:
+            return self._is_dir(ftp, list_dir)
+
+    def list_artifacts(self, path=None):
+        with self.get_ftp_client() as ftp:
+            artifact_dir = self.path
+            list_dir = posixpath.join(artifact_dir, path) if path else artifact_dir
+            if not self._is_dir(ftp, list_dir):
+                return []
+            artifact_files = ftp.nlst(list_dir)
+            # Make sure artifact_files is a list of file names because ftp.nlst
+            # may return absolute paths.
+            artifact_files = [os.path.basename(f) for f in artifact_files]
+            artifact_files = list(filter(lambda x: x != "." and x != "..", artifact_files))
+            infos = []
+            for file_name in artifact_files:
+                file_path = file_name if path is None else posixpath.join(path, file_name)
+                full_file_path = posixpath.join(list_dir, file_name)
+                if self._is_dir(ftp, full_file_path):
+                    infos.append(FileInfo(file_path, True, None))
+                else:
+                    size = self._size(ftp, full_file_path)
+                    infos.append(FileInfo(file_path, False, size))
+        return infos
+
+    def _download_file(self, remote_file_path, local_path):
+        remote_full_path = (
+            posixpath.join(self.path, remote_file_path) if remote_file_path else self.path
+        )
+        with self.get_ftp_client() as ftp:
+            with open(local_path, "wb") as f:
+                ftp.retrbinary("RETR " + remote_full_path, f.write)
+
+    def delete_artifacts(self, artifact_path=None):
+        raise MlflowException("Not implemented yet")

venv/lib/python3.9/site-packages/mlflow/store/artifact/gcs_artifact_repo.py

@@ -0,0 +1,292 @@
+import datetime
+import importlib.metadata
+import os
+import posixpath
+import urllib.parse
+from collections import namedtuple
+
+from packaging.version import Version
+
+from mlflow.entities import FileInfo
+from mlflow.entities.multipart_upload import (
+    CreateMultipartUploadResponse,
+    MultipartUploadCredential,
+)
+from mlflow.environment_variables import (
+    MLFLOW_ARTIFACT_UPLOAD_DOWNLOAD_TIMEOUT,
+    MLFLOW_GCS_DEFAULT_TIMEOUT,
+    MLFLOW_GCS_DOWNLOAD_CHUNK_SIZE,
+    MLFLOW_GCS_UPLOAD_CHUNK_SIZE,
+)
+from mlflow.exceptions import _UnsupportedMultipartUploadException
+from mlflow.store.artifact.artifact_repo import (
+    ArtifactRepository,
+    MultipartUploadMixin,
+    _retry_with_new_creds,
+)
+from mlflow.utils.file_utils import relative_path_to_artifact_path
+
+GCSMPUArguments = namedtuple("GCSMPUArguments", ["transport", "url", "headers", "content_type"])
+
+
+class GCSArtifactRepository(ArtifactRepository, MultipartUploadMixin):
+    """
+    Stores artifacts on Google Cloud Storage.
+
+    Args:
+        artifact_uri: URI of GCS bucket
+        client: Optional. The client to use for GCS operations; a default
+            client object will be created if unspecified, using default
+            credentials as described in https://google-cloud.readthedocs.io/en/latest/core/auth.html
+    """
+
+    def __init__(self, artifact_uri, client=None, credential_refresh_def=None):
+        super().__init__(artifact_uri)
+        from google.auth.exceptions import DefaultCredentialsError
+        from google.cloud import storage as gcs_storage
+        from google.cloud.storage.constants import _DEFAULT_TIMEOUT
+
+        self._GCS_DOWNLOAD_CHUNK_SIZE = MLFLOW_GCS_DOWNLOAD_CHUNK_SIZE.get()
+        self._GCS_UPLOAD_CHUNK_SIZE = MLFLOW_GCS_UPLOAD_CHUNK_SIZE.get()
+        self._GCS_DEFAULT_TIMEOUT = (
+            MLFLOW_ARTIFACT_UPLOAD_DOWNLOAD_TIMEOUT.get()
+            or MLFLOW_GCS_DEFAULT_TIMEOUT.get()
+            or _DEFAULT_TIMEOUT
+        )
+        # Method to use for refresh
+        self.credential_refresh_def = credential_refresh_def
+        # If the user-supplied timeout environment variable value is -1,
+        # use `None` for `self._GCS_DEFAULT_TIMEOUT`
+        # to use indefinite timeout
+        self._GCS_DEFAULT_TIMEOUT = (
+            None if self._GCS_DEFAULT_TIMEOUT == -1 else self._GCS_DEFAULT_TIMEOUT
+        )
+        if client is not None:
+            self.client = client
+        else:
+            try:
+                self.client = gcs_storage.Client()
+            except DefaultCredentialsError:
+                self.client = gcs_storage.Client.create_anonymous_client()
+
+    @staticmethod
+    def parse_gcs_uri(uri):
+        """Parse an GCS URI, returning (bucket, path)"""
+        parsed = urllib.parse.urlparse(uri)
+        if parsed.scheme != "gs":
+            raise Exception(f"Not a GCS URI: {uri}")
+        path = parsed.path
+        if path.startswith("/"):
+            path = path[1:]
+        return parsed.netloc, path
+
+    def _get_bucket(self, bucket):
+        return self.client.bucket(bucket)
+
+    def _refresh_credentials(self):
+        from google.cloud.storage import Client
+        from google.oauth2.credentials import Credentials
+
+        (bucket, _) = self.parse_gcs_uri(self.artifact_uri)
+        if not self.credential_refresh_def:
+            return self._get_bucket(bucket)
+        new_token = self.credential_refresh_def()
+        credentials = Credentials(new_token["oauth_token"])
+        self.client = Client(project="mlflow", credentials=credentials)
+        return self._get_bucket(bucket)
+
+    def log_artifact(self, local_file, artifact_path=None):
+        (bucket, dest_path) = self.parse_gcs_uri(self.artifact_uri)
+        if artifact_path:
+            dest_path = posixpath.join(dest_path, artifact_path)
+        dest_path = posixpath.join(dest_path, os.path.basename(local_file))
+
+        gcs_bucket = self._get_bucket(bucket)
+        blob = gcs_bucket.blob(dest_path, chunk_size=self._GCS_UPLOAD_CHUNK_SIZE)
+        blob.upload_from_filename(local_file, timeout=self._GCS_DEFAULT_TIMEOUT)
+
+    def log_artifacts(self, local_dir, artifact_path=None):
+        (bucket, dest_path) = self.parse_gcs_uri(self.artifact_uri)
+        if artifact_path:
+            dest_path = posixpath.join(dest_path, artifact_path)
+
+        local_dir = os.path.abspath(local_dir)
+
+        for root, _, filenames in os.walk(local_dir):
+            upload_path = dest_path
+            if root != local_dir:
+                rel_path = os.path.relpath(root, local_dir)
+                rel_path = relative_path_to_artifact_path(rel_path)
+                upload_path = posixpath.join(dest_path, rel_path)
+            for f in filenames:
+                gcs_bucket = self._get_bucket(bucket)
+                path = posixpath.join(upload_path, f)
+                # For large models, we need to speculatively retry a credential refresh
+                # and throw if it still fails.  We cannot use the built-in refresh because UC
+                # does not return a refresh token with the oauth token
+                file_name = os.path.join(root, f)
+
+                def try_func(gcs_bucket):
+                    gcs_bucket.blob(
+                        path, chunk_size=self._GCS_UPLOAD_CHUNK_SIZE
+                    ).upload_from_filename(file_name, timeout=self._GCS_DEFAULT_TIMEOUT)
+
+                _retry_with_new_creds(
+                    try_func=try_func, creds_func=self._refresh_credentials, orig_creds=gcs_bucket
+                )
+
+    def list_artifacts(self, path=None):
+        (bucket, artifact_path) = self.parse_gcs_uri(self.artifact_uri)
+        dest_path = artifact_path
+        if path:
+            dest_path = posixpath.join(dest_path, path)
+        prefix = dest_path if dest_path.endswith("/") else dest_path + "/"
+
+        bkt = self._get_bucket(bucket)
+
+        infos = self._list_folders(bkt, prefix, artifact_path)
+
+        results = bkt.list_blobs(prefix=prefix, delimiter="/")
+        for result in results:
+            # skip blobs matching current directory path as list_blobs api
+            # returns subdirectories as well
+            if result.name == prefix:
+                continue
+            blob_path = result.name[len(artifact_path) + 1 :]
+            infos.append(FileInfo(blob_path, False, result.size))
+
+        return sorted(infos, key=lambda f: f.path)
+
+    def _list_folders(self, bkt, prefix, artifact_path):
+        results = bkt.list_blobs(prefix=prefix, delimiter="/")
+        dir_paths = set()
+        for page in results.pages:
+            dir_paths.update(page.prefixes)
+
+        return [FileInfo(path[len(artifact_path) + 1 : -1], True, None) for path in dir_paths]
+
+    def _download_file(self, remote_file_path, local_path):
+        (bucket, remote_root_path) = self.parse_gcs_uri(self.artifact_uri)
+        remote_full_path = posixpath.join(remote_root_path, remote_file_path)
+        gcs_bucket = self._get_bucket(bucket)
+        gcs_bucket.blob(
+            remote_full_path, chunk_size=self._GCS_DOWNLOAD_CHUNK_SIZE
+        ).download_to_filename(local_path, timeout=self._GCS_DEFAULT_TIMEOUT)
+
+    def delete_artifacts(self, artifact_path=None):
+        (bucket_name, dest_path) = self.parse_gcs_uri(self.artifact_uri)
+        if artifact_path:
+            dest_path = posixpath.join(dest_path, artifact_path)
+
+        gcs_bucket = self._get_bucket(bucket_name)
+        blobs = gcs_bucket.list_blobs(prefix=f"{dest_path}")
+        for blob in blobs:
+            blob.delete()
+
+    @staticmethod
+    def _validate_support_mpu():
+        if Version(importlib.metadata.version("google-cloud-storage")) < Version(
+            "2.12.0"
+        ) or Version(importlib.metadata.version("google-resumable-media")) < Version("2.6.0"):
+            raise _UnsupportedMultipartUploadException()
+
+    @staticmethod
+    def _gcs_mpu_arguments(filename: str, blob) -> GCSMPUArguments:
+        """See :py:func:`google.cloud.storage.transfer_manager.upload_chunks_concurrently`"""
+        from google.cloud.storage.transfer_manager import _headers_from_metadata
+
+        bucket = blob.bucket
+        client = blob.client
+        transport = blob._get_transport(client)
+
+        hostname = client._connection.get_api_base_url_for_mtls()
+        url = f"{hostname}/{bucket.name}/{blob.name}"
+
+        base_headers, object_metadata, content_type = blob._get_upload_arguments(
+            client, None, filename=filename, command="tm.upload_sharded"
+        )
+        headers = {**base_headers, **_headers_from_metadata(object_metadata)}
+
+        if blob.user_project is not None:
+            headers["x-goog-user-project"] = blob.user_project
+
+        if blob.kms_key_name is not None and "cryptoKeyVersions" not in blob.kms_key_name:
+            headers["x-goog-encryption-kms-key-name"] = blob.kms_key_name
+
+        return GCSMPUArguments(
+            transport=transport, url=url, headers=headers, content_type=content_type
+        )
+
+    def create_multipart_upload(self, local_file, num_parts=1, artifact_path=None):
+        self._validate_support_mpu()
+        from google.resumable_media.requests import XMLMPUContainer
+
+        (bucket, dest_path) = self.parse_gcs_uri(self.artifact_uri)
+        if artifact_path:
+            dest_path = posixpath.join(dest_path, artifact_path)
+        dest_path = posixpath.join(dest_path, os.path.basename(local_file))
+
+        gcs_bucket = self._get_bucket(bucket)
+        blob = gcs_bucket.blob(dest_path)
+        args = self._gcs_mpu_arguments(local_file, blob)
+        container = XMLMPUContainer(args.url, local_file, headers=args.headers)
+        container.initiate(transport=args.transport, content_type=args.content_type)
+        upload_id = container.upload_id
+
+        credentials = []
+        for i in range(1, num_parts + 1):  # part number must be in [1, 10000]
+            signed_url = blob.generate_signed_url(
+                method="PUT",
+                version="v4",
+                expiration=datetime.timedelta(minutes=60),
+                query_parameters={
+                    "partNumber": i,
+                    "uploadId": upload_id,
+                },
+            )
+            credentials.append(
+                MultipartUploadCredential(
+                    url=signed_url,
+                    part_number=i,
+                    headers={},
+                )
+            )
+        return CreateMultipartUploadResponse(
+            credentials=credentials,
+            upload_id=upload_id,
+        )
+
+    def complete_multipart_upload(self, local_file, upload_id, parts=None, artifact_path=None):
+        self._validate_support_mpu()
+        from google.resumable_media.requests import XMLMPUContainer
+
+        (bucket, dest_path) = self.parse_gcs_uri(self.artifact_uri)
+        if artifact_path:
+            dest_path = posixpath.join(dest_path, artifact_path)
+        dest_path = posixpath.join(dest_path, os.path.basename(local_file))
+
+        gcs_bucket = self._get_bucket(bucket)
+        blob = gcs_bucket.blob(dest_path)
+        args = self._gcs_mpu_arguments(local_file, blob)
+        container = XMLMPUContainer(args.url, local_file, headers=args.headers)
+        container._upload_id = upload_id
+        for part in parts:
+            container.register_part(part.part_number, part.etag)
+
+        container.finalize(transport=args.transport)
+
+    def abort_multipart_upload(self, local_file, upload_id, artifact_path=None):
+        self._validate_support_mpu()
+        from google.resumable_media.requests import XMLMPUContainer
+
+        (bucket, dest_path) = self.parse_gcs_uri(self.artifact_uri)
+        if artifact_path:
+            dest_path = posixpath.join(dest_path, artifact_path)
+        dest_path = posixpath.join(dest_path, os.path.basename(local_file))
+
+        gcs_bucket = self._get_bucket(bucket)
+        blob = gcs_bucket.blob(dest_path)
+        args = self._gcs_mpu_arguments(local_file, blob)
+        container = XMLMPUContainer(args.url, local_file, headers=args.headers)
+        container._upload_id = upload_id
+        container.cancel(transport=args.transport)

venv/lib/python3.9/site-packages/mlflow/store/artifact/hdfs_artifact_repo.py

@@ -0,0 +1,208 @@
+import os
+import posixpath
+import urllib.parse
+from contextlib import contextmanager
+
+try:
+    from pyarrow.fs import FileSelector, FileType, HadoopFileSystem
+except ImportError:
+    pass
+
+from mlflow.entities import FileInfo
+from mlflow.environment_variables import (
+    MLFLOW_KERBEROS_TICKET_CACHE,
+    MLFLOW_KERBEROS_USER,
+    MLFLOW_PYARROW_EXTRA_CONF,
+)
+from mlflow.store.artifact.artifact_repo import ArtifactRepository
+from mlflow.utils.file_utils import relative_path_to_artifact_path
+
+
+class HdfsArtifactRepository(ArtifactRepository):
+    """
+    Stores artifacts on HDFS.
+
+    This repository is used with URIs of the form ``hdfs:/<path>``. The repository can only be used
+    together with the RestStore.
+    """
+
+    def __init__(self, artifact_uri):
+        self.scheme, self.host, self.port, self.path = _resolve_connection_params(artifact_uri)
+        super().__init__(artifact_uri)
+
+    def log_artifact(self, local_file, artifact_path=None):
+        """
+        Log artifact in hdfs.
+
+        Args:
+            local_file: Source file path.
+            artifact_path: When specified will attempt to write under artifact_uri/artifact_path.
+        """
+        hdfs_base_path = _resolve_base_path(self.path, artifact_path)
+
+        with hdfs_system(scheme=self.scheme, host=self.host, port=self.port) as hdfs:
+            _, file_name = os.path.split(local_file)
+            destination_path = posixpath.join(hdfs_base_path, file_name)
+            with open(local_file, "rb") as source:
+                with hdfs.open_output_stream(destination_path) as destination:
+                    destination.write(source.read())
+
+    def log_artifacts(self, local_dir, artifact_path=None):
+        """
+        Log artifacts in hdfs.
+        Missing remote sub-directories will be created if needed.
+
+        Args:
+            local_dir: Source dir path.
+            artifact_path: When specified will attempt to write under artifact_uri/artifact_path.
+        """
+        hdfs_base_path = _resolve_base_path(self.path, artifact_path)
+
+        with hdfs_system(scheme=self.scheme, host=self.host, port=self.port) as hdfs:
+            if not hdfs.get_file_info(hdfs_base_path).type == FileType.Directory:
+                hdfs.create_dir(hdfs_base_path, recursive=True)
+
+            for subdir_path, _, files in os.walk(local_dir):
+                relative_path = _relative_path_local(local_dir, subdir_path)
+
+                hdfs_subdir_path = (
+                    posixpath.join(hdfs_base_path, relative_path)
+                    if relative_path
+                    else hdfs_base_path
+                )
+
+                if not hdfs.get_file_info(hdfs_subdir_path).type == FileType.Directory:
+                    hdfs.create_dir(hdfs_subdir_path, recursive=True)
+
+                for each_file in files:
+                    source_path = os.path.join(subdir_path, each_file)
+                    destination_path = posixpath.join(hdfs_subdir_path, each_file)
+                    with open(source_path, "rb") as source:
+                        with hdfs.open_output_stream(destination_path) as destination:
+                            destination.write(source.read())
+
+    def list_artifacts(self, path=None):
+        """
+        Lists files and directories under artifacts directory for the current run_id.
+        (self.path contains the base path - hdfs:/some/path/run_id/artifacts)
+
+        Args:
+            path: Relative source path. Possible subdirectory existing under
+                hdfs:/some/path/run_id/artifacts
+
+        Returns:
+            List of FileInfos under given path
+        """
+        hdfs_base_path = _resolve_base_path(self.path, path)
+
+        with hdfs_system(scheme=self.scheme, host=self.host, port=self.port) as hdfs:
+            paths = []
+            base_info = hdfs.get_file_info(hdfs_base_path)
+            if base_info.type == FileType.Directory:
+                selector = FileSelector(hdfs_base_path)
+            elif base_info.type == FileType.File:
+                selector = [hdfs_base_path]
+            else:
+                return []
+
+            for file_detail in hdfs.get_file_info(selector):
+                file_name = file_detail.path
+
+                # file_name is hdfs_base_path and not a child of that path
+                if file_name == hdfs_base_path:
+                    continue
+
+                # Strip off anything that comes before the artifact root e.g. hdfs://name
+                offset = file_name.index(self.path)
+                rel_path = _relative_path_remote(self.path, file_name[offset:])
+                is_dir = file_detail.type == FileType.Directory
+                size = file_detail.size
+                paths.append(FileInfo(rel_path, is_dir=is_dir, file_size=size))
+            return sorted(paths, key=lambda f: paths)
+
+    def _is_directory(self, artifact_path):
+        hdfs_base_path = _resolve_base_path(self.path, artifact_path)
+        with hdfs_system(scheme=self.scheme, host=self.host, port=self.port) as hdfs:
+            return hdfs.get_file_info(hdfs_base_path).type == FileType.Directory
+
+    def _download_file(self, remote_file_path, local_path):
+        hdfs_base_path = _resolve_base_path(self.path, remote_file_path)
+        with hdfs_system(scheme=self.scheme, host=self.host, port=self.port) as hdfs:
+            with hdfs.open_input_stream(hdfs_base_path) as source:
+                with open(local_path, "wb") as destination:
+                    destination.write(source.read())
+
+    def delete_artifacts(self, artifact_path=None):
+        path = posixpath.join(self.path, artifact_path) if artifact_path else self.path
+        with hdfs_system(scheme=self.scheme, host=self.host, port=self.port) as hdfs:
+            file_info = hdfs.get_file_info(path)
+            if file_info.type == FileType.File:
+                hdfs.delete_file(path)
+            elif file_info.type == FileType.Directory:
+                hdfs.delete_dir_contents(path)
+
+
+@contextmanager
+def hdfs_system(scheme, host, port):
+    """
+    hdfs system context - Attempt to establish the connection to hdfs
+    and yields HadoopFileSystem
+
+    Args:
+        scheme: scheme or use hdfs:// as default
+        host: hostname or when relaying on the core-site.xml config use 'default'
+        port: port or when relaying on the core-site.xml config use 0
+    """
+    kerb_ticket = MLFLOW_KERBEROS_TICKET_CACHE.get()
+    kerberos_user = MLFLOW_KERBEROS_USER.get()
+    extra_conf = _parse_extra_conf(MLFLOW_PYARROW_EXTRA_CONF.get())
+
+    host = scheme + "://" + host if host else "default"
+
+    yield HadoopFileSystem(
+        host=host,
+        port=port or 0,
+        user=kerberos_user,
+        kerb_ticket=kerb_ticket,
+        extra_conf=extra_conf,
+    )
+
+
+def _resolve_connection_params(artifact_uri):
+    parsed = urllib.parse.urlparse(artifact_uri)
+
+    return parsed.scheme, parsed.hostname, parsed.port, parsed.path
+
+
+def _resolve_base_path(path, artifact_path):
+    if path == artifact_path:
+        return path
+    if artifact_path:
+        return posixpath.join(path, artifact_path)
+    return path
+
+
+def _relative_path(base_dir, subdir_path, path_module):
+    relative_path = path_module.relpath(subdir_path, base_dir)
+    return relative_path if relative_path != "." else None
+
+
+def _relative_path_local(base_dir, subdir_path):
+    rel_path = _relative_path(base_dir, subdir_path, os.path)
+    return relative_path_to_artifact_path(rel_path) if rel_path is not None else None
+
+
+def _relative_path_remote(base_dir, subdir_path):
+    return _relative_path(base_dir, subdir_path, posixpath)
+
+
+def _parse_extra_conf(extra_conf):
+    if extra_conf:
+
+        def as_pair(config):
+            key, val = config.split("=")
+            return key, val
+
+        list_of_key_val = [as_pair(conf) for conf in extra_conf.split(",")]
+        return dict(list_of_key_val)
+    return None

venv/lib/python3.9/site-packages/mlflow/store/artifact/http_artifact_repo.py

@@ -0,0 +1,218 @@
+import logging
+import os
+import posixpath
+
+import requests
+from requests import HTTPError
+
+from mlflow.entities import FileInfo
+from mlflow.entities.multipart_upload import (
+    CreateMultipartUploadResponse,
+    MultipartUploadCredential,
+    MultipartUploadPart,
+)
+from mlflow.environment_variables import (
+    MLFLOW_ENABLE_PROXY_MULTIPART_UPLOAD,
+    MLFLOW_MULTIPART_UPLOAD_CHUNK_SIZE,
+    MLFLOW_MULTIPART_UPLOAD_MINIMUM_FILE_SIZE,
+)
+from mlflow.exceptions import MlflowException, _UnsupportedMultipartUploadException
+from mlflow.store.artifact.artifact_repo import (
+    ArtifactRepository,
+    MultipartUploadMixin,
+    verify_artifact_path,
+)
+from mlflow.store.artifact.cloud_artifact_repo import _complete_futures, _compute_num_chunks
+from mlflow.utils.credentials import get_default_host_creds
+from mlflow.utils.file_utils import read_chunk, relative_path_to_artifact_path
+from mlflow.utils.mime_type_utils import _guess_mime_type
+from mlflow.utils.rest_utils import augmented_raise_for_status, http_request
+from mlflow.utils.uri import validate_path_is_safe
+
+_logger = logging.getLogger(__name__)
+
+
+class HttpArtifactRepository(ArtifactRepository, MultipartUploadMixin):
+    """Stores artifacts in a remote artifact storage using HTTP requests"""
+
+    @property
+    def _host_creds(self):
+        return get_default_host_creds(self.artifact_uri)
+
+    def log_artifact(self, local_file, artifact_path=None):
+        verify_artifact_path(artifact_path)
+
+        # Try to perform multipart upload if the file is large.
+        # If the server does not support, or if the upload failed, revert to normal upload.
+        if (
+            MLFLOW_ENABLE_PROXY_MULTIPART_UPLOAD.get()
+            and os.path.getsize(local_file) >= MLFLOW_MULTIPART_UPLOAD_MINIMUM_FILE_SIZE.get()
+        ):
+            try:
+                self._try_multipart_upload(local_file, artifact_path)
+                return
+            except _UnsupportedMultipartUploadException:
+                pass
+
+        file_name = os.path.basename(local_file)
+        mime_type = _guess_mime_type(file_name)
+        paths = (artifact_path, file_name) if artifact_path else (file_name,)
+        endpoint = posixpath.join("/", *paths)
+        extra_headers = {"Content-Type": mime_type}
+        with open(local_file, "rb") as f:
+            resp = http_request(
+                self._host_creds, endpoint, "PUT", data=f, extra_headers=extra_headers
+            )
+            augmented_raise_for_status(resp)
+
+    def log_artifacts(self, local_dir, artifact_path=None):
+        local_dir = os.path.abspath(local_dir)
+        for root, _, filenames in os.walk(local_dir):
+            if root == local_dir:
+                artifact_dir = artifact_path
+            else:
+                rel_path = os.path.relpath(root, local_dir)
+                rel_path = relative_path_to_artifact_path(rel_path)
+                artifact_dir = (
+                    posixpath.join(artifact_path, rel_path) if artifact_path else rel_path
+                )
+            for f in filenames:
+                self.log_artifact(os.path.join(root, f), artifact_dir)
+
+    def list_artifacts(self, path=None):
+        endpoint = "/mlflow-artifacts/artifacts"
+        url, tail = self.artifact_uri.split(endpoint, maxsplit=1)
+        root = tail.lstrip("/")
+        params = {"path": posixpath.join(root, path) if path else root}
+        host_creds = get_default_host_creds(url)
+        resp = http_request(host_creds, endpoint, "GET", params=params)
+        augmented_raise_for_status(resp)
+        file_infos = []
+        for f in resp.json().get("files", []):
+            validated_path = validate_path_is_safe(f["path"])
+            file_info = FileInfo(
+                posixpath.join(path, validated_path) if path else validated_path,
+                f["is_dir"],
+                int(f["file_size"]) if ("file_size" in f) else None,
+            )
+            file_infos.append(file_info)
+
+        return sorted(file_infos, key=lambda f: f.path)
+
+    def _download_file(self, remote_file_path, local_path):
+        endpoint = posixpath.join("/", remote_file_path)
+        resp = http_request(self._host_creds, endpoint, "GET", stream=True)
+        augmented_raise_for_status(resp)
+        with open(local_path, "wb") as f:
+            chunk_size = 1024 * 1024  # 1 MB
+            for chunk in resp.iter_content(chunk_size=chunk_size):
+                f.write(chunk)
+
+    def delete_artifacts(self, artifact_path=None):
+        endpoint = posixpath.join("/", artifact_path) if artifact_path else "/"
+        resp = http_request(self._host_creds, endpoint, "DELETE", stream=True)
+        augmented_raise_for_status(resp)
+
+    def _construct_mpu_uri_and_path(self, base_endpoint, artifact_path):
+        uri, path = self.artifact_uri.split("/mlflow-artifacts/artifacts", maxsplit=1)
+        path = path.strip("/")
+        endpoint = (
+            posixpath.join(base_endpoint, path, artifact_path)
+            if artifact_path
+            else posixpath.join(base_endpoint, path)
+        )
+        return uri, endpoint
+
+    def create_multipart_upload(self, local_file, num_parts=1, artifact_path=None):
+        uri, endpoint = self._construct_mpu_uri_and_path(
+            "/mlflow-artifacts/mpu/create", artifact_path
+        )
+        host_creds = get_default_host_creds(uri)
+        params = {
+            "path": local_file,
+            "num_parts": num_parts,
+        }
+        resp = http_request(host_creds, endpoint, "POST", json=params)
+        augmented_raise_for_status(resp)
+        return CreateMultipartUploadResponse.from_dict(resp.json())
+
+    def complete_multipart_upload(self, local_file, upload_id, parts=None, artifact_path=None):
+        uri, endpoint = self._construct_mpu_uri_and_path(
+            "/mlflow-artifacts/mpu/complete", artifact_path
+        )
+        host_creds = get_default_host_creds(uri)
+        params = {
+            "path": local_file,
+            "upload_id": upload_id,
+            "parts": [part.to_dict() for part in parts],
+        }
+        resp = http_request(host_creds, endpoint, "POST", json=params)
+        augmented_raise_for_status(resp)
+
+    def abort_multipart_upload(self, local_file, upload_id, artifact_path=None):
+        uri, endpoint = self._construct_mpu_uri_and_path(
+            "/mlflow-artifacts/mpu/abort", artifact_path
+        )
+        host_creds = get_default_host_creds(uri)
+        params = {
+            "path": local_file,
+            "upload_id": upload_id,
+        }
+        resp = http_request(host_creds, endpoint, "POST", json=params)
+        augmented_raise_for_status(resp)
+
+    @staticmethod
+    def _upload_part(credential: MultipartUploadCredential, local_file, size, start_byte):
+        data = read_chunk(local_file, size, start_byte)
+        response = requests.put(credential.url, data=data, headers=credential.headers)
+        augmented_raise_for_status(response)
+        return MultipartUploadPart(
+            part_number=credential.part_number,
+            etag=response.headers.get("ETag", ""),
+            url=credential.url,
+        )
+
+    def _try_multipart_upload(self, local_file, artifact_path=None):
+        """
+        Attempts to perform multipart upload to log an artifact.
+        Returns if the multipart upload is successful.
+        Raises UnsupportedMultipartUploadException if multipart upload is unsupported.
+        """
+        chunk_size = MLFLOW_MULTIPART_UPLOAD_CHUNK_SIZE.get()
+        num_parts = _compute_num_chunks(local_file, chunk_size)
+
+        try:
+            create = self.create_multipart_upload(local_file, num_parts, artifact_path)
+        except HTTPError as e:
+            # return False if server does not support multipart upload
+            error_message = e.response.json().get("message", "")
+            if isinstance(error_message, str) and error_message.startswith(
+                _UnsupportedMultipartUploadException.MESSAGE
+            ):
+                raise _UnsupportedMultipartUploadException()
+            raise
+
+        try:
+            futures = {}
+            for i, credential in enumerate(create.credentials):
+                future = self.thread_pool.submit(
+                    self._upload_part,
+                    credential=credential,
+                    local_file=local_file,
+                    size=chunk_size,
+                    start_byte=chunk_size * i,
+                )
+                futures[future] = credential.part_number
+
+            parts, errors = _complete_futures(futures, local_file)
+            if errors:
+                raise MlflowException(
+                    f"Failed to upload at least one part of {local_file}. Errors: {errors}"
+                )
+
+            parts = sorted(parts.values(), key=lambda part: part.part_number)
+            self.complete_multipart_upload(local_file, create.upload_id, parts, artifact_path)
+        except Exception as e:
+            self.abort_multipart_upload(local_file, create.upload_id, artifact_path)
+            _logger.warning(f"Failed to upload file {local_file} using multipart upload: {e}")
+            raise

venv/lib/python3.9/site-packages/mlflow/store/artifact/local_artifact_repo.py

@@ -0,0 +1,142 @@
+import os
+import shutil
+from typing import Any
+
+from mlflow.store.artifact.artifact_repo import (
+    ArtifactRepository,
+    try_read_trace_data,
+    verify_artifact_path,
+)
+from mlflow.tracing.artifact_utils import TRACE_DATA_FILE_NAME
+from mlflow.utils.file_utils import (
+    get_file_info,
+    list_all,
+    local_file_uri_to_path,
+    mkdir,
+    relative_path_to_artifact_path,
+)
+from mlflow.utils.uri import validate_path_is_safe
+
+
+class LocalArtifactRepository(ArtifactRepository):
+    """Stores artifacts as files in a local directory."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._artifact_dir = local_file_uri_to_path(self.artifact_uri)
+
+    @property
+    def artifact_dir(self):
+        return self._artifact_dir
+
+    def log_artifact(self, local_file, artifact_path=None):
+        verify_artifact_path(artifact_path)
+        # NOTE: The artifact_path is expected to be in posix format.
+        # Posix paths work fine on windows but just in case we normalize it here.
+        if artifact_path:
+            artifact_path = os.path.normpath(artifact_path)
+
+        artifact_dir = (
+            os.path.join(self.artifact_dir, artifact_path) if artifact_path else self.artifact_dir
+        )
+        if not os.path.exists(artifact_dir):
+            mkdir(artifact_dir)
+        try:
+            shutil.copy2(local_file, os.path.join(artifact_dir, os.path.basename(local_file)))
+        except shutil.SameFileError:
+            pass
+
+    def _is_directory(self, artifact_path):
+        # NOTE: The path is expected to be in posix format.
+        # Posix paths work fine on windows but just in case we normalize it here.
+        path = os.path.normpath(artifact_path) if artifact_path else ""
+        list_dir = os.path.join(self.artifact_dir, path) if path else self.artifact_dir
+        return os.path.isdir(list_dir)
+
+    def log_artifacts(self, local_dir, artifact_path=None):
+        verify_artifact_path(artifact_path)
+        # NOTE: The artifact_path is expected to be in posix format.
+        # Posix paths work fine on windows but just in case we normalize it here.
+        if artifact_path:
+            artifact_path = os.path.normpath(artifact_path)
+        artifact_dir = (
+            os.path.join(self.artifact_dir, artifact_path) if artifact_path else self.artifact_dir
+        )
+        if not os.path.exists(artifact_dir):
+            mkdir(artifact_dir)
+        shutil.copytree(src=local_dir, dst=artifact_dir, dirs_exist_ok=True)
+
+    def download_artifacts(self, artifact_path, dst_path=None):
+        """
+        Artifacts tracked by ``LocalArtifactRepository`` already exist on the local filesystem.
+        If ``dst_path`` is ``None``, the absolute filesystem path of the specified artifact is
+        returned. If ``dst_path`` is not ``None``, the local artifact is copied to ``dst_path``.
+
+        Args:
+            artifact_path: Relative source path to the desired artifacts.
+            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 absolute path of the local artifact will be returned.
+
+        Returns:
+            Absolute path of the local filesystem location containing the desired artifacts.
+        """
+        if dst_path:
+            return super().download_artifacts(artifact_path, dst_path)
+        # NOTE: The artifact_path is expected to be a relative path in posix format.
+        # Posix paths work fine on windows but just in case we normalize it here.
+        artifact_path = validate_path_is_safe(artifact_path)
+        local_artifact_path = os.path.join(self.artifact_dir, os.path.normpath(artifact_path))
+        if not os.path.exists(local_artifact_path):
+            raise OSError(f"No such file or directory: '{local_artifact_path}'")
+        return os.path.abspath(local_artifact_path)
+
+    def list_artifacts(self, path=None):
+        # NOTE: The path is expected to be in posix format.
+        # Posix paths work fine on windows but just in case we normalize it here.
+        if path:
+            path = os.path.normpath(path)
+        list_dir = os.path.join(self.artifact_dir, path) if path else self.artifact_dir
+        if os.path.isdir(list_dir):
+            artifact_files = list_all(list_dir, full_path=True)
+            infos = [
+                get_file_info(
+                    f, relative_path_to_artifact_path(os.path.relpath(f, self.artifact_dir))
+                )
+                for f in artifact_files
+            ]
+            return sorted(infos, key=lambda f: f.path)
+        else:
+            return []
+
+    def _download_file(self, remote_file_path, local_path):
+        # NOTE: The remote_file_path is expected to be a relative path in posix format.
+        # Posix paths work fine on windows but just in case we normalize it here.
+        remote_file_path = validate_path_is_safe(remote_file_path)
+        remote_file_path = os.path.join(self.artifact_dir, os.path.normpath(remote_file_path))
+        shutil.copy2(remote_file_path, local_path)
+
+    def delete_artifacts(self, artifact_path=None):
+        artifact_path = local_file_uri_to_path(
+            os.path.join(self._artifact_dir, artifact_path) if artifact_path else self._artifact_dir
+        )
+
+        if os.path.exists(artifact_path):
+            if os.path.isfile(artifact_path):
+                os.remove(artifact_path)
+            else:
+                shutil.rmtree(artifact_path)
+
+    def download_trace_data(self) -> dict[str, Any]:
+        """
+        Download the trace data.
+
+        Returns:
+            The trace data as a dictionary.
+
+        Raises:
+            - `MlflowTraceDataNotFound`: The trace data is not found.
+            - `MlflowTraceDataCorrupted`: The trace data is corrupted.
+        """
+        trace_data_path = os.path.join(self.artifact_dir, TRACE_DATA_FILE_NAME)
+        return try_read_trace_data(trace_data_path)

venv/lib/python3.9/site-packages/mlflow/store/artifact/mlflow_artifacts_repo.py

@@ -0,0 +1,92 @@
+import re
+from urllib.parse import urlparse, urlunparse
+
+from mlflow.exceptions import MlflowException
+from mlflow.store.artifact.http_artifact_repo import HttpArtifactRepository
+from mlflow.tracking._tracking_service.utils import get_tracking_uri
+
+
+def _check_if_host_is_numeric(hostname):
+    if hostname:
+        try:
+            float(hostname)
+            return True
+        except ValueError:
+            return False
+    else:
+        return False
+
+
+def _validate_port_mapped_to_hostname(uri_parse):
+    # This check is to catch an mlflow-artifacts uri that has a port designated but no
+    # hostname specified. `urllib.parse.urlparse` will treat such a uri as a filesystem
+    # definition, mapping the provided port as a hostname value if this condition is not
+    # validated.
+    if uri_parse.hostname and _check_if_host_is_numeric(uri_parse.hostname) and not uri_parse.port:
+        raise MlflowException(
+            "The mlflow-artifacts uri was supplied with a port number: "
+            f"{uri_parse.hostname}, but no host was defined."
+        )
+
+
+def _validate_uri_scheme(parsed_uri):
+    allowable_schemes = {"http", "https"}
+    if parsed_uri.scheme not in allowable_schemes:
+        raise MlflowException(
+            "When an mlflow-artifacts URI was supplied, the tracking URI must be a valid "
+            f"http or https URI, but it was currently set to {parsed_uri.geturl()}. "
+            "Perhaps you forgot to set the tracking URI to the running MLflow server. "
+            "To set the tracking URI, use either of the following methods:\n"
+            "1. Set the MLFLOW_TRACKING_URI environment variable to the desired tracking URI. "
+            "`export MLFLOW_TRACKING_URI=http://localhost:5000`\n"
+            "2. Set the tracking URI programmatically by calling `mlflow.set_tracking_uri`. "
+            "`mlflow.set_tracking_uri('http://localhost:5000')`"
+        )
+
+
+class MlflowArtifactsRepository(HttpArtifactRepository):
+    """Scheme wrapper around HttpArtifactRepository for mlflow-artifacts server functionality"""
+
+    def __init__(self, artifact_uri):
+        super().__init__(self.resolve_uri(artifact_uri, get_tracking_uri()))
+
+    @classmethod
+    def resolve_uri(cls, artifact_uri, tracking_uri):
+        base_url = "/api/2.0/mlflow-artifacts/artifacts"
+
+        track_parse = urlparse(tracking_uri)
+
+        uri_parse = urlparse(artifact_uri)
+
+        # Check to ensure that a port is present with no hostname
+        _validate_port_mapped_to_hostname(uri_parse)
+
+        # Check that tracking uri is http or https
+        _validate_uri_scheme(track_parse)
+
+        if uri_parse.path == "/":  # root directory; build simple path
+            resolved = f"{base_url}{uri_parse.path}"
+        elif uri_parse.path == base_url:  # for operations like list artifacts
+            resolved = base_url
+        else:
+            resolved = f"{track_parse.path}/{base_url}/{uri_parse.path}"
+        resolved = re.sub("//+", "/", resolved)
+
+        resolved_artifacts_uri = urlunparse(
+            (
+                # scheme
+                track_parse.scheme,
+                # netloc
+                uri_parse.netloc if uri_parse.netloc else track_parse.netloc,
+                # path
+                resolved,
+                # params
+                "",
+                # query
+                "",
+                # fragment
+                "",
+            )
+        )
+
+        return resolved_artifacts_uri.replace("///", "/").rstrip("/")

venv/lib/python3.9/site-packages/mlflow/store/artifact/models_artifact_repo.py

@@ -0,0 +1,232 @@
+import logging
+import os
+import urllib.parse
+
+import mlflow
+from mlflow.exceptions import MlflowException
+from mlflow.store.artifact.artifact_repo import ArtifactRepository
+from mlflow.store.artifact.databricks_models_artifact_repo import DatabricksModelsArtifactRepository
+from mlflow.store.artifact.unity_catalog_models_artifact_repo import (
+    UnityCatalogModelsArtifactRepository,
+)
+from mlflow.store.artifact.unity_catalog_oss_models_artifact_repo import (
+    UnityCatalogOSSModelsArtifactRepository,
+)
+from mlflow.store.artifact.utils.models import (
+    get_model_name_and_version,
+    is_using_databricks_registry,
+)
+from mlflow.utils.file_utils import write_yaml
+from mlflow.utils.uri import (
+    add_databricks_profile_info_to_artifact_uri,
+    get_databricks_profile_uri_from_artifact_uri,
+    is_databricks_unity_catalog_uri,
+    is_oss_unity_catalog_uri,
+)
+
+REGISTERED_MODEL_META_FILE_NAME = "registered_model_meta"
+
+_logger = logging.getLogger(__name__)
+
+
+class ModelsArtifactRepository(ArtifactRepository):
+    """
+    Handles artifacts associated with a model version in the model registry via URIs of the form:
+      - `models:/<model_name>/<model_version>`
+      - `models:/<model_name>/<stage>`  (refers to the latest model version in the given stage)
+      - `models:/<model_name>/latest` (refers to the latest of all model versions)
+    It is a light wrapper that resolves the artifact path to an absolute URI then instantiates
+    and uses the artifact repository for that URI.
+    """
+
+    def __init__(self, artifact_uri):
+        from mlflow.store.artifact.artifact_repository_registry import get_artifact_repository
+
+        super().__init__(artifact_uri)
+        registry_uri = mlflow.get_registry_uri()
+        if is_databricks_unity_catalog_uri(uri=registry_uri):
+            self.repo = UnityCatalogModelsArtifactRepository(
+                artifact_uri=artifact_uri, registry_uri=registry_uri
+            )
+            self.model_name = self.repo.model_name
+            self.model_version = self.repo.model_version
+        elif is_oss_unity_catalog_uri(uri=registry_uri):
+            self.repo = UnityCatalogOSSModelsArtifactRepository(
+                artifact_uri=artifact_uri, registry_uri=registry_uri
+            )
+            self.model_name = self.repo.model_name
+            self.model_version = self.repo.model_version
+        elif is_using_databricks_registry(artifact_uri):
+            # Use the DatabricksModelsArtifactRepository if a databricks profile is being used.
+            self.repo = DatabricksModelsArtifactRepository(artifact_uri)
+            self.model_name = self.repo.model_name
+            self.model_version = self.repo.model_version
+        else:
+            (
+                self.model_name,
+                self.model_version,
+                underlying_uri,
+            ) = ModelsArtifactRepository._get_model_uri_infos(artifact_uri)
+            self.repo = get_artifact_repository(underlying_uri)
+            # TODO: it may be nice to fall back to the source URI explicitly here if for some reason
+            #  we don't get a download URI here, or fail during the download itself.
+
+    @staticmethod
+    def is_models_uri(uri):
+        return urllib.parse.urlparse(uri).scheme == "models"
+
+    @staticmethod
+    def split_models_uri(uri):
+        """
+        Split 'models:/<name>/<version>/path/to/model' into
+        ('models:/<name>/<version>', 'path/to/model').
+        Split 'models://<scope>:<prefix>@databricks/<name>/<version>/path/to/model' into
+        ('models://<scope>:<prefix>@databricks/<name>/<version>', 'path/to/model').
+        Split 'models:/<name>@alias/path/to/model' into
+        ('models:/<name>@alias', 'path/to/model').
+        """
+        uri = uri.rstrip("/")
+        parsed_url = urllib.parse.urlparse(uri)
+        path = parsed_url.path
+        netloc = parsed_url.netloc
+        if path.count("/") >= 2 and not path.endswith("/"):
+            splits = path.split("/", 3)
+            cut_index = 2 if "@" in splits[1] else 3
+            model_name_and_version = splits[:cut_index]
+            artifact_path = "/".join(splits[cut_index:])
+            base_part = f"models://{netloc}" if netloc else "models:"
+            return base_part + "/".join(model_name_and_version), artifact_path
+        return uri, ""
+
+    @staticmethod
+    def _get_model_uri_infos(uri):
+        # Note: to support a registry URI that is different from the tracking URI here,
+        # we'll need to add setting of registry URIs via environment variables.
+
+        from mlflow import MlflowClient
+
+        databricks_profile_uri = (
+            get_databricks_profile_uri_from_artifact_uri(uri) or mlflow.get_registry_uri()
+        )
+        client = MlflowClient(registry_uri=databricks_profile_uri)
+        name, version = get_model_name_and_version(client, uri)
+        download_uri = client.get_model_version_download_uri(name, version)
+
+        return (
+            name,
+            version,
+            add_databricks_profile_info_to_artifact_uri(download_uri, databricks_profile_uri),
+        )
+
+    @staticmethod
+    def get_underlying_uri(uri):
+        _, _, underlying_uri = ModelsArtifactRepository._get_model_uri_infos(uri)
+
+        return underlying_uri
+
+    def log_artifact(self, local_file, artifact_path=None):
+        """
+        Log a local file as an artifact, optionally taking an ``artifact_path`` to place it in
+        within the run's artifacts. Run artifacts can be organized into directories, so you can
+        place the artifact in a directory this way.
+
+        Args:
+            local_file: Path to artifact to log.
+            artifact_path: Directory within the run's artifact directory in which to log the
+                artifact.
+        """
+        raise ValueError(
+            "log_artifact is not supported for models:/ URIs. Use register_model instead."
+        )
+
+    def log_artifacts(self, local_dir, artifact_path=None):
+        """
+        Log the files in the specified local directory as artifacts, optionally taking
+        an ``artifact_path`` to place them in within the run's artifacts.
+
+        Args:
+            local_dir: Directory of local artifacts to log.
+            artifact_path: Directory within the run's artifact directory in which to log the
+                artifacts.
+        """
+        raise ValueError(
+            "log_artifacts is not supported for models:/ URIs. Use register_model instead."
+        )
+
+    def list_artifacts(self, path):
+        """
+        Return all the artifacts for this run_id directly under path. If path is a file, returns
+        an empty list. Will error if path is neither a file nor directory.
+
+        Args:
+            path: Relative source path that contain desired artifacts.
+
+        Returns:
+            List of artifacts as FileInfo listed directly under path.
+        """
+        return self.repo.list_artifacts(path)
+
+    def _add_registered_model_meta_file(self, model_path):
+        write_yaml(
+            model_path,
+            REGISTERED_MODEL_META_FILE_NAME,
+            {
+                "model_name": self.model_name,
+                "model_version": self.model_version,
+            },
+            overwrite=True,
+            ensure_yaml_extension=False,
+        )
+
+    def download_artifacts(self, artifact_path, dst_path=None, lineage_header_info=None):
+        """
+        Download an artifact file or directory to a local directory if applicable, and return a
+        local path for it.
+        For registered models, when the artifact is downloaded, the model name and version
+        are saved in the "registered_model_meta" file on the caller's side.
+        The caller is responsible for managing the lifecycle of the downloaded artifacts.
+
+        Args:
+            artifact_path: Relative source path to the desired artifacts.
+            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.
+            lineage_header_info: Linear header information.
+
+        Returns:
+            Absolute path of the local filesystem location containing the desired artifacts.
+        """
+
+        from mlflow.models.model import MLMODEL_FILE_NAME
+
+        # Pass lineage header info if model is registered in UC
+        if isinstance(self.repo, UnityCatalogModelsArtifactRepository):
+            model_path = self.repo.download_artifacts(
+                artifact_path, dst_path, lineage_header_info=lineage_header_info
+            )
+        else:
+            model_path = self.repo.download_artifacts(artifact_path, dst_path)
+        # NB: only add the registered model metadata iff the artifact path is at the root model
+        # directory. For individual files or subdirectories within the model directory, do not
+        # create the metadata file.
+        if os.path.isdir(model_path) and MLMODEL_FILE_NAME in os.listdir(model_path):
+            self._add_registered_model_meta_file(model_path)
+
+        return model_path
+
+    def _download_file(self, remote_file_path, local_path):
+        """
+        Download the file at the specified relative remote path and saves
+        it at the specified local path.
+
+        Args:
+            remote_file_path: Source path to the remote file, relative to the root
+                directory of the artifact repository.
+            local_path: The path to which to save the downloaded file.
+        """
+        self.repo._download_file(remote_file_path, local_path)
+
+    def delete_artifacts(self, artifact_path=None):
+        raise MlflowException("Not implemented yet")

venv/lib/python3.9/site-packages/mlflow/store/artifact/optimized_s3_artifact_repo.py

@@ -0,0 +1,356 @@
+import json
+import logging
+import os
+import posixpath
+import urllib.parse
+from mimetypes import guess_type
+
+from mlflow.entities import FileInfo
+from mlflow.environment_variables import (
+    MLFLOW_ENABLE_MULTIPART_UPLOAD,
+    MLFLOW_MULTIPART_UPLOAD_CHUNK_SIZE,
+    MLFLOW_S3_UPLOAD_EXTRA_ARGS,
+)
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_artifacts_pb2 import ArtifactCredentialInfo
+from mlflow.store.artifact.artifact_repo import _retry_with_new_creds
+from mlflow.store.artifact.cloud_artifact_repo import (
+    CloudArtifactRepository,
+    _complete_futures,
+    _compute_num_chunks,
+    _validate_chunk_size_aws,
+)
+from mlflow.store.artifact.s3_artifact_repo import _get_s3_client
+from mlflow.utils.file_utils import read_chunk
+from mlflow.utils.request_utils import cloud_storage_http_request
+from mlflow.utils.rest_utils import augmented_raise_for_status
+
+_logger = logging.getLogger(__name__)
+_BUCKET_REGION = "BucketRegion"
+_RESPONSE_METADATA = "ResponseMetadata"
+_HTTP_HEADERS = "HTTPHeaders"
+_HTTP_HEADER_BUCKET_REGION = "x-amz-bucket-region"
+_BUCKET_LOCATION_NAME = "BucketLocationName"
+
+
+class OptimizedS3ArtifactRepository(CloudArtifactRepository):
+    """
+    An optimized version of the S3 Artifact Repository.
+
+    This class is used for uploading and downloading S3 artifacts for UC models. While it largely
+    copies the behavior of the S3ArtifactRepository, the `log_artifact`, `log_artifacts`, and
+    `_download_file` methods are optimized by replacing boto3 client operations with the use of
+    presigned URLs for both uploads and downloads.
+    """
+
+    def __init__(
+        self,
+        artifact_uri,
+        access_key_id=None,
+        secret_access_key=None,
+        session_token=None,
+        credential_refresh_def=None,
+        addressing_style=None,
+        s3_endpoint_url=None,
+        s3_upload_extra_args=None,
+    ):
+        super().__init__(artifact_uri)
+        self._access_key_id = access_key_id
+        self._secret_access_key = secret_access_key
+        self._session_token = session_token
+        self._credential_refresh_def = credential_refresh_def
+        self._addressing_style = addressing_style
+        self._s3_endpoint_url = s3_endpoint_url
+        self.bucket, self.bucket_path = self.parse_s3_compliant_uri(self.artifact_uri)
+        self._region_name = self._get_region_name()
+        self._s3_upload_extra_args = s3_upload_extra_args if s3_upload_extra_args else {}
+
+    def _refresh_credentials(self):
+        if not self._credential_refresh_def:
+            return self._get_s3_client()
+        new_creds = self._credential_refresh_def()
+        self._access_key_id = new_creds["access_key_id"]
+        self._secret_access_key = new_creds["secret_access_key"]
+        self._session_token = new_creds["session_token"]
+        self._s3_upload_extra_args = new_creds["s3_upload_extra_args"]
+        return self._get_s3_client()
+
+    def _get_region_name(self):
+        from botocore.exceptions import ClientError
+
+        temp_client = _get_s3_client(
+            addressing_style=self._addressing_style,
+            access_key_id=self._access_key_id,
+            secret_access_key=self._secret_access_key,
+            session_token=self._session_token,
+            s3_endpoint_url=self._s3_endpoint_url,
+        )
+        try:
+            head_bucket_resp = temp_client.head_bucket(Bucket=self.bucket)
+            # A normal response will have the region in the Bucket_Region field of the response
+            if _BUCKET_REGION in head_bucket_resp:
+                return head_bucket_resp[_BUCKET_REGION]
+            # If the bucket exists but the caller does not have permissions, the http headers
+            # are passed back as part of the metadata of a normal, non-throwing response.  In
+            # this case we use the x-amz-bucket-region field of the HTTP headers which should
+            # always be populated with the region.
+            if (
+                _RESPONSE_METADATA in head_bucket_resp
+                and _HTTP_HEADERS in head_bucket_resp[_RESPONSE_METADATA]
+                and _HTTP_HEADER_BUCKET_REGION
+                in head_bucket_resp[_RESPONSE_METADATA][_HTTP_HEADERS]
+            ):
+                return head_bucket_resp[_RESPONSE_METADATA][_HTTP_HEADERS][
+                    _HTTP_HEADER_BUCKET_REGION
+                ]
+            # Directory buckets do not have a Bucket_Region and instead have a
+            # Bucket_Location_Name.  This name cannot be used as the region name
+            # however, so we warn that this has happened and allow the exception
+            # at the end to be raised.
+            if _BUCKET_LOCATION_NAME in head_bucket_resp:
+                _logger.warning(
+                    f"Directory bucket {self.bucket} found with BucketLocationName "
+                    f"{head_bucket_resp[_BUCKET_LOCATION_NAME]}."
+                )
+            raise Exception(f"Unable to get the region name for bucket {self.bucket}.")
+        except ClientError as error:
+            # If a client error occurs, we check to see if the x-amz-bucket-region field is set
+            # in the response and return that.  If it is not present, this will raise due to the
+            # key not being present.
+            return error.response[_RESPONSE_METADATA][_HTTP_HEADERS][_HTTP_HEADER_BUCKET_REGION]
+
+    def _get_s3_client(self):
+        return _get_s3_client(
+            addressing_style=self._addressing_style,
+            access_key_id=self._access_key_id,
+            secret_access_key=self._secret_access_key,
+            session_token=self._session_token,
+            region_name=self._region_name,
+            s3_endpoint_url=self._s3_endpoint_url,
+        )
+
+    def parse_s3_compliant_uri(self, uri):
+        """Parse an S3 URI, returning (bucket, path)"""
+        parsed = urllib.parse.urlparse(uri)
+        if parsed.scheme != "s3":
+            raise Exception(f"Not an S3 URI: {uri}")
+        path = parsed.path
+        if path.startswith("/"):
+            path = path[1:]
+        return parsed.netloc, path
+
+    @staticmethod
+    def get_s3_file_upload_extra_args():
+        s3_file_upload_extra_args = MLFLOW_S3_UPLOAD_EXTRA_ARGS.get()
+        if s3_file_upload_extra_args:
+            return json.loads(s3_file_upload_extra_args)
+        else:
+            return None
+
+    def _upload_file(self, s3_client, local_file, bucket, key):
+        extra_args = {}
+        extra_args.update(self._s3_upload_extra_args)
+        guessed_type, guessed_encoding = guess_type(local_file)
+        if guessed_type is not None:
+            extra_args["ContentType"] = guessed_type
+        if guessed_encoding is not None:
+            extra_args["ContentEncoding"] = guessed_encoding
+        environ_extra_args = self.get_s3_file_upload_extra_args()
+        if environ_extra_args is not None:
+            extra_args.update(environ_extra_args)
+
+        def try_func(creds):
+            creds.upload_file(Filename=local_file, Bucket=bucket, Key=key, ExtraArgs=extra_args)
+
+        _retry_with_new_creds(
+            try_func=try_func, creds_func=self._refresh_credentials, orig_creds=s3_client
+        )
+
+    def log_artifact(self, local_file, artifact_path=None):
+        artifact_file_path = os.path.basename(local_file)
+        if artifact_path:
+            artifact_file_path = posixpath.join(artifact_path, artifact_file_path)
+        self._upload_to_cloud(
+            cloud_credential_info=self._get_s3_client(),
+            src_file_path=local_file,
+            artifact_file_path=artifact_file_path,
+        )
+
+    def _get_write_credential_infos(self, remote_file_paths):
+        """
+        Instead of returning ArtifactCredentialInfo objects, we instead return a list of initialized
+        S3 client. We do so because S3 clients cannot be instantiated within each thread.
+        """
+        return [self._get_s3_client() for _ in remote_file_paths]
+
+    def _upload_to_cloud(self, cloud_credential_info, src_file_path, artifact_file_path):
+        dest_path = posixpath.join(self.bucket_path, artifact_file_path)
+        key = posixpath.normpath(dest_path)
+        if (
+            MLFLOW_ENABLE_MULTIPART_UPLOAD.get()
+            and os.path.getsize(src_file_path) > MLFLOW_MULTIPART_UPLOAD_CHUNK_SIZE.get()
+        ):
+            self._multipart_upload(cloud_credential_info, src_file_path, self.bucket, key)
+        else:
+            self._upload_file(cloud_credential_info, src_file_path, self.bucket, key)
+
+    def _multipart_upload(self, cloud_credential_info, local_file, bucket, key):
+        # Create multipart upload
+        s3_client = cloud_credential_info
+        response = s3_client.create_multipart_upload(Bucket=bucket, Key=key)
+        upload_id = response["UploadId"]
+
+        num_parts = _compute_num_chunks(local_file, MLFLOW_MULTIPART_UPLOAD_CHUNK_SIZE.get())
+        _validate_chunk_size_aws(MLFLOW_MULTIPART_UPLOAD_CHUNK_SIZE.get())
+
+        # define helper functions for uploading data
+        def _upload_part(part_number, local_file, start_byte, size):
+            data = read_chunk(local_file, size, start_byte)
+
+            def try_func(creds):
+                # Create presigned URL for each part
+                presigned_url = creds.generate_presigned_url(
+                    "upload_part",
+                    Params={
+                        "Bucket": bucket,
+                        "Key": key,
+                        "UploadId": upload_id,
+                        "PartNumber": part_number,
+                    },
+                )
+                with cloud_storage_http_request("put", presigned_url, data=data) as response:
+                    augmented_raise_for_status(response)
+                    return response.headers["ETag"]
+
+            return _retry_with_new_creds(
+                try_func=try_func, creds_func=self._refresh_credentials, orig_creds=s3_client
+            )
+
+        try:
+            # Upload each part with retries
+            futures = {}
+            for index in range(num_parts):
+                part_number = index + 1
+                start_byte = index * MLFLOW_MULTIPART_UPLOAD_CHUNK_SIZE.get()
+                future = self.chunk_thread_pool.submit(
+                    _upload_part,
+                    part_number=part_number,
+                    local_file=local_file,
+                    start_byte=start_byte,
+                    size=MLFLOW_MULTIPART_UPLOAD_CHUNK_SIZE.get(),
+                )
+                futures[future] = part_number
+
+            results, errors = _complete_futures(futures, local_file)
+            if errors:
+                raise MlflowException(
+                    f"Failed to upload at least one part of {local_file}. Errors: {errors}"
+                )
+            parts = [
+                {"PartNumber": part_number, "ETag": results[part_number]}
+                for part_number in sorted(results)
+            ]
+
+            # Complete multipart upload
+            s3_client.complete_multipart_upload(
+                Bucket=bucket,
+                Key=key,
+                UploadId=upload_id,
+                MultipartUpload={"Parts": parts},
+            )
+        except Exception as e:
+            _logger.warning(
+                "Encountered an unexpected error during multipart upload: %s, aborting", e
+            )
+            s3_client.abort_multipart_upload(
+                Bucket=bucket,
+                Key=key,
+                UploadId=upload_id,
+            )
+            raise e
+
+    def list_artifacts(self, path=None):
+        artifact_path = self.bucket_path
+        dest_path = self.bucket_path
+        if path:
+            dest_path = posixpath.join(dest_path, path)
+        infos = []
+        dest_path = dest_path.rstrip("/") if dest_path else ""
+        prefix = dest_path + "/" if dest_path else ""
+        s3_client = self._get_s3_client()
+        paginator = s3_client.get_paginator("list_objects_v2")
+        results = paginator.paginate(Bucket=self.bucket, Prefix=prefix, Delimiter="/")
+        for result in results:
+            # Subdirectories will be listed as "common prefixes" due to the way we made the request
+            for obj in result.get("CommonPrefixes", []):
+                subdir_path = obj.get("Prefix")
+                self._verify_listed_object_contains_artifact_path_prefix(
+                    listed_object_path=subdir_path, artifact_path=artifact_path
+                )
+                subdir_rel_path = posixpath.relpath(path=subdir_path, start=artifact_path)
+                if subdir_rel_path.endswith("/"):
+                    subdir_rel_path = subdir_rel_path[:-1]
+                infos.append(FileInfo(subdir_rel_path, True, None))
+            # Objects listed directly will be files
+            for obj in result.get("Contents", []):
+                file_path = obj.get("Key")
+                self._verify_listed_object_contains_artifact_path_prefix(
+                    listed_object_path=file_path, artifact_path=artifact_path
+                )
+                file_rel_path = posixpath.relpath(path=file_path, start=artifact_path)
+                file_size = int(obj.get("Size"))
+                infos.append(FileInfo(file_rel_path, False, file_size))
+        return sorted(infos, key=lambda f: f.path)
+
+    @staticmethod
+    def _verify_listed_object_contains_artifact_path_prefix(listed_object_path, artifact_path):
+        if not listed_object_path.startswith(artifact_path):
+            raise MlflowException(
+                "The path of the listed S3 object does not begin with the specified"
+                f" artifact path. Artifact path: {artifact_path}. Object path:"
+                f" {listed_object_path}."
+            )
+
+    def _get_presigned_uri(self, remote_file_path):
+        s3_client = self._get_s3_client()
+        s3_full_path = posixpath.join(self.bucket_path, remote_file_path)
+        return s3_client.generate_presigned_url(
+            "get_object", Params={"Bucket": self.bucket, "Key": s3_full_path}
+        )
+
+    def _get_read_credential_infos(self, remote_file_paths):
+        return [
+            ArtifactCredentialInfo(signed_uri=self._get_presigned_uri(path))
+            for path in remote_file_paths
+        ]
+
+    def _download_from_cloud(self, remote_file_path, local_path):
+        s3_client = self._get_s3_client()
+        s3_full_path = posixpath.join(self.bucket_path, remote_file_path)
+
+        def try_func(creds):
+            creds.download_file(self.bucket, s3_full_path, local_path)
+
+        _retry_with_new_creds(
+            try_func=try_func, creds_func=self._refresh_credentials, orig_creds=s3_client
+        )
+
+    def delete_artifacts(self, artifact_path=None):
+        dest_path = self.bucket_path
+        if artifact_path:
+            dest_path = posixpath.join(dest_path, artifact_path)
+
+        dest_path = dest_path.rstrip("/") if dest_path else ""
+        s3_client = self._get_s3_client()
+        paginator = s3_client.get_paginator("list_objects_v2")
+        results = paginator.paginate(Bucket=self.bucket, Prefix=dest_path)
+        for result in results:
+            keys = []
+            for to_delete_obj in result.get("Contents", []):
+                file_path = to_delete_obj.get("Key")
+                self._verify_listed_object_contains_artifact_path_prefix(
+                    listed_object_path=file_path, artifact_path=dest_path
+                )
+                keys.append({"Key": file_path})
+            if keys:
+                s3_client.delete_objects(Bucket=self.bucket, Delete={"Objects": keys})

venv/lib/python3.9/site-packages/mlflow/store/artifact/presigned_url_artifact_repo.py

@@ -0,0 +1,169 @@
+import json
+import os
+import posixpath
+
+from mlflow.entities import FileInfo
+from mlflow.exceptions import RestException
+from mlflow.protos.databricks_artifacts_pb2 import ArtifactCredentialInfo
+from mlflow.protos.databricks_filesystem_service_pb2 import (
+    CreateDownloadUrlRequest,
+    CreateDownloadUrlResponse,
+    CreateUploadUrlRequest,
+    CreateUploadUrlResponse,
+    FilesystemService,
+    ListDirectoryResponse,
+)
+from mlflow.protos.databricks_pb2 import NOT_FOUND, ErrorCode
+from mlflow.store.artifact.artifact_repo import _retry_with_new_creds
+from mlflow.store.artifact.cloud_artifact_repo import CloudArtifactRepository
+from mlflow.utils.file_utils import download_file_using_http_uri
+from mlflow.utils.proto_json_utils import message_to_json
+from mlflow.utils.request_utils import augmented_raise_for_status, cloud_storage_http_request
+from mlflow.utils.rest_utils import (
+    _REST_API_PATH_PREFIX,
+    call_endpoint,
+    extract_api_info_for_service,
+)
+
+FILESYSTEM_METHOD_TO_INFO = extract_api_info_for_service(FilesystemService, _REST_API_PATH_PREFIX)
+DIRECTORIES_ENDPOINT = "/api/2.0/fs/directories"
+
+
+class PresignedUrlArtifactRepository(CloudArtifactRepository):
+    """
+    Stores and retrieves model artifacts using presigned URLs.
+    """
+
+    def __init__(self, db_creds, model_full_name, model_version):
+        artifact_uri = posixpath.join(
+            "/Models", model_full_name.replace(".", "/"), str(model_version)
+        )
+        super().__init__(artifact_uri)
+        self.db_creds = db_creds
+
+    def log_artifact(self, local_file, artifact_path=None):
+        artifact_file_path = os.path.basename(local_file)
+        if artifact_path:
+            artifact_file_path = posixpath.join(artifact_path, artifact_file_path)
+        cloud_credentials = self._get_write_credential_infos(
+            remote_file_paths=[artifact_file_path]
+        )[0]
+        self._upload_to_cloud(
+            cloud_credential_info=cloud_credentials,
+            src_file_path=local_file,
+            artifact_file_path=artifact_file_path,
+        )
+
+    def _get_write_credential_infos(self, remote_file_paths):
+        endpoint, method = FILESYSTEM_METHOD_TO_INFO[CreateUploadUrlRequest]
+        credential_infos = []
+        for relative_path in remote_file_paths:
+            fs_full_path = posixpath.join(self.artifact_uri, relative_path)
+            req_body = message_to_json(CreateUploadUrlRequest(path=fs_full_path))
+            response_proto = CreateUploadUrlResponse()
+            resp = call_endpoint(
+                host_creds=self.db_creds,
+                endpoint=endpoint,
+                method=method,
+                json_body=req_body,
+                response_proto=response_proto,
+            )
+            headers = [
+                ArtifactCredentialInfo.HttpHeader(name=header.name, value=header.value)
+                for header in resp.headers
+            ]
+            credential_infos.append(ArtifactCredentialInfo(signed_uri=resp.url, headers=headers))
+        return credential_infos
+
+    def _upload_to_cloud(self, cloud_credential_info, src_file_path, artifact_file_path=None):
+        # artifact_file_path is unused in this implementation because the presigned URL
+        # and local file path are sufficient for upload to cloud storage
+        def try_func(creds):
+            presigned_url = creds.signed_uri
+            headers = {header.name: header.value for header in creds.headers}
+            with open(src_file_path, "rb") as source_file:
+                data = source_file.read()
+                with cloud_storage_http_request(
+                    "put", presigned_url, data=data, headers=headers
+                ) as response:
+                    augmented_raise_for_status(response)
+
+        def creds_func():
+            return self._get_write_credential_infos(remote_file_paths=[artifact_file_path])[0]
+
+        _retry_with_new_creds(
+            try_func=try_func, creds_func=creds_func, orig_creds=cloud_credential_info
+        )
+
+    def list_artifacts(self, path=""):
+        infos = []
+        page_token = ""
+        while True:
+            endpoint = posixpath.join(DIRECTORIES_ENDPOINT, self.artifact_uri.lstrip("/"), path)
+            req_body = json.dumps({"page_token": page_token}) if page_token else None
+
+            response_proto = ListDirectoryResponse()
+
+            # If the path specified is not a directory, we return an empty list instead of raising
+            # an exception. This is due to this method being used in artifact_repo._is_directory
+            # to determine when a filepath is a directory.
+            try:
+                resp = call_endpoint(
+                    host_creds=self.db_creds,
+                    endpoint=endpoint,
+                    method="GET",
+                    json_body=req_body,
+                    response_proto=response_proto,
+                )
+            except RestException as e:
+                if e.error_code == ErrorCode.Name(NOT_FOUND):
+                    return []
+                else:
+                    raise e
+            for dir_entry in resp.contents:
+                rel_path = posixpath.relpath(dir_entry.path, self.artifact_uri)
+                if dir_entry.is_directory:
+                    infos.append(FileInfo(rel_path, True, None))
+                else:
+                    infos.append(FileInfo(rel_path, False, dir_entry.file_size))
+            page_token = resp.next_page_token
+            if not page_token:
+                break
+        return sorted(infos, key=lambda f: f.path)
+
+    def _get_read_credential_infos(self, remote_file_paths):
+        credential_infos = []
+        for remote_file_path in remote_file_paths:
+            resp = self._get_download_presigned_url_and_headers(remote_file_path)
+            headers = [
+                ArtifactCredentialInfo.HttpHeader(name=header.name, value=header.value)
+                for header in resp.headers
+            ]
+            credential_infos.append(ArtifactCredentialInfo(signed_uri=resp.url, headers=headers))
+        return credential_infos
+
+    def _download_from_cloud(self, remote_file_path, local_path):
+        def creds_func():
+            return self._get_download_presigned_url_and_headers(remote_file_path)
+
+        def try_func(creds):
+            presigned_url = creds.url
+            headers = {header.name: header.value for header in creds.headers}
+            download_file_using_http_uri(
+                http_uri=presigned_url, download_path=local_path, headers=headers
+            )
+
+        _retry_with_new_creds(try_func=try_func, creds_func=creds_func)
+
+    def _get_download_presigned_url_and_headers(self, remote_file_path):
+        remote_file_full_path = posixpath.join(self.artifact_uri, remote_file_path)
+        endpoint, method = FILESYSTEM_METHOD_TO_INFO[CreateDownloadUrlRequest]
+        req_body = message_to_json(CreateDownloadUrlRequest(path=remote_file_full_path))
+        response_proto = CreateDownloadUrlResponse()
+        return call_endpoint(
+            host_creds=self.db_creds,
+            endpoint=endpoint,
+            method=method,
+            json_body=req_body,
+            response_proto=response_proto,
+        )

venv/lib/python3.9/site-packages/mlflow/store/artifact/r2_artifact_repo.py

@@ -0,0 +1,70 @@
+from urllib.parse import urlparse
+
+from mlflow.store.artifact.optimized_s3_artifact_repo import OptimizedS3ArtifactRepository
+from mlflow.store.artifact.s3_artifact_repo import _get_s3_client
+
+
+class R2ArtifactRepository(OptimizedS3ArtifactRepository):
+    """Stores artifacts on Cloudflare R2."""
+
+    def __init__(
+        self,
+        artifact_uri,
+        access_key_id=None,
+        secret_access_key=None,
+        session_token=None,
+        credential_refresh_def=None,
+        s3_upload_extra_args=None,
+    ):
+        # setup Cloudflare R2 backend to be endpoint_url, otherwise all s3 requests
+        # will go to AWS S3 by default
+        s3_endpoint_url = self.convert_r2_uri_to_s3_endpoint_url(artifact_uri)
+        self._access_key_id = access_key_id
+        self._secret_access_key = secret_access_key
+        self._session_token = session_token
+        self._s3_endpoint_url = s3_endpoint_url
+        self.bucket, self.bucket_path = self.parse_s3_compliant_uri(artifact_uri)
+        super().__init__(
+            artifact_uri,
+            access_key_id=access_key_id,
+            secret_access_key=secret_access_key,
+            session_token=session_token,
+            credential_refresh_def=credential_refresh_def,
+            addressing_style="virtual",
+            s3_endpoint_url=s3_endpoint_url,
+            s3_upload_extra_args=s3_upload_extra_args,
+        )
+
+    # Cloudflare implementation of head_bucket is not the same as AWS's, so we
+    # temporarily use the old method of get_bucket_location until cloudflare
+    # updates their implementation
+    def _get_region_name(self):
+        # note: s3 client enforces path addressing style for get_bucket_location
+        temp_client = _get_s3_client(
+            addressing_style="path",
+            access_key_id=self._access_key_id,
+            secret_access_key=self._secret_access_key,
+            session_token=self._session_token,
+            s3_endpoint_url=self._s3_endpoint_url,
+        )
+        return temp_client.get_bucket_location(Bucket=self.bucket)["LocationConstraint"]
+
+    def parse_s3_compliant_uri(self, uri):
+        # r2 uri format(virtual): r2://<bucket-name>@<account-id>.r2.cloudflarestorage.com/<path>
+        parsed = urlparse(uri)
+        if parsed.scheme != "r2":
+            raise Exception(f"Not an R2 URI: {uri}")
+
+        host = parsed.netloc
+        path = parsed.path
+
+        bucket = host.split("@")[0]
+        if path.startswith("/"):
+            path = path[1:]
+        return bucket, path
+
+    @staticmethod
+    def convert_r2_uri_to_s3_endpoint_url(r2_uri):
+        host = urlparse(r2_uri).netloc
+        host_without_bucket = host.split("@")[-1]
+        return f"https://{host_without_bucket}"

venv/lib/python3.9/site-packages/mlflow/store/artifact/runs_artifact_repo.py

@@ -0,0 +1,147 @@
+import urllib.parse
+
+from mlflow.exceptions import MlflowException
+from mlflow.store.artifact.artifact_repo import ArtifactRepository
+from mlflow.utils.uri import (
+    add_databricks_profile_info_to_artifact_uri,
+    get_databricks_profile_uri_from_artifact_uri,
+)
+
+
+class RunsArtifactRepository(ArtifactRepository):
+    """
+    Handles artifacts associated with a Run via URIs of the form
+      `runs:/<run_id>/run-relative/path/to/artifact`.
+    It is a light wrapper that resolves the artifact path to an absolute URI then instantiates
+    and uses the artifact repository for that URI.
+
+    The relative path part of ``artifact_uri`` is expected to be in posixpath format, so Windows
+    users should take special care when constructing the URI.
+    """
+
+    def __init__(self, artifact_uri):
+        from mlflow.store.artifact.artifact_repository_registry import get_artifact_repository
+
+        super().__init__(artifact_uri)
+        uri = RunsArtifactRepository.get_underlying_uri(artifact_uri)
+        self.repo = get_artifact_repository(uri)
+
+    @staticmethod
+    def is_runs_uri(uri):
+        return urllib.parse.urlparse(uri).scheme == "runs"
+
+    @staticmethod
+    def get_underlying_uri(runs_uri):
+        from mlflow.tracking.artifact_utils import get_artifact_uri
+
+        (run_id, artifact_path) = RunsArtifactRepository.parse_runs_uri(runs_uri)
+        tracking_uri = get_databricks_profile_uri_from_artifact_uri(runs_uri)
+        uri = get_artifact_uri(run_id, artifact_path, tracking_uri)
+        assert not RunsArtifactRepository.is_runs_uri(uri)  # avoid an infinite loop
+        return add_databricks_profile_info_to_artifact_uri(uri, tracking_uri)
+
+    @staticmethod
+    def parse_runs_uri(run_uri):
+        parsed = urllib.parse.urlparse(run_uri)
+        if parsed.scheme != "runs":
+            raise MlflowException(
+                f"Not a proper runs:/ URI: {run_uri}. "
+                + "Runs URIs must be of the form 'runs:/<run_id>/run-relative/path/to/artifact'"
+            )
+
+        path = parsed.path
+        if not path.startswith("/") or len(path) <= 1:
+            raise MlflowException(
+                f"Not a proper runs:/ URI: {run_uri}. "
+                + "Runs URIs must be of the form 'runs:/<run_id>/run-relative/path/to/artifact'"
+            )
+        path = path[1:]
+
+        path_parts = path.split("/")
+        run_id = path_parts[0]
+        if run_id == "":
+            raise MlflowException(
+                f"Not a proper runs:/ URI: {run_uri}. "
+                + "Runs URIs must be of the form 'runs:/<run_id>/run-relative/path/to/artifact'"
+            )
+
+        artifact_path = "/".join(path_parts[1:]) if len(path_parts) > 1 else None
+        artifact_path = artifact_path if artifact_path != "" else None
+
+        return run_id, artifact_path
+
+    def log_artifact(self, local_file, artifact_path=None):
+        """
+        Log a local file as an artifact, optionally taking an ``artifact_path`` to place it in
+        within the run's artifacts. Run artifacts can be organized into directories, so you can
+        place the artifact in a directory this way.
+
+        Args:
+            local_file: Path to artifact to log.
+            artifact_path: Directory within the run's artifact directory in which to log the
+                artifact.
+        """
+        self.repo.log_artifact(local_file, artifact_path)
+
+    def log_artifacts(self, local_dir, artifact_path=None):
+        """
+        Log the files in the specified local directory as artifacts, optionally taking
+        an ``artifact_path`` to place them in within the run's artifacts.
+
+        Args:
+            local_dir: Directory of local artifacts to log.
+            artifact_path: Directory within the run's artifact directory in which to log the
+                artifacts.
+        """
+        self.repo.log_artifacts(local_dir, artifact_path)
+
+    def _is_directory(self, artifact_path):
+        return self.repo._is_directory(artifact_path)
+
+    def list_artifacts(self, path):
+        """
+        Return all the artifacts for this run_id directly under path. If path is a file, returns
+        an empty list. Will error if path is neither a file nor directory.
+
+        Args:
+            path: Relative source path that contain desired artifacts
+
+        Returns:
+            List of artifacts as FileInfo listed directly under path.
+        """
+        return self.repo.list_artifacts(path)
+
+    def download_artifacts(self, artifact_path, dst_path=None):
+        """
+        Download an artifact file or directory to a local directory if applicable, and return a
+        local path for it.
+        The caller is responsible for managing the lifecycle of the downloaded artifacts.
+
+        Args:
+            artifact_path: Relative source path to the desired artifacts.
+            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.
+
+        Returns:
+            Absolute path of the local filesystem location containing the desired artifacts.
+        """
+        return self.repo.download_artifacts(artifact_path, dst_path)
+
+    def _download_file(self, remote_file_path, local_path):
+        """
+        Download the file at the specified relative remote path and saves
+        it at the specified local path.
+
+        Args:
+            remote_file_path: Source path to the remote file, relative to the root
+                directory of the artifact repository.
+            local_path: The path to which to save the downloaded file.
+
+        """
+        self.repo._download_file(remote_file_path, local_path)
+
+    def delete_artifacts(self, artifact_path=None):
+        self.repo.delete_artifacts(artifact_path)

venv/lib/python3.9/site-packages/mlflow/store/artifact/s3_artifact_repo.py

@@ -0,0 +1,324 @@
+import json
+import os
+import posixpath
+import urllib.parse
+from datetime import datetime
+from functools import lru_cache
+from mimetypes import guess_type
+
+from mlflow.entities import FileInfo
+from mlflow.entities.multipart_upload import (
+    CreateMultipartUploadResponse,
+    MultipartUploadCredential,
+)
+from mlflow.environment_variables import (
+    MLFLOW_BOTO_CLIENT_ADDRESSING_STYLE,
+    MLFLOW_S3_ENDPOINT_URL,
+    MLFLOW_S3_IGNORE_TLS,
+    MLFLOW_S3_UPLOAD_EXTRA_ARGS,
+)
+from mlflow.exceptions import MlflowException
+from mlflow.store.artifact.artifact_repo import (
+    ArtifactRepository,
+    MultipartUploadMixin,
+)
+from mlflow.utils.file_utils import relative_path_to_artifact_path
+
+_MAX_CACHE_SECONDS = 300
+
+
+def _get_utcnow_timestamp():
+    return datetime.utcnow().timestamp()
+
+
+@lru_cache(maxsize=64)
+def _cached_get_s3_client(
+    signature_version,
+    addressing_style,
+    s3_endpoint_url,
+    verify,
+    timestamp,
+    access_key_id=None,
+    secret_access_key=None,
+    session_token=None,
+    region_name=None,
+):
+    """Returns a boto3 client, caching to avoid extra boto3 verify calls.
+
+    This method is outside of the S3ArtifactRepository as it is
+    agnostic and could be used by other instances.
+
+    `maxsize` set to avoid excessive memory consumption in the case
+    a user has dynamic endpoints (intentionally or as a bug).
+
+    Some of the boto3 endpoint urls, in very edge cases, might expire
+    after twelve hours as that is the current expiration time. To ensure
+    we throw an error on verification instead of using an expired endpoint
+    we utilise the `timestamp` parameter to invalidate cache.
+    """
+    import boto3
+    from botocore.client import Config
+
+    # Making it possible to access public S3 buckets
+    # Workaround for https://github.com/boto/botocore/issues/2442
+    if signature_version.lower() == "unsigned":
+        from botocore import UNSIGNED
+
+        signature_version = UNSIGNED
+
+    return boto3.client(
+        "s3",
+        config=Config(
+            signature_version=signature_version, s3={"addressing_style": addressing_style}
+        ),
+        endpoint_url=s3_endpoint_url,
+        verify=verify,
+        aws_access_key_id=access_key_id,
+        aws_secret_access_key=secret_access_key,
+        aws_session_token=session_token,
+        region_name=region_name,
+    )
+
+
+def _get_s3_client(
+    addressing_style=None,
+    access_key_id=None,
+    secret_access_key=None,
+    session_token=None,
+    region_name=None,
+    s3_endpoint_url=None,
+):
+    if not s3_endpoint_url:
+        s3_endpoint_url = MLFLOW_S3_ENDPOINT_URL.get()
+    do_verify = not MLFLOW_S3_IGNORE_TLS.get()
+
+    # The valid verify argument value is None/False/path to cert bundle file, See
+    # https://github.com/boto/boto3/blob/73865126cad3938ca80a2f567a1c79cb248169a7/
+    # boto3/session.py#L212
+    verify = None if do_verify else False
+
+    # NOTE: If you need to specify this env variable, please file an issue at
+    # https://github.com/mlflow/mlflow/issues so we know your use-case!
+    signature_version = os.environ.get("MLFLOW_EXPERIMENTAL_S3_SIGNATURE_VERSION", "s3v4")
+
+    # Invalidate cache every `_MAX_CACHE_SECONDS`
+    timestamp = int(_get_utcnow_timestamp() / _MAX_CACHE_SECONDS)
+
+    if not addressing_style:
+        addressing_style = MLFLOW_BOTO_CLIENT_ADDRESSING_STYLE.get()
+
+    return _cached_get_s3_client(
+        signature_version,
+        addressing_style,
+        s3_endpoint_url,
+        verify,
+        timestamp,
+        access_key_id=access_key_id,
+        secret_access_key=secret_access_key,
+        session_token=session_token,
+        region_name=region_name,
+    )
+
+
+class S3ArtifactRepository(ArtifactRepository, MultipartUploadMixin):
+    """Stores artifacts on Amazon S3."""
+
+    def __init__(
+        self, artifact_uri, access_key_id=None, secret_access_key=None, session_token=None
+    ):
+        super().__init__(artifact_uri)
+        self._access_key_id = access_key_id
+        self._secret_access_key = secret_access_key
+        self._session_token = session_token
+
+    def _get_s3_client(self):
+        return _get_s3_client(
+            access_key_id=self._access_key_id,
+            secret_access_key=self._secret_access_key,
+            session_token=self._session_token,
+        )
+
+    def parse_s3_compliant_uri(self, uri):
+        """Parse an S3 URI, returning (bucket, path)"""
+        parsed = urllib.parse.urlparse(uri)
+        if parsed.scheme != "s3":
+            raise Exception(f"Not an S3 URI: {uri}")
+        path = parsed.path
+        if path.startswith("/"):
+            path = path[1:]
+        return parsed.netloc, path
+
+    @staticmethod
+    def get_s3_file_upload_extra_args():
+        s3_file_upload_extra_args = MLFLOW_S3_UPLOAD_EXTRA_ARGS.get()
+        if s3_file_upload_extra_args:
+            return json.loads(s3_file_upload_extra_args)
+        else:
+            return None
+
+    def _upload_file(self, s3_client, local_file, bucket, key):
+        extra_args = {}
+        guessed_type, guessed_encoding = guess_type(local_file)
+        if guessed_type is not None:
+            extra_args["ContentType"] = guessed_type
+        if guessed_encoding is not None:
+            extra_args["ContentEncoding"] = guessed_encoding
+        environ_extra_args = self.get_s3_file_upload_extra_args()
+        if environ_extra_args is not None:
+            extra_args.update(environ_extra_args)
+        s3_client.upload_file(Filename=local_file, Bucket=bucket, Key=key, ExtraArgs=extra_args)
+
+    def log_artifact(self, local_file, artifact_path=None):
+        (bucket, dest_path) = self.parse_s3_compliant_uri(self.artifact_uri)
+        if artifact_path:
+            dest_path = posixpath.join(dest_path, artifact_path)
+        dest_path = posixpath.join(dest_path, os.path.basename(local_file))
+        self._upload_file(
+            s3_client=self._get_s3_client(), local_file=local_file, bucket=bucket, key=dest_path
+        )
+
+    def log_artifacts(self, local_dir, artifact_path=None):
+        (bucket, dest_path) = self.parse_s3_compliant_uri(self.artifact_uri)
+        if artifact_path:
+            dest_path = posixpath.join(dest_path, artifact_path)
+        s3_client = self._get_s3_client()
+        local_dir = os.path.abspath(local_dir)
+        for root, _, filenames in os.walk(local_dir):
+            upload_path = dest_path
+            if root != local_dir:
+                rel_path = os.path.relpath(root, local_dir)
+                rel_path = relative_path_to_artifact_path(rel_path)
+                upload_path = posixpath.join(dest_path, rel_path)
+
+            for f in filenames:
+                self._upload_file(
+                    s3_client=s3_client,
+                    local_file=os.path.join(root, f),
+                    bucket=bucket,
+                    key=posixpath.join(upload_path, f),
+                )
+
+    def list_artifacts(self, path=None):
+        (bucket, artifact_path) = self.parse_s3_compliant_uri(self.artifact_uri)
+        dest_path = artifact_path
+        if path:
+            dest_path = posixpath.join(dest_path, path)
+        dest_path = dest_path.rstrip("/") if dest_path else ""
+        infos = []
+        prefix = dest_path + "/" if dest_path else ""
+        s3_client = self._get_s3_client()
+        paginator = s3_client.get_paginator("list_objects_v2")
+        results = paginator.paginate(Bucket=bucket, Prefix=prefix, Delimiter="/")
+        for result in results:
+            # Subdirectories will be listed as "common prefixes" due to the way we made the request
+            for obj in result.get("CommonPrefixes", []):
+                subdir_path = obj.get("Prefix")
+                self._verify_listed_object_contains_artifact_path_prefix(
+                    listed_object_path=subdir_path, artifact_path=artifact_path
+                )
+                subdir_rel_path = posixpath.relpath(path=subdir_path, start=artifact_path)
+                if subdir_rel_path.endswith("/"):
+                    subdir_rel_path = subdir_rel_path[:-1]
+                infos.append(FileInfo(subdir_rel_path, True, None))
+            # Objects listed directly will be files
+            for obj in result.get("Contents", []):
+                file_path = obj.get("Key")
+                self._verify_listed_object_contains_artifact_path_prefix(
+                    listed_object_path=file_path, artifact_path=artifact_path
+                )
+                file_rel_path = posixpath.relpath(path=file_path, start=artifact_path)
+                file_size = int(obj.get("Size"))
+                infos.append(FileInfo(file_rel_path, False, file_size))
+        return sorted(infos, key=lambda f: f.path)
+
+    @staticmethod
+    def _verify_listed_object_contains_artifact_path_prefix(listed_object_path, artifact_path):
+        if not listed_object_path.startswith(artifact_path):
+            raise MlflowException(
+                "The path of the listed S3 object does not begin with the specified"
+                f" artifact path. Artifact path: {artifact_path}. Object path:"
+                f" {listed_object_path}."
+            )
+
+    def _download_file(self, remote_file_path, local_path):
+        (bucket, s3_root_path) = self.parse_s3_compliant_uri(self.artifact_uri)
+        s3_full_path = posixpath.join(s3_root_path, remote_file_path)
+        s3_client = self._get_s3_client()
+        s3_client.download_file(bucket, s3_full_path, local_path)
+
+    def delete_artifacts(self, artifact_path=None):
+        (bucket, dest_path) = self.parse_s3_compliant_uri(self.artifact_uri)
+        if artifact_path:
+            dest_path = posixpath.join(dest_path, artifact_path)
+
+        dest_path = dest_path.rstrip("/") if dest_path else ""
+        s3_client = self._get_s3_client()
+        paginator = s3_client.get_paginator("list_objects_v2")
+        results = paginator.paginate(Bucket=bucket, Prefix=dest_path)
+        for result in results:
+            keys = []
+            for to_delete_obj in result.get("Contents", []):
+                file_path = to_delete_obj.get("Key")
+                self._verify_listed_object_contains_artifact_path_prefix(
+                    listed_object_path=file_path, artifact_path=dest_path
+                )
+                keys.append({"Key": file_path})
+            if keys:
+                s3_client.delete_objects(Bucket=bucket, Delete={"Objects": keys})
+
+    def create_multipart_upload(self, local_file, num_parts=1, artifact_path=None):
+        (bucket, dest_path) = self.parse_s3_compliant_uri(self.artifact_uri)
+        if artifact_path:
+            dest_path = posixpath.join(dest_path, artifact_path)
+        dest_path = posixpath.join(dest_path, os.path.basename(local_file))
+        s3_client = self._get_s3_client()
+        create_response = s3_client.create_multipart_upload(
+            Bucket=bucket,
+            Key=dest_path,
+        )
+        upload_id = create_response["UploadId"]
+        credentials = []
+        for i in range(1, num_parts + 1):  # part number must be in [1, 10000]
+            url = s3_client.generate_presigned_url(
+                "upload_part",
+                Params={
+                    "Bucket": bucket,
+                    "Key": dest_path,
+                    "PartNumber": i,
+                    "UploadId": upload_id,
+                },
+            )
+            credentials.append(
+                MultipartUploadCredential(
+                    url=url,
+                    part_number=i,
+                    headers={},
+                )
+            )
+        return CreateMultipartUploadResponse(
+            credentials=credentials,
+            upload_id=upload_id,
+        )
+
+    def complete_multipart_upload(self, local_file, upload_id, parts=None, artifact_path=None):
+        (bucket, dest_path) = self.parse_s3_compliant_uri(self.artifact_uri)
+        if artifact_path:
+            dest_path = posixpath.join(dest_path, artifact_path)
+        dest_path = posixpath.join(dest_path, os.path.basename(local_file))
+        parts = [{"PartNumber": part.part_number, "ETag": part.etag} for part in parts]
+        s3_client = self._get_s3_client()
+        s3_client.complete_multipart_upload(
+            Bucket=bucket, Key=dest_path, UploadId=upload_id, MultipartUpload={"Parts": parts}
+        )
+
+    def abort_multipart_upload(self, local_file, upload_id, artifact_path=None):
+        (bucket, dest_path) = self.parse_s3_compliant_uri(self.artifact_uri)
+        if artifact_path:
+            dest_path = posixpath.join(dest_path, artifact_path)
+        dest_path = posixpath.join(dest_path, os.path.basename(local_file))
+        s3_client = self._get_s3_client()
+        s3_client.abort_multipart_upload(
+            Bucket=bucket,
+            Key=dest_path,
+            UploadId=upload_id,
+        )

venv/lib/python3.9/site-packages/mlflow/store/artifact/sftp_artifact_repo.py

@@ -0,0 +1,142 @@
+import os
+import posixpath
+import sys
+import urllib.parse
+from contextlib import contextmanager
+from queue import Queue
+
+from mlflow.entities import FileInfo
+from mlflow.store.artifact.artifact_repo import ArtifactRepository
+
+
+# Based on: https://stackoverflow.com/a/58466685
+def _put_r_for_windows(sftp, local_dir, remote_dir, preserve_mtime=False):
+    for entry in os.listdir(local_dir):
+        local_path = os.path.join(local_dir, entry)
+        remote_path = posixpath.join(remote_dir, entry)
+        if os.path.isdir(local_path):
+            sftp.mkdir(remote_path)
+            _put_r_for_windows(sftp, local_path, remote_path, preserve_mtime)
+        else:
+            sftp.put(local_path, remote_path, preserve_mtime=preserve_mtime)
+
+
+class _SftpPool:
+    def __init__(self, connections):
+        self._idle_connection_queue = Queue()
+        for c in connections:
+            self._idle_connection_queue.put(c)
+
+    @contextmanager
+    def get_sftp_connection(self):
+        c = self._idle_connection_queue.get(block=True)
+        try:
+            yield c
+        finally:
+            self._idle_connection_queue.put(c)
+
+
+class SFTPArtifactRepository(ArtifactRepository):
+    """Stores artifacts as files in a remote directory, via sftp."""
+
+    def __init__(self, artifact_uri):
+        self.uri = artifact_uri
+        parsed = urllib.parse.urlparse(artifact_uri)
+        self.config = {
+            "host": parsed.hostname,
+            "port": parsed.port,
+            "username": parsed.username,
+            "password": parsed.password,
+        }
+        self.path = parsed.path or "/"
+
+        import paramiko
+        import pysftp
+
+        if self.config["host"] is None:
+            self.config["host"] = "localhost"
+
+        ssh_config = paramiko.SSHConfig()
+        user_config_file = os.path.expanduser("~/.ssh/config")
+        if os.path.exists(user_config_file):
+            with open(user_config_file) as f:
+                ssh_config.parse(f)
+
+        user_config = ssh_config.lookup(self.config["host"])
+
+        if "hostname" in user_config:
+            self.config["host"] = user_config["hostname"]
+
+        if self.config.get("username", None) is None and "user" in user_config:
+            self.config["username"] = user_config["user"]
+
+        if self.config.get("port", None) is None:
+            if "port" in user_config:
+                self.config["port"] = int(user_config["port"])
+            else:
+                self.config["port"] = 22
+
+        if "identityfile" in user_config:
+            self.config["private_key"] = user_config["identityfile"][0]
+
+        connections = [pysftp.Connection(**self.config) for _ in range(self.max_workers)]
+        self.pool = _SftpPool(connections)
+
+        super().__init__(artifact_uri)
+
+    def log_artifact(self, local_file, artifact_path=None):
+        artifact_dir = posixpath.join(self.path, artifact_path) if artifact_path else self.path
+        with self.pool.get_sftp_connection() as sftp:
+            sftp.makedirs(artifact_dir)
+            sftp.put(local_file, posixpath.join(artifact_dir, os.path.basename(local_file)))
+
+    def log_artifacts(self, local_dir, artifact_path=None):
+        artifact_dir = posixpath.join(self.path, artifact_path) if artifact_path else self.path
+        with self.pool.get_sftp_connection() as sftp:
+            sftp.makedirs(artifact_dir)
+            if sys.platform == "win32":
+                _put_r_for_windows(sftp, local_dir, artifact_dir)
+            else:
+                sftp.put_r(local_dir, artifact_dir)
+
+    def _is_directory(self, artifact_path):
+        artifact_dir = self.path
+        path = posixpath.join(artifact_dir, artifact_path) if artifact_path else artifact_dir
+        with self.pool.get_sftp_connection() as sftp:
+            return sftp.isdir(path)
+
+    def list_artifacts(self, path=None):
+        artifact_dir = self.path
+        list_dir = posixpath.join(artifact_dir, path) if path else artifact_dir
+        with self.pool.get_sftp_connection() as sftp:
+            if not sftp.isdir(list_dir):
+                return []
+            artifact_files = sftp.listdir(list_dir)
+            infos = []
+            for file_name in artifact_files:
+                file_path = file_name if path is None else posixpath.join(path, file_name)
+                full_file_path = posixpath.join(list_dir, file_name)
+                if sftp.isdir(full_file_path):
+                    infos.append(FileInfo(file_path, True, None))
+                else:
+                    infos.append(FileInfo(file_path, False, sftp.stat(full_file_path).st_size))
+            return infos
+
+    def _download_file(self, remote_file_path, local_path):
+        remote_full_path = posixpath.join(self.path, remote_file_path)
+        with self.pool.get_sftp_connection() as sftp:
+            sftp.get(remote_full_path, local_path)
+
+    def delete_artifacts(self, artifact_path=None):
+        artifact_dir = posixpath.join(self.path, artifact_path) if artifact_path else self.path
+        with self.pool.get_sftp_connection() as sftp:
+            self._delete_inner(artifact_dir, sftp)
+
+    def _delete_inner(self, artifact_path, sftp):
+        if sftp.isdir(artifact_path):
+            with sftp.cd(artifact_path):
+                for element in sftp.listdir():
+                    self._delete_inner(element, sftp)
+            sftp.rmdir(artifact_path)
+        elif sftp.isfile(artifact_path):
+            sftp.remove(artifact_path)

venv/lib/python3.9/site-packages/mlflow/store/artifact/uc_volume_artifact_repo.py

@@ -0,0 +1,236 @@
+import os
+import posixpath
+from pathlib import Path
+from typing import Optional
+
+import mlflow.utils.databricks_utils
+from mlflow.entities import FileInfo
+from mlflow.environment_variables import MLFLOW_ENABLE_UC_VOLUME_FUSE_ARTIFACT_REPO
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
+from mlflow.store.artifact.artifact_repo import ArtifactRepository
+from mlflow.store.artifact.local_artifact_repo import LocalArtifactRepository
+from mlflow.utils.databricks_utils import get_databricks_host_creds
+from mlflow.utils.file_utils import relative_path_to_artifact_path
+from mlflow.utils.request_utils import augmented_raise_for_status
+from mlflow.utils.rest_utils import http_request
+from mlflow.utils.uri import (
+    get_databricks_profile_uri_from_artifact_uri,
+    is_databricks_model_registry_artifacts_uri,
+    is_valid_uc_volumes_uri,
+    remove_databricks_profile_info_from_artifact_uri,
+    strip_scheme,
+)
+
+# https://docs.databricks.com/api/workspace/files
+DIRECTORIES_API_ENDPOINT = "/api/2.0/fs/directories"
+FILES_API_ENDPOINT = "/api/2.0/fs/files"
+DOWNLOAD_CHUNK_SIZE = 1024
+
+
+class UCVolumesArtifactRepository(ArtifactRepository):
+    """
+    Stores artifacts on UC Volumes using the Files REST API.
+    """
+
+    def __init__(self, artifact_uri):
+        if not is_valid_uc_volumes_uri(artifact_uri):
+            raise MlflowException(
+                message=(
+                    f"UC volumes URI must be of the form "
+                    f"dbfs:/Volumes/<catalog>/<schema>/<volume>/<path>: {artifact_uri}"
+                ),
+                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))
+
+        # Absolute path to the root of the volume. For example, '/Volumes/my-volume' for
+        # 'dbfs:/Volumes/my-volume'.
+        self.root_path = "/" + strip_scheme(self.artifact_uri).strip("/")
+        self.databricks_profile_uri = (
+            get_databricks_profile_uri_from_artifact_uri(artifact_uri)
+            or mlflow.tracking.get_tracking_uri()
+        )
+
+    def _relative_to_root(self, path):
+        return posixpath.relpath(path, self.root_path)
+
+    def _api_request(self, endpoint, method, **kwargs):
+        creds = get_databricks_host_creds(self.databricks_profile_uri)
+        return http_request(host_creds=creds, endpoint=endpoint, method=method, **kwargs)
+
+    def _list_directory_contents(self, directory_path: str, page_token: Optional[str] = None):  # noqa: D417
+        """
+        Lists the contents of a directory.
+
+        Args:
+            directory_path: The absolute path of a directory.
+
+        Returns:
+            The response from the API.
+
+        See also:
+            https://docs.databricks.com/api/workspace/files/listdirectorycontents
+        """
+        endpoint = f"{DIRECTORIES_API_ENDPOINT}{directory_path}"
+        return self._api_request(
+            endpoint=endpoint,
+            method="GET",
+            params={"page_token": page_token} if page_token else None,
+        )
+
+    def _paginated_list_directory_contents(
+        self, directory_path: str, page_token: Optional[str] = None
+    ):
+        response = self._list_directory_contents(directory_path, page_token)
+        if response.status_code == 404:
+            return []
+
+        augmented_raise_for_status(response)
+
+        response_json = response.json()
+        contents = response_json.get("contents", [])
+        if next_page_token := response_json.get("next_page_token"):
+            next_contents = self._paginated_list_directory_contents(directory_path, next_page_token)
+            return contents + next_contents
+        return contents
+
+    def _download(self, output_path: str, file_path: str):
+        """
+        Downloads a file.
+
+        Args:
+            output_path: The local path to save the downloaded file.
+            file_path: The absolute path of the file to download.
+
+        Returns:
+            The response from the API.
+
+        See also:
+            https://docs.databricks.com/api/workspace/files/download
+        """
+        endpoint = f"{FILES_API_ENDPOINT}{file_path}"
+        with open(output_path, "wb") as f:
+            with self._api_request(endpoint=endpoint, method="GET", stream=True) as resp:
+                for content in resp.iter_content(chunk_size=DOWNLOAD_CHUNK_SIZE):
+                    f.write(content)
+                return resp
+
+    def _upload(self, local_file, file_path):
+        """
+        Uploads a file.
+
+        Args:
+            local_file: The local path of the file to upload.
+            file_path: The absolute path of the file to upload.
+
+        Returns:
+            The response from the API.
+
+        See also:
+            https://docs.databricks.com/api/workspace/files/upload
+        """
+        endpoint = f"{FILES_API_ENDPOINT}{file_path}"
+        with open(local_file, "rb") as f:
+            return self._api_request(endpoint=endpoint, method="PUT", data=f, allow_redirects=False)
+
+    def _get_path(self, artifact_path=None):
+        return (
+            posixpath.join(self.root_path, artifact_path.strip("/"))
+            if artifact_path
+            else self.root_path
+        )
+
+    def log_artifact(self, local_file, artifact_path=None):
+        basename = os.path.basename(local_file)
+        artifact_path = posixpath.join(artifact_path, basename) if artifact_path else basename
+        resp = self._upload(local_file, self._get_path(artifact_path))
+        augmented_raise_for_status(resp)
+
+    def log_artifacts(self, local_dir, artifact_path=None):
+        local_dir = Path(local_dir).resolve()
+        for local_path in local_dir.rglob("*"):
+            if local_path.is_file():
+                if local_path.parent == local_dir:
+                    artifact_subdir = artifact_path
+                else:
+                    rel_path = local_path.parent.relative_to(local_dir)
+                    posix_rel_path = relative_path_to_artifact_path(str(rel_path))
+                    artifact_subdir = (
+                        posixpath.join(artifact_path, posix_rel_path)
+                        if artifact_path
+                        else posix_rel_path
+                    )
+                self.log_artifact(local_path, artifact_subdir)
+
+    def list_artifacts(self, path=None):
+        # Response sample (https://docs.databricks.com/api/workspace/files/listdirectorycontents):
+        # {
+        #    "contents": [
+        #        {
+        #            "path": "string",
+        #            "is_directory": True,
+        #            "file_size": 0,
+        #            "last_modified": 0,
+        #            "name": "string",
+        #        }
+        #    ],
+        #    "next_page_token": "string",
+        # }
+        infos = []
+        for content in self._paginated_list_directory_contents(self._get_path(path)):
+            rel_path = self._relative_to_root(content["path"])
+            infos.append(FileInfo(rel_path, content["is_directory"], content.get("file_size")))
+        return sorted(infos, key=lambda f: f.path)
+
+    def _download_file(self, remote_file_path, local_path):
+        resp = self._download(output_path=local_path, file_path=self._get_path(remote_file_path))
+        augmented_raise_for_status(resp)
+
+    def delete_artifacts(self, artifact_path=None):
+        raise NotImplementedError("Not implemented yet")
+
+
+def uc_volume_artifact_repo_factory(artifact_uri):
+    """
+    Returns an ArtifactRepository subclass for storing artifacts on Volumes.
+
+    This factory method is used with URIs of the form ``dbfs:/Volumes/<path>``. Volume-backed
+    artifact storage can only be used together with the RestStore.
+
+    Args:
+        artifact_uri: Volume root artifact URI.
+
+    Returns:
+        Subclass of ArtifactRepository capable of storing artifacts on DBFS.
+    """
+    if not is_valid_uc_volumes_uri(artifact_uri):
+        raise MlflowException(
+            message=(
+                f"UC volumes URI must be of the form "
+                f"dbfs:/Volumes/<catalog>/<schema>/<volume>/<path>: {artifact_uri}"
+            ),
+            error_code=INVALID_PARAMETER_VALUE,
+        )
+
+    artifact_uri = artifact_uri.rstrip("/")
+    db_profile_uri = get_databricks_profile_uri_from_artifact_uri(artifact_uri)
+    if (
+        mlflow.utils.databricks_utils.is_uc_volume_fuse_available()
+        and MLFLOW_ENABLE_UC_VOLUME_FUSE_ARTIFACT_REPO.get()
+        and not is_databricks_model_registry_artifacts_uri(artifact_uri)
+        and (db_profile_uri is None or db_profile_uri == "databricks")
+    ):
+        # If the UC Volume FUSE mount is available, write artifacts directly to
+        # /Volumes/... 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 `UCVolumesArtifactRepository` to access
+        # the current workspace's Volumes should still work; it just may be slower.
+        uri_without_profile = remove_databricks_profile_info_from_artifact_uri(artifact_uri)
+        path = strip_scheme(uri_without_profile).lstrip("/")
+        return LocalArtifactRepository(f"file:///{path}")
+    return UCVolumesArtifactRepository(artifact_uri)

venv/lib/python3.9/site-packages/mlflow/store/artifact/unity_catalog_models_artifact_repo.py

@@ -0,0 +1,168 @@
+import base64
+
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
+from mlflow.protos.databricks_uc_registry_messages_pb2 import (
+    MODEL_VERSION_OPERATION_READ,
+    GenerateTemporaryModelVersionCredentialsRequest,
+    GenerateTemporaryModelVersionCredentialsResponse,
+    ModelVersionLineageDirection,
+    StorageMode,
+)
+from mlflow.protos.databricks_uc_registry_service_pb2 import UcModelRegistryService
+from mlflow.store._unity_catalog.lineage.constants import _DATABRICKS_LINEAGE_ID_HEADER
+from mlflow.store.artifact.artifact_repo import ArtifactRepository
+from mlflow.store.artifact.databricks_sdk_models_artifact_repo import (
+    DatabricksSDKModelsArtifactRepository,
+)
+from mlflow.store.artifact.presigned_url_artifact_repo import PresignedUrlArtifactRepository
+from mlflow.store.artifact.utils.models import (
+    get_model_name_and_version,
+)
+from mlflow.utils._spark_utils import _get_active_spark_session
+from mlflow.utils._unity_catalog_utils import (
+    emit_model_version_lineage,
+    get_artifact_repo_from_storage_info,
+    get_full_name_from_sc,
+    is_databricks_sdk_models_artifact_repository_enabled,
+)
+from mlflow.utils.databricks_utils import get_databricks_host_creds
+from mlflow.utils.proto_json_utils import message_to_json
+from mlflow.utils.rest_utils import (
+    _REST_API_PATH_PREFIX,
+    call_endpoint,
+    extract_api_info_for_service,
+)
+from mlflow.utils.uri import (
+    _DATABRICKS_UNITY_CATALOG_SCHEME,
+    get_databricks_profile_uri_from_artifact_uri,
+    get_db_info_from_uri,
+    is_databricks_unity_catalog_uri,
+)
+
+_METHOD_TO_INFO = extract_api_info_for_service(UcModelRegistryService, _REST_API_PATH_PREFIX)
+
+
+class UnityCatalogModelsArtifactRepository(ArtifactRepository):
+    """
+    Performs storage operations on artifacts controlled by a Unity Catalog model registry
+
+    Temporary scoped tokens for the appropriate cloud storage locations are fetched from the
+    remote backend and used to download model artifacts.
+
+    The artifact_uri is expected to be of the form `models:/<model_name>/<model_version>`
+
+    Note : This artifact repository is meant is to be instantiated by the ModelsArtifactRepository
+    when the client is pointing to a Unity Catalog model registry.
+    """
+
+    def __init__(self, artifact_uri, registry_uri):
+        if not is_databricks_unity_catalog_uri(registry_uri):
+            raise MlflowException(
+                message="Attempted to instantiate an artifact repo to access models in the "
+                f"Unity Catalog with non-Unity Catalog registry URI '{registry_uri}'. "
+                f"Please specify a Unity Catalog registry URI of the "
+                f"form '{_DATABRICKS_UNITY_CATALOG_SCHEME}[://profile]', e.g. by calling "
+                f"mlflow.set_registry_uri('{_DATABRICKS_UNITY_CATALOG_SCHEME}') if using the "
+                f"MLflow Python client",
+                error_code=INVALID_PARAMETER_VALUE,
+            )
+        super().__init__(artifact_uri)
+        from mlflow.tracking.client import MlflowClient
+
+        registry_uri_from_artifact_uri = get_databricks_profile_uri_from_artifact_uri(
+            artifact_uri, result_scheme=_DATABRICKS_UNITY_CATALOG_SCHEME
+        )
+        if registry_uri_from_artifact_uri is not None:
+            registry_uri = registry_uri_from_artifact_uri
+        _, key_prefix = get_db_info_from_uri(registry_uri)
+        if key_prefix is not None:
+            raise MlflowException(
+                "Remote model registry access via model URIs of the form "
+                "'models://<scope>@<prefix>/<model_name>/<version_or_stage>' is unsupported for "
+                "models in the Unity Catalog. We recommend that you access the Unity Catalog "
+                "from the current Databricks workspace instead."
+            )
+        self.registry_uri = registry_uri
+        self.client = MlflowClient(registry_uri=self.registry_uri)
+        try:
+            spark = _get_active_spark_session()
+        except Exception:
+            pass
+        model_name, self.model_version = get_model_name_and_version(self.client, artifact_uri)
+        self.model_name = get_full_name_from_sc(model_name, spark)
+
+    def _get_blob_storage_path(self):
+        return self.client.get_model_version_download_uri(self.model_name, self.model_version)
+
+    def _get_scoped_token(self, lineage_header_info=None):
+        extra_headers = {}
+        if lineage_header_info:
+            header_json = message_to_json(lineage_header_info)
+            header_base64 = base64.b64encode(header_json.encode())
+            extra_headers[_DATABRICKS_LINEAGE_ID_HEADER] = header_base64
+
+        db_creds = get_databricks_host_creds(self.registry_uri)
+        endpoint, method = _METHOD_TO_INFO[GenerateTemporaryModelVersionCredentialsRequest]
+        req_body = message_to_json(
+            GenerateTemporaryModelVersionCredentialsRequest(
+                name=self.model_name,
+                version=self.model_version,
+                operation=MODEL_VERSION_OPERATION_READ,
+            )
+        )
+        response_proto = GenerateTemporaryModelVersionCredentialsResponse()
+        return call_endpoint(
+            host_creds=db_creds,
+            endpoint=endpoint,
+            method=method,
+            json_body=req_body,
+            response_proto=response_proto,
+            extra_headers=extra_headers,
+        ).credentials
+
+    def _get_artifact_repo(self, lineage_header_info=None):
+        """
+        Get underlying ArtifactRepository instance for model version blob
+        storage
+        """
+        host_creds = get_databricks_host_creds(self.registry_uri)
+        if is_databricks_sdk_models_artifact_repository_enabled(host_creds):
+            entities = lineage_header_info.entities if lineage_header_info else []
+            emit_model_version_lineage(
+                host_creds,
+                self.model_name,
+                self.model_version,
+                entities,
+                ModelVersionLineageDirection.DOWNSTREAM,
+            )
+            return DatabricksSDKModelsArtifactRepository(self.model_name, self.model_version)
+        scoped_token = self._get_scoped_token(lineage_header_info=lineage_header_info)
+        if scoped_token.storage_mode == StorageMode.DEFAULT_STORAGE:
+            return PresignedUrlArtifactRepository(
+                get_databricks_host_creds(self.registry_uri), self.model_name, self.model_version
+            )
+
+        blob_storage_path = self._get_blob_storage_path()
+        return get_artifact_repo_from_storage_info(
+            storage_location=blob_storage_path,
+            scoped_token=scoped_token,
+            base_credential_refresh_def=self._get_scoped_token,
+        )
+
+    def list_artifacts(self, path=None):
+        return self._get_artifact_repo().list_artifacts(path=path)
+
+    def download_artifacts(self, artifact_path, dst_path=None, lineage_header_info=None):
+        return self._get_artifact_repo(lineage_header_info=lineage_header_info).download_artifacts(
+            artifact_path, dst_path
+        )
+
+    def log_artifact(self, local_file, artifact_path=None):
+        raise MlflowException("This repository does not support logging artifacts.")
+
+    def log_artifacts(self, local_dir, artifact_path=None):
+        raise MlflowException("This repository does not support logging artifacts.")
+
+    def delete_artifacts(self, artifact_path=None):
+        raise NotImplementedError("This artifact repository does not support deleting artifacts")

venv/lib/python3.9/site-packages/mlflow/store/artifact/unity_catalog_oss_models_artifact_repo.py

@@ -0,0 +1,168 @@
+import base64
+
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
+from mlflow.protos.databricks_uc_registry_service_pb2 import UcModelRegistryService
+from mlflow.protos.unity_catalog_oss_messages_pb2 import (
+    READ_MODEL_VERSION as MODEL_VERSION_OPERATION_READ_OSS,
+)
+from mlflow.protos.unity_catalog_oss_messages_pb2 import (
+    GenerateTemporaryModelVersionCredential as GenerateTemporaryModelVersionCredentialsOSS,
+)
+from mlflow.protos.unity_catalog_oss_messages_pb2 import (
+    TemporaryCredentials,
+)
+from mlflow.protos.unity_catalog_oss_service_pb2 import UnityCatalogService
+from mlflow.store._unity_catalog.lineage.constants import _DATABRICKS_LINEAGE_ID_HEADER
+from mlflow.store.artifact.artifact_repo import ArtifactRepository
+from mlflow.store.artifact.utils.models import (
+    get_model_name_and_version,
+)
+from mlflow.utils._spark_utils import _get_active_spark_session
+from mlflow.utils._unity_catalog_utils import (
+    get_artifact_repo_from_storage_info,
+    get_full_name_from_sc,
+)
+from mlflow.utils.oss_registry_utils import get_oss_host_creds
+from mlflow.utils.proto_json_utils import message_to_json
+from mlflow.utils.rest_utils import (
+    _REST_API_PATH_PREFIX,
+    _UC_OSS_REST_API_PATH_PREFIX,
+    call_endpoint,
+    extract_api_info_for_service,
+)
+from mlflow.utils.uri import (
+    _OSS_UNITY_CATALOG_SCHEME,
+    get_databricks_profile_uri_from_artifact_uri,
+    get_db_info_from_uri,
+    is_oss_unity_catalog_uri,
+)
+
+_METHOD_TO_INFO = extract_api_info_for_service(UcModelRegistryService, _REST_API_PATH_PREFIX)
+_METHOD_TO_INFO_OSS = extract_api_info_for_service(
+    UnityCatalogService, _UC_OSS_REST_API_PATH_PREFIX
+)
+
+import urllib.parse
+
+from mlflow.store.artifact.local_artifact_repo import LocalArtifactRepository
+from mlflow.utils.uri import is_file_uri
+
+
+class UnityCatalogOSSModelsArtifactRepository(ArtifactRepository):
+    """
+    Performs storage operations on artifacts controlled by a Unity Catalog model registry
+
+    Temporary scoped tokens for the appropriate cloud storage locations are fetched from the
+    remote backend and used to download model artifacts.
+
+    The artifact_uri is expected to be of the form `models:/<model_name>/<model_version>`
+
+    Note : This artifact repository is meant is to be instantiated by the ModelsArtifactRepository
+    when the client is pointing to a Unity Catalog model registry.
+    """
+
+    def __init__(self, artifact_uri, registry_uri):
+        if not is_oss_unity_catalog_uri(registry_uri):
+            raise MlflowException(
+                message="Attempted to instantiate an artifact repo to access models in the "
+                f"OSS Unity Catalog with non-Unity Catalog registry URI '{registry_uri}'. "
+                f"Please specify a Unity Catalog registry URI of the "
+                f"form '{_OSS_UNITY_CATALOG_SCHEME}[://profile]', e.g. by calling "
+                f"mlflow.set_registry_uri('{_OSS_UNITY_CATALOG_SCHEME}') if using the "
+                f"MLflow Python client",
+                error_code=INVALID_PARAMETER_VALUE,
+            )
+        super().__init__(artifact_uri)
+        from mlflow.tracking.client import MlflowClient
+
+        registry_uri_from_artifact_uri = get_databricks_profile_uri_from_artifact_uri(
+            artifact_uri, result_scheme=_OSS_UNITY_CATALOG_SCHEME
+        )
+        if registry_uri_from_artifact_uri is not None:
+            registry_uri = registry_uri_from_artifact_uri
+
+        _, key_prefix = get_db_info_from_uri(urllib.parse.urlparse(registry_uri).path)
+        if key_prefix is not None:
+            raise MlflowException(
+                "Remote model registry access via model URIs of the form "
+                "'models://<scope>@<prefix>/<model_name>/<version_or_stage>' is unsupported for "
+                "models in the Unity Catalog. We recommend that you access the Unity Catalog "
+                "from the current Databricks workspace instead."
+            )
+        self.registry_uri = registry_uri
+        self.client = MlflowClient(registry_uri=self.registry_uri)
+        try:
+            spark = _get_active_spark_session()
+        except Exception:
+            pass
+        model_name, self.model_version = get_model_name_and_version(self.client, artifact_uri)
+        self.model_name = get_full_name_from_sc(model_name, spark)
+
+    def _get_blob_storage_path(self):
+        return self.client.get_model_version_download_uri(self.model_name, self.model_version)
+
+    def _get_scoped_token(self, lineage_header_info=None):
+        extra_headers = {}
+        if lineage_header_info:
+            header_json = message_to_json(lineage_header_info)
+            header_base64 = base64.b64encode(header_json.encode())
+            extra_headers[_DATABRICKS_LINEAGE_ID_HEADER] = header_base64
+        oss_creds = get_oss_host_creds(
+            self.registry_uri
+        )  # Implement ENV variable the same way the databricks user/token is specified
+        oss_endpoint, oss_method = _METHOD_TO_INFO_OSS[GenerateTemporaryModelVersionCredentialsOSS]
+        [catalog_name, schema_name, model_name] = self.model_name.split(
+            "."
+        )  # self.model_name is actually the full name
+        oss_req_body = message_to_json(
+            GenerateTemporaryModelVersionCredentialsOSS(
+                catalog_name=catalog_name,
+                schema_name=schema_name,
+                model_name=model_name,
+                version=int(self.model_version),
+                operation=MODEL_VERSION_OPERATION_READ_OSS,
+            )
+        )
+        oss_response_proto = TemporaryCredentials()
+        return call_endpoint(
+            host_creds=oss_creds,
+            endpoint=oss_endpoint,
+            method=oss_method,
+            json_body=oss_req_body,
+            response_proto=oss_response_proto,
+            extra_headers=extra_headers,
+        )
+
+    def _get_artifact_repo(self, lineage_header_info=None):
+        """
+        Get underlying ArtifactRepository instance for model version blob
+        storage
+        """
+        blob_storage_path = self._get_blob_storage_path()
+        if is_file_uri(blob_storage_path):
+            return LocalArtifactRepository(artifact_uri=blob_storage_path)
+        scoped_token = self._get_scoped_token(lineage_header_info=lineage_header_info)
+        return get_artifact_repo_from_storage_info(
+            storage_location=blob_storage_path,
+            scoped_token=scoped_token,
+            base_credential_refresh_def=self._get_scoped_token,
+            is_oss=True,
+        )
+
+    def list_artifacts(self, path=None):
+        return self._get_artifact_repo().list_artifacts(path=path)
+
+    def download_artifacts(self, artifact_path, dst_path=None, lineage_header_info=None):
+        return self._get_artifact_repo(lineage_header_info=lineage_header_info).download_artifacts(
+            artifact_path, dst_path
+        )
+
+    def log_artifact(self, local_file, artifact_path=None):
+        raise MlflowException("This repository does not support logging artifacts.")
+
+    def log_artifacts(self, local_dir, artifact_path=None):
+        raise MlflowException("This repository does not support logging artifacts.")
+
+    def delete_artifacts(self, artifact_path=None):
+        raise NotImplementedError("This artifact repository does not support deleting artifacts")

venv/lib/python3.9/site-packages/mlflow/store/artifact/utils/models.py

@@ -0,0 +1,99 @@
+import urllib.parse
+from typing import NamedTuple, Optional
+
+import mlflow.tracking
+from mlflow.exceptions import MlflowException
+from mlflow.utils.uri import get_databricks_profile_uri_from_artifact_uri, is_databricks_uri
+
+_MODELS_URI_SUFFIX_LATEST = "latest"
+
+
+def is_using_databricks_registry(uri):
+    profile_uri = get_databricks_profile_uri_from_artifact_uri(uri) or mlflow.get_registry_uri()
+    return is_databricks_uri(profile_uri)
+
+
+def _improper_model_uri_msg(uri):
+    return (
+        f"Not a proper models:/ URI: {uri}. "
+        + "Models URIs must be of the form 'models:/model_name/suffix' "
+        + "or 'models:/model_name@alias' where suffix is a model version, stage, "
+        + f"or the string {_MODELS_URI_SUFFIX_LATEST!r} and where alias is a registered model "
+        + "alias. Only one of suffix or alias can be defined at a time."
+    )
+
+
+def _get_latest_model_version(client, name, stage):
+    """
+    Returns the latest version of the stage if stage is not None. Otherwise return the latest of all
+    versions.
+    """
+    latest = client.get_latest_versions(name, None if stage is None else [stage])
+    if len(latest) == 0:
+        stage_str = "" if stage is None else f" and stage '{stage}'"
+        raise MlflowException(f"No versions of model with name '{name}'{stage_str} found")
+    return max(int(x.version) for x in latest)
+
+
+class ParsedModelUri(NamedTuple):
+    name: str
+    version: Optional[str] = None
+    stage: Optional[str] = None
+    alias: Optional[str] = None
+
+
+def _parse_model_uri(uri):
+    """
+    Returns a ParsedModelUri tuple. Since a models:/ URI can only have one of
+    {version, stage, 'latest', alias}, it will return
+        - (name, version, None, None) to look for a specific version,
+        - (name, None, stage, None) to look for the latest version of a stage,
+        - (name, None, None, None) to look for the latest of all versions.
+        - (name, None, None, alias) to look for a registered model alias.
+    """
+    parsed = urllib.parse.urlparse(uri, allow_fragments=False)
+    if parsed.scheme != "models":
+        raise MlflowException(_improper_model_uri_msg(uri))
+    path = parsed.path
+    if not path.startswith("/") or len(path) <= 1:
+        raise MlflowException(_improper_model_uri_msg(uri))
+
+    parts = path.lstrip("/").split("/")
+    if len(parts) > 2 or parts[0].strip() == "":
+        raise MlflowException(_improper_model_uri_msg(uri))
+
+    if len(parts) == 2:
+        name, suffix = parts
+        if suffix.strip() == "":
+            raise MlflowException(_improper_model_uri_msg(uri))
+        # The URI is in the suffix format
+        if suffix.isdigit():
+            # The suffix is a specific version, e.g. "models:/AdsModel1/123"
+            return ParsedModelUri(name, version=suffix)
+        elif suffix.lower() == _MODELS_URI_SUFFIX_LATEST.lower():
+            # The suffix is the 'latest' string (case insensitive), e.g. "models:/AdsModel1/latest"
+            return ParsedModelUri(name)
+        else:
+            # The suffix is a specific stage (case insensitive), e.g. "models:/AdsModel1/Production"
+            return ParsedModelUri(name, stage=suffix)
+    else:
+        # The URI is an alias URI, e.g. "models:/AdsModel1@Champion"
+        alias_parts = parts[0].rsplit("@", 1)
+        if len(alias_parts) != 2 or alias_parts[1].strip() == "":
+            raise MlflowException(_improper_model_uri_msg(uri))
+        return ParsedModelUri(alias_parts[0], alias=alias_parts[1])
+
+
+def get_model_name_and_version(client, models_uri):
+    (model_name, model_version, model_stage, model_alias) = _parse_model_uri(models_uri)
+    if model_version is not None:
+        return model_name, model_version
+
+    # NB: Call get_model_version_by_alias of registry client directly to bypass prompt check
+    if isinstance(client, mlflow.MlflowClient):
+        client = client._get_registry_client()
+
+    if model_alias is not None:
+        mv = client.get_model_version_by_alias(model_name, model_alias)
+        return model_name, mv.version
+    return model_name, str(_get_latest_model_version(client, model_name, model_stage))