star

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

@@ -0,0 +1,22 @@
+"""
+An MLflow tracking server has two properties related to how data is stored: *backend store* to
+record ML experiments, runs, parameters, metrics, etc., and *artifact store* to store run
+artifacts like models, plots, images, etc.
+
+Several constants are used by multiple backend store implementations.
+"""
+
+# Path to default location for backend when using local FileStore or ArtifactStore.
+# Also used as default location for artifacts, when not provided, in non local file based backends
+# (eg MySQL)
+DEFAULT_LOCAL_FILE_AND_ARTIFACT_PATH = "./mlruns"
+# Used for defining the artifacts uri (`--default-artifact-root`) for the tracking server when
+# configuring the server to use the option `--serve-artifacts` mode. This default can be
+# overridden by specifying an override to `--default-artifact-root` for the MLflow tracking server.
+# When the server is not operating in `--serve-artifacts` configuration, the default artifact
+# storage location will be `DEFAULT_LOCAL_FILE_AND_ARTIFACT_PATH`.
+DEFAULT_ARTIFACTS_URI = "mlflow-artifacts:/"
+SEARCH_MAX_RESULTS_DEFAULT = 1000
+SEARCH_MAX_RESULTS_THRESHOLD = 50000
+GET_METRIC_HISTORY_MAX_RESULTS = 25000
+SEARCH_TRACES_DEFAULT_MAX_RESULTS = 100

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

@@ -0,0 +1,682 @@
+from abc import ABCMeta, abstractmethod
+from typing import Optional
+
+from mlflow.entities import (
+    DatasetInput,
+    TraceInfo,
+    ViewType,
+)
+from mlflow.entities.metric import MetricWithRunId
+from mlflow.entities.trace_status import TraceStatus
+from mlflow.exceptions import MlflowException
+from mlflow.store.entities.paged_list import PagedList
+from mlflow.store.tracking import SEARCH_MAX_RESULTS_DEFAULT, SEARCH_TRACES_DEFAULT_MAX_RESULTS
+from mlflow.utils.annotations import developer_stable
+from mlflow.utils.async_logging.async_logging_queue import AsyncLoggingQueue
+from mlflow.utils.async_logging.run_operations import RunOperations
+
+
+@developer_stable
+class AbstractStore:
+    """
+    Abstract class for Backend Storage.
+    This class defines the API interface for front ends to connect with various types of backends.
+    """
+
+    __metaclass__ = ABCMeta
+
+    def __init__(self):
+        """
+        Empty constructor for now. This is deliberately not marked as abstract, else every
+        derived class would be forced to create one.
+        """
+        self._async_logging_queue = AsyncLoggingQueue(logging_func=self.log_batch)
+
+    @abstractmethod
+    def search_experiments(
+        self,
+        view_type=ViewType.ACTIVE_ONLY,
+        max_results=SEARCH_MAX_RESULTS_DEFAULT,
+        filter_string=None,
+        order_by=None,
+        page_token=None,
+    ):
+        """
+        Search for experiments that match the specified search query.
+
+        Args:
+            view_type: One of enum values ``ACTIVE_ONLY``, ``DELETED_ONLY``, or ``ALL``
+                defined in :py:class:`mlflow.entities.ViewType`.
+            max_results: Maximum number of experiments desired. Certain server backend may apply
+                its own limit.
+            filter_string: Filter query string (e.g., ``"name = 'my_experiment'"``), defaults to
+                searching for all experiments. The following identifiers, comparators, and logical
+                operators are supported.
+
+                Identifiers
+                  - ``name``: Experiment name
+                  - ``creation_time``: Experiment creation time
+                  - ``last_update_time``: Experiment last update time
+                  - ``tags.<tag_key>``: Experiment tag. If ``tag_key`` contains
+                    spaces, it must be wrapped with backticks (e.g., ``"tags.`extra key`"``).
+
+                Comparators for string attributes and tags
+                  - ``=``: Equal to
+                  - ``!=``: Not equal to
+                  - ``LIKE``: Case-sensitive pattern match
+                  - ``ILIKE``: Case-insensitive pattern match
+
+                Comparators for numeric attributes
+                  - ``=``: Equal to
+                  - ``!=``: Not equal to
+                  - ``<``: Less than
+                  - ``<=``: Less than or equal to
+                  - ``>``: Greater than
+                  - ``>=``: Greater than or equal to
+
+                Logical operators
+                  - ``AND``: Combines two sub-queries and returns True if both of them are True.
+
+            order_by: List of columns to order by. The ``order_by`` column can contain an optional
+                ``DESC`` or ``ASC`` value (e.g., ``"name DESC"``). The default ordering is ``ASC``,
+                so ``"name"`` is equivalent to ``"name ASC"``. If unspecified, defaults to
+                ``["last_update_time DESC"]``, which lists experiments updated most recently first.
+                The following fields are supported:
+
+                - ``experiment_id``: Experiment ID
+                - ``name``: Experiment name
+                - ``creation_time``: Experiment creation time
+                - ``last_update_time``: Experiment last update time
+
+            page_token: Token specifying the next page of results. It should be obtained from
+                a ``search_experiments`` call.
+
+        Returns:
+            A :py:class:`PagedList <mlflow.store.entities.PagedList>` of
+            :py:class:`Experiment <mlflow.entities.Experiment>` objects. The pagination token
+            for the next page can be obtained via the ``token`` attribute of the object.
+
+        """
+
+    @abstractmethod
+    def create_experiment(self, name, artifact_location, tags):
+        """
+        Create a new experiment.
+        If an experiment with the given name already exists, throws exception.
+
+        Args:
+            name: Desired name for an experiment.
+            artifact_location: Base location for artifacts in runs. May be None.
+            tags: Experiment tags to set upon experiment creation
+
+        Returns:
+            experiment_id (string) for the newly created experiment if successful, else None.
+
+        """
+
+    @abstractmethod
+    def get_experiment(self, experiment_id):
+        """
+        Fetch the experiment by ID from the backend store.
+
+        Args:
+            experiment_id: String id for the experiment
+
+        Returns:
+            A single :py:class:`mlflow.entities.Experiment` object if it exists,
+            otherwise raises an exception.
+        """
+
+    def get_experiment_by_name(self, experiment_name):
+        """
+        Fetch the experiment by name from the backend store.
+
+        Args:
+            experiment_name: Name of experiment
+
+        Returns:
+            A single :py:class:`mlflow.entities.Experiment` object if it exists.
+        """
+
+    @abstractmethod
+    def delete_experiment(self, experiment_id):
+        """
+        Delete the experiment from the backend store. Deleted experiments can be restored until
+        permanently deleted.
+
+        Args:
+            experiment_id: String id for the experiment.
+        """
+
+    @abstractmethod
+    def restore_experiment(self, experiment_id):
+        """
+        Restore deleted experiment unless it is permanently deleted.
+
+        Args:
+            experiment_id: String id for the experiment.
+        """
+
+    @abstractmethod
+    def rename_experiment(self, experiment_id, new_name):
+        """
+        Update an experiment's name. The new name must be unique.
+
+        Args:
+            experiment_id: String id for the experiment.
+            new_name: New name for the experiment.
+        """
+
+    @abstractmethod
+    def get_run(self, run_id):
+        """
+        Fetch the run from backend store. The resulting :py:class:`Run <mlflow.entities.Run>`
+        contains a collection of run metadata - :py:class:`RunInfo <mlflow.entities.RunInfo>`,
+        as well as a collection of run parameters, tags, and metrics -
+        :py:class:`RunData <mlflow.entities.RunData>`. In the case where multiple metrics with the
+        same key are logged for the run, the :py:class:`RunData <mlflow.entities.RunData>` contains
+        the value at the latest timestamp for each metric. If there are multiple values with the
+        latest timestamp for a given metric, the maximum of these values is returned.
+
+        Args:
+            run_id: Unique identifier for the run.
+
+        Returns:
+            A single :py:class:`mlflow.entities.Run` object, if the run exists. Otherwise,
+            raises an exception.
+        """
+
+    @abstractmethod
+    def update_run_info(self, run_id, run_status, end_time, run_name):
+        """
+        Update the metadata of the specified run.
+
+        Returns:
+            mlflow.entities.RunInfo: Describing the updated run.
+        """
+
+    @abstractmethod
+    def create_run(self, experiment_id, user_id, start_time, tags, run_name):
+        """
+        Create a run under the specified experiment ID, setting the run's status to "RUNNING"
+        and the start time to the current time.
+
+        Args:
+            experiment_id: String id of the experiment for this run.
+            user_id: ID of the user launching this run.
+            start_time: Start time of the run.
+            tags: A dictionary of string keys and string values.
+            run_name: Name of the run.
+
+        Returns:
+            The created Run object
+        """
+
+    @abstractmethod
+    def delete_run(self, run_id):
+        """
+        Delete a run.
+
+        Args:
+            run_id: The ID of the run to delete.
+
+        """
+
+    @abstractmethod
+    def restore_run(self, run_id):
+        """
+        Restore a run.
+
+        Args:
+            run_id: The ID of the run to restore.
+
+        """
+
+    # TODO: rename this to create_trace_info
+    def start_trace(
+        self,
+        experiment_id: str,
+        timestamp_ms: int,
+        request_metadata: dict[str, str],
+        tags: dict[str, str],
+    ) -> TraceInfo:
+        """
+        Start an initial TraceInfo object in the backend store.
+
+        Args:
+            experiment_id: String id of the experiment for this run.
+            timestamp_ms: Start time of the trace, in milliseconds since the UNIX epoch.
+            request_metadata: Metadata of the trace.
+            tags: Tags of the trace.
+
+        Returns:
+            The created TraceInfo object.
+        """
+        raise NotImplementedError
+
+    # TODO: rename this to update_trace_info
+    # can we pass in execution_time_ms instead of timestamp_ms directly?
+    def end_trace(
+        self,
+        request_id: str,
+        timestamp_ms: int,
+        status: TraceStatus,
+        request_metadata: dict[str, str],
+        tags: dict[str, str],
+    ) -> TraceInfo:
+        """
+        Update the TraceInfo object in the backend store with the completed trace info.
+
+        Args:
+            request_id : Unique string identifier of the trace.
+            timestamp_ms: End time of the trace, in milliseconds. The execution time field
+                in the TraceInfo will be calculated by subtracting the start time from this.
+            status: Status of the trace.
+            request_metadata: Metadata of the trace. This will be merged with the existing
+                metadata logged during the start_trace call.
+            tags: Tags of the trace. This will be merged with the existing tags logged
+                during the start_trace or set_trace_tag calls.
+
+        Returns:
+            The updated TraceInfo object.
+        """
+        raise NotImplementedError
+
+    def delete_traces(
+        self,
+        experiment_id: str,
+        max_timestamp_millis: Optional[int] = None,
+        max_traces: Optional[int] = None,
+        request_ids: Optional[list[str]] = None,
+    ) -> int:
+        """
+        Delete traces based on the specified criteria.
+
+        - Either `max_timestamp_millis` or `request_ids` must be specified, but not both.
+        - `max_traces` can't be specified if `request_ids` is specified.
+
+        Args:
+            experiment_id: ID of the associated experiment.
+            max_timestamp_millis: The maximum timestamp in milliseconds since the UNIX epoch for
+                deleting traces. Traces older than or equal to this timestamp will be deleted.
+            max_traces: The maximum number of traces to delete. If max_traces is specified, and
+                it is less than the number of traces that would be deleted based on the
+                max_timestamp_millis, the oldest traces will be deleted first.
+            request_ids: A set of request IDs to delete.
+
+        Returns:
+            The number of traces deleted.
+        """
+        # request_ids can't be an empty list of string
+        if max_timestamp_millis is None and not request_ids:
+            raise MlflowException.invalid_parameter_value(
+                "Either `max_timestamp_millis` or `request_ids` must be specified.",
+            )
+        if max_timestamp_millis and request_ids:
+            raise MlflowException.invalid_parameter_value(
+                "Only one of `max_timestamp_millis` and `request_ids` can be specified.",
+            )
+        if request_ids and max_traces is not None:
+            raise MlflowException.invalid_parameter_value(
+                "`max_traces` can't be specified if `request_ids` is specified.",
+            )
+        if max_traces is not None and max_traces <= 0:
+            raise MlflowException.invalid_parameter_value(
+                f"`max_traces` must be a positive integer, received {max_traces}.",
+            )
+        return self._delete_traces(experiment_id, max_timestamp_millis, max_traces, request_ids)
+
+    def _delete_traces(
+        self,
+        experiment_id: str,
+        max_timestamp_millis: Optional[int] = None,
+        max_traces: Optional[int] = None,
+        request_ids: Optional[list[str]] = None,
+    ) -> int:
+        raise NotImplementedError
+
+    def get_trace_info(self, request_id: str) -> TraceInfo:
+        """
+        Get the trace matching the `request_id`.
+
+        Args:
+            request_id: String id of the trace to fetch.
+
+        Returns:
+            The fetched Trace object, of type ``mlflow.entities.TraceInfo``.
+        """
+        raise NotImplementedError
+
+    def search_traces(
+        self,
+        experiment_ids: list[str],
+        filter_string: Optional[str] = None,
+        max_results: int = SEARCH_TRACES_DEFAULT_MAX_RESULTS,
+        order_by: Optional[list[str]] = None,
+        page_token: Optional[str] = None,
+    ) -> tuple[list[TraceInfo], Optional[str]]:
+        """
+        Return traces that match the given list of search expressions within the experiments.
+
+        Args:
+            experiment_ids: List of experiment ids to scope the search.
+            filter_string: A search filter string.
+            max_results: Maximum number of traces desired.
+            order_by: List of order_by clauses.
+            page_token: Token specifying the next page of results. It should be obtained from
+                a ``search_traces`` call.
+
+        Returns:
+            A tuple of a list of :py:class:`TraceInfo <mlflow.entities.TraceInfo>` objects that
+            satisfy the search expressions and a pagination token for the next page of results.
+            If the underlying tracking store supports pagination, the token for the
+            next page may be obtained via the ``token`` attribute of the returned object; however,
+            some store implementations may not support pagination and thus the returned token would
+            not be meaningful in such cases.
+        """
+        raise NotImplementedError
+
+    def set_trace_tag(self, request_id: str, key: str, value: str):
+        """
+        Set a tag on the trace with the given request_id.
+
+        Args:
+            request_id: The ID of the trace.
+            key: The string key of the tag.
+            value: The string value of the tag.
+        """
+        raise NotImplementedError
+
+    def delete_trace_tag(self, request_id: str, key: str):
+        """
+        Delete a tag on the trace with the given request_id.
+
+        Args:
+            request_id: The ID of the trace.
+            key: The string key of the tag.
+        """
+        raise NotImplementedError
+
+    def log_metric(self, run_id, metric):
+        """
+        Log a metric for the specified run
+
+        Args:
+            run_id: String id for the run
+            metric: `mlflow.entities.Metric` instance to log
+        """
+        self.log_batch(run_id, metrics=[metric], params=[], tags=[])
+
+    def log_metric_async(self, run_id, metric) -> RunOperations:
+        """
+        Log a metric for the specified run in async fashion.
+
+        Args:
+            run_id: String id for the run
+            metric: `mlflow.entities.Metric` instance to log
+        """
+        return self.log_batch_async(run_id, metrics=[metric], params=[], tags=[])
+
+    def log_param(self, run_id, param):
+        """
+        Log a param for the specified run
+
+        Args:
+            run_id: String id for the run
+            param: :py:class:`mlflow.entities.Param` instance to log
+        """
+        self.log_batch(run_id, metrics=[], params=[param], tags=[])
+
+    def log_param_async(self, run_id, param) -> RunOperations:
+        """
+        Log a param for the specified run in async fashion.
+
+        Args:
+            run_id: String id for the run.
+            param: :py:class:`mlflow.entities.Param` instance to log.
+        """
+        return self.log_batch_async(run_id, metrics=[], params=[param], tags=[])
+
+    def set_experiment_tag(self, experiment_id, tag):
+        """
+        Set a tag for the specified experiment
+
+        Args:
+            experiment_id: String id for the experiment.
+            tag: :py:class:`mlflow.entities.ExperimentTag` instance to set.
+        """
+
+    def set_tag(self, run_id, tag):
+        """
+        Set a tag for the specified run
+
+        Args:
+            run_id: String id for the run.
+            tag: :py:class:`mlflow.entities.RunTag` instance to set.
+        """
+        self.log_batch(run_id, metrics=[], params=[], tags=[tag])
+
+    def set_tag_async(self, run_id, tag) -> RunOperations:
+        """
+        Set a tag for the specified run in async fashion.
+
+        Args:
+            run_id: String id for the run.
+            tag: :py:class:`mlflow.entities.RunTag` instance to set.
+        """
+        return self.log_batch_async(run_id, metrics=[], params=[], tags=[tag])
+
+    @abstractmethod
+    def get_metric_history(self, run_id, metric_key, max_results=None, page_token=None):
+        """
+        Return a list of metric objects corresponding to all values logged for a given metric
+        within a run.
+
+        Args:
+            run_id: Unique identifier for run.
+            metric_key: Metric name within the run.
+            max_results: Maximum number of metric history events (steps) to return per paged
+                query.
+            page_token: A Token specifying the next paginated set of results of metric history.
+                This value is obtained as a return value from a paginated call to GetMetricHistory.
+
+        Returns:
+            A list of :py:class:`mlflow.entities.Metric` entities if logged, else empty list.
+        """
+
+        # NB: Pagination for this API is not supported in FileStore or SQLAlchemyStore. The
+        # argument `max_results` is used as a pagination activation flag. If the `max_results`
+        # argument is not provided, this API will return a full metric history event collection
+        # without the paged queries to the backend store.
+
+    def get_metric_history_bulk_interval_from_steps(self, run_id, metric_key, steps, max_results):
+        """
+        Return a list of metric objects corresponding to all values logged
+        for a given metric within a run for the specified steps.
+
+        Args:
+            run_id: Unique identifier for run.
+            metric_key: Metric name within the run.
+            steps: List of steps for which to return metrics.
+            max_results: Maximum number of metric history events (steps) to return.
+
+        Returns:
+            A list of MetricWithRunId objects:
+                - key: Metric name within the run.
+                - value: Metric value.
+                - timestamp: Metric timestamp.
+                - step: Metric step.
+                - run_id: Unique identifier for run.
+        """
+        metrics_for_run = sorted(
+            [m for m in self.get_metric_history(run_id, metric_key) if m.step in steps],
+            key=lambda metric: (metric.step, metric.timestamp),
+        )[:max_results]
+        return [
+            MetricWithRunId(
+                run_id=run_id,
+                metric=metric,
+            )
+            for metric in metrics_for_run
+        ]
+
+    def search_runs(
+        self,
+        experiment_ids,
+        filter_string,
+        run_view_type,
+        max_results=SEARCH_MAX_RESULTS_DEFAULT,
+        order_by=None,
+        page_token=None,
+    ):
+        """
+        Return runs that match the given list of search expressions within the experiments.
+
+        Args:
+            experiment_ids: List of experiment ids to scope the search.
+            filter_string: A search filter string.
+            run_view_type: ACTIVE_ONLY, DELETED_ONLY, or ALL runs.
+            max_results: Maximum number of runs desired.
+            order_by: List of order_by clauses.
+            page_token: Token specifying the next page of results. It should be obtained from
+                a ``search_runs`` call.
+
+        Returns:
+            A :py:class:`PagedList <mlflow.store.entities.PagedList>` of
+            :py:class:`Run <mlflow.entities.Run>` objects that satisfy the search expressions.
+            If the underlying tracking store supports pagination, the token for the next page may
+            be obtained via the ``token`` attribute of the returned object; however, some store
+            implementations may not support pagination and thus the returned token would not be
+            meaningful in such cases.
+        """
+        runs, token = self._search_runs(
+            experiment_ids,
+            filter_string,
+            run_view_type,
+            max_results,
+            order_by,
+            page_token,
+        )
+        return PagedList(runs, token)
+
+    @abstractmethod
+    def _search_runs(
+        self,
+        experiment_ids,
+        filter_string,
+        run_view_type,
+        max_results,
+        order_by,
+        page_token,
+    ):
+        """
+        Return runs that match the given list of search expressions within the experiments, as
+        well as a pagination token (indicating where the next page should start). Subclasses of
+        ``AbstractStore`` should implement this method to support pagination instead of
+        ``search_runs``.
+
+        See ``search_runs`` for parameter descriptions.
+
+        Returns:
+            A tuple of ``runs`` and ``token`` where ``runs`` is a list of
+            :py:class:`mlflow.entities.Run` objects that satisfy the search expressions,
+            and ``token`` is the pagination token for the next page of results.
+        """
+
+    @abstractmethod
+    def log_batch(self, run_id, metrics, params, tags):
+        """
+        Log multiple metrics, params, and tags for the specified run
+
+        Args:
+            run_id: String id for the run
+            metrics: List of :py:class:`mlflow.entities.Metric` instances to log
+            params: List of :py:class:`mlflow.entities.Param` instances to log
+            tags: List of :py:class:`mlflow.entities.RunTag` instances to log
+
+        Returns:
+            None.
+        """
+
+    def log_batch_async(self, run_id, metrics, params, tags) -> RunOperations:
+        """
+        Log multiple metrics, params, and tags for the specified run in async fashion.
+        This API does not offer immediate consistency of the data. When API returns,
+        data is accepted but not persisted/processed by back end. Data would be processed
+        in near real time fashion.
+
+        Args:
+            run_id: String id for the run.
+            metrics: List of :py:class:`mlflow.entities.Metric` instances to log.
+            params: List of :py:class:`mlflow.entities.Param` instances to log.
+            tags: List of :py:class:`mlflow.entities.RunTag` instances to log.
+
+        Returns:
+            An :py:class:`mlflow.utils.async_logging.run_operations.RunOperations` instance
+            that represents future for logging operation.
+        """
+        if not self._async_logging_queue.is_active():
+            self._async_logging_queue.activate()
+
+        return self._async_logging_queue.log_batch_async(
+            run_id=run_id, metrics=metrics, params=params, tags=tags
+        )
+
+    def end_async_logging(self):
+        """
+        Ends the async logging queue. This method is a no-op if the queue is not active. This is
+        different from flush as it just stops the async logging queue from accepting
+        new data (moving the queue state TEAR_DOWN state), but flush will ensure all data
+        is processed before returning (moving the queue to IDLE state).
+        """
+        if self._async_logging_queue.is_active():
+            self._async_logging_queue.end_async_logging()
+
+    def flush_async_logging(self):
+        """
+        Flushes the async logging queue. This method is a no-op if the queue is already
+        at IDLE state. This methods also shutdown the logging worker threads.
+        After flushing, logging thread is setup again.
+        """
+        if not self._async_logging_queue.is_idle():
+            self._async_logging_queue.flush()
+
+    def shut_down_async_logging(self):
+        """
+        Shuts down the async logging queue. This method is a no-op if the queue is already
+        at IDLE state. This methods also shutdown the logging worker threads.
+        """
+        if not self._async_logging_queue.is_idle():
+            self._async_logging_queue.shut_down_async_logging()
+
+    @abstractmethod
+    def record_logged_model(self, run_id, mlflow_model):
+        """
+        Record logged model information with tracking store. The list of logged model infos is
+        maintained in a mlflow.models tag in JSON format.
+
+        Note: The actual models are logged as artifacts via artifact repository.
+
+        Args:
+            run_id: String id for the run.
+            mlflow_model: Model object to be recorded.
+
+        The default implementation is a no-op.
+
+        Returns:
+            None.
+        """
+
+    @abstractmethod
+    def log_inputs(self, run_id: str, datasets: Optional[list[DatasetInput]] = None):
+        """
+        Log inputs, such as datasets, to the specified run.
+
+        Args:
+            run_id: String id for the run
+            datasets: List of :py:class:`mlflow.entities.DatasetInput` instances to log
+                as inputs to the run.
+
+        Returns:
+            None.
+        """

venv/lib/python3.9/site-packages/mlflow/store/tracking/dbmodels/initial_models.py

@@ -0,0 +1,243 @@
+# Snapshot of MLflow DB models as of the 0.9.1 release, prior to the first database migration.
+# Used to standardize initial database state.
+# Copied with modifications from
+# https://github.com/mlflow/mlflow/blob/v0.9.1/mlflow/store/dbmodels/models.py, which
+# is the first database schema that users could be running. In particular, modifications have
+# been made to substitute constants from MLflow with hard-coded values (e.g. replacing
+# SourceType.to_string(SourceType.NOTEBOOK) with the constant "NOTEBOOK") and ensure
+# that all constraint names are unique. Note that pre-1.0 database schemas did not have unique
+# constraint names - we provided a one-time migration script for pre-1.0 users so that their
+# database schema matched the schema in this file.
+import time
+
+from sqlalchemy import (
+    BigInteger,
+    CheckConstraint,
+    Column,
+    Float,
+    ForeignKey,
+    Integer,
+    PrimaryKeyConstraint,
+    String,
+)
+from sqlalchemy.orm import backref, declarative_base, relationship
+
+Base = declarative_base()
+
+
+SourceTypes = [
+    "NOTEBOOK",
+    "JOB",
+    "LOCAL",
+    "UNKNOWN",
+    "PROJECT",
+]
+
+RunStatusTypes = [
+    "SCHEDULED",
+    "FAILED",
+    "FINISHED",
+    "RUNNING",
+]
+
+
+class SqlExperiment(Base):
+    """
+    DB model for :py:class:`mlflow.entities.Experiment`. These are recorded in ``experiment`` table.
+    """
+
+    __tablename__ = "experiments"
+
+    experiment_id = Column(Integer, autoincrement=True)
+    """
+    Experiment ID: `Integer`. *Primary Key* for ``experiment`` table.
+    """
+    name = Column(String(256), unique=True, nullable=False)
+    """
+    Experiment name: `String` (limit 256 characters). Defined as *Unique* and *Non null* in
+                     table schema.
+    """
+    artifact_location = Column(String(256), nullable=True)
+    """
+    Default artifact location for this experiment: `String` (limit 256 characters). Defined as
+                                                    *Non null* in table schema.
+    """
+    lifecycle_stage = Column(String(32), default="active")
+    """
+    Lifecycle Stage of experiment: `String` (limit 32 characters).
+                                    Can be either ``active`` (default) or ``deleted``.
+    """
+
+    __table_args__ = (
+        CheckConstraint(
+            lifecycle_stage.in_(["active", "deleted"]), name="experiments_lifecycle_stage"
+        ),
+        PrimaryKeyConstraint("experiment_id", name="experiment_pk"),
+    )
+
+    def __repr__(self):
+        return f"<SqlExperiment ({self.experiment_id}, {self.name})>"
+
+
+class SqlRun(Base):
+    """
+    DB model for :py:class:`mlflow.entities.Run`. These are recorded in ``runs`` table.
+    """
+
+    __tablename__ = "runs"
+
+    run_uuid = Column(String(32), nullable=False)
+    """
+    Run UUID: `String` (limit 32 characters). *Primary Key* for ``runs`` table.
+    """
+    name = Column(String(250))
+    """
+    Run name: `String` (limit 250 characters).
+    """
+    source_type = Column(String(20), default="LOCAL")
+    """
+    Source Type: `String` (limit 20 characters). Can be one of ``NOTEBOOK``, ``JOB``, ``PROJECT``,
+                 ``LOCAL`` (default), or ``UNKNOWN``.
+    """
+    source_name = Column(String(500))
+    """
+    Name of source recording the run: `String` (limit 500 characters).
+    """
+    entry_point_name = Column(String(50))
+    """
+    Entry-point name that launched the run run: `String` (limit 50 characters).
+    """
+    user_id = Column(String(256), nullable=True, default=None)
+    """
+    User ID: `String` (limit 256 characters). Defaults to ``null``.
+    """
+    status = Column(String(20), default="SCHEDULED")
+    """
+    Run Status: `String` (limit 20 characters). Can be one of ``RUNNING``, ``SCHEDULED`` (default),
+                ``FINISHED``, ``FAILED``.
+    """
+    start_time = Column(BigInteger, default=int(time.time()))
+    """
+    Run start time: `BigInteger`. Defaults to current system time.
+    """
+    end_time = Column(BigInteger, nullable=True, default=None)
+    """
+    Run end time: `BigInteger`.
+    """
+    source_version = Column(String(50))
+    """
+    Source version: `String` (limit 50 characters).
+    """
+    lifecycle_stage = Column(String(20), default="active")
+    """
+    Lifecycle Stage of run: `String` (limit 32 characters).
+                            Can be either ``active`` (default) or ``deleted``.
+    """
+    artifact_uri = Column(String(200), default=None)
+    """
+    Default artifact location for this run: `String` (limit 200 characters).
+    """
+    experiment_id = Column(Integer, ForeignKey("experiments.experiment_id"))
+    """
+    Experiment ID to which this run belongs to: *Foreign Key* into ``experiment`` table.
+    """
+    experiment = relationship("SqlExperiment", backref=backref("runs", cascade="all"))
+    """
+    SQLAlchemy relationship (many:one) with :py:class:`mlflow.store.dbmodels.models.SqlExperiment`.
+    """
+
+    __table_args__ = (
+        CheckConstraint(source_type.in_(SourceTypes), name="source_type"),
+        CheckConstraint(status.in_(RunStatusTypes), name="status"),
+        CheckConstraint(lifecycle_stage.in_(["active", "deleted"]), name="runs_lifecycle_stage"),
+        PrimaryKeyConstraint("run_uuid", name="run_pk"),
+    )
+
+
+class SqlTag(Base):
+    """
+    DB model for :py:class:`mlflow.entities.RunTag`. These are recorded in ``tags`` table.
+    """
+
+    __tablename__ = "tags"
+
+    key = Column(String(250))
+    """
+    Tag key: `String` (limit 250 characters). *Primary Key* for ``tags`` table.
+    """
+    value = Column(String(250), nullable=True)
+    """
+    Value associated with tag: `String` (limit 250 characters). Could be *null*.
+    """
+    run_uuid = Column(String(32), ForeignKey("runs.run_uuid"))
+    """
+    Run UUID to which this tag belongs to: *Foreign Key* into ``runs`` table.
+    """
+    run = relationship("SqlRun", backref=backref("tags", cascade="all"))
+    """
+    SQLAlchemy relationship (many:one) with :py:class:`mlflow.store.dbmodels.models.SqlRun`.
+    """
+
+    __table_args__ = (PrimaryKeyConstraint("key", "run_uuid", name="tag_pk"),)
+
+    def __repr__(self):
+        return f"<SqlRunTag({self.key}, {self.value})>"
+
+
+class SqlMetric(Base):
+    __tablename__ = "metrics"
+
+    key = Column(String(250))
+    """
+    Metric key: `String` (limit 250 characters). Part of *Primary Key* for ``metrics`` table.
+    """
+    value = Column(Float, nullable=False)
+    """
+    Metric value: `Float`. Defined as *Non-null* in schema.
+    """
+    timestamp = Column(BigInteger, default=lambda: int(time.time()))
+    """
+    Timestamp recorded for this metric entry: `BigInteger`. Part of *Primary Key* for
+                                               ``metrics`` table.
+    """
+    run_uuid = Column(String(32), ForeignKey("runs.run_uuid"))
+    """
+    Run UUID to which this metric belongs to: Part of *Primary Key* for ``metrics`` table.
+                                              *Foreign Key* into ``runs`` table.
+    """
+    run = relationship("SqlRun", backref=backref("metrics", cascade="all"))
+    """
+    SQLAlchemy relationship (many:one) with :py:class:`mlflow.store.dbmodels.models.SqlRun`.
+    """
+
+    __table_args__ = (PrimaryKeyConstraint("key", "timestamp", "run_uuid", name="metric_pk"),)
+
+    def __repr__(self):
+        return f"<SqlMetric({self.key}, {self.value}, {self.timestamp})>"
+
+
+class SqlParam(Base):
+    __tablename__ = "params"
+
+    key = Column(String(250))
+    """
+    Param key: `String` (limit 250 characters). Part of *Primary Key* for ``params`` table.
+    """
+    value = Column(String(250), nullable=False)
+    """
+    Param value: `String` (limit 250 characters). Defined as *Non-null* in schema.
+    """
+    run_uuid = Column(String(32), ForeignKey("runs.run_uuid"))
+    """
+    Run UUID to which this metric belongs to: Part of *Primary Key* for ``params`` table.
+                                              *Foreign Key* into ``runs`` table.
+    """
+    run = relationship("SqlRun", backref=backref("params", cascade="all"))
+    """
+    SQLAlchemy relationship (many:one) with :py:class:`mlflow.store.dbmodels.models.SqlRun`.
+    """
+
+    __table_args__ = (PrimaryKeyConstraint("key", "run_uuid", name="param_pk"),)
+
+    def __repr__(self):
+        return f"<SqlParam({self.key}, {self.value})>"

venv/lib/python3.9/site-packages/mlflow/store/tracking/dbmodels/models.py

@@ -0,0 +1,755 @@
+import sqlalchemy as sa
+from sqlalchemy import (
+    BigInteger,
+    Boolean,
+    CheckConstraint,
+    Column,
+    ForeignKey,
+    Index,
+    Integer,
+    PrimaryKeyConstraint,
+    String,
+    UnicodeText,
+)
+from sqlalchemy.orm import backref, relationship
+
+from mlflow.entities import (
+    Dataset,
+    Experiment,
+    ExperimentTag,
+    InputTag,
+    Metric,
+    Param,
+    Run,
+    RunData,
+    RunInfo,
+    RunStatus,
+    RunTag,
+    SourceType,
+    TraceInfo,
+    ViewType,
+)
+from mlflow.entities.lifecycle_stage import LifecycleStage
+from mlflow.entities.trace_info import TraceInfo
+from mlflow.entities.trace_status import TraceStatus
+from mlflow.store.db.base_sql_model import Base
+from mlflow.utils.mlflow_tags import _get_run_name_from_tags
+from mlflow.utils.time import get_current_time_millis
+
+SourceTypes = [
+    SourceType.to_string(SourceType.NOTEBOOK),
+    SourceType.to_string(SourceType.JOB),
+    SourceType.to_string(SourceType.LOCAL),
+    SourceType.to_string(SourceType.UNKNOWN),
+    SourceType.to_string(SourceType.PROJECT),
+]
+
+RunStatusTypes = [
+    RunStatus.to_string(RunStatus.SCHEDULED),
+    RunStatus.to_string(RunStatus.FAILED),
+    RunStatus.to_string(RunStatus.FINISHED),
+    RunStatus.to_string(RunStatus.RUNNING),
+    RunStatus.to_string(RunStatus.KILLED),
+]
+
+
+class SqlExperiment(Base):
+    """
+    DB model for :py:class:`mlflow.entities.Experiment`. These are recorded in ``experiment`` table.
+    """
+
+    __tablename__ = "experiments"
+
+    experiment_id = Column(Integer, autoincrement=True)
+    """
+    Experiment ID: `Integer`. *Primary Key* for ``experiment`` table.
+    """
+    name = Column(String(256), unique=True, nullable=False)
+    """
+    Experiment name: `String` (limit 256 characters). Defined as *Unique* and *Non null* in
+                     table schema.
+    """
+    artifact_location = Column(String(256), nullable=True)
+    """
+    Default artifact location for this experiment: `String` (limit 256 characters). Defined as
+                                                    *Non null* in table schema.
+    """
+    lifecycle_stage = Column(String(32), default=LifecycleStage.ACTIVE)
+    """
+    Lifecycle Stage of experiment: `String` (limit 32 characters).
+                                    Can be either ``active`` (default) or ``deleted``.
+    """
+    creation_time = Column(BigInteger(), default=get_current_time_millis)
+    """
+    Creation time of experiment: `BigInteger`.
+    """
+    last_update_time = Column(BigInteger(), default=get_current_time_millis)
+    """
+    Last Update time of experiment: `BigInteger`.
+    """
+
+    __table_args__ = (
+        CheckConstraint(
+            lifecycle_stage.in_(LifecycleStage.view_type_to_stages(ViewType.ALL)),
+            name="experiments_lifecycle_stage",
+        ),
+        PrimaryKeyConstraint("experiment_id", name="experiment_pk"),
+    )
+
+    def __repr__(self):
+        return f"<SqlExperiment ({self.experiment_id}, {self.name})>"
+
+    def to_mlflow_entity(self):
+        """
+        Convert DB model to corresponding MLflow entity.
+
+        Returns:
+            :py:class:`mlflow.entities.Experiment`.
+        """
+        return Experiment(
+            experiment_id=str(self.experiment_id),
+            name=self.name,
+            artifact_location=self.artifact_location,
+            lifecycle_stage=self.lifecycle_stage,
+            tags=[t.to_mlflow_entity() for t in self.tags],
+            creation_time=self.creation_time,
+            last_update_time=self.last_update_time,
+        )
+
+
+class SqlRun(Base):
+    """
+    DB model for :py:class:`mlflow.entities.Run`. These are recorded in ``runs`` table.
+    """
+
+    __tablename__ = "runs"
+
+    run_uuid = Column(String(32), nullable=False)
+    """
+    Run UUID: `String` (limit 32 characters). *Primary Key* for ``runs`` table.
+    """
+    name = Column(String(250))
+    """
+    Run name: `String` (limit 250 characters).
+    """
+    source_type = Column(String(20), default=SourceType.to_string(SourceType.LOCAL))
+    """
+    Source Type: `String` (limit 20 characters). Can be one of ``NOTEBOOK``, ``JOB``, ``PROJECT``,
+                 ``LOCAL`` (default), or ``UNKNOWN``.
+    """
+    source_name = Column(String(500))
+    """
+    Name of source recording the run: `String` (limit 500 characters).
+    """
+    entry_point_name = Column(String(50))
+    """
+    Entry-point name that launched the run run: `String` (limit 50 characters).
+    """
+    user_id = Column(String(256), nullable=True, default=None)
+    """
+    User ID: `String` (limit 256 characters). Defaults to ``null``.
+    """
+    status = Column(String(20), default=RunStatus.to_string(RunStatus.SCHEDULED))
+    """
+    Run Status: `String` (limit 20 characters). Can be one of ``RUNNING``, ``SCHEDULED`` (default),
+                ``FINISHED``, ``FAILED``.
+    """
+    start_time = Column(BigInteger, default=get_current_time_millis)
+    """
+    Run start time: `BigInteger`. Defaults to current system time.
+    """
+    end_time = Column(BigInteger, nullable=True, default=None)
+    """
+    Run end time: `BigInteger`.
+    """
+    deleted_time = Column(BigInteger, nullable=True, default=None)
+    """
+    Run deleted time: `BigInteger`. Timestamp of when run is deleted, defaults to none.
+    """
+    source_version = Column(String(50))
+    """
+    Source version: `String` (limit 50 characters).
+    """
+    lifecycle_stage = Column(String(20), default=LifecycleStage.ACTIVE)
+    """
+    Lifecycle Stage of run: `String` (limit 32 characters).
+                            Can be either ``active`` (default) or ``deleted``.
+    """
+    artifact_uri = Column(String(200), default=None)
+    """
+    Default artifact location for this run: `String` (limit 200 characters).
+    """
+    experiment_id = Column(Integer, ForeignKey("experiments.experiment_id"))
+    """
+    Experiment ID to which this run belongs to: *Foreign Key* into ``experiment`` table.
+    """
+    experiment = relationship("SqlExperiment", backref=backref("runs", cascade="all"))
+    """
+    SQLAlchemy relationship (many:one) with :py:class:`mlflow.store.dbmodels.models.SqlExperiment`.
+    """
+
+    __table_args__ = (
+        CheckConstraint(source_type.in_(SourceTypes), name="source_type"),
+        CheckConstraint(status.in_(RunStatusTypes), name="status"),
+        CheckConstraint(
+            lifecycle_stage.in_(LifecycleStage.view_type_to_stages(ViewType.ALL)),
+            name="runs_lifecycle_stage",
+        ),
+        PrimaryKeyConstraint("run_uuid", name="run_pk"),
+    )
+
+    @staticmethod
+    def get_attribute_name(mlflow_attribute_name):
+        """
+        Resolves an MLflow attribute name to a `SqlRun` attribute name.
+        """
+        # Currently, MLflow Search attributes defined in `SearchUtils.VALID_SEARCH_ATTRIBUTE_KEYS`
+        # share the same names as their corresponding `SqlRun` attributes. Therefore, this function
+        # returns the same attribute name
+        return {"run_name": "name", "run_id": "run_uuid"}.get(
+            mlflow_attribute_name, mlflow_attribute_name
+        )
+
+    def to_mlflow_entity(self):
+        """
+        Convert DB model to corresponding MLflow entity.
+
+        Returns:
+            mlflow.entities.Run: Description of the return value.
+        """
+        run_info = RunInfo(
+            run_uuid=self.run_uuid,
+            run_id=self.run_uuid,
+            run_name=self.name,
+            experiment_id=str(self.experiment_id),
+            user_id=self.user_id,
+            status=self.status,
+            start_time=self.start_time,
+            end_time=self.end_time,
+            lifecycle_stage=self.lifecycle_stage,
+            artifact_uri=self.artifact_uri,
+        )
+
+        tags = [t.to_mlflow_entity() for t in self.tags]
+        run_data = RunData(
+            metrics=[m.to_mlflow_entity() for m in self.latest_metrics],
+            params=[p.to_mlflow_entity() for p in self.params],
+            tags=tags,
+        )
+        if not run_info.run_name:
+            run_name = _get_run_name_from_tags(tags)
+            if run_name:
+                run_info._set_run_name(run_name)
+
+        return Run(run_info=run_info, run_data=run_data)
+
+
+class SqlExperimentTag(Base):
+    """
+    DB model for :py:class:`mlflow.entities.RunTag`.
+    These are recorded in ``experiment_tags`` table.
+    """
+
+    __tablename__ = "experiment_tags"
+
+    key = Column(String(250))
+    """
+    Tag key: `String` (limit 250 characters). *Primary Key* for ``tags`` table.
+    """
+    value = Column(String(5000), nullable=True)
+    """
+    Value associated with tag: `String` (limit 5000 characters). Could be *null*.
+    """
+    experiment_id = Column(Integer, ForeignKey("experiments.experiment_id"))
+    """
+    Experiment ID to which this tag belongs: *Foreign Key* into ``experiments`` table.
+    """
+    experiment = relationship("SqlExperiment", backref=backref("tags", cascade="all"))
+    """
+    SQLAlchemy relationship (many:one) with :py:class:`mlflow.store.dbmodels.models.SqlExperiment`.
+    """
+
+    __table_args__ = (PrimaryKeyConstraint("key", "experiment_id", name="experiment_tag_pk"),)
+
+    def __repr__(self):
+        return f"<SqlExperimentTag({self.key}, {self.value})>"
+
+    def to_mlflow_entity(self):
+        """
+        Convert DB model to corresponding MLflow entity.
+
+        Returns:
+            mlflow.entities.RunTag: Description of the return value.
+        """
+        return ExperimentTag(key=self.key, value=self.value)
+
+
+class SqlTag(Base):
+    """
+    DB model for :py:class:`mlflow.entities.RunTag`. These are recorded in ``tags`` table.
+    """
+
+    __tablename__ = "tags"
+    __table_args__ = (
+        PrimaryKeyConstraint("key", "run_uuid", name="tag_pk"),
+        Index(f"index_{__tablename__}_run_uuid", "run_uuid"),
+    )
+
+    key = Column(String(250))
+    """
+    Tag key: `String` (limit 250 characters). *Primary Key* for ``tags`` table.
+    """
+    value = Column(String(8000), nullable=True)
+    """
+    Value associated with tag: `String` (limit 8000 characters). Could be *null*.
+    """
+    run_uuid = Column(String(32), ForeignKey("runs.run_uuid"))
+    """
+    Run UUID to which this tag belongs to: *Foreign Key* into ``runs`` table.
+    """
+    run = relationship("SqlRun", backref=backref("tags", cascade="all"))
+    """
+    SQLAlchemy relationship (many:one) with :py:class:`mlflow.store.dbmodels.models.SqlRun`.
+    """
+
+    def __repr__(self):
+        return f"<SqlRunTag({self.key}, {self.value})>"
+
+    def to_mlflow_entity(self):
+        """
+        Convert DB model to corresponding MLflow entity.
+
+        Returns:
+            :py:class:`mlflow.entities.RunTag`.
+        """
+        return RunTag(key=self.key, value=self.value)
+
+
+class SqlMetric(Base):
+    __tablename__ = "metrics"
+    __table_args__ = (
+        PrimaryKeyConstraint(
+            "key", "timestamp", "step", "run_uuid", "value", "is_nan", name="metric_pk"
+        ),
+        Index(f"index_{__tablename__}_run_uuid", "run_uuid"),
+    )
+
+    key = Column(String(250))
+    """
+    Metric key: `String` (limit 250 characters). Part of *Primary Key* for ``metrics`` table.
+    """
+    value = Column(sa.types.Float(precision=53), nullable=False)
+    """
+    Metric value: `Float`. Defined as *Non-null* in schema.
+    """
+    timestamp = Column(BigInteger, default=get_current_time_millis)
+    """
+    Timestamp recorded for this metric entry: `BigInteger`. Part of *Primary Key* for
+                                               ``metrics`` table.
+    """
+    step = Column(BigInteger, default=0, nullable=False)
+    """
+    Step recorded for this metric entry: `BigInteger`.
+    """
+    is_nan = Column(Boolean(create_constraint=True), nullable=False, default=False)
+    """
+    True if the value is in fact NaN.
+    """
+    run_uuid = Column(String(32), ForeignKey("runs.run_uuid"))
+    """
+    Run UUID to which this metric belongs to: Part of *Primary Key* for ``metrics`` table.
+                                              *Foreign Key* into ``runs`` table.
+    """
+    run = relationship("SqlRun", backref=backref("metrics", cascade="all"))
+    """
+    SQLAlchemy relationship (many:one) with :py:class:`mlflow.store.dbmodels.models.SqlRun`.
+    """
+
+    def __repr__(self):
+        return f"<SqlMetric({self.key}, {self.value}, {self.timestamp}, {self.step})>"
+
+    def to_mlflow_entity(self):
+        """
+        Convert DB model to corresponding MLflow entity.
+
+        Returns:
+            mlflow.entities.Metric: Description of the return value.
+        """
+        return Metric(
+            key=self.key,
+            value=self.value if not self.is_nan else float("nan"),
+            timestamp=self.timestamp,
+            step=self.step,
+        )
+
+
+class SqlLatestMetric(Base):
+    __tablename__ = "latest_metrics"
+    __table_args__ = (
+        PrimaryKeyConstraint("key", "run_uuid", name="latest_metric_pk"),
+        Index(f"index_{__tablename__}_run_uuid", "run_uuid"),
+    )
+
+    key = Column(String(250))
+    """
+    Metric key: `String` (limit 250 characters). Part of *Primary Key* for ``latest_metrics`` table.
+    """
+    value = Column(sa.types.Float(precision=53), nullable=False)
+    """
+    Metric value: `Float`. Defined as *Non-null* in schema.
+    """
+    timestamp = Column(BigInteger, default=get_current_time_millis)
+    """
+    Timestamp recorded for this metric entry: `BigInteger`. Part of *Primary Key* for
+                                               ``latest_metrics`` table.
+    """
+    step = Column(BigInteger, default=0, nullable=False)
+    """
+    Step recorded for this metric entry: `BigInteger`.
+    """
+    is_nan = Column(Boolean(create_constraint=True), nullable=False, default=False)
+    """
+    True if the value is in fact NaN.
+    """
+    run_uuid = Column(String(32), ForeignKey("runs.run_uuid"))
+    """
+    Run UUID to which this metric belongs to: Part of *Primary Key* for ``latest_metrics`` table.
+                                              *Foreign Key* into ``runs`` table.
+    """
+    run = relationship("SqlRun", backref=backref("latest_metrics", cascade="all"))
+    """
+    SQLAlchemy relationship (many:one) with :py:class:`mlflow.store.dbmodels.models.SqlRun`.
+    """
+
+    def __repr__(self):
+        return f"<SqlLatestMetric({self.key}, {self.value}, {self.timestamp}, {self.step})>"
+
+    def to_mlflow_entity(self):
+        """
+        Convert DB model to corresponding MLflow entity.
+
+        Returns:
+            mlflow.entities.Metric: Description of the return value.
+        """
+        return Metric(
+            key=self.key,
+            value=self.value if not self.is_nan else float("nan"),
+            timestamp=self.timestamp,
+            step=self.step,
+        )
+
+
+class SqlParam(Base):
+    __tablename__ = "params"
+    __table_args__ = (
+        PrimaryKeyConstraint("key", "run_uuid", name="param_pk"),
+        Index(f"index_{__tablename__}_run_uuid", "run_uuid"),
+    )
+
+    key = Column(String(250))
+    """
+    Param key: `String` (limit 250 characters). Part of *Primary Key* for ``params`` table.
+    """
+    value = Column(String(8000), nullable=False)
+    """
+    Param value: `String` (limit 8000 characters). Defined as *Non-null* in schema.
+    """
+    run_uuid = Column(String(32), ForeignKey("runs.run_uuid"))
+    """
+    Run UUID to which this metric belongs to: Part of *Primary Key* for ``params`` table.
+                                              *Foreign Key* into ``runs`` table.
+    """
+    run = relationship("SqlRun", backref=backref("params", cascade="all"))
+    """
+    SQLAlchemy relationship (many:one) with :py:class:`mlflow.store.dbmodels.models.SqlRun`.
+    """
+
+    def __repr__(self):
+        return f"<SqlParam({self.key}, {self.value})>"
+
+    def to_mlflow_entity(self):
+        """
+        Convert DB model to corresponding MLflow entity.
+
+        Returns:
+            mlflow.entities.Param: Description of the return value.
+        """
+        return Param(key=self.key, value=self.value)
+
+
+class SqlDataset(Base):
+    __tablename__ = "datasets"
+    __table_args__ = (
+        PrimaryKeyConstraint("experiment_id", "name", "digest", name="dataset_pk"),
+        Index(f"index_{__tablename__}_dataset_uuid", "dataset_uuid"),
+        Index(
+            f"index_{__tablename__}_experiment_id_dataset_source_type",
+            "experiment_id",
+            "dataset_source_type",
+        ),
+    )
+
+    dataset_uuid = Column(String(36), nullable=False)
+    """
+    Dataset UUID: `String` (limit 36 characters). Defined as *Non-null* in schema.
+    Part of *Primary Key* for ``datasets`` table.
+    """
+    experiment_id = Column(Integer, ForeignKey("experiments.experiment_id", ondelete="CASCADE"))
+    """
+    Experiment ID to which this dataset belongs: *Foreign Key* into ``experiments`` table.
+    """
+    name = Column(String(500), nullable=False)
+    """
+    Param name: `String` (limit 500 characters). Defined as *Non-null* in schema.
+    Part of *Primary Key* for ``datasets`` table.
+    """
+    digest = Column(String(36), nullable=False)
+    """
+    Param digest: `String` (limit 500 characters). Defined as *Non-null* in schema.
+    Part of *Primary Key* for ``datasets`` table.
+    """
+    dataset_source_type = Column(String(36), nullable=False)
+    """
+    Param dataset_source_type: `String` (limit 36 characters). Defined as *Non-null* in schema.
+    """
+    dataset_source = Column(UnicodeText, nullable=False)
+    """
+    Param dataset_source: `UnicodeText`. Defined as *Non-null* in schema.
+    """
+    dataset_schema = Column(UnicodeText, nullable=True)
+    """
+    Param dataset_schema: `UnicodeText`.
+    """
+    dataset_profile = Column(UnicodeText, nullable=True)
+    """
+    Param dataset_profile: `UnicodeText`.
+    """
+
+    def __repr__(self):
+        return "<SqlDataset ({}, {}, {}, {}, {}, {}, {}, {})>".format(
+            self.dataset_uuid,
+            self.experiment_id,
+            self.name,
+            self.digest,
+            self.dataset_source_type,
+            self.dataset_source,
+            self.dataset_schema,
+            self.dataset_profile,
+        )
+
+    def to_mlflow_entity(self):
+        """
+        Convert DB model to corresponding MLflow entity.
+
+        Returns:
+            mlflow.entities.Dataset.
+        """
+        return Dataset(
+            name=self.name,
+            digest=self.digest,
+            source_type=self.dataset_source_type,
+            source=self.dataset_source,
+            schema=self.dataset_schema,
+            profile=self.dataset_profile,
+        )
+
+
+class SqlInput(Base):
+    __tablename__ = "inputs"
+    __table_args__ = (
+        PrimaryKeyConstraint(
+            "source_type", "source_id", "destination_type", "destination_id", name="inputs_pk"
+        ),
+        Index(f"index_{__tablename__}_input_uuid", "input_uuid"),
+        Index(
+            f"index_{__tablename__}_destination_type_destination_id_source_type",
+            "destination_type",
+            "destination_id",
+            "source_type",
+        ),
+    )
+
+    input_uuid = Column(String(36), nullable=False)
+    """
+    Input UUID: `String` (limit 36 characters). Defined as *Non-null* in schema.
+    """
+    source_type = Column(String(36), nullable=False)
+    """
+    Source type: `String` (limit 36 characters). Defined as *Non-null* in schema.
+    Part of *Primary Key* for ``inputs`` table.
+    """
+    source_id = Column(String(36), nullable=False)
+    """
+    Source Id: `String` (limit 36 characters). Defined as *Non-null* in schema.
+    Part of *Primary Key* for ``inputs`` table.
+    """
+    destination_type = Column(String(36), nullable=False)
+    """
+    Destination type: `String` (limit 36 characters). Defined as *Non-null* in schema.
+    Part of *Primary Key* for ``inputs`` table.
+    """
+    destination_id = Column(String(36), nullable=False)
+    """
+    Destination Id: `String` (limit 36 characters). Defined as *Non-null* in schema.
+    Part of *Primary Key* for ``inputs`` table.
+    """
+
+    def __repr__(self):
+        return "<SqlInput ({}, {}, {}, {}, {})>".format(
+            self.input_uuid,
+            self.source_type,
+            self.source_id,
+            self.destination_type,
+            self.destination_id,
+        )
+
+
+class SqlInputTag(Base):
+    __tablename__ = "input_tags"
+    __table_args__ = (PrimaryKeyConstraint("input_uuid", "name", name="input_tags_pk"),)
+
+    input_uuid = Column(String(36), ForeignKey("inputs.input_uuid"), nullable=False)
+    """
+    Input UUID: `String` (limit 36 characters). Defined as *Non-null* in schema.
+    *Foreign Key* into ``inputs`` table. Part of *Primary Key* for ``input_tags`` table.
+    """
+    name = Column(String(255), nullable=False)
+    """
+    Param name: `String` (limit 255 characters). Defined as *Non-null* in schema.
+    Part of *Primary Key* for ``input_tags`` table.
+    """
+    value = Column(String(500), nullable=False)
+    """
+    Param value: `String` (limit 500 characters). Defined as *Non-null* in schema.
+    Part of *Primary Key* for ``input_tags`` table.
+    """
+
+    def __repr__(self):
+        return f"<SqlInputTag ({self.input_uuid}, {self.name}, {self.value})>"
+
+    def to_mlflow_entity(self):
+        """
+        Convert DB model to corresponding MLflow entity.
+
+        Returns:
+            mlflow.entities.InputTag: Description of the return value.
+        """
+        return InputTag(key=self.name, value=self.value)
+
+
+#######################################################################################
+# Below are Tracing models. We may refactor them to be in a separate module in the future.
+#######################################################################################
+
+
+class SqlTraceInfo(Base):
+    __tablename__ = "trace_info"
+
+    request_id = Column(String(50), nullable=False)
+    """
+    Request ID: `String` (limit 50 characters). *Primary Key* for ``trace_info`` table.
+    """
+    experiment_id = Column(Integer, ForeignKey("experiments.experiment_id"), nullable=False)
+    """
+    Experiment ID to which this trace belongs: *Foreign Key* into ``experiments`` table.
+    """
+    timestamp_ms = Column(BigInteger, nullable=False)
+    """
+    Start time of the trace, in milliseconds.
+    """
+    execution_time_ms = Column(BigInteger, nullable=True)
+    """
+    Duration of the trace, in milliseconds. Could be *null* if the trace is still in progress
+    or not ended correctly for some reason.
+    """
+    status = Column(String(50), nullable=False)
+    """
+    Status of the trace. The values are defined in
+    :py:class:`mlflow.entities.trace_status.TraceStatus` enum but we don't enforce
+    constraint at DB level.
+    """
+
+    __table_args__ = (
+        PrimaryKeyConstraint("request_id", name="trace_info_pk"),
+        # The most frequent query will be get all traces in an experiment sorted by timestamp desc,
+        # which is the default view in the UI. Also every search query should have experiment_id(s)
+        # in the where clause.
+        Index(f"index_{__tablename__}_experiment_id_timestamp_ms", "experiment_id", "timestamp_ms"),
+    )
+
+    def to_mlflow_entity(self):
+        """
+        Convert DB model to corresponding MLflow entity.
+
+        Returns:
+            :py:class:`mlflow.entities.TraceInfo` object.
+        """
+        return TraceInfo(
+            request_id=self.request_id,
+            experiment_id=str(self.experiment_id),
+            timestamp_ms=self.timestamp_ms,
+            execution_time_ms=self.execution_time_ms,
+            status=TraceStatus(self.status),
+            tags={t.key: t.value for t in self.tags},
+            request_metadata={m.key: m.value for m in self.request_metadata},
+        )
+
+
+class SqlTraceTag(Base):
+    __tablename__ = "trace_tags"
+
+    key = Column(String(250))
+    """
+    Tag key: `String` (limit 250 characters).
+    """
+    value = Column(String(8000), nullable=True)
+    """
+    Value associated with tag: `String` (limit 250 characters). Could be *null*.
+    """
+    request_id = Column(
+        String(50), ForeignKey("trace_info.request_id", ondelete="CASCADE"), nullable=False
+    )
+    """
+    Request ID to which this tag belongs: *Foreign Key* into ``trace_info`` table.
+    """
+    trace_info = relationship("SqlTraceInfo", backref=backref("tags", cascade="all"))
+    """
+    SQLAlchemy relationship (many:one) with
+    :py:class:`mlflow.store.dbmodels.models.SqlTraceInfo`.
+    """
+
+    # Key is unique within a request_id
+    __table_args__ = (
+        PrimaryKeyConstraint("request_id", "key", name="trace_tag_pk"),
+        Index(f"index_{__tablename__}_request_id"),
+    )
+
+
+class SqlTraceRequestMetadata(Base):
+    __tablename__ = "trace_request_metadata"
+
+    key = Column(String(250))
+    """
+    Metadata key: `String` (limit 250 characters).
+    """
+    value = Column(String(8000), nullable=True)
+    """
+    Value associated with metadata: `String` (limit 250 characters). Could be *null*.
+    """
+    request_id = Column(
+        String(50), ForeignKey("trace_info.request_id", ondelete="CASCADE"), nullable=False
+    )
+    """
+    Request ID to which this metadata belongs: *Foreign Key* into ``trace_info`` table.
+    """
+    trace_info = relationship("SqlTraceInfo", backref=backref("request_metadata", cascade="all"))
+    """
+    SQLAlchemy relationship (many:one) with
+    :py:class:`mlflow.store.dbmodels.models.SqlTraceInfo`.
+    """
+
+    # Key is unique within a request_id
+    __table_args__ = (
+        PrimaryKeyConstraint("request_id", "key", name="trace_request_metadata_pk"),
+        Index(f"index_{__tablename__}_request_id"),
+    )

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

@@ -0,0 +1,680 @@
+import json
+import logging
+from typing import Optional
+
+from mlflow.entities import DatasetInput, Experiment, Metric, Run, RunInfo, TraceInfo, ViewType
+from mlflow.entities.assessment import Assessment, Expectation, Feedback
+from mlflow.entities.trace_status import TraceStatus
+from mlflow.exceptions import MlflowException
+from mlflow.protos import databricks_pb2
+from mlflow.protos.service_pb2 import (
+    CreateAssessment,
+    CreateExperiment,
+    CreateRun,
+    DeleteAssessment,
+    DeleteExperiment,
+    DeleteRun,
+    DeleteTag,
+    DeleteTraces,
+    DeleteTraceTag,
+    EndTrace,
+    GetExperiment,
+    GetExperimentByName,
+    GetMetricHistory,
+    GetRun,
+    GetTraceInfo,
+    GetTraceInfoV3,
+    LogBatch,
+    LogInputs,
+    LogMetric,
+    LogModel,
+    LogParam,
+    MlflowService,
+    RestoreExperiment,
+    RestoreRun,
+    SearchExperiments,
+    SearchRuns,
+    SearchTraces,
+    SetExperimentTag,
+    SetTag,
+    SetTraceTag,
+    StartTrace,
+    TraceRequestMetadata,
+    TraceTag,
+    UpdateAssessment,
+    UpdateExperiment,
+    UpdateRun,
+)
+from mlflow.store.entities.paged_list import PagedList
+from mlflow.store.tracking import SEARCH_TRACES_DEFAULT_MAX_RESULTS
+from mlflow.store.tracking.abstract_store import AbstractStore
+from mlflow.utils.proto_json_utils import message_to_json, set_pb_value
+from mlflow.utils.rest_utils import (
+    _REST_API_PATH_PREFIX,
+    call_endpoint,
+    extract_api_info_for_service,
+    get_create_assessment_endpoint,
+    get_set_trace_tag_endpoint,
+    get_single_assessment_endpoint,
+    get_single_trace_endpoint,
+    get_trace_assessment_endpoint,
+    get_trace_info_endpoint,
+)
+
+_METHOD_TO_INFO = extract_api_info_for_service(MlflowService, _REST_API_PATH_PREFIX)
+_logger = logging.getLogger(__name__)
+
+
+class RestStore(AbstractStore):
+    """
+    Client for a remote tracking server accessed via REST API calls
+
+    Args
+        get_host_creds: Method to be invoked prior to every REST request to get the
+            :py:class:`mlflow.rest_utils.MlflowHostCreds` for the request. Note that this
+            is a function so that we can obtain fresh credentials in the case of expiry.
+    """
+
+    def __init__(self, get_host_creds):
+        super().__init__()
+        self.get_host_creds = get_host_creds
+
+    def _call_endpoint(self, api, json_body, endpoint=None):
+        if endpoint:
+            # Allow customizing the endpoint for compatibility with dynamic endpoints, such as
+            # /mlflow/traces/{request_id}/info.
+            _, method = _METHOD_TO_INFO[api]
+        else:
+            endpoint, method = _METHOD_TO_INFO[api]
+        response_proto = api.Response()
+        return call_endpoint(self.get_host_creds(), endpoint, method, json_body, response_proto)
+
+    def search_experiments(
+        self,
+        view_type=ViewType.ACTIVE_ONLY,
+        max_results=None,
+        filter_string=None,
+        order_by=None,
+        page_token=None,
+    ):
+        req_body = message_to_json(
+            SearchExperiments(
+                view_type=view_type,
+                max_results=max_results,
+                page_token=page_token,
+                order_by=order_by,
+                filter=filter_string,
+            )
+        )
+        response_proto = self._call_endpoint(SearchExperiments, req_body)
+        experiments = [Experiment.from_proto(x) for x in response_proto.experiments]
+        token = (
+            response_proto.next_page_token if response_proto.HasField("next_page_token") else None
+        )
+        return PagedList(experiments, token)
+
+    def create_experiment(self, name, artifact_location=None, tags=None):
+        """
+        Create a new experiment.
+        If an experiment with the given name already exists, throws exception.
+
+        Args:
+            name: Desired name for an experiment.
+            artifact_location: Location to store run artifacts.
+            tags: A list of :py:class:`mlflow.entities.ExperimentTag` instances to set for the
+                experiment.
+
+        Returns:
+            experiment_id for the newly created experiment if successful, else None
+        """
+        tag_protos = [tag.to_proto() for tag in tags] if tags else []
+        req_body = message_to_json(
+            CreateExperiment(name=name, artifact_location=artifact_location, tags=tag_protos)
+        )
+        response_proto = self._call_endpoint(CreateExperiment, req_body)
+        return response_proto.experiment_id
+
+    def get_experiment(self, experiment_id):
+        """
+        Fetch the experiment from the backend store.
+
+        Args:
+            experiment_id: String id for the experiment
+
+        Returns:
+            A single :py:class:`mlflow.entities.Experiment` object if it exists,
+            otherwise raises an Exception.
+        """
+        req_body = message_to_json(GetExperiment(experiment_id=str(experiment_id)))
+        response_proto = self._call_endpoint(GetExperiment, req_body)
+        return Experiment.from_proto(response_proto.experiment)
+
+    def delete_experiment(self, experiment_id):
+        req_body = message_to_json(DeleteExperiment(experiment_id=str(experiment_id)))
+        self._call_endpoint(DeleteExperiment, req_body)
+
+    def restore_experiment(self, experiment_id):
+        req_body = message_to_json(RestoreExperiment(experiment_id=str(experiment_id)))
+        self._call_endpoint(RestoreExperiment, req_body)
+
+    def rename_experiment(self, experiment_id, new_name):
+        req_body = message_to_json(
+            UpdateExperiment(experiment_id=str(experiment_id), new_name=new_name)
+        )
+        self._call_endpoint(UpdateExperiment, req_body)
+
+    def get_run(self, run_id):
+        """
+        Fetch the run from backend store
+
+        Args:
+            run_id: Unique identifier for the run
+
+        Returns:
+            A single Run object if it exists, otherwise raises an Exception
+        """
+        req_body = message_to_json(GetRun(run_uuid=run_id, run_id=run_id))
+        response_proto = self._call_endpoint(GetRun, req_body)
+        return Run.from_proto(response_proto.run)
+
+    def update_run_info(self, run_id, run_status, end_time, run_name):
+        """Updates the metadata of the specified run."""
+        req_body = message_to_json(
+            UpdateRun(
+                run_uuid=run_id,
+                run_id=run_id,
+                status=run_status,
+                end_time=end_time,
+                run_name=run_name,
+            )
+        )
+        response_proto = self._call_endpoint(UpdateRun, req_body)
+        return RunInfo.from_proto(response_proto.run_info)
+
+    def create_run(self, experiment_id, user_id, start_time, tags, run_name):
+        """
+        Create a run under the specified experiment ID, setting the run's status to "RUNNING"
+        and the start time to the current time.
+
+        Args:
+            experiment_id: ID of the experiment for this run.
+            user_id: ID of the user launching this run.
+            start_time: timestamp of the initialization of the run.
+            tags: tags to apply to this run at initialization.
+            run_name: Name of this run.
+
+        Returns:
+            The created Run object.
+        """
+
+        tag_protos = [tag.to_proto() for tag in tags]
+        req_body = message_to_json(
+            CreateRun(
+                experiment_id=str(experiment_id),
+                user_id=user_id,
+                start_time=start_time,
+                tags=tag_protos,
+                run_name=run_name,
+            )
+        )
+        response_proto = self._call_endpoint(CreateRun, req_body)
+        return Run.from_proto(response_proto.run)
+
+    def start_trace(
+        self,
+        experiment_id: str,
+        timestamp_ms: int,
+        request_metadata: dict[str, str],
+        tags: dict[str, str],
+    ) -> TraceInfo:
+        """
+        Start an initial TraceInfo object in the backend store.
+
+        Args:
+            experiment_id: String id of the experiment for this run.
+            timestamp_ms: Start time of the trace, in milliseconds since the UNIX epoch.
+            request_metadata: Metadata of the trace.
+            tags: Tags of the trace.
+
+        Returns:
+            The created TraceInfo object.
+        """
+        request_metadata_proto = []
+        for key, value in request_metadata.items():
+            attr = TraceRequestMetadata()
+            attr.key = key
+            attr.value = str(value)
+            request_metadata_proto.append(attr)
+
+        tags_proto = []
+        for key, value in tags.items():
+            tag = TraceTag()
+            tag.key = key
+            tag.value = str(value)
+            tags_proto.append(tag)
+
+        req_body = message_to_json(
+            StartTrace(
+                experiment_id=str(experiment_id),
+                timestamp_ms=timestamp_ms,
+                request_metadata=request_metadata_proto,
+                tags=tags_proto,
+            )
+        )
+        response_proto = self._call_endpoint(StartTrace, req_body)
+        return TraceInfo.from_proto(response_proto.trace_info)
+
+    def end_trace(
+        self,
+        request_id: str,
+        timestamp_ms: int,
+        status: TraceStatus,
+        request_metadata: dict[str, str],
+        tags: dict[str, str],
+    ) -> TraceInfo:
+        """
+        Update the TraceInfo object in the backend store with the completed trace info.
+
+        Args:
+            request_id: Unique string identifier of the trace.
+            timestamp_ms: End time of the trace, in milliseconds. The execution time field
+                in the TraceInfo will be calculated by subtracting the start time from this.
+            status: Status of the trace.
+            request_metadata: Metadata of the trace. This will be merged with the existing
+                metadata logged during the start_trace call.
+            tags: Tags of the trace. This will be merged with the existing tags logged
+                during the start_trace or set_trace_tag calls.
+
+        Returns:
+            The updated TraceInfo object.
+        """
+        request_metadata_proto = []
+        for key, value in request_metadata.items():
+            attr = TraceRequestMetadata()
+            attr.key = key
+            attr.value = str(value)
+            request_metadata_proto.append(attr)
+
+        tags_proto = []
+        for key, value in tags.items():
+            tag = TraceTag()
+            tag.key = key
+            tag.value = str(value)
+            tags_proto.append(tag)
+
+        req_body = message_to_json(
+            EndTrace(
+                request_id=request_id,
+                timestamp_ms=timestamp_ms,
+                status=status.to_proto(),
+                request_metadata=request_metadata_proto,
+                tags=tags_proto,
+            )
+        )
+        # EndTrace endpoint is a dynamic path built with the request_id
+        endpoint = get_single_trace_endpoint(request_id)
+        response_proto = self._call_endpoint(EndTrace, req_body, endpoint=endpoint)
+        return TraceInfo.from_proto(response_proto.trace_info)
+
+    def _delete_traces(
+        self,
+        experiment_id: str,
+        max_timestamp_millis: Optional[int] = None,
+        max_traces: Optional[int] = None,
+        request_ids: Optional[list[str]] = None,
+    ) -> int:
+        req_body = message_to_json(
+            DeleteTraces(
+                experiment_id=experiment_id,
+                max_timestamp_millis=max_timestamp_millis,
+                max_traces=max_traces,
+                request_ids=request_ids,
+            )
+        )
+        res = self._call_endpoint(DeleteTraces, req_body)
+        return res.traces_deleted
+
+    def get_trace_info(self, request_id, should_query_v3: bool = False):
+        """
+        Get the trace matching the `request_id`.
+
+        Args:
+            request_id: String id of the trace to fetch.
+            should_query_v3: If True, the backend store will query the V3 API for the trace info.
+                TODO: Remove this flag once the V3 API is the default in OSS.
+
+        Returns:
+            The fetched Trace object, of type ``mlflow.entities.TraceInfo``.
+        """
+        req_body = message_to_json(GetTraceInfo(request_id=request_id))
+        endpoint = get_trace_info_endpoint(request_id)
+        response_proto = self._call_endpoint(GetTraceInfo, req_body, endpoint=endpoint)
+        assessments = None
+        if should_query_v3:
+            try:
+                tracev3_req_body = message_to_json(GetTraceInfoV3(trace_id=request_id))
+                tracev3_endpoint = get_trace_assessment_endpoint(request_id)
+                tracev3_response_proto = self._call_endpoint(
+                    GetTraceInfoV3, tracev3_req_body, endpoint=tracev3_endpoint
+                )
+                assessments = [
+                    Assessment.from_proto(a)
+                    for a in tracev3_response_proto.trace.trace_info.assessments
+                ]
+            except Exception:
+                # TraceV3 endpoint is not globally enabled yet; graceful fallback path.
+                pass
+        return TraceInfo.from_proto(response_proto.trace_info, assessments=assessments)
+
+    def search_traces(
+        self,
+        experiment_ids: list[str],
+        filter_string: Optional[str] = None,
+        max_results: int = SEARCH_TRACES_DEFAULT_MAX_RESULTS,
+        order_by: Optional[list[str]] = None,
+        page_token: Optional[str] = None,
+    ):
+        st = SearchTraces(
+            experiment_ids=experiment_ids,
+            filter=filter_string,
+            max_results=max_results,
+            order_by=order_by,
+            page_token=page_token,
+        )
+        req_body = message_to_json(st)
+        response_proto = self._call_endpoint(SearchTraces, req_body)
+        trace_infos = [TraceInfo.from_proto(t) for t in response_proto.traces]
+        return trace_infos, response_proto.next_page_token or None
+
+    def set_trace_tag(self, request_id: str, key: str, value: str):
+        """
+        Set a tag on the trace with the given request_id.
+
+        Args:
+            request_id: The ID of the trace.
+            key: The string key of the tag.
+            value: The string value of the tag.
+        """
+        req_body = message_to_json(SetTraceTag(key=key, value=value))
+        self._call_endpoint(SetTraceTag, req_body, endpoint=get_set_trace_tag_endpoint(request_id))
+
+    def delete_trace_tag(self, request_id: str, key: str):
+        """
+        Delete a tag on the trace with the given request_id.
+
+        Args:
+            request_id: The ID of the trace.
+            key: The string key of the tag.
+        """
+        req_body = message_to_json(DeleteTraceTag(key=key))
+        self._call_endpoint(
+            DeleteTraceTag, req_body, endpoint=get_set_trace_tag_endpoint(request_id)
+        )
+
+    def create_assessment(self, assessment: Assessment) -> Assessment:
+        """
+        Create an assessment entity in the backend store.
+
+        Args:
+            assessment: The assessment to log (without an assessment_id).
+
+        Returns:
+            The created Assessment object.
+        """
+        req_body = message_to_json(CreateAssessment(assessment=assessment.to_proto()))
+        response_proto = self._call_endpoint(
+            CreateAssessment,
+            req_body,
+            endpoint=get_create_assessment_endpoint(assessment.trace_id),
+        )
+        return Assessment.from_proto(response_proto.assessment)
+
+    def update_assessment(
+        self,
+        trace_id: str,
+        assessment_id: str,
+        name: Optional[str] = None,
+        expectation: Optional[Expectation] = None,
+        feedback: Optional[Feedback] = None,
+        rationale: Optional[str] = None,
+        metadata: Optional[dict[str, str]] = None,
+    ) -> Assessment:
+        """
+        Update an existing assessment entity in the backend store.
+
+        Args:
+            trace_id: The ID of the trace.
+            assessment_id: The ID of the assessment to update.
+            name: The updated name of the assessment.
+            expectation: The updated expectation value of the assessment.
+            feedback: The updated feedback value of the assessment.
+            rationale: The updated rationale of the feedback. Not applicable for expectations.
+            metadata: Additional metadata for the assessment.
+        """
+        if expectation is not None and feedback is not None:
+            raise MlflowException.invalid_parameter_value(
+                "Exactly one of `expectation` or `feedback` should be specified."
+            )
+
+        update = UpdateAssessment()
+
+        # The assessment object to be sent to the backend (only contains fields to update and IDs)
+        assessment = update.assessment
+        # Field mask specifies which fields to update.
+        mask = update.update_mask
+
+        assessment.assessment_id = assessment_id
+        assessment.trace_id = trace_id
+
+        if name is not None:
+            assessment.assessment_name = name
+            mask.paths.append("assessment_name")
+        if expectation is not None:
+            set_pb_value(assessment.expectation.value, expectation.value)
+            mask.paths.append("expectation")
+        if feedback is not None:
+            assessment.feedback.CopyFrom(feedback.to_proto())
+            mask.paths.append("feedback")
+        if rationale is not None:
+            assessment.rationale = rationale
+            mask.paths.append("rationale")
+        if metadata is not None:
+            assessment.metadata.update(metadata)
+            mask.paths.append("metadata")
+
+        req_body = message_to_json(update)
+        response_proto = self._call_endpoint(
+            UpdateAssessment,
+            req_body,
+            endpoint=get_single_assessment_endpoint(trace_id, assessment_id),
+        )
+        return Assessment.from_proto(response_proto.assessment)
+
+    def delete_assessment(self, trace_id: str, assessment_id: str):
+        """
+        Delete an assessment associated with a trace.
+
+        Args:
+            trace_id: String ID of the trace.
+            assessment_id: String ID of the assessment to delete.
+        """
+        req_body = message_to_json(DeleteAssessment(trace_id=trace_id, assessment_id=assessment_id))
+        self._call_endpoint(
+            DeleteAssessment,
+            req_body,
+            endpoint=get_single_assessment_endpoint(trace_id, assessment_id),
+        )
+
+    def log_metric(self, run_id: str, metric: Metric):
+        """
+        Log a metric for the specified run
+
+        Args:
+            run_id: String id for the run
+            metric: Metric instance to log
+        """
+        req_body = message_to_json(
+            LogMetric(
+                run_uuid=run_id,
+                run_id=run_id,
+                key=metric.key,
+                value=metric.value,
+                timestamp=metric.timestamp,
+                step=metric.step,
+            )
+        )
+        self._call_endpoint(LogMetric, req_body)
+
+    def log_param(self, run_id, param):
+        """
+        Log a param for the specified run
+
+        Args:
+            run_id: String id for the run
+            param: Param instance to log
+        """
+        req_body = message_to_json(
+            LogParam(run_uuid=run_id, run_id=run_id, key=param.key, value=param.value)
+        )
+        self._call_endpoint(LogParam, req_body)
+
+    def set_experiment_tag(self, experiment_id, tag):
+        """
+        Set a tag for the specified experiment
+
+        Args:
+            experiment_id: String ID of the experiment
+            tag: ExperimentRunTag instance to log
+        """
+        req_body = message_to_json(
+            SetExperimentTag(experiment_id=experiment_id, key=tag.key, value=tag.value)
+        )
+        self._call_endpoint(SetExperimentTag, req_body)
+
+    def set_tag(self, run_id, tag):
+        """
+        Set a tag for the specified run
+
+        Args:
+            run_id: String ID of the run
+            tag: RunTag instance to log
+        """
+        req_body = message_to_json(
+            SetTag(run_uuid=run_id, run_id=run_id, key=tag.key, value=tag.value)
+        )
+        self._call_endpoint(SetTag, req_body)
+
+    def delete_tag(self, run_id, key):
+        """
+        Delete a tag from a run. This is irreversible.
+
+        Args:
+            run_id: String ID of the run.
+            key: Name of the tag.
+        """
+        req_body = message_to_json(DeleteTag(run_id=run_id, key=key))
+        self._call_endpoint(DeleteTag, req_body)
+
+    def get_metric_history(self, run_id, metric_key, max_results=None, page_token=None):
+        """
+        Return all logged values for a given metric.
+
+        Args:
+            run_id: Unique identifier for run.
+            metric_key: Metric name within the run.
+            max_results: Maximum number of metric history events (steps) to return per paged
+                query. Only supported in 'databricks' backend.
+            page_token: A Token specifying the next paginated set of results of metric history.
+
+        Returns:
+            A PagedList of :py:class:`mlflow.entities.Metric` entities if a paginated request
+            is made by setting ``max_results`` to a value other than ``None``, a List of
+            :py:class:`mlflow.entities.Metric` entities if ``max_results`` is None, else, if no
+            metrics of the ``metric_key`` have been logged to the ``run_id``, an empty list.
+        """
+        req_body = message_to_json(
+            GetMetricHistory(
+                run_uuid=run_id,
+                run_id=run_id,
+                metric_key=metric_key,
+                max_results=max_results,
+                page_token=page_token,
+            )
+        )
+        response_proto = self._call_endpoint(GetMetricHistory, req_body)
+
+        metric_history = [Metric.from_proto(metric) for metric in response_proto.metrics]
+        return PagedList(metric_history, response_proto.next_page_token or None)
+
+    def _search_runs(
+        self, experiment_ids, filter_string, run_view_type, max_results, order_by, page_token
+    ):
+        experiment_ids = [str(experiment_id) for experiment_id in experiment_ids]
+        sr = SearchRuns(
+            experiment_ids=experiment_ids,
+            filter=filter_string,
+            run_view_type=ViewType.to_proto(run_view_type),
+            max_results=max_results,
+            order_by=order_by,
+            page_token=page_token,
+        )
+        req_body = message_to_json(sr)
+        response_proto = self._call_endpoint(SearchRuns, req_body)
+        runs = [Run.from_proto(proto_run) for proto_run in response_proto.runs]
+        # If next_page_token is not set, we will see it as "". We need to convert this to None.
+        next_page_token = None
+        if response_proto.next_page_token:
+            next_page_token = response_proto.next_page_token
+        return runs, next_page_token
+
+    def delete_run(self, run_id):
+        req_body = message_to_json(DeleteRun(run_id=run_id))
+        self._call_endpoint(DeleteRun, req_body)
+
+    def restore_run(self, run_id):
+        req_body = message_to_json(RestoreRun(run_id=run_id))
+        self._call_endpoint(RestoreRun, req_body)
+
+    def get_experiment_by_name(self, experiment_name):
+        try:
+            req_body = message_to_json(GetExperimentByName(experiment_name=experiment_name))
+            response_proto = self._call_endpoint(GetExperimentByName, req_body)
+            return Experiment.from_proto(response_proto.experiment)
+        except MlflowException as e:
+            if e.error_code == databricks_pb2.ErrorCode.Name(
+                databricks_pb2.RESOURCE_DOES_NOT_EXIST
+            ):
+                return None
+            else:
+                raise
+
+    def log_batch(self, run_id, metrics, params, tags):
+        metric_protos = [metric.to_proto() for metric in metrics]
+        param_protos = [param.to_proto() for param in params]
+        tag_protos = [tag.to_proto() for tag in tags]
+        req_body = message_to_json(
+            LogBatch(metrics=metric_protos, params=param_protos, tags=tag_protos, run_id=run_id)
+        )
+        self._call_endpoint(LogBatch, req_body)
+
+    def record_logged_model(self, run_id, mlflow_model):
+        req_body = message_to_json(
+            LogModel(run_id=run_id, model_json=json.dumps(mlflow_model.get_tags_dict()))
+        )
+        self._call_endpoint(LogModel, req_body)
+
+    def log_inputs(self, run_id: str, datasets: Optional[list[DatasetInput]] = None):
+        """
+        Log inputs, such as datasets, to the specified run.
+
+        Args:
+            run_id: String id for the run
+            datasets: List of :py:class:`mlflow.entities.DatasetInput` instances to log
+                as inputs to the run.
+
+        Returns:
+            None.
+        """
+        datasets_protos = [dataset.to_proto() for dataset in datasets]
+        req_body = message_to_json(LogInputs(run_id=run_id, datasets=datasets_protos))
+        self._call_endpoint(LogInputs, req_body)