diff --git a/tensorrt_llm/executor/executor.py b/tensorrt_llm/executor/executor.py index 7e8b565683a8..6738aa4986a3 100644 --- a/tensorrt_llm/executor/executor.py +++ b/tensorrt_llm/executor/executor.py @@ -337,6 +337,16 @@ def is_shutdown(self) -> bool: """Return True if the executor is shutting down or fatally errored.""" return self.doing_shutdown or self._fatal_error is not None + @property + def fatal_error(self) -> Optional[BaseException]: + """The unrecoverable engine error recorded by ``_set_fatal_error``, or None.""" + return self._fatal_error + + @property + def engine_dead(self) -> bool: + """Whether the engine is known dead (an unrecoverable error was recorded).""" + return self._fatal_error is not None + def check_health(self) -> bool: """Check whether the executor is healthy and able to process requests. diff --git a/tensorrt_llm/executor/proxy.py b/tensorrt_llm/executor/proxy.py index 3e0f3f919ecc..765e9a001bef 100644 --- a/tensorrt_llm/executor/proxy.py +++ b/tensorrt_llm/executor/proxy.py @@ -480,6 +480,15 @@ def check_health(self) -> bool: return True + @property + def engine_dead(self) -> bool: + """Whether the engine is known dead (a worker died or a fatal error was recorded). + + ``_mark_engine_dead`` can set the sticky flag before the fatal error + is recorded, so check both. + """ + return self._engine_dead or self._fatal_error is not None + def _set_fatal_error(self, error: BaseException) -> None: """Record the fatal error, then unblock every pending request. diff --git a/tensorrt_llm/grpc/openengine/control.py b/tensorrt_llm/grpc/openengine/control.py index 47278788dc33..3275880862e8 100644 --- a/tensorrt_llm/grpc/openengine/control.py +++ b/tensorrt_llm/grpc/openengine/control.py @@ -140,23 +140,19 @@ async def _get_kv_cache_load(self) -> dict: def _engine_is_healthy(self) -> bool: """Readiness via the predicate the engine's own HTTP /health uses. - `_check_health` is stronger than the `is_shutdown()` guard generate_async + `check_health()` is stronger than the `is_shutdown()` guard generate_async raises on: it drains the executor's error queue, so a fatal error queued but not yet promoted is caught, and the proxy polls MPI worker liveness. - It also covers encode-only engines, where `_executor` is None. + It also covers encode-only engines, where there is no generation executor. """ - check = getattr(self._llm, "_check_health", None) - if callable(check): - try: - return bool(check()) - except (RuntimeError, ValueError, OSError) as error: - logger.warning(f"OpenEngine health check reported a failure: {error}") - return False - except (TypeError, AttributeError) as error: - logger.error(f"OpenEngine health check contract failed: {error}") - return False - executor = getattr(self._llm, "_executor", None) - return executor is not None and not executor.is_shutdown() + try: + return bool(self._llm.check_health()) + except (RuntimeError, ValueError, OSError) as error: + logger.warning(f"OpenEngine health check reported a failure: {error}") + return False + except (TypeError, AttributeError) as error: + logger.error(f"OpenEngine health check contract failed: {error}") + return False # -- Runtime metadata --------------------------------------------------- diff --git a/tensorrt_llm/grpc/openengine/server.py b/tensorrt_llm/grpc/openengine/server.py index 97ccb3210637..b495db5f6c1d 100644 --- a/tensorrt_llm/grpc/openengine/server.py +++ b/tensorrt_llm/grpc/openengine/server.py @@ -215,11 +215,7 @@ async def _refresh_inference_health(self) -> None: except (RuntimeError, ValueError, OSError, CoordinationError) as error: logger.warning(f"OpenEngine inference health check failed: {error}") self._engine_healthy = False - executor = getattr(self._llm, "_executor", None) - if executor is not None and ( - getattr(executor, "_fatal_error", None) is not None - or getattr(executor, "_engine_dead", False) - ): + if self._llm.engine_dead: self._fatal_engine_failure = True if self._stopping: return @@ -412,11 +408,7 @@ def healthy() -> bool: ) def terminal_engine_failure() -> bool: - executor = getattr(llm, "_executor", None) - return executor is not None and ( - getattr(executor, "_fatal_error", None) is not None - or getattr(executor, "_engine_dead", False) - ) + return llm.engine_dead async def supervise() -> None: while not stop_event.is_set(): @@ -428,7 +420,7 @@ async def supervise() -> None: return if server is None: engine_ready = await asyncio.wait_for( - asyncio.to_thread(llm._check_health), + asyncio.to_thread(llm.check_health), timeout=_HEALTH_CHECK_TIMEOUT_SECONDS, ) if server is None and not engine_ready: diff --git a/tensorrt_llm/grpc/smg/request_manager.py b/tensorrt_llm/grpc/smg/request_manager.py index 460d19804b0b..01a527691302 100644 --- a/tensorrt_llm/grpc/smg/request_manager.py +++ b/tensorrt_llm/grpc/smg/request_manager.py @@ -168,20 +168,21 @@ async def health_check(self) -> Tuple[bool, str]: if self.llm is None: return False, "LLM not initialized" - # Check executor health (includes error queue, MPI worker - # liveness, and fatal error state — not just doing_shutdown). - if hasattr(self.llm, "_executor"): - if self.llm._executor is None: - return False, "Executor is not available" - if not self.llm._executor.check_health(): - error_msg = "Executor is unhealthy" - if self.llm._executor._fatal_error is not None: - exc = self.llm._executor._fatal_error - lines = str(exc).splitlines() - short = (lines[0] if lines else type(exc).__name__)[:200] - error_msg = f"{type(exc).__name__}: {short}" - logger.error(f"Health check fatal error: {repr(exc)}") - return False, error_msg + # An encode-only LLM has no generation executor; SMG serves + # generate requests only. + if self.llm.args.encode_only is True: + return False, "Executor is not available" + # check_health() drains queued errors and checks MPI worker + # liveness and the fatal error state, not just shutdown. + if not self.llm.check_health(): + error_msg = "Executor is unhealthy" + exc = self.llm.fatal_error + if exc is not None: + lines = str(exc).splitlines() + short = (lines[0] if lines else type(exc).__name__)[:200] + error_msg = f"{type(exc).__name__}: {short}" + logger.error(f"Health check fatal error: {repr(exc)}") + return False, error_msg return True, "OK" except Exception as e: diff --git a/tensorrt_llm/llmapi/llm.py b/tensorrt_llm/llmapi/llm.py index 2d9dabd427d9..985e84978550 100644 --- a/tensorrt_llm/llmapi/llm.py +++ b/tensorrt_llm/llmapi/llm.py @@ -339,6 +339,9 @@ def _multimodal_params_have_encoder_features( llm_id (str): The unique ID of the LLM instance. disaggregated_params (dict): The disaggregated parameters of the LLM instance. startup_metrics (dict): The startup metrics reported by worker rank 0. + fatal_error (BaseException, optional): The unrecoverable engine error the executor recorded, or None. + shutdown_started (bool): Whether shutdown of the generation executor has begun. + engine_dead (bool): Whether the engine is known dead (a worker died or a fatal error was recorded). """ @@ -1785,11 +1788,20 @@ def _record_usage_shutdown(self) -> None: except Exception as exc: logger.debug(f"Usage telemetry shutdown tracking failed: {exc}") - def _check_health(self) -> bool: - """Check if the LLM is healthy. + @set_api_status("prototype") + def check_health(self) -> bool: + """Check whether the LLM can serve requests. + + For ``encode_only=True`` the LLM is healthy while its encoder executor + exists. Otherwise this delegates to the generation executor's + ``check_health()``, which drains queued background errors (a fatal one + is recorded, see ``fatal_error``, and starts executor shutdown) and, + for the multi-process proxy, polls worker liveness. Without a + generation executor (never built, or after ``shutdown()``) the LLM is + unhealthy. Returns: - bool: True if the executor is running and healthy, False otherwise. + bool: True if the LLM is healthy, False otherwise. """ if self._encode_only: return (hasattr(self, "_encoder_executor") @@ -1799,6 +1811,46 @@ def _check_health(self) -> bool: return False + # Kept for one release for out-of-tree callers that look this name up + # (ai-dynamo components/src/dynamo/trtllm/engine.py); use check_health(). + _check_health = check_health + + @property + @set_api_status("prototype") + def fatal_error(self) -> Optional[BaseException]: + """The unrecoverable engine error the executor recorded, if any. + + None when no fatal error was recorded or the LLM has no generation + executor (``encode_only=True``, or after ``shutdown()``). + """ + executor = getattr(self, "_executor", None) + return None if executor is None else executor.fatal_error + + @property + @set_api_status("prototype") + def shutdown_started(self) -> bool: + """Whether shutdown of the generation executor has begun. + + True once ``shutdown()`` was called or the executor started shutting + down after a fatal error. False for ``encode_only=True``, which has no + generation executor. + """ + executor = getattr(self, "_executor", None) + if executor is None: + return not getattr(self, "_encode_only", False) + return executor.doing_shutdown + + @property + @set_api_status("prototype") + def engine_dead(self) -> bool: + """Whether the engine is known dead (a worker died or a fatal error was recorded). + + False when the LLM has no generation executor (``encode_only=True``, + or after ``shutdown()``). + """ + executor = getattr(self, "_executor", None) + return False if executor is None else executor.engine_dead + @staticmethod def _shutdown_wrapper(self_ref): # Retrieve the instance if it still exists diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index 5079b9b073e4..d8adde6eafd7 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -1227,9 +1227,11 @@ def _check_health(self) -> bool: batcher = self.embedding_batcher if batcher is not None and not batcher.is_alive(): return False - if hasattr(self.generator, '_check_health'): - return self.generator._check_health() - return True + if self._is_visual_gen: + # VisualGen has no check_health(), so its /health always + # answers 200. + return True + return self.generator.check_health() def register_routes(self): self.app.add_api_route("/health", self.health, methods=["GET"]) @@ -1727,15 +1729,14 @@ async def health(self) -> Response: # its graceful shutdown sequence. The doing_shutdown guard # ensures we only send SIGINT once — repeated health probes # won't stack signals during an in-progress teardown. - executor = getattr(self.generator, '_executor', None) - if executor is not None and getattr(executor, '_fatal_error', - None) is not None: - if not getattr(executor, 'doing_shutdown', True): - logger.error( - "Health check detected fatal engine error, initiating " - f"server shutdown: {executor._fatal_error}") - _record_generator_termination(self.generator) - signal.raise_signal(signal.SIGINT) + fatal_error = (None if self._is_visual_gen else + self.generator.fatal_error) + if fatal_error is not None and not self.generator.shutdown_started: + logger.error( + "Health check detected fatal engine error, initiating " + f"server shutdown: {fatal_error}") + _record_generator_termination(self.generator) + signal.raise_signal(signal.SIGINT) return Response( status_code=503, content= diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 550796e3fa0c..69af33f55635 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -236,6 +236,8 @@ l0_cpu: - unittest/grpc/openengine/test_openengine.py - unittest/grpc/openengine/test_metrics.py - unittest/grpc/openengine/test_servicer.py + - unittest/grpc/openengine/test_private_access_lint_openengine.py + - unittest/grpc/smg/test_smg_health.py - unittest/grpc/smg/test_smg.py::test_streaming_deltas_preserve_sequences_logprobs_and_final_tokens - unittest/grpc/smg/test_smg.py::test_streaming_empty_token_history - unittest/grpc/test_grpc_optional.py::test_protobuf_warning_in_fresh_process @@ -295,6 +297,7 @@ l0_cpu: - unittest/llmapi/test_kv_cache_dtype_override.py - unittest/llmapi/test_llm_args.py - unittest/llmapi/test_gms_args.py + - unittest/llmapi/test_llm_public_runtime_api.py - unittest/llmapi/test_mx_args.py - unittest/llmapi/test_custom_tokenizer_aliases.py # Reads the Muse Spark Open tokenizer files from LLM_MODELS_ROOT. @@ -334,6 +337,9 @@ l0_cpu: - unittest/quantization/test_mode.py - unittest/scaffolding - unittest/scripts + - unittest/serve/api/visual_gen/test_private_access_lint_visual_gen.py + - unittest/serve/test_private_access_lint.py + - unittest/serve/test_serve_health.py - unittest/tools - unittest/usage - unittest/utils/test_logger.py diff --git a/tests/unittest/api_stability/references/llm.yaml b/tests/unittest/api_stability/references/llm.yaml index c75fe8fe084b..b1be771de618 100644 --- a/tests/unittest/api_stability/references/llm.yaml +++ b/tests/unittest/api_stability/references/llm.yaml @@ -528,6 +528,10 @@ methods: default: 2 return_annotation: tensorrt_llm.executor.result.IterationResult status: beta + check_health: + parameters: {} + return_annotation: bool + status: prototype get_data_transceiver_state: parameters: {} return_annotation: bytes @@ -584,3 +588,12 @@ properties: startup_metrics: annotation: dict default: None + fatal_error: + annotation: Optional[BaseException] + default: inspect._empty + shutdown_started: + annotation: bool + default: inspect._empty + engine_dead: + annotation: bool + default: inspect._empty diff --git a/tests/unittest/executor/test_fatal_error_health_check.py b/tests/unittest/executor/test_fatal_error_health_check.py index 3f8f7396a6ff..77a4b2698275 100644 --- a/tests/unittest/executor/test_fatal_error_health_check.py +++ b/tests/unittest/executor/test_fatal_error_health_check.py @@ -20,9 +20,11 @@ - GenerationExecutor.check_health(), _set_fatal_error(), is_shutdown() - GenerationExecutorProxy.check_health() with MPI worker liveness - GenerationExecutorProxy._error_monitor_loop() background detection -- GrpcRequestManager.health_check() integration -- OpenAIServer /health endpoint with fatal error shutdown -- BaseLLM._check_health() delegation + +Frontend health is tested against production code in +tests/unittest/llmapi/test_llm_public_runtime_api.py, +tests/unittest/serve/test_serve_health.py and +tests/unittest/grpc/smg/test_smg_health.py. """ # Import classify_error directly from the source file to avoid triggering @@ -30,14 +32,13 @@ import importlib.util import logging import pathlib -import signal import subprocess import sys import threading import time from concurrent.futures import Future from queue import Empty, Queue -from unittest.mock import MagicMock, Mock, patch +from unittest.mock import Mock, patch import pytest @@ -827,157 +828,6 @@ def test_stops_on_shutdown_flag(self): assert not t.is_alive() -# --------------------------------------------------------------------------- -# gRPC health check -# --------------------------------------------------------------------------- -class TestGrpcHealthCheck: - """Tests for gRPC health check integration with fatal error state.""" - - @staticmethod - async def _health_check(self): - """Mirrors GrpcRequestManager.health_check logic.""" - try: - if self.llm is None: - return False, "LLM not initialized" - if hasattr(self.llm, "_executor"): - if self.llm._executor is None: - return False, "Executor is not available" - if not self.llm._executor.check_health(): - error_msg = "Executor is unhealthy" - if self.llm._executor._fatal_error is not None: - exc = self.llm._executor._fatal_error - short = str(exc).splitlines()[0][:200] - error_msg = f"{type(exc).__name__}: {short}" - return False, error_msg - return True, "OK" - except Exception as e: - return False, f"Error: {e}" - - def _make_manager(self, executor=None): - """Create a mock gRPC manager with the given executor.""" - m = MagicMock() - m.llm = MagicMock() - m.llm._executor = executor - m.health_check = self._health_check.__get__(m) - return m - - @pytest.mark.asyncio - @pytest.mark.parametrize( - "setup,expected_healthy,msg_substr", - [ - ("healthy", True, "OK"), - ("fatal", False, "CUDA OOM"), - ("no_executor", False, "not available"), - ("no_llm", False, "not initialized"), - ("shutdown", False, "unhealthy"), - ], - ) - async def test_health_check(self, setup, expected_healthy, msg_substr): - """Verify gRPC health check returns correct status and message.""" - if setup == "no_llm": - m = MagicMock() - m.llm = None - m.health_check = self._health_check.__get__(m) - else: - ex = ConcreteExecutor() - if setup == "fatal": - ex._set_fatal_error(RuntimeError("CUDA OOM")) - elif setup == "no_executor": - ex = None - elif setup == "shutdown": - ex.doing_shutdown = True - m = self._make_manager(ex) - - healthy, msg = await m.health_check() - assert healthy is expected_healthy - assert msg_substr in msg - - -# --------------------------------------------------------------------------- -# OpenAI /health endpoint -# --------------------------------------------------------------------------- -class TestOpenAIHealthEndpoint: - """Tests for OpenAI /health endpoint with fatal error shutdown.""" - - @pytest.mark.asyncio - @pytest.mark.parametrize( - "check_healthy,fatal_error,expect_code,sigint", - [ - (True, None, 200, False), - (False, None, 503, False), - (False, RuntimeError("CUDA OOM"), 503, True), - ], - ) - async def test_health(self, check_healthy, fatal_error, expect_code, sigint): - """Verify status code and SIGINT behavior for different health states.""" - from starlette.responses import Response - - server = MagicMock() - server._check_health = Mock(return_value=check_healthy) - executor = Mock() - executor._fatal_error = fatal_error - executor.doing_shutdown = False - server.generator = Mock() - server.generator._executor = executor - - with patch.object(signal, "raise_signal") as mock_sig: - if server._check_health(): - response = Response(status_code=200) - else: - ex = getattr(server.generator, "_executor", None) - if ( - ex is not None - and getattr(ex, "_fatal_error", None) is not None - and not getattr(ex, "doing_shutdown", True) - ): - signal.raise_signal(signal.SIGINT) - response = Response(status_code=503) - - assert response.status_code == expect_code - if sigint: - mock_sig.assert_called_once_with(signal.SIGINT) - else: - mock_sig.assert_not_called() - - -# --------------------------------------------------------------------------- -# BaseLLM._check_health delegation -# --------------------------------------------------------------------------- -class TestBaseLLMCheckHealth: - """Tests for BaseLLM._check_health delegation to executor.""" - - @staticmethod - def _check_health(llm) -> bool: - """Mirrors BaseLLM._check_health logic.""" - if hasattr(llm, "_executor") and llm._executor is not None: - return llm._executor.check_health() - return False - - @pytest.mark.parametrize( - "has_executor,fatal,expected", - [ - (True, False, True), - (True, True, False), - (False, False, False), - ], - ) - def test_delegation(self, has_executor, fatal, expected): - """_check_health delegates to executor.check_health correctly.""" - llm = Mock() - if has_executor: - ex = ConcreteExecutor() - if fatal: - ex._set_fatal_error(RuntimeError("x")) - llm._executor = ex - else: - llm._executor = None - assert self._check_health(llm) is expected - - def test_no_executor_attr(self): - """Returns False when the object has no _executor attribute.""" - assert self._check_health(object()) is False - - # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- diff --git a/tests/unittest/grpc/openengine/test_control.py b/tests/unittest/grpc/openengine/test_control.py index 8723dde4d6eb..38149c34e7b1 100644 --- a/tests/unittest/grpc/openengine/test_control.py +++ b/tests/unittest/grpc/openengine/test_control.py @@ -128,7 +128,7 @@ def _llm(**overrides): supports_token_ids_input=overrides.get("_supports_token_ids_input", True) ), _executor=executor, - _check_health=lambda: executor.check_health(), + check_health=lambda: executor.check_health(), ) @@ -610,7 +610,7 @@ async def test_health_reports_not_ready_for_a_crashed_engine_that_is_not_shut_do async def test_health_survives_a_raising_health_check(error_type): """A probe reports state; it must not turn an engine fault into an RPC error.""" llm = _llm() - llm._check_health = lambda: (_ for _ in ()).throw(error_type("engine gone")) + llm.check_health = lambda: (_ for _ in ()).throw(error_type("engine gone")) servicer = OpenEngineControlServicer(llm, MODEL, _inference(), kv_transfer_backend="") response = await servicer.Health(lifecycle_pb2.HealthRequest(), FakeServicerContext()) diff --git a/tests/unittest/grpc/openengine/test_coordinator.py b/tests/unittest/grpc/openengine/test_coordinator.py index 22cade427513..01a83bfb8a83 100644 --- a/tests/unittest/grpc/openengine/test_coordinator.py +++ b/tests/unittest/grpc/openengine/test_coordinator.py @@ -202,6 +202,7 @@ async def __aiter__(self): handle = ResultHandle() llm = SimpleNamespace( + engine_dead=False, llm_id="shared-engine", args=SimpleNamespace(guided_decoding_backend=None), tokenizer=None, @@ -252,10 +253,11 @@ async def cleaned_up(): async def test_standard_health_tracks_frontend_group_readiness(): async with group() as (coordinator, (frontend, _), _): llm = SimpleNamespace( + engine_dead=False, llm_id="shared-engine", args=SimpleNamespace(guided_decoding_backend=None), tokenizer=None, - _check_health=lambda: True, + check_health=lambda: True, ) server = OpenEngineServer("127.0.0.1", 0, llm, "test-model", frontend=frontend) await server.start() @@ -314,10 +316,11 @@ async def aresult(self): return None llm = SimpleNamespace( + engine_dead=False, llm_id="shared-engine", args=SimpleNamespace(guided_decoding_backend=None), tokenizer=None, - _check_health=lambda: True, + check_health=lambda: True, generate_async=lambda *args, **kwargs: ProbeHandle(), ) original_request = frontend.request @@ -354,6 +357,7 @@ async def test_oversized_request_id_does_not_stop_group(request_id): """A rejected public ID must not kill unrelated requests or strand ownership.""" async with group() as (_, (first, second), failures): llm = SimpleNamespace( + engine_dead=False, llm_id="shared-engine", args=SimpleNamespace(guided_decoding_backend=None), tokenizer=None, diff --git a/tests/unittest/grpc/openengine/test_frontend_lifecycle.py b/tests/unittest/grpc/openengine/test_frontend_lifecycle.py index 6732fd5b76f7..ca4ef9a222af 100644 --- a/tests/unittest/grpc/openengine/test_frontend_lifecycle.py +++ b/tests/unittest/grpc/openengine/test_frontend_lifecycle.py @@ -176,9 +176,10 @@ def on_term(signum, frame): self.llm_id = "shared-engine" self.args = SimpleNamespace(guided_decoding_backend=None) self.tokenizer = None + self.engine_dead = False self._executor = Mock(spec=GenerationExecutorProxy) self._executor.multi_frontend_attach_info.return_value = {"mode": "classic"} - def _check_health(self): + def check_health(self): return not ( frontend == 0 and sys.argv[3] == "child-startup-engine-failure" and (root / "initializing").exists() diff --git a/tests/unittest/grpc/openengine/test_metrics.py b/tests/unittest/grpc/openengine/test_metrics.py index 3e97cc47d4b5..32dd0edb402e 100644 --- a/tests/unittest/grpc/openengine/test_metrics.py +++ b/tests/unittest/grpc/openengine/test_metrics.py @@ -151,7 +151,7 @@ async def __aiter__(self): llm = _FakeLlm([]) llm.args = _args() llm.llm_id = "test-engine" - llm._check_health = lambda: True + llm.check_health = lambda: True reads = [] def get_stats(timeout): diff --git a/tests/unittest/grpc/openengine/test_openengine.py b/tests/unittest/grpc/openengine/test_openengine.py index faf6e92ace17..4e6456a8cfe4 100644 --- a/tests/unittest/grpc/openengine/test_openengine.py +++ b/tests/unittest/grpc/openengine/test_openengine.py @@ -100,6 +100,7 @@ def test_openengine_server_serves_the_control_contract() -> None: answering rather than falling through to the generated base. """ llm = SimpleNamespace( + engine_dead=False, args=SimpleNamespace( max_seq_len=2048, max_batch_size=8, @@ -113,7 +114,7 @@ def test_openengine_server_serves_the_control_contract() -> None: ), llm_id="test-instance", tokenizer=object(), - _check_health=lambda: True, + check_health=lambda: True, ) async def exercise_server() -> None: @@ -160,10 +161,11 @@ def test_health_and_control_are_reachable_on_pod_address() -> None: pytest.skip("No non-loopback IPv4 address is available") llm = SimpleNamespace( + engine_dead=False, args=SimpleNamespace(guided_decoding_backend=None), llm_id="test-instance", tokenizer=None, - _check_health=lambda: True, + check_health=lambda: True, ) async def exercise_server() -> None: @@ -194,10 +196,11 @@ def test_standard_grpc_health_check_and_watch(monkeypatch: pytest.MonkeyPatch) - monkeypatch.setattr(oe_server, "_HEALTH_CHECK_INTERVAL_SECONDS", 0.01) engine = {"healthy": False} llm = SimpleNamespace( + engine_dead=False, args=SimpleNamespace(guided_decoding_backend=None), llm_id="test-instance", tokenizer=None, - _check_health=lambda: engine["healthy"], + check_health=lambda: engine["healthy"], ) async def exercise_server() -> None: @@ -264,10 +267,11 @@ async def check(service: str) -> int: def test_standard_grpc_health_watch_reports_shutdown() -> None: """Shutdown withdraws inference and finishes watchers without failing liveness.""" llm = SimpleNamespace( + engine_dead=False, args=SimpleNamespace(guided_decoding_backend=None), llm_id="test-instance", tokenizer=None, - _check_health=lambda: True, + check_health=lambda: True, ) async def exercise_server() -> None: @@ -319,10 +323,11 @@ def check_health() -> bool: return True llm = SimpleNamespace( + engine_dead=False, args=SimpleNamespace(guided_decoding_backend=None), llm_id="test-instance", tokenizer=None, - _check_health=check_health, + check_health=check_health, ) async def exercise_server() -> None: @@ -373,13 +378,22 @@ def test_fatal_engine_failure_fails_liveness_after_temporary_unavailability( monkeypatch.setattr(oe_server, "_HEALTH_CHECK_INTERVAL_SECONDS", 0.01) engine = {"healthy": True} executor = SimpleNamespace(_fatal_error=None, _engine_dead=False) - llm = SimpleNamespace( - args=SimpleNamespace(guided_decoding_backend=None), - llm_id="test-instance", - tokenizer=None, - _executor=executor, - _check_health=lambda: engine["healthy"] and executor._fatal_error is None, - ) + + class EngineFailureLLM: + """LLM double whose engine_dead derives from the executor's two fatal fields.""" + + args = SimpleNamespace(guided_decoding_backend=None) + llm_id = "test-instance" + tokenizer = None + + def check_health(self) -> bool: + return engine["healthy"] and executor._fatal_error is None + + @property + def engine_dead(self) -> bool: + return executor._fatal_error is not None or bool(executor._engine_dead) + + llm = EngineFailureLLM() async def exercise_server() -> None: server = OpenEngineServer("127.0.0.1", 0, llm, "test-model") @@ -450,10 +464,11 @@ def check_health() -> bool: return True llm = SimpleNamespace( + engine_dead=False, args=SimpleNamespace(guided_decoding_backend=None), llm_id="test-instance", tokenizer=None, - _check_health=check_health, + check_health=check_health, ) async def exercise_server() -> None: @@ -504,6 +519,7 @@ def test_event_startup_requires_loaded_publisher(publisher_enabled: bool) -> Non from tensorrt_llm.grpc.openengine.kv_events import KvEventsUnavailableError llm = SimpleNamespace( + engine_dead=False, args=SimpleNamespace( guided_decoding_backend=None, kv_cache_config=SimpleNamespace( diff --git a/tests/unittest/grpc/openengine/test_private_access_lint_openengine.py b/tests/unittest/grpc/openengine/test_private_access_lint_openengine.py new file mode 100644 index 000000000000..cee2100effc5 --- /dev/null +++ b/tests/unittest/grpc/openengine/test_private_access_lint_openengine.py @@ -0,0 +1,33 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""OpenEngine part of the private-access lint. + +Pull requests that change only tensorrt_llm/grpc/openengine/ are claimed by CBTS +Tier 1 and run only the test-db entries under unittest/grpc/openengine/. This file +applies the rule and baseline of tests/unittest/serve/test_private_access_lint.py +to tensorrt_llm/grpc/openengine/, so those pull requests are still checked. +""" + +import importlib.util +from pathlib import Path + +import pytest + +pytestmark = pytest.mark.cpu_only + +_MAIN = Path(__file__).resolve().parents[2] / "serve" / "test_private_access_lint.py" +_spec = importlib.util.spec_from_file_location("_serve_private_access_lint", _MAIN) +lint = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(lint) + + +def test_no_new_cross_package_private_access_in_openengine(): + assert (lint.REPO_ROOT / lint.OPENENGINE_SCOPE).is_dir(), ( + f"{lint.OPENENGINE_SCOPE} does not exist. If OpenEngine moved, point OPENENGINE_SCOPE " + "in tests/unittest/serve/test_private_access_lint.py at the new directory, and move " + "this file and its l0_cpu.yml entry to the test path the CBTS OpenEngine rule keeps " + "(jenkins/scripts/cbts/cbts/rules/openengine.py)." + ) + rows = lint.scan(prefixes=(lint.OPENENGINE_SCOPE,)) + errors = lint._compare(rows, lint.load_baseline(), prefix=lint.OPENENGINE_SCOPE) + assert not errors, "\n".join(errors) diff --git a/tests/unittest/grpc/openengine/test_servicer.py b/tests/unittest/grpc/openengine/test_servicer.py index a953bceba01f..178fdc51fbaa 100644 --- a/tests/unittest/grpc/openengine/test_servicer.py +++ b/tests/unittest/grpc/openengine/test_servicer.py @@ -95,6 +95,7 @@ def __init__(self, results: list[Any]) -> None: self.tokenizer = _FakeTokenizer() self.result_handle = _FakeResultHandle(results) self.generate_kwargs = None + self.engine_dead = False def generate_async(self, **kwargs: Any) -> _FakeResultHandle: self.generate_kwargs = kwargs diff --git a/tests/unittest/grpc/smg/test_smg_health.py b/tests/unittest/grpc/smg/test_smg_health.py new file mode 100644 index 000000000000..33467eadca8c --- /dev/null +++ b/tests/unittest/grpc/smg/test_smg_health.py @@ -0,0 +1,71 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Production tests of GrpcRequestManager.health_check. + +The (healthy, message) pair feeds HealthCheckResponse.status, so every message +string is wire-visible and pinned byte for byte, including the first-line rule +and the 200-character cut for fatal errors. +""" + +import asyncio +from types import SimpleNamespace + +import pytest + +pytestmark = pytest.mark.cpu_only + +pytest.importorskip("smg_grpc_proto") + +from tensorrt_llm.grpc.smg.request_manager import GrpcRequestManager # noqa: E402 + + +class _FakeLLM: + """LLM double with exactly the public health surface health_check reads.""" + + def __init__(self, healthy=True, fatal_error=None, encode_only=False, check_raises=None): + self.args = SimpleNamespace(encode_only=encode_only) + self.fatal_error = fatal_error + self._healthy = healthy + self._check_raises = check_raises + + def check_health(self): + if self._check_raises is not None: + raise self._check_raises + return self._healthy + + +def _health_check(llm): + return asyncio.run(GrpcRequestManager(llm).health_check()) + + +def test_healthy_llm_answers_ok(): + assert _health_check(_FakeLLM(healthy=True)) == (True, "OK") + + +def test_missing_llm_answers_not_initialized(): + assert _health_check(None) == (False, "LLM not initialized") + + +def test_encode_only_llm_answers_executor_not_available(): + assert _health_check(_FakeLLM(encode_only=True)) == (False, "Executor is not available") + + +def test_unhealthy_llm_without_fatal_error_answers_executor_unhealthy(): + assert _health_check(_FakeLLM(healthy=False)) == (False, "Executor is unhealthy") + + +def test_fatal_error_message_names_the_exception_type(): + result = _health_check(_FakeLLM(healthy=False, fatal_error=RuntimeError("CUDA OOM"))) + assert result == (False, "RuntimeError: CUDA OOM") + + +def test_fatal_error_message_keeps_the_first_line_cut_to_200_characters(): + long_first_line = "x" * 300 + result = _health_check( + _FakeLLM(healthy=False, fatal_error=RuntimeError(f"{long_first_line}\nsecond line")) + ) + assert result == (False, f"RuntimeError: {'x' * 200}") + + +def test_raising_health_check_answers_error(): + assert _health_check(_FakeLLM(check_raises=ValueError("boom"))) == (False, "Error: boom") diff --git a/tests/unittest/llmapi/test_llm_encode.py b/tests/unittest/llmapi/test_llm_encode.py index 44d9db731e1c..5714f41a2979 100644 --- a/tests/unittest/llmapi/test_llm_encode.py +++ b/tests/unittest/llmapi/test_llm_encode.py @@ -237,8 +237,8 @@ def test_encode_add_special_tokens_false(bert_encode_llm): def test_check_health_encoder_only(bert_encode_llm): - """_check_health() returns True for a live encoder-only LLM.""" - assert bert_encode_llm._check_health() is True + """check_health() returns True for a live encoder-only LLM.""" + assert bert_encode_llm.check_health() is True # --------------------------------------------------------------------------- # diff --git a/tests/unittest/llmapi/test_llm_public_runtime_api.py b/tests/unittest/llmapi/test_llm_public_runtime_api.py new file mode 100644 index 000000000000..e1f9b321d1ee --- /dev/null +++ b/tests/unittest/llmapi/test_llm_public_runtime_api.py @@ -0,0 +1,181 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Production tests of the public runtime health API the serving frontends use. + +Covers BaseLLM.check_health(), BaseLLM.fatal_error, BaseLLM.shutdown_started, +BaseLLM.engine_dead and the executor-side GenerationExecutor.fatal_error and +engine_dead, against the real classes (no mirrors; the LLM objects are built +with object.__new__ plus exactly the attributes BaseLLM.__init__ sets). +""" + +import pytest + +from tensorrt_llm.executor.executor import GenerationExecutor +from tensorrt_llm.executor.proxy import GenerationExecutorProxy +from tensorrt_llm.llmapi.llm import LLM + +pytestmark = pytest.mark.cpu_only + + +class _NoOpExecutor(GenerationExecutor): + """Minimal concrete GenerationExecutor; health logic comes from the base.""" + + def submit(self, request): + raise NotImplementedError("health tests never submit requests") + + def abort_request(self, request_id): + pass + + def shutdown(self): + self.doing_shutdown = True + + +def _llm(executor, encode_only=False, encoder_executor=None): + """An LLM with exactly the attributes BaseLLM.__init__ sets for health.""" + llm = object.__new__(LLM) + llm._executor = executor + llm._encode_only = encode_only + llm._encoder_executor = encoder_executor + return llm + + +def _proxy(engine_dead=False, fatal_error=None): + """A proxy with exactly the two fields its engine_dead property reads.""" + proxy = object.__new__(GenerationExecutorProxy) + proxy.workers_started = False + proxy._engine_dead = engine_dead + proxy._fatal_error = fatal_error + return proxy + + +# --------------------------------------------------------------------------- +# BaseLLM.check_health() +# --------------------------------------------------------------------------- +def test_check_health_encode_only_with_encoder_executor(): + assert _llm(None, encode_only=True, encoder_executor=_NoOpExecutor()).check_health() is True + + +def test_check_health_encode_only_without_encoder_executor(): + assert _llm(None, encode_only=True, encoder_executor=None).check_health() is False + + +def test_check_health_healthy_executor(): + assert _llm(_NoOpExecutor()).check_health() is True + + +def test_check_health_fatal_error(): + executor = _NoOpExecutor() + executor._set_fatal_error(RuntimeError("CUDA OOM")) + assert _llm(executor).check_health() is False + + +def test_check_health_doing_shutdown(): + executor = _NoOpExecutor() + executor.doing_shutdown = True + assert _llm(executor).check_health() is False + + +def test_check_health_without_executor(): + assert _llm(None).check_health() is False + + +def test_check_health_drains_queued_fatal_error(): + """A fatal error still in the error queue is promoted by check_health().""" + executor = _NoOpExecutor() + queued = RuntimeError("engine died in a background thread") + executor._error_queue.put(queued) + + llm = _llm(executor) + assert llm.check_health() is False + assert llm.fatal_error is queued + assert llm.shutdown_started is True + + +def test_check_health_private_alias_is_the_same_function(): + """Kept for one release for out-of-tree callers that look the name up.""" + assert LLM._check_health is LLM.check_health + + +# --------------------------------------------------------------------------- +# BaseLLM.fatal_error +# --------------------------------------------------------------------------- +def test_fatal_error_none_without_executor(): + assert _llm(None).fatal_error is None + + +def test_fatal_error_none_on_healthy_executor(): + assert _llm(_NoOpExecutor()).fatal_error is None + + +def test_fatal_error_returns_recorded_exception_first_one_wins(): + executor = _NoOpExecutor() + first = RuntimeError("first fatal error") + executor._set_fatal_error(first) + executor._set_fatal_error(RuntimeError("second fatal error")) + assert _llm(executor).fatal_error is first + + +# --------------------------------------------------------------------------- +# BaseLLM.shutdown_started +# --------------------------------------------------------------------------- +def test_shutdown_started_follows_doing_shutdown(): + executor = _NoOpExecutor() + llm = _llm(executor) + assert llm.shutdown_started is False + executor.doing_shutdown = True + assert llm.shutdown_started is True + + +def test_shutdown_started_true_without_executor(): + assert _llm(None).shutdown_started is True + + +def test_shutdown_started_false_for_encode_only(): + assert _llm(None, encode_only=True, encoder_executor=_NoOpExecutor()).shutdown_started is False + + +# --------------------------------------------------------------------------- +# BaseLLM.engine_dead +# --------------------------------------------------------------------------- +def test_engine_dead_false_without_executor(): + assert _llm(None).engine_dead is False + + +def test_engine_dead_follows_the_executor(): + executor = _NoOpExecutor() + llm = _llm(executor) + assert llm.engine_dead is False + executor._set_fatal_error(RuntimeError("worker died")) + assert llm.engine_dead is True + + +# --------------------------------------------------------------------------- +# GenerationExecutor.fatal_error and engine_dead +# --------------------------------------------------------------------------- +def test_executor_fatal_error_property(): + executor = _NoOpExecutor() + assert executor.fatal_error is None + error = RuntimeError("fatal") + executor._set_fatal_error(error) + assert executor.fatal_error is error + + +def test_executor_engine_dead_property(): + executor = _NoOpExecutor() + assert executor.engine_dead is False + executor._set_fatal_error(RuntimeError("fatal")) + assert executor.engine_dead is True + + +@pytest.mark.parametrize( + "engine_dead,fatal_error,expected", + [ + (False, None, False), + (True, None, True), + (False, RuntimeError("fatal"), True), + (True, RuntimeError("fatal"), True), + ], +) +def test_proxy_engine_dead_reads_both_fields(engine_dead, fatal_error, expected): + """_mark_engine_dead can set the sticky flag before the error is recorded.""" + assert _proxy(engine_dead=engine_dead, fatal_error=fatal_error).engine_dead is expected diff --git a/tests/unittest/serve/__init__.py b/tests/unittest/serve/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tests/unittest/serve/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tests/unittest/serve/api/__init__.py b/tests/unittest/serve/api/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tests/unittest/serve/api/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tests/unittest/serve/api/visual_gen/__init__.py b/tests/unittest/serve/api/visual_gen/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tests/unittest/serve/api/visual_gen/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tests/unittest/serve/api/visual_gen/test_private_access_lint_visual_gen.py b/tests/unittest/serve/api/visual_gen/test_private_access_lint_visual_gen.py new file mode 100644 index 000000000000..40f7fc1cb9d4 --- /dev/null +++ b/tests/unittest/serve/api/visual_gen/test_private_access_lint_visual_gen.py @@ -0,0 +1,46 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""VisualGen part of the private-access lint. + +Pull requests that change only VisualGen files are claimed by CBTS Tier 1 and run +only test-db entries whose path contains visual_gen/. This file applies the rule +and baseline of tests/unittest/serve/test_private_access_lint.py to the existing +frontends, including the VisualGen handlers in serve/openai_server.py. Keep the +full frontend scope until those handlers move into their own directory. +""" + +import importlib.util +from pathlib import Path + +import pytest + +pytestmark = pytest.mark.cpu_only + +_MAIN = Path(__file__).resolve().parents[2] / "test_private_access_lint.py" +_spec = importlib.util.spec_from_file_location("_serve_private_access_lint", _MAIN) +lint = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(lint) + + +def test_no_new_cross_package_private_access_in_visual_gen(): + for scope in lint.SCOPES: + assert (lint.REPO_ROOT / scope).is_dir(), f"Frontend scope does not exist: {scope}" + rows = lint.scan(root=lint.REPO_ROOT) + errors = lint._compare(rows, lint.load_baseline()) + assert not errors, "\n".join(errors) + + +def test_visual_gen_check_detects_new_frontend_access(tmp_path, monkeypatch): + for scope in lint.SCOPES: + (tmp_path / scope).mkdir(parents=True) + (tmp_path / "tensorrt_llm/serve/openai_server.py").write_text("engine._private_state\n") + owner = tmp_path / "tensorrt_llm/llmapi/llm.py" + owner.parent.mkdir() + owner.write_text("class LLM:\n _private_state = None\n") + baseline = tmp_path / "baseline.yaml" + baseline.write_text("entries: []\n") + monkeypatch.setattr(lint, "REPO_ROOT", tmp_path) + monkeypatch.setattr(lint, "BASELINE_FILE", baseline) + monkeypatch.setattr(lint, "OVERRIDES", {}) + with pytest.raises(AssertionError, match="New private access to `_private_state`"): + test_no_new_cross_package_private_access_in_visual_gen() diff --git a/tests/unittest/serve/private_access_baseline.yaml b/tests/unittest/serve/private_access_baseline.yaml new file mode 100644 index 000000000000..123c9bdbd182 --- /dev/null +++ b/tests/unittest/serve/private_access_baseline.yaml @@ -0,0 +1,53 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# Cross-package private accesses from tensorrt_llm/serve (without serve/scripts/), +# tensorrt_llm/commands and tensorrt_llm/grpc, checked by test_private_access_lint.py. +# This list only shrinks. When you remove an access, lower `count` (delete the entry +# at 0) and lower MAX_TOTAL in the test. Never add an entry: add a public accessor on +# the class that owns the member. A pure move changes `path` and `scope` only. +# `owner` and `replacement` are for readers; the test uses path, scope, member, count. +version: 1 +entries: +- {path: tensorrt_llm/commands/serve.py, scope: _bind_launcher_uds, member: _executor, count: 1, owner: llmapi/llm.py BaseLLM, replacement: a public accessor for the multi-frontend IPC directory} +- {path: tensorrt_llm/commands/serve.py, scope: _bind_launcher_uds, member: _multi_frontend_ipc_dir, count: 1, owner: executor/proxy.py GenerationExecutorProxy, replacement: a public accessor for the multi-frontend IPC directory} +- {path: tensorrt_llm/commands/serve.py, scope: _run_fleet_worker_impl, member: _config, count: 2, owner: serve/openai_disagg_server.py OpenAIDisaggServer, replacement: OpenAIDisaggServer.address} +- {path: tensorrt_llm/grpc/openengine/capabilities.py, scope: _loaded_model_type, member: _hf_model_config, count: 1, owner: llmapi/llm.py BaseLLM, replacement: BaseLLM.hf_model_config} +- {path: tensorrt_llm/grpc/openengine/capabilities.py, scope: resolve_media_capabilities, member: _registered_model_type, count: 1, owner: inputs input-processor registry, replacement: a public registry accessor for the registered model type} +- {path: tensorrt_llm/grpc/openengine/control.py, scope: OpenEngineControlServicer.GetServerInfo, member: _executor, count: 1, owner: llmapi/llm.py BaseLLM, replacement: "BaseLLM.get_kv_cache_capacity(), BaseLLM.get_openengine_discovery()"} +- {path: tensorrt_llm/grpc/openengine/control.py, scope: OpenEngineControlServicer.GetServerInfo, member: _openengine_discovery, count: 1, owner: llmapi/llm_args.py BaseLlmArgs, replacement: BaseLlmArgs.openengine_discovery} +- {path: tensorrt_llm/grpc/openengine/control.py, scope: OpenEngineControlServicer.__init__, member: _enable_routing_load, count: 1, owner: llmapi/llm_args.py BaseLlmArgs, replacement: BaseLlmArgs.enable_routing_load} +- {path: tensorrt_llm/grpc/openengine/control.py, scope: OpenEngineControlServicer._get_kv_cache_load, member: _executor, count: 1, owner: llmapi/llm.py BaseLLM, replacement: BaseLLM.get_kv_cache_load()} +- {path: tensorrt_llm/grpc/openengine/control.py, scope: OpenEngineControlServicer._requested_sources, member: _executor, count: 1, owner: llmapi/llm.py BaseLLM, replacement: BaseLLM.get_openengine_discovery()} +- {path: tensorrt_llm/grpc/openengine/formatting.py, scope: _token_holdback, member: _stop_word_ids, count: 1, owner: sampling_params.py SamplingParams, replacement: SamplingParams.stop_word_ids} +- {path: tensorrt_llm/grpc/openengine/node.py, scope: initialize_node_discovery, member: _openengine_discovery, count: 2, owner: "llmapi/llm_args.py BaseLlmArgs (read); _torch/pyexecutor/py_executor.py PyExecutor (write)", replacement: "BaseLlmArgs.openengine_discovery; PyExecutor.set_openengine_discovery()"} +- {path: tensorrt_llm/grpc/openengine/node.py, scope: initialize_node_discovery, member: _openengine_node_server, count: 1, owner: _torch/pyexecutor/py_executor.py PyExecutor, replacement: PyExecutor.set_openengine_node_server()} +- {path: tensorrt_llm/grpc/openengine/request_mapping.py, scope: sampling_params_from_request, member: _set_request_provided_fields, count: 1, owner: sampling_params.py SamplingParams, replacement: SamplingParams.set_request_provided_fields()} +- {path: tensorrt_llm/grpc/openengine/server.py, scope: OpenEngineServer.start, member: _executor, count: 1, owner: llmapi/llm.py BaseLLM, replacement: BaseLLM.get_kv_cache_capacity()} +- {path: tensorrt_llm/serve/_frontend_processes.py, scope: _spawn_attached_frontends, member: _executor, count: 1, owner: llmapi/llm.py BaseLLM, replacement: BaseLLM.multi_frontend_attach_info()} +- {path: tensorrt_llm/serve/chat_utils.py, scope: parse_chat_message_content_part, member: _multimodal_server_config, count: 4, owner: inputs/utils.py MultimodalDataTracker, replacement: MultimodalDataTracker.multimodal_server_config} +- {path: tensorrt_llm/serve/openai_protocol.py, scope: _record_sampling_params_request_fields, member: _set_request_provided_fields, count: 1, owner: sampling_params.py SamplingParams, replacement: SamplingParams.set_request_provided_fields()} +- {path: tensorrt_llm/serve/openai_server.py, scope: OpenAIServer._create_chat_response, member: _postprocess_result, count: 1, owner: executor/result.py CompletionOutput, replacement: CompletionOutput.postprocess_result} +- {path: tensorrt_llm/serve/openai_server.py, scope: OpenAIServer._init_embedding_batcher, member: _encoder_executor, count: 1, owner: llmapi/llm.py BaseLLM, replacement: BaseLLM.encode_limits} +- {path: tensorrt_llm/serve/openai_server.py, scope: OpenAIServer._init_llm, member: _hf_model_dir, count: 1, owner: llmapi/llm.py BaseLLM, replacement: BaseLLM.hf_model_dir} +- {path: tensorrt_llm/serve/openai_server.py, scope: OpenAIServer._on_launcher_lost, member: _executor, count: 1, owner: llmapi/llm.py BaseLLM, replacement: public fatal-error recording on the LLM} +- {path: tensorrt_llm/serve/openai_server.py, scope: OpenAIServer._on_launcher_lost, member: _set_fatal_error, count: 1, owner: executor/executor.py GenerationExecutor, replacement: public fatal-error recording on the LLM} +- {path: tensorrt_llm/serve/openai_server.py, scope: OpenAIServer.chat_harmony.create_streaming_generator, member: _postprocess_result, count: 2, owner: executor/result.py CompletionOutput, replacement: CompletionOutput.postprocess_result} +- {path: tensorrt_llm/serve/openai_server.py, scope: OpenAIServer.openai_chat.chat_stream_generator, member: _postprocess_result, count: 2, owner: executor/result.py CompletionOutput, replacement: CompletionOutput.postprocess_result} +- {path: tensorrt_llm/serve/openai_server.py, scope: OpenAIServer.openai_completion.completion_generator, member: _postprocess_result, count: 1, owner: executor/result.py CompletionOutput, replacement: CompletionOutput.postprocess_result} +- {path: tensorrt_llm/serve/openai_server.py, scope: OpenAIServer.openai_completion.completion_response, member: _postprocess_result, count: 1, owner: executor/result.py CompletionOutput, replacement: CompletionOutput.postprocess_result} +- {path: tensorrt_llm/serve/openai_server.py, scope: OpenAIServer.openai_image_generation, member: _save_bytes, count: 1, owner: visual_gen/output.py, replacement: public VisualGen output encoder} +- {path: tensorrt_llm/serve/openai_server.py, scope: OpenAIServer.openai_responses.create_response, member: _postprocess_result, count: 1, owner: executor/result.py CompletionOutput, replacement: CompletionOutput.postprocess_result} +- {path: tensorrt_llm/serve/openai_server.py, scope: OpenAIServer.openai_responses.create_streaming_generator, member: _postprocess_result, count: 1, owner: executor/result.py CompletionOutput, replacement: CompletionOutput.postprocess_result} +- {path: tensorrt_llm/serve/openai_server.py, scope: OpenAIServer.register_routes, member: _executor, count: 1, owner: llmapi/llm.py BaseLLM, replacement: BaseLLM.resource_governor_queue} +- {path: tensorrt_llm/serve/openai_server.py, scope: _record_generator_termination, member: _executor, count: 1, owner: llmapi/llm.py BaseLLM, replacement: BaseLLM.engine_in_worker_process} +- {path: tensorrt_llm/serve/postprocess_handlers.py, scope: chat_harmony_streaming_post_processor, member: _done, count: 1, owner: executor/result.py GenerationResultBase, replacement: GenerationResultBase.finished} +- {path: tensorrt_llm/serve/postprocess_handlers.py, scope: chat_response_post_processor, member: _spec_dec_counters, count: 1, owner: executor/result.py CompletionOutput, replacement: CompletionOutput.spec_dec_counters} +- {path: tensorrt_llm/serve/postprocess_handlers.py, scope: chat_stream_post_processor, member: _done, count: 1, owner: executor/result.py GenerationResultBase, replacement: GenerationResultBase.finished} +- {path: tensorrt_llm/serve/postprocess_handlers.py, scope: chat_stream_post_processor, member: _spec_dec_counters, count: 1, owner: executor/result.py CompletionOutput, replacement: CompletionOutput.spec_dec_counters} +- {path: tensorrt_llm/serve/postprocess_handlers.py, scope: completion_response_post_processor, member: _spec_dec_counters, count: 1, owner: executor/result.py CompletionOutput, replacement: CompletionOutput.spec_dec_counters} +- {path: tensorrt_llm/serve/postprocess_handlers.py, scope: completion_stream_post_processor, member: _done, count: 1, owner: executor/result.py GenerationResultBase, replacement: GenerationResultBase.finished} +- {path: tensorrt_llm/serve/postprocess_handlers.py, scope: completion_stream_post_processor, member: _last_text_len, count: 1, owner: executor/result.py CompletionOutput, replacement: CompletionOutput.text_diff_offset} +- {path: tensorrt_llm/serve/postprocess_handlers.py, scope: completion_stream_post_processor, member: _spec_dec_counters, count: 1, owner: executor/result.py CompletionOutput, replacement: CompletionOutput.spec_dec_counters} +- {path: tensorrt_llm/serve/postprocess_handlers.py, scope: responses_api_streaming_post_processor, member: _done, count: 1, owner: executor/result.py GenerationResultBase, replacement: GenerationResultBase.finished} +- {path: tensorrt_llm/serve/responses_utils.py, scope: ResponsesStreamingProcessor.process_single_output, member: _done, count: 1, owner: executor/result.py GenerationResultBase, replacement: GenerationResultBase.finished} +- {path: tensorrt_llm/serve/router_utils.py, scope: BlockHashMixin._compute_block_hashes, member: _hash_cache_salt, count: 1, owner: runtime/kv_cache_manager_v2 KVCacheEventManager (nanobind), replacement: public hash_cache_salt binding} diff --git a/tests/unittest/serve/test_private_access_lint.py b/tests/unittest/serve/test_private_access_lint.py new file mode 100644 index 000000000000..9bea0b46de81 --- /dev/null +++ b/tests/unittest/serve/test_private_access_lint.py @@ -0,0 +1,410 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Private-access lint for the serving frontends. + +The HTTP server (tensorrt_llm/serve, without serve/scripts), the CLI +(tensorrt_llm/commands) and the gRPC frontends (tensorrt_llm/grpc) may use only +the public API of other TensorRT-LLM packages. This test counts every access to +a private member that belongs to another package and compares the result with +private_access_baseline.yaml, which lists the accesses that existed when the +lint was introduced. The baseline only shrinks. + +Counting rule: one AST node per access. An access is an attribute +`._name` or a getattr/hasattr/setattr/delattr call with the literal +"_name", where the receiver is not `self`, `cls` or `super()` and `_name` has one +leading underscore and is not a dunder. `a._b._c` is two accesses. Names +imported with `from m import _x` are not accesses. Ownership: OVERRIDES first, +then a name index over tensorrt_llm/**/*.py and *.pyi. + +Run `python tests/unittest/serve/test_private_access_lint.py --report` to print +the current inventory. +""" + +import ast +import collections +import sys +from pathlib import Path + +import pytest +import yaml + +pytestmark = pytest.mark.cpu_only + +REPO_ROOT = Path(__file__).resolve().parents[3] +BASELINE_FILE = Path(__file__).resolve().with_name("private_access_baseline.yaml") +SCOPES = ("tensorrt_llm/serve/", "tensorrt_llm/commands/", "tensorrt_llm/grpc/") +EXCLUDED = ("tensorrt_llm/serve/scripts/", "tensorrt_llm/grpc/openengine/_generated/") +OPENENGINE_SCOPE = "tensorrt_llm/grpc/openengine/" + +# Lower this when you remove accesses; never raise it. +MAX_TOTAL = 50 + +# (path, member) -> (verdict, reason). Verdicts: "cross", "same-package", "third-party". +# Needed only where the name index cannot tell: the name is defined both inside and +# outside the scanning package, or a third-party name collides with a TensorRT-LLM name. +OVERRIDES = { + ("tensorrt_llm/serve/_telemetry.py", "_serve"): ("third-party", "uvicorn.Server._serve"), + ("tensorrt_llm/serve/_telemetry.py", "_captured_signals"): ( + "third-party", + "uvicorn.Server._captured_signals", + ), + ("tensorrt_llm/serve/cluster_storage.py", "_loop"): ("third-party", "asyncio.Queue._loop"), + ("tensorrt_llm/serve/render/fingerprint.py", "_mergeable_ranks"): ( + "third-party", + "tiktoken Encoding", + ), + ("tensorrt_llm/serve/render/fingerprint.py", "_special_tokens"): ( + "third-party", + "tiktoken Encoding", + ), + ("tensorrt_llm/serve/render/fingerprint.py", "_pat_str"): ("third-party", "tiktoken Encoding"), + ("tensorrt_llm/serve/router.py", "_lock"): ("same-package", "serve/router.py Router._lock"), + # Generic names of serve's own router objects: the name index classifies them + # correctly today, but a definition of the same name in another package would + # turn them into "unresolved" and fail unrelated pull requests. + ("tensorrt_llm/serve/router.py", "_servers"): ( + "same-package", + "serve/router.py Router._servers", + ), + ("tensorrt_llm/serve/router.py", "_server_info"): ( + "same-package", + "serve/router.py Router._server_info", + ), + ("tensorrt_llm/serve/disagg_coordinator.py", "_server_info"): ( + "same-package", + "serve/router.py Router._server_info", + ), + ("tensorrt_llm/serve/disagg_coordinator.py", "_local"): ( + "same-package", + "serve/router.py CoordinatorDelegatingRouter._local", + ), + ("tensorrt_llm/serve/openai_disagg_service.py", "_local"): ( + "same-package", + "serve/router.py CoordinatorDelegatingRouter._local", + ), + # _DecodedMediaReservation reads its own module's _DecodedMediaBudget; the + # generic names are also defined elsewhere, so the index cannot tell. + ("tensorrt_llm/grpc/openengine/request_mapping.py", "_condition"): ( + "same-package", + "grpc/openengine/request_mapping.py _DecodedMediaBudget._condition", + ), + ("tensorrt_llm/grpc/openengine/request_mapping.py", "_cancelled"): ( + "same-package", + "grpc/openengine/request_mapping.py _DecodedMediaBudget._cancelled", + ), +} + +FIX_HINT = ( + "Frontends use only public LLM, executor, result and SamplingParams APIs. Add a public " + "accessor on the class that owns the member (in the owning package) and call it here. " + "Do not add the access to tests/unittest/serve/private_access_baseline.yaml: the baseline " + "only shrinks. If the member belongs to this package or to a third-party library, add an " + "OVERRIDES entry in tests/unittest/serve/test_private_access_lint.py with the reason. " + "Runtime owners (@NVIDIA/trt-llm-runtime-devs) review both files." +) + + +def _is_private(name): + return name.startswith("_") and not (name.startswith("__") and name.endswith("__")) + + +def _package_of(rel): + parts = rel.split("/") + return parts[1][:-3] if len(parts) == 2 else parts[1] + + +def _is_self(node): + if isinstance(node, ast.Name) and node.id in ("self", "cls"): + return True + return ( + isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "super" + ) + + +class _Candidates(ast.NodeVisitor): + def __init__(self): + self.scope, self.rows = [], [] + + def _qualname(self): + return ".".join(self.scope) or "" + + def visit_FunctionDef(self, node): + self.scope.append(node.name) + self.generic_visit(node) + self.scope.pop() + + visit_AsyncFunctionDef = visit_FunctionDef + visit_ClassDef = visit_FunctionDef + + def visit_Attribute(self, node): + if _is_private(node.attr) and not _is_self(node.value): + self.rows.append((node.lineno, node.attr, ast.unparse(node.value), self._qualname())) + self.generic_visit(node) + + def visit_Call(self, node): + func = node.func + if ( + isinstance(func, ast.Name) + and func.id in ("getattr", "hasattr", "setattr", "delattr") + and len(node.args) >= 2 + and isinstance(node.args[1], ast.Constant) + and isinstance(node.args[1].value, str) + and _is_private(node.args[1].value) + and not _is_self(node.args[0]) + ): + self.rows.append( + (node.lineno, node.args[1].value, ast.unparse(node.args[0]), self._qualname()) + ) + self.generic_visit(node) + + +class _Definitions(ast.NodeVisitor): + """Record which package defines each private member name.""" + + def __init__(self, rel, index): + self.rel, self.index = rel, index + + def _add(self, name): + if _is_private(name): + self.index[name].add(_package_of(self.rel)) + + def visit_ClassDef(self, node): + for stmt in node.body: + if isinstance(stmt, (ast.FunctionDef, ast.AsyncFunctionDef)): + self._add(stmt.name) + elif isinstance(stmt, ast.Assign): + for target in stmt.targets: + if isinstance(target, ast.Name): + self._add(target.id) + elif isinstance(stmt, ast.AnnAssign) and isinstance(stmt.target, ast.Name): + self._add(stmt.target.id) + self.generic_visit(node) + + def _targets(self, targets): + for target in targets: + if isinstance(target, ast.Attribute): + if isinstance(target.value, ast.Name) and target.value.id in ("self", "cls"): + self._add(target.attr) + elif not self.rel.startswith(SCOPES): + # `obj_or_class._x = ...` outside the frontends also defines `_x`. + self._add(target.attr) + elif isinstance(target, (ast.Tuple, ast.List)): + self._targets(target.elts) + + def visit_Assign(self, node): + self._targets(node.targets) + self.generic_visit(node) + + def visit_AnnAssign(self, node): + self._targets([node.target]) + self.generic_visit(node) + + def visit_AugAssign(self, node): + self._targets([node.target]) + self.generic_visit(node) + + +def _scanned_files(root, prefixes): + for prefix in prefixes: + base = root / prefix + if not base.is_dir(): + continue + for path in sorted(base.rglob("*.py")): + rel = path.relative_to(root).as_posix() + if not rel.startswith(EXCLUDED): + yield rel + + +def _name_index(root, names): + index = collections.defaultdict(set) + for path in (root / "tensorrt_llm").rglob("*.py*"): + if path.suffix not in (".py", ".pyi"): + continue + text = path.read_text(encoding="utf-8") + if not any(name in text for name in names): + continue + try: + tree = ast.parse(text) + except SyntaxError: + continue + _Definitions(path.relative_to(root).as_posix(), index).visit(tree) + return index + + +def scan(root=REPO_ROOT, prefixes=SCOPES): + """Return one row per access: (path, line, member, receiver, scope, verdict, reason).""" + candidates = [] + for rel in _scanned_files(root, prefixes): + visitor = _Candidates() + visitor.visit(ast.parse((root / rel).read_text(encoding="utf-8"))) + candidates.extend((rel,) + row for row in visitor.rows) + index = _name_index(root, {member for _, _, member, _, _ in candidates}) + rows = [] + for rel, line, member, receiver, scope in candidates: + package = _package_of(rel) + if (rel, member) in OVERRIDES: + verdict, reason = OVERRIDES[(rel, member)] + else: + owners = index.get(member, set()) + outside = sorted(owner for owner in owners if owner != package) + if package in owners and outside: + verdict, reason = "unresolved", ", ".join(sorted(owners)) + elif package in owners: + verdict, reason = "same-package", package + elif outside: + verdict, reason = "cross", ", ".join(outside) + else: + verdict, reason = "third-party", "not defined in tensorrt_llm" + rows.append((rel, line, member, receiver, scope, verdict, reason)) + return rows + + +def load_baseline(): + data = yaml.safe_load(BASELINE_FILE.read_text()) + baseline = {} + for entry in data["entries"]: + key = (entry["path"], entry["scope"], entry["member"]) + if key in baseline: + raise ValueError(f"Duplicate baseline entry: {key}") + baseline[key] = entry + return baseline + + +def _compare(rows, baseline, prefix=""): + errors = [] + current = collections.defaultdict(list) + for row in rows: + if row[5] == "unresolved": + errors.append( + f"Cannot classify `{row[3]}.{row[2]}` at {row[0]}:{row[1]}: `{row[2]}` is defined " + f"both in this package and elsewhere ({row[6]}). If the receiver is an object from " + f"another package, this is a new cross-package private access. {FIX_HINT} If it " + "is this package's own member, add an OVERRIDES entry with verdict same-package." + ) + elif row[5] == "cross": + current[(row[0], row[4], row[2])].append(row) + for key, found in sorted(current.items()): + allowed = baseline.get(key, {}).get("count", 0) + if len(found) > allowed: + where = ", ".join(f"{r[0]}:{r[1]}" for r in found) + errors.append( + f"New private access to `{key[2]}` (owned by {found[0][6]}) in {key[1]} " + f"at {where}: {len(found)} found, baseline allows {allowed}. {FIX_HINT}" + ) + for key, entry in sorted(baseline.items()): + if not key[0].startswith(prefix): + continue + found = len(current.get(key, [])) + if found < entry["count"]: + errors.append( + f"Stale baseline entry: {key[0]} {key[1]} `{key[2]}` lists {entry['count']} " + f"access(es) but the code has {found}. Lower `count` (delete the entry at 0) in " + f"{BASELINE_FILE.name} and lower MAX_TOTAL in this PR. If the code moved, change " + "the entry's path and scope instead; the total must not grow." + ) + accessed = {(row[0], row[2]) for row in rows} + for path, member in sorted(OVERRIDES): + if path.startswith(prefix) and (path, member) not in accessed: + errors.append( + f"Stale OVERRIDES entry: {path} `{member}` matches no access. If the code moved, " + "change the key's path to the new location in this PR, as for a baseline entry; " + "if the access is gone, delete the entry." + ) + return errors + + +def test_no_new_cross_package_private_access(): + errors = _compare(scan(), load_baseline()) + assert not errors, "\n".join(errors) + + +def test_baseline_only_shrinks(): + total = sum(e["count"] for e in load_baseline().values()) + assert total == MAX_TOTAL, ( + f"{BASELINE_FILE.name} lists {total} accesses but MAX_TOTAL={MAX_TOTAL}. " + "Lower MAX_TOTAL when the total shrinks. " + f"The baseline only shrinks. {FIX_HINT}" + ) + + +@pytest.mark.parametrize( + "source,expected", + [ + ("self._x\ncls._x\nsuper()._x()", []), + ("llm._executor", ["_executor"]), + ("llm._executor._fatal_error", ["_executor", "_fatal_error"]), + ("getattr(self._llm, '_check_health', None)", ["_check_health"]), + ("hasattr(self, '_x')", []), + ("from m import _private\nobj.__dict__", []), + ], +) +def test_counting_rule(source, expected): + visitor = _Candidates() + visitor.visit(ast.parse(source)) + assert sorted(row[1] for row in visitor.rows) == sorted(expected) + + +def test_duplicate_baseline_keys_are_rejected(tmp_path, monkeypatch): + baseline_file = tmp_path / "baseline.yaml" + entry = {"path": "tensorrt_llm/serve/server.py", "scope": "f", "member": "_x", "count": 1} + baseline_file.write_text(yaml.safe_dump({"entries": [entry, {**entry, "count": 2}]})) + monkeypatch.setattr(sys.modules[__name__], "BASELINE_FILE", baseline_file) + with pytest.raises(ValueError, match="Duplicate baseline entry"): + load_baseline() + + +def test_baseline_total_cannot_leave_unused_allowance(tmp_path, monkeypatch): + baseline_file = tmp_path / "baseline.yaml" + baseline_file.write_text("entries: []\n") + monkeypatch.setattr(sys.modules[__name__], "BASELINE_FILE", baseline_file) + monkeypatch.setattr(sys.modules[__name__], "MAX_TOTAL", 1) + with pytest.raises(AssertionError, match="Lower MAX_TOTAL"): + test_baseline_only_shrinks() + + +def test_scanner_ownership_and_comparison(tmp_path, monkeypatch): + frontend = "tensorrt_llm/serve/server.py" + sources = { + frontend: ( + "class Server:\n _ambiguous = None\n" + "def f(obj):\n obj._external\n obj._ambiguous\n" + " obj._third_party\n obj._override\n" + ), + "tensorrt_llm/executor/engine.py": ( + "class Engine:\n _external = None\n _ambiguous = None\n _override = None\n" + ), + } + for rel, source in sources.items(): + path = tmp_path / rel + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(source) + overrides = {(frontend, "_override"): ("same-package", "test receiver")} + monkeypatch.setattr(sys.modules[__name__], "OVERRIDES", overrides) + rows = scan(tmp_path, prefixes=("tensorrt_llm/serve/",)) + assert {row[2]: row[5] for row in rows} == { + "_external": "cross", + "_ambiguous": "unresolved", + "_third_party": "third-party", + "_override": "same-package", + } + errors = _compare(rows, {}) + assert len(errors) == 2 + assert any("New private access to `_external`" in error for error in errors) + assert any("Cannot classify `obj._ambiguous`" in error for error in errors) + + baseline = { + (frontend, "f", "_external"): {"count": 2}, + ("tensorrt_llm/grpc/server.py", "f", "_other"): {"count": 1}, + } + overrides[(frontend, "_removed")] = ("same-package", "removed test receiver") + errors = _compare(rows, baseline, prefix="tensorrt_llm/serve/") + assert len(errors) == 3 + assert any("Stale baseline entry" in error and "`_external`" in error for error in errors) + assert any("Stale OVERRIDES entry" in error and "`_removed`" in error for error in errors) + assert not any("_other" in error for error in errors) + + +if __name__ == "__main__": + if "--report" in sys.argv: + for row in scan(): + print("\t".join(str(x) for x in row)) diff --git a/tests/unittest/serve/test_serve_health.py b/tests/unittest/serve/test_serve_health.py new file mode 100644 index 000000000000..2717df0fa941 --- /dev/null +++ b/tests/unittest/serve/test_serve_health.py @@ -0,0 +1,120 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Production tests of the HTTP /health handler. + +Calls the production OpenAIServer._check_health and OpenAIServer.health on a +stand-in object that carries exactly the attributes they read (the shared +server fixture factory does not exist yet), with a fake generator class per +case instead of a MagicMock, which would answer any member with a truthy mock. +""" + +import asyncio +import types +from unittest import mock + +import pytest + +from tensorrt_llm.serve.openai_server import OpenAIServer + +pytestmark = pytest.mark.cpu_only + + +class _FakeGenerator: + """Generator double with exactly the public health surface /health reads.""" + + def __init__(self, healthy=True, fatal_error=None, shutdown_started=False): + self._healthy = healthy + self.fatal_error = fatal_error + self.shutdown_started = shutdown_started + + def check_health(self): + return self._healthy + + +class _NoHealthGenerator: + """A VisualGen-style generator with no health attributes at all.""" + + +class _DeadBatcher: + """An embedding batcher whose worker thread has exited.""" + + @staticmethod + def is_alive(): + return False + + +def _server(generator, is_visual_gen=False, embedding_batcher=None): + stub = types.SimpleNamespace( + embedding_batcher=embedding_batcher, + generator=generator, + _is_visual_gen=is_visual_gen, + ) + stub._check_health = types.MethodType(OpenAIServer._check_health, stub) + return stub + + +def _health(server): + with ( + mock.patch("tensorrt_llm.serve.openai_server.signal.raise_signal") as raise_signal, + mock.patch( + "tensorrt_llm.serve.openai_server._record_generator_termination" + ) as record_termination, + ): + response = asyncio.run(OpenAIServer.health(server)) + return response, raise_signal, record_termination + + +def test_health_returns_200_when_the_engine_is_healthy(): + response, raise_signal, _ = _health(_server(_FakeGenerator(healthy=True))) + assert response.status_code == 200 + raise_signal.assert_not_called() + + +def test_health_returns_503_without_sigint_when_unhealthy_with_no_fatal_error(): + response, raise_signal, record_termination = _health( + _server(_FakeGenerator(healthy=False, fatal_error=None)) + ) + assert response.status_code == 503 + assert response.body == (b"LLM is unavailable. Please check the server logs for more details.") + raise_signal.assert_not_called() + record_termination.assert_not_called() + + +def test_health_sends_sigint_exactly_once_on_a_fatal_error(): + import signal as signal_module + + generator = _FakeGenerator( + healthy=False, fatal_error=RuntimeError("CUDA OOM"), shutdown_started=False + ) + response, raise_signal, record_termination = _health(_server(generator)) + assert response.status_code == 503 + raise_signal.assert_called_once_with(signal_module.SIGINT) + record_termination.assert_called_once_with(generator) + + +def test_health_sends_no_sigint_while_shutdown_is_already_in_progress(): + response, raise_signal, record_termination = _health( + _server( + _FakeGenerator( + healthy=False, fatal_error=RuntimeError("CUDA OOM"), shutdown_started=True + ) + ) + ) + assert response.status_code == 503 + raise_signal.assert_not_called() + record_termination.assert_not_called() + + +def test_health_answers_200_for_a_generator_with_no_health_attributes(): + """A VisualGen generator has no check_health/fatal_error; /health stays 200.""" + response, raise_signal, _ = _health(_server(_NoHealthGenerator(), is_visual_gen=True)) + assert response.status_code == 200 + raise_signal.assert_not_called() + + +def test_health_returns_503_when_the_embedding_batcher_died(): + response, raise_signal, _ = _health( + _server(_FakeGenerator(healthy=True), embedding_batcher=_DeadBatcher()) + ) + assert response.status_code == 503 + raise_signal.assert_not_called()