feat(inference): add validate-smoke-full sequence orchestrator

Author: Tobias Buck

docs/inference_workflows.rst

@@ -388,6 +388,14 @@ Real Data Runbook
       python scripts/generate_ifu_experiment_report.py \
         --output-dir outputs/ifu_realdata_smoke
 
+For a one-command staged run, use the sequence orchestrator:
+
+.. code-block:: bash
+
+   python scripts/run_ifu_experiment_sequence.py \
+     --config rubix/config/inference_experiment_realdata_scaffold.yml \
+     --output-root-dir outputs/ifu_sequence
+
 Benchmarking Full-IFU Optimization
 ----------------------------------
 

rubix/inference/__init__.py

@@ -20,6 +20,7 @@
     generate_ifu_experiment_report,
     normalize_experiment_config,
     run_ifu_experiment,
+    run_ifu_experiment_sequence,
     save_ifu_experiment_outputs,
     validate_ifu_experiment_inputs,
 )
@@ -136,6 +137,7 @@
     "summarize_predictive_cube_samples",
     "save_checkpoint",
     "run_ifu_experiment",
+    "run_ifu_experiment_sequence",
     "save_ifu_experiment_outputs",
     "validate_ifu_experiment_inputs",
     "value_and_grad",

rubix/inference/experiment.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import copy
 import csv
 import hashlib
 import json
@@ -1317,3 +1318,84 @@ def generate_ifu_experiment_report(output_dir: str) -> dict[str, Any]:
             report["artifacts"]["residual_products_keys"] = list(data.files)
 
     return report
+
+
+def run_ifu_experiment_sequence(
+    config: ExperimentConfigInput,
+    pipeline_factory: PipelineFactory = make_inference_pipeline,
+    run_validate: bool = True,
+    run_smoke: bool = True,
+    run_full: bool = True,
+    output_root_dir: Optional[str] = None,
+) -> dict[str, Any]:
+    """Run validate -> smoke -> full IFU workflow sequence.
+
+    Args:
+        config (ExperimentConfigInput): Mapping or YAML path following the
+            experiment schema consumed by :func:`normalize_experiment_config`.
+        pipeline_factory (PipelineFactory, optional): Pipeline builder used to
+            instantiate and prepare a Rubix pipeline. Defaults to
+            :func:`make_inference_pipeline`.
+        run_validate (bool, optional): Whether to run input validation phase.
+            Defaults to ``True``.
+        run_smoke (bool, optional): Whether to run smoke-only phase.
+            Defaults to ``True``.
+        run_full (bool, optional): Whether to run full optimization/VI phase.
+            Defaults to ``True``.
+        output_root_dir (Optional[str], optional): Optional root output
+            directory for phase artifacts. Defaults to ``None`` (use config).
+
+    Raises:
+        RuntimeError: If requested validation phase fails.
+
+    Returns:
+        dict[str, Any]: Per-phase outputs/statuses and output directories.
+    """
+    raw_cfg = read_yaml(config) if isinstance(config, str) else dict(config)
+    base_cfg = normalize_experiment_config(raw_cfg)
+
+    if output_root_dir is None:
+        output_root = Path(str(base_cfg["run"]["output_dir"]))
+    else:
+        output_root = Path(output_root_dir)
+    output_root.mkdir(parents=True, exist_ok=True)
+
+    sequence: dict[str, Any] = {
+        "output_root_dir": str(output_root),
+        "validate": None,
+        "smoke": None,
+        "full": None,
+    }
+
+    if run_validate:
+        report = validate_ifu_experiment_inputs(
+            config=base_cfg,
+            pipeline_factory=pipeline_factory,
+        )
+        (output_root / "validate_report.json").write_text(
+            json.dumps(_to_jsonable(report), indent=2),
+            encoding="utf-8",
+        )
+        sequence["validate"] = {"ok": bool(report["ok"]), "report": report}
+        if not report["ok"]:
+            raise RuntimeError("validation phase failed; see validate_report.json")
+
+    if run_smoke:
+        smoke_cfg = copy.deepcopy(base_cfg)
+        smoke_cfg["run"]["smoke_only"] = True
+        smoke_cfg["optimization"]["enabled"] = False
+        smoke_cfg["variational"]["enabled"] = False
+        smoke_out = run_ifu_experiment(smoke_cfg, pipeline_factory=pipeline_factory)
+        smoke_dir = output_root / "smoke"
+        save_ifu_experiment_outputs(smoke_out, str(smoke_dir))
+        sequence["smoke"] = {"output_dir": str(smoke_dir), "outputs": smoke_out}
+
+    if run_full:
+        full_cfg = copy.deepcopy(base_cfg)
+        full_cfg["run"]["smoke_only"] = False
+        full_out = run_ifu_experiment(full_cfg, pipeline_factory=pipeline_factory)
+        full_dir = output_root / "full"
+        save_ifu_experiment_outputs(full_out, str(full_dir))
+        sequence["full"] = {"output_dir": str(full_dir), "outputs": full_out}
+
+    return sequence

scripts/run_ifu_experiment_sequence.py

@@ -0,0 +1,50 @@
+#!/usr/bin/env python
+import argparse
+import json
+
+from rubix.inference.experiment import run_ifu_experiment_sequence
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Run IFU workflow sequence: validate -> smoke -> full."
+    )
+    parser.add_argument("--config", type=str, required=True)
+    parser.add_argument(
+        "--output-root-dir",
+        type=str,
+        default=None,
+        help="Optional root directory for validate/smoke/full outputs.",
+    )
+    parser.add_argument(
+        "--skip-validate",
+        action="store_true",
+        help="Skip validation phase.",
+    )
+    parser.add_argument(
+        "--skip-smoke",
+        action="store_true",
+        help="Skip smoke-only phase.",
+    )
+    parser.add_argument(
+        "--skip-full",
+        action="store_true",
+        help="Skip full optimization/VI phase.",
+    )
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    result = run_ifu_experiment_sequence(
+        config=args.config,
+        run_validate=not args.skip_validate,
+        run_smoke=not args.skip_smoke,
+        run_full=not args.skip_full,
+        output_root_dir=args.output_root_dir,
+    )
+    print(json.dumps(result, indent=2, default=str))
+
+
+if __name__ == "__main__":
+    main()

tests/test_inference_experiment.py

@@ -13,6 +13,7 @@
     generate_ifu_experiment_report,
     normalize_experiment_config,
     run_ifu_experiment,
+    run_ifu_experiment_sequence,
     save_ifu_experiment_outputs,
     validate_ifu_experiment_inputs,
 )
@@ -496,3 +497,49 @@ def test_generate_ifu_experiment_report_requires_summary(tmp_path):
         assert "summary.json" in str(exc)
     else:
         raise AssertionError("Expected FileNotFoundError when summary.json is missing")
+
+
+def test_run_ifu_experiment_sequence_writes_phase_outputs(tmp_path):
+    cube = np.ones((2, 2, 4), dtype=np.float32)
+    target_path = tmp_path / "target.npy"
+    np.save(target_path, cube)
+    np.save(tmp_path / "mask.npy", np.ones_like(cube))
+    np.save(tmp_path / "ivar.npy", np.ones_like(cube))
+    config_path = _write_config(tmp_path, str(target_path), str(tmp_path / "ckpt"))
+
+    def pipeline_factory(_cfg, _mode):
+        return PreparedSyntheticPipeline(jnp.asarray(cube))
+
+    result = run_ifu_experiment_sequence(
+        config=str(config_path),
+        pipeline_factory=pipeline_factory,
+        output_root_dir=str(tmp_path / "sequence"),
+    )
+    assert result["validate"]["ok"] is True
+    assert (tmp_path / "sequence" / "validate_report.json").exists()
+    assert (tmp_path / "sequence" / "smoke" / "summary.json").exists()
+    assert (tmp_path / "sequence" / "full" / "summary.json").exists()
+
+
+def test_run_ifu_experiment_sequence_raises_on_failed_validation(tmp_path):
+    cube = np.ones((2, 2, 4), dtype=np.float32)
+    target_path = tmp_path / "target.npy"
+    np.save(target_path, cube)
+    # Intentional mismatch for validation failure
+    np.save(tmp_path / "mask.npy", np.ones((2, 2, 5), dtype=np.float32))
+    np.save(tmp_path / "ivar.npy", np.ones_like(cube))
+    config_path = _write_config(tmp_path, str(target_path), str(tmp_path / "ckpt"))
+
+    def pipeline_factory(_cfg, _mode):
+        return PreparedSyntheticPipeline(jnp.asarray(cube))
+
+    try:
+        run_ifu_experiment_sequence(
+            config=str(config_path),
+            pipeline_factory=pipeline_factory,
+            output_root_dir=str(tmp_path / "sequence"),
+        )
+    except RuntimeError as exc:
+        assert "validation phase failed" in str(exc)
+    else:
+        raise AssertionError("Expected RuntimeError for failed validation phase")