fix(inference): AUT-2539 reply to a KV handoff import that fails before admission

megatron/core/inference/disaggregation/inference_state_handoff.py

@@ -1194,6 +1194,13 @@ def _admit_pending_kv_imports(self) -> int:
                 if not pending.future.done():
                     pending.future.set_exception(exc)
                 logging.exception("DISAGG_DECODE_PULL_FAILED request_id=%d", pending.request_id)
+                if pending.request_id not in self.requests:
+                    # The import never reached add_request(), so no request record can
+                    # carry an ENGINE_REPLY and the coordinator would keep its routing
+                    # entry forever. Every MP rank applies the same broadcast completion
+                    # decision to mirrored request state, so all ranks register the same
+                    # failed placeholder and only the MP coordinator sends the reply.
+                    self._fail_submission(pending.request_id, pending.sampling_params, exc)
                 if failed:
                     continue
                 remaining.extend(self._pending_kv_imports)

megatron/core/inference/engines/dynamic_engine.py

@@ -1827,7 +1827,8 @@ def _fail_submission(
         produces a client-visible failure reply.
 
         Called from the SUBMIT_REQUEST handler when image preprocessing or
-        add_request raises. Registering a placeholder record with
+        add_request raises, and when a KV handoff import fails before its
+        request is admitted. Registering a placeholder record with
         Status.FAILED lets ``_handle_failed_request`` publish the ENGINE_REPLY
         without leaving the client hanging or killing the engine loop.
         """

tests/unit_tests/inference/test_disagg_handoff_lifecycle.py

@@ -156,6 +156,7 @@ def __init__(self, loop, *, hybrid=False, available=0, extra_eos_token_id_set=()
         self.requests = {}
         self.blocks_to_bind = {}
         self.precomputed_hashes = {}
+        self.failed_submissions = []
         self.partial_admissions = set()
         self.track_generated_token_events = False
         self.use_coordinator = False
@@ -180,6 +181,9 @@ def add_request(self, request_id, prompt, sampling_params, precomputed_block_has
     def get_request(self, request_id):
         return self.requests[request_id]
 
+    def _fail_submission(self, request_id, sampling_params, exc):
+        self.failed_submissions.append((request_id, sampling_params, exc))
+
     def _terminating_token_ids(self, request):
         termination_id = request.sampling_params.termination_id
         explicit = frozenset(request.sampling_params.stop_token_ids or ())
@@ -663,6 +667,8 @@ def test_peer_failure_quarantines_an_unfinished_local_transfer(handoff_loop):
         "Quarantining request %d cache storage after an incomplete handoff", 4
     )
     log_exception.assert_called_once_with("DISAGG_DECODE_PULL_FAILED request_id=%d", 4)
+    # The import never reached add_request(), so the failure itself must reply.
+    assert engine.failed_submissions == [(4, pending.sampling_params, pending.future.exception())]
 
 
 def test_completed_handoff_keeps_notification_while_decode_is_full(handoff_loop):
@@ -1075,6 +1081,16 @@ def test_handoff_submission_failure_is_reported_to_model_parallel_coordinator(ha
     engine._handoff_completion_tracker.report.assert_called_once_with(8, True)
 
     engine._record_handoff_completion_notification(8, failed=True)
+    # Run the engine's real failure-reply path; only the coordinator socket is faked.
+    engine.rank = 1
+    engine.use_coordinator = True
+    engine.is_mp_coordinator = True
+    engine.failed_request_ids = []
+    engine.context.remove_vlm_request_data = mock.Mock()
+    engine._complete_request = DynamicInferenceEngine._complete_request
+    engine._handle_failed_request = DynamicInferenceEngine._handle_failed_request.__get__(engine)
+    engine._fail_submission = DynamicInferenceEngine._fail_submission.__get__(engine)
+    engine._send_requests_to_coordinator = mock.Mock()
     with mock.patch(
         "megatron.core.inference.disaggregation.inference_state_handoff.logging.exception"
     ) as log_exception:
@@ -1084,6 +1100,12 @@ def test_handoff_submission_failure_is_reported_to_model_parallel_coordinator(ha
     assert engine.context.kv_block_allocator.releases == [[10, 11, 12]]
     assert not engine._pending_kv_imports
     log_exception.assert_called_once_with("DISAGG_DECODE_PULL_FAILED request_id=%d", 8)
+    # The import never reached add_request(), so the failure itself replies exactly once.
+    assert engine.failed_request_ids == [8]
+    engine._send_requests_to_coordinator.assert_called_once()
+    (replies,) = engine._send_requests_to_coordinator.call_args.args
+    assert [(reply.request_id, reply.status) for reply in replies] == [(8, Status.FAILED)]
+    assert engine.requests[8].future.result() is replies[0]
 
 
 def test_nixl_handoff_trims_pipeline_stage_block_lists(handoff_loop):