star

venv/lib/python3.9/site-packages/mlflow/tracking/__init__.py

@@ -0,0 +1,30 @@
+"""
+The ``mlflow.tracking`` module provides a Python CRUD interface to MLflow experiments
+and runs. This is a lower level API that directly translates to MLflow
+`REST API <../rest-api.html>`_ calls.
+For a higher level API for managing an "active run", use the :py:mod:`mlflow` module.
+"""
+
+from mlflow.tracking._model_registry.utils import (
+    get_registry_uri,
+    set_registry_uri,
+)
+from mlflow.tracking._tracking_service.utils import (
+    _get_artifact_repo,
+    _get_store,
+    get_tracking_uri,
+    is_tracking_uri_set,
+    set_tracking_uri,
+)
+from mlflow.tracking.client import MlflowClient
+
+__all__ = [
+    "MlflowClient",
+    "get_tracking_uri",
+    "set_tracking_uri",
+    "is_tracking_uri_set",
+    "_get_store",
+    "get_registry_uri",
+    "set_registry_uri",
+    "_get_artifact_repo",
+]

venv/lib/python3.9/site-packages/mlflow/tracking/_model_registry/__init__.py

@@ -0,0 +1 @@
+DEFAULT_AWAIT_MAX_SLEEP_SECONDS = 5 * 60

venv/lib/python3.9/site-packages/mlflow/tracking/_model_registry/client.py

@@ -0,0 +1,433 @@
+"""
+Internal package providing a Python CRUD interface to MLflow models and versions.
+This is a lower level API than the :py:mod:`mlflow.tracking.fluent` module, and is
+exposed in the :py:mod:`mlflow.tracking` module.
+"""
+
+import logging
+
+from mlflow.entities.model_registry import ModelVersionTag, RegisteredModelTag
+from mlflow.exceptions import MlflowException
+from mlflow.prompt.registry_utils import (
+    add_prompt_filter_string,
+    is_prompt_supported_registry,
+)
+from mlflow.store.model_registry import (
+    SEARCH_MODEL_VERSION_MAX_RESULTS_DEFAULT,
+    SEARCH_REGISTERED_MODEL_MAX_RESULTS_DEFAULT,
+)
+from mlflow.tracking._model_registry import DEFAULT_AWAIT_MAX_SLEEP_SECONDS, utils
+from mlflow.utils.arguments_utils import _get_arg_names
+
+_logger = logging.getLogger(__name__)
+
+
+class ModelRegistryClient:
+    """
+    Client of an MLflow Model Registry Server that creates and manages registered
+    models and model versions.
+    """
+
+    def __init__(self, registry_uri, tracking_uri):
+        """
+        Args:
+            registry_uri: Address of local or remote model registry server.
+            tracking_uri: Address of local or remote tracking server.
+        """
+        self.registry_uri = registry_uri
+        self.tracking_uri = tracking_uri
+        # NB: Fetch the tracking store (`self.store`) upon client initialization to ensure that
+        # the tracking URI is valid and the store can be properly resolved. We define `store` as a
+        # property method to ensure that the client is serializable, even if the store is not
+        self.store
+
+    @property
+    def store(self):
+        return utils._get_store(self.registry_uri, self.tracking_uri)
+
+    # Registered Model Methods
+
+    def create_registered_model(self, name, tags=None, description=None):
+        """Create a new registered model in backend store.
+
+        Args:
+            name: Name of the new model. This is expected to be unique in the backend store.
+            tags: A dictionary of key-value pairs that are converted into
+                :py:class:`mlflow.entities.model_registry.RegisteredModelTag` objects.
+            description: Description of the model.
+
+        Returns:
+            A single object of :py:class:`mlflow.entities.model_registry.RegisteredModel`
+            created by backend.
+
+        """
+        # TODO: Do we want to validate the name is legit here - non-empty without "/" and ":" ?
+        #       Those are constraints applicable to any backend, given the model URI format.
+        tags = tags if tags else {}
+        tags = [RegisteredModelTag(key, str(value)) for key, value in tags.items()]
+        return self.store.create_registered_model(name, tags, description)
+
+    def update_registered_model(self, name, description):
+        """Updates description for RegisteredModel entity.
+
+        Backend raises exception if a registered model with given name does not exist.
+
+        Args:
+            name: Name of the registered model to update.
+            description: New description.
+
+        Returns:
+            A single updated :py:class:`mlflow.entities.model_registry.RegisteredModel` object.
+
+        """
+        return self.store.update_registered_model(name=name, description=description)
+
+    def rename_registered_model(self, name, new_name):
+        """Update registered model name.
+
+        Args:
+            name: Name of the registered model to update.
+            new_name: New proposed name for the registered model.
+
+        Returns:
+            A single updated :py:class:`mlflow.entities.model_registry.RegisteredModel` object.
+
+        """
+        if new_name.strip() == "":
+            raise MlflowException("The name must not be an empty string.")
+        return self.store.rename_registered_model(name=name, new_name=new_name)
+
+    def delete_registered_model(self, name):
+        """Delete registered model.
+        Backend raises exception if a registered model with given name does not exist.
+
+        Args:
+            name: Name of the registered model to delete.
+        """
+        self.store.delete_registered_model(name)
+
+    def search_registered_models(
+        self,
+        filter_string=None,
+        max_results=SEARCH_REGISTERED_MODEL_MAX_RESULTS_DEFAULT,
+        order_by=None,
+        page_token=None,
+    ):
+        """Search for registered models in backend that satisfy the filter criteria.
+
+        Args:
+            filter_string: Filter query string, defaults to searching all registered models.
+            max_results: Maximum number of registered models desired.
+            order_by: List of column names with ASC|DESC annotation, to be used for ordering
+                matching search results.
+            page_token: Token specifying the next page of results. It should be obtained from
+                a ``search_registered_models`` call.
+
+        Returns:
+            A PagedList of :py:class:`mlflow.entities.model_registry.RegisteredModel` objects
+            that satisfy the search expressions. The pagination token for the next page can be
+            obtained via the ``token`` attribute of the object.
+
+        """
+        if is_prompt_supported_registry(self.registry_uri):
+            # Adjust filter string to include or exclude prompts
+            filter_string = add_prompt_filter_string(filter_string, False)
+
+        return self.store.search_registered_models(filter_string, max_results, order_by, page_token)
+
+    def get_registered_model(self, name):
+        """
+        Args:
+            name: Name of the registered model to get.
+
+        Returns:
+            A single :py:class:`mlflow.entities.model_registry.RegisteredModel` object.
+        """
+        return self.store.get_registered_model(name)
+
+    def get_latest_versions(self, name, stages=None):
+        """Latest version models for each requests stage. If no ``stages`` provided, returns the
+        latest version for each stage.
+
+        Args:
+            name: Name of the registered model from which to get the latest versions.
+            stages: List of desired stages. If input list is None, return latest versions for
+                'Staging' and 'Production' stages.
+
+        Returns:
+            List of :py:class:`mlflow.entities.model_registry.ModelVersion` objects.
+
+        """
+        return self.store.get_latest_versions(name, stages)
+
+    def set_registered_model_tag(self, name, key, value):
+        """Set a tag for the registered model.
+
+        Args:
+            name: Registered model name.
+            key: Tag key to log.
+            value: Tag value log.
+
+        Returns:
+            None
+        """
+        self.store.set_registered_model_tag(name, RegisteredModelTag(key, str(value)))
+
+    def delete_registered_model_tag(self, name, key):
+        """Delete a tag associated with the registered model.
+
+        Args:
+            name: Registered model name.
+            key: Registered model tag key.
+
+        Returns:
+            None
+        """
+        self.store.delete_registered_model_tag(name, key)
+
+    # Model Version Methods
+
+    def create_model_version(
+        self,
+        name,
+        source,
+        run_id=None,
+        tags=None,
+        run_link=None,
+        description=None,
+        await_creation_for=DEFAULT_AWAIT_MAX_SLEEP_SECONDS,
+        local_model_path=None,
+    ):
+        """Create a new model version from given source.
+
+        Args:
+            name: Name of the containing registered model.
+            source: URI indicating the location of the model artifacts.
+            run_id: Run ID from MLflow tracking server that generated the model.
+            tags: A dictionary of key-value pairs that are converted into
+                :py:class:`mlflow.entities.model_registry.ModelVersionTag` objects.
+            run_link: Link to the run from an MLflow tracking server that generated this model.
+            description: Description of the version.
+            await_creation_for: Number of seconds to wait for the model version to finish being
+                created and is in ``READY`` status. By default, the function
+                waits for five minutes. Specify 0 or None to skip waiting.
+            local_model_path: Local path to the MLflow model, if it's already accessible on the
+                local filesystem. Can be used by AbstractStores that upload model version files
+                to the model registry to avoid a redundant download from the source location when
+                logging and registering a model via a single
+                mlflow.<flavor>.log_model(..., registered_model_name) call.
+
+        Returns:
+            Single :py:class:`mlflow.entities.model_registry.ModelVersion` object created by
+            backend.
+
+        """
+        tags = tags if tags else {}
+        tags = [ModelVersionTag(key, str(value)) for key, value in tags.items()]
+        arg_names = _get_arg_names(self.store.create_model_version)
+        if "local_model_path" in arg_names:
+            mv = self.store.create_model_version(
+                name,
+                source,
+                run_id,
+                tags,
+                run_link,
+                description,
+                local_model_path=local_model_path,
+            )
+        else:
+            # Fall back to calling create_model_version without
+            # local_model_path since old model registry store implementations may not
+            # support the local_model_path argument.
+            mv = self.store.create_model_version(name, source, run_id, tags, run_link, description)
+        if await_creation_for and await_creation_for > 0:
+            self.store._await_model_version_creation(mv, await_creation_for)
+        return mv
+
+    def copy_model_version(self, src_mv, dst_name):
+        """Copy a model version from one registered model to another as a new model version.
+
+        Args:
+            src_mv: A :py:class:`mlflow.entities.model_registry.ModelVersion` object representing
+                the source model version.
+            dst_name: The name of the registered model to copy the model version to. If a
+                registered model with this name does not exist, it will be created.
+
+        Returns:
+            Single :py:class:`mlflow.entities.model_registry.ModelVersion` object representing
+            the cloned model version.
+
+        """
+        return self.store.copy_model_version(src_mv=src_mv, dst_name=dst_name)
+
+    def update_model_version(self, name, version, description):
+        """Update metadata associated with a model version in backend.
+
+        Args:
+            name: Name of the containing registered model.
+            version: Version number of the model version.
+            description: New description.
+        """
+        return self.store.update_model_version(name=name, version=version, description=description)
+
+    def transition_model_version_stage(self, name, version, stage, archive_existing_versions=False):
+        """Update model version stage.
+
+        Args:
+            name: Registered model name.
+            version: Registered model version.
+            stage: New desired stage for this model version.
+            archive_existing_versions: If this flag is set to ``True``, all existing model
+                versions in the stage will be automatically moved to the "archived" stage. Only
+                valid when ``stage`` is ``"staging"`` or ``"production"`` otherwise an error will be
+                raised.
+
+        Returns:
+            A single :py:class:`mlflow.entities.model_registry.ModelVersion` object.
+
+        """
+        if stage.strip() == "":
+            raise MlflowException("The stage must not be an empty string.")
+        return self.store.transition_model_version_stage(
+            name=name,
+            version=version,
+            stage=stage,
+            archive_existing_versions=archive_existing_versions,
+        )
+
+    def get_model_version(self, name, version):
+        """
+        Args:
+            name: Name of the containing registered model.
+            version: Version number of the model version.
+
+        Returns:
+            A single :py:class:`mlflow.entities.model_registry.ModelVersion` object.
+        """
+        return self.store.get_model_version(name, version)
+
+    def delete_model_version(self, name, version):
+        """Delete model version in backend.
+
+        Args:
+            name: Name of the containing registered model.
+            version: Version number of the model version.
+
+        """
+        self.store.delete_model_version(name, version)
+
+    def get_model_version_download_uri(self, name, version):
+        """Get the download location in Model Registry for this model version.
+
+        Args:
+            name: Name of the containing registered model.
+            version: Version number of the model version.
+
+        Returns:
+            A single URI location that allows reads for downloading.
+
+        """
+        return self.store.get_model_version_download_uri(name, version)
+
+    def search_model_versions(
+        self,
+        filter_string=None,
+        max_results=SEARCH_MODEL_VERSION_MAX_RESULTS_DEFAULT,
+        order_by=None,
+        page_token=None,
+    ):
+        """Search for model versions in backend that satisfy the filter criteria.
+
+        .. warning:
+
+            The model version search results may not have aliases populated for performance reasons.
+
+        Args:
+            filter_string: A filter string expression. Currently supports a single filter
+                condition either name of model like ``name = 'model_name'`` or
+                ``run_id = '...'``.
+            max_results: Maximum number of model versions desired.
+            order_by: List of column names with ASC|DESC annotation, to be used for ordering
+                matching search results.
+            page_token: Token specifying the next page of results. It should be obtained from
+                a ``search_model_versions`` call.
+
+        Returns:
+            A PagedList of :py:class:`mlflow.entities.model_registry.ModelVersion`
+            objects that satisfy the search expressions. The pagination token for the next
+            page can be obtained via the ``token`` attribute of the object.
+
+        """
+        return self.store.search_model_versions(filter_string, max_results, order_by, page_token)
+
+    def get_model_version_stages(self, name, version):
+        """
+        Returns:
+            A list of valid stages.
+        """
+        return self.store.get_model_version_stages(name, version)
+
+    def set_model_version_tag(self, name, version, key, value):
+        """Set a tag for the model version.
+
+        Args:
+            name: Registered model name.
+            version: Registered model version.
+            key: Tag key to log.
+            value: Tag value to log.
+
+        Returns:
+            None
+        """
+        self.store.set_model_version_tag(name, version, ModelVersionTag(key, str(value)))
+
+    def delete_model_version_tag(self, name, version, key):
+        """Delete a tag associated with the model version.
+
+        Args:
+            name: Registered model name.
+            version: Registered model version.
+            key: Tag key.
+
+        Returns:
+            None
+        """
+        self.store.delete_model_version_tag(name, version, key)
+
+    def set_registered_model_alias(self, name, alias, version):
+        """Set a registered model alias pointing to a model version.
+
+        Args:
+            name: Registered model name.
+            alias: Name of the alias.
+            version: Registered model version number.
+
+        Returns:
+            None
+        """
+        self.store.set_registered_model_alias(name, alias, version)
+
+    def delete_registered_model_alias(self, name, alias):
+        """Delete an alias associated with a registered model.
+
+        Args:
+            name: Registered model name.
+            alias: Name of the alias.
+
+        Returns:
+            None
+        """
+        self.store.delete_registered_model_alias(name, alias)
+
+    def get_model_version_by_alias(self, name, alias):
+        """Get the model version instance by name and alias.
+
+        Args:
+            name: Registered model name.
+            alias: Name of the alias.
+
+        Returns:
+            A single :py:class:`mlflow.entities.model_registry.ModelVersion` object.
+
+        """
+        return self.store.get_model_version_by_alias(name, alias)

venv/lib/python3.9/site-packages/mlflow/tracking/_model_registry/fluent.py

@@ -0,0 +1,525 @@
+from typing import Any, Optional
+
+from mlflow.entities.model_registry import ModelVersion, Prompt, RegisteredModel
+from mlflow.exceptions import MlflowException
+from mlflow.prompt.registry_utils import require_prompt_registry
+from mlflow.protos.databricks_pb2 import ALREADY_EXISTS, RESOURCE_ALREADY_EXISTS, ErrorCode
+from mlflow.store.artifact.runs_artifact_repo import RunsArtifactRepository
+from mlflow.store.model_registry import (
+    SEARCH_MODEL_VERSION_MAX_RESULTS_DEFAULT,
+    SEARCH_REGISTERED_MODEL_MAX_RESULTS_DEFAULT,
+)
+from mlflow.tracking._model_registry import DEFAULT_AWAIT_MAX_SLEEP_SECONDS
+from mlflow.tracking.client import MlflowClient
+from mlflow.tracking.fluent import active_run
+from mlflow.utils import get_results_from_paginated_fn
+from mlflow.utils.annotations import experimental
+from mlflow.utils.logging_utils import eprint
+
+
+def register_model(
+    model_uri,
+    name,
+    await_registration_for=DEFAULT_AWAIT_MAX_SLEEP_SECONDS,
+    *,
+    tags: Optional[dict[str, Any]] = None,
+) -> ModelVersion:
+    """Create a new model version in model registry for the model files specified by ``model_uri``.
+
+    Note that this method assumes the model registry backend URI is the same as that of the
+    tracking backend.
+
+    Args:
+        model_uri: URI referring to the MLmodel directory. Use a ``runs:/`` URI if you want to
+            record the run ID with the model in model registry (recommended), or pass the
+            local filesystem path of the model if registering a locally-persisted MLflow
+            model that was previously saved using ``save_model``.
+            ``models:/`` URIs are currently not supported.
+        name: Name of the registered model under which to create a new model version. If a
+            registered model with the given name does not exist, it will be created
+            automatically.
+        await_registration_for: Number of seconds to wait for the model version to finish
+            being created and is in ``READY`` status. By default, the function
+            waits for five minutes. Specify 0 or None to skip waiting.
+        tags: A dictionary of key-value pairs that are converted into
+            :py:class:`mlflow.entities.model_registry.ModelVersionTag` objects.
+
+    Returns:
+        Single :py:class:`mlflow.entities.model_registry.ModelVersion` object created by
+        backend.
+
+    .. code-block:: python
+        :test:
+        :caption: Example
+
+        import mlflow.sklearn
+        from mlflow.models import infer_signature
+        from sklearn.datasets import make_regression
+        from sklearn.ensemble import RandomForestRegressor
+
+        mlflow.set_tracking_uri("sqlite:////tmp/mlruns.db")
+        params = {"n_estimators": 3, "random_state": 42}
+        X, y = make_regression(n_features=4, n_informative=2, random_state=0, shuffle=False)
+        # Log MLflow entities
+        with mlflow.start_run() as run:
+            rfr = RandomForestRegressor(**params).fit(X, y)
+            signature = infer_signature(X, rfr.predict(X))
+            mlflow.log_params(params)
+            mlflow.sklearn.log_model(rfr, artifact_path="sklearn-model", signature=signature)
+        model_uri = f"runs:/{run.info.run_id}/sklearn-model"
+        mv = mlflow.register_model(model_uri, "RandomForestRegressionModel")
+        print(f"Name: {mv.name}")
+        print(f"Version: {mv.version}")
+
+    .. code-block:: text
+        :caption: Output
+
+        Name: RandomForestRegressionModel
+        Version: 1
+    """
+    return _register_model(
+        model_uri=model_uri, name=name, await_registration_for=await_registration_for, tags=tags
+    )
+
+
+def _register_model(
+    model_uri,
+    name,
+    await_registration_for=DEFAULT_AWAIT_MAX_SLEEP_SECONDS,
+    *,
+    tags: Optional[dict[str, Any]] = None,
+    local_model_path=None,
+) -> ModelVersion:
+    client = MlflowClient()
+    try:
+        create_model_response = client.create_registered_model(name)
+        eprint(f"Successfully registered model '{create_model_response.name}'.")
+    except MlflowException as e:
+        if e.error_code in (
+            ErrorCode.Name(RESOURCE_ALREADY_EXISTS),
+            ErrorCode.Name(ALREADY_EXISTS),
+        ):
+            eprint(
+                f"Registered model {name!r} already exists. Creating a new version of this model..."
+            )
+        else:
+            raise e
+
+    run_id = None
+    source = model_uri
+    if RunsArtifactRepository.is_runs_uri(model_uri):
+        source = RunsArtifactRepository.get_underlying_uri(model_uri)
+        (run_id, _) = RunsArtifactRepository.parse_runs_uri(model_uri)
+
+    create_version_response = client._create_model_version(
+        name=name,
+        source=source,
+        run_id=run_id,
+        tags=tags,
+        await_creation_for=await_registration_for,
+        local_model_path=local_model_path,
+    )
+    eprint(
+        f"Created version '{create_version_response.version}' of model "
+        f"'{create_version_response.name}'."
+    )
+    return create_version_response
+
+
+def search_registered_models(
+    max_results: Optional[int] = None,
+    filter_string: Optional[str] = None,
+    order_by: Optional[list[str]] = None,
+) -> list[RegisteredModel]:
+    """Search for registered models that satisfy the filter criteria.
+
+    Args:
+        max_results: If passed, specifies the maximum number of models desired. If not
+            passed, all models will be returned.
+        filter_string: Filter query string (e.g., "name = 'a_model_name' and tag.key = 'value1'"),
+            defaults to searching for all registered models. The following identifiers, comparators,
+            and logical operators are supported.
+
+            Identifiers
+              - "name": registered model name.
+              - "tags.<tag_key>": registered model tag. If "tag_key" contains spaces, it must be
+                wrapped with backticks (e.g., "tags.`extra key`").
+
+            Comparators
+              - "=": Equal to.
+              - "!=": Not equal to.
+              - "LIKE": Case-sensitive pattern match.
+              - "ILIKE": Case-insensitive pattern match.
+
+            Logical operators
+              - "AND": Combines two sub-queries and returns True if both of them are True.
+
+        order_by: List of column names with ASC|DESC annotation, to be used for ordering
+            matching search results.
+
+    Returns:
+        A list of :py:class:`mlflow.entities.model_registry.RegisteredModel` objects
+        that satisfy the search expressions.
+
+    .. code-block:: python
+        :test:
+        :caption: Example
+
+        import mlflow
+        from sklearn.linear_model import LogisticRegression
+
+        with mlflow.start_run():
+            mlflow.sklearn.log_model(
+                LogisticRegression(),
+                "Cordoba",
+                registered_model_name="CordobaWeatherForecastModel",
+            )
+            mlflow.sklearn.log_model(
+                LogisticRegression(),
+                "Boston",
+                registered_model_name="BostonWeatherForecastModel",
+            )
+
+        # Get search results filtered by the registered model name
+        filter_string = "name = 'CordobaWeatherForecastModel'"
+        results = mlflow.search_registered_models(filter_string=filter_string)
+        print("-" * 80)
+        for res in results:
+            for mv in res.latest_versions:
+                print(f"name={mv.name}; run_id={mv.run_id}; version={mv.version}")
+
+        # Get search results filtered by the registered model name that matches
+        # prefix pattern
+        filter_string = "name LIKE 'Boston%'"
+        results = mlflow.search_registered_models(filter_string=filter_string)
+        print("-" * 80)
+        for res in results:
+            for mv in res.latest_versions:
+                print(f"name={mv.name}; run_id={mv.run_id}; version={mv.version}")
+
+        # Get all registered models and order them by ascending order of the names
+        results = mlflow.search_registered_models(order_by=["name ASC"])
+        print("-" * 80)
+        for res in results:
+            for mv in res.latest_versions:
+                print(f"name={mv.name}; run_id={mv.run_id}; version={mv.version}")
+
+    .. code-block:: text
+        :caption: Output
+
+        --------------------------------------------------------------------------------
+        name=CordobaWeatherForecastModel; run_id=248c66a666744b4887bdeb2f9cf7f1c6; version=1
+        --------------------------------------------------------------------------------
+        name=BostonWeatherForecastModel; run_id=248c66a666744b4887bdeb2f9cf7f1c6; version=1
+        --------------------------------------------------------------------------------
+        name=BostonWeatherForecastModel; run_id=248c66a666744b4887bdeb2f9cf7f1c6; version=1
+        name=CordobaWeatherForecastModel; run_id=248c66a666744b4887bdeb2f9cf7f1c6; version=1
+    """
+
+    def pagination_wrapper_func(number_to_get, next_page_token):
+        return MlflowClient().search_registered_models(
+            max_results=number_to_get,
+            filter_string=filter_string,
+            order_by=order_by,
+            page_token=next_page_token,
+        )
+
+    return get_results_from_paginated_fn(
+        pagination_wrapper_func,
+        SEARCH_REGISTERED_MODEL_MAX_RESULTS_DEFAULT,
+        max_results,
+    )
+
+
+def search_model_versions(
+    max_results: Optional[int] = None,
+    filter_string: Optional[str] = None,
+    order_by: Optional[list[str]] = None,
+) -> list[ModelVersion]:
+    """Search for model versions that satisfy the filter criteria.
+
+    .. warning:
+
+        The model version search results may not have aliases populated for performance reasons.
+
+    Args:
+        max_results: If passed, specifies the maximum number of models desired. If not
+            passed, all models will be returned.
+        filter_string: Filter query string
+            (e.g., ``"name = 'a_model_name' and tag.key = 'value1'"``),
+            defaults to searching for all model versions. The following identifiers, comparators,
+            and logical operators are supported.
+
+            Identifiers
+              - ``name``: model name.
+              - ``source_path``: model version source path.
+              - ``run_id``: The id of the mlflow run that generates the model version.
+              - ``tags.<tag_key>``: model version tag. If ``tag_key`` contains spaces, it must be
+                wrapped with backticks (e.g., ``"tags.`extra key`"``).
+
+            Comparators
+              - ``=``: Equal to.
+              - ``!=``: Not equal to.
+              - ``LIKE``: Case-sensitive pattern match.
+              - ``ILIKE``: Case-insensitive pattern match.
+              - ``IN``: In a value list. Only ``run_id`` identifier supports ``IN`` comparator.
+
+            Logical operators
+              - ``AND``: Combines two sub-queries and returns True if both of them are True.
+
+        order_by: List of column names with ASC|DESC annotation, to be used for ordering
+            matching search results.
+
+    Returns:
+        A list of :py:class:`mlflow.entities.model_registry.ModelVersion` objects
+            that satisfy the search expressions.
+
+    .. code-block:: python
+        :test:
+        :caption: Example
+
+        import mlflow
+        from sklearn.linear_model import LogisticRegression
+
+        for _ in range(2):
+            with mlflow.start_run():
+                mlflow.sklearn.log_model(
+                    LogisticRegression(),
+                    "Cordoba",
+                    registered_model_name="CordobaWeatherForecastModel",
+                )
+
+        # Get all versions of the model filtered by name
+        filter_string = "name = 'CordobaWeatherForecastModel'"
+        results = mlflow.search_model_versions(filter_string=filter_string)
+        print("-" * 80)
+        for res in results:
+            print(f"name={res.name}; run_id={res.run_id}; version={res.version}")
+
+        # Get the version of the model filtered by run_id
+        filter_string = "run_id = 'ae9a606a12834c04a8ef1006d0cff779'"
+        results = mlflow.search_model_versions(filter_string=filter_string)
+        print("-" * 80)
+        for res in results:
+            print(f"name={res.name}; run_id={res.run_id}; version={res.version}")
+
+    .. code-block:: text
+        :caption: Output
+
+        --------------------------------------------------------------------------------
+        name=CordobaWeatherForecastModel; run_id=ae9a606a12834c04a8ef1006d0cff779; version=2
+        name=CordobaWeatherForecastModel; run_id=d8f028b5fedf4faf8e458f7693dfa7ce; version=1
+        --------------------------------------------------------------------------------
+        name=CordobaWeatherForecastModel; run_id=ae9a606a12834c04a8ef1006d0cff779; version=2
+    """
+
+    def pagination_wrapper_func(number_to_get, next_page_token):
+        return MlflowClient().search_model_versions(
+            max_results=number_to_get,
+            filter_string=filter_string,
+            order_by=order_by,
+            page_token=next_page_token,
+        )
+
+    return get_results_from_paginated_fn(
+        paginated_fn=pagination_wrapper_func,
+        max_results_per_page=SEARCH_MODEL_VERSION_MAX_RESULTS_DEFAULT,
+        max_results=max_results,
+    )
+
+
+@experimental
+@require_prompt_registry
+def register_prompt(
+    name: str,
+    template: str,
+    commit_message: Optional[str] = None,
+    version_metadata: Optional[dict[str, str]] = None,
+    tags: Optional[dict[str, str]] = None,
+) -> Prompt:
+    """
+    Register a new :py:class:`Prompt <mlflow.entities.Prompt>` in the MLflow Prompt Registry.
+
+    A :py:class:`Prompt <mlflow.entities.Prompt>` is a pair of name and
+    template text at minimum. With MLflow Prompt Registry, you can create, manage, and
+    version control prompts with the MLflow's robust model tracking framework.
+
+    If there is no registered prompt with the given name, a new prompt will be created.
+    Otherwise, a new version of the existing prompt will be created.
+
+
+    Args:
+        name: The name of the prompt.
+        template: The template text of the prompt. It can contain variables enclosed in
+            double curly braces, e.g. {variable}, which will be replaced with actual values
+            by the `format` method.
+
+            .. note::
+
+                If you want to use the prompt with a framework that uses single curly braces
+                e.g. LangChain, you can use the `to_single_brace_format` method to convert the
+                loaded prompt to a format that uses single curly braces.
+
+                .. code-block:: python
+
+                    prompt = client.load_prompt("my_prompt")
+                    langchain_format = prompt.to_single_brace_format()
+
+        commit_message: A message describing the changes made to the prompt, similar to a
+            Git commit message. Optional.
+        version_metadata: A dictionary of metadata associated with the **prompt version**.
+            This is useful for storing version-specific information, such as the author of
+            the changes. Optional.
+        tags: A dictionary of tags associated with the entire prompt. This is different from
+            the `version_metadata` as it is not tied to a specific version of the prompt,
+            but to the prompt as a whole. For example, you can use tags to add an application
+            name for which the prompt is created. Since the application uses the prompt in
+            multiple versions, it makes sense to use tags instead of version-specific metadata.
+            Optional.
+
+    Returns:
+        A :py:class:`Prompt <mlflow.entities.Prompt>` object that was created.
+
+    Example:
+
+    .. code-block:: python
+
+        import mlflow
+
+        # Register a new prompt
+        mlflow.register_prompt(
+            name="my_prompt",
+            template="Respond to the user's message as a {{style}} AI.",
+            version_metadata={"author": "Alice"},
+        )
+
+        # Load the prompt from the registry
+        prompt = mlflow.load_prompt("my_prompt")
+
+        # Use the prompt in your application
+        import openai
+
+        openai_client = openai.OpenAI()
+        openai_client.chat.completion.create(
+            model="gpt-4o-mini",
+            messages=[
+                {"role": "system", "content": prompt.format(style="friendly")},
+                {"role": "user", "content": "Hello, how are you?"},
+            ],
+        )
+
+        # Update the prompt with a new version
+        prompt = mlflow.register_prompt(
+            name="my_prompt",
+            template="Respond to the user's message as a {{style}} AI. {{greeting}}",
+            commit_message="Add a greeting to the prompt.",
+            version_metadata={"author": "Bob"},
+        )
+    """
+    return MlflowClient().register_prompt(
+        name=name,
+        template=template,
+        commit_message=commit_message,
+        tags=tags,
+        version_metadata=version_metadata,
+    )
+
+
+@experimental
+@require_prompt_registry
+def load_prompt(name_or_uri: str, version: Optional[int] = None) -> Prompt:
+    """
+    Load a :py:class:`Prompt <mlflow.entities.Prompt>` from the MLflow Prompt Registry.
+
+    The prompt can be specified by name and version, or by URI.
+
+    Args:
+        name_or_uri: The name of the prompt, or the URI in the format "prompts:/name/version".
+        version: The version of the prompt. If not specified, the latest version will be loaded.
+
+    Example:
+
+    .. code-block:: python
+
+        import mlflow
+
+        # Load the latest version of the prompt
+        prompt = mlflow.load_prompt("my_prompt")
+
+        # Load a specific version of the prompt
+        prompt = mlflow.load_prompt("my_prompt", version=1)
+
+        # Load a specific version of the prompt by URI
+        prompt = mlflow.load_prompt(uri="prompts:/my_prompt/1")
+
+        # Load a prompt version with an alias "production"
+        prompt = mlflow.load_prompt("prompts:/my_prompt@production")
+
+    """
+    client = MlflowClient()
+    prompt = client.load_prompt(name_or_uri=name_or_uri, version=version)
+
+    # If there is an active MLflow run, associate the prompt with the run
+    if run := active_run():
+        client.log_prompt(run.info.run_id, f"prompts:/{prompt.name}/{prompt.version}")
+
+    return prompt
+
+
+@experimental
+@require_prompt_registry
+def delete_prompt(name: str, version: int) -> Prompt:
+    """
+    Delete a :py:class:`Prompt <mlflow.entities.Prompt>` from the MLflow Prompt Registry.
+
+    Args:
+        name: The name of the prompt.
+        version: The version of the prompt to delete.
+    """
+    return MlflowClient().delete_prompt(name=name, version=version)
+
+
+@experimental
+@require_prompt_registry
+def set_prompt_alias(name: str, alias: str, version: int) -> Prompt:
+    """
+    Set an alias for a :py:class:`Prompt <mlflow.entities.Prompt>` in the MLflow Prompt Registry.
+
+    Args:
+        name: The name of the prompt.
+        alias: The alias to set for the prompt.
+        version: The version of the prompt.
+
+    Example:
+
+    .. code-block:: python
+
+        import mlflow
+
+        # Set an alias for the prompt
+        mlflow.set_prompt_alias(name="my_prompt", version=1, alias="production")
+
+        # Load the prompt by alias (use "@" to specify the alias)
+        prompt = mlflow.load_prompt("prompts:/my_prompt@production")
+
+        # Switch the alias to a new version of the prompt
+        mlflow.set_prompt_alias(name="my_prompt", version=2, alias="production")
+
+        # Delete the alias
+        mlflow.delete_prompt_alias(name="my_prompt", alias="production")
+    """
+
+    return MlflowClient().set_prompt_alias(name=name, version=version, alias=alias)
+
+
+@experimental
+@require_prompt_registry
+def delete_prompt_alias(name: str, alias: str) -> Prompt:
+    """
+    Delete an alias for a :py:class:`Prompt <mlflow.entities.Prompt>` in the MLflow Prompt Registry.
+
+    Args:
+        name: The name of the prompt.
+        alias: The alias to delete for the prompt.
+    """
+    return MlflowClient().delete_prompt_alias(name=name, alias=alias)

venv/lib/python3.9/site-packages/mlflow/tracking/_model_registry/registry.py

@@ -0,0 +1,67 @@
+import inspect
+import threading
+from functools import lru_cache
+
+from mlflow.tracking.registry import StoreRegistry
+
+_building_store_lock = threading.Lock()
+
+
+class ModelRegistryStoreRegistry(StoreRegistry):
+    """Scheme-based registry for model registry store implementations
+
+    This class allows the registration of a function or class to provide an
+    implementation for a given scheme of `store_uri` through the `register`
+    methods. Implementations declared though the entrypoints
+    `mlflow.registry_store` group can be automatically registered through the
+    `register_entrypoints` method.
+
+    When instantiating a store through the `get_store` method, the scheme of
+    the store URI provided (or inferred from environment) will be used to
+    select which implementation to instantiate, which will be called with same
+    arguments passed to the `get_store` method.
+    """
+
+    def __init__(self):
+        super().__init__("mlflow.model_registry_store")
+
+    def get_store(self, store_uri=None, tracking_uri=None):
+        """Get a store from the registry based on the scheme of store_uri
+
+        Args:
+            store_uri: The store URI. If None, it will be inferred from the environment. This URI
+                is used to select which tracking store implementation to instantiate and
+                is passed to the constructor of the implementation.
+            tracking_uri: The optional string tracking URI to use for any MLflow tracking-related
+                operations in the registry client, e.g. downloading source run
+                artifacts in order to re-upload them to the model registry location.
+
+        Returns:
+            An instance of `mlflow.store.model_registry.AbstractStore` that fulfills the
+            store URI requirements.
+
+        """
+        from mlflow.tracking._model_registry.utils import _resolve_registry_uri
+        from mlflow.tracking._tracking_service.utils import _resolve_tracking_uri
+
+        resolved_store_uri = _resolve_registry_uri(store_uri)
+        resolved_tracking_uri = _resolve_tracking_uri(tracking_uri)
+        return self._get_store_with_resolved_uri(resolved_store_uri, resolved_tracking_uri)
+
+    @lru_cache(maxsize=100)
+    def _get_store_with_resolved_uri(self, resolved_store_uri, resolved_tracking_uri):
+        """
+        Retrieve the store associated with a resolved (non-None) store URI and an artifact URI.
+        Caching is done on resolved URIs because the meaning of an unresolved (None) URI may change
+        depending on external configuration, such as environment variables
+        """
+        with _building_store_lock:
+            builder = self.get_store_builder(resolved_store_uri)
+            builder_param_names = set(inspect.signature(builder).parameters.keys())
+            if "store_uri" in builder_param_names and "tracking_uri" in builder_param_names:
+                return builder(store_uri=resolved_store_uri, tracking_uri=resolved_tracking_uri)
+            else:
+                # Not all model registry stores accept a tracking_uri parameter
+                # (e.g. old plugins may not recognize it), so we fall back to
+                # passing just the registry URI
+                return builder(store_uri=resolved_store_uri)

venv/lib/python3.9/site-packages/mlflow/tracking/_model_registry/utils.py

@@ -0,0 +1,190 @@
+from functools import partial
+
+from mlflow.environment_variables import MLFLOW_REGISTRY_URI
+from mlflow.store.db.db_types import DATABASE_ENGINES
+from mlflow.store.model_registry.databricks_workspace_model_registry_rest_store import (
+    DatabricksWorkspaceModelRegistryRestStore,
+)
+from mlflow.store.model_registry.file_store import FileStore
+from mlflow.store.model_registry.rest_store import RestStore
+from mlflow.tracking._model_registry.registry import ModelRegistryStoreRegistry
+from mlflow.tracking._tracking_service.utils import (
+    _resolve_tracking_uri,
+    get_tracking_uri,
+)
+from mlflow.utils._spark_utils import _get_active_spark_session
+from mlflow.utils.credentials import get_default_host_creds
+from mlflow.utils.databricks_utils import (
+    get_databricks_host_creds,
+    is_in_databricks_serverless_runtime,
+    warn_on_deprecated_cross_workspace_registry_uri,
+)
+from mlflow.utils.uri import _DATABRICKS_UNITY_CATALOG_SCHEME, _OSS_UNITY_CATALOG_SCHEME
+
+# NOTE: in contrast to tracking, we do not support the following ways to specify
+# the model registry URI:
+#  - via environment variables like MLFLOW_TRACKING_URI, MLFLOW_TRACKING_USERNAME, ...
+# We do support specifying it
+#  - via the ``model_registry_uri`` parameter when creating an ``MlflowClient`` or
+#    ``ModelRegistryClient``.
+#  - via a utility method ``mlflow.set_registry_uri``
+#  - by not specifying anything: in this case we assume the model registry store URI is
+#    the same as the tracking store URI. This means Tracking and Model Registry are
+#    backed by the same backend DB/Rest server. However, note that we access them via
+#    different ``Store`` classes (e.g. ``mlflow.store.tracking.SQLAlchemyStore`` &
+#    ``mlflow.store.model_registry.SQLAlchemyStore``).
+# This means the following combinations are not supported:
+#  - Tracking RestStore & Model Registry RestStore that use different credentials.
+
+_registry_uri = None
+
+
+def set_registry_uri(uri: str) -> None:
+    """Set the registry server URI. This method is especially useful if you have a registry server
+    that's different from the tracking server.
+
+    Args:
+        uri: An empty string, or a local file path, prefixed with ``file:/``. Data is stored
+            locally at the provided file (or ``./mlruns`` if empty). An HTTP URI like
+            ``https://my-tracking-server:5000`` or ``http://my-oss-uc-server:8080``. A Databricks
+            workspace, provided as the string "databricks" or, to use a Databricks CLI
+            `profile <https://github.com/databricks/databricks-cli#installation>`_,
+            "databricks://<profileName>".
+
+    .. code-block:: python
+        :caption: Example
+
+        import mflow
+
+        # Set model registry uri, fetch the set uri, and compare
+        # it with the tracking uri. They should be different
+        mlflow.set_registry_uri("sqlite:////tmp/registry.db")
+        mr_uri = mlflow.get_registry_uri()
+        print(f"Current registry uri: {mr_uri}")
+        tracking_uri = mlflow.get_tracking_uri()
+        print(f"Current tracking uri: {tracking_uri}")
+
+        # They should be different
+        assert tracking_uri != mr_uri
+
+    .. code-block:: text
+        :caption: Output
+
+        Current registry uri: sqlite:////tmp/registry.db
+        Current tracking uri: file:///.../mlruns
+
+    """
+    global _registry_uri
+    _registry_uri = uri
+    if uri:
+        # Set 'MLFLOW_REGISTRY_URI' environment variable
+        # so that subprocess can inherit it.
+        MLFLOW_REGISTRY_URI.set(_registry_uri)
+
+
+def _get_registry_uri_from_spark_session():
+    session = _get_active_spark_session()
+    if session is None:
+        return None
+
+    if is_in_databricks_serverless_runtime():
+        # Connected to Serverless
+        return "databricks-uc"
+
+    return session.conf.get("spark.mlflow.modelRegistryUri", None)
+
+
+def _get_registry_uri_from_context():
+    if _registry_uri is not None:
+        return _registry_uri
+    elif (uri := MLFLOW_REGISTRY_URI.get()) or (uri := _get_registry_uri_from_spark_session()):
+        return uri
+    return _registry_uri
+
+
+def get_registry_uri() -> str:
+    """Get the current registry URI. If none has been specified, defaults to the tracking URI.
+
+    Returns:
+        The registry URI.
+
+    .. code-block:: python
+
+        # Get the current model registry uri
+        mr_uri = mlflow.get_registry_uri()
+        print(f"Current model registry uri: {mr_uri}")
+
+        # Get the current tracking uri
+        tracking_uri = mlflow.get_tracking_uri()
+        print(f"Current tracking uri: {tracking_uri}")
+
+        # They should be the same
+        assert mr_uri == tracking_uri
+
+    .. code-block:: text
+
+        Current model registry uri: file:///.../mlruns
+        Current tracking uri: file:///.../mlruns
+
+    """
+    return _get_registry_uri_from_context() or get_tracking_uri()
+
+
+def _resolve_registry_uri(registry_uri=None, tracking_uri=None):
+    return registry_uri or _get_registry_uri_from_context() or _resolve_tracking_uri(tracking_uri)
+
+
+def _get_sqlalchemy_store(store_uri):
+    from mlflow.store.model_registry.sqlalchemy_store import SqlAlchemyStore
+
+    return SqlAlchemyStore(store_uri)
+
+
+def _get_rest_store(store_uri, **_):
+    return RestStore(partial(get_default_host_creds, store_uri))
+
+
+def _get_databricks_rest_store(store_uri, **_):
+    warn_on_deprecated_cross_workspace_registry_uri(registry_uri=store_uri)
+    return DatabricksWorkspaceModelRegistryRestStore(partial(get_databricks_host_creds, store_uri))
+
+
+# We define the global variable as `None` so that instantiating the store does not lead to circular
+# dependency issues.
+_model_registry_store_registry = None
+
+
+def _get_file_store(store_uri, **_):
+    return FileStore(store_uri)
+
+
+def _get_store_registry():
+    global _model_registry_store_registry
+    from mlflow.store._unity_catalog.registry.rest_store import UcModelRegistryStore
+    from mlflow.store._unity_catalog.registry.uc_oss_rest_store import UnityCatalogOssStore
+
+    if _model_registry_store_registry is not None:
+        return _model_registry_store_registry
+
+    _model_registry_store_registry = ModelRegistryStoreRegistry()
+    _model_registry_store_registry.register("databricks", _get_databricks_rest_store)
+    # Register a placeholder function that raises if users pass a registry URI with scheme
+    # "databricks-uc"
+    _model_registry_store_registry.register(_DATABRICKS_UNITY_CATALOG_SCHEME, UcModelRegistryStore)
+    _model_registry_store_registry.register(_OSS_UNITY_CATALOG_SCHEME, UnityCatalogOssStore)
+
+    for scheme in ["http", "https"]:
+        _model_registry_store_registry.register(scheme, _get_rest_store)
+
+    for scheme in DATABASE_ENGINES:
+        _model_registry_store_registry.register(scheme, _get_sqlalchemy_store)
+
+    for scheme in ["", "file"]:
+        _model_registry_store_registry.register(scheme, _get_file_store)
+
+    _model_registry_store_registry.register_entrypoints()
+    return _model_registry_store_registry
+
+
+def _get_store(store_uri=None, tracking_uri=None):
+    return _get_store_registry().get_store(store_uri, tracking_uri)

venv/lib/python3.9/site-packages/mlflow/tracking/_tracking_service/registry.py

@@ -0,0 +1,56 @@
+import threading
+from functools import lru_cache
+
+from mlflow.tracking.registry import StoreRegistry
+
+_building_store_lock = threading.Lock()
+
+
+class TrackingStoreRegistry(StoreRegistry):
+    """Scheme-based registry for tracking store implementations
+
+    This class allows the registration of a function or class to provide an
+    implementation for a given scheme of `store_uri` through the `register`
+    methods. Implementations declared though the entrypoints
+    `mlflow.tracking_store` group can be automatically registered through the
+    `register_entrypoints` method.
+
+    When instantiating a store through the `get_store` method, the scheme of
+    the store URI provided (or inferred from environment) will be used to
+    select which implementation to instantiate, which will be called with same
+    arguments passed to the `get_store` method.
+    """
+
+    def __init__(self):
+        super().__init__("mlflow.tracking_store")
+
+    def get_store(self, store_uri=None, artifact_uri=None):
+        """Get a store from the registry based on the scheme of store_uri
+
+        Args:
+            store_uri: The store URI. If None, it will be inferred from the environment. This URI
+                is used to select which tracking store implementation to instantiate and
+                is passed to the constructor of the implementation.
+            artifact_uri: Artifact repository URI. Passed through to the tracking store
+                implementation.
+
+        Returns:
+            An instance of `mlflow.store.tracking.AbstractStore` that fulfills the store URI
+            requirements.
+
+        """
+        from mlflow.tracking._tracking_service import utils
+
+        resolved_store_uri = utils._resolve_tracking_uri(store_uri)
+        return self._get_store_with_resolved_uri(resolved_store_uri, artifact_uri)
+
+    @lru_cache(maxsize=100)
+    def _get_store_with_resolved_uri(self, resolved_store_uri, artifact_uri=None):
+        """
+        Retrieve the store associated with a resolved (non-None) store URI and an artifact URI.
+        Caching is done on resolved URIs because the meaning of an unresolved (None) URI may change
+        depending on external configuration, such as environment variables
+        """
+        with _building_store_lock:
+            builder = self.get_store_builder(resolved_store_uri)
+            return builder(store_uri=resolved_store_uri, artifact_uri=artifact_uri)

venv/lib/python3.9/site-packages/mlflow/tracking/_tracking_service/utils.py

@@ -0,0 +1,261 @@
+import logging
+import os
+from collections import OrderedDict
+from contextlib import contextmanager
+from functools import partial
+from pathlib import Path
+from typing import Generator, Union
+
+from mlflow.environment_variables import MLFLOW_TRACKING_URI
+from mlflow.store.db.db_types import DATABASE_ENGINES
+from mlflow.store.tracking import DEFAULT_LOCAL_FILE_AND_ARTIFACT_PATH
+from mlflow.store.tracking.file_store import FileStore
+from mlflow.store.tracking.rest_store import RestStore
+from mlflow.tracing.provider import reset
+from mlflow.tracking._tracking_service.registry import TrackingStoreRegistry
+from mlflow.utils.credentials import get_default_host_creds
+from mlflow.utils.databricks_utils import get_databricks_host_creds
+from mlflow.utils.file_utils import path_to_local_file_uri
+from mlflow.utils.uri import _DATABRICKS_UNITY_CATALOG_SCHEME, _OSS_UNITY_CATALOG_SCHEME
+
+_logger = logging.getLogger(__name__)
+_tracking_uri = None
+
+
+def is_tracking_uri_set():
+    """Returns True if the tracking URI has been set, False otherwise."""
+    if _tracking_uri or MLFLOW_TRACKING_URI.get():
+        return True
+    return False
+
+
+def set_tracking_uri(uri: Union[str, Path]) -> None:
+    """
+    Set the tracking server URI. This does not affect the
+    currently active run (if one exists), but takes effect for successive runs.
+
+    Args:
+        uri:
+
+            - An empty string, or a local file path, prefixed with ``file:/``. Data is stored
+              locally at the provided file (or ``./mlruns`` if empty).
+            - An HTTP URI like ``https://my-tracking-server:5000``.
+            - A Databricks workspace, provided as the string "databricks" or, to use a Databricks
+              CLI `profile <https://github.com/databricks/databricks-cli#installation>`_,
+              "databricks://<profileName>".
+            - A :py:class:`pathlib.Path` instance
+
+    .. code-block:: python
+        :test:
+        :caption: Example
+
+        import mlflow
+
+        mlflow.set_tracking_uri("file:///tmp/my_tracking")
+        tracking_uri = mlflow.get_tracking_uri()
+        print(f"Current tracking uri: {tracking_uri}")
+
+    .. code-block:: text
+        :caption: Output
+
+        Current tracking uri: file:///tmp/my_tracking
+    """
+    if isinstance(uri, Path):
+        # On Windows with Python3.8 (https://bugs.python.org/issue38671)
+        # .resolve() doesn't return the absolute path if the directory doesn't exist
+        # so we're calling .absolute() first to get the absolute path on Windows,
+        # then .resolve() to clean the path
+        uri = uri.absolute().resolve().as_uri()
+    global _tracking_uri
+
+    if _tracking_uri != uri:
+        _tracking_uri = uri
+        if _tracking_uri is not None:
+            # Set 'MLFLOW_TRACKING_URI' environment variable
+            # so that subprocess can inherit it.
+            MLFLOW_TRACKING_URI.set(_tracking_uri)
+        else:
+            MLFLOW_TRACKING_URI.unset()
+
+        # Tracer provider uses tracking URI to determine where to export traces.
+        # Tracer provider stores the URI as its state so we need to reset
+        # it explicitly when the global tracking URI changes.
+        reset()
+
+
+@contextmanager
+def _use_tracking_uri(uri: str) -> Generator[None, None, None]:
+    """Temporarily use the specified tracking URI.
+
+    Args:
+        uri: The tracking URI to use.
+
+    """
+    old_tracking_uri = _tracking_uri
+    try:
+        set_tracking_uri(uri)
+        yield
+    finally:
+        set_tracking_uri(old_tracking_uri)
+
+
+def _resolve_tracking_uri(tracking_uri=None):
+    return tracking_uri or get_tracking_uri()
+
+
+def get_tracking_uri() -> str:
+    """Get the current tracking URI. This may not correspond to the tracking URI of
+    the currently active run, since the tracking URI can be updated via ``set_tracking_uri``.
+
+    Returns:
+        The tracking URI.
+
+    .. code-block:: python
+
+        import mlflow
+
+        # Get the current tracking uri
+        tracking_uri = mlflow.get_tracking_uri()
+        print(f"Current tracking uri: {tracking_uri}")
+
+    .. code-block:: text
+
+        Current tracking uri: file:///.../mlruns
+    """
+    if _tracking_uri is not None:
+        return _tracking_uri
+    elif uri := MLFLOW_TRACKING_URI.get():
+        return uri
+    else:
+        return path_to_local_file_uri(os.path.abspath(DEFAULT_LOCAL_FILE_AND_ARTIFACT_PATH))
+
+
+def _get_file_store(store_uri, **_):
+    return FileStore(store_uri, store_uri)
+
+
+def _get_sqlalchemy_store(store_uri, artifact_uri):
+    from mlflow.store.tracking.sqlalchemy_store import SqlAlchemyStore
+
+    if artifact_uri is None:
+        artifact_uri = DEFAULT_LOCAL_FILE_AND_ARTIFACT_PATH
+    return SqlAlchemyStore(store_uri, artifact_uri)
+
+
+def _get_rest_store(store_uri, **_):
+    return RestStore(partial(get_default_host_creds, store_uri))
+
+
+def _get_databricks_rest_store(store_uri, **_):
+    return RestStore(partial(get_databricks_host_creds, store_uri))
+
+
+def _get_databricks_uc_rest_store(store_uri, **_):
+    from mlflow.exceptions import MlflowException
+    from mlflow.version import VERSION
+
+    supported_schemes = [
+        scheme
+        for scheme in _tracking_store_registry._registry
+        if scheme not in {_DATABRICKS_UNITY_CATALOG_SCHEME, _OSS_UNITY_CATALOG_SCHEME}
+    ]
+    raise MlflowException(
+        f"Detected Unity Catalog tracking URI '{store_uri}'. "
+        "Setting the tracking URI to a Unity Catalog backend is not supported in the current "
+        f"version of the MLflow client ({VERSION}). "
+        "Please specify a different tracking URI via mlflow.set_tracking_uri, with "
+        "one of the supported schemes: "
+        f"{supported_schemes}. If you're trying to access models in the Unity "
+        "Catalog, please upgrade to the latest version of the MLflow Python "
+        "client, then specify a Unity Catalog model registry URI via "
+        f"mlflow.set_registry_uri('{_DATABRICKS_UNITY_CATALOG_SCHEME}') or "
+        f"mlflow.set_registry_uri('{_DATABRICKS_UNITY_CATALOG_SCHEME}://profile_name') where "
+        "'profile_name' is the name of the Databricks CLI profile to use for "
+        "authentication. A OSS Unity Catalog model registry URI can also be specified via "
+        f"mlflow.set_registry_uri('{_OSS_UNITY_CATALOG_SCHEME}:http://localhost:8080')."
+        "Be sure to leave the registry URI configured to use one of the supported"
+        "schemes listed above."
+    )
+
+
+_tracking_store_registry = TrackingStoreRegistry()
+
+
+def _register_tracking_stores():
+    _tracking_store_registry.register("", _get_file_store)
+    _tracking_store_registry.register("file", _get_file_store)
+    _tracking_store_registry.register("databricks", _get_databricks_rest_store)
+    _tracking_store_registry.register(
+        _DATABRICKS_UNITY_CATALOG_SCHEME, _get_databricks_uc_rest_store
+    )
+    _tracking_store_registry.register(_OSS_UNITY_CATALOG_SCHEME, _get_databricks_uc_rest_store)
+
+    for scheme in ["http", "https"]:
+        _tracking_store_registry.register(scheme, _get_rest_store)
+
+    for scheme in DATABASE_ENGINES:
+        _tracking_store_registry.register(scheme, _get_sqlalchemy_store)
+
+    _tracking_store_registry.register_entrypoints()
+
+
+def _register(scheme, builder):
+    _tracking_store_registry.register(scheme, builder)
+
+
+_register_tracking_stores()
+
+
+def _get_store(store_uri=None, artifact_uri=None):
+    return _tracking_store_registry.get_store(store_uri, artifact_uri)
+
+
+_artifact_repos_cache = OrderedDict()
+
+
+def _get_artifact_repo(run_id):
+    return _artifact_repos_cache.get(run_id)
+
+
+# TODO(sueann): move to a projects utils module
+def _get_git_url_if_present(uri):
+    """Return the path git_uri#sub_directory if the URI passed is a local path that's part of
+    a Git repo, or returns the original URI otherwise.
+
+    Args:
+        uri: The expanded uri.
+
+    Returns:
+        The git_uri#sub_directory if the uri is part of a Git repo, otherwise return the original
+        uri.
+    """
+    if "#" in uri:
+        # Already a URI in git repo format
+        return uri
+    try:
+        from git import GitCommandNotFound, InvalidGitRepositoryError, NoSuchPathError, Repo
+    except ImportError as e:
+        _logger.warning(
+            "Failed to import Git (the git executable is probably not on your PATH),"
+            " so Git SHA is not available. Error: %s",
+            e,
+        )
+        return uri
+    try:
+        # Check whether this is part of a git repo
+        repo = Repo(uri, search_parent_directories=True)
+
+        # Repo url
+        repo_url = f"file://{repo.working_tree_dir}"
+
+        # Sub directory
+        rlpath = uri.replace(repo.working_tree_dir, "")
+        if rlpath == "":
+            git_path = repo_url
+        elif rlpath[0] == "/":
+            git_path = repo_url + "#" + rlpath[1:]
+        else:
+            git_path = repo_url + "#" + rlpath
+        return git_path
+    except (InvalidGitRepositoryError, GitCommandNotFound, ValueError, NoSuchPathError):
+        return uri

venv/lib/python3.9/site-packages/mlflow/tracking/artifact_utils.py

@@ -0,0 +1,164 @@
+"""
+Utilities for dealing with artifacts in the context of a Run.
+"""
+
+import os
+import pathlib
+import posixpath
+import tempfile
+import urllib.parse
+import uuid
+
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
+from mlflow.store.artifact.artifact_repository_registry import get_artifact_repository
+from mlflow.store.artifact.dbfs_artifact_repo import DbfsRestArtifactRepository
+from mlflow.store.artifact.models_artifact_repo import ModelsArtifactRepository
+from mlflow.tracking._tracking_service.utils import _get_store
+from mlflow.utils.file_utils import path_to_local_file_uri
+from mlflow.utils.os import is_windows
+from mlflow.utils.uri import add_databricks_profile_info_to_artifact_uri, append_to_uri_path
+
+
+def get_artifact_uri(run_id, artifact_path=None, tracking_uri=None):
+    """Get the absolute URI of the specified artifact in the specified run. If `path` is not
+    specified the artifact root URI of the specified run will be returned; calls to ``log_artifact``
+    and ``log_artifacts`` write artifact(s) to subdirectories of the artifact root URI.
+
+    Args:
+        run_id: The ID of the run for which to obtain an absolute artifact URI.
+        artifact_path: The run-relative artifact path. For example,
+            ``path/to/artifact``. If unspecified, the artifact root URI for the
+            specified run will be returned.
+        tracking_uri: The tracking URI from which to get the run and its artifact location. If
+            not given, the current default tracking URI is used.
+
+    Returns:
+        An *absolute* URI referring to the specified artifact or the specified run's artifact
+        root. For example, if an artifact path is provided and the specified run uses an
+        S3-backed  store, this may be a uri of the form
+        ``s3://<bucket_name>/path/to/artifact/root/path/to/artifact``. If an artifact path
+        is not provided and the specified run uses an S3-backed store, this may be a URI of
+        the form ``s3://<bucket_name>/path/to/artifact/root``.
+
+    """
+    if not run_id:
+        raise MlflowException(
+            message="A run_id must be specified in order to obtain an artifact uri!",
+            error_code=INVALID_PARAMETER_VALUE,
+        )
+
+    store = _get_store(tracking_uri)
+    run = store.get_run(run_id)
+    # Maybe move this method to RunsArtifactRepository so the circular dependency is clearer.
+    assert urllib.parse.urlparse(run.info.artifact_uri).scheme != "runs"  # avoid an infinite loop
+    if artifact_path is None:
+        return run.info.artifact_uri
+    else:
+        return append_to_uri_path(run.info.artifact_uri, artifact_path)
+
+
+# TODO: This would be much simpler if artifact_repo.download_artifacts could take the absolute path
+# or no path.
+def _get_root_uri_and_artifact_path(artifact_uri):
+    """Parse the artifact_uri to get the root_uri and artifact_path.
+
+    Args:
+        artifact_uri: The *absolute* URI of the artifact.
+    """
+    if os.path.exists(artifact_uri):
+        if not is_windows():
+            # If we're dealing with local files, just reference the direct pathing.
+            # non-nt-based file systems can directly reference path information, while nt-based
+            # systems need to url-encode special characters in directory listings to be able to
+            # resolve them (i.e., spaces converted to %20 within a file name or path listing)
+            root_uri = os.path.dirname(artifact_uri)
+            artifact_path = os.path.basename(artifact_uri)
+            return root_uri, artifact_path
+        else:  # if we're dealing with nt-based systems, we need to utilize pathname2url to encode.
+            artifact_uri = path_to_local_file_uri(artifact_uri)
+
+    parsed_uri = urllib.parse.urlparse(str(artifact_uri))
+    prefix = ""
+    if parsed_uri.scheme and not parsed_uri.path.startswith("/"):
+        # relative path is a special case, urllib does not reconstruct it properly
+        prefix = parsed_uri.scheme + ":"
+        parsed_uri = parsed_uri._replace(scheme="")
+
+    # For models:/ URIs, it doesn't make sense to initialize a ModelsArtifactRepository with only
+    # the model name portion of the URI, then call download_artifacts with the version info.
+    if ModelsArtifactRepository.is_models_uri(artifact_uri):
+        root_uri, artifact_path = ModelsArtifactRepository.split_models_uri(artifact_uri)
+    else:
+        artifact_path = posixpath.basename(parsed_uri.path)
+        parsed_uri = parsed_uri._replace(path=posixpath.dirname(parsed_uri.path))
+        root_uri = prefix + urllib.parse.urlunparse(parsed_uri)
+    return root_uri, artifact_path
+
+
+def _download_artifact_from_uri(artifact_uri, output_path=None, lineage_header_info=None):
+    """
+    Args:
+        artifact_uri: The *absolute* URI of the artifact to download.
+        output_path: The local filesystem path to which to download the artifact. If unspecified,
+            a local output path will be created.
+        lineage_header_info: The model lineage header info to be consumed by lineage services.
+    """
+    root_uri, artifact_path = _get_root_uri_and_artifact_path(artifact_uri)
+    repo = get_artifact_repository(artifact_uri=root_uri)
+
+    if isinstance(repo, ModelsArtifactRepository):
+        return repo.download_artifacts(
+            artifact_path=artifact_path,
+            dst_path=output_path,
+            lineage_header_info=lineage_header_info,
+        )
+    return repo.download_artifacts(artifact_path=artifact_path, dst_path=output_path)
+
+
+def _upload_artifact_to_uri(local_path, artifact_uri):
+    """Uploads a local artifact (file) to a specified URI.
+
+    Args:
+        local_path: The local path of the file to upload.
+        artifact_uri: The *absolute* URI of the path to upload the artifact to.
+
+    """
+    root_uri, artifact_path = _get_root_uri_and_artifact_path(artifact_uri)
+    get_artifact_repository(artifact_uri=root_uri).log_artifact(local_path, artifact_path)
+
+
+def _upload_artifacts_to_databricks(
+    source, run_id, source_host_uri=None, target_databricks_profile_uri=None
+):
+    """Copy the artifacts from ``source`` to the destination Databricks workspace (DBFS) given by
+    ``databricks_profile_uri`` or the current tracking URI.
+
+    Args:
+        source: Source location for the artifacts to copy.
+        run_id: Run ID to associate the artifacts with.
+        source_host_uri: Specifies the source artifact's host URI (e.g. Databricks tracking URI)
+            if applicable. If not given, defaults to the current tracking URI.
+        target_databricks_profile_uri: Specifies the destination Databricks host. If not given,
+            defaults to the current tracking URI.
+
+    Returns:
+        The DBFS location in the target Databricks workspace the model files have been
+        uploaded to.
+    """
+
+    with tempfile.TemporaryDirectory() as local_dir:
+        source_with_profile = add_databricks_profile_info_to_artifact_uri(source, source_host_uri)
+        _download_artifact_from_uri(source_with_profile, local_dir)
+        dest_root = "dbfs:/databricks/mlflow/tmp-external-source/"
+        dest_root_with_profile = add_databricks_profile_info_to_artifact_uri(
+            dest_root, target_databricks_profile_uri
+        )
+        dest_repo = DbfsRestArtifactRepository(dest_root_with_profile)
+        dest_artifact_path = run_id if run_id else uuid.uuid4().hex
+        # Allow uploading from the same run id multiple times by randomizing a suffix
+        if len(dest_repo.list_artifacts(dest_artifact_path)) > 0:
+            dest_artifact_path = dest_artifact_path + "-" + uuid.uuid4().hex[0:4]
+        dest_repo.log_artifacts(local_dir, artifact_path=dest_artifact_path)
+        dirname = pathlib.PurePath(source).name  # innermost directory name
+        return posixpath.join(dest_root, dest_artifact_path, dirname)  # new source

venv/lib/python3.9/site-packages/mlflow/tracking/context/abstract_context.py

@@ -0,0 +1,35 @@
+from abc import ABCMeta, abstractmethod
+
+from mlflow.utils.annotations import developer_stable
+
+
+@developer_stable
+class RunContextProvider:
+    """
+    Abstract base class for context provider objects specifying custom tags at run-creation time
+    (e.g. tags specifying the git repo with which the run is associated).
+
+    When a run is created via the fluent ``mlflow.start_run`` method, MLflow iterates through all
+    registered RunContextProviders. For each context provider where ``in_context`` returns ``True``,
+    MLflow calls the ``tags`` method on the context provider to compute context tags for the run.
+    All context tags are then merged together and set on the newly-created run.
+    """
+
+    __metaclass__ = ABCMeta
+
+    @abstractmethod
+    def in_context(self):
+        """Determine if MLflow is running in this context.
+
+        Returns:
+            bool indicating if in this context.
+
+        """
+
+    @abstractmethod
+    def tags(self):
+        """Generate context-specific tags.
+
+        Returns:
+            dict of tags.
+        """

venv/lib/python3.9/site-packages/mlflow/tracking/context/databricks_cluster_context.py

@@ -0,0 +1,15 @@
+from mlflow.tracking.context.abstract_context import RunContextProvider
+from mlflow.utils import databricks_utils
+from mlflow.utils.mlflow_tags import MLFLOW_DATABRICKS_CLUSTER_ID
+
+
+class DatabricksClusterRunContext(RunContextProvider):
+    def in_context(self):
+        return databricks_utils.is_in_cluster()
+
+    def tags(self):
+        cluster_id = databricks_utils.get_cluster_id()
+        tags = {}
+        if cluster_id is not None:
+            tags[MLFLOW_DATABRICKS_CLUSTER_ID] = cluster_id
+        return tags

venv/lib/python3.9/site-packages/mlflow/tracking/context/databricks_command_context.py

@@ -0,0 +1,15 @@
+from mlflow.tracking.context.abstract_context import RunContextProvider
+from mlflow.utils import databricks_utils
+from mlflow.utils.mlflow_tags import MLFLOW_DATABRICKS_NOTEBOOK_COMMAND_ID
+
+
+class DatabricksCommandRunContext(RunContextProvider):
+    def in_context(self):
+        return databricks_utils.get_job_group_id() is not None
+
+    def tags(self):
+        job_group_id = databricks_utils.get_job_group_id()
+        tags = {}
+        if job_group_id is not None:
+            tags[MLFLOW_DATABRICKS_NOTEBOOK_COMMAND_ID] = job_group_id
+        return tags

venv/lib/python3.9/site-packages/mlflow/tracking/context/databricks_job_context.py

@@ -0,0 +1,49 @@
+from mlflow.entities import SourceType
+from mlflow.tracking.context.abstract_context import RunContextProvider
+from mlflow.utils import databricks_utils
+from mlflow.utils.mlflow_tags import (
+    MLFLOW_DATABRICKS_JOB_ID,
+    MLFLOW_DATABRICKS_JOB_RUN_ID,
+    MLFLOW_DATABRICKS_JOB_TYPE,
+    MLFLOW_DATABRICKS_WEBAPP_URL,
+    MLFLOW_DATABRICKS_WORKSPACE_ID,
+    MLFLOW_DATABRICKS_WORKSPACE_URL,
+    MLFLOW_SOURCE_NAME,
+    MLFLOW_SOURCE_TYPE,
+)
+
+
+class DatabricksJobRunContext(RunContextProvider):
+    def in_context(self):
+        return databricks_utils.is_in_databricks_job()
+
+    def tags(self):
+        job_id = databricks_utils.get_job_id()
+        job_run_id = databricks_utils.get_job_run_id()
+        job_type = databricks_utils.get_job_type()
+        webapp_url = databricks_utils.get_webapp_url()
+        workspace_url = databricks_utils.get_workspace_url()
+        workspace_url_fallback, workspace_id = databricks_utils.get_workspace_info_from_dbutils()
+        tags = {
+            MLFLOW_SOURCE_NAME: (
+                f"jobs/{job_id}/run/{job_run_id}"
+                if job_id is not None and job_run_id is not None
+                else None
+            ),
+            MLFLOW_SOURCE_TYPE: SourceType.to_string(SourceType.JOB),
+        }
+        if job_id is not None:
+            tags[MLFLOW_DATABRICKS_JOB_ID] = job_id
+        if job_run_id is not None:
+            tags[MLFLOW_DATABRICKS_JOB_RUN_ID] = job_run_id
+        if job_type is not None:
+            tags[MLFLOW_DATABRICKS_JOB_TYPE] = job_type
+        if webapp_url is not None:
+            tags[MLFLOW_DATABRICKS_WEBAPP_URL] = webapp_url
+        if workspace_url is not None:
+            tags[MLFLOW_DATABRICKS_WORKSPACE_URL] = workspace_url
+        elif workspace_url_fallback is not None:
+            tags[MLFLOW_DATABRICKS_WORKSPACE_URL] = workspace_url_fallback
+        if workspace_id is not None:
+            tags[MLFLOW_DATABRICKS_WORKSPACE_ID] = workspace_id
+        return tags

venv/lib/python3.9/site-packages/mlflow/tracking/context/databricks_notebook_context.py

@@ -0,0 +1,41 @@
+from mlflow.entities import SourceType
+from mlflow.tracking.context.abstract_context import RunContextProvider
+from mlflow.utils import databricks_utils
+from mlflow.utils.mlflow_tags import (
+    MLFLOW_DATABRICKS_NOTEBOOK_ID,
+    MLFLOW_DATABRICKS_NOTEBOOK_PATH,
+    MLFLOW_DATABRICKS_WEBAPP_URL,
+    MLFLOW_DATABRICKS_WORKSPACE_ID,
+    MLFLOW_DATABRICKS_WORKSPACE_URL,
+    MLFLOW_SOURCE_NAME,
+    MLFLOW_SOURCE_TYPE,
+)
+
+
+class DatabricksNotebookRunContext(RunContextProvider):
+    def in_context(self):
+        return databricks_utils.is_in_databricks_notebook()
+
+    def tags(self):
+        notebook_id = databricks_utils.get_notebook_id()
+        notebook_path = databricks_utils.get_notebook_path()
+        webapp_url = databricks_utils.get_webapp_url()
+        workspace_url = databricks_utils.get_workspace_url()
+        workspace_url_fallback, workspace_id = databricks_utils.get_workspace_info_from_dbutils()
+        tags = {
+            MLFLOW_SOURCE_NAME: notebook_path,
+            MLFLOW_SOURCE_TYPE: SourceType.to_string(SourceType.NOTEBOOK),
+        }
+        if notebook_id is not None:
+            tags[MLFLOW_DATABRICKS_NOTEBOOK_ID] = notebook_id
+        if notebook_path is not None:
+            tags[MLFLOW_DATABRICKS_NOTEBOOK_PATH] = notebook_path
+        if webapp_url is not None:
+            tags[MLFLOW_DATABRICKS_WEBAPP_URL] = webapp_url
+        if workspace_url is not None:
+            tags[MLFLOW_DATABRICKS_WORKSPACE_URL] = workspace_url
+        elif workspace_url_fallback is not None:
+            tags[MLFLOW_DATABRICKS_WORKSPACE_URL] = workspace_url_fallback
+        if workspace_id is not None:
+            tags[MLFLOW_DATABRICKS_WORKSPACE_ID] = workspace_id
+        return tags

venv/lib/python3.9/site-packages/mlflow/tracking/context/databricks_repo_context.py

@@ -0,0 +1,43 @@
+from mlflow.tracking.context.abstract_context import RunContextProvider
+from mlflow.utils import databricks_utils
+from mlflow.utils.mlflow_tags import (
+    MLFLOW_DATABRICKS_GIT_REPO_COMMIT,
+    MLFLOW_DATABRICKS_GIT_REPO_PROVIDER,
+    MLFLOW_DATABRICKS_GIT_REPO_REFERENCE,
+    MLFLOW_DATABRICKS_GIT_REPO_REFERENCE_TYPE,
+    MLFLOW_DATABRICKS_GIT_REPO_RELATIVE_PATH,
+    MLFLOW_DATABRICKS_GIT_REPO_STATUS,
+    MLFLOW_DATABRICKS_GIT_REPO_URL,
+)
+
+
+class DatabricksRepoRunContext(RunContextProvider):
+    def in_context(self):
+        return databricks_utils.is_in_databricks_repo()
+
+    def tags(self):
+        tags = {}
+        git_repo_url = databricks_utils.get_git_repo_url()
+        git_repo_provider = databricks_utils.get_git_repo_provider()
+        git_repo_commit = databricks_utils.get_git_repo_commit()
+        git_repo_relative_path = databricks_utils.get_git_repo_relative_path()
+        git_repo_reference = databricks_utils.get_git_repo_reference()
+        git_repo_reference_type = databricks_utils.get_git_repo_reference_type()
+        git_repo_status = databricks_utils.get_git_repo_status()
+
+        if git_repo_url is not None:
+            tags[MLFLOW_DATABRICKS_GIT_REPO_URL] = git_repo_url
+        if git_repo_provider is not None:
+            tags[MLFLOW_DATABRICKS_GIT_REPO_PROVIDER] = git_repo_provider
+        if git_repo_commit is not None:
+            tags[MLFLOW_DATABRICKS_GIT_REPO_COMMIT] = git_repo_commit
+        if git_repo_relative_path is not None:
+            tags[MLFLOW_DATABRICKS_GIT_REPO_RELATIVE_PATH] = git_repo_relative_path
+        if git_repo_reference is not None:
+            tags[MLFLOW_DATABRICKS_GIT_REPO_REFERENCE] = git_repo_reference
+        if git_repo_reference_type is not None:
+            tags[MLFLOW_DATABRICKS_GIT_REPO_REFERENCE_TYPE] = git_repo_reference_type
+        if git_repo_status is not None:
+            tags[MLFLOW_DATABRICKS_GIT_REPO_STATUS] = git_repo_status
+
+        return tags

venv/lib/python3.9/site-packages/mlflow/tracking/context/default_context.py

@@ -0,0 +1,51 @@
+import getpass
+import sys
+
+from mlflow.entities import SourceType
+from mlflow.tracking.context.abstract_context import RunContextProvider
+from mlflow.utils.credentials import read_mlflow_creds
+from mlflow.utils.mlflow_tags import (
+    MLFLOW_SOURCE_NAME,
+    MLFLOW_SOURCE_TYPE,
+    MLFLOW_USER,
+)
+
+_DEFAULT_USER = "unknown"
+
+
+def _get_user():
+    """Get the current computer username."""
+    try:
+        return getpass.getuser()
+    except ImportError:
+        return _DEFAULT_USER
+
+
+def _get_main_file():
+    if len(sys.argv) > 0:
+        return sys.argv[0]
+    return None
+
+
+def _get_source_name():
+    main_file = _get_main_file()
+    if main_file is not None:
+        return main_file
+    return "<console>"
+
+
+def _get_source_type():
+    return SourceType.LOCAL
+
+
+class DefaultRunContext(RunContextProvider):
+    def in_context(self):
+        return True
+
+    def tags(self):
+        creds = read_mlflow_creds()
+        return {
+            MLFLOW_USER: creds.username or _get_user(),
+            MLFLOW_SOURCE_NAME: _get_source_name(),
+            MLFLOW_SOURCE_TYPE: SourceType.to_string(_get_source_type()),
+        }

venv/lib/python3.9/site-packages/mlflow/tracking/context/git_context.py

@@ -0,0 +1,32 @@
+import logging
+
+from mlflow.tracking.context.abstract_context import RunContextProvider
+from mlflow.tracking.context.default_context import _get_main_file
+from mlflow.utils.git_utils import get_git_commit
+from mlflow.utils.mlflow_tags import MLFLOW_GIT_COMMIT
+
+_logger = logging.getLogger(__name__)
+
+
+def _get_source_version():
+    main_file = _get_main_file()
+    if main_file is not None:
+        return get_git_commit(main_file)
+    return None
+
+
+class GitRunContext(RunContextProvider):
+    def __init__(self):
+        self._cache = {}
+
+    @property
+    def _source_version(self):
+        if "source_version" not in self._cache:
+            self._cache["source_version"] = _get_source_version()
+        return self._cache["source_version"]
+
+    def in_context(self):
+        return self._source_version is not None
+
+    def tags(self):
+        return {MLFLOW_GIT_COMMIT: self._source_version}

venv/lib/python3.9/site-packages/mlflow/tracking/context/registry.py

@@ -0,0 +1,98 @@
+import logging
+import warnings
+from typing import Optional
+
+from mlflow.tracking.context.abstract_context import RunContextProvider
+from mlflow.tracking.context.databricks_cluster_context import DatabricksClusterRunContext
+from mlflow.tracking.context.databricks_command_context import DatabricksCommandRunContext
+from mlflow.tracking.context.databricks_job_context import DatabricksJobRunContext
+from mlflow.tracking.context.databricks_notebook_context import DatabricksNotebookRunContext
+from mlflow.tracking.context.databricks_repo_context import DatabricksRepoRunContext
+from mlflow.tracking.context.default_context import DefaultRunContext
+from mlflow.tracking.context.git_context import GitRunContext
+from mlflow.tracking.context.system_environment_context import SystemEnvironmentContext
+from mlflow.utils.plugins import get_entry_points
+
+_logger = logging.getLogger(__name__)
+
+
+class RunContextProviderRegistry:
+    """Registry for run context provider implementations
+
+    This class allows the registration of a run context provider which can be used to infer meta
+    information about the context of an MLflow experiment run. Implementations declared though the
+    entrypoints `mlflow.run_context_provider` group can be automatically registered through the
+    `register_entrypoints` method.
+
+    Registered run context providers can return tags that override those implemented in the core
+    library, however the order in which plugins are resolved is undefined.
+    """
+
+    def __init__(self):
+        self._registry = []
+
+    def register(self, run_context_provider_cls):
+        self._registry.append(run_context_provider_cls())
+
+    def register_entrypoints(self):
+        """Register tracking stores provided by other packages"""
+        for entrypoint in get_entry_points("mlflow.run_context_provider"):
+            try:
+                self.register(entrypoint.load())
+            except (AttributeError, ImportError) as exc:
+                warnings.warn(
+                    'Failure attempting to register context provider "{}": {}'.format(
+                        entrypoint.name, str(exc)
+                    ),
+                    stacklevel=2,
+                )
+
+    def __iter__(self):
+        return iter(self._registry)
+
+
+_run_context_provider_registry = RunContextProviderRegistry()
+_run_context_provider_registry.register(DefaultRunContext)
+_run_context_provider_registry.register(GitRunContext)
+_run_context_provider_registry.register(DatabricksNotebookRunContext)
+_run_context_provider_registry.register(DatabricksJobRunContext)
+_run_context_provider_registry.register(DatabricksClusterRunContext)
+_run_context_provider_registry.register(DatabricksCommandRunContext)
+_run_context_provider_registry.register(DatabricksRepoRunContext)
+_run_context_provider_registry.register(SystemEnvironmentContext)
+
+_run_context_provider_registry.register_entrypoints()
+
+
+def resolve_tags(tags=None, ignore: Optional[list[RunContextProvider]] = None):
+    """Generate a set of tags for the current run context. Tags are resolved in the order,
+    contexts are registered. Argument tags are applied last.
+
+    This function iterates through all run context providers in the registry. Additional context
+    providers can be registered as described in
+    :py:class:`mlflow.tracking.context.RunContextProvider`.
+
+    Args:
+        tags: A dictionary of tags to override. If specified, tags passed in this argument will
+            override those inferred from the context.
+        ignore: A list of RunContextProvider classes to exclude from the resolution.
+
+    Returns:
+        A dictionary of resolved tags.
+    """
+    ignore = ignore or []
+    all_tags = {}
+    for provider in _run_context_provider_registry:
+        if any(isinstance(provider, ig) for ig in ignore):
+            continue
+
+        try:
+            if provider.in_context():
+                all_tags.update(provider.tags())
+        except Exception as e:
+            _logger.warning("Encountered unexpected error during resolving tags: %s", e)
+
+    if tags is not None:
+        all_tags.update(tags)
+
+    return all_tags

venv/lib/python3.9/site-packages/mlflow/tracking/context/system_environment_context.py

@@ -0,0 +1,15 @@
+import json
+
+from mlflow.environment_variables import MLFLOW_RUN_CONTEXT
+from mlflow.tracking.context.abstract_context import RunContextProvider
+
+# The constant MLFLOW_RUN_CONTEXT_ENV_VAR is marked as @developer_stable
+MLFLOW_RUN_CONTEXT_ENV_VAR = MLFLOW_RUN_CONTEXT.name
+
+
+class SystemEnvironmentContext(RunContextProvider):
+    def in_context(self):
+        return MLFLOW_RUN_CONTEXT.defined
+
+    def tags(self):
+        return json.loads(MLFLOW_RUN_CONTEXT.get())

venv/lib/python3.9/site-packages/mlflow/tracking/default_experiment/__init__.py

@@ -0,0 +1 @@
+DEFAULT_EXPERIMENT_ID = "0"

venv/lib/python3.9/site-packages/mlflow/tracking/default_experiment/abstract_context.py

@@ -0,0 +1,43 @@
+from abc import ABCMeta, abstractmethod
+
+from mlflow.utils.annotations import developer_stable
+
+
+@developer_stable
+class DefaultExperimentProvider:
+    """
+    Abstract base class for objects that provide the ID of an MLflow Experiment based on the
+    current client context. For example, when the MLflow client is running in a Databricks Job,
+    a provider is used to obtain the ID of the MLflow Experiment associated with the Job.
+
+    Usually the experiment_id is set explicitly by the user, but if the experiment is not set,
+    MLflow computes a default experiment id based on different contexts.
+    When an experiment is created via the fluent ``mlflow.start_run`` method, MLflow iterates
+    through the registered ``DefaultExperimentProvider``s until it finds one whose
+    ``in_context()`` method returns ``True``; MLflow then calls the provider's
+    ``get_experiment_id()`` method and uses the resulting experiment ID for Tracking operations.
+    """
+
+    __metaclass__ = ABCMeta
+
+    @abstractmethod
+    def in_context(self):
+        """Determine if the MLflow client is running in a context where this provider can
+        identify an associated MLflow Experiment ID.
+
+        Returns:
+            True if the MLflow client is running in a context where the provider
+            can identify an associated MLflow Experiment ID. False otherwise.
+
+        """
+
+    @abstractmethod
+    def get_experiment_id(self):
+        """Provide the MLflow Experiment ID for the current MLflow client context.
+
+        Assumes that ``in_context()`` is ``True``.
+
+        Returns:
+            The ID of the MLflow Experiment associated with the current context.
+
+        """

venv/lib/python3.9/site-packages/mlflow/tracking/default_experiment/databricks_notebook_experiment_provider.py

@@ -0,0 +1,44 @@
+from mlflow.exceptions import MlflowException
+from mlflow.protos import databricks_pb2
+from mlflow.tracking.client import MlflowClient
+from mlflow.tracking.default_experiment.abstract_context import DefaultExperimentProvider
+from mlflow.utils import databricks_utils
+from mlflow.utils.mlflow_tags import MLFLOW_EXPERIMENT_SOURCE_ID, MLFLOW_EXPERIMENT_SOURCE_TYPE
+
+
+class DatabricksNotebookExperimentProvider(DefaultExperimentProvider):
+    _resolved_notebook_experiment_id = None
+
+    def in_context(self):
+        return databricks_utils.is_in_databricks_notebook()
+
+    def get_experiment_id(self):
+        if DatabricksNotebookExperimentProvider._resolved_notebook_experiment_id:
+            return DatabricksNotebookExperimentProvider._resolved_notebook_experiment_id
+
+        source_notebook_id = databricks_utils.get_notebook_id()
+        source_notebook_name = databricks_utils.get_notebook_path()
+        tags = {
+            MLFLOW_EXPERIMENT_SOURCE_ID: source_notebook_id,
+        }
+
+        if databricks_utils.is_in_databricks_repo_notebook():
+            tags[MLFLOW_EXPERIMENT_SOURCE_TYPE] = "REPO_NOTEBOOK"
+
+        # With the presence of the source id, the following is a get or create in which it will
+        # return the corresponding experiment if one exists for the repo notebook.
+        # For non-repo notebooks, it will raise an exception and we will use source_notebook_id
+        try:
+            experiment_id = MlflowClient().create_experiment(source_notebook_name, None, tags)
+        except MlflowException as e:
+            if e.error_code == databricks_pb2.ErrorCode.Name(
+                databricks_pb2.INVALID_PARAMETER_VALUE
+            ):
+                # If determined that it is not a repo notebook
+                experiment_id = source_notebook_id
+            else:
+                raise e
+
+        DatabricksNotebookExperimentProvider._resolved_notebook_experiment_id = experiment_id
+
+        return experiment_id

venv/lib/python3.9/site-packages/mlflow/tracking/default_experiment/registry.py

@@ -0,0 +1,75 @@
+import logging
+import warnings
+
+from mlflow.tracking.default_experiment import DEFAULT_EXPERIMENT_ID
+from mlflow.tracking.default_experiment.databricks_notebook_experiment_provider import (
+    DatabricksNotebookExperimentProvider,
+)
+from mlflow.utils.plugins import get_entry_points
+
+_logger = logging.getLogger(__name__)
+# Listed below are the list of providers, which are used to provide MLflow Experiment IDs based on
+# the current context where the MLflow client is running when the user has not explicitly set
+# an experiment. The order below is the order in which the these providers are registered.
+_EXPERIMENT_PROVIDERS = (DatabricksNotebookExperimentProvider,)
+
+
+class DefaultExperimentProviderRegistry:
+    """Registry for default experiment provider implementations
+
+    This class allows the registration of default experiment providers, which are used to provide
+    MLflow Experiment IDs based on the current context where the MLflow client is running when
+    the user has not explicitly set an experiment. Implementations declared though the entrypoints
+    `mlflow.default_experiment_provider` group can be automatically registered through the
+    `register_entrypoints` method.
+    """
+
+    def __init__(self):
+        self._registry = []
+
+    def register(self, default_experiment_provider_cls):
+        self._registry.append(default_experiment_provider_cls())
+
+    def register_entrypoints(self):
+        """Register tracking stores provided by other packages"""
+        for entrypoint in get_entry_points("mlflow.default_experiment_provider"):
+            try:
+                self.register(entrypoint.load())
+            except (AttributeError, ImportError) as exc:
+                warnings.warn(
+                    "Failure attempting to register default experiment"
+                    + f'context provider "{entrypoint.name}": {exc}',
+                    stacklevel=2,
+                )
+
+    def __iter__(self):
+        return iter(self._registry)
+
+
+_default_experiment_provider_registry = DefaultExperimentProviderRegistry()
+for exp_provider in _EXPERIMENT_PROVIDERS:
+    _default_experiment_provider_registry.register(exp_provider)
+
+_default_experiment_provider_registry.register_entrypoints()
+
+
+def get_experiment_id():
+    """Get an experiment ID for the current context.
+
+    The experiment ID is fetched by querying providers, in the order that they were registered.
+    This function iterates through all default experiment context providers in the registry.
+
+    Returns:
+        An experiment_id.
+    """
+
+    experiment_id = DEFAULT_EXPERIMENT_ID
+    for provider in _default_experiment_provider_registry:
+        try:
+            if provider.in_context():
+                experiment_id = provider.get_experiment_id()
+                break
+        except Exception as e:
+            _logger.warning("Encountered unexpected error while getting experiment_id: %s", e)
+
+    return experiment_id

venv/lib/python3.9/site-packages/mlflow/tracking/metric_value_conversion_utils.py

@@ -0,0 +1,93 @@
+import sys
+
+from mlflow.exceptions import INVALID_PARAMETER_VALUE, MlflowException
+
+
+def _is_module_imported(module_name: str) -> bool:
+    return module_name in sys.modules
+
+
+def _try_get_item(x):
+    try:
+        return x.item()
+    except Exception as e:
+        raise MlflowException(
+            f"Failed to convert metric value to float: {e}",
+            error_code=INVALID_PARAMETER_VALUE,
+        )
+
+
+def _converter_requires(module_name: str):
+    """Wrapper function that checks if specified `module_name` is already imported before
+    invoking wrapped function.
+    """
+
+    def decorator(func):
+        def wrapper(x):
+            if not _is_module_imported(module_name):
+                return x
+
+            return func(x)
+
+        return wrapper
+
+    return decorator
+
+
+def convert_metric_value_to_float_if_possible(x) -> float:
+    if x is None or type(x) == float:
+        return x
+
+    converter_fns_to_try = [
+        convert_metric_value_to_float_if_ndarray,
+        convert_metric_value_to_float_if_tensorflow_tensor,
+        convert_metric_value_to_float_if_torch_tensor,
+    ]
+
+    for converter_fn in converter_fns_to_try:
+        possible_float = converter_fn(x)
+        if type(possible_float) == float:
+            return possible_float
+
+    try:
+        return float(x)
+    except ValueError:
+        return x  # let backend handle conversion if possible
+
+
+@_converter_requires("numpy")
+def convert_metric_value_to_float_if_ndarray(x):
+    import numpy as np
+
+    if isinstance(x, np.ndarray):
+        return float(_try_get_item(x))
+
+    return x
+
+
+@_converter_requires("torch")
+def convert_metric_value_to_float_if_torch_tensor(x):
+    import torch
+
+    if isinstance(x, torch.Tensor):
+        extracted_tensor_val = x.detach().cpu()
+
+        return float(_try_get_item(extracted_tensor_val))
+
+    return x
+
+
+@_converter_requires("tensorflow")
+def convert_metric_value_to_float_if_tensorflow_tensor(x):
+    import tensorflow as tf
+
+    if isinstance(x, tf.Tensor):
+        try:
+            return float(x)
+        except Exception as e:
+            raise MlflowException(
+                f"Failed to convert metric value to float: {e!r}",
+                error_code=INVALID_PARAMETER_VALUE,
+            )
+
+    return x

venv/lib/python3.9/site-packages/mlflow/tracking/multimedia.py

@@ -0,0 +1,206 @@
+"""
+Internal module implementing multi-media objects and utilities in MLflow. Multi-media objects are
+exposed to users at the top-level :py:mod:`mlflow` module.
+"""
+
+import warnings
+from typing import TYPE_CHECKING, Optional, Union
+
+if TYPE_CHECKING:
+    import numpy
+    import PIL
+
+
+COMPRESSED_IMAGE_SIZE = 256
+
+
+def compress_image_size(
+    image: "PIL.Image.Image", max_size: Optional[int] = COMPRESSED_IMAGE_SIZE
+) -> "PIL.Image.Image":
+    """
+    Scale the image to fit within a square with length `max_size` while maintaining
+    the aspect ratio.
+    """
+    # scale the image to max(width, height) <= compressed_file_max_size
+    width, height = image.size
+    if width > height:
+        new_width = max_size
+        new_height = int(height * (new_width / width))
+    else:
+        new_height = max_size
+        new_width = int(width * (new_height / height))
+    return image.resize((new_width, new_height))
+
+
+def convert_to_pil_image(image: Union["numpy.ndarray", list]) -> "PIL.Image.Image":
+    """
+    Convert a numpy array to a PIL image.
+    """
+    import numpy as np
+
+    try:
+        from PIL import Image
+    except ImportError as exc:
+        raise ImportError(
+            "Pillow is required to serialize a numpy array as an image. "
+            "Please install it via: `pip install Pillow`"
+        ) from exc
+
+    def _normalize_to_uint8(x):
+        is_int = np.issubdtype(x.dtype, np.integer)
+        low = 0
+        high = 255 if is_int else 1
+        if x.min() < low or x.max() > high:
+            if is_int:
+                raise ValueError(
+                    "Integer pixel values out of acceptable range [0, 255]. "
+                    f"Found minimum value {x.min()} and maximum value {x.max()}. "
+                    "Ensure all pixel values are within the specified range."
+                )
+            else:
+                warnings.warn(
+                    "Float pixel values out of acceptable range [0.0, 1.0]. "
+                    f"Found minimum value {x.min()} and maximum value {x.max()}. "
+                    "Rescaling values to [0.0, 1.0] with min/max scaler.",
+                    stacklevel=2,
+                )
+                # Min-max scaling
+                x = (x - x.min()) / (x.max() - x.min())
+
+        # float or bool
+        if not is_int:
+            x = x * 255
+
+        return x.astype(np.uint8)
+
+    # Ref.: https://numpy.org/doc/stable/reference/generated/numpy.dtype.kind.html#numpy-dtype-kind
+    valid_data_types = {
+        "b": "bool",
+        "i": "signed integer",
+        "u": "unsigned integer",
+        "f": "floating",
+    }
+
+    if image.dtype.kind not in valid_data_types:
+        raise TypeError(
+            f"Invalid array data type: '{image.dtype}'. "
+            f"Must be one of {list(valid_data_types.values())}"
+        )
+
+    if image.ndim not in [2, 3]:
+        raise ValueError(f"`image` must be a 2D or 3D array but got image shape: {image.shape}")
+
+    if (image.ndim == 3) and (image.shape[2] not in [1, 3, 4]):
+        raise ValueError(f"Invalid channel length: {image.shape[2]}. Must be one of [1, 3, 4]")
+
+    # squeeze a 3D grayscale image since `Image.fromarray` doesn't accept it.
+    if image.ndim == 3 and image.shape[2] == 1:
+        image = image[:, :, 0]
+
+    image = _normalize_to_uint8(image)
+    return Image.fromarray(image)
+
+
+# MLflow media object: Image
+class Image:
+    """
+    `mlflow.Image` is an image media object that provides a lightweight option
+    for handling images in MLflow.
+    The image can be a numpy array, a PIL image, or a file path to an image. The image is
+    stored as a PIL image and can be logged to MLflow using `mlflow.log_image` or
+    `mlflow.log_table`.
+
+    Args:
+        image: Image can be a numpy array, a PIL image, or a file path to an image.
+
+    .. code-block:: python
+        :caption: Example
+
+        import mlflow
+        import numpy as np
+        from PIL import Image
+
+        # Create an image as a numpy array
+        image = np.zeros((100, 100, 3), dtype=np.uint8)
+        image[:, :50] = [255, 128, 0]
+        # Create an Image object
+        image_obj = mlflow.Image(image)
+        # Convert the Image object to a list of pixel values
+        pixel_values = image_obj.to_list()
+    """
+
+    def __init__(self, image: Union["numpy.ndarray", "PIL.Image.Image", str, list]):
+        import numpy as np
+
+        try:
+            from PIL import Image
+        except ImportError as exc:
+            raise ImportError(
+                "`mlflow.Image` requires Pillow to serialize a numpy array as an image. "
+                "Please install it via: `pip install Pillow`."
+            ) from exc
+
+        if isinstance(image, str):
+            self.image = Image.open(image)
+        elif isinstance(image, (list, np.ndarray)):
+            self.image = convert_to_pil_image(np.array(image))
+        elif isinstance(image, Image.Image):
+            self.image = image
+        else:
+            raise TypeError(
+                f"Unsupported image object type: {type(image)}. "
+                "`image` must be one of numpy.ndarray, "
+                "PIL.Image.Image, or a filepath to an image."
+            )
+        self.size = self.image.size
+
+    def to_list(self):
+        """
+        Convert the image to a list of pixel values.
+
+        Returns:
+            List of pixel values.
+        """
+        return list(self.image.getdata())
+
+    def to_array(self):
+        """
+        Convert the image to a numpy array.
+
+        Returns:
+            Numpy array of pixel values.
+        """
+        import numpy as np
+
+        return np.array(self.image)
+
+    def to_pil(self):
+        """
+        Convert the image to a PIL image.
+
+        Returns:
+            PIL image.
+        """
+        return self.image
+
+    def save(self, path: str):
+        """
+        Save the image to a file.
+
+        Args:
+            path: File path to save the image.
+        """
+        self.image.save(path)
+
+    def resize(self, size: tuple[int, int]):
+        """
+        Resize the image to the specified size.
+
+        Args:
+            size: Size to resize the image to.
+
+        Returns:
+            A copy of the resized image object.
+        """
+        image = self.image.resize(size)
+        return Image(image)

venv/lib/python3.9/site-packages/mlflow/tracking/registry.py

@@ -0,0 +1,86 @@
+import warnings
+from abc import ABCMeta
+
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
+from mlflow.utils.plugins import get_entry_points
+from mlflow.utils.uri import get_uri_scheme
+
+
+class UnsupportedModelRegistryStoreURIException(MlflowException):
+    """Exception thrown when building a model registry store with an unsupported URI"""
+
+    def __init__(self, unsupported_uri, supported_uri_schemes):
+        message = (
+            " Model registry functionality is unavailable; got unsupported URI"
+            f" '{unsupported_uri}' for model registry data storage. Supported URI schemes are:"
+            f" {supported_uri_schemes}."
+            " See https://www.mlflow.org/docs/latest/tracking.html#storage for how to run"
+            " an MLflow server against one of the supported backend storage locations."
+        )
+        super().__init__(message, error_code=INVALID_PARAMETER_VALUE)
+        self.supported_uri_schemes = supported_uri_schemes
+
+
+class StoreRegistry:
+    """
+    Abstract class defining a scheme-based registry for store implementations.
+
+    This class allows the registration of a function or class to provide an
+    implementation for a given scheme of `store_uri` through the `register`
+    methods. Implementations declared though the entrypoints can be automatically
+    registered through the `register_entrypoints` method.
+
+    When instantiating a store through the `get_store` method, the scheme of
+    the store URI provided (or inferred from environment) will be used to
+    select which implementation to instantiate, which will be called with same
+    arguments passed to the `get_store` method.
+    """
+
+    __metaclass__ = ABCMeta
+
+    def __init__(self, group_name):
+        self._registry = {}
+        self.group_name = group_name
+
+    def register(self, scheme, store_builder):
+        self._registry[scheme] = store_builder
+
+    def register_entrypoints(self):
+        """Register tracking stores provided by other packages"""
+        for entrypoint in get_entry_points(self.group_name):
+            try:
+                self.register(entrypoint.name, entrypoint.load())
+            except (AttributeError, ImportError) as exc:
+                warnings.warn(
+                    'Failure attempting to register store for scheme "{}": {}'.format(
+                        entrypoint.name, str(exc)
+                    ),
+                    stacklevel=2,
+                )
+
+    def get_store_builder(self, store_uri):
+        """Get a store from the registry based on the scheme of store_uri
+
+        Args:
+            store_uri: The store URI. If None, it will be inferred from the environment. This
+                URI is used to select which tracking store implementation to instantiate
+                and is passed to the constructor of the implementation.
+
+        Returns:
+            A function that returns an instance of
+            ``mlflow.store.{tracking|model_registry}.AbstractStore`` that fulfills the store
+            URI requirements.
+        """
+        scheme = (
+            store_uri
+            if store_uri in {"databricks", "databricks-uc", "uc"}
+            else get_uri_scheme(store_uri)
+        )
+        try:
+            store_builder = self._registry[scheme]
+        except KeyError:
+            raise UnsupportedModelRegistryStoreURIException(
+                unsupported_uri=store_uri, supported_uri_schemes=list(self._registry.keys())
+            )
+        return store_builder

venv/lib/python3.9/site-packages/mlflow/tracking/request_auth/abstract_request_auth_provider.py

@@ -0,0 +1,34 @@
+from abc import ABC, abstractmethod
+
+from mlflow.utils.annotations import developer_stable
+
+
+@developer_stable
+class RequestAuthProvider(ABC):
+    """
+    Abstract base class for specifying custom request auth to add to outgoing requests
+
+    When a request is sent, MLflow will iterate through all registered RequestAuthProviders.
+    For each provider where ``get_name`` matches auth provider name, MLflow calls the ``get_auth``
+    method on the provider to compute request auth.
+
+    The resulting request auth will then be added and sent with the request.
+    """
+
+    @abstractmethod
+    def get_name(self):
+        """Get the name of the request auth provider.
+
+        Returns:
+            str of request auth provider name.
+        """
+
+    @abstractmethod
+    def get_auth(self):
+        """
+        Generate request auth object (e.g., `requests.auth import HTTPBasicAuth`). See
+        https://requests.readthedocs.io/en/latest/user/authentication/ for more details.
+
+        Returns:
+            request auth object.
+        """

venv/lib/python3.9/site-packages/mlflow/tracking/request_auth/registry.py

@@ -0,0 +1,60 @@
+import warnings
+
+from mlflow.utils.plugins import get_entry_points
+
+REQUEST_AUTH_PROVIDER_ENTRYPOINT = "mlflow.request_auth_provider"
+
+
+class RequestAuthProviderRegistry:
+    def __init__(self):
+        self._registry = []
+
+    def register(self, request_auth_provider):
+        self._registry.append(request_auth_provider())
+
+    def register_entrypoints(self):
+        for entrypoint in get_entry_points(REQUEST_AUTH_PROVIDER_ENTRYPOINT):
+            try:
+                self.register(entrypoint.load())
+            except (AttributeError, ImportError) as exc:
+                warnings.warn(
+                    'Failure attempting to register request auth provider "{}": {}'.format(
+                        entrypoint.name, str(exc)
+                    ),
+                    stacklevel=2,
+                )
+
+    def __iter__(self):
+        return iter(self._registry)
+
+
+_request_auth_provider_registry = RequestAuthProviderRegistry()
+_request_auth_provider_registry.register_entrypoints()
+
+
+def fetch_auth(request_auth):
+    """
+    Find the request auth from registered providers based on the auth provider's name.
+    The auth provider's name can be provided through environment variable `MLFLOW_TRACKING_AUTH`.
+
+    This function iterates through all request auth providers in the registry. Additional context
+    providers can be registered as described in
+    :py:class:`mlflow.tracking.request_auth.RequestAuthProvider`.
+
+    Args:
+        request_auth: The name of request auth provider.
+
+    Returns:
+        The auth object.
+    """
+
+    for auth_provider in _request_auth_provider_registry:
+        if auth_provider.get_name() == request_auth:
+            return auth_provider.get_auth()
+
+    warnings.warn(
+        f"Could not find any registered plugin for {request_auth}. "
+        "No authentication header will be added. Please check your "
+        "provider documentation for installing the right plugin or "
+        "correct provider name."
+    )

venv/lib/python3.9/site-packages/mlflow/tracking/request_header/abstract_request_header_provider.py

@@ -0,0 +1,36 @@
+from abc import ABCMeta, abstractmethod
+
+from mlflow.utils.annotations import developer_stable
+
+
+@developer_stable
+class RequestHeaderProvider:
+    """
+    Abstract base class for specifying custom request headers to add to outgoing requests
+    (e.g. request headers specifying the environment from which mlflow is running).
+
+    When a request is sent, MLflow will iterate through all registered RequestHeaderProviders.
+    For each provider where ``in_context`` returns ``True``, MLflow calls the ``request_headers``
+    method on the provider to compute request headers.
+
+    All resulting request headers will then be merged together and sent with the request.
+    """
+
+    __metaclass__ = ABCMeta
+
+    @abstractmethod
+    def in_context(self):
+        """Determine if MLflow is running in this context.
+
+        Returns:
+            bool indicating if in this context.
+
+        """
+
+    @abstractmethod
+    def request_headers(self):
+        """Generate context-specific request headers.
+
+        Returns:
+            dict of request headers.
+        """

venv/lib/python3.9/site-packages/mlflow/tracking/request_header/databricks_request_header_provider.py

@@ -0,0 +1,38 @@
+from mlflow.tracking.request_header.abstract_request_header_provider import RequestHeaderProvider
+from mlflow.utils import databricks_utils
+
+
+class DatabricksRequestHeaderProvider(RequestHeaderProvider):
+    """
+    Provides request headers indicating the type of Databricks environment from which a request
+    was made.
+    """
+
+    def in_context(self):
+        return (
+            databricks_utils.is_in_cluster()
+            or databricks_utils.is_in_databricks_notebook()
+            or databricks_utils.is_in_databricks_job()
+        )
+
+    def request_headers(self):
+        request_headers = {}
+        if databricks_utils.is_in_databricks_notebook():
+            request_headers["notebook_id"] = databricks_utils.get_notebook_id()
+        if databricks_utils.is_in_databricks_job():
+            request_headers["job_id"] = databricks_utils.get_job_id()
+            request_headers["job_run_id"] = databricks_utils.get_job_run_id()
+            request_headers["job_type"] = databricks_utils.get_job_type()
+        if databricks_utils.is_in_cluster():
+            request_headers["cluster_id"] = databricks_utils.get_cluster_id()
+        command_run_id = databricks_utils.get_command_run_id()
+        if command_run_id is not None:
+            request_headers["command_run_id"] = command_run_id
+        workload_id = databricks_utils.get_workload_id()
+        workload_class = databricks_utils.get_workload_class()
+        if workload_id is not None:
+            request_headers["workload_id"] = workload_id
+        if workload_class is not None:
+            request_headers["workload_class"] = workload_class
+
+        return request_headers

venv/lib/python3.9/site-packages/mlflow/tracking/request_header/default_request_header_provider.py

@@ -0,0 +1,17 @@
+from mlflow import __version__
+from mlflow.tracking.request_header.abstract_request_header_provider import RequestHeaderProvider
+
+_USER_AGENT = "User-Agent"
+_DEFAULT_HEADERS = {_USER_AGENT: f"mlflow-python-client/{__version__}"}
+
+
+class DefaultRequestHeaderProvider(RequestHeaderProvider):
+    """
+    Provides default request headers for outgoing request.
+    """
+
+    def in_context(self):
+        return True
+
+    def request_headers(self):
+        return dict(**_DEFAULT_HEADERS)

venv/lib/python3.9/site-packages/mlflow/tracking/request_header/registry.py

@@ -0,0 +1,79 @@
+import logging
+import warnings
+
+from mlflow.tracking.request_header.databricks_request_header_provider import (
+    DatabricksRequestHeaderProvider,
+)
+from mlflow.tracking.request_header.default_request_header_provider import (
+    DefaultRequestHeaderProvider,
+)
+from mlflow.utils.plugins import get_entry_points
+
+_logger = logging.getLogger(__name__)
+
+
+class RequestHeaderProviderRegistry:
+    def __init__(self):
+        self._registry = []
+
+    def register(self, request_header_provider):
+        self._registry.append(request_header_provider())
+
+    def register_entrypoints(self):
+        """Register tracking stores provided by other packages"""
+        for entrypoint in get_entry_points("mlflow.request_header_provider"):
+            try:
+                self.register(entrypoint.load())
+            except (AttributeError, ImportError) as exc:
+                warnings.warn(
+                    'Failure attempting to register request header provider "{}": {}'.format(
+                        entrypoint.name, str(exc)
+                    ),
+                    stacklevel=2,
+                )
+
+    def __iter__(self):
+        return iter(self._registry)
+
+
+_request_header_provider_registry = RequestHeaderProviderRegistry()
+_request_header_provider_registry.register(DatabricksRequestHeaderProvider)
+_request_header_provider_registry.register(DefaultRequestHeaderProvider)
+
+_request_header_provider_registry.register_entrypoints()
+
+
+def resolve_request_headers(request_headers=None):
+    """Generate a set of request headers from registered providers.
+
+    Request headers are resolved in the order that providers are registered. Argument headers are
+    applied last. This function iterates through all request header providers in the registry.
+    Additional context providers can be registered as described in
+    :py:class:`mlflow.tracking.request_header.RequestHeaderProvider`.
+
+    Args:
+        request_headers: A dictionary of request headers to override. If specified, headers passed
+            in this argument will override those inferred from the context.
+
+    Returns:
+        A dictionary of resolved headers.
+    """
+
+    all_request_headers = {}
+    for provider in _request_header_provider_registry:
+        try:
+            if provider.in_context():
+                # all_request_headers.update(provider.request_headers())
+                for header, value in provider.request_headers().items():
+                    all_request_headers[header] = (
+                        f"{all_request_headers[header]} {value}"
+                        if header in all_request_headers
+                        else value
+                    )
+        except Exception as e:
+            _logger.warning("Encountered unexpected error during resolving request headers: %s", e)
+
+    if request_headers is not None:
+        all_request_headers.update(request_headers)
+
+    return all_request_headers