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

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