diff --git a/e2e_test/router/test_pd_responses.py b/e2e_test/router/test_pd_responses.py new file mode 100644 index 000000000..3f4677a13 --- /dev/null +++ b/e2e_test/router/test_pd_responses.py @@ -0,0 +1,103 @@ +"""Responses API tests for PD (Prefill-Decode) disaggregated gRPC routing. + +/v1/responses rides the same mode-parameterized gRPC pipeline as chat +completions, so every create below exercises prefill/decode pair selection, +bootstrap injection, and dual dispatch. + +Backends: +- "pd_grpc": gRPC mode (both SGLang and vLLM) + +Requirements: + - SGLang: sgl_kernel package + - vLLM: NIXL or Mooncake KV transfer support + - GPUs: num_prefill + num_decode (default: 2 GPUs for 1+1) + +Usage: + pytest e2e_test/router/test_pd_responses.py -v + + # vLLM + E2E_RUNTIME=vllm pytest e2e_test/router/test_pd_responses.py -v +""" + +from __future__ import annotations + +import logging + +import openai +import pytest +import smg_client + +logger = logging.getLogger(__name__) + + +@pytest.mark.engine("sglang", "vllm") +@pytest.mark.gpu(2) +@pytest.mark.model("meta-llama/Llama-3.1-8B-Instruct") +@pytest.mark.e2e +@pytest.mark.gateway(extra_args=["--history-backend", "memory"]) +@pytest.mark.parametrize("setup_backend", ["pd_grpc"], indirect=True) +@pytest.mark.parametrize("api_client", ["openai", "smg"], indirect=True) +class TestPDResponsesGrpc: + """Responses API tests using PD disaggregation (gRPC mode).""" + + def test_basic_response_creation(self, model, api_client): + """Test basic response creation.""" + resp = api_client.responses.create(model=model, input="What is 2+2?") + + assert resp.id is not None + assert resp.error is None + assert resp.status == "completed" + assert len(resp.output_text) > 0 + assert resp.usage is not None + + def test_streaming_response(self, model, api_client): + """Test streaming response.""" + resp = api_client.responses.create( + model=model, input="Count to 5", stream=True, max_output_tokens=50 + ) + + events = list(resp) + created_events = [e for e in events if e.type == "response.created"] + assert len(created_events) > 0 + + delta_events = [e for e in events if e.type == "response.output_text.delta"] + assert len(delta_events) > 0 + + completed_events = [e for e in events if e.type == "response.completed"] + assert len(completed_events) == 1 + + def test_previous_response_id_chaining(self, model, api_client): + """Test chaining responses using previous_response_id.""" + # First response + resp1 = api_client.responses.create( + model=model, input="My name is Alice and my friend is Bob. Remember it." + ) + assert resp1.error is None + assert resp1.status == "completed" + + # Second response referencing first + resp2 = api_client.responses.create( + model=model, input="What is my name", previous_response_id=resp1.id + ) + assert resp2.error is None + assert resp2.status == "completed" + assert "Alice" in resp2.output_text + + # Third response referencing second + resp3 = api_client.responses.create( + model=model, + input="What is my friend name?", + previous_response_id=resp2.id, + ) + assert resp3.error is None + assert resp3.status == "completed" + assert "Bob" in resp3.output_text + + def test_store_false_not_retrievable(self, model, api_client): + """Test that store=false responses cannot be retrieved.""" + resp = api_client.responses.create(model=model, input="Hello", store=False) + assert resp.id is not None + assert resp.status == "completed" + + with pytest.raises((openai.NotFoundError, smg_client.NotFoundError)): + api_client.responses.retrieve(response_id=resp.id) diff --git a/model_gateway/src/routers/factory.rs b/model_gateway/src/routers/factory.rs index a525ef5c8..ac5554173 100644 --- a/model_gateway/src/routers/factory.rs +++ b/model_gateway/src/routers/factory.rs @@ -324,7 +324,7 @@ mod grpc_router_type_tests { /// Build an `AppContext` for a gRPC `RoutingMode` with the components /// `GrpcRouter::new` needs (parser factories always; an initialized MCP - /// orchestrator for Regular, which is the only mode that reads it). + /// orchestrator for Regular/PD, which build the responses contexts). async fn grpc_ctx(mode: RoutingMode) -> Arc { let config = RouterConfig::builder() .mode(mode) diff --git a/model_gateway/src/routers/grpc/router.rs b/model_gateway/src/routers/grpc/router.rs index d37231cb6..327e5bc38 100644 --- a/model_gateway/src/routers/grpc/router.rs +++ b/model_gateway/src/routers/grpc/router.rs @@ -282,8 +282,8 @@ fn transcription_response(format: TranscriptionResponseFormat, text: String) -> } } -/// `501 NOT_IMPLEMENTED`, returned by Regular-only endpoints when this router is -/// in PD/EPD mode (matching the `RouterTrait` default). +/// `501 NOT_IMPLEMENTED`, returned by endpoints this router's mode doesn't +/// serve (matching the `RouterTrait` default). fn not_implemented(message: &'static str) -> Response { (StatusCode::NOT_IMPLEMENTED, message).into_response() } @@ -293,10 +293,11 @@ fn not_implemented(message: &'static str) -> Response { /// A single `Mode`-parameterized router serving Regular, PrefillDecode, and /// EncodePrefillDecode. `mode` selects the disaggregation params baked into /// every pipeline and drives the per-mode retry-metric labels, `Debug` output, -/// and `router_type`. The Regular-only members (`harmony_pipeline`, -/// `embedding_pipeline`, `classify_pipeline`, `responses_context`, -/// `harmony_responses_context`) are `Some` only in `Mode::Regular`; PD/EPD leave -/// them `None` and 501 the corresponding endpoints. +/// and `router_type`. Optional members are `None` when the mode doesn't serve +/// them and the corresponding endpoints 501: `embedding_pipeline` and +/// `classify_pipeline` are Regular-only; `harmony_pipeline`, +/// `responses_context`, and `harmony_responses_context` exist in every mode +/// but EPD. #[derive(Clone)] pub struct GrpcRouter { worker_registry: Arc, @@ -314,9 +315,9 @@ pub struct GrpcRouter { } impl GrpcRouter { - /// Only `Mode::Regular` builds the Harmony, embedding, classify, and - /// responses members and requires the MCP orchestrator; PD/EPD leave them - /// `None` and 501 those endpoints. + /// Regular and PD build the Harmony pipeline and responses contexts and + /// require the MCP orchestrator; EPD leaves them `None` and 501s those + /// endpoints. Embedding/classify pipelines are Regular-only. pub fn new(ctx: &Arc, mode: Mode) -> Result { // Get tokenizer registry (no longer requires pre-loaded tokenizer) let tokenizer_registry = ctx.tokenizer_registry.clone(); @@ -374,14 +375,16 @@ impl GrpcRouter { let completion_pipeline = RequestPipeline::build(Endpoint::Completion, mode, &pair_deps) .ok_or_else(|| format!("gRPC router: no completion pipeline for mode {mode:?}"))?; - // Regular-only pipelines; `None` in PD/EPD (which 501 these endpoints). + // `None` when the (endpoint, mode) combo is unsupported; those endpoints 501. let harmony_pipeline = RequestPipeline::build(Endpoint::Harmony, mode, &configured_deps); let embedding_pipeline = RequestPipeline::build(Endpoint::Embeddings, mode, &pair_deps); let classify_pipeline = RequestPipeline::build(Endpoint::Classify, mode, &pair_deps); - // Responses contexts are Regular-only and are the sole consumer of the MCP - // orchestrator; PD/EPD skip both (they don't serve /v1/responses). - let (responses_context, harmony_responses_context) = if mode == Mode::Regular { + // Responses contexts are the sole consumer of the MCP orchestrator; EPD + // builds neither (it doesn't serve /v1/responses). + let (responses_context, harmony_responses_context) = if mode == Mode::EncodePrefillDecode { + (None, None) + } else { let mcp_orchestrator = ctx .mcp_orchestrator .get() @@ -410,11 +413,9 @@ impl GrpcRouter { .as_ref() .map(&create_responses_context) .ok_or_else(|| { - "gRPC router: regular mode must build a harmony pipeline".to_string() + format!("gRPC router: mode {mode:?} must build a harmony pipeline") })?; (Some(responses_context), Some(harmony_responses_context)) - } else { - (None, None) }; Ok(GrpcRouter { @@ -482,8 +483,8 @@ impl GrpcRouter { return *response; } - // Harmony routing is Regular-only: PD/EPD have no Harmony pipeline, so all - // chat requests use the single chat/generate pipeline. + // EPD has no Harmony pipeline, so its chat requests all use the single + // chat/generate pipeline. let is_harmony = self.harmony_pipeline.is_some() && HarmonyDetector::is_harmony_model_in_registry(&self.worker_registry, &body.model); @@ -909,8 +910,7 @@ impl RouterTrait for GrpcRouter { audio: AudioFile, model_id: &str, ) -> Response { - // Routes through the regular chat pipeline (`execute_chat_for_responses`), - // which PD/EPD do not serve. + // Qwen3-ASR transcription is Regular-only. if self.mode != Mode::Regular { return not_implemented("Audio transcriptions not implemented"); } @@ -1194,7 +1194,7 @@ mod tests { } #[cfg(test)] -mod pd_retry_tests { +mod pd_tests { use std::sync::{Arc, OnceLock}; use llm_tokenizer::registry::TokenizerRegistry; @@ -1202,25 +1202,42 @@ mod pd_retry_tests { use smg_data_connector::{ MemoryConversationItemStorage, MemoryConversationStorage, MemoryResponseStorage, }; + use smg_mcp::{McpConfig, McpOrchestrator}; use tool_parser::ParserFactory as ToolParserFactory; use super::*; use crate::{ config::{PolicyConfig, RouterConfig, RoutingMode}, policies::PolicyRegistry, + tenant::TenantKey, worker::WorkerRegistry, }; - /// Minimal `AppContext` for constructing a gRPC PD router. PD/EPD don't read - /// the MCP orchestrator, so an empty `OnceLock` suffices. - fn pd_ctx() -> Arc { + fn pd_routing_mode() -> RoutingMode { + RoutingMode::PrefillDecode { + prefill_urls: vec![], + decode_urls: vec![], + prefill_policy: None, + decode_policy: None, + } + } + + fn epd_routing_mode() -> RoutingMode { + RoutingMode::EncodePrefillDecode { + encode_urls: vec![], + prefill_urls: vec![], + decode_urls: vec![], + encode_policy: None, + prefill_policy: None, + decode_policy: None, + } + } + + /// Minimal `AppContext` for constructing a disaggregated gRPC router. PD + /// serves /v1/responses, so the MCP orchestrator must be initialized. + async fn grpc_ctx(mode: RoutingMode) -> Arc { let config = RouterConfig::builder() - .mode(RoutingMode::PrefillDecode { - prefill_urls: vec![], - decode_urls: vec![], - prefill_policy: None, - decode_policy: None, - }) + .mode(mode) .grpc_connection() .policy(PolicyConfig::Random) .host("127.0.0.1") @@ -1233,6 +1250,15 @@ mod pd_retry_tests { .queue_timeout_secs(60) .build_unchecked(); + let mcp_orchestrator = Arc::new(OnceLock::new()); + mcp_orchestrator + .set(Arc::new( + McpOrchestrator::new(McpConfig::default()) + .await + .expect("mcp orchestrator"), + )) + .ok(); + Arc::new( AppContext::builder() .router_config(config.clone()) @@ -1247,17 +1273,24 @@ mod pd_retry_tests { .conversation_item_storage(Arc::new(MemoryConversationItemStorage::new())) .worker_job_queue(Arc::new(OnceLock::new())) .workflow_engines(Arc::new(OnceLock::new())) - .mcp_orchestrator(Arc::new(OnceLock::new())) + .mcp_orchestrator(mcp_orchestrator) .build() .expect("app context"), ) } + fn responses_request(model: &str) -> ResponsesRequest { + serde_json::from_value(json!({"model": model, "input": "hi"})).expect("responses request") + } + /// PD-mode completion must honor a per-model retry override, not the router /// default. - #[test] - fn pd_completion_honors_per_model_retry_override() { - let ctx = pd_ctx(); + // Multi-thread runtime: router construction eagerly builds the Harmony + // pipeline, whose one-shot encoding load uses `block_in_place` (which + // panics on a current-thread runtime). + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn pd_completion_honors_per_model_retry_override() { + let ctx = grpc_ctx(pd_routing_mode()).await; let router = GrpcRouter::new(&ctx, Mode::PrefillDecode).expect("pd router"); // Router default differs from the override so the assertion is meaningful. @@ -1279,4 +1312,39 @@ mod pd_retry_tests { let fallback = router.resolve_retry_config("model-without-override"); assert_eq!(fallback.max_retries, default_retries); } + + /// PD serves /v1/responses: an unknown model is rejected per-request (404), + /// not gated behind a blanket 501, and cancel reaches storage. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn pd_router_serves_responses_and_cancel() { + let ctx = grpc_ctx(pd_routing_mode()).await; + let router = GrpcRouter::new(&ctx, Mode::PrefillDecode).expect("pd router"); + let tenant_meta = TenantRequestMeta::new(TenantKey::new("test-tenant")); + + let request = responses_request("missing-model"); + let response = router + .route_responses(None, &tenant_meta, &request, "missing-model") + .await; + assert_eq!(response.status(), StatusCode::NOT_FOUND); + + let cancel = router.cancel_response(None, "resp_missing").await; + assert_eq!(cancel.status(), StatusCode::NOT_FOUND); + } + + /// EPD still 501s /v1/responses and cancel. + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn epd_router_501s_responses_and_cancel() { + let ctx = grpc_ctx(epd_routing_mode()).await; + let router = GrpcRouter::new(&ctx, Mode::EncodePrefillDecode).expect("epd router"); + let tenant_meta = TenantRequestMeta::new(TenantKey::new("test-tenant")); + + let request = responses_request("missing-model"); + let response = router + .route_responses(None, &tenant_meta, &request, "missing-model") + .await; + assert_eq!(response.status(), StatusCode::NOT_IMPLEMENTED); + + let cancel = router.cancel_response(None, "resp_missing").await; + assert_eq!(cancel.status(), StatusCode::NOT_IMPLEMENTED); + } }