star

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)