genesis-flow 1.0.0__py3-none-any.whl

mlflow/store/tracking/rest_store.py

@@ -0,0 +1,1131 @@
+import json
+import logging
+from typing import Any, Optional
+
+from mlflow.entities import (
+    DatasetInput,
+    Experiment,
+    LoggedModel,
+    LoggedModelInput,
+    LoggedModelOutput,
+    LoggedModelParameter,
+    LoggedModelStatus,
+    LoggedModelTag,
+    Metric,
+    Run,
+    RunInfo,
+    ViewType,
+)
+from mlflow.entities.assessment import Assessment, Expectation, Feedback
+from mlflow.entities.trace import Trace
+from mlflow.entities.trace_data import TraceData
+from mlflow.entities.trace_info import TraceInfo
+from mlflow.entities.trace_info_v2 import TraceInfoV2
+from mlflow.entities.trace_location import TraceLocation
+from mlflow.entities.trace_status import TraceStatus
+from mlflow.environment_variables import (
+    _MLFLOW_CREATE_LOGGED_MODEL_PARAMS_BATCH_SIZE,
+    _MLFLOW_LOG_LOGGED_MODEL_PARAMS_BATCH_SIZE,
+    MLFLOW_ASYNC_TRACE_LOGGING_RETRY_TIMEOUT,
+)
+from mlflow.exceptions import MlflowException
+from mlflow.protos import databricks_pb2
+from mlflow.protos.service_pb2 import (
+    CreateAssessment,
+    CreateExperiment,
+    CreateLoggedModel,
+    CreateRun,
+    DeleteAssessment,
+    DeleteExperiment,
+    DeleteLoggedModel,
+    DeleteLoggedModelTag,
+    DeleteRun,
+    DeleteTag,
+    DeleteTraces,
+    DeleteTraceTag,
+    EndTrace,
+    FinalizeLoggedModel,
+    GetAssessmentRequest,
+    GetExperiment,
+    GetExperimentByName,
+    GetLoggedModel,
+    GetMetricHistory,
+    GetOnlineTraceDetails,
+    GetRun,
+    GetTraceInfo,
+    GetTraceInfoV3,
+    LogBatch,
+    LogInputs,
+    LogLoggedModelParamsRequest,
+    LogMetric,
+    LogModel,
+    LogOutputs,
+    LogParam,
+    MlflowService,
+    RestoreExperiment,
+    RestoreRun,
+    SearchExperiments,
+    SearchLoggedModels,
+    SearchRuns,
+    SearchTraces,
+    SearchTracesV3,
+    SearchUnifiedTraces,
+    SetExperimentTag,
+    SetLoggedModelTags,
+    SetTag,
+    SetTraceTag,
+    StartTrace,
+    StartTraceV3,
+    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
+from mlflow.utils.rest_utils import (
+    _REST_API_PATH_PREFIX,
+    _V3_TRACE_REST_API_PATH_PREFIX,
+    call_endpoint,
+    extract_api_info_for_service,
+    get_logged_model_endpoint,
+    get_single_assessment_endpoint,
+    get_single_trace_endpoint,
+    get_trace_tag_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=None,
+        endpoint=None,
+        retry_timeout_seconds=None,
+    ):
+        if endpoint:
+            # Allow customizing the endpoint for compatibility with dynamic endpoints, such as
+            # /mlflow/traces/{trace_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,
+            retry_timeout_seconds=retry_timeout_seconds,
+        )
+
+    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, trace_info: TraceInfo) -> TraceInfo:
+        """
+        Create a new trace using the V3 API format.
+
+        NB: The backend API is named "StartTraceV3" for some internal reason, but actually
+        it is supposed to be called at the end of the trace.
+
+        Args:
+            trace_info: The TraceInfo object to create in the backend.
+
+        Returns:
+            The returned TraceInfo object from the backend.
+        """
+        # NB: The Databricks backend expects a Trace object, not a TraceInfo object, although
+        # it doesn't use the data field at all. Trace data increases the payload size significantly,
+        # so we create a Trace object with an empty data field here.
+        trace = Trace(info=trace_info, data=TraceData(spans=[]))
+        req_body = message_to_json(StartTraceV3(trace=trace.to_proto()))
+
+        try:
+            response_proto = self._call_endpoint(
+                # NB: _call_endpoint doesn't handle versioning between v2 and v3 endpoint
+                # yet, so manually passing the v3 endpoint here.
+                StartTraceV3,
+                req_body,
+                endpoint=_V3_TRACE_REST_API_PATH_PREFIX,
+                retry_timeout_seconds=MLFLOW_ASYNC_TRACE_LOGGING_RETRY_TIMEOUT.get(),
+            )
+            return TraceInfo.from_proto(response_proto.trace.trace_info)
+        except MlflowException as e:
+            if e.error_code == databricks_pb2.ErrorCode.Name(databricks_pb2.ENDPOINT_NOT_FOUND):
+                _logger.debug(
+                    "Server does not support StartTraceV3 API yet. Falling back to V2 API."
+                )
+                return self._create_trace_v2_fallback(trace_info)
+            raise
+
+    def _create_trace_v2_fallback(self, trace_info: TraceInfo) -> TraceInfo:
+        """
+        Create a new trace using the V2 API format. This is a fallback for the case where the
+        client is v3 but the tracking server does not support v3 yet(<= 3.2.0).
+        """
+        trace_info_v2 = self.deprecated_start_trace_v2(
+            experiment_id=trace_info.experiment_id,
+            timestamp_ms=trace_info.request_time,
+            request_metadata=trace_info.trace_metadata,
+            tags=trace_info.tags,
+        )
+        self.deprecated_end_trace_v2(
+            request_id=trace_info_v2.request_id,
+            timestamp_ms=trace_info.request_time + trace_info.execution_duration,
+            status=trace_info.status,
+            request_metadata=trace_info.trace_metadata,
+            tags=trace_info.tags,
+        )
+        return trace_info_v2.to_v3()
+
+    def _delete_traces(
+        self,
+        experiment_id: str,
+        max_timestamp_millis: Optional[int] = None,
+        max_traces: Optional[int] = None,
+        trace_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=trace_ids,
+            )
+        )
+        res = self._call_endpoint(DeleteTraces, req_body)
+        return res.traces_deleted
+
+    def get_trace_info(self, trace_id: str) -> TraceInfo:
+        """
+        Get the trace matching the `trace_id`.
+
+        Args:
+            trace_id: String id of the trace to fetch.
+
+        Returns:
+            The fetched Trace object, of type ``mlflow.entities.TraceInfo``.
+        """
+        trace_v3_req_body = message_to_json(GetTraceInfoV3(trace_id=trace_id))
+        trace_v3_endpoint = get_single_trace_endpoint(trace_id)
+        try:
+            trace_v3_response_proto = self._call_endpoint(
+                GetTraceInfoV3, trace_v3_req_body, endpoint=trace_v3_endpoint
+            )
+            return TraceInfo.from_proto(trace_v3_response_proto.trace.trace_info)
+        except MlflowException as e:
+            # If the tracking server does not support V3 trace API yet, fallback to V2 API.
+            if e.error_code != databricks_pb2.ErrorCode.Name(databricks_pb2.ENDPOINT_NOT_FOUND):
+                raise
+            _logger.debug("Server does not support GetTraceInfoV3 API yet. Falling back to V2 API.")
+
+        req_body = message_to_json(GetTraceInfo(request_id=trace_id))
+        endpoint = get_single_trace_endpoint(trace_id, use_v3=False)
+        response_proto = self._call_endpoint(GetTraceInfo, req_body, endpoint=endpoint)
+        return TraceInfoV2.from_proto(response_proto.trace_info).to_v3()
+
+    def get_online_trace_details(
+        self,
+        trace_id: str,
+        sql_warehouse_id: str,
+        source_inference_table: str,
+        source_databricks_request_id: str,
+    ):
+        req = GetOnlineTraceDetails(
+            trace_id=trace_id,
+            sql_warehouse_id=sql_warehouse_id,
+            source_inference_table=source_inference_table,
+            source_databricks_request_id=source_databricks_request_id,
+        )
+        req_body = message_to_json(req)
+        response_proto = self._call_endpoint(GetOnlineTraceDetails, req_body)
+        return response_proto.trace_data
+
+    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,
+        model_id: Optional[str] = None,
+        sql_warehouse_id: Optional[str] = None,
+    ):
+        if sql_warehouse_id is None:
+            # Create trace_locations from experiment_ids for the V3 API
+            trace_locations = []
+            for exp_id in experiment_ids:
+                try:
+                    location = TraceLocation.from_experiment_id(exp_id)
+                    proto_location = location.to_proto()
+                    trace_locations.append(proto_location)
+                except Exception as e:
+                    raise MlflowException(
+                        f"Invalid experiment ID format: {exp_id}. Error: {e!s}"
+                    ) from e
+
+            # Create V3 request message using protobuf
+            request = SearchTracesV3(
+                locations=trace_locations,
+                filter=filter_string,
+                max_results=max_results,
+                order_by=order_by,
+                page_token=page_token,
+            )
+
+            req_body = message_to_json(request)
+            v3_endpoint = f"{_V3_TRACE_REST_API_PATH_PREFIX}/search"
+
+            try:
+                response_proto = self._call_endpoint(SearchTracesV3, req_body, v3_endpoint)
+            except MlflowException as e:
+                if e.error_code == databricks_pb2.ErrorCode.Name(databricks_pb2.ENDPOINT_NOT_FOUND):
+                    _logger.debug(
+                        "Server does not support SearchTracesV3 API yet. Falling back to V2 API."
+                    )
+                    response_proto = self._call_endpoint(SearchTraces, req_body)
+                else:
+                    raise
+
+            trace_infos = [TraceInfo.from_proto(t) for t in response_proto.traces]
+        else:
+            response_proto = self._search_unified_traces(
+                model_id=model_id,
+                sql_warehouse_id=sql_warehouse_id,
+                experiment_ids=experiment_ids,
+                filter_string=filter_string,
+                max_results=max_results,
+                order_by=order_by,
+                page_token=page_token,
+            )
+            # Convert TraceInfo (v2) objects to TraceInfoV3 objects for consistency
+            trace_infos = [TraceInfo.from_proto(t) for t in response_proto.traces]
+        return trace_infos, response_proto.next_page_token or None
+
+    def _search_unified_traces(
+        self,
+        model_id: str,
+        sql_warehouse_id: str,
+        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,
+    ):
+        request = SearchUnifiedTraces(
+            model_id=model_id,
+            sql_warehouse_id=sql_warehouse_id,
+            experiment_ids=experiment_ids,
+            filter=filter_string,
+            max_results=max_results,
+            order_by=order_by,
+            page_token=page_token,
+        )
+        req_body = message_to_json(request)
+        return self._call_endpoint(SearchUnifiedTraces, req_body)
+
+    def set_trace_tag(self, trace_id: str, key: str, value: str):
+        """
+        Set a tag on the trace with the given trace_id.
+
+        Args:
+            trace_id: The ID of the trace.
+            key: The string key of the tag.
+            value: The string value of the tag.
+        """
+        # Always use v2 endpoint
+        req_body = message_to_json(SetTraceTag(key=key, value=value))
+        self._call_endpoint(SetTraceTag, req_body, endpoint=get_trace_tag_endpoint(trace_id))
+
+    def delete_trace_tag(self, trace_id: str, key: str):
+        """
+        Delete a tag on the trace with the given trace_id.
+
+        Args:
+            trace_id: The ID of the trace.
+            key: The string key of the tag.
+        """
+        # Always use v2 endpoint
+        req_body = message_to_json(DeleteTraceTag(key=key))
+        self._call_endpoint(DeleteTraceTag, req_body, endpoint=get_trace_tag_endpoint(trace_id))
+
+    def get_assessment(self, trace_id: str, assessment_id: str) -> Assessment:
+        """
+        Get an assessment entity from the backend store.
+        """
+        req_body = message_to_json(
+            GetAssessmentRequest(trace_id=trace_id, assessment_id=assessment_id)
+        )
+        response_proto = self._call_endpoint(
+            GetAssessmentRequest,
+            req_body,
+            endpoint=get_single_assessment_endpoint(trace_id, assessment_id),
+        )
+        return Assessment.from_proto(response_proto.assessment)
+
+    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=f"{_V3_TRACE_REST_API_PATH_PREFIX}/{assessment.trace_id}/assessments",
+        )
+        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:
+            assessment.expectation.CopyFrom(expectation.to_proto())
+            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,
+                model_id=metric.model_id,
+                dataset_name=metric.dataset_name,
+                dataset_digest=metric.dataset_digest,
+            )
+        )
+        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 create_logged_model(
+        self,
+        experiment_id: str,
+        name: Optional[str] = None,
+        source_run_id: Optional[str] = None,
+        tags: Optional[list[LoggedModelTag]] = None,
+        params: Optional[list[LoggedModelParameter]] = None,
+        model_type: Optional[str] = None,
+    ) -> LoggedModel:
+        """
+        Create a new logged model.
+
+        Args:
+            experiment_id: ID of the experiment to which the model belongs.
+            name: Name of the model. If not specified, a random name will be generated.
+            source_run_id: ID of the run that produced the model.
+            tags: Tags to set on the model.
+            params: Parameters to set on the model.
+            model_type: Type of the model.
+
+        Returns:
+            The created model.
+        """
+        # Include the first 100 params in the initial request
+        initial_params = []
+        remaining_params = []
+        if params:
+            initial_batch_size = _MLFLOW_CREATE_LOGGED_MODEL_PARAMS_BATCH_SIZE.get()
+            initial_params = params[:initial_batch_size]
+            remaining_params = params[initial_batch_size:]
+
+        req_body = message_to_json(
+            CreateLoggedModel(
+                experiment_id=experiment_id,
+                name=name,
+                model_type=model_type,
+                source_run_id=source_run_id,
+                params=[p.to_proto() for p in initial_params],
+                tags=[t.to_proto() for t in tags or []],
+            )
+        )
+
+        response_proto = self._call_endpoint(CreateLoggedModel, req_body)
+        model = LoggedModel.from_proto(response_proto.model)
+
+        # Log remaining params if there are any
+        if remaining_params:
+            self.log_logged_model_params(model_id=model.model_id, params=remaining_params)
+            model = self.get_logged_model(model_id=model.model_id)
+
+        return model
+
+    def log_logged_model_params(self, model_id: str, params: list[LoggedModelParameter]) -> None:
+        """
+        Log parameters for a logged model in batches of 100.
+
+        Args:
+            model_id: ID of the model to log parameters for.
+            params: List of parameters to log.
+
+        Returns:
+            None
+        """
+        # Process params in batches to avoid exceeding per-request backend limits
+        batch_size = _MLFLOW_LOG_LOGGED_MODEL_PARAMS_BATCH_SIZE.get()
+        endpoint = get_logged_model_endpoint(model_id)
+        for i in range(0, len(params), batch_size):
+            batch = params[i : i + batch_size]
+            req_body = message_to_json(
+                LogLoggedModelParamsRequest(
+                    model_id=model_id,
+                    params=[p.to_proto() for p in batch],
+                )
+            )
+            self._call_endpoint(
+                LogLoggedModelParamsRequest, json_body=req_body, endpoint=f"{endpoint}/params"
+            )
+
+    def get_logged_model(self, model_id: str) -> LoggedModel:
+        """
+        Fetch the logged model with the specified ID.
+
+        Args:
+            model_id: ID of the model to fetch.
+
+        Returns:
+            The fetched model.
+        """
+        endpoint = get_logged_model_endpoint(model_id)
+        response_proto = self._call_endpoint(GetLoggedModel, endpoint=endpoint)
+        return LoggedModel.from_proto(response_proto.model)
+
+    def delete_logged_model(self, model_id) -> None:
+        request = DeleteLoggedModel(model_id=model_id)
+        endpoint = get_logged_model_endpoint(model_id)
+        self._call_endpoint(
+            DeleteLoggedModel, endpoint=endpoint, json_body=message_to_json(request)
+        )
+
+    def search_logged_models(
+        self,
+        experiment_ids: list[str],
+        filter_string: Optional[str] = None,
+        datasets: Optional[list[dict[str, Any]]] = None,
+        max_results: Optional[int] = None,
+        order_by: Optional[list[dict[str, Any]]] = None,
+        page_token: Optional[str] = None,
+    ) -> PagedList[LoggedModel]:
+        """
+        Search for logged models that match the specified search criteria.
+
+        Args:
+            experiment_ids: List of experiment ids to scope the search.
+            filter_string: A search filter string.
+            datasets: List of dictionaries to specify datasets on which to apply metrics filters.
+                The following fields are supported:
+
+                dataset_name (str): Required. Name of the dataset.
+                dataset_digest (str): Optional. Digest of the dataset.
+            max_results: Maximum number of logged models desired.
+            order_by: List of dictionaries to specify the ordering of the search results.
+                The following fields are supported:
+
+                field_name (str): Required. Name of the field to order by, e.g. "metrics.accuracy".
+                ascending: (bool): Optional. Whether the order is ascending or not.
+                dataset_name: (str): Optional. If ``field_name`` refers to a metric, this field
+                    specifies the name of the dataset associated with the metric. Only metrics
+                    associated with the specified dataset name will be considered for ordering.
+                    This field may only be set if ``field_name`` refers to a metric.
+                dataset_digest (str): Optional. If ``field_name`` refers to a metric, this field
+                    specifies the digest of the dataset associated with the metric. Only metrics
+                    associated with the specified dataset name and digest will be considered for
+                    ordering. This field may only be set if ``dataset_name`` is also set.
+            page_token: Token specifying the next page of results.
+
+        Returns:
+            A :py:class:`PagedList <mlflow.store.entities.PagedList>` of
+            :py:class:`LoggedModel <mlflow.entities.LoggedModel>` objects.
+        """
+        req_body = message_to_json(
+            SearchLoggedModels(
+                experiment_ids=experiment_ids,
+                filter=filter_string,
+                datasets=[
+                    SearchLoggedModels.Dataset(
+                        dataset_name=d["dataset_name"],
+                        dataset_digest=d.get("dataset_digest"),
+                    )
+                    for d in datasets or []
+                ],
+                max_results=max_results,
+                order_by=[
+                    SearchLoggedModels.OrderBy(
+                        field_name=d["field_name"],
+                        ascending=d.get("ascending", True),
+                        dataset_name=d.get("dataset_name"),
+                        dataset_digest=d.get("dataset_digest"),
+                    )
+                    for d in order_by or []
+                ],
+                page_token=page_token,
+            )
+        )
+        response_proto = self._call_endpoint(SearchLoggedModels, req_body)
+        models = [LoggedModel.from_proto(x) for x in response_proto.models]
+        return PagedList(models, response_proto.next_page_token or None)
+
+    def finalize_logged_model(self, model_id: str, status: LoggedModelStatus) -> LoggedModel:
+        """
+        Finalize a model by updating its status.
+
+        Args:
+            model_id: ID of the model to finalize.
+            status: Final status to set on the model.
+
+        Returns:
+            The updated model.
+        """
+        endpoint = get_logged_model_endpoint(model_id)
+        json_body = message_to_json(
+            FinalizeLoggedModel(model_id=model_id, status=status.to_proto())
+        )
+        response_proto = self._call_endpoint(
+            FinalizeLoggedModel, json_body=json_body, endpoint=endpoint
+        )
+        return LoggedModel.from_proto(response_proto.model)
+
+    def set_logged_model_tags(self, model_id: str, tags: list[LoggedModelTag]) -> None:
+        """
+        Set tags on the specified logged model.
+
+        Args:
+            model_id: ID of the model.
+            tags: Tags to set on the model.
+
+        Returns:
+            None
+        """
+        endpoint = get_logged_model_endpoint(model_id)
+        json_body = message_to_json(SetLoggedModelTags(tags=[tag.to_proto() for tag in tags]))
+        self._call_endpoint(SetLoggedModelTags, json_body=json_body, endpoint=f"{endpoint}/tags")
+
+    def delete_logged_model_tag(self, model_id: str, key: str) -> None:
+        """
+        Delete a tag from the specified logged model.
+
+        Args:
+            model_id: ID of the model.
+            key: Key of the tag to delete.
+
+        Returns:
+            The model with the specified tag removed.
+        """
+        endpoint = get_logged_model_endpoint(model_id)
+        self._call_endpoint(DeleteLoggedModelTag, endpoint=f"{endpoint}/tags/{key}")
+
+    def log_inputs(
+        self,
+        run_id: str,
+        datasets: Optional[list[DatasetInput]] = None,
+        models: Optional[list[LoggedModelInput]] = 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.
+            models: List of :py:class:`mlflow.entities.LoggedModelInput` instances to log.
+
+        Returns:
+            None.
+        """
+        datasets_protos = [dataset.to_proto() for dataset in datasets or []]
+        models_protos = [model.to_proto() for model in models or []]
+        req_body = message_to_json(
+            LogInputs(
+                run_id=run_id,
+                datasets=datasets_protos,
+                models=models_protos,
+            )
+        )
+        self._call_endpoint(LogInputs, req_body)
+
+    def log_outputs(self, run_id: str, models: list[LoggedModelOutput]):
+        """
+        Log outputs, such as models, to the specified run.
+
+        Args:
+            run_id: String id for the run
+            models: List of :py:class:`mlflow.entities.LoggedModelOutput` instances to log
+                as outputs of the run.
+
+        Returns:
+            None.
+        """
+        req_body = message_to_json(LogOutputs(run_id=run_id, models=[m.to_proto() for m in models]))
+        self._call_endpoint(LogOutputs, req_body)
+
+    ############################################################################################
+    # Deprecated MLflow Tracing APIs. Kept for backward compatibility but do not use.
+    ############################################################################################
+    def deprecated_start_trace_v2(
+        self,
+        experiment_id: str,
+        timestamp_ms: int,
+        request_metadata: dict[str, str],
+        tags: dict[str, str],
+    ) -> TraceInfoV2:
+        """
+        DEPRECATED. DO NOT USE.
+
+        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 TraceInfoV2.from_proto(response_proto.trace_info)
+
+    def deprecated_end_trace_v2(
+        self,
+        request_id: str,
+        timestamp_ms: int,
+        status: TraceStatus,
+        request_metadata: dict[str, str],
+        tags: dict[str, str],
+    ) -> TraceInfoV2:
+        """
+        DEPRECATED. DO NOT USE.
+
+        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
+        # Always use v2 endpoint (not v3) for this endpoint to maintain compatibility
+        endpoint = f"{_REST_API_PATH_PREFIX}/mlflow/traces/{request_id}"
+        response_proto = self._call_endpoint(EndTrace, req_body, endpoint=endpoint)
+        return TraceInfoV2.from_proto(response_proto.trace_info)