[https://nvbugs/6490049][test] Unwaive GB300 Kimi disagg e2e

jenkins/scripts/perf/submit.py

@@ -33,9 +33,12 @@
 """
 
 import argparse
+import heapq
+import json
 import math
 import os
 import re
+import shlex
 import sys
 
 import yaml
@@ -63,26 +66,166 @@ def _import_precheck_config(llm_src):
 # --------------------------------------------------------------------------- #
 # Test list parsing
 # --------------------------------------------------------------------------- #
-def parse_test_case_name(test_list_path, llm_src, split_group=0):
-    """Parse the selected line of the test list.
-
-    Returns (config_yaml_path, server_name, benchmark_mode, runtime_mode).
-    See the module docstring for the supported test name shapes.
-    """
+def _read_test_list_lines(test_list_path):
     with open(test_list_path, "r") as f:
         lines = [line.strip() for line in f if line.strip()]
-
     if not lines:
         raise ValueError(f"Test list is empty: {test_list_path}")
+    return lines
+
+
+def _pytest_command_tokens(script_prefix_lines):
+    pytest_command_line = next(
+        (line for line in script_prefix_lines if "export pytestCommand=" in line), ""
+    )
+    if not pytest_command_line:
+        return []
+    command = pytest_command_line.split("=", 1)[1].strip()
+    if len(command) >= 2 and command[0] == command[-1] and command[0] in ('"', "'"):
+        command = command[1:-1]
+    return shlex.split(command)
+
+
+def _pytest_option(tokens, option):
+    for index, token in enumerate(tokens):
+        if token == option:
+            return tokens[index + 1] if index + 1 < len(tokens) else None
+        if token.startswith(f"{option}="):
+            return token.split("=", 1)[1]
+    return None
+
+
+def _test_nodeid(test_line):
+    """Strip test-list markers from a line, matching pytest's selected nodeid."""
+    return re.split(
+        r"\s+(?:XFAIL|SKIP|UNSTABLE|TIMEOUT)(?:\s|$)",
+        test_line,
+        maxsplit=1,
+    )[0]
+
+
+def _load_pytest_split_durations(tokens, llm_src):
+    durations_option = _pytest_option(tokens, "--durations-path")
+    if durations_option:
+        durations_path = durations_option
+        if not os.path.exists(durations_path):
+            durations_path = os.path.join(
+                llm_src,
+                "tests",
+                "integration",
+                "defs",
+                os.path.basename(durations_option),
+            )
+    else:
+        durations_path = os.path.join(llm_src, "tests", "integration", "defs", ".test_durations")
+
+    try:
+        with open(durations_path, "r") as durations_file:
+            durations = json.load(durations_file)
+    except FileNotFoundError:
+        durations = {}
+    if isinstance(durations, list):
+        durations = dict(durations)
+    if not isinstance(durations, dict):
+        raise ValueError(f"Invalid pytest-split durations file: {durations_path}")
+    return durations, durations_path
+
+
+def _select_least_duration_group(lines, durations, splits, group):
+    """Mirror pytest-split's LeastDurationAlgorithm exactly."""
+    if splits < 1:
+        raise ValueError(f"pytest --splits must be >= 1, got {splits}")
+    if group < 1 or group > splits:
+        raise ValueError(f"pytest --group must be in [1, {splits}], got {group}")
+
+    nodeids = [_test_nodeid(line) for line in lines]
+    relevant_durations = {
+        nodeid: float(durations[nodeid]) for nodeid in nodeids if nodeid in durations
+    }
+    average_duration = (
+        sum(relevant_durations.values()) / len(relevant_durations) if relevant_durations else 1.0
+    )
+    items = [
+        (line, nodeid, relevant_durations.get(nodeid, average_duration), original_index)
+        for original_index, (line, nodeid) in enumerate(zip(lines, nodeids))
+    ]
 
-    if split_group > 0:
+    # pytest-split first sorts by item name, then performs a stable descending
+    # duration sort. It greedily places each item in the least-loaded group.
+    items.sort(key=lambda item: item[1])
+    items.sort(key=lambda item: item[2], reverse=True)
+    selected = [[] for _ in range(splits)]
+    group_heap = [(0.0, group_index) for group_index in range(splits)]
+    heapq.heapify(group_heap)
+    for line, _nodeid, duration, original_index in items:
+        group_duration, group_index = heapq.heappop(group_heap)
+        selected[group_index].append((original_index, line))
+        heapq.heappush(group_heap, (group_duration + duration, group_index))
+
+    return [line for _original_index, line in sorted(selected[group - 1], key=lambda item: item[0])]
+
+
+def select_test_case_line(test_list_path, llm_src, script_prefix_lines, split_group=0):
+    """Select the same test as the pytest-split shard in ``pytestCommand``."""
+    lines = _read_test_list_lines(test_list_path)
+    if split_group <= 0:
+        return lines[0]
+
+    tokens = _pytest_command_tokens(script_prefix_lines)
+    splits_option = _pytest_option(tokens, "--splits")
+    group_option = _pytest_option(tokens, "--group")
+    algorithm = _pytest_option(tokens, "--splitting-algorithm")
+    if splits_option is None or group_option is None:
         if split_group > len(lines):
             raise ValueError(
                 f"split_group {split_group} exceeds number of tests in test list ({len(lines)})"
             )
-        line = lines[split_group - 1]
+        return lines[split_group - 1]
+    if algorithm != "least_duration":
+        raise ValueError(
+            "Multi-node perf launcher only supports pytest-split's least_duration "
+            f"algorithm, got {algorithm!r}"
+        )
+
+    splits = int(splits_option)
+    pytest_group = int(group_option)
+    if pytest_group != split_group:
+        raise ValueError(
+            f"submit.py split_group={split_group} disagrees with pytest --group={pytest_group}"
+        )
+    durations, durations_path = _load_pytest_split_durations(tokens, llm_src)
+    selected = _select_least_duration_group(lines, durations, splits, pytest_group)
+    if len(selected) != 1:
+        raise ValueError(
+            "Multi-node perf launch requires exactly one test in each pytest-split "
+            f"group, but group {pytest_group}/{splits} selected {len(selected)} tests "
+            f"using {durations_path}: {selected}"
+        )
+    print(
+        f"Selected pytest-split group {pytest_group}/{splits} test using "
+        f"{durations_path}: {_test_nodeid(selected[0])}"
+    )
+    return selected[0]
+
+
+def parse_test_case_name(test_list_path, llm_src, split_group=0, selected_line=None):
+    """Parse the selected line of the test list.
+
+    Returns (config_yaml_path, server_name, benchmark_mode, runtime_mode).
+    See the module docstring for the supported test name shapes.
+    """
+    if selected_line is not None:
+        line = selected_line
     else:
-        line = lines[0]
+        lines = _read_test_list_lines(test_list_path)
+        if split_group > 0:
+            if split_group > len(lines):
+                raise ValueError(
+                    f"split_group {split_group} exceeds number of tests in test list ({len(lines)})"
+                )
+            line = lines[split_group - 1]
+        else:
+            line = lines[0]
 
     if "[" not in line or "]" not in line:
         raise ValueError(f"Invalid test list format. Expected name with brackets: {line}")
@@ -484,7 +627,9 @@ def main():
         "--split-group",
         type=int,
         default=0,
-        help="1-indexed split group id. Selects the N-th test from the test list.",
+        help=(
+            "1-indexed pytest-split group id. Selects the same duration-balanced test as pytest."
+        ),
     )
     parser.add_argument("--stage-name", default="", help="Stage name (for logging / GPU detect)")
     parser.add_argument(
@@ -497,19 +642,29 @@ def main():
 
     args = parser.parse_args()
 
+    with open(args.script_prefix, "r") as f:
+        script_prefix_content = f.read()
+    script_prefix_lines = script_prefix_content.split("\n")
+
+    selected_test_line = select_test_case_line(
+        args.test_list,
+        args.llm_src,
+        script_prefix_lines,
+        args.split_group,
+    )
     config_yaml, server_name, benchmark_mode, runtime_mode = parse_test_case_name(
-        args.test_list, args.llm_src, args.split_group
+        args.test_list,
+        args.llm_src,
+        args.split_group,
+        selected_line=selected_test_line,
     )
 
     with open(config_yaml, "r") as f:
         config = yaml.safe_load(f)
 
-    # Recover test_case_name (the bracketed pytest test id) for the per-test
-    # output dir — same line/split logic as parse_test_case_name.
-    with open(args.test_list, "r") as f:
-        lines = [ln.strip() for ln in f if ln.strip()]
-    sel = lines[args.split_group - 1] if args.split_group > 0 else lines[0]
-    test_case_name = sel.split("[")[-1].split("]")[0] if "[" in sel else ""
+    test_case_name = (
+        selected_test_line.split("[")[-1].split("]")[0] if "[" in selected_test_line else ""
+    )
 
     hardware_config = get_hardware_config(config, runtime_mode, benchmark_mode, server_name)
     env_config = get_env_config(config, runtime_mode, benchmark_mode, server_name)
@@ -522,10 +677,6 @@ def main():
     print(f"Environment configuration: {env_config}")
     print(f"Benchmark configuration: {benchmark_config}")
 
-    with open(args.script_prefix, "r") as f:
-        script_prefix_content = f.read()
-    script_prefix_lines = script_prefix_content.split("\n")
-
     with open(args.srun_args, "r") as f:
         srun_args_content = f.read()
     srun_args_lines = srun_args_content.split()

tests/integration/defs/perf/test_perf_sanity.py

@@ -104,6 +104,10 @@ def ensure_bench_serving_repo() -> str:
 # once EVERY ctx/gen worker has finished model load + autotune + warmup.
 AGG_SERVER_READY_TIMEOUT = 1800
 DISAGG_SERVER_READY_TIMEOUT = 3600
+# GEN workers normally reap within seconds after benchmark_status is written.
+# Keep this well below the whole-test timeout so a stuck multi-node srun cannot
+# turn the optional log-flush synchronization into a pytest/Slurm cancellation.
+GEN_LOG_SENTINEL_TIMEOUT = 120
 
 
 def server_ready_timeout(default: int, mode: str) -> int:
@@ -1348,7 +1352,11 @@ def wait_for_benchmark_ready(
                 )
             time.sleep(10)
 
-    def wait_for_gen_log_sentinels(self, poll_interval: float = 2.0) -> bool:
+    def wait_for_gen_log_sentinels(
+        self,
+        timeout: float = GEN_LOG_SENTINEL_TIMEOUT,
+        poll_interval: float = 2.0,
+    ) -> bool:
         """Block until every gen worker signals that its log is fully written.
 
         Each gen worker's srun in slurm_launch_draft.sh redirects all of its
@@ -1357,25 +1365,27 @@ def wait_for_gen_log_sentinels(self, poll_interval: float = 2.0) -> bool:
         flushed). The benchmark writes benchmark_status *before* calling this,
         which is what lets the gen srun exit — so this is not circular.
 
-        Returns True once all sentinels exist, or False if self.timeout is
-        reached first. On False the caller still parses whatever is on disk:
-        the sentinel is a correctness optimization against reading a
-        mid-flush log (nvbugs 6487036 / 6487040), never a hang risk for CI.
+        Returns True once all sentinels exist, or False if the dedicated
+        sentinel timeout is reached first. On False the caller still parses
+        whatever is on disk: the sentinel is a correctness optimization
+        against reading a mid-flush log, not a reason to consume the whole-test
+        timeout and trigger Slurm's kill-on-bad-exit cascade (nvbugs 6487036 /
+        6487040 / 6487038).
         """
         sentinels = [
             os.path.join(self.test_output_dir, f"gen_server_{i}.done")
             for i in range(self.num_gen_servers)
         ]
-        start_time = time.time()
+        start_time = time.monotonic()
         while True:
             missing = [p for p in sentinels if not os.path.exists(p)]
             if not missing:
                 print_info("All gen worker log sentinels present; log flush complete.")
                 return True
-            elapsed_time = time.time() - start_time
-            if elapsed_time > self.timeout:
+            elapsed_time = time.monotonic() - start_time
+            if elapsed_time > timeout:
                 print_info(
-                    f"Timeout ({self.timeout}s) waiting for gen worker log "
+                    f"Timeout ({timeout}s) waiting for gen worker log "
                     f"sentinels {missing}; parsing current log contents."
                 )
                 return False
@@ -1504,6 +1514,9 @@ def run_cmd(self, server_idx: int) -> List[str]:
             # the loop (as before) could read a truncated / not-yet-flushed log
             # and report a wrong mean (nvbugs 6487036 / 6487040).
             pending_device_step_time: List[dict] = []
+            collect_device_step_time = (
+                configs_for_idx is not None and configs_for_idx[2].benchmark_mode == "gen_only"
+            )
             try:
                 disagg_server_hostname, disagg_server_port = (
                     self._get_disagg_server_hostname_and_port(server_idx)
@@ -1535,11 +1548,15 @@ def run_cmd(self, server_idx: int) -> List[str]:
                         )
                         print_info(f"Starting benchmark. cmd is {client_cmd_with_port}")
 
-                        # Snapshot gen_server log sizes so the per-client
-                        # average covers only iterations driven by this client.
-                        gen_log_start_offsets = gen_worker_log_sizes(
-                            self.test_output_dir, self.num_gen_servers
-                        )
+                        # Snapshot gen_server log sizes so the gen_only
+                        # per-client average covers only iterations driven by
+                        # this client. Other modes do not emit this metric and
+                        # must not wait for the GEN teardown sentinel.
+                        gen_log_start_offsets = None
+                        if collect_device_step_time:
+                            gen_log_start_offsets = gen_worker_log_sizes(
+                                self.test_output_dir, self.num_gen_servers
+                            )
 
                         bench_env = copy.deepcopy(os.environ)
                         if client_config:
@@ -1554,16 +1571,17 @@ def run_cmd(self, server_idx: int) -> List[str]:
                             benchmark_ctx.write(output)
 
                         outputs.append(output)
-                        # Defer the gen-worker device-step-time parse until the
-                        # gen logs are flushed (see below); remember where to
-                        # write the summary back.
-                        pending_device_step_time.append(
-                            {
-                                "output_index": len(outputs) - 1,
-                                "benchmark_file_path": benchmark_file_path,
-                                "start_offsets": gen_log_start_offsets,
-                            }
-                        )
+                        if collect_device_step_time:
+                            # Defer the gen-worker device-step-time parse until
+                            # the gen logs are flushed (see below); remember
+                            # where to write the summary back.
+                            pending_device_step_time.append(
+                                {
+                                    "output_index": len(outputs) - 1,
+                                    "benchmark_file_path": benchmark_file_path,
+                                    "start_offsets": gen_log_start_offsets,
+                                }
+                            )
                     else:
                         print_info(
                             f"Skipping perf benchmark for client {client_idx}: "
@@ -1599,11 +1617,10 @@ def run_cmd(self, server_idx: int) -> List[str]:
 
             # benchmark_status is written, so the gen workers can now stop and
             # their srun will exit and drop gen_server_{i}.done. Wait once for
-            # those sentinels (bounded by self.timeout), then parse each
-            # benchmark client's gen-worker device step time a single time: the
-            # flushed log is complete, so no settle polling is needed. Only
-            # gen_only runs emit prev_device_step_time; other modes parse to
-            # None and skip the summary line.
+            # those sentinels (bounded independently of the whole-test timeout),
+            # then parse each benchmark client's gen-worker device step time a
+            # single time. Only gen_only runs populate this queue; other modes
+            # skip both the sentinel wait and device-step-time parsing.
             if pending_device_step_time:
                 self.wait_for_gen_log_sentinels()
                 for record in pending_device_step_time:

tests/integration/test_lists/waives.txt

@@ -341,7 +341,6 @@ perf/test_perf_sanity.py::test_e2e[aggr_upload-qwen3_5_397b_fp4_blackwell-qwen3_
 perf/test_perf_sanity.py::test_e2e[aggr_upload-qwen3_5_397b_fp4_blackwell-qwen3_5_397b_fp4_tp4_8k1k] SKIP (https://nvbugs/6535767)
 perf/test_perf_sanity.py::test_e2e[aggr_upload-super_ad_blackwell-super_ad_ws1_1k1k] SKIP (https://nvbugs/6153575)
 perf/test_perf_sanity.py::test_e2e[disagg_upload-e2e-gb200_deepseek-r1-fp4_128k8k_con128_ctx1_pp8_gen1_dep16_eplb0_mtp1_ccb-NIXL] SKIP (https://nvbugs/6426890)
-perf/test_perf_sanity.py::test_e2e[disagg_upload-e2e-gb300_kimi-k25-thinking-fp4_8k1k_con4096_ctx1_dep4_gen1_dep16_eplb0_mtp0_ccb-NIXL] SKIP (https://nvbugs/6490049)
 perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb200_deepseek-v32-fp4_32k4k_con2048_ctx1_dep4_gen1_dep32_eplb288_mtp1_ccb-NIXL] SKIP (https://nvbugs/6374872)
 perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb200_deepseek-v32-fp4_8k1k_con4096_ctx1_dep4_gen1_dep32_eplb256_mtp0_ccb-NIXL] SKIP (https://nvbugs/6490049)
 perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb200_gpt-oss-120b-fp4_8k1k_con4_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL] SKIP (https://nvbugs/6490049)