[https://nvbugs/6487040][test] Wait for gen-log end-of-write sentinel before parsing per-iter step time

jenkins/scripts/perf/disaggregated/slurm_launch_draft.sh

@@ -55,12 +55,28 @@ for i in $(seq 0 $((numGenServers - 1))); do
     gen_world_size=$((nodesPerGenServer * gpusPerNodePerGenServer))
     export DISAGG_SERVING_TYPE="GEN_$i"
     export pytestCommand="$pytestCommandGENWorker"
-    srun "${srunArgs[@]}" --mpi=pmix --kill-on-bad-exit=1 \
-        -N $nodesPerGenServer \
-        -w "${genNodeLists[$i]}" \
-        --ntasks=$gen_world_size \
-        --ntasks-per-node=$gpusPerNodePerGenServer \
-        $runScript &> $testOutputDir/gen_server_$i.log &
+    # End-of-write sentinel: gen_server_$i.log is the srun's &> aggregate of
+    # every gen-worker rank (the per-iter prev_device_step_time lines the
+    # benchmark parses live only here, not in trtllm-serve.GEN_*.log). The
+    # file descriptor is owned by this srun, so the log is only guaranteed
+    # fully flushed once the srun is reaped. Run srun in the foreground of a
+    # backgrounded subshell and touch gen_server_$i.done immediately after it
+    # returns: the benchmark srun blocks on that sentinel before parsing, so
+    # it never reads a truncated / not-yet-flushed log (nvbugs 6487036 /
+    # 6487040). A stale sentinel from a re-run output dir is removed first.
+    # Note: srun is foreground inside the subshell (not `srun ... &` + a
+    # `kill -0` poll) so `touch` runs strictly after reap, with no
+    # late-zombie race that could either skip or prematurely fire the signal.
+    rm -f "$testOutputDir/gen_server_$i.done"
+    (
+        srun "${srunArgs[@]}" --mpi=pmix --kill-on-bad-exit=1 \
+            -N $nodesPerGenServer \
+            -w "${genNodeLists[$i]}" \
+            --ntasks=$gen_world_size \
+            --ntasks-per-node=$gpusPerNodePerGenServer \
+            $runScript &> $testOutputDir/gen_server_$i.log
+        touch "$testOutputDir/gen_server_$i.done"
+    ) &
     echo "Started gen server $i on ${genNodeLists[$i]}"
     sleep 5  # Wait for pyxis container namespace initialization to avoid race condition
 done

tests/integration/defs/perf/test_perf_sanity.py

@@ -213,20 +213,26 @@ def _scan_gen_worker_device_step_time(
     output_dir: str,
     num_gen_servers: int,
     start_offsets: Optional[List[int]] = None,
-) -> Tuple[List[Dict[int, Tuple[int, float]]], int]:
+) -> Tuple[List[Tuple[Dict[int, Tuple[int, float]], int, float]], int]:
     """Single-pass scan of the gen logs.
 
-    Returns (per_file_by_ngen, total_count):
-      - per_file_by_ngen: one dict per file that produced >=1 usable line,
-        mapping num_generation_tokens -> (count, Welford mean of
-        prev_device_step_time) over rows with iter >= 5 and a numeric
-        prev_device_step_time. Rows lacking num_generation_tokens on the same
-        line are skipped for the mean but still counted for settle detection.
+    Returns (per_file_scans, total_count):
+      - per_file_scans: one entry per file that produced >=1 usable row, each
+        a tuple (by_ngen, all_count, all_mean):
+          * by_ngen maps num_generation_tokens -> (count, Welford mean of
+            prev_device_step_time) over rows with iter >= 5, a numeric
+            prev_device_step_time, and a parseable num_generation_tokens on
+            the same line.
+          * all_count / all_mean are the count and Welford mean of
+            prev_device_step_time over ALL iter >= 5 numeric rows in the file,
+            including those whose num_generation_tokens did not parse. This is
+            the fallback aggregate used when a worker never emits a parseable
+            num_generation_tokens (nvbugs 6487036 / 6487040): PR #16298 began
+            requiring num_generation_tokens on every line, so a worker whose
+            states dict renders it as e.g. tensor(256) would drop to no
+            buckets and the metric would wrongly parse to None.
       - total_count: the number of iter >= 5 rows with a numeric
-        prev_device_step_time across all files. This is monotonic as new
-        lines flush across NFS (rows only get appended) so the caller can use
-        it as the settle signal without worrying that changes to a per-ngen
-        filter can make it drop.
+        prev_device_step_time across all files.
 
     Memory is O(distinct num_generation_tokens per file), a small constant
     in practice (steady-state plus a shrinking tail).
@@ -235,7 +241,7 @@ def _scan_gen_worker_device_step_time(
     (model load) write partial multibyte sequences that would otherwise raise
     UnicodeDecodeError mid-scan.
     """
-    per_file_by_ngen: List[Dict[int, Tuple[int, float]]] = []
+    per_file_scans: List[Tuple[Dict[int, Tuple[int, float]], int, float]] = []
     total_count = 0
     for i in range(num_gen_servers):
         log_path = os.path.join(output_dir, f"gen_server_{i}.log")
@@ -249,6 +255,8 @@ def _scan_gen_worker_device_step_time(
         )
 
         by_ngen: Dict[int, Tuple[int, float]] = {}
+        all_count = 0
+        all_mean = 0.0
         with open(log_path, errors="replace") as f:
             if seek_to:
                 f.seek(seek_to)
@@ -259,39 +267,48 @@ def _scan_gen_worker_device_step_time(
                 if int(m.group(1)) < 5:
                     continue
                 total_count += 1
+                dt = float(m.group(2))
+                # All-iter fallback aggregate (every usable row).
+                all_count += 1
+                all_mean += (dt - all_mean) / all_count
+                # Per-ngen bucket (only rows with a parseable ngen).
                 ngen_m = _NUM_GEN_TOKENS_RE.search(line)
                 if ngen_m is None:
                     continue
                 ngen = int(ngen_m.group(1))
-                dt = float(m.group(2))
                 count, mean = by_ngen.get(ngen, (0, 0.0))
                 count += 1
                 mean += (dt - mean) / count
                 by_ngen[ngen] = (count, mean)
-        if by_ngen:
-            per_file_by_ngen.append(by_ngen)
-    return per_file_by_ngen, total_count
+        if all_count:
+            per_file_scans.append((by_ngen, all_count, all_mean))
+    return per_file_scans, total_count
 
 
 def _mean_at_mode_ngen(
-    per_file_by_ngen: List[Dict[int, Tuple[int, float]]],
+    per_file_scans: List[Tuple[Dict[int, Tuple[int, float]], int, float]],
 ) -> Optional[float]:
-    """Aggregate per-file per-ngen buckets into a single mean.
+    """Aggregate per-file scans into a single mean.
 
     Within each file pick the num_generation_tokens value with the most
     iterations (the mode) and take its Welford mean; ties break to the
     largest ngen because the steady-state plateau is the upper of any tied
     clusters. Mode is more robust than strict == max — a one-off spike where
     a single iter's ngen briefly exceeds the sustained batch would otherwise
-    collapse the mean to 1-2 samples. Then average the per-file means across
-    workers. Returns None if no file had a usable row.
+    collapse the mean to 1-2 samples. When a file produced usable rows but no
+    parseable num_generation_tokens on any of them, fall back to the file's
+    all-iter mean so a present metric is never lost (nvbugs 6487036 /
+    6487040). Then average the per-file means across workers. Returns None if
+    no file had a usable row.
     """
     means: List[float] = []
-    for by_ngen in per_file_by_ngen:
-        if not by_ngen:
-            continue
-        _mode_ngen, (_count, mean) = max(by_ngen.items(), key=lambda kv: (kv[1][0], kv[0]))
-        means.append(mean)
+    for by_ngen, _all_count, all_mean in per_file_scans:
+        if by_ngen:
+            _mode_ngen, (_count, mean) = max(by_ngen.items(), key=lambda kv: (kv[1][0], kv[0]))
+            means.append(mean)
+        else:
+            # No parseable ngen anywhere in this worker; use the all-iter mean.
+            means.append(all_mean)
     if not means:
         return None
     return sum(means) / len(means)
@@ -301,8 +318,6 @@ def parse_gen_worker_device_step_time(
     output_dir: str,
     num_gen_servers: int,
     start_offsets: Optional[List[int]] = None,
-    settle_timeout: float = 90.0,
-    poll_interval: float = 3.0,
 ) -> Optional[float]:
     """Mean per-iter prev_device_step_time (ms) across all gen workers.
 
@@ -315,44 +330,26 @@ def parse_gen_worker_device_step_time(
     below the steady-state cost. Using the mode (rather than strict == max)
     is robust against a single iter whose ngen briefly spikes above the
     sustained batch, which would otherwise collapse the mean to 1-2 samples.
-    Returns None if no usable line is found in any file.
+    A worker whose num_generation_tokens never parses falls back to its
+    all-iter mean rather than being dropped to None. Returns None only if no
+    usable line is found in any file.
 
     When start_offsets is provided, only the bytes from start_offsets[i] to
     end-of-file are considered for gen_server_{i}.log — used to slice out a
     single client's iteration segment.
 
-    The gen worker writes gen_server_{i}.log on a different node than the
-    benchmark/pytest process, and the worker is kept alive (waiting on the
-    benchmark_status file) when this runs — so when the client returns, the
-    decode iterations are done but their log lines may still be flushing across
-    NFS. Reading once immediately can see zero iter>=5 lines and wrongly return
-    None. So poll the slice until the iter>=5 row count is non-zero AND
-    stable across two consecutive reads (flush drained), bounded by
-    settle_timeout. The settle signal is the raw iter>=5 row count (not the
-    mode-bucket count) because raw rows are monotonic across polls, whereas
-    the mode ngen — and therefore its bucket size — can shift while the tail
-    is still flushing.
+    The log is read exactly once. The caller (DisaggTestCmds.run_cmd) blocks
+    on the gen_server_{i}.done sentinels before calling this, so every gen
+    srun has already exited and its &> aggregate log is fully flushed — there
+    is no partially-written tail to poll for. This replaces the earlier
+    settle-poll heuristic, which could return a mean over a truncated prefix
+    when it accepted the first repeated row count while the log was still
+    flushing across NFS (nvbugs 6487036 / 6487040).
     """
-    deadline = time.time() + settle_timeout
-    prev_count = -1
-    while True:
-        per_file_by_ngen, total_count = _scan_gen_worker_device_step_time(
-            output_dir, num_gen_servers, start_offsets
-        )
-        # Non-empty and unchanged since the last poll → the flush has settled.
-        if total_count > 0 and total_count == prev_count:
-            return _mean_at_mode_ngen(per_file_by_ngen)
-        if time.time() >= deadline:
-            if per_file_by_ngen:
-                print_info(
-                    f"parse_gen_worker_device_step_time: settle_timeout "
-                    f"({settle_timeout}s) reached with {total_count} line(s); "
-                    "returning current mean."
-                )
-                return _mean_at_mode_ngen(per_file_by_ngen)
-            return None
-        prev_count = total_count
-        time.sleep(poll_interval)
+    per_file_scans, _total_count = _scan_gen_worker_device_step_time(
+        output_dir, num_gen_servers, start_offsets
+    )
+    return _mean_at_mode_ngen(per_file_scans)
 
 
 def add_perf_metric_value(
@@ -1351,6 +1348,42 @@ def wait_for_benchmark_ready(
                 )
             time.sleep(10)
 
+    def wait_for_gen_log_sentinels(self, 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
+        ranks' stdout to gen_server_{i}.log via `&>` and touches
+        gen_server_{i}.done only after that srun is reaped (fd closed, log
+        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.
+        """
+        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()
+        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:
+                print_info(
+                    f"Timeout ({self.timeout}s) waiting for gen worker log "
+                    f"sentinels {missing}; parsing current log contents."
+                )
+                return False
+            print_info(
+                f"Waiting for gen worker log sentinels {missing}, elapsed time: {elapsed_time:.0f}s"
+            )
+            time.sleep(poll_interval)
+
     def get_server_logs(self, server_idx: int) -> List[str]:
         server_logs = []
         for i in range(self.num_ctx_servers):
@@ -1463,6 +1496,14 @@ def run_cmd(self, server_idx: int) -> List[str]:
                 disagg_server_proc.wait()
 
         elif self.disagg_serving_type == "BENCHMARK":
+            # Perf-benchmark clients whose gen-worker device step time must be
+            # parsed once the gen logs are flushed. The parse is deferred out of
+            # the client loop because gen_server_*.log keeps being written until
+            # the gen srun exits, and the gen srun only exits after
+            # benchmark_status is written in the finally below. Parsing inside
+            # 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] = []
             try:
                 disagg_server_hostname, disagg_server_port = (
                     self._get_disagg_server_hostname_and_port(server_idx)
@@ -1512,20 +1553,17 @@ def run_cmd(self, server_idx: int) -> List[str]:
                         with open(benchmark_file_path, "w") as benchmark_ctx:
                             benchmark_ctx.write(output)
 
-                        # Only gen_only emits prev_device_step_time; other
-                        # modes yield None and we skip writing the line.
-                        device_step_time_mean = parse_gen_worker_device_step_time(
-                            self.test_output_dir,
-                            self.num_gen_servers,
-                            start_offsets=gen_log_start_offsets,
-                        )
-                        if device_step_time_mean is not None:
-                            summary_line = f"Average Per Iter Device Step Time (ms): {device_step_time_mean:.2f}"
-                            with open(benchmark_file_path, "a") as benchmark_ctx:
-                                benchmark_ctx.write(f"\n{summary_line}\n")
-                            output = f"{output}\n{summary_line}\n"
-
                         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,
+                            }
+                        )
                     else:
                         print_info(
                             f"Skipping perf benchmark for client {client_idx}: "
@@ -1559,6 +1597,30 @@ def run_cmd(self, server_idx: int) -> List[str]:
                 with open(benchmark_status_file, "w") as status_file:
                     status_file.write("Done")
 
+            # 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.
+            if pending_device_step_time:
+                self.wait_for_gen_log_sentinels()
+                for record in pending_device_step_time:
+                    device_step_time_mean = parse_gen_worker_device_step_time(
+                        self.test_output_dir,
+                        self.num_gen_servers,
+                        start_offsets=record["start_offsets"],
+                    )
+                    if device_step_time_mean is not None:
+                        summary_line = (
+                            f"Average Per Iter Device Step Time (ms): {device_step_time_mean:.2f}"
+                        )
+                        with open(record["benchmark_file_path"], "a") as benchmark_ctx:
+                            benchmark_ctx.write(f"\n{summary_line}\n")
+                        idx = record["output_index"]
+                        outputs[idx] = f"{outputs[idx]}\n{summary_line}\n"
+
         return outputs
 
     def get_cmd_str(self, server_idx: int) -> List[str]: