[Bugfix][NIXL] Skip heteroblocksize post-processing for failed recv requests in get_finished

When a handshake fails, the remote engine is never registered in kv_topo.remote_block_size (that only happens on successful add_remote_agent). However, get_finished() was calling block_size_ratio_from_engine_id() unconditionally for all requests in done_recving, including those that failed during handshake. This caused a KeyError on the remote engine ID, breaking test_handshake_failure_returns_finished. The bug was introduced by PR #26759 (heterogeneous block_size support), which added the block_size_ratio post-processing step without guarding against requests that failed before remote engine registration. Fix: save the set of failed requests before clearing _failed_recv_reqs, then skip the host-buffer sync and block_size_ratio post-processing for those requests since no data was ever transferred. Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com> Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
2026-03-20 21:45:08 +00:00
1 changed files with 8 additions and 0 deletions
@@ -2034,6 +2034,7 @@ class NixlConnectorWorker:
        done_recving = self._pop_done_transfers(self._recving_transfers)

        # add requests that skipped transfer to done_recving
+        failed_recv_reqs = set(self._failed_recv_reqs)
        done_recving.update(self._failed_recv_reqs)
        self._failed_recv_reqs.clear()

@@ -2052,6 +2053,13 @@ class NixlConnectorWorker:
            meta = self._recving_metadata.pop(req_id, None)
            assert meta is not None, f"{req_id} not found in recving_metadata list"
            assert meta.remote is not None
+
+            # Skip post-processing for failed requests: handshake failure means
+            # no data was transferred and the remote engine was never registered
+            # in kv_topo, so block_size_ratio_from_engine_id would KeyError.
+            if req_id in failed_recv_reqs:
+                continue
+
            if self.use_host_buffer:
                self.sync_recved_kv_to_device(req_id, meta)