Implement MLPerf v6.1 RCP logging and deferred offline evaluation for DeepSWE

Author: uwlei

tests/examples/deepswe_mllog_utils_test.py

@@ -22,6 +22,7 @@
 import shutil
 import tempfile
 import types
+from typing import Any
 from unittest import mock
 
 import numpy as np
@@ -30,6 +31,15 @@
 from tunix.utils import mllog_utils
 
 
+def _read_mllog_events(path: str) -> list[dict[str, Any]]:
+  with open(path, "r", encoding="utf-8") as f:
+    return [
+        json.loads(line.split(":::MLLOG ", 1)[1])
+        for line in f
+        if ":::MLLOG " in line
+    ]
+
+
 @absltest.skipIf(mllog_utils.mllogger is None, "mlperf_logging is not installed")
 class MllogUtilsTest(absltest.TestCase):
 
@@ -602,7 +612,150 @@ def test_train_stop(self):
       content = f.read()
 
     self.assertIn('"key": "block_stop"', content)
-    self.assertIn('"key": "run_stop"', content)
+    # run_stop is emitted by the offline evaluator, not by train_stop.
+    self.assertNotIn('"key": "run_stop"', content)
+
+  def test_compute_val_start_step(self):
+    self.assertEqual(mllog_utils.compute_val_start_step(256), 18)
+    self.assertEqual(mllog_utils.compute_val_start_step(512), 10)
+    self.assertEqual(mllog_utils.compute_val_start_step(1024), 7)
+    self.assertEqual(mllog_utils.compute_val_start_step(256, 1), 1)
+    self.assertEqual(mllog_utils.compute_val_start_step(256, 0), 18)
+    with self.assertRaises(ValueError):
+      mllog_utils.compute_val_start_step(0)
+
+  def test_append_checkpoint_manifest_upserts_sorted_records(self):
+    manifest_path = os.path.join(
+        self.test_dir, "mllog", "eval_checkpoints.jsonl"
+    )
+    for step, ts_ms in ((19, 2000), (18, 1000), (18, 1500)):
+      mllog_utils.append_checkpoint_manifest(
+          manifest_path,
+          {
+              "step": step,
+              "checkpoint_path": f"gs://ckpt/{step}/model_params",
+              "timestamp_ms": ts_ms,
+              "samples_count": step * 256,
+              "val_start_at": 18,
+          },
+      )
+    with open(manifest_path, "r", encoding="utf-8") as f:
+      records = [json.loads(line) for line in f]
+    self.assertEqual([r["step"] for r in records], [18, 19])
+    self.assertEqual([r["timestamp_ms"] for r in records], [1500, 2000])
+
+  def test_configure_logger_downloads_existing_gcs_log_before_config(self):
+    calls = mock.MagicMock()
+    fake_mllogger = mock.MagicMock()
+    fake_mllogger.logger.handlers = []
+    with (
+        mock.patch.object(mllog_utils, "mllog", calls.mllog),
+        mock.patch.object(mllog_utils, "mllogger", fake_mllogger),
+        mock.patch.object(mllog_utils, "_is_master_process", return_value=True),
+        mock.patch.object(mllog_utils, "_gcs_target_path", None),
+        mock.patch.object(mllog_utils, "_local_log_path", None),
+        mock.patch.object(
+            mllog_utils, "_download_from_gcs_if_exists", calls.download
+        ),
+    ):
+      mllog_utils.configure_logger(metric_logger_dir="gs://b/mllog", seed=42)
+      local_path = mllog_utils._local_log_path  # pylint: disable=protected-access
+
+    self.assertEqual(
+        [c[0] for c in calls.mock_calls], ["download", "mllog.config"]
+    )
+    calls.download.assert_called_once_with("gs://b/mllog/seed_42.out", local_path)
+    self.assertEqual(
+        os.path.abspath(calls.mllog.config.call_args.kwargs["filename"]),
+        local_path,
+    )
+
+  def test_offline_eval_rcp_sequence_converged(self):
+    log_dir = self.test_dir
+    args = types.SimpleNamespace(batch_size=16, num_generations=16)
+    mllog_utils.configure_logger(metric_logger_dir=log_dir, seed=42)
+    mllog_utils.train_stop(args, step=19, time_ms=2000)
+    self.assertFalse(
+        mllog_utils.log_offline_eval_step(
+            step=18,
+            samples_count=4608,
+            eval_accuracy=0.65,
+            checkpoint_timestamp_ms=1000,
+        )
+    )
+    self.assertTrue(
+        mllog_utils.log_offline_eval_step(
+            step=19,
+            samples_count=4864,
+            eval_accuracy=0.70,
+            checkpoint_timestamp_ms=2000,
+            is_last_checkpoint=True,
+        )
+    )
+
+    events = _read_mllog_events(os.path.join(log_dir, "seed_42.out"))
+    self.assertEqual(
+        [e["key"] for e in events],
+        ["block_stop"] + ["eval_start", "eval_accuracy", "eval_stop"] * 2
+        + ["run_stop"],
+    )
+    self.assertEqual(events[0]["time_ms"], 2000)
+    self.assertEqual(events[0]["metadata"]["step"], 19)
+    self.assertEqual([events[2]["value"], events[5]["value"]], [0.65, 0.70])
+    run_stop = events[-1]
+    self.assertEqual(run_stop["time_ms"], 2000)
+    self.assertEqual(run_stop["metadata"]["status"], "success")
+    self.assertEqual(run_stop["metadata"]["samples_count"], 4864)
+
+  def test_offline_eval_rcp_sequence_not_converged(self):
+    log_dir = self.test_dir
+    mllog_utils.configure_logger(metric_logger_dir=log_dir, seed=42)
+    for step, ts_ms in ((18, 1000), (19, 2000), (20, 3000)):
+      self.assertFalse(
+          mllog_utils.log_offline_eval_step(
+              step=step,
+              samples_count=step * 256,
+              eval_accuracy=0.5,
+              checkpoint_timestamp_ms=ts_ms,
+              is_last_checkpoint=step == 20,
+          )
+      )
+
+    events = _read_mllog_events(os.path.join(log_dir, "seed_42.out"))
+    run_stops = [e for e in events if e["key"] == "run_stop"]
+    self.assertLen(run_stops, 1)
+    self.assertEqual(events[-1]["key"], "run_stop")
+    self.assertEqual(run_stops[0]["time_ms"], 3000)
+    self.assertEqual(run_stops[0]["metadata"]["status"], "aborted")
+    self.assertEqual(run_stops[0]["metadata"]["samples_count"], 5120)
+
+  def test_mlperf_6_1_0_init_print_disclosures(self):
+    fake_mllogger = mock.MagicMock()
+    with (
+        mock.patch.object(mllog_utils, "mllogger", fake_mllogger),
+        mock.patch.object(mllog_utils, "_is_master_process", return_value=True),
+        mock.patch.object(mllog_utils, "_flush_to_gcs_if_needed"),
+    ):
+      args = types.SimpleNamespace(
+          batch_size=16,
+          num_generations=16,
+          max_prompt_length=4096,
+          max_response_length=61440,
+      )
+      mllog_utils.init_print(args)
+    emitted = {
+        c.kwargs["key"]: c.kwargs["value"]
+        for c in fake_mllogger.event.call_args_list
+    }
+    self.assertEqual(emitted["eval_samples"], 251)
+    self.assertEqual(emitted["max_sequence_length"], 65536)
+    for key in (
+        "lowest_numerical_precision_in_linear",
+        "lowest_numerical_precision_in_attn",
+        "lowest_numerical_precision_in_comm",
+    ):
+      self.assertEqual(emitted[key], "bfloat16")
+    self.assertEqual(emitted["config_filename"], "qwen35_397b_grpo")
 
 
 if __name__ == "__main__":

tests/experimental/examples/deepswe_dist/eval_deepswe_test.py

@@ -168,6 +168,21 @@ def test_missing_and_failed_attempts_are_not_dropped(self):
     with self.assertRaisesRegex(ValueError, "duplicate"):
       eval_lib.summarize(rows + [rows[0]], ["a", "b"], 4)
 
+  def test_pass_at_4_is_reported_with_more_attempts(self):
+    rows = [
+        dict(
+            instance_id="a",
+            attempt=i,
+            reward=float(i == 5),
+            resolved=i == 5,
+            status="SUCCEEDED",
+        )
+        for i in range(8)
+    ]
+    summary = eval_lib.summarize(rows, ["a"], 8)
+    # Unbiased pass@k = 1 - C(n - c, k) / C(n, k) with n=8, c=1.
+    self.assertEqual(summary["pass_at_k"], {"1": 1 / 8, "4": 1 / 2, "8": 1.0})
+
   def test_reward_comes_from_trajectory_not_completion_status(self):
     response = types.SimpleNamespace(
         error=None,

tunix/experimental/examples/deepswe_dist/eval_deepswe.py

@@ -128,6 +128,59 @@ def parse_args(argv=None):
   )
   p.add_argument("--max_warmpool_size", type=int, default=1)
   p.add_argument("--output_dir", default="eval_results")
+  p.add_argument(
+      "--rcp_logging",
+      type=boolean,
+      nargs="?",
+      const=True,
+      default=os.environ.get("RCP_LOGGING", "0").lower() in ("1", "true"),
+      help="Enable MLPerf RCP (mllog) compliance logging.",
+  )
+  p.add_argument(
+      "--metric_logger_dir",
+      default=os.environ.get("METRIC_LOGGER_DIR", ""),
+      help="Directory or GCS URI for MLPerf RCP output (seed_<seed>.out).",
+  )
+  p.add_argument(
+      "--target_accuracy",
+      type=float,
+      default=float(os.environ.get("TARGET_ACCURACY", "0.69")),
+      help="Target evaluation accuracy for MLPerf RCP compliance logging.",
+  )
+  p.add_argument(
+      "--checkpoint_step",
+      type=int,
+      default=int(os.environ.get("CHECKPOINT_STEP", "0")),
+      help="Optimizer step corresponding to the evaluated checkpoint.",
+  )
+  p.add_argument(
+      "--checkpoint_timestamp_ms",
+      type=int,
+      default=(
+          int(os.environ["CHECKPOINT_TIMESTAMP_MS"])
+          if os.environ.get("CHECKPOINT_TIMESTAMP_MS", "").strip()
+          else None
+      ),
+      help=(
+          "Training epoch timestamp (ms) when the checkpoint weights were "
+          "updated, used to backdate run_stop."
+      ),
+  )
+  p.add_argument(
+      "--samples_count",
+      type=int,
+      default=int(os.environ.get("SAMPLES_COUNT", "0")),
+      help="Cumulative training samples at checkpoint_step.",
+  )
+  p.add_argument(
+      "--is_last_checkpoint",
+      type=boolean,
+      nargs="?",
+      const=True,
+      default=os.environ.get("IS_LAST_CHECKPOINT", "0").lower()
+      in ("1", "true"),
+      help="Whether this checkpoint is the final checkpoint in the manifest.",
+  )
   a = p.parse_args(argv)
   for name in (
       "mesh_fsdp",
@@ -268,7 +321,10 @@ def summarize(rows, instance_ids, attempts):
   total = len(instance_ids) * attempts
   solved = sum(row["resolved"] for row in rows)
   pass_at_k = {}
-  for k in sorted({1, attempts}):
+  ks = {1, attempts}
+  if attempts >= 4:
+    ks.add(4)
+  for k in sorted(ks):
     values = []
     for group in grouped.values():
       c = sum(row["resolved"] for row in group)
@@ -423,6 +479,13 @@ def load_entries(a):
 
 async def run_controller(a):
   from tunix.experimental.worker import remote_execution
+  from tunix.utils import mllog_utils
+
+  if getattr(a, "rcp_logging", False) and getattr(a, "metric_logger_dir", ""):
+    mllog_utils.configure_logger(
+        metric_logger_dir=a.metric_logger_dir,
+        seed=getattr(a, "seed", 42),
+    )
 
   entries = load_entries(a)
   run_id = (
@@ -458,6 +521,8 @@ async def run_controller(a):
   failure = None
   fleet = None
   entry_stream = entries
+  t_eval_start = None
+  eval_start_time_ms = None
 
   def record(row):
     writer.record(row)
@@ -479,6 +544,8 @@ async def ready(handle):
         raise ValueError(f"Worker/controller model settings differ: {profile}")
 
     await asyncio.gather(*(ready(h) for h in handles))
+    t_eval_start = time.monotonic()
+    eval_start_time_ms = time.time_ns() // 1_000_000
     if a.use_agent_sandbox:
       from examples.deepswe import sandbox_utils  # pylint: disable=import-outside-toplevel
 
@@ -535,12 +602,60 @@ async def ready(handle):
       from examples.deepswe import sandbox_utils  # pylint: disable=import-outside-toplevel
 
       await asyncio.to_thread(sandbox_utils.teardown_global_fleet)
+    validation_time = (
+        time.monotonic() - t_eval_start if t_eval_start is not None else None
+    )
     summary = summarize(
         all_rows,
         [str(e["instance_id"]) for e in entries],
         a.num_rollouts_per_instance,
     )
     summary["fatal_error"] = failure
+    expected_attempts = int(summary.get("expected_attempts") or 0)
+    error_attempts = int(summary.get("error_attempts") or 0)
+    eval_ok = (
+        failure is None
+        and bool(summary.get("complete"))
+        and (expected_attempts > 0 and error_attempts < expected_attempts)
+    )
+    pass_at_k = summary.get("pass_at_k") or {}
+    eval_accuracy = float(
+        pass_at_k.get(
+            "4",
+            pass_at_k.get(
+                str(getattr(a, "num_rollouts_per_instance", 4)),
+                summary.get("mean_reward", 0.0),
+            ),
+        )
+    )
+    target_acc = float(getattr(a, "target_accuracy", 0.69))
+    target_reached = bool(eval_ok and eval_accuracy >= target_acc)
+    rcp_logged = False
+    if getattr(a, "rcp_logging", False) and eval_ok:
+      mllog_utils.start_eval(
+          step=int(getattr(a, "checkpoint_step", 0)),
+          samples_count=int(getattr(a, "samples_count", 0)),
+          time_ms=eval_start_time_ms,
+      )
+      target_reached = mllog_utils.log_offline_eval_step(
+          step=int(getattr(a, "checkpoint_step", 0)),
+          samples_count=int(getattr(a, "samples_count", 0)),
+          eval_accuracy=eval_accuracy,
+          target_accuracy=target_acc,
+          checkpoint_timestamp_ms=getattr(a, "checkpoint_timestamp_ms", None),
+          is_last_checkpoint=bool(getattr(a, "is_last_checkpoint", False)),
+          validation_time=validation_time,
+          emit_start_eval=False,
+      )
+      rcp_logged = True
+    summary["rcp_logged"] = rcp_logged
+    summary["target_accuracy"] = target_acc
+    summary["target_reached"] = bool(target_reached)
+    summary["checkpoint_step"] = int(getattr(a, "checkpoint_step", 0))
+    summary["checkpoint_timestamp_ms"] = getattr(
+        a, "checkpoint_timestamp_ms", None
+    )
+    summary["samples_count"] = int(getattr(a, "samples_count", 0))
     try:
       writer.write("summary.json", summary)
     finally: