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

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)
@@ -543,7 +493,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,
@@ -656,7 +606,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,
@@ -756,7 +706,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,

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__)
 
@@ -236,6 +236,9 @@ def test_request_migration_trtllm_aggregated(
                     immediate_kill=immediate_kill,
                     use_chat_completion=(request_api == "chat"),
                     stream=stream,
+                    graceful_shutdown=lambda worker: graceful_worker_shutdown(
+                        frontend, worker
+                    ),
                 )
 
 
@@ -466,4 +469,7 @@ def test_request_migration_trtllm_decode(
                         use_chat_completion=(request_api == "chat"),
                         stream=stream,
                         wait_for_new_response_before_stop=True,
+                        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,
@@ -444,6 +445,9 @@ def test_request_migration_vllm_aggregated(
                 stream=stream,
                 max_tokens=AGGREGATED_MAX_TOKENS,
                 expected_ongoing_request_count=1,
+                graceful_shutdown=lambda worker: graceful_worker_shutdown(
+                    frontend, worker
+                ),
             )
 
 
@@ -533,6 +537,9 @@ def test_request_migration_vllm_kv_transfer(
                 use_long_prompt=True,
                 long_prompt_repetitions=KV_TRANSFER_PROMPT_REPETITIONS,
                 expected_ongoing_request_count=1,
+                graceful_shutdown=lambda worker: graceful_worker_shutdown(
+                    frontend, worker
+                ),
             )
 
 
@@ -618,4 +625,7 @@ def test_request_migration_vllm_decode(
                 max_tokens=DECODE_MAX_TOKENS,
                 wait_for_new_response_before_stop=True,
                 expected_ongoing_request_count=1,
+                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
@@ -22,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(
@@ -463,6 +470,69 @@ def wait_for_endpoint_instance_reduction(
     )
 
 
+@contextmanager
+def graceful_worker_shutdown(
+    frontend: DynamoFrontendProcess,
+    worker: ManagedProcess,
+) -> 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, 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
+        worker: Worker to shut down
+    """
+    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")) == BACKEND_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:
+                pass
+    except psutil.NoSuchProcess:
+        pass
+
+    try:
+        parent.terminate()
+        wait_for_endpoint_instance_reduction(
+            frontend.frontend_port,
+            BACKEND_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,