Add MLflow logging to correlations and eval stages

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

rules/correlations.smk

@@ -27,6 +27,8 @@ rule correlations:
         circuits=" ".join(config.get("correlation", {}).get("circuits", ["QK"])),
         metrics=" ".join(config.get("correlation", {}).get("metrics", ["frob_cosine"])),
         cache_dir=config["cache_dir"],
+        mlflow_uri=config["mlflow_uri"],
+        mlflow_experiment=config["mlflow_experiment"],
     shell:
         """
         python scripts/run_correlations.py \
@@ -35,5 +37,7 @@ rule correlations:
             --out-dir correlations \
             --cache-dir {params.cache_dir} \
             --circuits {params.circuits} \
-            --metrics {params.metrics}
+            --metrics {params.metrics} \
+            --mlflow-uri {params.mlflow_uri} \
+            --mlflow-experiment {params.mlflow_experiment}
         """

rules/eval.smk

@@ -40,6 +40,8 @@ rule eval:
     params:
         p=_eval_params,
         cache_dir=config["cache_dir"],
+        mlflow_uri=config["mlflow_uri"],
+        mlflow_experiment=config["mlflow_experiment"],
     shell:
         """
         python scripts/run_eval.py \
@@ -49,6 +51,8 @@ rule eval:
             --pile-tokens {params.p[pile_tokens]} \
             --out {output} \
             --cache {params.cache_dir} \
+            --mlflow-uri {params.mlflow_uri} \
+            --mlflow-experiment {params.mlflow_experiment} \
             {params.p[device_flag]} \
             {params.p[pile_cache_flag]}
         """

scripts/run_correlations.py

@@ -10,6 +10,10 @@ Examples:
 """
 
 import argparse
+import subprocess
+import time
+
+import mlflow
 
 from transformer_analysis.pair_pipeline import run_multi_circuit_analysis
 
@@ -24,23 +28,50 @@ def parse_args():
     p.add_argument("--metrics", nargs="+", default=["frob_cosine", "pearson_corr", "hist_jensen_shannon"])
     p.add_argument("--device", default=None)
     p.add_argument("--max-workers", type=int, default=4)
+    p.add_argument("--mlflow-uri", default="file:./mlruns")
+    p.add_argument("--mlflow-experiment", default="production")
     return p.parse_args()
 
 
+def _git_sha():
+    try:
+        return subprocess.check_output(["git", "rev-parse", "--short", "HEAD"], text=True).strip()
+    except Exception:
+        return "unknown"
+
+
 def main():
     args = parse_args()
     revision = args.revision or None
 
-    run_multi_circuit_analysis(
-        model_name=args.model,
-        revision=revision,
-        circuits=tuple(args.circuits),
-        metrics=tuple(args.metrics),
-        cache_dir=args.cache_dir,
-        out_dir=args.out_dir,
-        device=args.device,
-        max_workers=args.max_workers,
-    )
+    mlflow.set_tracking_uri(args.mlflow_uri)
+    mlflow.set_experiment(args.mlflow_experiment)
+
+    rev_label = revision or "main"
+    run_name = f"{args.model}-{rev_label}-correlations"
+
+    with mlflow.start_run(run_name=run_name):
+        mlflow.log_params({
+            "model": args.model,
+            "revision": rev_label,
+            "circuits": ",".join(args.circuits),
+            "metrics": ",".join(args.metrics),
+            "max_workers": args.max_workers,
+            "git_sha": _git_sha(),
+        })
+
+        t0 = time.time()
+        run_multi_circuit_analysis(
+            model_name=args.model,
+            revision=revision,
+            circuits=tuple(args.circuits),
+            metrics=tuple(args.metrics),
+            cache_dir=args.cache_dir,
+            out_dir=args.out_dir,
+            device=args.device,
+            max_workers=args.max_workers,
+        )
+        mlflow.log_metric("wall_time_s", time.time() - t0)
 
 
 if __name__ == "__main__":

scripts/run_eval.py

@@ -4,8 +4,10 @@
 import argparse
 import os
 import sys
+import time
 from pathlib import Path
 
+import mlflow
 import pandas as pd
 
 sys.path.insert(0, str(Path(__file__).parent.parent / "src"))
@@ -25,22 +27,47 @@ def main():
     parser.add_argument("--out", type=str, required=True)
     parser.add_argument("--cache", type=str, default="model_cache")
     parser.add_argument("--device", type=str, default=None, choices=["cuda", "mps", "cpu"])
+    parser.add_argument("--mlflow-uri", default="file:./mlruns")
+    parser.add_argument("--mlflow-experiment", default="production")
 
     args = parser.parse_args()
+    revision = args.revision or None
+    rev_label = revision or "main"
 
-    result = evaluate_model(
-        model_name=args.model, revision=args.revision or None,
-        corpus=args.corpus, pile_tokens=args.pile_tokens,
-        cache_dir=args.cache, device_str=args.device,
-        stride=args.stride, max_tokens=args.max_tokens,
-        pile_cache=args.pile_cache,
-    )
-    print(f"  ppl={result['perplexity']:.2f}  bpb={result['bpb']:.4f}")
-
-    df = pd.DataFrame(to_long_format(result))
-    os.makedirs(os.path.dirname(args.out) or ".", exist_ok=True)
-    df.to_parquet(args.out, index=False)
-    print(f"  Saved {len(df)} rows → {args.out}")
+    mlflow.set_tracking_uri(args.mlflow_uri)
+    mlflow.set_experiment(args.mlflow_experiment)
+
+    with mlflow.start_run(run_name=f"{args.model}-{rev_label}-eval"):
+        mlflow.log_params({
+            "model": args.model,
+            "revision": rev_label,
+            "corpus": args.corpus,
+            "pile_tokens": args.pile_tokens,
+            "stride": args.stride,
+        })
+
+        t0 = time.time()
+        result = evaluate_model(
+            model_name=args.model, revision=revision,
+            corpus=args.corpus, pile_tokens=args.pile_tokens,
+            cache_dir=args.cache, device_str=args.device,
+            stride=args.stride, max_tokens=args.max_tokens,
+            pile_cache=args.pile_cache,
+        )
+        wall_time = time.time() - t0
+
+        print(f"  ppl={result['perplexity']:.2f}  bpb={result['bpb']:.4f}")
+        mlflow.log_metrics({
+            "perplexity": result["perplexity"],
+            "nll": result["nll"],
+            "bpb": result["bpb"],
+            "wall_time_s": wall_time,
+        })
+
+        df = pd.DataFrame(to_long_format(result))
+        os.makedirs(os.path.dirname(args.out) or ".", exist_ok=True)
+        df.to_parquet(args.out, index=False)
+        print(f"  Saved {len(df)} rows → {args.out}")
 
 
 if __name__ == "__main__":