test(fault_tolerance): share a SIGTERM-only graceful shutdown context

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
+                        ),
                     )

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
+                ),
             )

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,14 @@ 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 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
             cleanup after the request outcome is known.
@@ -762,6 +834,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 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(
@@ -821,11 +910,12 @@ 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 completes its admitted request; otherwise
+    # migration succeeds only when enabled and under the seq-len cap, else fails.
     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,
@@ -850,15 +940,29 @@ 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.
+    # 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
+
     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
 
     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,
     )