move runs_on + rdma into runner_configs.yml (#25264)

scripts/ci/runner_configs.py

@@ -1,14 +1,24 @@
-"""Emit a runner_config's setup details (install / artifact_version /
-install_timeout) in $GITHUB_OUTPUT format. Reads scripts/ci/runner_configs.yml.
-Called by .github/workflows/_pr-test-stage.yml.
+"""Emit runner_config setup for GitHub Actions $GITHUB_OUTPUT.
+
+runner_configs.py <runner_config>
+    Per-field `key=value` lines (install / artifact_version /
+    install_timeout / rdma_devices). `runs_on` is intentionally omitted —
+    it carries the `$b200_runner` sentinel and is resolved via --map.
+    Called per stage by _pr-test-stage.yml.
+
+runner_configs.py --map <b200_runner_label>
+    `runs_on_map={json}` — flat dict {runner_config: runs_on}, with
+    `$b200_runner` substituted. Called once by _pr-test-check-changes.yml.
 """
 
+import json
 import os
 import sys
 
 import yaml
 
 _YAML_PATH = os.path.join(os.path.dirname(__file__), "runner_configs.yml")
+_B200_SENTINEL = "$b200_runner"
 
 
 def load() -> dict:
@@ -16,12 +26,35 @@ def load() -> dict:
         return yaml.safe_load(f)["runner_configs"]
 
 
-if __name__ == "__main__":
-    if len(sys.argv) != 2:
-        sys.exit("usage: runner_configs.py <runner_config>")
-    rc = sys.argv[1]
-    config = load().get(rc)
-    if config is None:
+def _emit_single(rc: str) -> None:
+    # runs_on goes through --map (resolves $b200_runner). Suppress it here so a
+    # consumer can't accidentally read the raw sentinel value.
+    cfg = load().get(rc)
+    if cfg is None:
         sys.exit(f"unknown runner_config: {rc!r}")
-    for key, value in config.items():
+    for key, value in cfg.items():
+        if key == "runs_on":
+            continue
         print(f"{key}={value}")
+
+
+def _emit_map(b200_runner: str) -> None:
+    runs_on = {
+        name: (b200_runner if cfg.get("runs_on") == _B200_SENTINEL else cfg["runs_on"])
+        for name, cfg in load().items()
+    }
+    print(f"runs_on_map={json.dumps(runs_on, separators=(',', ':'))}")
+
+
+if __name__ == "__main__":
+    args = sys.argv[1:]
+    if len(args) == 1:
+        _emit_single(args[0])
+    elif len(args) == 2 and args[0] == "--map":
+        _emit_map(args[1])
+    else:
+        sys.exit(
+            "usage:\n"
+            "  runner_configs.py <runner_config>\n"
+            "  runner_configs.py --map <b200_runner_label>"
+        )