From f20a813baad9ba208aa6facddf87bbdf39f67bbf Mon Sep 17 00:00:00 2001 From: tanmayv25 Date: Wed, 4 Nov 2026 14:44:24 -0800 Subject: [PATCH] test(sidecar): cover clean TensorRT-LLM liveness shutdown Signed-off-by: tanmayv25 --- .../tests/sidecar_mocker_integration.rs | 39 +++++++++++++++ lib/sidecar/testkit/tests/support/trtllm.rs | 48 ++++++++++++++++++- 2 files changed, 85 insertions(+), 2 deletions(-) diff --git a/lib/sidecar/testkit/tests/sidecar_mocker_integration.rs b/lib/sidecar/testkit/tests/sidecar_mocker_integration.rs index 00e09958f0c5..80edb0e45717 100644 --- a/lib/sidecar/testkit/tests/sidecar_mocker_integration.rs +++ b/lib/sidecar/testkit/tests/sidecar_mocker_integration.rs @@ -934,6 +934,45 @@ async fn trtllm_engine_lost_on_liveness_not_serving() { finish(&mut fixture, &engine).await; } +#[tokio::test] +async fn trtllm_clean_liveness_end_preserves_in_flight_generation() { + let control = Controller::default(); + let mut fixture = + support::trtllm::Fixture::start(control.clone(), FixtureConfig::default()).await; + let engine = fixture.engine().await; + engine.start(0).await.unwrap(); + + let ctx = mock_context(); + let handle = control.request(ctx.id(), after_token_responses(1, StreamAction::Continue)); + let mut stream = engine + .generate( + request("mocker-model", vec![11, 22, 33], 3), + GenerateContext::new(ctx, None), + ) + .await + .unwrap(); + let mut outputs = checkpoint_outputs(&mut stream, &handle).await; + fixture.end_liveness_watch().await; + let error = bounded("engine lost on clean liveness EOF", engine.engine_lost()).await; + assert_eq!( + error.error_type(), + ErrorType::Backend(BackendError::EngineShutdown) + ); + assert!( + error.to_string().contains("liveness stream ended"), + "{error}" + ); + + // Native graceful stop ends Watch before draining its active Generate RPCs. + assert!(poll!(stream.next()).is_pending()); + handle.release(); + outputs.extend(bounded("in-flight generation drains", stream.collect::()).await); + terminal(outputs, &handle.tokens(), 3, FinishReason::Length); + bounded("drained RPC released", handle.wait(Event::Dropped)).await; + fixture.scheduler_idle().await; + finish(&mut fixture, &engine).await; +} + fn http_request() -> PreprocessedRequest { let mut req = request("mocker-model", vec![11, 22, 33], 2); req.extra_args = Some(json!({"sglang_tito": {"sampling_params": {"max_new_tokens": 2}}})); diff --git a/lib/sidecar/testkit/tests/support/trtllm.rs b/lib/sidecar/testkit/tests/support/trtllm.rs index 5420166f3427..cb5a4285f24a 100644 --- a/lib/sidecar/testkit/tests/support/trtllm.rs +++ b/lib/sidecar/testkit/tests/support/trtllm.rs @@ -9,6 +9,7 @@ use std::sync::{Arc, Mutex}; use dynamo_backend_common::{BackendError, DisaggregationMode, PreprocessedRequest}; use dynamo_llm::model_card::ModelDeploymentCard; use dynamo_mocker::common::protocols::EngineType; +use dynamo_sidecar_testkit::bounded; use dynamo_sidecar_testkit::control::{Controller, Protocol, RequestHandle}; use dynamo_sidecar_testkit::fixtures::Outputs; use dynamo_sidecar_testkit::server::TestServer; @@ -35,6 +36,7 @@ pub struct Fixture { scripted: Arc>>>, pub discovery: Arc, health: tonic_health::server::HealthReporter, + liveness_started: Arc, } #[derive(Default)] @@ -49,6 +51,11 @@ pub struct Discovery { } impl Fixture { + pub async fn end_liveness_watch(&mut self) { + bounded("liveness Watch opened", self.liveness_started.notified()).await; + self.health.clear_service_status("liveness").await; + } + /// Sets what `Watch` reports for TensorRT-LLM's `liveness` health. pub async fn set_liveness_serving(&self, is_serving: bool) { let status = if is_serving { @@ -134,13 +141,20 @@ impl SidecarFixture for Fixture { inner: service.clone(), discovery: discovery.clone(), }; - let (health, health_service) = tonic_health::server::health_reporter(); + let health = tonic_health::server::HealthReporter::new(); + let liveness_started = Arc::new(tokio::sync::Notify::new()); + let health_service = ControlledHealth { + inner: tonic_health::server::HealthService::from_health_reporter(health.clone()), + liveness_started: Arc::clone(&liveness_started), + }; health .set_service_status("liveness", tonic_health::ServingStatus::Serving) .await; let server = TestServer::start(move |listener, shutdown| async move { tonic::transport::Server::builder() - .add_service(health_service) + .add_service(tonic_health::pb::health_server::HealthServer::new( + health_service, + )) .add_service(InferenceServer::new(controlled)) .add_service(ControlServer::new(control_service)) .serve_with_incoming_shutdown(TcpListenerStream::new(listener), async { @@ -158,6 +172,7 @@ impl SidecarFixture for Fixture { scripted, discovery, health, + liveness_started, } } @@ -189,6 +204,35 @@ impl SidecarFixture for Fixture { } } +struct ControlledHealth { + inner: tonic_health::server::HealthService, + liveness_started: Arc, +} + +#[tonic::async_trait] +impl tonic_health::pb::health_server::Health for ControlledHealth { + type WatchStream = tonic_health::server::WatchStream; + + async fn check( + &self, + request: Request, + ) -> Result, Status> { + self.inner.check(request).await + } + + async fn watch( + &self, + request: Request, + ) -> Result, Status> { + let liveness = request.get_ref().service == "liveness"; + let response = self.inner.watch(request).await?; + if liveness { + self.liveness_started.notify_one(); + } + Ok(response) + } +} + struct ControlledControl { inner: TrtllmMockerService, discovery: Arc,