From fdf0ffa96959237aaa869293b203071b15672fe1 Mon Sep 17 00:00:00 2001 From: hallelx2 Date: Sun, 25 Oct 2026 08:45:15 +0100 Subject: [PATCH] fix(mlflow): nest the logged runs under the pipeline's own run `cancerag.pipeline` opens an MLflow run for the pipeline as a whole. The `advanced` step then calls `log_all_runs`, which opened its own top-level run per CV fold, and MLflow refuses a second active top-level run: Run with UUID ... is already active. To start a new run, first end the current run with mlflow.end_run(). To start a nested run, call start_run with nested=True So the step died on its final phase and advanced_summary.json was never written -- while optuna_trials.csv and the tuned models sat on disk, which made it look like the step had simply not run yet rather than failed. It is the second terminal failure in the chain, independent of the figures one. All five call sites now go through a `_record` helper that nests when a run is active and stays top-level when none is. Nesting is the right relationship rather than a workaround: these per-fold and per-experiment records genuinely are children of the pipeline run. Run standalone, nothing is active and the behaviour is unchanged. The tracking URI and experiment are no longer set under an active run either. They are already fixed by whoever opened it, and a nested run belongs to its parent's experiment regardless, so setting them only misreported where the records land. The tests were checked against the unfixed code: reverting the helper reproduces the exact production exception. --- src/cancerag/ml/mlflow_logging.py | 38 ++++++++--- tests/unit/test_mlflow_logging.py | 104 ++++++++++++++++++++++++++++++ 2 files changed, 133 insertions(+), 9 deletions(-) create mode 100644 tests/unit/test_mlflow_logging.py diff --git a/src/cancerag/ml/mlflow_logging.py b/src/cancerag/ml/mlflow_logging.py index e051257..313ba86 100644 --- a/src/cancerag/ml/mlflow_logging.py +++ b/src/cancerag/ml/mlflow_logging.py @@ -14,8 +14,6 @@ import json import logging -import os -from datetime import datetime from pathlib import Path import pandas as pd @@ -41,13 +39,34 @@ def _safe_log_param(mlflow, key, value): logger.debug("log_param(%s, %s) skipped: %s", key, value, exc) +def _record(mlflow, **kw): + """Open an MLflow run, nesting it when one is already active. + + `cancerag.pipeline` opens a run for the pipeline as a whole, so a plain + `start_run` here raises "Run with UUID ... is already active" and takes + the `advanced` step down with it -- which is why `advanced_summary.json` + was never written even though Optuna had finished and its artefacts were + on disk. + + Nesting is the right relationship, not a workaround: these per-fold and + per-experiment records genuinely are children of the pipeline run. Run + standalone (`python -m cancerag.ml.mlflow_logging`) nothing is active and + they stay top-level, as before. + """ + return mlflow.start_run(nested=mlflow.active_run() is not None, **kw) + + def log_all_runs(*, mlruns_dir: Path | str = MLRUNS_DIR) -> dict: """Walk all known result files and emit MLflow runs.""" import mlflow mlruns_dir = Path(mlruns_dir).resolve() mlruns_dir.mkdir(parents=True, exist_ok=True) - mlflow.set_tracking_uri(f"file://{mlruns_dir}") - mlflow.set_experiment("cancerag") + # Under an active run the tracking URI and experiment are already fixed by + # whoever opened it, and a nested run belongs to its parent's experiment + # regardless -- setting them here would only misreport where these land. + if mlflow.active_run() is None: + mlflow.set_tracking_uri(f"file://{mlruns_dir}") + mlflow.set_experiment("cancerag") n_runs = 0 # 1. Bake-off CV results — one MLflow run per (model, split, fold, seed) @@ -55,7 +74,7 @@ def log_all_runs(*, mlruns_dir: Path | str = MLRUNS_DIR) -> dict: if cv_path.exists(): cv = pd.read_csv(cv_path) for _, row in cv.iterrows(): - with mlflow.start_run(run_name=f"{row['model']}|{row['split']}|fold{int(row['fold'])}|s{int(row['seed'])}"): + with _record(mlflow, run_name=f"{row['model']}|{row['split']}|fold{int(row['fold'])}|s{int(row['seed'])}"): _safe_log_param(mlflow, "model", row["model"]) _safe_log_param(mlflow, "split", row["split"]) _safe_log_param(mlflow, "fold", int(row["fold"])) @@ -74,7 +93,7 @@ def log_all_runs(*, mlruns_dir: Path | str = MLRUNS_DIR) -> dict: cb_path = ML_MODELS / "baselines/chemberta/chemberta_summary.json" if cb_path.exists(): cb = json.loads(cb_path.read_text()) - with mlflow.start_run(run_name="chemberta_baseline_summary"): + with _record(mlflow, run_name="chemberta_baseline_summary"): _safe_log_param(mlflow, "model", "chemberta_logreg") _safe_log_param(mlflow, "checkpoint", cb.get("checkpoint")) _safe_log_param(mlflow, "embedding_dim", cb.get("embedding_dim")) @@ -97,7 +116,8 @@ def log_all_runs(*, mlruns_dir: Path | str = MLRUNS_DIR) -> dict: if sa_path.exists(): sa = pd.read_csv(sa_path) for _, row in sa.iterrows(): - with mlflow.start_run( + with _record( + mlflow, run_name=f"selector_ablation|{row['selector']}|{row['model']}|f{int(row['fold'])}" ): _safe_log_param(mlflow, "selector", row["selector"]) @@ -119,7 +139,7 @@ def log_all_runs(*, mlruns_dir: Path | str = MLRUNS_DIR) -> dict: if not path.exists(): continue meta = json.loads(path.read_text()) - with mlflow.start_run(run_name=name): + with _record(mlflow, run_name=name): _safe_log_param(mlflow, "experiment", name) for k, v in meta.items(): if isinstance(v, (int, float)): @@ -146,7 +166,7 @@ def log_all_runs(*, mlruns_dir: Path | str = MLRUNS_DIR) -> dict: continue d = pd.read_csv(path) for _, row in d.iterrows(): - with mlflow.start_run(run_name=f"{name}|row{int(row.name)}"): + with _record(mlflow, run_name=f"{name}|row{int(row.name)}"): _safe_log_param(mlflow, "experiment", name) for col in d.columns: val = row[col] diff --git a/tests/unit/test_mlflow_logging.py b/tests/unit/test_mlflow_logging.py new file mode 100644 index 0000000..6c6a486 --- /dev/null +++ b/tests/unit/test_mlflow_logging.py @@ -0,0 +1,104 @@ +"""`log_all_runs` must work inside the pipeline's own MLflow run. + +HAL-4159. `cancerag.pipeline` opens a run for the pipeline as a whole, then +the `advanced` step calls `log_all_runs`, which opened its own *top-level* +run per CV fold. MLflow refuses a second active top-level run: + + Exception: Run with UUID ... is already active. To start a new run, first + end the current run with mlflow.end_run(). To start a nested run, call + start_run with nested=True + +So `advanced` died on its final phase and `advanced_summary.json` was never +written -- while `optuna_trials.csv` and the tuned models sat on disk, which +made the step look like it had simply not run. +""" + +from __future__ import annotations + +from pathlib import Path + +import pandas as pd +import pytest + +from cancerag.ml import mlflow_logging as ml + + +@pytest.fixture +def results(tmp_path, monkeypatch) -> Path: + """A minimal ml_models tree with one CV row and one extras row.""" + models = tmp_path / "ml_models" + (models / "extras").mkdir(parents=True) + pd.DataFrame([{ + "model": "lightgbm", "split": "scaffold_kfold", "fold": 0, "seed": 42, + "macro_f1": 0.5, "macro_f1_ci_lo": 0.4, "macro_f1_ci_hi": 0.6, + "balanced_accuracy": 0.5, "accuracy": 0.5, "n_train": 400, "n_test": 78, + }]).to_csv(models / "cv_results_long.csv", index=False) + pd.DataFrame([{"n_train": 100, "macro_f1": 0.3}]).to_csv( + models / "extras" / "learning_curve.csv", index=False) + monkeypatch.setattr(ml, "ML_MODELS", models) + return tmp_path + + +@pytest.mark.unit +def test_logs_runs_standalone(results): + """Run directly, nothing active: the runs stay top-level, as before.""" + import mlflow + + assert mlflow.active_run() is None + out = ml.log_all_runs(mlruns_dir=results / "mlruns") + assert out["n_runs_logged"] == 2 + assert mlflow.active_run() is None, "must not leave a run open" + + +@pytest.mark.unit +def test_logs_runs_inside_an_active_run(results): + """The case that killed the advanced step. + + Before the fix this raised "Run with UUID ... is already active" on the + first CV row. + """ + import mlflow + + mlflow.set_tracking_uri(f"file://{(results / 'mlruns').resolve()}") + mlflow.set_experiment("cancerag-pipeline") + with mlflow.start_run(run_name="pipeline"): + parent = mlflow.active_run().info.run_id + out = ml.log_all_runs(mlruns_dir=results / "mlruns") + assert out["n_runs_logged"] == 2 + still = mlflow.active_run() + assert still is not None, "the outer run must survive" + assert still.info.run_id == parent, ( + "log_all_runs must not close or replace the pipeline's run" + ) + assert mlflow.active_run() is None + + +@pytest.mark.unit +def test_nested_flag_follows_whether_a_run_is_active(monkeypatch): + """`_record` decides nesting from the live state, not a constant.""" + seen = {} + + class FakeMlflow: + def __init__(self, active): + self._active = active + + def active_run(self): + return object() if self._active else None + + def start_run(self, **kw): + seen.update(kw) + return kw + + ml._record(FakeMlflow(active=False), run_name="x") + assert seen["nested"] is False + ml._record(FakeMlflow(active=True), run_name="x") + assert seen["nested"] is True + + +@pytest.mark.unit +def test_every_call_site_goes_through_the_helper(): + """A sixth `mlflow.start_run` added later would reintroduce the bug.""" + src = Path(ml.__file__).read_text() + assert "with mlflow.start_run(" not in src, ( + "open runs through _record so they nest under the pipeline's run" + )