[benchmarking] Multiple fixes to stabilize the nightly benchmark suite

benchmarking/nightly-benchmark.yaml

@@ -24,6 +24,12 @@ paths:
 # datasets_path: "/path/to/datasets"
 # model_weights_path: "/path/to/model_weights"
 
+# Background per-GPU stats recorder. While each entry runs, polls all visible GPUs
+# (independent of CUDA_VISIBLE_DEVICES) at `interval_s` and writes one row per
+# (timestamp, GPU) to {session_entry_path}/gpustats.csv. Set interval_s to 0 to disable.
+gpu_stats_recorder:
+  interval_s: 1.0
+
 datasets:
   - name: "tinystories"
     formats:
@@ -173,9 +179,13 @@ entries:
         min_value: 2565
       - metric: number_of_domains_predicted
         exact_value: 26
-      # We choose two domains randomly and ensure that the counts are as expected
+      # We choose two domains randomly and ensure that the counts are as expected.
+      # domain_label_games_count is given a +/-1% tolerance window around the
+      # observed 149816 because the classifier output drifts run-to-run;
+      # domain_label_news_count remains exact pending further investigation.
       - metric: domain_label_games_count
-        exact_value: 149816
+        min_value: 148318
+        max_value: 151314
       - metric: domain_label_news_count
         exact_value: 2817
 
@@ -204,9 +214,13 @@ entries:
         min_value: 2755
       - metric: number_of_domains_predicted
         exact_value: 26
-      # We choose two domains randomly and ensure that the counts are as expected
+      # We choose two domains randomly and ensure that the counts are as expected.
+      # domain_label_games_count is given a +/-1% tolerance window around the
+      # observed 149816 because the classifier output drifts run-to-run;
+      # domain_label_news_count remains exact pending further investigation.
       - metric: domain_label_games_count
-        exact_value: 149816
+        min_value: 148318
+        max_value: 151314
       - metric: domain_label_news_count
         exact_value: 2817
 
@@ -273,7 +287,7 @@ entries:
       --input-filetype=parquet
       --text-field=raw_content
       --input-blocksize=3GiB
-    timeout_s: 500
+    timeout_s: 1500
     sink_data:
       - name: slack
         additional_metrics:
@@ -369,7 +383,7 @@ entries:
       --id-field=_curator_dedup_id
       --duplicate-id-field=_curator_dedup_id
       --blocksize=1.5GiB
-    timeout_s: 1100
+    timeout_s: 1800
     sink_data:
       - name: slack
         additional_metrics:
@@ -398,7 +412,7 @@ entries:
       --id-field=_curator_dedup_id
       --duplicate-id-field=_curator_dedup_id
       --blocksize=1.5GiB
-    timeout_s: 1500
+    timeout_s: 1800
     sink_data:
       - name: slack
         additional_metrics:
@@ -1019,7 +1033,7 @@ entries:
       --executor=ray_data
       --engine-kwargs='{"tensor_parallel_size": 1}'
       --autoscaling-config='{"min_replicas": 4, "max_replicas": 4}'
-    timeout_s: 700
+    timeout_s: 1200
     ray:
       num_cpus: 16
     sink_data:

benchmarking/run.py

@@ -22,6 +22,7 @@
 import time
 import traceback
 from collections.abc import Mapping
+from contextlib import nullcontext
 from pathlib import Path
 from typing import Any
 
@@ -42,6 +43,7 @@
 from runner.datasets import DatasetResolver
 from runner.entry import Entry
 from runner.env_capture import dump_env
+from runner.gpu_stats_recorder import GPUStatsRecorder
 from runner.path_resolver import PathResolver
 from runner.process import run_command_with_timeout
 from runner.ray_cluster import (
@@ -148,12 +150,13 @@ def check_requirements_update_results(result_data: dict[str, Any], requirements:
     return meets_requirements
 
 
-def run_entry(
+def run_entry(  # noqa: PLR0913
     entry: Entry,
     path_resolver: PathResolver,
     dataset_resolver: DatasetResolver,
     session_entry_path: Path,
     result_data: dict[str, Any],
+    gpu_stats_recorder_interval_s: float = 1.0,
 ) -> bool:
     # session_entry_path : This is the directory where benchmark results are stored
     # scratch_path : This is the directory provided to users for saving scratch/temp data; it'll be cleaned up after the entry is done if delete_scratch is True
@@ -212,14 +215,25 @@ def run_entry(
             warning_threshold_msg="used before benchmark started",
         )
         logger.info(f"\tRunning command {' '.join(cmd) if isinstance(cmd, list) else cmd}")
-        started_exec = time.time()
-        run_data = run_command_with_timeout(
-            command=cmd,
-            timeout=entry.timeout_s,
-            stdouterr_path=stdouterr_path,
-            run_id=run_id,
-            fancy=os.environ.get("CURATOR_BENCHMARKING_DEBUG", "0") == "0",
+        # Background poller writes per-GPU stats (utilization, memory, temperature, processes) to
+        # gpustats.csv every gpu_stats_recorder_interval_s seconds. Set to 0 in YAML to disable.
+        gpu_stats_recorder_ctx = (
+            GPUStatsRecorder(
+                output_path=session_entry_path / "gpustats.csv",
+                interval_s=gpu_stats_recorder_interval_s,
+            )
+            if gpu_stats_recorder_interval_s > 0
+            else nullcontext()
         )
+        started_exec = time.time()
+        with gpu_stats_recorder_ctx:
+            run_data = run_command_with_timeout(
+                command=cmd,
+                timeout=entry.timeout_s,
+                stdouterr_path=stdouterr_path,
+                run_id=run_id,
+                fancy=os.environ.get("CURATOR_BENCHMARKING_DEBUG", "0") == "0",
+            )
         ended_exec = time.time()
         logger.info("\tGPU stats (after):")
         warnings += log_gpu_stats(
@@ -281,7 +295,7 @@ def run_entry(
             shutil.rmtree(scratch_path, ignore_errors=True)
 
 
-def main() -> int:  # noqa: C901
+def main() -> int:  # noqa: C901, PLR0912, PLR0915
     parser = argparse.ArgumentParser(description="Runs the benchmarking application")
     parser.add_argument(
         "--config",
@@ -307,12 +321,50 @@ def main() -> int:  # noqa: C901
             "enabled entries will be run."
         ),
     )
+    parser.add_argument(
+        "--entries-exact",
+        default=None,
+        help=(
+            "Comma-separated list of exact entry names to run. Unlike --entries (a pytest "
+            "'-k' style substring expression), names here must match entry names exactly. "
+            "Every supplied name must correspond to a configured (enabled) entry; otherwise "
+            "the run fails with an error listing the unknown names. Useful for both "
+            "automated callers (e.g. CI per-job invocations) and users targeting a specific "
+            "set of entries by exact name. Mutually exclusive with --entries."
+        ),
+    )
     parser.add_argument(
         "--list",
         default=False,
         action="store_true",
         help="List entries to run and exit.",
     )
+    parser.add_argument(
+        "--strict-config-check",
+        default=False,
+        action="store_true",
+        help=(
+            "If set, fail with an error when an environment variable referenced in the "
+            "config is undefined or empty. By default, undefined env var references are "
+            "replaced with an empty string and a warning is logged."
+        ),
+    )
+    parser.add_argument(
+        "--viewer-url",
+        default=None,
+        help=(
+            "Run-viewer URL to surface in sinks (e.g. Slack parent message footer). "
+            "When set, the Slack sink renders a 'Results viewer' section linking to this URL."
+        ),
+    )
+    parser.add_argument(
+        "--run-reason",
+        default=None,
+        help=(
+            "Free-text reason for this run, recorded in env.json and surfaced in the Slack "
+            "environment block. Useful for audit trails on ad-hoc runs."
+        ),
+    )
     args = parser.parse_args()
 
     # Consolidate the configuration from all YAML files into a single dict
@@ -322,13 +374,36 @@ def main() -> int:  # noqa: C901
     try:
         assert_valid_config_dict(config_dict)
         config_dict = remove_disabled_blocks(config_dict)
-        config_dict = resolve_env_vars(config_dict)
+        config_dict = resolve_env_vars(config_dict, strict=args.strict_config_check)
     except ValueError as e:
         logger.error(f"Invalid configuration: {e}")
         return 1
 
+    if args.entries is not None and args.entries_exact is not None:
+        logger.error("--entries and --entries-exact are mutually exclusive")
+        return 1
+
+    entries_exact_list: list[str] | None = None
+    if args.entries_exact is not None:
+        entries_exact_list = [name.strip() for name in args.entries_exact.split(",") if name.strip()]
+        if not entries_exact_list:
+            logger.error("--entries-exact must contain at least one non-empty name")
+            return 1
+
     # Now that all YAML config files have been read, merged, and processed, create the Session object.
-    session = Session.from_dict(config_dict, args.entries)
+    try:
+        session = Session.from_dict(
+            config_dict,
+            entry_filter_expr=args.entries,
+            entries_exact=entries_exact_list,
+        )
+    except ValueError as e:
+        logger.error(str(e))
+        return 1
+
+    # GPU stats recorder config: polls every interval_s seconds while each entry runs.
+    # Default 1.0 (1 Hz). Set to 0 to disable.
+    gpu_stats_recorder_interval_s = float(config_dict.get("gpu_stats_recorder", {}).get("interval_s", 1.0))
 
     if args.list:
         for entry in session.entries:
@@ -344,6 +419,18 @@ def main() -> int:  # noqa: C901
     logger.info(f"Started session {session_name}...")
     env_dict = dump_env(session_obj=session, output_path=session_path)
 
+    # Record an optional free-text reason for the run (e.g. "regression check after MR !2442").
+    # Appears in env.json and the Slack environment block. No-op when unset.
+    if args.run_reason:
+        env_dict["run_reason"] = args.run_reason
+
+    # Surface an optional run-viewer URL in the Slack sink. Patch sink_config in-process
+    # so we don't have to teach the YAML config loader about a per-launch viewer URL.
+    if args.viewer_url:
+        for sink in session.sinks:
+            if getattr(sink, "name", None) == "slack":
+                sink.sink_config["viewer_url"] = args.viewer_url
+
     for sink in session.sinks:
         sink.initialize(session_name=session_name, session=session, env_dict=env_dict)
 
@@ -380,6 +467,7 @@ def main() -> int:  # noqa: C901
                 dataset_resolver=session.dataset_resolver,
                 session_entry_path=session_entry_path,
                 result_data=result_data,
+                gpu_stats_recorder_interval_s=gpu_stats_recorder_interval_s,
             )
 
         except Exception as e: