[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(
-                        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}")
 

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)

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

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