From 142ed710b79e58785f7971268214a2c5d38734ff Mon Sep 17 00:00:00 2001
From: Liangsheng Yin <hnyls2002@gmail.com>
Date: Wed, 13 May 2026 23:53:49 -0700
Subject: [PATCH] [CI] Support new-style `register_cuda_ci(stage=,
 runner_config=)` in slash handler + est-time updater (#25248)

---
 scripts/ci/update_est_time.py             | 18 ++++++++--
 scripts/ci/utils/slash_command_handler.py | 44 ++++++++++++++---------
 2 files changed, 44 insertions(+), 18 deletions(-)

diff --git a/scripts/ci/update_est_time.py b/scripts/ci/update_est_time.py
index 62bb00e9e..ce19e7113 100755
--- a/scripts/ci/update_est_time.py
+++ b/scripts/ci/update_est_time.py
@@ -215,11 +215,25 @@ def update_est_times(p90s, dry_run=False):
 
         for suite, backend, p90 in entries:
             # Match registration calls with this specific backend and suite.
-            # Handles: register_cuda_ci(est_time=300, suite="stage-c-test-4-gpu-h100")
-            pattern = re.compile(
+            # Two styles:
+            #   legacy: register_X_ci(est_time=N, suite="stage-Y-test-Z")
+            #   new:    register_X_ci(est_time=N, stage="stage-Y", runner_config="Z")
+            # New-style files all use the canonical `stage=` then `runner_config=` order.
+            legacy_pattern = re.compile(
                 rf"(register_{backend}_ci\(est_time=)(\d+)"
                 rf'(,\s*suite="{re.escape(suite)}")'
             )
+            pattern = legacy_pattern if legacy_pattern.search(new_content) else None
+            if pattern is None and "-test-" in suite:
+                stage, _, rc = suite.partition("-test-")
+                new_style_pattern = re.compile(
+                    rf"(register_{backend}_ci\(est_time=)(\d+)"
+                    rf'(,\s*stage="{re.escape(stage)}",\s*runner_config="{re.escape(rc)}")'
+                )
+                if new_style_pattern.search(new_content):
+                    pattern = new_style_pattern
+            if pattern is None:
+                continue
             match = pattern.search(new_content)
             if not match:
                 continue
diff --git a/scripts/ci/utils/slash_command_handler.py b/scripts/ci/utils/slash_command_handler.py
index 1c08a60f3..4c29a996f 100644
--- a/scripts/ci/utils/slash_command_handler.py
+++ b/scripts/ci/utils/slash_command_handler.py
@@ -789,6 +789,30 @@ def detect_multimodal_suite(file_path):
     return MULTIMODAL_DEFAULT_RUNNER, None
 
 
+def _extract_suite(content, func_name):
+    """Pull a suite name out of a `register_{cuda,cpu}_ci(...)` call.
+
+    Two styles are supported:
+      1. legacy: register_cuda_ci(..., suite="stage-X-test-Y")
+      2. new:    register_cuda_ci(..., stage="stage-X", runner_config="Y")
+                 -> suite = f"{stage}-test-{runner_config}"
+    """
+    legacy = re.search(
+        rf'^[^#\n]*{func_name}\([^)]*suite\s*=\s*["\']([^"\']+)["\']',
+        content,
+        re.MULTILINE,
+    )
+    if legacy:
+        return legacy.group(1)
+    args = re.search(rf"^[^#\n]*{func_name}\(([^)]*)\)", content, re.MULTILINE)
+    if args:
+        stage_m = re.search(r'stage\s*=\s*["\']([^"\']+)["\']', args.group(1))
+        rc_m = re.search(r'runner_config\s*=\s*["\']([^"\']+)["\']', args.group(1))
+        if stage_m and rc_m:
+            return f"{stage_m.group(1)}-test-{rc_m.group(1)}"
+    return None
+
+
 def detect_suite(file_path_from_test):
     """
     Read a test file and extract the suite from register_cuda_ci or register_cpu_ci.
@@ -799,14 +823,8 @@ def detect_suite(file_path_from_test):
     with open(full_path, "r") as f:
         content = f.read()
 
-    # Try CUDA first
-    match = re.search(
-        r'^[^#\n]*register_cuda_ci\([^)]*suite\s*=\s*["\']([^"\']+)["\']',
-        content,
-        re.MULTILINE,
-    )
-    if match:
-        suite = match.group(1)
+    suite = _extract_suite(content, "register_cuda_ci")
+    if suite:
         runner = CUDA_SUITE_TO_RUNNER.get(suite)
         if not runner:
             known = ", ".join(f"`{s}`" for s in sorted(CUDA_SUITE_TO_RUNNER))
@@ -823,14 +841,8 @@ def detect_suite(file_path_from_test):
         use_deepep = suite in DEEPEP_SUITES
         return suite, runner, use_deepep, False, None
 
-    # Try CPU
-    match = re.search(
-        r'^[^#\n]*register_cpu_ci\([^)]*suite\s*=\s*["\']([^"\']+)["\']',
-        content,
-        re.MULTILINE,
-    )
-    if match:
-        suite = match.group(1)
+    suite = _extract_suite(content, "register_cpu_ci")
+    if suite:
         return suite, "ubuntu-latest", False, True, None
 
     return (