/mlflow/store/tracking/file_store.py

from __future__ import annotations

import hashlib
import json
import logging
import os
import shutil
import sys
import time
import uuid
from collections import defaultdict
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, NamedTuple, TypedDict

from mlflow.entities import (
    Assessment,
    Dataset,
    DatasetInput,
    Expectation,
    Experiment,
    ExperimentTag,
    Feedback,
    InputTag,
    LoggedModel,
    LoggedModelInput,
    LoggedModelOutput,
    LoggedModelParameter,
    LoggedModelStatus,
    LoggedModelTag,
    Metric,
    Param,
    Run,
    RunData,
    RunInfo,
    RunInputs,
    RunOutputs,
    RunStatus,
    RunTag,
    SourceType,
    TraceInfo,
    ViewType,
    _DatasetSummary,
)
from mlflow.entities.lifecycle_stage import LifecycleStage
from mlflow.entities.run_info import check_run_is_active
from mlflow.entities.trace_info_v2 import TraceInfoV2
from mlflow.entities.trace_status import TraceStatus
from mlflow.environment_variables import MLFLOW_ALLOW_FILE_STORE, MLFLOW_TRACKING_DIR
from mlflow.exceptions import MissingConfigException, MlflowException
from mlflow.protos import databricks_pb2
from mlflow.protos.databricks_pb2 import (
    INTERNAL_ERROR,
    INVALID_PARAMETER_VALUE,
    RESOURCE_DOES_NOT_EXIST,
)
from mlflow.protos.internal_pb2 import InputVertexType, OutputVertexType
from mlflow.store.entities.paged_list import PagedList
from mlflow.store.model_registry.file_store import FileStore as ModelRegistryFileStore
from mlflow.store.tracking import (
    DEFAULT_LOCAL_FILE_AND_ARTIFACT_PATH,
    SEARCH_LOGGED_MODEL_MAX_RESULTS_DEFAULT,
    SEARCH_MAX_RESULTS_DEFAULT,
    SEARCH_MAX_RESULTS_THRESHOLD,
    SEARCH_TRACES_DEFAULT_MAX_RESULTS,
)
from mlflow.store.tracking._sql_backend_utils import filestore_not_supported
from mlflow.store.tracking.abstract_store import AbstractStore
from mlflow.tracing.utils import (
    generate_assessment_id,
    generate_request_id_v2,
)
from mlflow.utils import get_results_from_paginated_fn
from mlflow.utils.file_utils import (
    append_to,
    exists,
    find,
    get_parent_dir,
    is_directory,
    list_all,
    list_subdirs,
    local_file_uri_to_path,
    make_containing_dirs,
    mkdir,
    mv,
    path_to_local_file_uri,
    read_file,
    read_file_lines,
    write_to,
)
from mlflow.utils.mlflow_tags import (
    MLFLOW_ARTIFACT_LOCATION,
    MLFLOW_DATASET_CONTEXT,
    MLFLOW_LOGGED_MODELS,
    MLFLOW_RUN_NAME,
    _get_run_name_from_tags,
)
from mlflow.utils.name_utils import _generate_random_name, _generate_unique_integer_id
from mlflow.utils.search_utils import (
    SearchExperimentsUtils,
    SearchLoggedModelsUtils,
    SearchTraceUtils,
    SearchUtils,
)
from mlflow.utils.string_utils import is_string_type
from mlflow.utils.time import get_current_time_millis
from mlflow.utils.uri import (
    append_to_uri_path,
    resolve_uri_if_local,
)
from mlflow.utils.validation import (
    _resolve_experiment_ids_and_locations,
    _validate_batch_log_data,
    _validate_batch_log_limits,
    _validate_experiment_artifact_location_length,
    _validate_experiment_id,
    _validate_experiment_name,
    _validate_experiment_tag,
    _validate_logged_model_name,
    _validate_metric,
    _validate_metric_name,
    _validate_param,
    _validate_param_keys_unique,
    _validate_param_name,
    _validate_run_id,
    _validate_tag_name,
)
from mlflow.utils.yaml_utils import overwrite_yaml, read_yaml, write_yaml

if TYPE_CHECKING:
    from mlflow.entities.model_registry.prompt_version import PromptVersion

_logger = logging.getLogger(__name__)


def _default_root_dir():
    return MLFLOW_TRACKING_DIR.get() or os.path.abspath(DEFAULT_LOCAL_FILE_AND_ARTIFACT_PATH)


def _read_persisted_experiment_dict(experiment_dict):
    dict_copy = experiment_dict.copy()

    # 'experiment_id' was changed from int to string, so we must cast to string
    # when reading legacy experiments
    if isinstance(dict_copy["experiment_id"], int):
        dict_copy["experiment_id"] = str(dict_copy["experiment_id"])
    return Experiment.from_dictionary(dict_copy)


def _make_persisted_run_info_dict(run_info):
    # 'tags' was moved from RunInfo to RunData, so we must keep storing it in the meta.yaml for
    # old mlflow versions to read
    run_info_dict = dict(run_info)
    run_info_dict["tags"] = []
    if "status" in run_info_dict:
        # 'status' is stored as an integer enum in meta file, but RunInfo.status field is a string.
        # Convert from string to enum/int before storing.
        run_info_dict["status"] = RunStatus.from_string(run_info.status)
    else:
        run_info_dict["status"] = RunStatus.RUNNING
    run_info_dict["source_type"] = SourceType.LOCAL
    run_info_dict["source_name"] = ""
    run_info_dict["entry_point_name"] = ""
    run_info_dict["source_version"] = ""
    return run_info_dict


def _read_persisted_run_info_dict(run_info_dict):
    dict_copy = run_info_dict.copy()
    if "lifecycle_stage" not in dict_copy:
        dict_copy["lifecycle_stage"] = LifecycleStage.ACTIVE
    # 'status' is stored as an integer enum in meta file, but RunInfo.status field is a string.
    # converting to string before hydrating RunInfo.
    # If 'status' value not recorded in files, mark it as 'RUNNING' (default)
    dict_copy["status"] = RunStatus.to_string(run_info_dict.get("status", RunStatus.RUNNING))

    # 'experiment_id' was changed from int to string, so we must cast to string
    # when reading legacy run_infos
    if isinstance(dict_copy["experiment_id"], int):
        dict_copy["experiment_id"] = str(dict_copy["experiment_id"])
    return RunInfo.from_dictionary(dict_copy)


class DatasetFilter(TypedDict, total=False):
    """
    Dataset filter used for search_logged_models.
    """

    dataset_name: str
    dataset_digest: str


class FileStore(AbstractStore):
    TRASH_FOLDER_NAME = ".trash"
    ARTIFACTS_FOLDER_NAME = "artifacts"
    METRICS_FOLDER_NAME = "metrics"
    PARAMS_FOLDER_NAME = "params"
    TAGS_FOLDER_NAME = "tags"
    EXPERIMENT_TAGS_FOLDER_NAME = "tags"
    DATASETS_FOLDER_NAME = "datasets"
    INPUTS_FOLDER_NAME = "inputs"
    OUTPUTS_FOLDER_NAME = "outputs"
    META_DATA_FILE_NAME = "meta.yaml"
    DEFAULT_EXPERIMENT_ID = "0"
    TRACE_INFO_FILE_NAME = "trace_info.yaml"
    TRACES_FOLDER_NAME = "traces"
    TRACE_TAGS_FOLDER_NAME = "tags"
    ASSESSMENTS_FOLDER_NAME = "assessments"
    # "request_metadata" field is renamed to "trace_metadata" in V3,
    # but we keep the old name for backward compatibility
    TRACE_TRACE_METADATA_FOLDER_NAME = "request_metadata"
    MODELS_FOLDER_NAME = "models"
    RESERVED_EXPERIMENT_FOLDERS = [
        EXPERIMENT_TAGS_FOLDER_NAME,
        DATASETS_FOLDER_NAME,
        TRACES_FOLDER_NAME,
        MODELS_FOLDER_NAME,
    ]

    def __init__(self, root_directory=None, artifact_root_uri=None):
        """
        Create a new FileStore with the given root directory and a given default artifact root URI.
        """
        super().__init__()
        if not MLFLOW_ALLOW_FILE_STORE.get():
            raise MlflowException(
                "The filesystem tracking backend (e.g., './mlruns') is in maintenance mode "
                "and will not receive further updates. Please migrate to a "
                "database backend (e.g., 'sqlite:///mlflow.db') to access the latest MLflow "
                "features. The `mlflow migrate-filestore` tool migrates your existing data "
                "losslessly. See "
                "https://mlflow.org/docs/latest/self-hosting/migrate-from-file-store "
                "for migration guidance. If the filesystem backend is required for your "
                "workflow, set `MLFLOW_ALLOW_FILE_STORE=true` to opt out of this exception.",
                error_code=INVALID_PARAMETER_VALUE,
            )
        self.root_directory = local_file_uri_to_path(root_directory or _default_root_dir())
        if not artifact_root_uri:
            self.artifact_root_uri = path_to_local_file_uri(self.root_directory)
        else:
            self.artifact_root_uri = resolve_uri_if_local(artifact_root_uri)
        self.trash_folder = os.path.join(self.root_directory, FileStore.TRASH_FOLDER_NAME)
        # Create root directory if needed
        if not exists(self.root_directory):
            self._create_default_experiment()
        # Create trash folder if needed
        if not exists(self.trash_folder):
            mkdir(self.trash_folder)

    def _create_default_experiment(self):
        mkdir(self.root_directory)
        self._create_experiment_with_id(
            name=Experiment.DEFAULT_EXPERIMENT_NAME,
            experiment_id=FileStore.DEFAULT_EXPERIMENT_ID,
            artifact_uri=None,
            tags=None,
        )

    def _check_root_dir(self):
        """
        Run checks before running directory operations.
        """
        if not exists(self.root_directory):
            raise Exception(f"'{self.root_directory}' does not exist.")
        if not is_directory(self.root_directory):
            raise Exception(f"'{self.root_directory}' is not a directory.")

    def _get_experiment_path(self, experiment_id, view_type=ViewType.ALL, assert_exists=False):
        parents = []
        if view_type in (ViewType.ACTIVE_ONLY, ViewType.ALL):
            parents.append(self.root_directory)
        if view_type in (ViewType.DELETED_ONLY, ViewType.ALL):
            parents.append(self.trash_folder)
        for parent in parents:
            exp_list = find(parent, experiment_id, full_path=True)
            if len(exp_list) > 0:
                return exp_list[0]
        if assert_exists:
            raise MlflowException(
                f"Experiment {experiment_id} does not exist.",
                databricks_pb2.RESOURCE_DOES_NOT_EXIST,
            )
        return None

    def _get_run_dir(self, experiment_id, run_uuid):
        _validate_run_id(run_uuid)
        if not self._has_experiment(experiment_id):
            return None
        return os.path.join(self._get_experiment_path(experiment_id, assert_exists=True), run_uuid)

    def _get_metric_path(self, experiment_id, run_uuid, metric_key):
        _validate_run_id(run_uuid)
        _validate_metric_name(metric_key, "name")
        return os.path.join(
            self._get_run_dir(experiment_id, run_uuid),
            FileStore.METRICS_FOLDER_NAME,
            metric_key,
        )

    def _get_model_metric_path(self, experiment_id: str, model_id: str, metric_key: str) -> str:
        _validate_metric_name(metric_key)
        return os.path.join(
            self._get_model_dir(experiment_id, model_id), FileStore.METRICS_FOLDER_NAME, metric_key
        )

    def _get_param_path(self, experiment_id, run_uuid, param_name):
        _validate_run_id(run_uuid)
        _validate_param_name(param_name)
        return os.path.join(
            self._get_run_dir(experiment_id, run_uuid),
            FileStore.PARAMS_FOLDER_NAME,
            param_name,
        )

    def _get_experiment_tag_path(self, experiment_id, tag_name):
        _validate_experiment_id(experiment_id)
        _validate_tag_name(tag_name)
        if not self._has_experiment(experiment_id):
            return None
        return os.path.join(
            self._get_experiment_path(experiment_id, assert_exists=True),
            FileStore.TAGS_FOLDER_NAME,
            tag_name,
        )

    def _get_tag_path(self, experiment_id, run_uuid, tag_name):
        _validate_run_id(run_uuid)
        _validate_tag_name(tag_name)
        return os.path.join(
            self._get_run_dir(experiment_id, run_uuid),
            FileStore.TAGS_FOLDER_NAME,
            tag_name,
        )

    def _get_artifact_dir(self, experiment_id, run_uuid):
        _validate_run_id(run_uuid)
        return append_to_uri_path(
            self.get_experiment(experiment_id).artifact_location,
            run_uuid,
            FileStore.ARTIFACTS_FOLDER_NAME,
        )

    def _get_active_experiments(self, full_path=False):
        exp_list = list_subdirs(self.root_directory, full_path)
        return [
            exp
            for exp in exp_list
            if not exp.endswith(FileStore.TRASH_FOLDER_NAME)
            and exp != ModelRegistryFileStore.MODELS_FOLDER_NAME
        ]

    def _get_deleted_experiments(self, full_path=False):
        return list_subdirs(self.trash_folder, full_path)

    def search_experiments(
        self,
        view_type=ViewType.ACTIVE_ONLY,
        max_results=SEARCH_MAX_RESULTS_DEFAULT,
        filter_string=None,
        order_by=None,
        page_token=None,
    ):
        if not isinstance(max_results, int) or max_results < 1:
            raise MlflowException(
                f"Invalid value {max_results} for parameter 'max_results' supplied. It must be "
                f"a positive integer",
                INVALID_PARAMETER_VALUE,
            )
        if max_results > SEARCH_MAX_RESULTS_THRESHOLD:
            raise MlflowException(
                f"Invalid value {max_results} for parameter 'max_results' supplied. It must be at "
                f"most {SEARCH_MAX_RESULTS_THRESHOLD}",
                INVALID_PARAMETER_VALUE,
            )

        self._check_root_dir()
        experiment_ids = []
        if view_type in (ViewType.ACTIVE_ONLY, ViewType.ALL):
            experiment_ids += self._get_active_experiments(full_path=False)
        if view_type in (ViewType.DELETED_ONLY, ViewType.ALL):
            experiment_ids += self._get_deleted_experiments(full_path=False)

        experiments = []
        for exp_id in experiment_ids:
            try:
                # trap and warn known issues, will raise unexpected exceptions to caller
                exp = self._get_experiment(exp_id, view_type)
                if exp is not None:
                    experiments.append(exp)
            except MissingConfigException as e:
                logging.warning(
                    f"Malformed experiment '{exp_id}'. Detailed error {e}",
                    exc_info=True,
                )
        filtered = SearchExperimentsUtils.filter(experiments, filter_string)
        sorted_experiments = SearchExperimentsUtils.sort(
            filtered, order_by or ["creation_time DESC", "experiment_id ASC"]
        )
        experiments, next_page_token = SearchUtils.paginate(
            sorted_experiments, page_token, max_results
        )
        return PagedList(experiments, next_page_token)

    def get_experiment_by_name(self, experiment_name):
        def pagination_wrapper_func(number_to_get, next_page_token):
            return self.search_experiments(
                view_type=ViewType.ALL,
                max_results=number_to_get,
                filter_string=f"name = '{experiment_name}'",
                page_token=next_page_token,
            )

        experiments = get_results_from_paginated_fn(
            paginated_fn=pagination_wrapper_func,
            max_results_per_page=SEARCH_MAX_RESULTS_THRESHOLD,
            max_results=None,
        )
        return experiments[0] if len(experiments) > 0 else None

    def _create_experiment_with_id(self, name, experiment_id, artifact_uri, tags):
        if not artifact_uri:
            resolved_artifact_uri = append_to_uri_path(self.artifact_root_uri, str(experiment_id))
        else:
            resolved_artifact_uri = resolve_uri_if_local(artifact_uri)
        meta_dir = mkdir(self.root_directory, str(experiment_id))
        creation_time = get_current_time_millis()
        experiment = Experiment(
            experiment_id,
            name,
            resolved_artifact_uri,
            LifecycleStage.ACTIVE,
            creation_time=creation_time,
            last_update_time=creation_time,
        )
        experiment_dict = dict(experiment)
        # tags are added to the file system and are not written to this dict on write
        # As such, we should not include them in the meta file.
        del experiment_dict["tags"]
        write_yaml(meta_dir, FileStore.META_DATA_FILE_NAME, experiment_dict)
        if tags is not None:
            for tag in tags:
                self.set_experiment_tag(experiment_id, tag)
        return experiment_id

    def _validate_experiment_does_not_exist(self, name):
        experiment = self.get_experiment_by_name(name)
        if experiment is not None:
            if experiment.lifecycle_stage == LifecycleStage.DELETED:
                raise MlflowException(
                    f"Experiment {experiment.name!r} already exists in deleted state. "
                    "You can restore the experiment, or permanently delete the experiment "
                    "from the .trash folder (under tracking server's root folder) in order to "
                    "use this experiment name again.",
                    databricks_pb2.RESOURCE_ALREADY_EXISTS,
                )
            else:
                raise MlflowException(
                    f"Experiment '{experiment.name}' already exists.",
                    databricks_pb2.RESOURCE_ALREADY_EXISTS,
                )

    def create_experiment(self, name, artifact_location=None, tags=None):
        self._check_root_dir()
        _validate_experiment_name(name)

        if artifact_location:
            _validate_experiment_artifact_location_length(artifact_location)
        if tags:
            for tag in tags:
                _validate_experiment_tag(tag.key, tag.value)

        self._validate_experiment_does_not_exist(name)
        experiment_id = _generate_unique_integer_id()
        return self._create_experiment_with_id(name, str(experiment_id), artifact_location, tags)

    def _has_experiment(self, experiment_id):
        return self._get_experiment_path(experiment_id) is not None

    def _get_experiment(self, experiment_id, view_type=ViewType.ALL):
        self._check_root_dir()
        _validate_experiment_id(experiment_id)
        experiment_dir = self._get_experiment_path(experiment_id, view_type)
        if experiment_dir is None:
            raise MlflowException(
                f"Could not find experiment with ID {experiment_id}",
                databricks_pb2.RESOURCE_DOES_NOT_EXIST,
            )
        meta = FileStore._read_yaml(experiment_dir, FileStore.META_DATA_FILE_NAME)
        if meta is None:
            raise MissingConfigException(
                f"Experiment {experiment_id} is invalid with empty "
                f"{FileStore.META_DATA_FILE_NAME} in directory '{experiment_dir}'."
            )

        meta["tags"] = self.get_all_experiment_tags(experiment_id)
        experiment = _read_persisted_experiment_dict(meta)
        if experiment_id != experiment.experiment_id:
            logging.warning(
                "Experiment ID mismatch for exp %s. ID recorded as '%s' in meta data. "
                "Experiment will be ignored.",
                experiment_id,
                experiment.experiment_id,
                exc_info=True,
            )
            return None
        return experiment

    def get_experiment(self, experiment_id):
        """
        Fetch the experiment.
        Note: This API will search for active as well as deleted experiments.

        Args:
            experiment_id: Integer id for the experiment

        Returns:
            A single Experiment object if it exists, otherwise raises an Exception.
        """
        experiment_id = FileStore.DEFAULT_EXPERIMENT_ID if experiment_id is None else experiment_id
        experiment = self._get_experiment(experiment_id)
        if experiment is None:
            raise MlflowException(
                f"Experiment '{experiment_id}' does not exist.",
                databricks_pb2.RESOURCE_DOES_NOT_EXIST,
            )
        return experiment

    def delete_experiment(self, experiment_id):
        if str(experiment_id) == str(FileStore.DEFAULT_EXPERIMENT_ID):
            raise MlflowException(
                "Cannot delete the default experiment "
                f"'{FileStore.DEFAULT_EXPERIMENT_ID}'. This is an internally "
                f"reserved experiment."
            )
        experiment_dir = self._get_experiment_path(experiment_id, ViewType.ACTIVE_ONLY)
        if experiment_dir is None:
            raise MlflowException(
                f"Could not find experiment with ID {experiment_id}",
                databricks_pb2.RESOURCE_DOES_NOT_EXIST,
            )
        experiment = self._get_experiment(experiment_id)
        experiment._lifecycle_stage = LifecycleStage.DELETED
        deletion_time = get_current_time_millis()
        experiment._set_last_update_time(deletion_time)
        runs = self._list_run_infos(experiment_id, view_type=ViewType.ACTIVE_ONLY)
        for run_info in runs:
            if run_info is not None:
                new_info = run_info._copy_with_overrides(lifecycle_stage=LifecycleStage.DELETED)
                self._overwrite_run_info(new_info, deleted_time=deletion_time)
            else:
                logging.warning("Run metadata is in invalid state.")
        meta_dir = os.path.join(self.root_directory, experiment_id)
        overwrite_yaml(
            root=meta_dir,
            file_name=FileStore.META_DATA_FILE_NAME,
            data=dict(experiment),
        )
        mv(experiment_dir, self.trash_folder)

    def _hard_delete_experiment(self, experiment_id):
        """
        Permanently delete an experiment.
        This is used by the ``mlflow gc`` command line and is not intended to be used elsewhere.
        """
        experiment_dir = self._get_experiment_path(experiment_id, ViewType.DELETED_ONLY)
        shutil.rmtree(experiment_dir)

    def restore_experiment(self, experiment_id):
        experiment_dir = self._get_experiment_path(experiment_id, ViewType.DELETED_ONLY)
        if experiment_dir is None:
            raise MlflowException(
                f"Could not find deleted experiment with ID {experiment_id}",
                databricks_pb2.RESOURCE_DOES_NOT_EXIST,
            )
        conflict_experiment = self._get_experiment_path(experiment_id, ViewType.ACTIVE_ONLY)
        if conflict_experiment is not None:
            raise MlflowException(
                f"Cannot restore experiment with ID {experiment_id}. "
                "An experiment with same ID already exists.",
                databricks_pb2.RESOURCE_ALREADY_EXISTS,
            )
        mv(experiment_dir, self.root_directory)
        experiment = self._get_experiment(experiment_id)
        meta_dir = os.path.join(self.root_directory, experiment_id)
        experiment._lifecycle_stage = LifecycleStage.ACTIVE
        experiment._set_last_update_time(get_current_time_millis())
        runs = self._list_run_infos(experiment_id, view_type=ViewType.DELETED_ONLY)
        for run_info in runs:
            if run_info is not None:
                new_info = run_info._copy_with_overrides(lifecycle_stage=LifecycleStage.ACTIVE)
                self._overwrite_run_info(new_info, deleted_time=None)
            else:
                logging.warning("Run metadata is in invalid state.")
        overwrite_yaml(
            root=meta_dir,
            file_name=FileStore.META_DATA_FILE_NAME,
            data=dict(experiment),
        )

    def rename_experiment(self, experiment_id, new_name):
        _validate_experiment_name(new_name)
        meta_dir = os.path.join(self.root_directory, experiment_id)
        # if experiment is malformed, will raise error
        experiment = self._get_experiment(experiment_id)
        if experiment is None:
            raise MlflowException(
                f"Experiment '{experiment_id}' does not exist.",
                databricks_pb2.RESOURCE_DOES_NOT_EXIST,
            )
        self._validate_experiment_does_not_exist(new_name)
        experiment._set_name(new_name)
        experiment._set_last_update_time(get_current_time_millis())
        if experiment.lifecycle_stage != LifecycleStage.ACTIVE:
            raise Exception(
                "Cannot rename experiment in non-active lifecycle stage."
                f" Current stage: {experiment.lifecycle_stage}"
            )
        overwrite_yaml(
            root=meta_dir,
            file_name=FileStore.META_DATA_FILE_NAME,
            data=dict(experiment),
        )

    def delete_run(self, run_id):
        run_info = self._get_run_info(run_id)
        if run_info is None:
            raise MlflowException(
                f"Run '{run_id}' metadata is in invalid state.",
                databricks_pb2.INVALID_STATE,
            )
        new_info = run_info._copy_with_overrides(lifecycle_stage=LifecycleStage.DELETED)
        self._overwrite_run_info(new_info, deleted_time=get_current_time_millis())

    def _hard_delete_run(self, run_id):
        """
        Permanently delete a run (metadata and metrics, tags, parameters).
        This is used by the ``mlflow gc`` command line and is not intended to be used elsewhere.
        """
        # NB: Skip validation here since artifacts may have already been deleted
        # by gc before calling this method. The run_id was already validated
        # by search_runs/get_run before reaching this point.
        _, run_dir = self._find_run_root(run_id, validate_structure=False)
        shutil.rmtree(run_dir)

    def _get_deleted_runs(self, older_than=0):
        """
        Get all deleted run ids.

        Args:
            older_than: get runs that is older than this variable in number of milliseconds.
                        defaults to 0 ms to get all deleted runs.
        """
        current_time = get_current_time_millis()
        experiment_ids = self._get_active_experiments() + self._get_deleted_experiments()
        deleted_runs = self.search_runs(
            experiment_ids=experiment_ids,
            filter_string="",
            run_view_type=ViewType.DELETED_ONLY,
        )
        deleted_run_ids = []
        for deleted_run in deleted_runs:
            _, run_dir = self._find_run_root(deleted_run.info.run_id)
            meta = read_yaml(run_dir, FileStore.META_DATA_FILE_NAME)
            if "deleted_time" not in meta or current_time - int(meta["deleted_time"]) >= older_than:
                deleted_run_ids.append(deleted_run.info.run_id)

        return deleted_run_ids

    def restore_run(self, run_id):
        run_info = self._get_run_info(run_id)
        if run_info is None:
            raise MlflowException(
                f"Run '{run_id}' metadata is in invalid state.",
                databricks_pb2.INVALID_STATE,
            )
        new_info = run_info._copy_with_overrides(lifecycle_stage=LifecycleStage.ACTIVE)
        self._overwrite_run_info(new_info, deleted_time=None)

    def _find_experiment_folder(self, run_path):
        """
        Given a run path, return the parent directory for its experiment.
        """
        parent = get_parent_dir(run_path)
        if os.path.basename(parent) == FileStore.TRASH_FOLDER_NAME:
            return get_parent_dir(parent)
        return parent

    def _is_valid_run_directory(self, run_dir):
        # Defense in depth: ensure we're not inside an artifacts folder
        path_parts = os.path.normpath(run_dir).split(os.sep)
        if FileStore.ARTIFACTS_FOLDER_NAME in path_parts[:-1]:
            return False

        required_subdirs = [
            FileStore.METRICS_FOLDER_NAME,
            FileStore.PARAMS_FOLDER_NAME,
            FileStore.ARTIFACTS_FOLDER_NAME,
        ]
        return all(is_directory(os.path.join(run_dir, subdir)) for subdir in required_subdirs)

    def _find_run_root(self, run_uuid, validate_structure=True):
        _validate_run_id(run_uuid)
        self._check_root_dir()
        all_experiments = self._get_active_experiments(True) + self._get_deleted_experiments(True)
        for experiment_dir in all_experiments:
            runs = find(experiment_dir, run_uuid, full_path=True)
            if len(runs) == 0:
                continue
            run_dir = runs[0]
            # NB: Validate run directory structure to prevent path traversal via malicious
            # meta.yaml in artifact folders (ZDI-CAN-26649)
            if validate_structure and not self._is_valid_run_directory(run_dir):
                continue
            return os.path.basename(os.path.abspath(experiment_dir)), run_dir
        return None, None

    def update_run_info(self, run_id, run_status, end_time, run_name):
        _validate_run_id(run_id)
        run_info = self._get_run_info(run_id)
        check_run_is_active(run_info)
        new_info = run_info._copy_with_overrides(run_status, end_time, run_name=run_name)
        if run_name:
            self._set_run_tag(run_info, RunTag(MLFLOW_RUN_NAME, run_name))
        self._overwrite_run_info(new_info)
        return new_info

    def create_run(self, experiment_id, user_id, start_time, tags, run_name):
        """
        Creates a run with the specified attributes.
        """
        experiment_id = FileStore.DEFAULT_EXPERIMENT_ID if experiment_id is None else experiment_id
        experiment = self.get_experiment(experiment_id)
        if experiment is None:
            raise MlflowException(
                f"Could not create run under experiment with ID {experiment_id} - no such "
                "experiment exists.",
                databricks_pb2.RESOURCE_DOES_NOT_EXIST,
            )
        if experiment.lifecycle_stage != LifecycleStage.ACTIVE:
            raise MlflowException(
                f"Could not create run under non-active experiment with ID {experiment_id}.",
                databricks_pb2.INVALID_STATE,
            )
        tags = tags or []
        run_name_tag = _get_run_name_from_tags(tags)
        if run_name and run_name_tag and run_name != run_name_tag:
            raise MlflowException(
                "Both 'run_name' argument and 'mlflow.runName' tag are specified, but with "
                f"different values (run_name='{run_name}', mlflow.runName='{run_name_tag}').",
                INVALID_PARAMETER_VALUE,
            )
        run_name = run_name or run_name_tag or _generate_random_name()
        if not run_name_tag:
            tags.append(RunTag(key=MLFLOW_RUN_NAME, value=run_name))
        run_uuid = uuid.uuid4().hex
        artifact_uri = self._get_artifact_dir(experiment_id, run_uuid)
        run_info = RunInfo(
            run_id=run_uuid,
            run_name=run_name,
            experiment_id=experiment_id,
            artifact_uri=artifact_uri,
            user_id=user_id,
            status=RunStatus.to_string(RunStatus.RUNNING),
            start_time=start_time,
            end_time=None,
            lifecycle_stage=LifecycleStage.ACTIVE,
        )
        # Persist run metadata and create directories for logging metrics, parameters, artifacts
        run_dir = self._get_run_dir(run_info.experiment_id, run_info.run_id)
        mkdir(run_dir)
        run_info_dict = _make_persisted_run_info_dict(run_info)
        run_info_dict["deleted_time"] = None
        write_yaml(run_dir, FileStore.META_DATA_FILE_NAME, run_info_dict)
        mkdir(run_dir, FileStore.METRICS_FOLDER_NAME)
        mkdir(run_dir, FileStore.PARAMS_FOLDER_NAME)
        mkdir(run_dir, FileStore.ARTIFACTS_FOLDER_NAME)
        for tag in tags:
            self.set_tag(run_uuid, tag)
        return self.get_run(run_id=run_uuid)

    def get_run(self, run_id):
        """
        Note: Will get both active and deleted runs.
        """
        _validate_run_id(run_id)
        run_info = self._get_run_info(run_id)
        if run_info is None:
            raise MlflowException(
                f"Run '{run_id}' metadata is in invalid state.",
                databricks_pb2.INVALID_STATE,
            )
        return self._get_run_from_info(run_info)

    def _get_run_from_info(self, run_info):
        metrics = self._get_all_metrics(run_info)
        params = self._get_all_params(run_info)
        tags = self._get_all_tags(run_info)
        inputs: RunInputs = self._get_all_inputs(run_info)
        outputs: RunOutputs = self._get_all_outputs(run_info)
        if not run_info.run_name:
            if run_name := _get_run_name_from_tags(tags):
                run_info._set_run_name(run_name)
        return Run(run_info, RunData(metrics, params, tags), inputs, outputs)

    def _get_run_info(self, run_uuid):
        """
        Note: Will get both active and deleted runs.
        """
        exp_id, run_dir = self._find_run_root(run_uuid)
        if run_dir is None:
            raise MlflowException(
                f"Run '{run_uuid}' not found", databricks_pb2.RESOURCE_DOES_NOT_EXIST
            )
        run_info = self._get_run_info_from_dir(run_dir)
        if run_info.experiment_id != exp_id:
            raise MlflowException(
                f"Run '{run_uuid}' metadata is in invalid state.",
                databricks_pb2.INVALID_STATE,
            )
        # Defense in depth: verify run_id in meta.yaml matches the directory name
        if run_info.run_id != os.path.basename(run_dir):
            raise MlflowException(
                f"Run '{run_uuid}' metadata is in invalid state.",
                databricks_pb2.INVALID_STATE,
            )
        return run_info

    def _get_run_info_from_dir(self, run_dir):
        meta = FileStore._read_yaml(run_dir, FileStore.META_DATA_FILE_NAME)
        return _read_persisted_run_info_dict(meta)

    def _get_run_files(self, run_info, resource_type):
        run_dir = self._get_run_dir(run_info.experiment_id, run_info.run_id)
        # run_dir exists since run validity has been confirmed above.
        if resource_type == "metric":
            subfolder_name = FileStore.METRICS_FOLDER_NAME
        elif resource_type == "param":
            subfolder_name = FileStore.PARAMS_FOLDER_NAME
        elif resource_type == "tag":
            subfolder_name = FileStore.TAGS_FOLDER_NAME
        else:
            raise Exception("Looking for unknown resource under run.")
        return self._get_resource_files(run_dir, subfolder_name)

    def _get_experiment_files(self, experiment_id):
        _validate_experiment_id(experiment_id)
        experiment_dir = self._get_experiment_path(experiment_id, assert_exists=True)
        return self._get_resource_files(experiment_dir, FileStore.EXPERIMENT_TAGS_FOLDER_NAME)

    def _get_resource_files(self, root_dir, subfolder_name):
        source_dirs = find(root_dir, subfolder_name, full_path=True)
        if len(source_dirs) == 0:
            return root_dir, []
        file_names = []
        for root, _, files in os.walk(source_dirs[0]):
            for name in files:
                abspath = os.path.join(root, name)
                file_names.append(os.path.relpath(abspath, source_dirs[0]))
        if sys.platform == "win32":
            # Turn metric relative path into metric name.
            # Metrics can have '/' in the name. On windows, '/' is interpreted as a separator.
            # When the metric is read back the path will use '\' for separator.
            # We need to translate the path into posix path.
            from mlflow.utils.file_utils import relative_path_to_artifact_path

            file_names = [relative_path_to_artifact_path(x) for x in file_names]
        return source_dirs[0], file_names

    @staticmethod
    def _get_metric_from_file(
        parent_path: str, metric_name: str, run_id: str, exp_id: str
    ) -> Metric:
        _validate_metric_name(metric_name)
        metric_objs = [
            FileStore._get_metric_from_line(run_id, metric_name, line, exp_id)
            for line in read_file_lines(parent_path, metric_name)
        ]
        if len(metric_objs) == 0:
            raise ValueError(f"Metric '{metric_name}' is malformed. No data found.")
        # Python performs element-wise comparison of equal-length tuples, ordering them
        # based on their first differing element. Therefore, we use max() operator to find the
        # largest value at the largest timestamp. For more information, see
        # https://docs.python.org/3/reference/expressions.html#value-comparisons
        return max(metric_objs, key=lambda m: (m.step, m.timestamp, m.value))

    def get_all_metrics(self, run_uuid):
        _validate_run_id(run_uuid)
        run_info = self._get_run_info(run_uuid)
        return self._get_all_metrics(run_info)

    def _get_all_metrics(self, run_info):
        parent_path, metric_files = self._get_run_files(run_info, "metric")
        return [
            self._get_metric_from_file(
                parent_path, metric_file, run_info.run_id, run_info.experiment_id
            )
            for metric_file in metric_files
        ]

    @staticmethod
    def _get_metric_from_line(
        run_id: str, metric_name: str, metric_line: str, exp_id: str
    ) -> Metric:
        metric_parts = metric_line.strip().split(" ")
        if len(metric_parts) != 2 and len(metric_parts) != 3 and len(metric_parts) != 5:
            raise MlflowException(
                f"Metric '{metric_name}' is malformed; persisted metric data contained "
                f"{len(metric_parts)} fields. Expected 2, 3, or 5 fields. "
                f"Experiment id: {exp_id}",
                databricks_pb2.INTERNAL_ERROR,
            )
        ts = int(metric_parts[0])
        val = float(metric_parts[1])
        step = int(metric_parts[2]) if len(metric_parts) == 3 else 0
        dataset_name = str(metric_parts[3]) if len(metric_parts) == 5 else None
        dataset_digest = str(metric_parts[4]) if len(metric_parts) == 5 else None
        return Metric(
            key=metric_name,
            value=val,
            timestamp=ts,
            step=step,
            dataset_name=dataset_name,
            dataset_digest=dataset_digest,
            run_id=run_id,
        )

    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: An indicator for paginated results.
            page_token: Token indicating the page of metric history to fetch.

        Returns:
            A :py:class:`mlflow.store.entities.paged_list.PagedList` of
            :py:class:`mlflow.entities.Metric` entities if ``metric_key`` values
            have been logged to the ``run_id``, else an empty list.

        """
        _validate_run_id(run_id)
        _validate_metric_name(metric_key)
        run_info = self._get_run_info(run_id)

        parent_path, metric_files = self._get_run_files(run_info, "metric")
        if metric_key not in metric_files:
            return PagedList([], None)

        all_lines = read_file_lines(parent_path, metric_key)

        all_metrics = [
            FileStore._get_metric_from_line(run_id, metric_key, line, run_info.experiment_id)
            for line in all_lines
        ]

        if max_results is None:
            # If no max_results specified, return all metrics but handle page_token if provided
            offset = SearchUtils.parse_start_offset_from_page_token(page_token)
            metrics = all_metrics[offset:]
            next_page_token = None
        else:
            metrics, next_page_token = SearchUtils.paginate(all_metrics, page_token, max_results)

        return PagedList(metrics, next_page_token)

    @staticmethod
    def _get_param_from_file(parent_path, param_name):
        _validate_param_name(param_name)
        value = read_file(parent_path, param_name)
        return Param(param_name, value)

    def get_all_params(self, run_uuid):
        _validate_run_id(run_uuid)
        run_info = self._get_run_info(run_uuid)
        return self._get_all_params(run_info)

    def _get_all_params(self, run_info):
        parent_path, param_files = self._get_run_files(run_info, "param")
        return [self._get_param_from_file(parent_path, param_file) for param_file in param_files]

    @staticmethod
    def _get_experiment_tag_from_file(parent_path, tag_name):
        _validate_tag_name(tag_name)
        tag_data = read_file(parent_path, tag_name)
        return ExperimentTag(tag_name, tag_data)

    def get_all_experiment_tags(self, exp_id):
        parent_path, tag_files = self._get_experiment_files(exp_id)
        return [self._get_experiment_tag_from_file(parent_path, tag_file) for tag_file in tag_files]

    @staticmethod
    def _get_tag_from_file(parent_path, tag_name):
        _validate_tag_name(tag_name)
        tag_data = read_file(parent_path, tag_name)
        return RunTag(tag_name, tag_data)

    def get_all_tags(self, run_uuid):
        _validate_run_id(run_uuid)
        run_info = self._get_run_info(run_uuid)
        return self._get_all_tags(run_info)

    def _get_all_tags(self, run_info):
        parent_path, tag_files = self._get_run_files(run_info, "tag")
        return [self._get_tag_from_file(parent_path, tag_file) for tag_file in tag_files]

    def _list_run_infos(self, experiment_id, view_type):
        self._check_root_dir()
        if not self._has_experiment(experiment_id):
            return []
        experiment_dir = self._get_experiment_path(experiment_id, assert_exists=True)
        run_dirs = list_all(
            experiment_dir,
            filter_func=lambda x: (
                all(
                    os.path.basename(os.path.normpath(x)) != reservedFolderName
                    for reservedFolderName in FileStore.RESERVED_EXPERIMENT_FOLDERS
                )
                and os.path.isdir(x)
            ),
            full_path=True,
        )
        run_infos = []
        for r_dir in run_dirs:
            try:
                # trap and warn known issues, will raise unexpected exceptions to caller
                run_info = self._get_run_info_from_dir(r_dir)
                if run_info.experiment_id != experiment_id:
                    logging.warning(
                        "Wrong experiment ID (%s) recorded for run '%s'. "
                        "It should be %s. Run will be ignored.",
                        str(run_info.experiment_id),
                        str(run_info.run_id),
                        str(experiment_id),
                        exc_info=True,
                    )
                    continue
                if LifecycleStage.matches_view_type(view_type, run_info.lifecycle_stage):
                    run_infos.append(run_info)
            except MissingConfigException as rnfe:
                # trap malformed run exception and log
                # this is at debug level because if the same store is used for
                # artifact storage, it's common the folder is not a run folder
                r_id = os.path.basename(r_dir)
                logging.debug(
                    "Malformed run '%s'. Detailed error %s",
                    r_id,
                    str(rnfe),
                    exc_info=True,
                )
        return run_infos

    def _search_runs(
        self,
        experiment_ids,
        filter_string,
        run_view_type,
        max_results,
        order_by,
        page_token,
    ):
        if max_results > SEARCH_MAX_RESULTS_THRESHOLD:
            raise MlflowException(
                "Invalid value for request parameter max_results. It must be at "
                f"most {SEARCH_MAX_RESULTS_THRESHOLD}, but got value {max_results}",
                databricks_pb2.INVALID_PARAMETER_VALUE,
            )
        runs = []
        for experiment_id in experiment_ids:
            run_infos = self._list_run_infos(experiment_id, run_view_type)
            runs.extend(self._get_run_from_info(r) for r in run_infos)
        filtered = SearchUtils.filter(runs, filter_string)
        sorted_runs = SearchUtils.sort(filtered, order_by)
        runs, next_page_token = SearchUtils.paginate(sorted_runs, page_token, max_results)
        return runs, next_page_token

    def log_metric(self, run_id: str, metric: Metric):
        _validate_run_id(run_id)
        _validate_metric(metric.key, metric.value, metric.timestamp, metric.step)
        run_info = self._get_run_info(run_id)
        check_run_is_active(run_info)
        self._log_run_metric(run_info, metric)
        if metric.model_id is not None:
            self._log_model_metric(
                experiment_id=run_info.experiment_id,
                model_id=metric.model_id,
                run_id=run_id,
                metric=metric,
            )

    def _log_run_metric(self, run_info, metric):
        metric_path = self._get_metric_path(run_info.experiment_id, run_info.run_id, metric.key)
        make_containing_dirs(metric_path)
        if metric.dataset_name is not None and metric.dataset_digest is not None:
            append_to(
                metric_path,
                f"{metric.timestamp} {metric.value} {metric.step} {metric.dataset_name} "
                f"{metric.dataset_digest}\n",
            )
        else:
            append_to(metric_path, f"{metric.timestamp} {metric.value} {metric.step}\n")

    def _log_model_metric(self, experiment_id: str, model_id: str, run_id: str, metric: Metric):
        metric_path = self._get_model_metric_path(
            experiment_id=experiment_id, model_id=model_id, metric_key=metric.key
        )
        make_containing_dirs(metric_path)
        if metric.dataset_name is not None and metric.dataset_digest is not None:
            append_to(
                metric_path,
                f"{metric.timestamp} {metric.value} {metric.step} {run_id} {metric.dataset_name} "
                f"{metric.dataset_digest}\n",
            )
        else:
            append_to(metric_path, f"{metric.timestamp} {metric.value} {metric.step} {run_id}\n")

    def _writeable_value(self, tag_value):
        if tag_value is None:
            return ""
        elif is_string_type(tag_value):
            return tag_value
        else:
            return str(tag_value)

    def log_param(self, run_id, param):
        _validate_run_id(run_id)
        param = _validate_param(param.key, param.value)
        run_info = self._get_run_info(run_id)
        check_run_is_active(run_info)
        self._log_run_param(run_info, param)

    def _log_run_param(self, run_info, param):
        param_path = self._get_param_path(run_info.experiment_id, run_info.run_id, param.key)
        writeable_param_value = self._writeable_value(param.value)
        if os.path.exists(param_path):
            self._validate_new_param_value(
                param_path=param_path,
                param_key=param.key,
                run_id=run_info.run_id,
                new_value=writeable_param_value,
            )
        make_containing_dirs(param_path)
        write_to(param_path, writeable_param_value)

    def _validate_new_param_value(self, param_path, param_key, run_id, new_value):
        """
        When logging a parameter with a key that already exists, this function is used to
        enforce immutability by verifying that the specified parameter value matches the existing
        value.
        :raises: py:class:`mlflow.exceptions.MlflowException` if the specified new parameter value
                 does not match the existing parameter value.
        """
        with open(param_path) as param_file:
            current_value = param_file.read()
        if current_value != new_value:
            raise MlflowException(
                f"Changing param values is not allowed. Param with key='{param_key}' was already"
                f" logged with value='{current_value}' for run ID='{run_id}'. Attempted logging"
                f" new value '{new_value}'.",
                databricks_pb2.INVALID_PARAMETER_VALUE,
            )

    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
        """
        _validate_experiment_tag(tag.key, tag.value)
        experiment = self.get_experiment(experiment_id)
        if experiment.lifecycle_stage != LifecycleStage.ACTIVE:
            raise MlflowException(
                f"The experiment {experiment.experiment_id} must be in the 'active' "
                "lifecycle_stage to set tags",
                error_code=databricks_pb2.INVALID_PARAMETER_VALUE,
            )
        tag_path = self._get_experiment_tag_path(experiment_id, tag.key)
        make_containing_dirs(tag_path)
        write_to(tag_path, self._writeable_value(tag.value))

    def delete_experiment_tag(self, experiment_id, key):
        """
        Delete a tag from the specified experiment

        Args:
            experiment_id: String ID of the experiment
            key: String name of the tag to be deleted
        """
        experiment = self.get_experiment(experiment_id)
        if experiment.lifecycle_stage != LifecycleStage.ACTIVE:
            raise MlflowException(
                f"The experiment {experiment.experiment_id} must be in the 'active' "
                "lifecycle_stage to delete tags",
                error_code=databricks_pb2.INVALID_PARAMETER_VALUE,
            )
        tag_path = self._get_experiment_tag_path(experiment_id, key)
        if not exists(tag_path):
            raise MlflowException(
                f"No tag with name: {key} in experiment with id {experiment_id}",
                error_code=RESOURCE_DOES_NOT_EXIST,
            )
        os.remove(tag_path)

    def set_tag(self, run_id, tag):
        _validate_run_id(run_id)
        _validate_tag_name(tag.key)
        run_info = self._get_run_info(run_id)
        check_run_is_active(run_info)
        self._set_run_tag(run_info, tag)
        if tag.key == MLFLOW_RUN_NAME:
            run_status = RunStatus.from_string(run_info.status)
            self.update_run_info(run_id, run_status, run_info.end_time, tag.value)

    def _set_run_tag(self, run_info, tag):
        tag_path = self._get_tag_path(run_info.experiment_id, run_info.run_id, tag.key)
        make_containing_dirs(tag_path)
        # Don't add trailing newline
        write_to(tag_path, self._writeable_value(tag.value))

    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.
        """
        _validate_run_id(run_id)
        run_info = self._get_run_info(run_id)
        check_run_is_active(run_info)
        tag_path = self._get_tag_path(run_info.experiment_id, run_id, key)
        if not exists(tag_path):
            raise MlflowException(
                f"No tag with name: {key} in run with id {run_id}",
                error_code=RESOURCE_DOES_NOT_EXIST,
            )
        os.remove(tag_path)

    def _overwrite_run_info(self, run_info, deleted_time=None):
        run_dir = self._get_run_dir(run_info.experiment_id, run_info.run_id)
        run_info_dict = _make_persisted_run_info_dict(run_info)
        if deleted_time is not None:
            run_info_dict["deleted_time"] = deleted_time
        write_yaml(run_dir, FileStore.META_DATA_FILE_NAME, run_info_dict, overwrite=True)

    def log_batch(self, run_id, metrics, params, tags):
        _validate_run_id(run_id)
        metrics, params, tags = _validate_batch_log_data(metrics, params, tags)
        _validate_batch_log_limits(metrics, params, tags)
        _validate_param_keys_unique(params)
        run_info = self._get_run_info(run_id)
        check_run_is_active(run_info)
        try:
            for param in params:
                self._log_run_param(run_info, param)
            for metric in metrics:
                self._log_run_metric(run_info, metric)
                if metric.model_id is not None:
                    self._log_model_metric(
                        experiment_id=run_info.experiment_id,
                        model_id=metric.model_id,
                        run_id=run_id,
                        metric=metric,
                    )
            for tag in tags:
                # NB: If the tag run name value is set, update the run info to assure
                # synchronization.
                if tag.key == MLFLOW_RUN_NAME:
                    run_status = RunStatus.from_string(run_info.status)
                    self.update_run_info(run_id, run_status, run_info.end_time, tag.value)
                self._set_run_tag(run_info, tag)
        except Exception as e:
            raise MlflowException(e, INTERNAL_ERROR)

    def record_logged_model(self, run_id, mlflow_model):
        from mlflow.models import Model

        if not isinstance(mlflow_model, Model):
            raise TypeError(
                f"Argument 'mlflow_model' should be mlflow.models.Model, got '{type(mlflow_model)}'"
            )
        _validate_run_id(run_id)
        run_info = self._get_run_info(run_id)
        check_run_is_active(run_info)
        model_dict = mlflow_model.get_tags_dict()
        run_info = self._get_run_info(run_id)
        path = self._get_tag_path(run_info.experiment_id, run_info.run_id, MLFLOW_LOGGED_MODELS)
        if os.path.exists(path):
            with open(path) as f:
                model_list = json.loads(f.read())
        else:
            model_list = []
        tag = RunTag(MLFLOW_LOGGED_MODELS, json.dumps(model_list + [model_dict]))

        try:
            self._set_run_tag(run_info, tag)
        except Exception as e:
            raise MlflowException(e, INTERNAL_ERROR)

    def log_inputs(
        self,
        run_id: str,
        datasets: list[DatasetInput] | None = None,
        models: list[LoggedModelInput] | None = None,
    ):
        """
        Log inputs, such as datasets and models, 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
                as inputs to the run.

        Returns:
            None.
        """
        _validate_run_id(run_id)
        run_info = self._get_run_info(run_id)
        check_run_is_active(run_info)

        if datasets is None and models is None:
            return

        experiment_dir = self._get_experiment_path(run_info.experiment_id, assert_exists=True)
        run_dir = self._get_run_dir(run_info.experiment_id, run_id)

        for dataset_input in datasets or []:
            dataset = dataset_input.dataset
            dataset_id = FileStore._get_dataset_id(
                dataset_name=dataset.name, dataset_digest=dataset.digest
            )
            dataset_dir = os.path.join(experiment_dir, FileStore.DATASETS_FOLDER_NAME, dataset_id)
            if not os.path.exists(dataset_dir):
                os.makedirs(dataset_dir, exist_ok=True)
                write_yaml(dataset_dir, FileStore.META_DATA_FILE_NAME, dict(dataset))

            input_id = FileStore._get_dataset_input_id(dataset_id=dataset_id, run_id=run_id)
            input_dir = os.path.join(run_dir, FileStore.INPUTS_FOLDER_NAME, input_id)
            if not os.path.exists(input_dir):
                os.makedirs(input_dir, exist_ok=True)
                fs_input = FileStore._FileStoreInput(
                    source_type=InputVertexType.DATASET,
                    source_id=dataset_id,
                    destination_type=InputVertexType.RUN,
                    destination_id=run_id,
                    tags={tag.key: tag.value for tag in dataset_input.tags},
                )
                fs_input.write_yaml(input_dir, FileStore.META_DATA_FILE_NAME)

        for model_input in models or []:
            model_id = model_input.model_id
            input_id = FileStore._get_model_input_id(model_id=model_id, run_id=run_id)
            input_dir = os.path.join(run_dir, FileStore.INPUTS_FOLDER_NAME, input_id)
            if not os.path.exists(input_dir):
                os.makedirs(input_dir, exist_ok=True)
                fs_input = FileStore._FileStoreInput(
                    source_type=InputVertexType.MODEL,
                    source_id=model_id,
                    destination_type=InputVertexType.RUN,
                    destination_id=run_id,
                    tags={},
                )
                fs_input.write_yaml(input_dir, FileStore.META_DATA_FILE_NAME)

    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.
        """
        _validate_run_id(run_id)
        run_info = self._get_run_info(run_id)
        check_run_is_active(run_info)

        if models is None:
            return

        run_dir = self._get_run_dir(run_info.experiment_id, run_id)

        for model_output in models:
            model_id = model_output.model_id
            output_dir = os.path.join(run_dir, FileStore.OUTPUTS_FOLDER_NAME, model_id)
            if not os.path.exists(output_dir):
                os.makedirs(output_dir, exist_ok=True)
                fs_output = FileStore._FileStoreOutput(
                    source_type=OutputVertexType.RUN_OUTPUT,
                    source_id=model_id,
                    destination_type=OutputVertexType.MODEL_OUTPUT,
                    destination_id=run_id,
                    tags={},
                    step=model_output.step,
                )
                fs_output.write_yaml(output_dir, FileStore.META_DATA_FILE_NAME)

    @staticmethod
    def _get_dataset_id(dataset_name: str, dataset_digest: str) -> str:
        md5 = hashlib.md5(dataset_name.encode("utf-8"), usedforsecurity=False)
        md5.update(dataset_digest.encode("utf-8"))
        return md5.hexdigest()

    @staticmethod
    def _get_dataset_input_id(dataset_id: str, run_id: str) -> str:
        md5 = hashlib.md5(dataset_id.encode("utf-8"), usedforsecurity=False)
        md5.update(run_id.encode("utf-8"))
        return md5.hexdigest()

    @staticmethod
    def _get_model_input_id(model_id: str, run_id: str) -> str:
        md5 = hashlib.md5(model_id.encode("utf-8"), usedforsecurity=False)
        md5.update(run_id.encode("utf-8"))
        return md5.hexdigest()

    class _FileStoreInput(NamedTuple):
        source_type: int
        source_id: str
        destination_type: int
        destination_id: str
        tags: dict[str, str]

        def write_yaml(self, root: str, file_name: str):
            dict_for_yaml = {
                "source_type": InputVertexType.Name(self.source_type),
                "source_id": self.source_id,
                "destination_type": InputVertexType.Name(self.destination_type),
                "destination_id": self.source_id,
                "tags": self.tags,
            }
            write_yaml(root, file_name, dict_for_yaml)

        @classmethod
        def from_yaml(cls, root, file_name):
            dict_from_yaml = FileStore._read_yaml(root, file_name)
            return cls(
                source_type=InputVertexType.Value(dict_from_yaml["source_type"]),
                source_id=dict_from_yaml["source_id"],
                destination_type=InputVertexType.Value(dict_from_yaml["destination_type"]),
                destination_id=dict_from_yaml["destination_id"],
                tags=dict_from_yaml["tags"],
            )

    class _FileStoreOutput(NamedTuple):
        source_type: int
        source_id: str
        destination_type: int
        destination_id: str
        tags: dict[str, str]
        step: int

        def write_yaml(self, root: str, file_name: str):
            dict_for_yaml = {
                "source_type": OutputVertexType.Name(self.source_type),
                "source_id": self.source_id,
                "destination_type": OutputVertexType.Name(self.destination_type),
                "destination_id": self.source_id,
                "tags": self.tags,
                "step": self.step,
            }
            write_yaml(root, file_name, dict_for_yaml)

        @classmethod
        def from_yaml(cls, root, file_name):
            dict_from_yaml = FileStore._read_yaml(root, file_name)
            return cls(
                source_type=OutputVertexType.Value(dict_from_yaml["source_type"]),
                source_id=dict_from_yaml["source_id"],
                destination_type=OutputVertexType.Value(dict_from_yaml["destination_type"]),
                destination_id=dict_from_yaml["destination_id"],
                tags=dict_from_yaml["tags"],
                step=dict_from_yaml["step"],
            )

    def _get_all_inputs(self, run_info: RunInfo) -> RunInputs:
        run_dir = self._get_run_dir(run_info.experiment_id, run_info.run_id)
        inputs_parent_path = os.path.join(run_dir, FileStore.INPUTS_FOLDER_NAME)
        if not os.path.exists(inputs_parent_path):
            return RunInputs(dataset_inputs=[], model_inputs=[])

        experiment_dir = self._get_experiment_path(run_info.experiment_id, assert_exists=True)
        dataset_inputs = self._get_dataset_inputs(run_info, inputs_parent_path, experiment_dir)
        model_inputs = self._get_model_inputs(inputs_parent_path, experiment_dir)
        return RunInputs(dataset_inputs=dataset_inputs, model_inputs=model_inputs)

    def _get_dataset_inputs(
        self, run_info: RunInfo, inputs_parent_path: str, experiment_dir_path: str
    ) -> list[DatasetInput]:
        datasets_parent_path = os.path.join(experiment_dir_path, FileStore.DATASETS_FOLDER_NAME)
        if not os.path.exists(datasets_parent_path):
            return []

        dataset_dirs = os.listdir(datasets_parent_path)
        dataset_inputs = []
        for input_dir in os.listdir(inputs_parent_path):
            input_dir_full_path = os.path.join(inputs_parent_path, input_dir)
            fs_input = FileStore._FileStoreInput.from_yaml(
                input_dir_full_path, FileStore.META_DATA_FILE_NAME
            )
            if fs_input.source_type != InputVertexType.DATASET:
                continue

            matching_dataset_dirs = [d for d in dataset_dirs if d == fs_input.source_id]
            if not matching_dataset_dirs:
                logging.warning(
                    f"Failed to find dataset with ID '{fs_input.source_id}' referenced as an input"
                    f" of the run with ID '{run_info.run_id}'. Skipping."
                )
                continue
            elif len(matching_dataset_dirs) > 1:
                logging.warning(
                    f"Found multiple datasets with ID '{fs_input.source_id}'. Using the first one."
                )

            dataset_dir = matching_dataset_dirs[0]
            dataset = FileStore._get_dataset_from_dir(datasets_parent_path, dataset_dir)
            dataset_input = DatasetInput(
                dataset=dataset,
                tags=[InputTag(key=key, value=value) for key, value in fs_input.tags.items()],
            )
            dataset_inputs.append(dataset_input)

        return dataset_inputs

    def _get_model_inputs(
        self, inputs_parent_path: str, experiment_dir_path: str
    ) -> list[LoggedModelInput]:
        model_inputs = []
        for input_dir in os.listdir(inputs_parent_path):
            input_dir_full_path = os.path.join(inputs_parent_path, input_dir)
            fs_input = FileStore._FileStoreInput.from_yaml(
                input_dir_full_path, FileStore.META_DATA_FILE_NAME
            )
            if fs_input.source_type != InputVertexType.MODEL:
                continue

            model_input = LoggedModelInput(model_id=fs_input.source_id)
            model_inputs.append(model_input)

        return model_inputs

    def _get_all_outputs(self, run_info: RunInfo) -> RunOutputs:
        run_dir = self._get_run_dir(run_info.experiment_id, run_info.run_id)
        outputs_parent_path = os.path.join(run_dir, FileStore.OUTPUTS_FOLDER_NAME)
        if not os.path.exists(outputs_parent_path):
            return RunOutputs(model_outputs=[])

        experiment_dir = self._get_experiment_path(run_info.experiment_id, assert_exists=True)
        model_outputs = self._get_model_outputs(outputs_parent_path, experiment_dir)
        return RunOutputs(model_outputs=model_outputs)

    def _get_model_outputs(
        self, outputs_parent_path: str, experiment_dir: str
    ) -> list[LoggedModelOutput]:
        model_outputs = []
        for output_dir in os.listdir(outputs_parent_path):
            output_dir_full_path = os.path.join(outputs_parent_path, output_dir)
            fs_output = FileStore._FileStoreOutput.from_yaml(
                output_dir_full_path, FileStore.META_DATA_FILE_NAME
            )
            if fs_output.destination_type != OutputVertexType.MODEL_OUTPUT:
                continue

            model_output = LoggedModelOutput(model_id=fs_output.destination_id, step=fs_output.step)
            model_outputs.append(model_output)

        return model_outputs

    def _search_datasets(self, experiment_ids) -> list[_DatasetSummary]:
        """
        Return all dataset summaries associated to the given experiments.

        Args:
            experiment_ids: List of experiment ids to scope the search

        Returns:
            A List of :py:class:`mlflow.entities.DatasetSummary` entities.

        """

        @dataclass(frozen=True)
        class _SummaryTuple:
            experiment_id: str
            name: str
            digest: str
            context: str

        MAX_DATASET_SUMMARIES_RESULTS = 1000
        summaries = set()
        for experiment_id in experiment_ids:
            experiment_dir = self._get_experiment_path(experiment_id, assert_exists=True)
            run_dirs = list_all(
                experiment_dir,
                filter_func=lambda x: (
                    all(
                        os.path.basename(os.path.normpath(x)) != reservedFolderName
                        for reservedFolderName in FileStore.RESERVED_EXPERIMENT_FOLDERS
                    )
                    and os.path.isdir(x)
                ),
                full_path=True,
            )
            for run_dir in run_dirs:
                run_info = self._get_run_info_from_dir(run_dir)
                run_inputs = self._get_all_inputs(run_info)
                for dataset_input in run_inputs.dataset_inputs:
                    context = None
                    for input_tag in dataset_input.tags:
                        if input_tag.key == MLFLOW_DATASET_CONTEXT:
                            context = input_tag.value
                            break
                    dataset = dataset_input.dataset
                    summaries.add(
                        _SummaryTuple(experiment_id, dataset.name, dataset.digest, context)
                    )
                    # If we reached MAX_DATASET_SUMMARIES_RESULTS entries, then return right away.
                    if len(summaries) == MAX_DATASET_SUMMARIES_RESULTS:
                        return [
                            _DatasetSummary(
                                experiment_id=summary.experiment_id,
                                name=summary.name,
                                digest=summary.digest,
                                context=summary.context,
                            )
                            for summary in summaries
                        ]

        return [
            _DatasetSummary(
                experiment_id=summary.experiment_id,
                name=summary.name,
                digest=summary.digest,
                context=summary.context,
            )
            for summary in summaries
        ]

    @staticmethod
    def _get_dataset_from_dir(parent_path, dataset_dir) -> Dataset:
        dataset_dict = FileStore._read_yaml(
            os.path.join(parent_path, dataset_dir), FileStore.META_DATA_FILE_NAME
        )
        return Dataset.from_dictionary(dataset_dict)

    @staticmethod
    def _read_yaml(root, file_name, retries=2):
        """
        Read data from yaml file and return as dictionary, retrying up to
        a specified number of times if the file contents are unexpectedly
        empty due to a concurrent write.

        Args:
            root: Directory name.
            file_name: File name. Expects to have '.yaml' extension.
            retries: The number of times to retry for unexpected empty content.

        Returns:
            Data in yaml file as dictionary.
        """

        def _read_helper(root, file_name, attempts_remaining=2):
            result = read_yaml(root, file_name)
            if result is not None or attempts_remaining == 0:
                return result
            else:
                time.sleep(0.1 * (3 - attempts_remaining))
                return _read_helper(root, file_name, attempts_remaining - 1)

        return _read_helper(root, file_name, attempts_remaining=retries)

    def _get_traces_artifact_dir(self, experiment_id, trace_id):
        return append_to_uri_path(
            self.get_experiment(experiment_id).artifact_location,
            FileStore.TRACES_FOLDER_NAME,
            trace_id,
            FileStore.ARTIFACTS_FOLDER_NAME,
        )

    def _save_trace_info(self, trace_info: TraceInfo, trace_dir, overwrite=False):
        """
        TraceInfo is saved into `traces` folder under the experiment, each trace
        is saved in the folder named by its trace_id.
        `request_metadata` and `tags` folder store their key-value pairs such that each
        key is the file name, and value is written as the string value.
        Detailed directories structure is as below:
        | - experiment_id
        |   - traces
        |     - trace_id1
        |       - trace_info.yaml
        |       - request_metadata
        |         - key
        |       - tags
        |     - trace_id2
        |     - ...
        |   - run_id1 ...
        |   - run_id2 ...
        """
        # Save basic trace info to TRACE_INFO_FILE_NAME
        trace_info_dict = self._convert_trace_info_to_dict(trace_info)
        write_yaml(
            trace_dir,
            FileStore.TRACE_INFO_FILE_NAME,
            trace_info_dict,
            overwrite=overwrite,
        )
        # Save trace_metadata to its own folder
        self._write_dict_to_trace_sub_folder(
            trace_dir,
            FileStore.TRACE_TRACE_METADATA_FOLDER_NAME,
            trace_info.trace_metadata,
        )
        # Save tags to its own folder
        self._write_dict_to_trace_sub_folder(
            trace_dir, FileStore.TRACE_TAGS_FOLDER_NAME, trace_info.tags
        )
        # Save assessments to its own folder
        for assessment in trace_info.assessments:
            self.create_assessment(assessment)

    def _convert_trace_info_to_dict(self, trace_info: TraceInfo):
        """
        Convert trace info to a dictionary for persistence.
        Drop request_metadata and tags as they're saved into separate files.
        """
        trace_info_dict = trace_info.to_dict()
        trace_info_dict.pop("trace_metadata", None)
        trace_info_dict.pop("tags", None)
        return trace_info_dict

    def _write_dict_to_trace_sub_folder(self, trace_dir, sub_folder, dictionary):
        mkdir(trace_dir, sub_folder)
        for key, value in dictionary.items():
            # always validate as tag name to make sure the file name is valid
            _validate_tag_name(key)
            tag_path = os.path.join(trace_dir, sub_folder, key)
            # value are written as strings
            write_to(tag_path, self._writeable_value(value))

    def _get_dict_from_trace_sub_folder(self, trace_dir, sub_folder):
        parent_path, files = self._get_resource_files(trace_dir, sub_folder)
        dictionary = {}
        for file_name in files:
            _validate_tag_name(file_name)
            value = read_file(parent_path, file_name)
            dictionary[file_name] = value
        return dictionary

    def start_trace(self, trace_info: TraceInfo) -> TraceInfo:
        """
        Create a trace using the V3 API format with a complete Trace object.

        Args:
            trace_info: The TraceInfo object to create in the backend.

        Returns:
            The created TraceInfo object from the backend.
        """
        _validate_experiment_id(trace_info.experiment_id)
        experiment_dir = self._get_experiment_path(
            trace_info.experiment_id, view_type=ViewType.ACTIVE_ONLY, assert_exists=True
        )

        # Create traces directory structure
        mkdir(experiment_dir, FileStore.TRACES_FOLDER_NAME)
        traces_dir = os.path.join(experiment_dir, FileStore.TRACES_FOLDER_NAME)
        mkdir(traces_dir, trace_info.trace_id)
        trace_dir = os.path.join(traces_dir, trace_info.trace_id)

        # Add artifact location to tags
        artifact_uri = self._get_traces_artifact_dir(trace_info.experiment_id, trace_info.trace_id)
        tags = dict(trace_info.tags)
        tags[MLFLOW_ARTIFACT_LOCATION] = artifact_uri

        # Create updated TraceInfo with artifact location tag
        trace_info.tags.update(tags)
        self._save_trace_info(trace_info, trace_dir)
        return trace_info

    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``.
        """
        return self._get_trace_info_and_dir(trace_id)[0]

    def _get_trace_info_and_dir(self, trace_id: str) -> tuple[TraceInfo, str]:
        trace_dir = self._find_trace_dir(trace_id, assert_exists=True)
        trace_info = self._get_trace_info_from_dir(trace_dir)
        if trace_info and trace_info.trace_id != trace_id:
            raise MlflowException(
                f"Trace with ID '{trace_id}' metadata is in invalid state.",
                databricks_pb2.INVALID_STATE,
            )
        return trace_info, trace_dir

    def _find_trace_dir(self, trace_id, assert_exists=False):
        self._check_root_dir()
        all_experiments = self._get_active_experiments(True) + self._get_deleted_experiments(True)
        for experiment_dir in all_experiments:
            traces_dir = os.path.join(experiment_dir, FileStore.TRACES_FOLDER_NAME)
            if exists(traces_dir):
                if traces := find(traces_dir, trace_id, full_path=True):
                    return traces[0]
        if assert_exists:
            raise MlflowException(
                f"Trace with ID '{trace_id}' not found",
                RESOURCE_DOES_NOT_EXIST,
            )

    def _get_trace_info_from_dir(self, trace_dir) -> TraceInfo | None:
        if not os.path.exists(os.path.join(trace_dir, FileStore.TRACE_INFO_FILE_NAME)):
            return None
        trace_info_dict = FileStore._read_yaml(trace_dir, FileStore.TRACE_INFO_FILE_NAME)
        trace_info = TraceInfo.from_dict(trace_info_dict)
        trace_info.trace_metadata = self._get_dict_from_trace_sub_folder(
            trace_dir, FileStore.TRACE_TRACE_METADATA_FOLDER_NAME
        )
        trace_info.tags = self._get_dict_from_trace_sub_folder(
            trace_dir, FileStore.TRACE_TAGS_FOLDER_NAME
        )
        trace_info.assessments = self._load_assessments(trace_info.trace_id)
        return trace_info

    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.
        """
        trace_dir = self._find_trace_dir(trace_id, assert_exists=True)
        self._write_dict_to_trace_sub_folder(
            trace_dir, FileStore.TRACE_TAGS_FOLDER_NAME, {key: value}
        )

    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.
        """
        _validate_tag_name(key)
        trace_dir = self._find_trace_dir(trace_id, assert_exists=True)
        tag_path = os.path.join(trace_dir, FileStore.TRACE_TAGS_FOLDER_NAME, key)
        if not exists(tag_path):
            raise MlflowException(
                f"No tag with name: {key} in trace with ID {trace_id}.",
                RESOURCE_DOES_NOT_EXIST,
            )
        os.remove(tag_path)

    def _get_assessments_dir(self, trace_id: str) -> str:
        trace_dir = self._find_trace_dir(trace_id, assert_exists=True)
        return os.path.join(trace_dir, FileStore.ASSESSMENTS_FOLDER_NAME)

    def _get_assessment_path(self, trace_id: str, assessment_id: str) -> str:
        assessments_dir = self._get_assessments_dir(trace_id)
        return os.path.join(assessments_dir, f"{assessment_id}.yaml")

    def _save_assessment(self, assessment: Assessment) -> None:
        assessment_path = self._get_assessment_path(assessment.trace_id, assessment.assessment_id)
        make_containing_dirs(assessment_path)

        assessment_dict = assessment.to_dictionary()
        write_yaml(
            root=os.path.dirname(assessment_path),
            file_name=os.path.basename(assessment_path),
            data=assessment_dict,
            overwrite=True,
        )

    def _load_assessments(self, trace_id: str) -> list[Assessment]:
        assessments_dir = self._get_assessments_dir(trace_id)
        if not exists(assessments_dir):
            return []
        assessment_paths = os.listdir(assessments_dir)
        return [
            self._load_assessment(trace_id, assessment_path.split(".")[0])
            for assessment_path in assessment_paths
        ]

    def _load_assessment(self, trace_id: str, assessment_id: str) -> Assessment:
        assessment_path = self._get_assessment_path(trace_id, assessment_id)

        if not exists(assessment_path):
            raise MlflowException(
                f"Assessment with ID '{assessment_id}' not found for trace '{trace_id}'",
                RESOURCE_DOES_NOT_EXIST,
            )

        try:
            assessment_dict = FileStore._read_yaml(
                root=os.path.dirname(assessment_path), file_name=os.path.basename(assessment_path)
            )
            return Assessment.from_dictionary(assessment_dict)
        except Exception as e:
            raise MlflowException(
                f"Failed to load assessment with ID '{assessment_id}' for trace '{trace_id}': {e}",
                INTERNAL_ERROR,
            ) from e

    def get_assessment(self, trace_id: str, assessment_id: str) -> Assessment:
        """
        Retrieves a specific assessment associated with a trace from the file store.

        Args:
            trace_id: The unique identifier of the trace containing the assessment.
            assessment_id: The unique identifier of the assessment to retrieve.

        Returns:
            Assessment: The requested assessment object (either Expectation or Feedback).

        Raises:
            MlflowException: If the trace_id is not found, if no assessment with the
                specified assessment_id exists for the trace, or if the stored
                assessment data cannot be deserialized.
        """

        return self._load_assessment(trace_id, assessment_id)

    def create_assessment(self, assessment: Assessment) -> Assessment:
        """
        Creates a new assessment record associated with a specific trace.

        Args:
            assessment: The assessment object to create. The assessment will be modified
                    in-place to include the generated assessment_id and timestamps.

        Returns:
            Assessment: The input assessment object updated with backend-generated metadata.

        Raises:
            MlflowException: If the trace doesn't exist or there's an error saving the assessment.
        """

        assessment_id = generate_assessment_id()
        creation_timestamp = int(time.time() * 1000)

        assessment.assessment_id = assessment_id
        assessment.create_time_ms = creation_timestamp
        assessment.last_update_time_ms = creation_timestamp
        assessment.valid = True

        if assessment.overrides:
            original_assessment = self.get_assessment(assessment.trace_id, assessment.overrides)
            original_assessment.valid = False
            self._save_assessment(original_assessment)

        self._save_assessment(assessment)
        return assessment

    def update_assessment(
        self,
        trace_id: str,
        assessment_id: str,
        name: str | None = None,
        expectation: Expectation | None = None,
        feedback: Feedback | None = None,
        rationale: str | None = None,
        metadata: dict[str, str] | None = None,
    ) -> Assessment:
        """
        Updates an existing assessment with new values while preserving immutable fields.

        `source` and `span_id` are immutable and cannot be changed.
        The last_update_time_ms will always be updated to the current timestamp.
        Metadata will be merged with the new metadata taking precedence.

        Args:
            trace_id: The unique identifier of the trace containing the assessment.
            assessment_id: The unique identifier of the assessment to update.
            name: The updated name of the assessment. If None, preserves existing name.
            expectation: Updated expectation value for expectation assessments.
            feedback: Updated feedback value for feedback assessments.
            rationale: Updated rationale text. If None, preserves existing rationale.
            metadata: Updated metadata dict. Will be merged with existing metadata.

        Returns:
            Assessment: The updated assessment object with new last_update_time_ms.

        Raises:
            MlflowException: If the assessment doesn't exist, if immutable fields have
                            changed, or if there's an error saving the assessment.
        """
        existing_assessment = self.get_assessment(trace_id, assessment_id)

        if expectation is not None and feedback is not None:
            raise MlflowException.invalid_parameter_value(
                "Cannot specify both `expectation` and `feedback` parameters."
            )

        if expectation is not None and not isinstance(existing_assessment, Expectation):
            raise MlflowException.invalid_parameter_value(
                "Cannot update expectation value on a Feedback assessment."
            )

        if feedback is not None and not isinstance(existing_assessment, Feedback):
            raise MlflowException.invalid_parameter_value(
                "Cannot update feedback value on an Expectation assessment."
            )

        merged_metadata = None
        if existing_assessment.metadata or metadata:
            merged_metadata = (existing_assessment.metadata or {}).copy()
            if metadata:
                merged_metadata.update(metadata)

        updated_timestamp = int(time.time() * 1000)

        if isinstance(existing_assessment, Expectation):
            new_value = expectation.value if expectation is not None else existing_assessment.value

            updated_assessment = Expectation(
                name=name if name is not None else existing_assessment.name,
                value=new_value,
                source=existing_assessment.source,
                trace_id=trace_id,
                metadata=merged_metadata,
                span_id=existing_assessment.span_id,
                create_time_ms=existing_assessment.create_time_ms,
                last_update_time_ms=updated_timestamp,
            )
        else:
            if feedback is not None:
                new_value = feedback.value
                new_error = feedback.error
            else:
                new_value = existing_assessment.value
                new_error = existing_assessment.error

            updated_assessment = Feedback(
                name=name if name is not None else existing_assessment.name,
                value=new_value,
                error=new_error,
                source=existing_assessment.source,
                trace_id=trace_id,
                metadata=merged_metadata,
                span_id=existing_assessment.span_id,
                create_time_ms=existing_assessment.create_time_ms,
                last_update_time_ms=updated_timestamp,
                rationale=rationale if rationale is not None else existing_assessment.rationale,
            )

        updated_assessment.assessment_id = existing_assessment.assessment_id
        updated_assessment.valid = existing_assessment.valid
        updated_assessment.overrides = existing_assessment.overrides

        if hasattr(existing_assessment, "run_id"):
            updated_assessment.run_id = existing_assessment.run_id

        self._save_assessment(updated_assessment)

        return updated_assessment

    def delete_assessment(self, trace_id: str, assessment_id: str) -> None:
        """
        Delete an assessment from a trace.

        If the deleted assessment was overriding another assessment, the overridden
        assessment will be restored to valid=True.

        Args:
            trace_id: The ID of the trace containing the assessment.
            assessment_id: The ID of the assessment to delete.

        Raises:
            MlflowException: If the trace_id is not found or deletion fails.
        """
        # First validate that the trace exists (this will raise if trace not found)
        self._find_trace_dir(trace_id, assert_exists=True)

        assessment_path = self._get_assessment_path(trace_id, assessment_id)

        # Early return if assessment doesn't exist (idempotent behavior)
        if not exists(assessment_path):
            return

        # Get override info before deletion
        overrides_assessment_id = None
        try:
            assessment_to_delete = self._load_assessment(trace_id, assessment_id)
            overrides_assessment_id = assessment_to_delete.overrides
        except Exception:
            pass

        assessment_path = self._get_assessment_path(trace_id, assessment_id)
        try:
            os.remove(assessment_path)

            # Clean up empty assessments directory if no more assessments exist
            assessments_dir = self._get_assessments_dir(trace_id)
            if exists(assessments_dir) and not os.listdir(assessments_dir):
                os.rmdir(assessments_dir)

        except OSError as e:
            raise MlflowException(
                f"Failed to delete assessment with ID '{assessment_id}' "
                f"for trace '{trace_id}': {e}",
                INTERNAL_ERROR,
            ) from e

        # If this assessment was overriding another assessment, restore the original
        if overrides_assessment_id:
            try:
                original_assessment = self.get_assessment(trace_id, overrides_assessment_id)
                original_assessment.valid = True
                original_assessment.last_update_time_ms = int(time.time() * 1000)
                self._save_assessment(original_assessment)
            except MlflowException:
                pass

    def _delete_traces(
        self,
        experiment_id: str,
        max_timestamp_millis: int | None = None,
        max_traces: int | None = None,
        trace_ids: list[str] | None = None,
    ) -> int:
        """
        Delete traces based on the specified criteria.

        - Either `max_timestamp_millis` or `trace_ids` must be specified, but not both.
        - `max_traces` can't be specified if `trace_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.
            trace_ids: A set of trace IDs to delete.

        Returns:
            The number of traces deleted.
        """
        experiment_path = self._get_experiment_path(experiment_id, assert_exists=True)
        traces_path = os.path.join(experiment_path, FileStore.TRACES_FOLDER_NAME)
        deleted_traces = 0
        if max_timestamp_millis is not None:
            trace_paths = list_all(traces_path, lambda x: os.path.isdir(x), full_path=True)
            trace_info_and_paths = []
            for trace_path in trace_paths:
                try:
                    trace_info = self._get_trace_info_from_dir(trace_path)
                    if trace_info and trace_info.timestamp_ms <= max_timestamp_millis:
                        trace_info_and_paths.append((trace_info, trace_path))
                except MissingConfigException as e:
                    # trap malformed trace exception and log warning
                    trace_id = os.path.basename(trace_path)
                    _logger.warning(
                        f"Malformed trace with ID '{trace_id}'. Detailed error {e}",
                        exc_info=_logger.isEnabledFor(logging.DEBUG),
                    )
            trace_info_and_paths.sort(key=lambda x: x[0].timestamp_ms)
            # if max_traces is not None then it must > 0
            deleted_traces = min(len(trace_info_and_paths), max_traces or len(trace_info_and_paths))
            trace_info_and_paths = trace_info_and_paths[:deleted_traces]
            for _, trace_path in trace_info_and_paths:
                shutil.rmtree(trace_path)
            return deleted_traces
        if trace_ids:
            for trace_id in trace_ids:
                trace_path = os.path.join(traces_path, trace_id)
                # Do not throw if the trace doesn't exist
                if exists(trace_path):
                    shutil.rmtree(trace_path)
                    deleted_traces += 1
            return deleted_traces

    def search_traces(
        self,
        experiment_ids: list[str] | None = None,
        filter_string: str | None = None,
        max_results: int = SEARCH_TRACES_DEFAULT_MAX_RESULTS,
        order_by: list[str] | None = None,
        page_token: str | None = None,
        model_id: str | None = None,
        locations: list[str] | None = None,
    ) -> tuple[list[TraceInfo], str | None]:
        """
        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. Supported filter keys are `name`,
                           `status`, `timestamp_ms` and `tags`.
            max_results: Maximum number of traces desired.
            order_by: List of order_by clauses. Supported sort key is `timestamp_ms`. By default
                      we sort by timestamp_ms DESC.
            page_token: Token specifying the next page of results. It should be obtained from
                a ``search_traces`` call.
            model_id: If specified, return traces associated with the model ID.
            locations: A list of locations to search over. To search over experiments, provide
                a list of experiment IDs.

        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.
        """
        locations = _resolve_experiment_ids_and_locations(experiment_ids, locations)
        if max_results > SEARCH_MAX_RESULTS_THRESHOLD:
            raise MlflowException(
                "Invalid value for request parameter max_results. It must be at "
                f"most {SEARCH_MAX_RESULTS_THRESHOLD}, but got value {max_results}",
                INVALID_PARAMETER_VALUE,
            )
        traces = []
        for experiment_id in locations:
            trace_infos = self._list_trace_infos(experiment_id)
            traces.extend(trace_infos)
        filtered = SearchTraceUtils.filter(traces, filter_string)
        sorted_traces = SearchTraceUtils.sort(filtered, order_by)
        traces, next_page_token = SearchTraceUtils.paginate(sorted_traces, page_token, max_results)
        return traces, next_page_token

    def _list_trace_infos(self, experiment_id):
        experiment_path = self._get_experiment_path(experiment_id, assert_exists=True)
        traces_path = os.path.join(experiment_path, FileStore.TRACES_FOLDER_NAME)
        if not os.path.exists(traces_path):
            return []
        trace_paths = list_all(traces_path, lambda x: os.path.isdir(x), full_path=True)
        trace_infos = []
        for trace_path in trace_paths:
            try:
                if trace_info := self._get_trace_info_from_dir(trace_path):
                    trace_infos.append(trace_info)
            except MissingConfigException as e:
                # trap malformed trace exception and log warning
                trace_id = os.path.basename(trace_path)
                logging.warning(
                    f"Malformed trace with ID '{trace_id}'. Detailed error {e}",
                    exc_info=_logger.isEnabledFor(logging.DEBUG),
                )
        return trace_infos

    def create_logged_model(
        self,
        experiment_id: str = DEFAULT_EXPERIMENT_ID,
        name: str | None = None,
        source_run_id: str | None = None,
        tags: list[LoggedModelTag] | None = None,
        params: list[LoggedModelParameter] | None = None,
        model_type: str | None = 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.
        """
        _validate_logged_model_name(name)
        experiment = self.get_experiment(experiment_id)
        if experiment is None:
            raise MlflowException(
                f"Could not create model under experiment with ID {experiment_id} - no such "
                "experiment exists." % experiment_id,
                databricks_pb2.RESOURCE_DOES_NOT_EXIST,
            )
        if experiment.lifecycle_stage != LifecycleStage.ACTIVE:
            raise MlflowException(
                f"Could not create model under non-active experiment with ID {experiment_id}.",
                databricks_pb2.INVALID_STATE,
            )
        for param in params or []:
            _validate_param(param.key, param.value)

        name = name or _generate_random_name()
        model_id = f"m-{str(uuid.uuid4()).replace('-', '')}"
        artifact_location = self._get_model_artifact_dir(experiment_id, model_id)
        creation_timestamp = int(time.time() * 1000)
        model = LoggedModel(
            experiment_id=experiment_id,
            model_id=model_id,
            name=name,
            artifact_location=artifact_location,
            creation_timestamp=creation_timestamp,
            last_updated_timestamp=creation_timestamp,
            source_run_id=source_run_id,
            status=LoggedModelStatus.PENDING,
            tags=tags,
            params=params,
            model_type=model_type,
        )

        # Persist model metadata and create directories for logging metrics, tags
        model_dir = self._get_model_dir(experiment_id, model_id)
        mkdir(model_dir)
        model_info_dict: dict[str, Any] = self._make_persisted_model_dict(model)
        model_info_dict["lifecycle_stage"] = LifecycleStage.ACTIVE
        write_yaml(model_dir, FileStore.META_DATA_FILE_NAME, model_info_dict)
        mkdir(model_dir, FileStore.METRICS_FOLDER_NAME)
        mkdir(model_dir, FileStore.PARAMS_FOLDER_NAME)
        self.log_logged_model_params(model_id=model_id, params=params or [])
        self.set_logged_model_tags(model_id=model_id, tags=tags or [])

        return self.get_logged_model(model_id=model_id)

    def log_logged_model_params(self, model_id: str, params: list[LoggedModelParameter]):
        """
        Set parameters on the specified logged model.

        Args:
            model_id: ID of the model.
            params: Parameters to set on the model.

        Returns:
            None
        """
        for param in params or []:
            _validate_param(param.key, param.value)

        model = self.get_logged_model(model_id)
        for param in params:
            param_path = os.path.join(
                self._get_model_dir(model.experiment_id, model.model_id),
                FileStore.PARAMS_FOLDER_NAME,
                param.key,
            )
            make_containing_dirs(param_path)
            # Don't add trailing newline
            write_to(param_path, self._writeable_value(param.value))

    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.
        """
        model_dict = self._get_model_dict(model_id)
        model = LoggedModel.from_dictionary(model_dict)
        model.status = status
        model.last_updated_timestamp = int(time.time() * 1000)
        model_dir = self._get_model_dir(model.experiment_id, model.model_id)
        model_info_dict = self._make_persisted_model_dict(model)
        write_yaml(model_dir, FileStore.META_DATA_FILE_NAME, model_info_dict, overwrite=True)
        return self.get_logged_model(model_id)

    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
        """
        model = self.get_logged_model(model_id)
        for tag in tags:
            _validate_tag_name(tag.key)
            tag_path = os.path.join(
                self._get_model_dir(model.experiment_id, model.model_id),
                FileStore.TAGS_FOLDER_NAME,
                tag.key,
            )
            make_containing_dirs(tag_path)
            # Don't add trailing newline
            write_to(tag_path, self._writeable_value(tag.value))

    def delete_logged_model_tag(self, model_id: str, key: str) -> None:
        """
        Delete a tag on the specified logged model.

        Args:
            model_id: ID of the model.
            key: The string key of the tag.

        Returns:
            None
        """
        _validate_tag_name(key)
        model = self.get_logged_model(model_id)
        tag_path = os.path.join(
            self._get_model_dir(model.experiment_id, model.model_id),
            FileStore.TAGS_FOLDER_NAME,
            key,
        )
        if not exists(tag_path):
            raise MlflowException(
                f"No tag with key {key!r} found for model with ID {model_id!r}.",
                RESOURCE_DOES_NOT_EXIST,
            )
        os.remove(tag_path)

    def get_logged_model(self, model_id: str, allow_deleted: bool = False) -> LoggedModel:
        """
        Fetch the logged model with the specified ID.

        Args:
            model_id: ID of the model to fetch.
            allow_deleted: If ``True``, allow fetching logged models in the deleted lifecycle
                stage. Defaults to ``False``.

        Returns:
            The fetched model.
        """
        if not allow_deleted:
            return LoggedModel.from_dictionary(self._get_model_dict(model_id))

        exp_id, model_dir = self._find_model_root(model_id)
        if model_dir is None:
            raise MlflowException(
                f"Model '{model_id}' not found", databricks_pb2.RESOURCE_DOES_NOT_EXIST
            )

        model = self._get_model_from_dir(model_dir)
        if model.experiment_id != exp_id:
            raise MlflowException(
                f"Model '{model_id}' metadata is in invalid state.", databricks_pb2.INVALID_STATE
            )
        return model

    def delete_logged_model(self, model_id: str) -> None:
        model = self.get_logged_model(model_id)
        model.last_updated_timestamp = get_current_time_millis()
        model_dict = self._make_persisted_model_dict(model)
        model_dict["lifecycle_stage"] = LifecycleStage.DELETED
        model_dir = self._get_model_dir(model.experiment_id, model.model_id)
        write_yaml(
            model_dir,
            FileStore.META_DATA_FILE_NAME,
            model_dict,
            overwrite=True,
        )

    def _hard_delete_logged_model(self, model_id: str) -> None:
        model = self.get_logged_model(model_id, allow_deleted=True)
        model_dir = self._get_model_dir(model.experiment_id, model.model_id)
        shutil.rmtree(model_dir)

    def _get_deleted_logged_models(self, older_than=0) -> list[str]:
        current_time = get_current_time_millis()
        experiment_ids = self._get_active_experiments(False) + self._get_deleted_experiments(False)
        deleted_models = []
        for exp_id in experiment_ids:
            experiment_dir = self._get_experiment_path(exp_id, assert_exists=True)
            models_folder = os.path.join(experiment_dir, FileStore.MODELS_FOLDER_NAME)
            if not exists(models_folder):
                continue
            model_dirs = list_all(
                models_folder,
                filter_func=lambda path: (
                    all(
                        os.path.basename(os.path.normpath(path)) != reservedFolderName
                        for reservedFolderName in FileStore.RESERVED_EXPERIMENT_FOLDERS
                    )
                    and os.path.isdir(path)
                ),
                full_path=True,
            )
            for m_dir in model_dirs:
                try:
                    m_dict = self._get_model_info_from_dir(m_dir)
                except MissingConfigException:
                    continue
                if (
                    m_dict.get("lifecycle_stage") == LifecycleStage.DELETED
                    and m_dict.get("last_updated_timestamp", 0) <= current_time - older_than
                ):
                    deleted_models.append(m_dict["model_id"])
        return deleted_models

    def _get_model_artifact_dir(self, experiment_id: str, model_id: str) -> str:
        return append_to_uri_path(
            self.get_experiment(experiment_id).artifact_location,
            FileStore.MODELS_FOLDER_NAME,
            model_id,
            FileStore.ARTIFACTS_FOLDER_NAME,
        )

    def _make_persisted_model_dict(self, model: LoggedModel) -> dict[str, Any]:
        model_dict = model.to_dictionary()
        for field in ("tags", "params", "metrics"):
            model_dict.pop(field, None)
        return model_dict

    def _get_model_dict(self, model_id: str) -> dict[str, Any]:
        exp_id, model_dir = self._find_model_root(model_id)
        if model_dir is None:
            raise MlflowException(
                f"Model '{model_id}' not found", databricks_pb2.RESOURCE_DOES_NOT_EXIST
            )
        model_dict: dict[str, Any] = self._get_model_info_from_dir(model_dir)
        if model_dict.get("lifecycle_stage") == LifecycleStage.DELETED:
            raise MlflowException(
                f"Model '{model_id}' not found", databricks_pb2.RESOURCE_DOES_NOT_EXIST
            )

        if model_dict["experiment_id"] != exp_id:
            raise MlflowException(
                f"Model '{model_id}' metadata is in invalid state.", databricks_pb2.INVALID_STATE
            )
        return model_dict

    def _get_model_dir(self, experiment_id: str, model_id: str) -> str:
        if not self._has_experiment(experiment_id):
            return None
        return os.path.join(
            self._get_experiment_path(experiment_id, assert_exists=True),
            FileStore.MODELS_FOLDER_NAME,
            model_id,
        )

    def _find_model_root(self, model_id):
        self._check_root_dir()
        all_experiments = self._get_active_experiments(False) + self._get_deleted_experiments(False)
        for experiment_dir in all_experiments:
            models_dir_path = os.path.join(
                self.root_directory, experiment_dir, FileStore.MODELS_FOLDER_NAME
            )
            if not os.path.exists(models_dir_path):
                continue
            models = find(models_dir_path, model_id, full_path=True)
            if len(models) == 0:
                continue
            return os.path.basename(os.path.dirname(os.path.abspath(models_dir_path))), models[0]
        return None, None

    def _get_model_from_dir(self, model_dir: str) -> LoggedModel:
        return LoggedModel.from_dictionary(self._get_model_info_from_dir(model_dir))

    def _get_model_info_from_dir(self, model_dir: str) -> dict[str, Any]:
        model_dict = FileStore._read_yaml(model_dir, FileStore.META_DATA_FILE_NAME)
        model_dict["tags"] = self._get_all_model_tags(model_dir)
        model_dict["params"] = {p.key: p.value for p in self._get_all_model_params(model_dir)}
        model_dict["metrics"] = self._get_all_model_metrics(
            model_id=model_dict["model_id"], model_dir=model_dir
        )
        return model_dict

    def _get_all_model_tags(self, model_dir: str) -> list[LoggedModelTag]:
        parent_path, tag_files = self._get_resource_files(model_dir, FileStore.TAGS_FOLDER_NAME)
        return [self._get_tag_from_file(parent_path, tag_file) for tag_file in tag_files]

    def _get_all_model_params(self, model_dir: str) -> list[LoggedModelParameter]:
        parent_path, param_files = self._get_resource_files(model_dir, FileStore.PARAMS_FOLDER_NAME)
        return [self._get_param_from_file(parent_path, param_file) for param_file in param_files]

    def _get_all_model_metrics(self, model_id: str, model_dir: str) -> list[Metric]:
        parent_path, metric_files = self._get_resource_files(
            model_dir, FileStore.METRICS_FOLDER_NAME
        )
        metrics = []
        for metric_file in metric_files:
            metrics.extend(
                FileStore._get_model_metrics_from_file(
                    model_id=model_id, parent_path=parent_path, metric_name=metric_file
                )
            )
        return metrics

    @staticmethod
    def _get_model_metrics_from_file(
        model_id: str, parent_path: str, metric_name: str
    ) -> list[Metric]:
        _validate_metric_name(metric_name)
        metric_objs = [
            FileStore._get_model_metric_from_line(model_id, metric_name, line)
            for line in read_file_lines(parent_path, metric_name)
        ]
        if len(metric_objs) == 0:
            raise ValueError(f"Metric '{metric_name}' is malformed. No data found.")

        # Group metrics by (dataset_name, dataset_digest)
        grouped_metrics = defaultdict(list)
        for metric in metric_objs:
            key = (metric.dataset_name, metric.dataset_digest)
            grouped_metrics[key].append(metric)

        # Compute the max for each group
        return [
            max(group, key=lambda m: (m.step, m.timestamp, m.value))
            for group in grouped_metrics.values()
        ]

    @staticmethod
    def _get_model_metric_from_line(model_id: str, metric_name: str, metric_line: str) -> Metric:
        metric_parts = metric_line.strip().split(" ")
        if len(metric_parts) not in [4, 6]:
            raise MlflowException(
                f"Metric '{metric_name}' is malformed; persisted metric data contained "
                f"{len(metric_parts)} fields. Expected 4 or 6 fields.",
                databricks_pb2.INTERNAL_ERROR,
            )
        ts = int(metric_parts[0])
        val = float(metric_parts[1])
        step = int(metric_parts[2])
        run_id = str(metric_parts[3])
        dataset_name = str(metric_parts[4]) if len(metric_parts) == 6 else None
        dataset_digest = str(metric_parts[5]) if len(metric_parts) == 6 else None
        # TODO: Read run ID from the metric file and pass it to the Metric constructor
        return Metric(
            key=metric_name,
            value=val,
            timestamp=ts,
            step=step,
            model_id=model_id,
            dataset_name=dataset_name,
            dataset_digest=dataset_digest,
            run_id=run_id,
        )

    def search_logged_models(
        self,
        experiment_ids: list[str],
        filter_string: str | None = None,
        datasets: list[DatasetFilter] | None = None,
        max_results: int | None = None,
        order_by: list[dict[str, Any]] | None = None,
        page_token: str | None = 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. Default is 100.
            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.
        """
        if datasets and not all(d.get("dataset_name") for d in datasets):
            raise MlflowException(
                "`dataset_name` in the `datasets` clause must be specified.",
                INVALID_PARAMETER_VALUE,
            )
        max_results = max_results or SEARCH_LOGGED_MODEL_MAX_RESULTS_DEFAULT
        all_models = []
        for experiment_id in experiment_ids:
            models = self._list_models(experiment_id)
            all_models.extend(models)
        filtered = SearchLoggedModelsUtils.filter_logged_models(all_models, filter_string, datasets)
        sorted_logged_models = SearchLoggedModelsUtils.sort(filtered, order_by)
        logged_models, next_page_token = SearchLoggedModelsUtils.paginate(
            sorted_logged_models, page_token, max_results
        )
        return PagedList(logged_models, next_page_token)

    def _list_models(self, experiment_id: str) -> list[LoggedModel]:
        self._check_root_dir()
        if not self._has_experiment(experiment_id):
            return []
        experiment_dir = self._get_experiment_path(experiment_id, assert_exists=True)
        models_folder = os.path.join(experiment_dir, FileStore.MODELS_FOLDER_NAME)
        if not exists(models_folder):
            return []
        model_dirs = list_all(
            models_folder,
            filter_func=lambda x: (
                all(
                    os.path.basename(os.path.normpath(x)) != reservedFolderName
                    for reservedFolderName in FileStore.RESERVED_EXPERIMENT_FOLDERS
                )
                and os.path.isdir(x)
            ),
            full_path=True,
        )
        models = []
        for m_dir in model_dirs:
            try:
                # trap and warn known issues, will raise unexpected exceptions to caller
                m_dict = self._get_model_info_from_dir(m_dir)
                if m_dict.get("lifecycle_stage") == LifecycleStage.DELETED:
                    continue
                model = LoggedModel.from_dictionary(m_dict)
                if model.experiment_id != experiment_id:
                    logging.warning(
                        "Wrong experiment ID (%s) recorded for model '%s'. "
                        "It should be %s. Model will be ignored.",
                        str(model.experiment_id),
                        str(model.model_id),
                        str(experiment_id),
                        exc_info=True,
                    )
                    continue
                models.append(model)
            except MissingConfigException as exc:
                # trap malformed model exception and log
                # this is at debug level because if the same store is used for
                # artifact storage, it's common the folder is not a run folder
                m_id = os.path.basename(m_dir)
                logging.debug(
                    "Malformed model '%s'. Detailed error %s", m_id, str(exc), exc_info=True
                )
        return models

    #######################################################################################
    # Below are legacy V2 Tracing APIs. DO NOT USE. Use the V3 APIs instead.
    #######################################################################################
    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_id = generate_request_id_v2()
        _validate_experiment_id(experiment_id)
        experiment_dir = self._get_experiment_path(
            experiment_id, view_type=ViewType.ACTIVE_ONLY, assert_exists=True
        )
        mkdir(experiment_dir, FileStore.TRACES_FOLDER_NAME)
        traces_dir = os.path.join(experiment_dir, FileStore.TRACES_FOLDER_NAME)
        mkdir(traces_dir, request_id)
        trace_dir = os.path.join(traces_dir, request_id)
        artifact_uri = self._get_traces_artifact_dir(experiment_id, request_id)
        tags.update({MLFLOW_ARTIFACT_LOCATION: artifact_uri})
        trace_info = TraceInfoV2(
            request_id=request_id,
            experiment_id=experiment_id,
            timestamp_ms=timestamp_ms,
            execution_time_ms=None,
            status=TraceStatus.IN_PROGRESS,
            request_metadata=request_metadata,
            tags=tags,
        )
        self._save_trace_info(trace_info.to_v3(), trace_dir)
        return 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.
        """
        trace_info, trace_dir = self._get_trace_info_and_dir(request_id)
        trace_info.execution_duration = timestamp_ms - trace_info.request_time
        trace_info.state = status.to_state()
        trace_info.trace_metadata.update(request_metadata)
        trace_info.tags.update(tags)
        self._save_trace_info(trace_info, trace_dir, overwrite=True)
        return TraceInfoV2.from_v3(trace_info)

    # Evaluation Dataset APIs - Not supported in FileStore

    @filestore_not_supported
    def create_dataset(
        self,
        name: str,
        tags: dict[str, Any] | None = None,
        experiment_ids: list[str] | None = None,
    ):
        pass

    @filestore_not_supported
    def get_dataset(self, dataset_id):
        pass

    @filestore_not_supported
    def delete_dataset(self, dataset_id):
        pass

    @filestore_not_supported
    def search_datasets(
        self,
        experiment_ids=None,
        filter_string=None,
        max_results=1000,
        order_by=None,
        page_token=None,
    ):
        pass

    @filestore_not_supported
    def upsert_dataset_records(self, dataset_id, records):
        pass

    @filestore_not_supported
    def set_dataset_tags(self, dataset_id, tags):
        pass

    @filestore_not_supported
    def get_dataset_experiment_ids(self, dataset_id):
        pass

    @filestore_not_supported
    def delete_dataset_tag(self, dataset_id, key):
        pass

    @filestore_not_supported
    def add_dataset_to_experiments(self, dataset_id, experiment_ids):
        pass

    @filestore_not_supported
    def remove_dataset_from_experiments(self, dataset_id, experiment_ids):
        pass

    def link_traces_to_run(self, trace_ids: list[str], run_id: str) -> None:
        """
        Link multiple traces to a run by creating entity associations.

        Note: This feature is not supported in FileStore.

        Args:
            trace_ids: List of trace IDs to link to the run.
            run_id: ID of the run to link traces to.

        Raises:
            MlflowException: Always raised as this operation is not supported in FileStore.
        """
        raise MlflowException(
            "Linking traces to runs is not supported in FileStore. "
            "Please use a database-backed store (e.g., SQLAlchemy store) for this feature.",
            error_code=databricks_pb2.INVALID_PARAMETER_VALUE,
        )

    def link_prompts_to_trace(self, trace_id: str, prompt_versions: list[PromptVersion]) -> None:
        """
        Link multiple prompt versions to a trace by creating entity associations.

        Args:
            trace_id: ID of the trace to link prompt versions to.
            prompt_versions: List of PromptVersion objects to link.
        """
        raise MlflowException(
            "Linking prompts to traces is not supported in FileStore. "
            "Please use a database-backed store (e.g., SQLAlchemy store) for this feature.",
            error_code=databricks_pb2.INVALID_PARAMETER_VALUE,
        )

    # Trace metrics API is not supported in FileStore, override the
    # abstract method to raise an explicit error.

    @filestore_not_supported
    def query_trace_metrics(self, *args, **kwargs):
        pass