From 76a9065befb4086a7e56abd58ba3b3b7249f9ffa Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Sun, 20 Sep 2026 17:35:27 -0700 Subject: [PATCH] [Fix] Raise on undelivered embeddings in `send_with_url`, fix broken tests (#40502) --- .../srt/disaggregation/encoder/server.py | 27 ++++++++-- .../disaggregation/test_epd_disaggregation.py | 40 ++++++++------ .../test_encoder_server_metrics.py | 54 ++++++++++++------- .../test_modelopt_nvfp4_moe_dispatch.py | 3 ++ 4 files changed, 82 insertions(+), 42 deletions(-) diff --git a/python/sglang/srt/disaggregation/encoder/server.py b/python/sglang/srt/disaggregation/encoder/server.py index 150a76580..b4b2aef1e 100644 --- a/python/sglang/srt/disaggregation/encoder/server.py +++ b/python/sglang/srt/disaggregation/encoder/server.py @@ -2313,10 +2313,13 @@ class MMEncoder: start_time = asyncio.get_running_loop().time() timeout = self.send_timeout cond = await _get_receive_condition(req_id) + failure: Optional[str] = None + failure_code = HTTPStatus.BAD_GATEWAY try: while True: if state.release_requested: + # An upstream abort, not a delivery failure. break async with rid_lock: @@ -2345,9 +2348,11 @@ class MMEncoder: break remaining = timeout - (asyncio.get_running_loop().time() - start_time) if remaining <= 0: - logger.error( - f"[{req_id}] Timeout! Sent {len(sent_urls)}/{expected_count}" + failure = ( + f"timed out after {timeout}s with " + f"{len(sent_urls)}/{expected_count} destination(s) initiated" ) + failure_code = HTTPStatus.GATEWAY_TIMEOUT break async with cond: @@ -2363,13 +2368,25 @@ class MMEncoder: tasks_only = [t[0] for t in all_tasks] results = await asyncio.gather(*tasks_only, return_exceptions=True) - # Process results and log errors + failed = [] for i, result in enumerate(results): url = all_tasks[i][1] # Retrieve URL associated with the task - if isinstance(result, Exception): - logger.error(f"Failed to send to {url}: {result}") + # A cancelled send delivered nothing, and CancelledError + # is not an Exception; outer cancellation re-raises out of + # gather rather than landing here. + if isinstance(result, BaseException): + logger.error(f"Failed to send to {url}: {result!r}") + failed.append(url) else: logger.debug(f"Successfully sent to {url}") + if failed and failure is None: + failure = f"delivery failed for {failed}" + + if failure is not None: + raise MMError( + f"[{req_id}] embedding delivery failed: {failure}", + code=failure_code, + ) logger.info(f"All tasks completed for req_id: {req_id}") diff --git a/test/registered/disaggregation/test_epd_disaggregation.py b/test/registered/disaggregation/test_epd_disaggregation.py index ecfb7a075..506db9a46 100644 --- a/test/registered/disaggregation/test_epd_disaggregation.py +++ b/test/registered/disaggregation/test_epd_disaggregation.py @@ -5,6 +5,7 @@ import subprocess import threading import time import unittest +from concurrent.futures import ThreadPoolExecutor import grpc import openai @@ -1522,23 +1523,28 @@ class TestEPDDisaggregationGrpcEncoderOnly(PDDisaggregationServerBase): image_path = os.path.abspath("examples/assets/example_image.png") try: - stub.SchedulerReceiveUrl( - sglang_encoder_pb2.SchedulerReceiveUrlRequest( - req_id=req_id, - receive_url=f"{self.base_host}:{recv_port}", - receive_count=1, - ), - timeout=60, - ) - stub.Encode( - sglang_encoder_pb2.EncodeRequest( - mm_items=[image_path], - req_id=req_id, - num_parts=1, - part_idx=0, - ), - timeout=300, - ) + # A scheduler registers concurrently with Encode, never before it: + # the request state only exists once Encode dispatches. + with ThreadPoolExecutor(max_workers=1) as pool: + registration = pool.submit( + stub.SchedulerReceiveUrl, + sglang_encoder_pb2.SchedulerReceiveUrlRequest( + req_id=req_id, + receive_url=f"{self.base_host}:{recv_port}", + receive_count=1, + ), + timeout=60, + ) + stub.Encode( + sglang_encoder_pb2.EncodeRequest( + mm_items=[image_path], + req_id=req_id, + num_parts=1, + part_idx=0, + ), + timeout=300, + ) + registration.result(timeout=60) poller = zmq.Poller() poller.register(recv_socket, zmq.POLLIN) diff --git a/test/registered/observability/test_encoder_server_metrics.py b/test/registered/observability/test_encoder_server_metrics.py index 5bc7a84b8..4a6bc2679 100644 --- a/test/registered/observability/test_encoder_server_metrics.py +++ b/test/registered/observability/test_encoder_server_metrics.py @@ -2,6 +2,7 @@ import unittest import uuid +from concurrent.futures import ThreadPoolExecutor from typing import Dict, List from urllib.parse import urlparse @@ -61,26 +62,39 @@ class TestEncoderServerMetrics(CustomTestCase): self.assertEqual(health.status_code, 200) req_id = f"metrics-probe-{uuid.uuid4().hex}" - requests.post( - f"{DEFAULT_URL_FOR_TEST}/scheduler_receive_url", - json={ - "req_id": req_id, - "receive_url": f"{base_host}:{recv_port}", - "receive_count": 1, - }, - ) - response = requests.post( - f"{DEFAULT_URL_FOR_TEST}/encode", - json={ - "req_id": req_id, - "modality": "IMAGE", - "mm_items": [f"data:image/png;base64,{MINIMUM_PNG_PICTURE_BASE64}"], - "num_parts": 1, - "part_idx": 0, - "embedding_port": None, - }, - timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - ) + # A scheduler registers concurrently with /encode, never before + # it: the request state only exists once /encode dispatches. + with ThreadPoolExecutor(max_workers=1) as pool: + registration = pool.submit( + requests.post, + f"{DEFAULT_URL_FOR_TEST}/scheduler_receive_url", + json={ + "req_id": req_id, + "receive_url": f"{base_host}:{recv_port}", + "receive_count": 1, + }, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + ) + response = requests.post( + f"{DEFAULT_URL_FOR_TEST}/encode", + json={ + "req_id": req_id, + "modality": "IMAGE", + "mm_items": [ + f"data:image/png;base64,{MINIMUM_PNG_PICTURE_BASE64}" + ], + "num_parts": 1, + "part_idx": 0, + "embedding_port": None, + }, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + ) + registration_response = registration.result( + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH + ) + # A 200 from /encode alone does not prove the embedding reached + # anyone; the registration is the other half of that contract. + self.assertEqual(registration_response.status_code, 200) self.assertEqual(response.status_code, 200) metrics_response = requests.get(f"{DEFAULT_URL_FOR_TEST}/metrics") diff --git a/test/registered/unit/layers/quantization/test_modelopt_nvfp4_moe_dispatch.py b/test/registered/unit/layers/quantization/test_modelopt_nvfp4_moe_dispatch.py index 99c203176..e14477155 100644 --- a/test/registered/unit/layers/quantization/test_modelopt_nvfp4_moe_dispatch.py +++ b/test/registered/unit/layers/quantization/test_modelopt_nvfp4_moe_dispatch.py @@ -83,6 +83,9 @@ def _trtllm_prepared_layer() -> SimpleNamespace: num_local_experts=NUM_EXPERTS, moe_ep_rank=0, intermediate_size_per_partition=INTERMEDIATE, + # FusedMoE.__init__ sets this; apply() reads it to reject the fused + # fallback for MegaMoE experts. + _mega_moe_nvfp4=False, )