NVIDIA · chienchunhung · Aug 5, 2026 · Jul 31, 2026
@@ -967,13 +967,15 @@ def main():
                     hardware_config.get("gpus_per_ctx_server", 0) or 0,
                     hardware_config.get("gpus_per_gen_server", 0) or 0,
                 ),
+                llm_models_root=args.llm_models_root,
             )
         )
 
         # Add srun args for disagg
         srun_args_lines.extend(
             [
                 "--container-env=DISAGG_SERVING_TYPE",
+                "--container-env=LLM_MODELS_ROOT",
                 "--container-env=pytestCommand",
             ]
         )

@@ -588,6 +588,34 @@ def get_test_output_dir(script_prefix_lines, test_case_name):
     return os.path.join(output_dir, test_case_name) if test_case_name else output_dir
 
 
+def extract_pytest_command_env(script_prefix_lines, name):
+    """Read a leading environment assignment from the exported pytest command."""
+    line = next((ln for ln in script_prefix_lines if "export pytestCommand=" in ln), None)
+    if line is None:
+        raise ValueError("launch prefix does not export pytestCommand")
+    try:
+        outer_tokens = shlex.split(line)
+    except ValueError as e:
+        raise ValueError(f"cannot parse exported pytestCommand: {e}") from e
+    command_assignment = next(
+        (token for token in outer_tokens if token.startswith("pytestCommand=")), None
+    )
+    if command_assignment is None:
+        raise ValueError("launch prefix has a malformed pytestCommand export")
+    command = command_assignment.partition("=")[2]
+    try:
+        command_tokens = shlex.split(command)
+    except ValueError as e:
+        raise ValueError(f"cannot parse pytestCommand payload: {e}") from e
+    for token in command_tokens:
+        if not re.match(r"^[A-Za-z_][A-Za-z0-9_]*=", token):
+            break
+        key, value = token.split("=", 1)
+        if key == name:
+            return value
+    raise ValueError(f"pytestCommand does not set leading environment variable {name}")
+
+
 def remove_whitespace_lines(lines):
     return [line.strip() for line in lines if line.strip()]
 
@@ -791,6 +819,7 @@ def main():
         # Enable/kill-switch policy and timeouts live in precheck_config
         # (single owner, shared with the local flow).
         pcfg = _import_precheck_config(args.llm_src)
+        llm_models_root = extract_pytest_command_env(script_prefix_lines, "LLM_MODELS_ROOT")
         script_prefix_lines.extend(
             pcfg.precheck_prefix_lines(
                 config,
@@ -801,12 +830,14 @@ def main():
                     hardware_config["gpus_per_ctx_server"],
                     hardware_config["gpus_per_gen_server"],
                 ),
+                llm_models_root=llm_models_root,
                 stage_name=args.stage_name,
             )
         )
         srun_args_lines.extend(
             [
                 "--container-env=DISAGG_SERVING_TYPE",
+                "--container-env=LLM_MODELS_ROOT",
                 "--container-env=pytestCommand",
             ]
         )

@@ -79,6 +79,10 @@
 # Number of worker threads for KV transfer queues (default: 1)
 KV_TRANSFER_NUM_THREADS = int(os.environ.get("TRTLLM_KV_TRANSFER_NUM_THREADS", "1"))
 
+# Keep standalone TxSession waits responsive to cancellation even when callers
+# do not configure a sender-future wait slice.
+_FALLBACK_TX_WAIT_SLICE_S = 1.0
+
 
 @dataclass
 class RecvReqInfo:
@@ -1324,11 +1328,13 @@ def has_transferring_tasks(self) -> bool:
     def wait_complete(self, blocking: bool = True) -> Optional[WaitResult]:
         """Poll or block until KV (and optionally aux) transfer finishes.
 
-        With blocking=True (default): waits up to _timeout_s for each task.
+        With blocking=True (default): retries bounded wait slices until a
+        successful transfer finishes. Errors and cancellation remain terminal;
+        callers must not interpret them as proof that peer writes quiesced.
         With blocking=False: polls non-blockingly; returns None if any KV task
         or aux is not yet done.
         """
-        if self.status in (SessionStatus.ERROR, SessionStatus.CANCELLED):
+        if self.has_failed():
             return WaitResult.FAILED
         if not self.kv_tasks:
             return None
@@ -1350,17 +1356,42 @@ def wait_complete(self, blocking: bool = True) -> Optional[WaitResult]:
                     return None
             return WaitResult.COMPLETED
 
+        # ``_timeout_s`` bounds one scheduler wait slice; it is not a transfer
+        # deadline. A successful blockAll must not return merely because one
+        # slice expired: NIXL may still be reading the request's KV pages.
+        wait_slice_s = self._timeout_s
+        if wait_slice_s is None or wait_slice_s <= 0:
+            wait_slice_s = _FALLBACK_TX_WAIT_SLICE_S
+        # The loop preserves standalone block-until-terminal behavior while a
+        # bounded slice keeps cancellation and sibling failure observable.
         for task in self.kv_tasks:
-            if not task.wait(timeout=self._timeout_s):
-                return WaitResult.TIMEOUT
+            while not task.wait(timeout=wait_slice_s):
+                # cancel() leaves a TRANSFERRING task's event unset until the
+                # physical writer finishes. Preserve the bounded-slice control
+                # point so blockAll can still report terminal cancellation.
+                # Check every task: a sibling slice can fail while this one is
+                # still pending without setting the session terminal status.
+                if self.has_failed():
+                    return WaitResult.FAILED
             if task.status == TaskStatus.ERROR:
                 return WaitResult.FAILED
-        if self._need_aux and self.aux_task is not None:
-            if not self.aux_task.wait(timeout=self._timeout_s):
-                return WaitResult.TIMEOUT
+        if self._need_aux:
+            if self.aux_task is None:
+                return (
+                    WaitResult.FAILED
+                    if self.status in (SessionStatus.ERROR, SessionStatus.CANCELLED)
+                    else None
+                )
+            while not self.aux_task.wait(timeout=wait_slice_s):
+                if self.has_failed():
+                    return WaitResult.FAILED
             if self.aux_task.status == TaskStatus.ERROR:
                 return WaitResult.FAILED
-        return WaitResult.COMPLETED
+        return (
+            WaitResult.FAILED
+            if self.status in (SessionStatus.ERROR, SessionStatus.CANCELLED)
+            else WaitResult.COMPLETED
+        )
 
     def set_exception(self, reason: str = ""):
         msg = f"TxSession {self.disagg_request_id} exception"

@@ -703,8 +703,9 @@ def check_context_transfer_status(
             elif result is None:
                 continue
             elif result == WaitResult.TIMEOUT:
-                logger.warning(
-                    f"TxSession rid={session.disagg_request_id} timed out after {self._sender_future_timeout_ms}ms"
+                logger.debug(
+                    f"TxSession rid={session.disagg_request_id} not ready after "
+                    f"{self._sender_future_timeout_ms}ms wait slice; keeping it in progress"
                 )
                 timed_out.append(rid)
             else:

@@ -16,7 +16,7 @@ starts.
 | Same UCX env vars (incl. the `unset UCX_TLS` cases) | `jenkins/scripts/perf/submit.py` builds the precheck commands from the **same** `ucx_tls_cmd` + `$CTX/GEN_WORKER_ENV_VARS` strings as the worker steps; `slurm_precheck_run.sh` sources the same `slurm_env_setup.sh` (the `UCX_TLS=tcp` fixup) as `slurm_run.sh`. |
 | Same instance count / parallelism | One precheck `srun` per ctx/gen server with the same `-N/--ntasks/--ntasks-per-node/--mpi=pmix` and the same node slices (`-w`) as the real server steps (`slurm_launch_draft.sh`). TP/PP/CP/attention-DP come from the same `worker_config`. |
 | Same transceiver config | `CacheTransceiverConfig(**yaml["worker_config"][role]["cache_transceiver_config"])` — the yaml block is passed through verbatim (backend, `max_tokens_in_buffer`, timeouts, ...). |
-| Same KV cache manager version + transceiver runtime | Explicit per-side `kv_cache_config.use_kv_cache_manager_v2` wins; absent means "auto" and resolves against the model class's `get_model_defaults()`, and `transceiver_runtime: auto` resolves via `get_preferred_transceiver_runtime()` (NIXL-gated) — both through the same llm_utils code serving uses. V2 requires the Python transceiver (the C++ one only supports V1); a V2+CPP combination fails fast with INIT_ERROR. |
+| Same KV cache manager version + transceiver runtime | The launch generator forwards the real test's `LLM_MODELS_ROOT` into every precheck process. Explicit per-side `kv_cache_config.use_kv_cache_manager_v2` wins; absent means "auto" and requires a registered model class, then resolves against its optional `get_model_defaults()` hook (no hook means empty defaults, matching serving). `transceiver_runtime: auto` resolves via `get_preferred_transceiver_runtime()` (NIXL-gated) — both through the same llm_utils code serving uses. An unresolved manager-version `auto` setting fails with INIT_ERROR instead of silently assuming V1. V2 requires the Python transceiver (the C++ one only supports V1); a V2+CPP combination also fails fast. |
 
 Asymmetric layouts (ctx dep4 → gen dep16, ctx pp8 → gen tp32, ...) are
 supported: data is seeded per (request, **global** layer) and constant along
@@ -77,13 +77,13 @@ csv/ctx_<i>/<uuid>_<rank>.csv       # Python transceiver per-task perf
 
 ```bash
 # Inspect what a yaml resolves to (no GPU needed):
-python3 run_precheck.py --role gen --server-idx 0 --dry-run \
+LLM_MODELS_ROOT='<models>' python3 run_precheck.py --role gen --server-idx 0 --dry-run \
     --config ../disaggregated/<test>.yaml --work-dir /tmp/ct --llm-src <repo>
 
 # On a SLURM allocation: one srun per instance, e.g. ctx dep4 + gen dep8:
-srun -N1 --ntasks=4 --mpi=pmix python3 run_precheck.py --role ctx --server-idx 0 \
-    --config <yaml> --work-dir <shared-dir> --llm-src <repo> &
-srun -N2 --ntasks=8 --mpi=pmix python3 run_precheck.py --role gen --server-idx 0 \
+LLM_MODELS_ROOT='<models>' srun -N1 --ntasks=4 --mpi=pmix python3 run_precheck.py \
+    --role ctx --server-idx 0 --config <yaml> --work-dir <shared-dir> --llm-src <repo> &
+LLM_MODELS_ROOT='<models>' srun -N2 --ntasks=8 --mpi=pmix python3 run_precheck.py --role gen --server-idx 0 \
     --config <yaml> --work-dir <shared-dir> --llm-src <repo> &
 wait
 ```

@@ -26,6 +26,7 @@
 
 import json
 import os
+import shlex
 
 # Optional per-yaml overrides live under a `cache_transceiver_precheck:` block.
 PRECHECK_DEFAULTS = {
@@ -61,8 +62,9 @@
     "verify_data": True,
 }
 
-# Fallback KV shape when the model directory cannot be resolved: the precheck
-# still exercises the exact network path, just with a synthetic cache shape.
+# Fallback KV shape for dry-runs and explicitly selected manager versions when
+# the model directory cannot be resolved. Manager-version "auto" resolution
+# fails fast instead of silently pairing this shape with V1.
 FALLBACK_KV_SHAPE = {
     "num_layers": 32,
     "num_kv_heads": 8,
@@ -128,7 +130,13 @@ def default_step_timeout_s(max_world):
 
 
 def precheck_prefix_lines(
-    cfg, benchmark_mode, config_path_expr, ucx_tls_cmd, max_world, stage_name=""
+    cfg,
+    benchmark_mode,
+    config_path_expr,
+    ucx_tls_cmd,
+    max_world,
+    llm_models_root,
+    stage_name="",
 ):
     """Launch-script export lines wiring the precheck gate.
 
@@ -164,6 +172,9 @@ def precheck_prefix_lines(
         f"--benchmark-mode {benchmark_mode} --llm-src $llmSrcNode"
     )
     lines = [
+        # Keep this as a top-level assignment. shlex.quote() is not safe when
+        # nested inside the double-quoted pytestCommand exports below.
+        f"export LLM_MODELS_ROOT={shlex.quote(llm_models_root)}",
         f"export ctPrecheckEnabled={int(enabled)}",
         # The external srun timeout must cover the driver's first-rep NIXL
         # wire-up allowance; the default derives from the same formula the