From d661a269f674a05ebe4cd9c1df8900331cabb8dd Mon Sep 17 00:00:00 2001
From: "svc-glamr@nvidia.com" <svc-glamr@nvidia.com>
Date: Thu, 3 Sep 2026 18:05:14 +0000
Subject: [PATCH 1/6] test(fault_tolerance): assert drain, not migration, on
 graceful shutdown

The graceful_shutdown arm of the request-migration tests asserted that a
SIGTERM'd worker migrates its in-flight request. That contradicts Dynamo's
documented graceful-shutdown contract: a worker that receives SIGTERM stops
accepting new work and drains the requests it has already admitted, so no
migration should occur at all.

The tests only saw a migration because run_migration_test's default
non-immediate-kill path calls terminate_process_tree with timeout=2, which
escalates to SIGKILL two seconds after SIGTERM and cuts the drain short.
That is a delayed worker failure, not a graceful shutdown.

- utils.py: promote the SIGTERM-only shutdown context (previously private to
  test_sglang.py) to a public, backend-agnostic graceful_worker_shutdown
  helper, parameterized by the (component, endpoint) pair the faulted worker
  registers so prefill workers can be drained too.
- utils.py: add an opt-in expect_drain flag to run_migration_test. It requires
  the request to succeed regardless of migration settings and pins every
  migration counter to exactly zero (exact_counts=True), because a lower-bound
  check against zero asserts nothing.
- test_vllm.py, test_trtllm.py: drive the graceful_shutdown arm through the
  shared context with expect_drain, replacing the unconditional
  expected_ongoing_request_count=1 in the vLLM tests. The worker_failure arm
  is unchanged.

test_sglang.py is deliberately untouched; its collected parametrization ids
are byte-identical before and after this change.

Signed-off-by: svc-glamr@nvidia.com <svc-glamr@nvidia.com>
---
 .../fault_tolerance/migration/test_trtllm.py  |  18 ++-
 tests/fault_tolerance/migration/test_vllm.py  |  19 ++-
 tests/fault_tolerance/migration/utils.py      | 111 +++++++++++++++++-
 3 files changed, 139 insertions(+), 9 deletions(-)

diff --git a/tests/fault_tolerance/migration/test_trtllm.py b/tests/fault_tolerance/migration/test_trtllm.py
index c183baac6d78..fcc6da69aa29 100644
--- a/tests/fault_tolerance/migration/test_trtllm.py
+++ b/tests/fault_tolerance/migration/test_trtllm.py
@@ -21,7 +21,7 @@
 from tests.utils.port_utils import allocate_port, deallocate_port
 
 # Customized utils for migration tests
-from .utils import DynamoFrontendProcess, run_migration_test
+from .utils import DynamoFrontendProcess, graceful_worker_shutdown, run_migration_test
 
 logger = logging.getLogger(__name__)
 
@@ -250,6 +250,10 @@ def test_request_migration_trtllm_aggregated(
                     immediate_kill=immediate_kill,
                     use_chat_completion=(request_api == "chat"),
                     stream=stream,
+                    expect_drain=not immediate_kill,
+                    graceful_shutdown=lambda worker: graceful_worker_shutdown(
+                        frontend, worker
+                    ),
                 )
 
 
@@ -327,6 +331,10 @@ def test_request_migration_trtllm_prefill(
                         use_chat_completion=(request_api == "chat"),
                         stream=stream,
                         use_long_prompt=True,
+                        expect_drain=not immediate_kill,
+                        graceful_shutdown=lambda worker: graceful_worker_shutdown(
+                            frontend, worker, endpoint=("prefill", "generate")
+                        ),
                     )
 
 
@@ -402,6 +410,10 @@ def test_request_migration_trtllm_kv_transfer(
                         use_chat_completion=(request_api == "chat"),
                         stream=stream,
                         use_long_prompt=True,
+                        expect_drain=not immediate_kill,
+                        graceful_shutdown=lambda worker: graceful_worker_shutdown(
+                            frontend, worker
+                        ),
                     )
 
 
@@ -480,4 +492,8 @@ def test_request_migration_trtllm_decode(
                         use_chat_completion=(request_api == "chat"),
                         stream=stream,
                         wait_for_new_response_before_stop=True,
+                        expect_drain=not immediate_kill,
+                        graceful_shutdown=lambda worker: graceful_worker_shutdown(
+                            frontend, worker
+                        ),
                     )
diff --git a/tests/fault_tolerance/migration/test_vllm.py b/tests/fault_tolerance/migration/test_vllm.py
index 9e5c3b1bec0d..9fd380d269da 100644
--- a/tests/fault_tolerance/migration/test_vllm.py
+++ b/tests/fault_tolerance/migration/test_vllm.py
@@ -24,6 +24,7 @@
 # Customized utils for migration tests
 from .utils import (
     DynamoFrontendProcess,
+    graceful_worker_shutdown,
     managed_processes_concurrently,
     run_migration_test,
     wait_for_endpoint_instances,
@@ -459,7 +460,11 @@ def test_request_migration_vllm_aggregated(
                 use_chat_completion=(request_api == "chat"),
                 stream=stream,
                 max_tokens=AGGREGATED_MAX_TOKENS,
-                expected_ongoing_request_count=1,
+                expected_ongoing_request_count=1 if immediate_kill else None,
+                expect_drain=not immediate_kill,
+                graceful_shutdown=lambda worker: graceful_worker_shutdown(
+                    frontend, worker
+                ),
             )
 
 
@@ -548,7 +553,11 @@ def test_request_migration_vllm_kv_transfer(
                 max_tokens=KV_TRANSFER_MAX_TOKENS,
                 use_long_prompt=True,
                 long_prompt_repetitions=KV_TRANSFER_PROMPT_REPETITIONS,
-                expected_ongoing_request_count=1,
+                expected_ongoing_request_count=1 if immediate_kill else None,
+                expect_drain=not immediate_kill,
+                graceful_shutdown=lambda worker: graceful_worker_shutdown(
+                    frontend, worker
+                ),
             )
 
 
@@ -633,5 +642,9 @@ def test_request_migration_vllm_decode(
                 stream=True,
                 max_tokens=DECODE_MAX_TOKENS,
                 wait_for_new_response_before_stop=True,
-                expected_ongoing_request_count=1,
+                expected_ongoing_request_count=1 if immediate_kill else None,
+                expect_drain=not immediate_kill,
+                graceful_shutdown=lambda worker: graceful_worker_shutdown(
+                    frontend, worker
+                ),
             )
diff --git a/tests/fault_tolerance/migration/utils.py b/tests/fault_tolerance/migration/utils.py
index 84db34852f62..f5313ed175ca 100644
--- a/tests/fault_tolerance/migration/utils.py
+++ b/tests/fault_tolerance/migration/utils.py
@@ -2,13 +2,16 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import logging
+import os
 import re
+import signal
 import threading
 import time
 from collections.abc import Callable, Iterator
 from concurrent.futures import ThreadPoolExecutor
 from contextlib import AbstractContextManager, ExitStack, contextmanager, nullcontext
 
+import psutil
 import pytest
 import requests
 from openai import APIError, OpenAI
@@ -463,6 +466,66 @@ def wait_for_endpoint_instance_reduction(
     )
 
 
+@contextmanager
+def graceful_worker_shutdown(
+    frontend: DynamoFrontendProcess,
+    worker: ManagedProcess,
+    endpoint: tuple[str, str] = ("backend", "generate"),
+) -> Iterator[None]:
+    """Send SIGTERM only, and keep the worker alive until the outcome is known.
+
+    `terminate_process_tree` escalates to SIGKILL a fixed number of seconds
+    after SIGTERM, so a worker that is still draining a generation is killed
+    mid-stream and the request migrates instead. This context sends SIGTERM,
+    waits for the worker to leave frontend discovery, yields while the request
+    outcome is observed, and only then force-kills the worker's process groups.
+    Teardown still cannot leak engine processes that pin the GPU.
+
+    Args:
+        frontend: Frontend whose `/health` view reports endpoint instances
+        worker: Worker to shut down
+        endpoint: The (component, endpoint) pair this worker registers.
+            Prefill workers register ("prefill", "generate").
+    """
+    response = requests.get(
+        f"http://localhost:{frontend.frontend_port}/health",
+        timeout=1,
+    )
+    response.raise_for_status()
+    previous_count = sum(
+        1
+        for instance in response.json().get("instances", [])
+        if (instance.get("component"), instance.get("endpoint")) == endpoint
+    )
+
+    pid = worker.get_pid()
+    parent = psutil.Process(pid)
+    process_groups = {os.getpgid(pid)}
+    try:
+        for child in parent.children(recursive=True):
+            try:
+                process_groups.add(os.getpgid(child.pid))
+            except (ProcessLookupError, OSError):
+                pass
+    except (psutil.AccessDenied, psutil.NoSuchProcess):
+        pass
+
+    try:
+        parent.terminate()
+        wait_for_endpoint_instance_reduction(
+            frontend.frontend_port,
+            endpoint,
+            previous_count,
+        )
+        yield
+    finally:
+        for process_group in process_groups:
+            try:
+                os.killpg(process_group, signal.SIGKILL)
+            except ProcessLookupError:
+                pass
+
+
 def wait_for_response(
     response_list: list[tuple[str | None | Exception, float]],
     num_responses: int = 5,
@@ -723,6 +786,7 @@ def run_migration_test(
     long_prompt_repetitions: int = 8_000,
     wait_for_new_response_before_stop: bool = False,
     expected_ongoing_request_count: int | None = None,
+    expect_drain: bool = False,
     graceful_shutdown: Callable[[ManagedProcess], AbstractContextManager[None]]
     | None = None,
     verify_replacement_worker: bool = False,
@@ -749,6 +813,12 @@ def run_migration_test(
         expected_ongoing_request_count: Exact expected count for callers that
             opt into strict metric validation. When omitted, preserve the
             shared helper's historical backend-agnostic lower-bound behavior.
+        expect_drain: Assert the graceful-shutdown contract instead of the
+            migration contract: a SIGTERM'd worker finishes the requests it has
+            already admitted, so the request must succeed and every migration
+            counter must be exactly zero. Requires immediate_kill=False and a
+            graceful_shutdown context, and is mutually exclusive with
+            expected_ongoing_request_count.
         graceful_shutdown: Optional backend-specific context that initiates
             graceful shutdown before response validation and performs final
             cleanup after the request outcome is known.
@@ -762,6 +832,23 @@ def run_migration_test(
             max_tokens budget so state-based fault synchronization cannot race
             an early EOS.
     """
+    if expect_drain:
+        if immediate_kill:
+            raise ValueError(
+                "expect_drain asserts the graceful-shutdown drain contract and "
+                "is incompatible with immediate_kill=True"
+            )
+        if graceful_shutdown is None:
+            raise ValueError(
+                "expect_drain requires a graceful_shutdown context; the default "
+                "shutdown path escalates to SIGKILL and cuts the drain short"
+            )
+        if expected_ongoing_request_count is not None:
+            raise ValueError(
+                "expect_drain already pins every migration counter to zero; "
+                "pass expected_ongoing_request_count only without it"
+            )
+
     # Step 1: Send the request
     if use_chat_completion:
         request_thread, response_list = start_chat_completion_request(
@@ -821,11 +908,14 @@ def run_migration_test(
             shutdown_context = graceful_shutdown(worker)
 
     # Step 5: Validate the request outcome via its response (the user-facing
-    # contract). Migration is expected to succeed only when it is enabled and the
-    # request does not exceed the migration seq-len cap; otherwise the in-flight
-    # request must fail.
+    # contract). A drained worker must complete the request it already admitted,
+    # whatever the migration settings are. Otherwise migration is expected to
+    # succeed only when it is enabled and the request does not exceed the
+    # migration seq-len cap; otherwise the in-flight request must fail.
     with shutdown_context:
-        if migration_limit > 0 and migration_max_seq_len != 1:
+        if expect_drain:
+            validate_response(request_thread, response_list)
+        elif migration_limit > 0 and migration_max_seq_len != 1:
             if verify_replacement_worker:
                 wait_for_worker_request_id(
                     replacement_worker,
@@ -851,7 +941,18 @@ def run_migration_test(
     # stream, including an attempt that cannot retry because migration_limit is
     # zero. It is the structured equivalent of the old "Stream disconnected,
     # recreating stream" log assertion. `max_seq_len_exceeded` records hitting
-    # the migration seq-len cap.
+    # the migration seq-len cap. A drain migrates nothing at all, so every
+    # counter must be exactly zero -- lower-bound checks would pass vacuously.
+    if expect_drain:
+        verify_migration_metrics(
+            frontend.frontend_port,
+            expected_ongoing_request_count=0,
+            expected_new_request_count=0,
+            expected_max_seq_len_exceeded_count=0,
+            exact_counts=True,
+        )
+        return
+
     exact_metric_counts = expected_ongoing_request_count is not None
     if expected_ongoing_request_count is None:
         expected_ongoing_request_count = 1 if migration_limit > 0 else 0

From f8645345c87450cbef6d673333deb32f3f21fc8a Mon Sep 17 00:00:00 2001
From: "svc-glamr@nvidia.com" <svc-glamr@nvidia.com>
Date: Thu, 3 Sep 2026 18:45:33 +0000
Subject: [PATCH 2/6] test(fault_tolerance): keep the seq-cap counter out of
 the drain contract

The expect_drain branch pinned expected_max_seq_len_exceeded_count to 0,
but that counter is driven by ordinary token accounting in
RetryManager::exceed_max_seq_len, not by a worker fault. It fires once at
request build time whenever the prompt already exceeds the configured
migration seq-len cap, so rows with migration_max_seq_len=1 record one
event even when the worker drains perfectly.

Share the existing expression with the non-drain path so both arms expect
the same seq-cap value, and document why that counter is not part of the
drain contract.

Signed-off-by: svc-glamr@nvidia.com <svc-glamr@nvidia.com>
---
 tests/fault_tolerance/migration/utils.py | 32 +++++++++++++++++-------
 1 file changed, 23 insertions(+), 9 deletions(-)

diff --git a/tests/fault_tolerance/migration/utils.py b/tests/fault_tolerance/migration/utils.py
index f5313ed175ca..210824642836 100644
--- a/tests/fault_tolerance/migration/utils.py
+++ b/tests/fault_tolerance/migration/utils.py
@@ -815,9 +815,11 @@ def run_migration_test(
             shared helper's historical backend-agnostic lower-bound behavior.
         expect_drain: Assert the graceful-shutdown contract instead of the
             migration contract: a SIGTERM'd worker finishes the requests it has
-            already admitted, so the request must succeed and every migration
-            counter must be exactly zero. Requires immediate_kill=False and a
-            graceful_shutdown context, and is mutually exclusive with
+            already admitted, so the request must succeed and both migration
+            counters must be exactly zero. The max_seq_len_exceeded counter is
+            not a migration outcome and keeps its usual expectation -- see
+            Step 6. Requires immediate_kill=False and a graceful_shutdown
+            context, and is mutually exclusive with
             expected_ongoing_request_count.
         graceful_shutdown: Optional backend-specific context that initiates
             graceful shutdown before response validation and performs final
@@ -845,7 +847,7 @@ def run_migration_test(
             )
         if expected_ongoing_request_count is not None:
             raise ValueError(
-                "expect_drain already pins every migration counter to zero; "
+                "expect_drain already pins both migration counters to zero; "
                 "pass expected_ongoing_request_count only without it"
             )
 
@@ -940,15 +942,27 @@ def run_migration_test(
     # log strings. `ongoing_request` counts an error from an established
     # stream, including an attempt that cannot retry because migration_limit is
     # zero. It is the structured equivalent of the old "Stream disconnected,
-    # recreating stream" log assertion. `max_seq_len_exceeded` records hitting
-    # the migration seq-len cap. A drain migrates nothing at all, so every
-    # counter must be exactly zero -- lower-bound checks would pass vacuously.
+    # recreating stream" log assertion.
+    #
+    # `max_seq_len_exceeded` is deliberately NOT part of the drain contract. It
+    # is not a migration outcome: RetryManager increments it from ordinary token
+    # accounting in `exceed_max_seq_len` (lib/llm/src/migration.rs), which runs
+    # once at request build time against the prompt and then per response chunk
+    # from `track_response`. No fault, stream break, or migration is involved,
+    # and the counter fires at most once because it zeroes `retries_left`. So a
+    # request that drains perfectly with migration_max_seq_len=1 still records
+    # exactly one seq-cap event, and both arms share the same expectation.
+    expected_max_seq_len_exceeded_count = 1 if migration_max_seq_len == 1 else 0
+
     if expect_drain:
+        # A drain migrates nothing, so both migration counters must be exactly
+        # zero -- under the default lower-bound mode a zero expectation asserts
+        # nothing at all, hence exact_counts=True.
         verify_migration_metrics(
             frontend.frontend_port,
             expected_ongoing_request_count=0,
             expected_new_request_count=0,
-            expected_max_seq_len_exceeded_count=0,
+            expected_max_seq_len_exceeded_count=expected_max_seq_len_exceeded_count,
             exact_counts=True,
         )
         return
@@ -960,6 +974,6 @@ def run_migration_test(
     verify_migration_metrics(
         frontend.frontend_port,
         expected_ongoing_request_count=expected_ongoing_request_count,
-        expected_max_seq_len_exceeded_count=1 if migration_max_seq_len == 1 else 0,
+        expected_max_seq_len_exceeded_count=expected_max_seq_len_exceeded_count,
         exact_counts=exact_metric_counts,
     )

From 6efa8a9d7414082bbfb9c2a90c7c585d96f85245 Mon Sep 17 00:00:00 2001
From: "svc-glamr@nvidia.com" <svc-glamr@nvidia.com>
Date: Thu, 3 Sep 2026 19:17:15 +0000
Subject: [PATCH 3/6] chore: remove factory working notes

Signed-off-by: svc-glamr@nvidia.com <svc-glamr@nvidia.com>
---
 tests/fault_tolerance/migration/utils.py | 23 ++++++-----------------
 1 file changed, 6 insertions(+), 17 deletions(-)

diff --git a/tests/fault_tolerance/migration/utils.py b/tests/fault_tolerance/migration/utils.py
index 210824642836..1fa7f0e84754 100644
--- a/tests/fault_tolerance/migration/utils.py
+++ b/tests/fault_tolerance/migration/utils.py
@@ -910,10 +910,8 @@ def run_migration_test(
             shutdown_context = graceful_shutdown(worker)
 
     # Step 5: Validate the request outcome via its response (the user-facing
-    # contract). A drained worker must complete the request it already admitted,
-    # whatever the migration settings are. Otherwise migration is expected to
-    # succeed only when it is enabled and the request does not exceed the
-    # migration seq-len cap; otherwise the in-flight request must fail.
+    # contract). A drained worker completes its admitted request; otherwise
+    # migration succeeds only when enabled and under the seq-len cap, else fails.
     with shutdown_context:
         if expect_drain:
             validate_response(request_thread, response_list)
@@ -942,22 +940,13 @@ def run_migration_test(
     # log strings. `ongoing_request` counts an error from an established
     # stream, including an attempt that cannot retry because migration_limit is
     # zero. It is the structured equivalent of the old "Stream disconnected,
-    # recreating stream" log assertion.
-    #
-    # `max_seq_len_exceeded` is deliberately NOT part of the drain contract. It
-    # is not a migration outcome: RetryManager increments it from ordinary token
-    # accounting in `exceed_max_seq_len` (lib/llm/src/migration.rs), which runs
-    # once at request build time against the prompt and then per response chunk
-    # from `track_response`. No fault, stream break, or migration is involved,
-    # and the counter fires at most once because it zeroes `retries_left`. So a
-    # request that drains perfectly with migration_max_seq_len=1 still records
-    # exactly one seq-cap event, and both arms share the same expectation.
+    # recreating stream" log assertion. `max_seq_len_exceeded` records the
+    # seq-len cap from token accounting, not migration, so both arms share it.
     expected_max_seq_len_exceeded_count = 1 if migration_max_seq_len == 1 else 0
 
     if expect_drain:
-        # A drain migrates nothing, so both migration counters must be exactly
-        # zero -- under the default lower-bound mode a zero expectation asserts
-        # nothing at all, hence exact_counts=True.
+        # A drain migrates nothing. Under the default lower-bound mode a zero
+        # expectation asserts nothing, hence exact_counts=True.
         verify_migration_metrics(
             frontend.frontend_port,
             expected_ongoing_request_count=0,

From 5507c6eb827ccdd07cf6c842a1e436a6f9a0a241 Mon Sep 17 00:00:00 2001
From: "svc-glamr@nvidia.com" <svc-glamr@nvidia.com>
Date: Thu, 3 Sep 2026 19:39:48 +0000
Subject: [PATCH 4/6] test(fault_tolerance): scope drain assertions to runnable
 TRT-LLM tests

Keep the graceful-drain contract on the TRT-LLM aggregated and decode
tests, which run in CI, and leave the unconditionally skipped prefill and
KV-transfer tests unchanged so no unexercised assertion ships.

With no prefill caller left, graceful_worker_shutdown resolves the worker
endpoint from a module constant instead of an unused parameter, and it
suppresses only the disappearing-process errors so a permission failure
cannot silently drop a live child process group.

Signed-off-by: svc-glamr@nvidia.com <svc-glamr@nvidia.com>
---
 .../fault_tolerance/migration/test_trtllm.py  |  8 -------
 tests/fault_tolerance/migration/utils.py      | 24 ++++++++-----------
 2 files changed, 10 insertions(+), 22 deletions(-)

diff --git a/tests/fault_tolerance/migration/test_trtllm.py b/tests/fault_tolerance/migration/test_trtllm.py
index fcc6da69aa29..4bba7b660333 100644
--- a/tests/fault_tolerance/migration/test_trtllm.py
+++ b/tests/fault_tolerance/migration/test_trtllm.py
@@ -331,10 +331,6 @@ def test_request_migration_trtllm_prefill(
                         use_chat_completion=(request_api == "chat"),
                         stream=stream,
                         use_long_prompt=True,
-                        expect_drain=not immediate_kill,
-                        graceful_shutdown=lambda worker: graceful_worker_shutdown(
-                            frontend, worker, endpoint=("prefill", "generate")
-                        ),
                     )
 
 
@@ -410,10 +406,6 @@ def test_request_migration_trtllm_kv_transfer(
                         use_chat_completion=(request_api == "chat"),
                         stream=stream,
                         use_long_prompt=True,
-                        expect_drain=not immediate_kill,
-                        graceful_shutdown=lambda worker: graceful_worker_shutdown(
-                            frontend, worker
-                        ),
                     )
 
 
diff --git a/tests/fault_tolerance/migration/utils.py b/tests/fault_tolerance/migration/utils.py
index 1fa7f0e84754..8e55aabc3d3b 100644
--- a/tests/fault_tolerance/migration/utils.py
+++ b/tests/fault_tolerance/migration/utils.py
@@ -25,6 +25,10 @@
 
 logger = logging.getLogger(__name__)
 
+# The (component, endpoint) pair a decode or aggregated worker registers with
+# the frontend.
+BACKEND_ENDPOINT = ("backend", "generate")
+
 
 @contextmanager
 def managed_processes_concurrently(
@@ -470,7 +474,6 @@ def wait_for_endpoint_instance_reduction(
 def graceful_worker_shutdown(
     frontend: DynamoFrontendProcess,
     worker: ManagedProcess,
-    endpoint: tuple[str, str] = ("backend", "generate"),
 ) -> Iterator[None]:
     """Send SIGTERM only, and keep the worker alive until the outcome is known.
 
@@ -484,8 +487,6 @@ def graceful_worker_shutdown(
     Args:
         frontend: Frontend whose `/health` view reports endpoint instances
         worker: Worker to shut down
-        endpoint: The (component, endpoint) pair this worker registers.
-            Prefill workers register ("prefill", "generate").
     """
     response = requests.get(
         f"http://localhost:{frontend.frontend_port}/health",
@@ -495,7 +496,7 @@ def graceful_worker_shutdown(
     previous_count = sum(
         1
         for instance in response.json().get("instances", [])
-        if (instance.get("component"), instance.get("endpoint")) == endpoint
+        if (instance.get("component"), instance.get("endpoint")) == BACKEND_ENDPOINT
     )
 
     pid = worker.get_pid()
@@ -505,16 +506,16 @@ def graceful_worker_shutdown(
         for child in parent.children(recursive=True):
             try:
                 process_groups.add(os.getpgid(child.pid))
-            except (ProcessLookupError, OSError):
+            except ProcessLookupError:
                 pass
-    except (psutil.AccessDenied, psutil.NoSuchProcess):
+    except psutil.NoSuchProcess:
         pass
 
     try:
         parent.terminate()
         wait_for_endpoint_instance_reduction(
             frontend.frontend_port,
-            endpoint,
+            BACKEND_ENDPOINT,
             previous_count,
         )
         yield
@@ -814,13 +815,8 @@ def run_migration_test(
             opt into strict metric validation. When omitted, preserve the
             shared helper's historical backend-agnostic lower-bound behavior.
         expect_drain: Assert the graceful-shutdown contract instead of the
-            migration contract: a SIGTERM'd worker finishes the requests it has
-            already admitted, so the request must succeed and both migration
-            counters must be exactly zero. The max_seq_len_exceeded counter is
-            not a migration outcome and keeps its usual expectation -- see
-            Step 6. Requires immediate_kill=False and a graceful_shutdown
-            context, and is mutually exclusive with
-            expected_ongoing_request_count.
+            migration contract: a SIGTERM'd worker finishes the request it has
+            already admitted, so the request succeeds and nothing migrates.
         graceful_shutdown: Optional backend-specific context that initiates
             graceful shutdown before response validation and performs final
             cleanup after the request outcome is known.

From 5f7c4f36f4af3b1dcb2aefee46c59db6c029f57d Mon Sep 17 00:00:00 2001
From: "svc-glamr@nvidia.com" <svc-glamr@nvidia.com>
Date: Thu, 10 Sep 2026 18:21:10 +0000
Subject: [PATCH 5/6] test(fault_tolerance): share one graceful-shutdown
 context across backends

_sglang_graceful_shutdown duplicated graceful_worker_shutdown line for
line: same /health baseline, same process-group capture, same SIGTERM
then wait-for-discovery-removal, same SIGKILL by process group in a
finally. Import the shared helper from .utils and drop the local copy so
the three backend suites share one implementation.

The shared helper lets an inspection failure other than a missing process
propagate instead of being swallowed, which the local copy suppressed via
psutil.AccessDenied and OSError.

Signed-off-by: svc-glamr@nvidia.com <svc-glamr@nvidia.com>
---
 .../fault_tolerance/migration/test_sglang.py  | 58 ++-----------------
 1 file changed, 4 insertions(+), 54 deletions(-)

diff --git a/tests/fault_tolerance/migration/test_sglang.py b/tests/fault_tolerance/migration/test_sglang.py
index aec2d0cd8a1e..01b4250e357f 100644
--- a/tests/fault_tolerance/migration/test_sglang.py
+++ b/tests/fault_tolerance/migration/test_sglang.py
@@ -3,14 +3,10 @@
 
 import logging
 import os
-import signal
 import threading
 import time
-from collections.abc import Iterator
-from contextlib import contextmanager
 from pathlib import Path
 
-import psutil
 import pytest
 import requests
 
@@ -24,9 +20,9 @@
 # Customized utils for migration tests
 from .utils import (
     DynamoFrontendProcess,
+    graceful_worker_shutdown,
     managed_processes_concurrently,
     run_migration_test,
-    wait_for_endpoint_instance_reduction,
     wait_for_endpoint_instances,
 )
 
@@ -45,52 +41,6 @@
 SGLANG_MIGRATION_TEST_TIMEOUT_S = 780
 
 
-@contextmanager
-def _sglang_graceful_shutdown(
-    frontend: DynamoFrontendProcess,
-    worker: ManagedProcess,
-) -> Iterator[None]:
-    """Keep the failed worker alive until the request outcome is observable."""
-    endpoint = ("backend", "generate")
-    response = requests.get(
-        f"http://localhost:{frontend.frontend_port}/health",
-        timeout=1,
-    )
-    response.raise_for_status()
-    previous_count = sum(
-        1
-        for instance in response.json().get("instances", [])
-        if (instance.get("component"), instance.get("endpoint")) == endpoint
-    )
-
-    pid = worker.get_pid()
-    parent = psutil.Process(pid)
-    process_groups = {os.getpgid(pid)}
-    try:
-        for child in parent.children(recursive=True):
-            try:
-                process_groups.add(os.getpgid(child.pid))
-            except (ProcessLookupError, OSError):
-                pass
-    except (psutil.AccessDenied, psutil.NoSuchProcess):
-        pass
-
-    try:
-        parent.terminate()
-        wait_for_endpoint_instance_reduction(
-            frontend.frontend_port,
-            endpoint,
-            previous_count,
-        )
-        yield
-    finally:
-        for process_group in process_groups:
-            try:
-                os.killpg(process_group, signal.SIGKILL)
-            except ProcessLookupError:
-                pass
-
-
 def get_sglang_kv_transfer_metrics(worker_system_port: int) -> tuple[float, float]:
     """Return the completed SGLang KV-transfer count and total size."""
     response = requests.get(f"http://localhost:{worker_system_port}/metrics", timeout=1)
@@ -557,7 +507,7 @@ def test_request_migration_sglang_aggregated(
                 stream=stream,
                 max_tokens=AGGREGATED_MAX_TOKENS,
                 expected_ongoing_request_count=1,
-                graceful_shutdown=lambda worker: _sglang_graceful_shutdown(
+                graceful_shutdown=lambda worker: graceful_worker_shutdown(
                     frontend, worker
                 ),
                 verify_replacement_worker=True,
@@ -670,7 +620,7 @@ def wait_for_initial_transfer() -> None:
                 use_long_prompt=True,
                 long_prompt_repetitions=KV_TRANSFER_PROMPT_REPETITIONS,
                 expected_ongoing_request_count=1,
-                graceful_shutdown=lambda worker: _sglang_graceful_shutdown(
+                graceful_shutdown=lambda worker: graceful_worker_shutdown(
                     frontend, worker
                 ),
                 verify_replacement_worker=True,
@@ -770,7 +720,7 @@ def test_request_migration_sglang_decode(
                 max_tokens=DECODE_MAX_TOKENS,
                 wait_for_new_response_before_stop=True,
                 expected_ongoing_request_count=1,
-                graceful_shutdown=lambda worker: _sglang_graceful_shutdown(
+                graceful_shutdown=lambda worker: graceful_worker_shutdown(
                     frontend, worker
                 ),
                 verify_replacement_worker=True,

From 6f7cb164fe9729e87d0d02a77fe0b9dfcb6e4ade Mon Sep 17 00:00:00 2001
From: "svc-glamr@nvidia.com" <svc-glamr@nvidia.com>
Date: Fri, 11 Sep 2026 21:08:14 +0000
Subject: [PATCH 6/6] test(fault_tolerance): keep the migration expectation on
 the graceful arm

The backend does interrupt an admitted request once its graceful-shutdown
period expires. `graceful_shutdown_with_discovery` sets the shared
`shutdown_event` after the grace sleep and the pre-shutdown callback but
before engine cleanup, and each handler's per-request abort monitor is
waiting on it: it calls `engine_client.abort(request_id)` and raises
`EngineShutdown`, which reaches the frontend as
`ErrorType::Backend(BackendError::EngineShutdown)` and is migratable.

So a SIGTERM'd worker does not drain its in-flight request, and the drain
assertion this branch added was asserting the wrong contract. Drop
`expect_drain` and restore the migration expectation on every graceful row,
which also makes the vLLM and TRT-LLM suites agree with SGLang instead of
contradicting it.

`graceful_worker_shutdown` stays. Removing the harness's two-second SIGKILL
escalation is what makes these rows meaningful: with it, a severed
connection produced the migration, so the row passed without ever exercising
the backend's shutdown event, abort monitor or error propagation. With
SIGTERM alone, that path is the only thing left that can produce it.

Signed-off-by: GLAMR <svc-glamr@nvidia.com>
Signed-off-by: svc-glamr@nvidia.com <svc-glamr@nvidia.com>
---
 .../fault_tolerance/migration/test_trtllm.py  |  2 -
 tests/fault_tolerance/migration/test_vllm.py  |  9 +--
 tests/fault_tolerance/migration/utils.py      | 66 +++++--------------
 3 files changed, 21 insertions(+), 56 deletions(-)

diff --git a/tests/fault_tolerance/migration/test_trtllm.py b/tests/fault_tolerance/migration/test_trtllm.py
index 4bba7b660333..41cb558dc49c 100644
--- a/tests/fault_tolerance/migration/test_trtllm.py
+++ b/tests/fault_tolerance/migration/test_trtllm.py
@@ -250,7 +250,6 @@ def test_request_migration_trtllm_aggregated(
                     immediate_kill=immediate_kill,
                     use_chat_completion=(request_api == "chat"),
                     stream=stream,
-                    expect_drain=not immediate_kill,
                     graceful_shutdown=lambda worker: graceful_worker_shutdown(
                         frontend, worker
                     ),
@@ -484,7 +483,6 @@ def test_request_migration_trtllm_decode(
                         use_chat_completion=(request_api == "chat"),
                         stream=stream,
                         wait_for_new_response_before_stop=True,
-                        expect_drain=not immediate_kill,
                         graceful_shutdown=lambda worker: graceful_worker_shutdown(
                             frontend, worker
                         ),
diff --git a/tests/fault_tolerance/migration/test_vllm.py b/tests/fault_tolerance/migration/test_vllm.py
index 9fd380d269da..b4c0ee780bad 100644
--- a/tests/fault_tolerance/migration/test_vllm.py
+++ b/tests/fault_tolerance/migration/test_vllm.py
@@ -460,8 +460,7 @@ def test_request_migration_vllm_aggregated(
                 use_chat_completion=(request_api == "chat"),
                 stream=stream,
                 max_tokens=AGGREGATED_MAX_TOKENS,
-                expected_ongoing_request_count=1 if immediate_kill else None,
-                expect_drain=not immediate_kill,
+                expected_ongoing_request_count=1,
                 graceful_shutdown=lambda worker: graceful_worker_shutdown(
                     frontend, worker
                 ),
@@ -553,8 +552,7 @@ def test_request_migration_vllm_kv_transfer(
                 max_tokens=KV_TRANSFER_MAX_TOKENS,
                 use_long_prompt=True,
                 long_prompt_repetitions=KV_TRANSFER_PROMPT_REPETITIONS,
-                expected_ongoing_request_count=1 if immediate_kill else None,
-                expect_drain=not immediate_kill,
+                expected_ongoing_request_count=1,
                 graceful_shutdown=lambda worker: graceful_worker_shutdown(
                     frontend, worker
                 ),
@@ -642,8 +640,7 @@ def test_request_migration_vllm_decode(
                 stream=True,
                 max_tokens=DECODE_MAX_TOKENS,
                 wait_for_new_response_before_stop=True,
-                expected_ongoing_request_count=1 if immediate_kill else None,
-                expect_drain=not immediate_kill,
+                expected_ongoing_request_count=1,
                 graceful_shutdown=lambda worker: graceful_worker_shutdown(
                     frontend, worker
                 ),
diff --git a/tests/fault_tolerance/migration/utils.py b/tests/fault_tolerance/migration/utils.py
index 8e55aabc3d3b..6c89472b07f2 100644
--- a/tests/fault_tolerance/migration/utils.py
+++ b/tests/fault_tolerance/migration/utils.py
@@ -478,11 +478,17 @@ def graceful_worker_shutdown(
     """Send SIGTERM only, and keep the worker alive until the outcome is known.
 
     `terminate_process_tree` escalates to SIGKILL a fixed number of seconds
-    after SIGTERM, so a worker that is still draining a generation is killed
-    mid-stream and the request migrates instead. This context sends SIGTERM,
-    waits for the worker to leave frontend discovery, yields while the request
-    outcome is observed, and only then force-kills the worker's process groups.
-    Teardown still cannot leak engine processes that pin the GPU.
+    after SIGTERM, which severs the worker's streams before its own shutdown
+    path can report anything. A migration observed that way shows only that the
+    frontend reacts to a dead connection, not that the backend interrupted the
+    request at grace expiry and reported a migratable error. Sending SIGTERM
+    alone leaves the backend's shutdown event, abort monitor and error
+    propagation as the only thing that can produce the migration.
+
+    This context sends SIGTERM, waits for the worker to leave frontend
+    discovery, yields while the request outcome is observed, and only then
+    force-kills the worker's process groups, so teardown still cannot leak
+    engine processes that pin the GPU.
 
     Args:
         frontend: Frontend whose `/health` view reports endpoint instances
@@ -787,7 +793,6 @@ def run_migration_test(
     long_prompt_repetitions: int = 8_000,
     wait_for_new_response_before_stop: bool = False,
     expected_ongoing_request_count: int | None = None,
-    expect_drain: bool = False,
     graceful_shutdown: Callable[[ManagedProcess], AbstractContextManager[None]]
     | None = None,
     verify_replacement_worker: bool = False,
@@ -814,9 +819,6 @@ def run_migration_test(
         expected_ongoing_request_count: Exact expected count for callers that
             opt into strict metric validation. When omitted, preserve the
             shared helper's historical backend-agnostic lower-bound behavior.
-        expect_drain: Assert the graceful-shutdown contract instead of the
-            migration contract: a SIGTERM'd worker finishes the request it has
-            already admitted, so the request succeeds and nothing migrates.
         graceful_shutdown: Optional backend-specific context that initiates
             graceful shutdown before response validation and performs final
             cleanup after the request outcome is known.
@@ -830,23 +832,6 @@ def run_migration_test(
             max_tokens budget so state-based fault synchronization cannot race
             an early EOS.
     """
-    if expect_drain:
-        if immediate_kill:
-            raise ValueError(
-                "expect_drain asserts the graceful-shutdown drain contract and "
-                "is incompatible with immediate_kill=True"
-            )
-        if graceful_shutdown is None:
-            raise ValueError(
-                "expect_drain requires a graceful_shutdown context; the default "
-                "shutdown path escalates to SIGKILL and cuts the drain short"
-            )
-        if expected_ongoing_request_count is not None:
-            raise ValueError(
-                "expect_drain already pins both migration counters to zero; "
-                "pass expected_ongoing_request_count only without it"
-            )
-
     # Step 1: Send the request
     if use_chat_completion:
         request_thread, response_list = start_chat_completion_request(
@@ -906,12 +891,11 @@ def run_migration_test(
             shutdown_context = graceful_shutdown(worker)
 
     # Step 5: Validate the request outcome via its response (the user-facing
-    # contract). A drained worker completes its admitted request; otherwise
-    # migration succeeds only when enabled and under the seq-len cap, else fails.
+    # contract). Migration is expected to succeed only when it is enabled and the
+    # request does not exceed the migration seq-len cap; otherwise the in-flight
+    # request must fail.
     with shutdown_context:
-        if expect_drain:
-            validate_response(request_thread, response_list)
-        elif migration_limit > 0 and migration_max_seq_len != 1:
+        if migration_limit > 0 and migration_max_seq_len != 1:
             if verify_replacement_worker:
                 wait_for_worker_request_id(
                     replacement_worker,
@@ -936,22 +920,8 @@ def run_migration_test(
     # log strings. `ongoing_request` counts an error from an established
     # stream, including an attempt that cannot retry because migration_limit is
     # zero. It is the structured equivalent of the old "Stream disconnected,
-    # recreating stream" log assertion. `max_seq_len_exceeded` records the
-    # seq-len cap from token accounting, not migration, so both arms share it.
-    expected_max_seq_len_exceeded_count = 1 if migration_max_seq_len == 1 else 0
-
-    if expect_drain:
-        # A drain migrates nothing. Under the default lower-bound mode a zero
-        # expectation asserts nothing, hence exact_counts=True.
-        verify_migration_metrics(
-            frontend.frontend_port,
-            expected_ongoing_request_count=0,
-            expected_new_request_count=0,
-            expected_max_seq_len_exceeded_count=expected_max_seq_len_exceeded_count,
-            exact_counts=True,
-        )
-        return
-
+    # recreating stream" log assertion. `max_seq_len_exceeded` records hitting
+    # the migration seq-len cap.
     exact_metric_counts = expected_ongoing_request_count is not None
     if expected_ongoing_request_count is None:
         expected_ongoing_request_count = 1 if migration_limit > 0 else 0
@@ -959,6 +929,6 @@ def run_migration_test(
     verify_migration_metrics(
         frontend.frontend_port,
         expected_ongoing_request_count=expected_ongoing_request_count,
-        expected_max_seq_len_exceeded_count=expected_max_seq_len_exceeded_count,
+        expected_max_seq_len_exceeded_count=1 if migration_max_seq_len == 1 else 0,
         exact_counts=exact_metric_counts,
     )