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" + )