[Fix] Raise on undelivered embeddings in send_with_url, fix broken tests (#40502)

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(
+                    failure = (
-                        f"[{req_id}] Timeout! Sent {len(sent_urls)}/{expected_count}"
+                        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):
+                    # A cancelled send delivered nothing, and CancelledError
-                        logger.error(f"Failed to send to {url}: {result}")
+                    # 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}")
 

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,7 +1523,11 @@ class TestEPDDisaggregationGrpcEncoderOnly(PDDisaggregationServerBase):
         image_path = os.path.abspath("examples/assets/example_image.png")
 
         try:
-            stub.SchedulerReceiveUrl(
+            # 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}",
@@ -1539,6 +1544,7 @@ class TestEPDDisaggregationGrpcEncoderOnly(PDDisaggregationServerBase):
                     ),
                     timeout=300,
                 )
+                registration.result(timeout=60)
 
             poller = zmq.Poller()
             poller.register(recv_socket, zmq.POLLIN)

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(
+            # 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}"],
+                        "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")

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,
     )