mlflow--mlflow/mlflow/optuna/storage.py

import copy
import datetime
import json
import threading
import time
import uuid
import weakref
from collections.abc import Container, Sequence
from typing import Any

from mlflow import MlflowClient
from mlflow.entities import Metric, Param, RunTag
from mlflow.utils.mlflow_tags import MLFLOW_PARENT_RUN_ID

try:
    from optuna._typing import JSONSerializable
    from optuna.distributions import (
        BaseDistribution,
        check_distribution_compatibility,
        distribution_to_json,
        json_to_distribution,
    )
    from optuna.storages import BaseStorage
    from optuna.storages._base import DEFAULT_STUDY_NAME_PREFIX
    from optuna.study import StudyDirection
    from optuna.study._frozen import FrozenStudy
    from optuna.trial import FrozenTrial, TrialState
except ImportError as e:
    raise ImportError("Install optuna to use `mlflow.optuna` module") from e

optuna_mlflow_status_map = {
    TrialState.RUNNING: "RUNNING",
    TrialState.COMPLETE: "FINISHED",
    TrialState.PRUNED: "KILLED",
    TrialState.FAIL: "FAILED",
    TrialState.WAITING: "SCHEDULED",
}

mlflow_optuna_status_map = {
    "RUNNING": TrialState.RUNNING,
    "FINISHED": TrialState.COMPLETE,
    "KILLED": TrialState.PRUNED,
    "FAILED": TrialState.FAIL,
    "SCHEDULED": TrialState.WAITING,
}


def _periodic_flush_worker(
    ref: "weakref.ref[MlflowStorage]",
    stop_event: threading.Event,
    flush_interval: float,
) -> None:
    """Background thread that periodically flushes batched data.

    Uses a weak reference so the thread does not prevent garbage collection
    of the owning MlflowStorage instance.  When the instance is collected
    the weak reference returns None and the loop exits.
    """
    while not stop_event.wait(timeout=min(0.1, flush_interval / 10)):
        obj = ref()
        if obj is None:
            break
        try:
            current_time = time.time()
            if current_time - obj._last_flush_time >= obj._batch_flush_interval:
                obj.flush_all_batches()
                obj._last_flush_time = current_time
        except Exception:
            # Sleep longer on error to avoid tight retry loop
            stop_event.wait(timeout=1.0)
        finally:
            del obj


class MlflowStorage(BaseStorage):
    """
    MLflow based storage class with batch processing to avoid REST API throttling.
    """

    def __init__(
        self,
        experiment_id: str,
        name: str | None = None,
        batch_flush_interval: float = 1.0,
        batch_size_threshold: int = 100,
    ):
        """
        Initialize MLFlowStorage with batching capabilities.

        Parameters
        ----------
        experiment_id : str
            MLflow experiment ID
        name : Optional[str]
            Optional name for the storage
        batch_flush_interval : float
            Time in seconds between automatic batch flushes (default: 1.0)
        batch_size_threshold : int
            Maximum number of items in batch before triggering a flush (default: 100)
        """
        if not experiment_id:
            raise Exception("No experiment_id provided. MLFlowStorage cannot create experiments.")

        self._experiment_id = experiment_id
        self._mlflow_client = MlflowClient()
        self._name = name

        # Batching configuration
        self._batch_flush_interval = batch_flush_interval
        self._batch_size_threshold = batch_size_threshold

        # Batching queues for metrics, parameters, and tags
        self._batch_queue = {}  # Dictionary of run_id -> {'metrics': [], 'params': [], 'tags': []}
        self._batch_lock = threading.RLock()
        self._last_flush_time = time.time()

        # Event to signal the worker thread to stop
        self._stop_event = threading.Event()

        # Register a weak-ref callback that fires the stop event when this
        # instance is garbage-collected, so the flush thread exits promptly.
        # Pass the Event object directly (not a bound method on self) to avoid
        # preventing GC via a reference cycle.
        stop = self._stop_event
        weakref.finalize(self, stop.set)

        # Start a background thread for periodic flushing using a weak reference
        # so the thread does not prevent garbage collection of this instance.
        self._flush_thread = threading.Thread(
            target=_periodic_flush_worker,
            args=(weakref.ref(self), self._stop_event, self._batch_flush_interval),
            daemon=True,
            name=f"mlflow_optuna_batch_flush_worker_{uuid.uuid4().hex[:8]}",
        )
        self._flush_thread.start()

    def __getstate__(self):
        """
        Prepare the object for serialization by removing non-picklable components.
        This is called when the object is being pickled.
        """
        state = self.__dict__.copy()

        # Remove thread-related attributes that can't be pickled
        state.pop("_batch_lock", None)
        state.pop("_flush_thread", None)
        state.pop("_stop_event", None)

        # Store the configuration but not the actual lock/thread
        state["_thread_running"] = hasattr(self, "_flush_thread") and (
            self._flush_thread is not None and self._flush_thread.is_alive()
        )

        return state

    def __setstate__(self, state):
        """
        Restore the object after deserialization by recreating non-picklable components.
        This is called when the object is being unpickled.
        """
        # First, update the instance with the pickled state
        self.__dict__.update(state)

        # Recreate the lock
        self._batch_lock = threading.RLock()

        # Don't automatically restart the thread on workers - this would create too many threads
        # Instead, we'll use a manual flush approach in distributed contexts
        self._flush_thread = None

        # If we're on a worker node, we should disable automatic background flushing
        # because it could cause issues with multiple threads trying to write to MLflow
        self._stop_event = threading.Event()
        self._stop_event.set()  # Don't start flushing on deserialized workers

    def _queue_batch_operation(
        self,
        run_id: str,
        metrics: list[Metric] | None = None,
        params: list[Param] | None = None,
        tags: list[RunTag] | None = None,
    ):
        """Queue metrics, parameters, or tags for batched processing."""
        with self._batch_lock:
            if run_id not in self._batch_queue:
                self._batch_queue[run_id] = {"metrics": [], "params": [], "tags": []}

            batch = self._batch_queue[run_id]

            if metrics:
                batch["metrics"].extend(metrics)
            if params:
                batch["params"].extend(params)
            if tags:
                batch["tags"].extend(tags)

            # Check if we've reached the batch size threshold for this run
            batch_size = len(batch["metrics"]) + len(batch["params"]) + len(batch["tags"])
            if batch_size >= self._batch_size_threshold:
                self._flush_batch(run_id)

    def _flush_batch(self, run_id: str):
        """Flush the batch for a specific run_id to MLflow."""
        with self._batch_lock:
            if run_id not in self._batch_queue:
                return

            batch = self._batch_queue[run_id]

            # Only make the API call if there's something to flush
            if batch["metrics"] or batch["params"] or batch["tags"]:
                try:
                    self._mlflow_client.log_batch(
                        run_id, metrics=batch["metrics"], params=batch["params"], tags=batch["tags"]
                    )
                except Exception as e:
                    # If the run doesn't exist, propagate the error
                    if "Run with id=" in str(e) and "not found" in str(e):
                        raise
                    # Otherwise, handle or log the error as needed

                # Clear the batch
                batch["metrics"] = []
                batch["params"] = []
                batch["tags"] = []

    def flush_all_batches(self):
        """Flush all pending batches to MLflow."""
        with self._batch_lock:
            run_ids = list(self._batch_queue.keys())

        # Flush each run's batch
        for run_id in run_ids:
            self._flush_batch(run_id)

    def _search_runs_by_name(self, run_name: str):
        filter_string = f"tags.mlflow.runName = '{run_name}'"
        return self._mlflow_client.search_runs(
            experiment_ids=[self._experiment_id],
            filter_string=filter_string,
            order_by=["attributes.start_time DESC"],
        )

    def create_new_study(
        self, directions: Sequence[StudyDirection], study_name: str | None = None
    ) -> int:
        """Create a new study as a mlflow run."""
        study_name = study_name or DEFAULT_STUDY_NAME_PREFIX + str(uuid.uuid4())
        tags = {
            "mlflow.runName": study_name,
            "optuna.study_direction": ",".join(direction.name for direction in directions),
        }
        study_run = self._mlflow_client.create_run(experiment_id=self._experiment_id, tags=tags)
        return study_run.info.run_id

    def delete_study(self, study_id) -> None:
        """Delete a study."""
        # Ensure any pending changes are saved before deletion
        self._flush_batch(study_id)
        self._mlflow_client.delete_run(study_id)

    def set_study_user_attr(self, study_id, key: str, value: JSONSerializable) -> None:
        """Register a user-defined attribute as mlflow run tags to a study run."""
        # Verify the run exists first to fail fast if it doesn't
        self._mlflow_client.get_run(study_id)

        # Queue the tag if the run exists
        self._queue_batch_operation(study_id, tags=[RunTag(f"user_{key}", json.dumps(value))])

    def set_study_system_attr(self, study_id, key: str, value: JSONSerializable) -> None:
        """Register a optuna-internal attribute as mlflow run tags to a study run."""
        # Verify the run exists first to fail fast if it doesn't
        self._mlflow_client.get_run(study_id)

        # Queue the tag if the run exists
        self._queue_batch_operation(study_id, tags=[RunTag(f"sys_{key}", json.dumps(value))])

    def get_study_id_from_name(self, study_name: str) -> int:
        # Flush all batches to ensure we have the latest data
        self.flush_all_batches()

        runs = self._search_runs_by_name(study_name)
        if len(runs):
            return runs[0].info.run_id
        else:
            raise Exception(f"Study {study_name} not found")

    def get_study_id_by_name_if_exists(self, study_name: str) -> str | None:
        """Get study ID from name if it exists, otherwise return None.

        Args:
            study_name: The name of the study to look for

        Returns:
            Study ID if found, None otherwise
        """
        # Flush all batches to ensure we have the latest data
        self.flush_all_batches()

        if runs := self._search_runs_by_name(study_name):
            return runs[0].info.run_id
        else:
            return None

    def get_study_name_from_id(self, study_id) -> str:
        # Flush the batch for this study to ensure we have the latest data
        self._flush_batch(study_id)

        run = self._mlflow_client.get_run(study_id)
        return run.data.tags["mlflow.runName"]

    def get_study_directions(self, study_id) -> list[StudyDirection]:
        # Flush the batch for this study to ensure we have the latest data
        self._flush_batch(study_id)

        run = self._mlflow_client.get_run(study_id)
        directions_str = run.data.tags["optuna.study_direction"]
        return [StudyDirection[name] for name in directions_str.split(",")]

    def get_study_user_attrs(self, study_id) -> dict[str, Any]:
        # Flush the batch for this study to ensure we have the latest data
        self._flush_batch(study_id)

        run = self._mlflow_client.get_run(study_id)
        user_attrs = {}
        for key, value in run.data.tags.items():
            if key.startswith("user_"):
                user_attrs[key[5:]] = json.loads(value)
        return user_attrs

    def get_study_system_attrs(self, study_id) -> dict[str, Any]:
        # Flush the batch for this study to ensure we have the latest data
        self._flush_batch(study_id)

        run = self._mlflow_client.get_run(study_id)
        system_attrs = {}
        for key, value in run.data.tags.items():
            if key.startswith("sys_"):
                system_attrs[key[4:]] = json.loads(value)
        return system_attrs

    def get_all_studies(self) -> list[FrozenStudy]:
        # Flush all batches to ensure we have the latest data
        self.flush_all_batches()

        runs = self._mlflow_client.search_runs(experiment_ids=[self._experiment_id])
        studies = []
        for run in runs:
            study_id = run.info.run_id
            study_name = run.data.tags["mlflow.runName"]
            directions_str = run.data.tags["optuna.study_direction"]
            directions = [StudyDirection[name] for name in directions_str.split(",")]
            studies.append(
                FrozenStudy(
                    study_name=study_name,
                    direction=None,
                    directions=directions,
                    user_attrs=self.get_study_user_attrs(study_id),
                    system_attrs=self.get_study_system_attrs(study_id),
                    study_id=study_id,
                )
            )
        return studies

    def create_new_trial(self, study_id, template_trial: FrozenTrial | None = None) -> int:
        # Ensure study batch is flushed before creating a new trial
        self._flush_batch(study_id)

        if template_trial:
            frozen = copy.deepcopy(template_trial)
        else:
            frozen = FrozenTrial(
                trial_id=-1,  # dummy value.
                number=-1,  # dummy value.
                state=TrialState.RUNNING,
                params={},
                distributions={},
                user_attrs={},
                system_attrs={},
                value=None,
                intermediate_values={},
                datetime_start=datetime.datetime.now(),
                datetime_complete=None,
            )

        distribution_json = {
            k: distribution_to_json(dist) for k, dist in frozen.distributions.items()
        }
        distribution_str = json.dumps(distribution_json)
        tags = {"param_directions": distribution_str}

        trial_run = self._mlflow_client.create_run(experiment_id=self._experiment_id, tags=tags)
        trial_id = trial_run.info.run_id

        # Add parent run ID tag
        self._queue_batch_operation(trial_id, tags=[RunTag(MLFLOW_PARENT_RUN_ID, study_id)])

        # Log trial_id metric to study
        hash_id = float(hash(trial_id))
        self._queue_batch_operation(
            study_id, metrics=[Metric("trial_id", hash_id, int(time.time() * 1000), 1)]
        )

        # Ensure study batch is flushed to get accurate metric history
        self._flush_batch(study_id)

        trial_ids = self._mlflow_client.get_metric_history(study_id, "trial_id")
        index = next((i for i, obj in enumerate(trial_ids) if obj.value == hash_id), -1)

        self._queue_batch_operation(trial_id, tags=[RunTag("numbers", str(index))])

        # Set trial state
        state = frozen.state
        if state.is_finished():
            self._mlflow_client.set_terminated(trial_id, status=optuna_mlflow_status_map[state])
        else:
            self._mlflow_client.update_run(trial_id, status=optuna_mlflow_status_map[state])

        timestamp = int(time.time() * 1000)
        metrics = []
        params = []
        tags = []

        # Add metrics
        if frozen.values is not None:
            if len(frozen.values) > 1:
                metrics.extend([
                    Metric(f"value_{idx}", val, timestamp, 1)
                    for idx, val in enumerate(frozen.values)
                ])
            else:
                metrics.append(Metric("value", frozen.values[0], timestamp, 1))
        elif frozen.value is not None:
            metrics.append(Metric("value", frozen.value, timestamp, 1))

        # Add intermediate values
        metrics.extend([
            Metric("intermediate_value", val, timestamp, int(k))
            for k, val in frozen.intermediate_values.items()
        ])

        # Add params
        params.extend([Param(k, param) for k, param in frozen.params.items()])

        # Add tags
        tags.extend([
            RunTag(f"user_{key}", json.dumps(value)) for key, value in frozen.user_attrs.items()
        ])
        tags.extend([
            RunTag(f"sys_{key}", json.dumps(value)) for key, value in frozen.system_attrs.items()
        ])
        tags.extend([
            RunTag(
                f"param_internal_val_{k}",
                json.dumps(frozen.distributions[k].to_internal_repr(param)),
            )
            for k, param in frozen.params.items()
        ])

        # Queue all the data to be sent in batches
        self._queue_batch_operation(trial_id, metrics=metrics, params=params, tags=tags)

        return trial_id

    def set_trial_param(
        self,
        trial_id,
        param_name: str,
        param_value_internal: float,
        distribution: BaseDistribution,
    ) -> None:
        # Flush the batch for this trial to ensure we have the latest data
        self._flush_batch(trial_id)

        trial_run = self._mlflow_client.get_run(trial_id)
        distributions_dict = json.loads(trial_run.data.tags["param_directions"])
        self.check_trial_is_updatable(trial_id, mlflow_optuna_status_map[trial_run.info.status])

        if param_name in trial_run.data.params:
            param_distribution = json_to_distribution(distributions_dict[param_name])
            check_distribution_compatibility(param_distribution, distribution)

        # Queue parameter update
        self._queue_batch_operation(
            trial_id,
            params=[Param(param_name, distribution.to_external_repr(param_value_internal))],
            tags=[RunTag(f"param_internal_val_{param_name}", json.dumps(param_value_internal))],
        )

        distributions_dict[param_name] = distribution_to_json(distribution)
        self._queue_batch_operation(
            trial_id, tags=[RunTag("param_directions", json.dumps(distributions_dict))]
        )

    def get_trial_id_from_study_id_trial_number(self, study_id, trial_number: int) -> int:
        raise NotImplementedError("This method is not supported in MLflow backend.")

    def get_trial_number_from_id(self, trial_id) -> int:
        # Flush the batch for this trial to ensure we have the latest data
        self._flush_batch(trial_id)

        trial_run = self._mlflow_client.get_run(trial_id)
        return int(trial_run.data.tags.get("numbers", 0))

    def get_trial_param(self, trial_id, param_name: str) -> float:
        # Flush the batch for this trial to ensure we have the latest data
        self._flush_batch(trial_id)

        trial_run = self._mlflow_client.get_run(trial_id)
        param_value = trial_run.data.tags[f"param_internal_val_{param_name}"]

        return float(json.loads(param_value))

    def set_trial_state_values(
        self, trial_id, state: TrialState, values: Sequence[float] | None = None
    ) -> bool:
        # Update trial state
        if state.is_finished():
            self._mlflow_client.set_terminated(trial_id, status=optuna_mlflow_status_map[state])
        else:
            self._mlflow_client.update_run(trial_id, status=optuna_mlflow_status_map[state])

        # Queue value metrics if provided
        timestamp = int(time.time() * 1000)
        if values is not None:
            metrics = []
            if len(values) > 1:
                metrics = [
                    Metric(f"value_{idx}", val, timestamp, 1) for idx, val in enumerate(values)
                ]
            else:
                metrics = [Metric("value", values[0], timestamp, 1)]

            self._queue_batch_operation(trial_id, metrics=metrics)

        if state == TrialState.RUNNING and state != TrialState.WAITING:
            return False
        return True

    def set_trial_intermediate_value(self, trial_id, step: int, intermediate_value: float) -> None:
        # Queue intermediate value metric
        self._queue_batch_operation(
            trial_id,
            metrics=[
                Metric("intermediate_value", intermediate_value, int(time.time() * 1000), step)
            ],
        )

    def set_trial_user_attr(self, trial_id, key: str, value: Any) -> None:
        # Queue user attribute tag
        self._queue_batch_operation(trial_id, tags=[RunTag(f"user_{key}", json.dumps(value))])

    def set_trial_system_attr(self, trial_id, key: str, value: Any) -> None:
        # Queue system attribute tag
        self._queue_batch_operation(trial_id, tags=[RunTag(f"sys_{key}", json.dumps(value))])

    def get_trial(self, trial_id) -> FrozenTrial:
        # Flush the batch for this trial to ensure we have the latest data
        self._flush_batch(trial_id)

        trial_run = self._mlflow_client.get_run(trial_id)
        param_directions = trial_run.data.tags["param_directions"]
        try:
            distributions_dict = json.loads(param_directions)
        except json.decoder.JSONDecodeError as e:
            raise ValueError(f"error with param_directions = {param_directions!r}") from e

        distributions = {
            k: json_to_distribution(distribution) for k, distribution in distributions_dict.items()
        }
        params = {}
        for key, value in trial_run.data.tags.items():
            if key.startswith("param_internal_val_"):
                param_name = key[19:]
                param_value = json.loads(value)
                params[param_name] = distributions[param_name].to_external_repr(float(param_value))

        metrics = trial_run.data.metrics
        values = None
        if "value" in metrics:
            values = [metrics["value"]]
        if "value_0" in metrics:
            values = [metrics[f"value_{idx}"] for idx in range(len(metrics))]

        run_number = int(trial_run.data.tags.get("numbers", 0))

        start_time = datetime.datetime.fromtimestamp(trial_run.info.start_time / 1000)
        if trial_run.info.end_time:
            end_time = datetime.datetime.fromtimestamp(trial_run.info.end_time / 1000)
        else:
            end_time = None
        return FrozenTrial(
            trial_id=trial_id,
            number=run_number,
            state=mlflow_optuna_status_map[trial_run.info.status],
            value=None,
            values=values,
            datetime_start=start_time,
            datetime_complete=end_time,
            params=params,
            distributions=distributions,
            user_attrs=self.get_trial_user_attrs(trial_id),
            system_attrs=self.get_trial_system_attrs(trial_id),
            intermediate_values={
                v.step: v.value
                for idx, v in enumerate(
                    self._mlflow_client.get_metric_history(trial_id, "intermediate_value")
                )
            },
        )

    def get_trial_user_attrs(self, trial_id) -> dict[str, Any]:
        # Flush the batch for this trial to ensure we have the latest data
        self._flush_batch(trial_id)

        run = self._mlflow_client.get_run(trial_id)
        user_attrs = {}
        for key, value in run.data.tags.items():
            if key.startswith("user_"):
                user_attrs[key[5:]] = json.loads(value)
        return user_attrs

    def get_trial_system_attrs(self, trial_id) -> dict[str, Any]:
        # Flush the batch for this trial to ensure we have the latest data
        self._flush_batch(trial_id)

        run = self._mlflow_client.get_run(trial_id)
        system_attrs = {}
        for key, value in run.data.tags.items():
            if key.startswith("sys_"):
                system_attrs[key[4:]] = json.loads(value)
        return system_attrs

    def get_all_trials(
        self,
        study_id,
        deepcopy: bool = True,
        states: Container[TrialState] | None = None,
    ) -> list[FrozenTrial]:
        # Flush all batches to ensure we have the latest data
        self.flush_all_batches()

        runs = self._mlflow_client.search_runs(
            experiment_ids=[self._experiment_id],
            filter_string=f"tags.mlflow.parentRunId='{study_id}'",
        )
        trials = [self.get_trial(run.info.run_id) for run in runs]

        frozen_trials: list[FrozenTrial] = [
            trial for trial in trials if states is None or trial.state in states
        ]
        return frozen_trials

    def get_n_trials(self, study_id, states=None) -> int:
        # Flush all batches to ensure we have the latest data
        self.flush_all_batches()

        runs = self._mlflow_client.search_runs(
            experiment_ids=[self._experiment_id],
            filter_string=f"tags.mlflow.parentRunId='{study_id}'",
        )
        return len(runs)