@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 91f7d9d0b6

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 15, 04:45 UTC

[P/D] Asynchronously do _nixl_handshake (#19836)

Signed-off-by: Linkun Chen <github@lkchen.net> Signed-off-by: Nick Hill <nhill@redhat.com> Co-authored-by: Nick Hill <nhill@redhat.com>

lkchen committed Jun 24, 2025 at 12:46 UTC 91f7d9d0b698a3488d62717246abb1d1e2354a68

2 files changed +264 -96

tests/v1/kv_connector/unit/test_nixl_connector.py

+168 -58

@@ -7,13 +7,6 @@ from collections import defaultdict

7 from typing import Optional

8 from unittest.mock import patch

10 -import pytest

11 -

12 -try:

13 - from nixl._api import nixl_agent as NixlWrapper

14 -except ImportError:

15 - NixlWrapper = None

16 -

10 from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import (

11 KVConnectorRole, NixlAgentMetadata, NixlConnector, NixlConnectorMetadata,

12 NixlConnectorWorker)

@@ -92,7 +85,8 @@ def test_prompt_less_than_block_size():

85 class FakeNixlWrapper:

86 """Mock implementation of NixlWrapper for testing.

95 - We don't inherit from NixlWrapper because NixlWrapper could be None.

88 + We don't inherit from nixl._api.nixl_agent because nixl may not be

89 + installed.

90 """

92 AGENT_METADATA = b"fake_agent_metadata"

@@ -167,7 +161,7 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):

161 super().__init__(*args, **kwargs)

162 self._hand_shake_latency = hand_shake_latency

163

170 - def _nixl_handshake(self, host: str, port: int):

164 + def _nixl_handshake(self, host: str, port: int) -> dict[int, str]:

165 # Mimic slow _nixl_handshake, as well as bypass zmq communication.

166 time.sleep(self._hand_shake_latency)

167 # These should've been done in register_kv_caches(), called by

@@ -177,7 +171,7 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):

171 self.num_blocks = 1

172 self.dst_num_blocks[self.engine_id] = self.num_blocks

173

180 - self.add_remote_agent(

174 + remote_agent_name = self.add_remote_agent(

175 NixlAgentMetadata(

176 engine_id=self.REMOTE_ENGINE_ID,

177 agent_metadata=FakeNixlWrapper.AGENT_METADATA,

@@ -187,40 +181,101 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):

181 block_len=self.block_len,

182 attn_backend_name=self.backend_name,

183 ))

190 -

191 -

192 -@pytest.mark.skipif(NixlWrapper is None, reason="nixl not installed")

193 -@patch(

194 - "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper",

195 - FakeNixlWrapper)

196 -def test_multi_xfer_one_engine(

197 - # dist_init is a fixture that initializes the distributed environment.

198 - dist_init):

199 - """Test case where multiple xfers are initiated to the same engine.

200 -

201 - This test triggers the connector to load remote KV for the same

202 - `request_id`. The transfer is not done immediately due to

203 - `set_cycles_before_xfer_done`, so there is a state where there are multiple

204 - transfer states for the same `request_id`, and `get_finished` should handle

205 - it correctly (wait for all transfers to be done).

206 - """

207 - vllm_config = create_vllm_config()

208 -

209 - request_id = "req_id"

210 -

211 - # Test worker role in decode server.

212 - connector = NixlConnector(vllm_config, KVConnectorRole.WORKER)

213 - connector.connector_worker = FakeNixlConnectorWorker(vllm_config,

214 - connector.engine_id,

215 - hand_shake_latency=0)

216 - assert isinstance(connector.connector_worker.nixl_wrapper, FakeNixlWrapper)

217 - connector.connector_worker.nixl_wrapper.set_cycles_before_xfer_done(3)

218 - for i in range(4):

184 + return {0: remote_agent_name}

185 +

186 +

187 +class TestNixlHandshake:

188 +

189 + @patch(

190 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper",

191 + FakeNixlWrapper)

192 + def test_multi_xfer_one_engine(

193 + self,

194 + # dist_init is a fixture that initializes the distributed environment.

195 + dist_init):

196 + """Test case where multiple xfers are initiated to the same engine.

197 +

198 + This test triggers the connector to load remote KV for the same

199 + `request_id`. The transfer is not done immediately due to

200 + `set_cycles_before_xfer_done`, so there is a state where there are

201 + multiple transfer states for the same `request_id`, and `get_finished`

202 + should handle it correctly (wait for all transfers to be done).

203 + """

204 + vllm_config = create_vllm_config()

205 +

206 + request_id = "req_id"

207 +

208 + # Test worker role in decode server.

209 + connector = NixlConnector(vllm_config, KVConnectorRole.WORKER)

210 + connector.connector_worker = FakeNixlConnectorWorker(

211 + vllm_config, connector.engine_id, hand_shake_latency=0)

212 + assert isinstance(connector.connector_worker.nixl_wrapper,

213 + FakeNixlWrapper)

214 + connector.connector_worker.nixl_wrapper.set_cycles_before_xfer_done(3)

215 + num_xfers = 4

216 + while True:

217 + # For the same request_id, initiate multiple xfers across different

218 + # round of `execute_model` calls.

219 + metadata = NixlConnectorMetadata()

220 + if num_xfers > 0:

221 + num_xfers -= 1

222 + metadata.add_new_req(

223 + request_id=request_id,

224 + local_block_ids=[

225 + num_xfers + 1, num_xfers + 2, num_xfers + 3

226 + ],

227 + kv_transfer_params={

228 + "remote_block_ids":

229 + [num_xfers + 4, num_xfers + 5, num_xfers + 6],

230 + "remote_engine_id":

231 + FakeNixlConnectorWorker.REMOTE_ENGINE_ID,

232 + "remote_host":

233 + "localhost",

234 + "remote_port":

235 + 1234,

236 + })

237 + connector.bind_connector_metadata(metadata)

238 +

239 + # Mimic maybe_setup_kv_connector in gpu_model_runner.

240 + dummy_ctx = ForwardContext(

241 + no_compile_layers={},

242 + attn_metadata={},

243 + virtual_engine=0,

244 + )

245 + _before_load = time.perf_counter()

246 + connector.start_load_kv(dummy_ctx)

247 + _after_load = time.perf_counter()

248 + assert _after_load - _before_load < 0.1, "start_load_kv took " \

249 + f"{_after_load - _before_load} seconds"

250 +

251 + # Mimic get_finished_kv_transfers in gpu_model_runner.

252 + _, done_recving = connector.get_finished(finished_req_ids=set())

253 + if len(done_recving) > 0:

254 + assert request_id in done_recving

255 + break

256 +

257 + connector.clear_connector_metadata()

258 +

259 + @patch(

260 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper",

261 + FakeNixlWrapper)

262 + def test_async_load_kv(

263 + self,

264 + # dist_init is a fixture that initializes the distributed environment.

265 + dist_init):

266 + """Test that NixlConnector's start_load_kv should be non-blocking."""

267 +

268 + vllm_config = create_vllm_config()

269 +

270 + # Test worker role in decode server.

271 + connector = NixlConnector(vllm_config, KVConnectorRole.WORKER)

272 + connector.connector_worker = FakeNixlConnectorWorker(

273 + vllm_config, connector.engine_id)

274 metadata = NixlConnectorMetadata()

220 - metadata.add_new_req(request_id=request_id,

221 - local_block_ids=[i + 1, i + 2, i + 3],

275 + metadata.add_new_req(request_id="id",

276 + local_block_ids=[1, 2, 3],

277 kv_transfer_params={

223 - "remote_block_ids": [i + 4, i + 5, i + 6],

278 + "remote_block_ids": [4, 5, 6],

279 "remote_engine_id":

280 FakeNixlConnectorWorker.REMOTE_ENGINE_ID,

281 "remote_host": "localhost",

@@ -228,19 +283,74 @@ def test_multi_xfer_one_engine(

283 })

284 connector.bind_connector_metadata(metadata)

285

231 - dummy_ctx = ForwardContext(

232 - no_compile_layers={},

233 - attn_metadata={},

234 - virtual_engine=0,

235 - )

236 - _before_load = time.perf_counter()

237 - connector.start_load_kv(dummy_ctx)

238 - _after_load = time.perf_counter()

239 - assert _after_load - _before_load < 0.1, "start_load_kv took " \

240 - f"{_after_load - _before_load} seconds"

241 -

242 - while True:

243 - _, done_recving = connector.get_finished(finished_req_ids=set())

244 - if len(done_recving) > 0:

245 - assert request_id in done_recving

246 - break

286 + timeout = 2.5

287 + start = time.perf_counter()

288 + while time.perf_counter() - start < timeout:

289 + dummy_ctx = ForwardContext(

290 + no_compile_layers={},

291 + attn_metadata={},

292 + virtual_engine=0,

293 + )

294 + _before_load = time.perf_counter()

295 + connector.start_load_kv(dummy_ctx)

296 + _after_load = time.perf_counter()

297 + assert _after_load - _before_load < 0.1, "start_load_kv took " \

298 + f"{_after_load - _before_load} seconds"

299 + time.sleep(0.5) # backoff for the async handshake to complete.

300 + connector.bind_connector_metadata(NixlConnectorMetadata())

301 + _, done_recving = connector.get_finished(finished_req_ids=set())

302 + if len(done_recving) > 0:

303 + return

304 + raise TimeoutError("Took too long to complete async handshake.")

305 +

306 + @patch(

307 + "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper",

308 + FakeNixlWrapper)

309 + def test_concurrent_load_kv(

310 + self,

311 + # dist_init is a fixture that initializes the distributed environment.

312 + dist_init):

313 + """Test that multiple start_load_kv calls should occur concurrently."""

314 +

315 + vllm_config = create_vllm_config()

316 +

317 + # Test worker role in decode server.

318 + connector = NixlConnector(vllm_config, KVConnectorRole.WORKER)

319 + connector.connector_worker = FakeNixlConnectorWorker(

320 + vllm_config, connector.engine_id)

321 + metadata = NixlConnectorMetadata()

322 + total_reqs = 5

323 + for i in range(total_reqs):

324 + metadata.add_new_req(request_id=f"id_{i}",

325 + local_block_ids=[1, 2, 3],

326 + kv_transfer_params={

327 + "remote_block_ids": [4, 5, 6],

328 + "remote_engine_id":

329 + FakeNixlConnectorWorker.REMOTE_ENGINE_ID,

330 + "remote_host": "localhost",

331 + "remote_port": 1234,

332 + })

333 + connector.bind_connector_metadata(metadata)

334 +

335 + timeout = 2.5 * total_reqs

336 + cnt_finished_reqs = 0

337 + start = time.perf_counter()

338 + while time.perf_counter() - start < timeout:

339 + dummy_ctx = ForwardContext(

340 + no_compile_layers={},

341 + attn_metadata={},

342 + virtual_engine=0,

343 + )

344 + _before_load = time.perf_counter()

345 + connector.start_load_kv(dummy_ctx)

346 + _after_load = time.perf_counter()

347 + assert _after_load - _before_load < 0.1, "start_load_kv took " \

348 + f"{_after_load - _before_load} seconds"

349 + time.sleep(0.5) # backoff for the async handshake to complete.

350 + connector.bind_connector_metadata(NixlConnectorMetadata())

351 + _, done_recving = connector.get_finished(finished_req_ids=set())

352 + if len(done_recving) > 0:

353 + cnt_finished_reqs += len(done_recving)

354 + if cnt_finished_reqs == total_reqs:

355 + return

356 + raise TimeoutError("Took too long to complete async handshake.")

vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py

+96 -38

@@ -2,11 +2,13 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 import contextlib

4 import math

5 +import queue

6 import threading

7 import time

8 import uuid

9 from collections import defaultdict

10 from collections.abc import Iterator

11 +from concurrent.futures import Future, ThreadPoolExecutor

12 from dataclasses import dataclass

13 from typing import TYPE_CHECKING, Any, Optional

@@ -23,6 +25,7 @@ from vllm.distributed.parallel_state import (

25 get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size,

26 get_tp_group)

27 from vllm.distributed.utils import divide

28 +from vllm.forward_context import ForwardContext

29 from vllm.logger import init_logger

30 from vllm.platforms import _Backend

31 from vllm.utils import make_zmq_path, make_zmq_socket, round_down

@@ -31,7 +34,6 @@ from vllm.v1.request import RequestStatus

35 if TYPE_CHECKING:

36 from vllm.attention.backends.abstract import AttentionMetadata

34 - from vllm.forward_context import ForwardContext

37 from vllm.v1.core.kv_cache_manager import KVCacheBlocks

38 from vllm.v1.request import Request

@@ -71,7 +73,7 @@ class ReqMeta:

73 remote_block_ids: list[int]

74 remote_host: str

75 remote_port: int

74 - remote_engine_id: str

76 + remote_engine_id: EngineId

79 class NixlConnectorMetadata(KVConnectorMetadata):

@@ -81,7 +83,7 @@ class NixlConnectorMetadata(KVConnectorMetadata):

84 def add_new_req(

85 self,

84 - request_id: str,

86 + request_id: ReqId,

87 local_block_ids: list[int],

88 kv_transfer_params: dict[str, Any],

89 ):

@@ -102,7 +104,7 @@ class NixlConnector(KVConnectorBase_V1):

104 self.engine_id: EngineId = vllm_config.kv_transfer_config.engine_id

105

106 if role == KVConnectorRole.SCHEDULER:

105 - self.connector_scheduler : Optional[NixlConnectorScheduler] = \

107 + self.connector_scheduler: Optional[NixlConnectorScheduler] = \

108 NixlConnectorScheduler(vllm_config, self.engine_id)

109 self.connector_worker: Optional[NixlConnectorWorker] = None

110 elif role == KVConnectorRole.WORKER:

@@ -186,7 +188,7 @@ class NixlConnectorScheduler:

188 self.side_channel_host = envs.VLLM_NIXL_SIDE_CHANNEL_HOST

189 self.side_channel_port = (

190 envs.VLLM_NIXL_SIDE_CHANNEL_PORT +

189 - vllm_config.parallel_config.data_parallel_rank_local *

191 + vllm_config.parallel_config.data_parallel_rank *

192 vllm_config.parallel_config.tensor_parallel_size)

193 logger.info("Initializing NIXL Scheduler %s", engine_id)

194

@@ -343,7 +345,7 @@ class NixlConnectorWorker:

345 # Each TP rank listens/queries on the base_port + tp_rank.

346 self.side_channel_port: int = (

347 envs.VLLM_NIXL_SIDE_CHANNEL_PORT +

346 - vllm_config.parallel_config.data_parallel_rank_local *

348 + vllm_config.parallel_config.data_parallel_rank *

349 vllm_config.parallel_config.tensor_parallel_size)

350

351 # Metadata.

@@ -386,8 +388,17 @@ class NixlConnectorWorker:

388 self._done_sending_count: defaultdict[ReqId,

389 int] = defaultdict(lambda: 0)

390

389 - # Background thread for establishing new connections.

391 + # Background thread for handling new handshake requests.

392 self._nixl_handshake_listener_t: Optional[threading.Thread] = None

393 + # Background thread for initializing new NIXL handshakes.

394 + self._handshake_initiation_executor = ThreadPoolExecutor(

395 + # NIXL is not guaranteed to be thread-safe, limit 1 worker.

396 + max_workers=1,

397 + thread_name_prefix="vllm-nixl-handshake-initiator")

398 + self._ready_requests = queue.Queue[tuple[ReqId, ReqMeta]]()

399 + self._handshake_futures: dict[EngineId, Future[dict[int, str]]] = {}

400 + # Protects _handshake_futures and _remote_agents.

401 + self._handshake_lock = threading.RLock()

402

403 self.vllm_config = vllm_config

404 self.block_size = vllm_config.cache_config.block_size

@@ -416,6 +427,12 @@ class NixlConnectorWorker:

427 # finish reading before safely freeing the blocks.

428 self.consumer_notification_counts_by_req = defaultdict[ReqId, int](int)

429

430 + def __del__(self):

431 + """Cleanup background threads on destruction."""

432 + self._handshake_initiation_executor.shutdown(wait=False)

433 + if self._nixl_handshake_listener_t:

434 + self._nixl_handshake_listener_t.join(timeout=0)

435 +

436 @staticmethod

437 def _nixl_handshake_listener(metadata: NixlAgentMetadata,

438 ready_event: threading.Event, base_port: int,

@@ -443,7 +460,7 @@ class NixlConnectorWorker:

460 "Connection listener got unexpected message %s", msg)

461 sock.send_multipart((identity, b"", encoded_data))

462

446 - def _nixl_handshake(self, host: str, port: int):

463 + def _nixl_handshake(self, host: str, port: int) -> dict[int, str]:

464 """Do a NIXL handshake with a remote instance."""

465

466 start_time = time.perf_counter()

@@ -452,7 +469,7 @@ class NixlConnectorWorker:

469 # a hack to keep us moving. We will switch when moving to etcd

470 # or where we have a single ZMQ socket in the scheduler.

471

455 - def handshake(path: str, rank: int) -> NixlAgentMetadata:

472 + def handshake(path: str, rank: int) -> tuple[NixlAgentMetadata, str]:

473 # Send query for the request.

474 with zmq_ctx(zmq.REQ, path) as sock:

475 sock.send(GET_META_MSG)

@@ -462,19 +479,20 @@ class NixlConnectorWorker:

479 got_metadata_time = time.perf_counter()

480

481 # Register Remote agent.

465 - self.add_remote_agent(metadata, rank)

482 + remote_agent_name = self.add_remote_agent(metadata, rank)

483 setup_agent_time = time.perf_counter()

484

485 logger.debug("NIXL handshake: get metadata took: %s",

486 got_metadata_time - start_time)

487 logger.debug("NIXL handshake: add agent took: %s",

488 setup_agent_time - got_metadata_time)

472 - return metadata

489 + return metadata, remote_agent_name

490

491 # Handshake with remote agent-rank0 first to get the tp_size of remote

492 path = make_zmq_path("tcp", host, port)

493 logger.debug("Querying master rank metadata on path: %s", path)

477 - metadata = handshake(path, 0)

494 + rank_to_agent_name: dict[int, str] = {}

495 + metadata, rank_to_agent_name[0] = handshake(path, 0)

496

497 # Handshake only with the other TP remote the current local rank will

498 # pull from. With homogeneous TP it happens to be the same rank_i.

@@ -484,7 +502,10 @@ class NixlConnectorWorker:

502 path = make_zmq_path("tcp", host, port + p_remote_rank)

503 logger.debug("Querying metadata on path: %s at remote rank %s",

504 path, p_remote_rank)

487 - _ = handshake(path, p_remote_rank)

505 + _, rank_to_agent_name[p_remote_rank] = handshake(

506 + path, p_remote_rank)

507 +

508 + return rank_to_agent_name

509

510 def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):

511 """Register the KV Cache data in nixl."""

@@ -621,11 +642,11 @@ class NixlConnectorWorker:

642 daemon=True,

643 name="nixl_handshake_listener")

644 self._nixl_handshake_listener_t.start()

624 - ready_event.wait()

645 + ready_event.wait() # Wait for listener ZMQ socket to be ready.

646

647 def add_remote_agent(self,

648 nixl_agent_meta: NixlAgentMetadata,

628 - remote_tp_rank: int = 0):

649 + remote_tp_rank: int = 0) -> str:

650 """

651 Add the remote NIXL agent and prepare the descriptors for reading cache

652 blocks from remote.

@@ -666,8 +687,8 @@ class NixlConnectorWorker:

687 """ # noqa: E501

688 engine_id = nixl_agent_meta.engine_id

689 # TODO re-evaluate refreshing for scaling/recovery

669 - if remote_tp_rank in self._remote_agents.get(engine_id, ()):

670 - return

690 + if remote_tp_rank in self._remote_agents.get(engine_id, {}):

691 + return self._remote_agents[engine_id][remote_tp_rank]

692

693 if engine_id in self._tp_size:

694 assert self._tp_size[engine_id] == nixl_agent_meta.tp_size

@@ -677,9 +698,8 @@ class NixlConnectorWorker:

698 # layout and close outputs.

699 assert nixl_agent_meta.attn_backend_name == self.backend_name

700

680 - self._remote_agents[engine_id][

681 - remote_tp_rank] = self.nixl_wrapper.add_remote_agent(

682 - nixl_agent_meta.agent_metadata)

701 + remote_agent_name = self.nixl_wrapper.add_remote_agent(

702 + nixl_agent_meta.agent_metadata)

703

704 # Number of D TP workers reading from a single P TP worker. This is

705 # 1 when P and D `--tensor-parallel-size` match.

@@ -708,8 +728,9 @@ class NixlConnectorWorker:

728 "local_kv_heads*tp_ratio, block_size, head_dim] and same dtype."

729 )

730

711 - assert self.block_size == remote_block_size, "Remote P worker with " \

712 - "different block size is not supported"

731 + assert self.block_size == remote_block_size, (

732 + "Remote P worker with different block size is not supported "

733 + f"{self.block_size=} {remote_block_size=}")

734

735 # Create dst descs and xfer side handles. TP workers have same #blocks.

736 if engine_id in self.dst_num_blocks:

@@ -748,7 +769,9 @@ class NixlConnectorWorker:

769 descs = self.nixl_wrapper.get_xfer_descs(blocks_data, "VRAM")

770 self.dst_xfer_side_handles[

771 engine_id] = self.nixl_wrapper.prep_xfer_dlist(

751 - self._remote_agents[engine_id][remote_tp_rank], descs)

772 + remote_agent_name, descs)

773 +

774 + return remote_agent_name

775

776 def get_finished(self) -> tuple[set[str], set[str]]:

777 """

@@ -866,33 +889,68 @@ class NixlConnectorWorker:

889 We check for these trnxs to complete in each step().

890 """

891 for req_id, meta in metadata.requests.items():

892 + remote_engine_id = meta.remote_engine_id

893 logger.debug(

894 "start_load_kv for request %s from remote engine %s. "

895 "Num local_block_ids: %s. Num remote_block_ids: %s. ", req_id,

872 - meta.remote_engine_id, len(meta.local_block_ids),

896 + remote_engine_id, len(meta.local_block_ids),

897 len(meta.remote_block_ids))

874 - self._read_blocks(

875 - request_id=req_id,

876 - dst_engine_id=meta.remote_engine_id,

877 - local_block_ids=meta.local_block_ids,

878 - remote_block_ids=meta.remote_block_ids,

879 - remote_host=meta.remote_host,

880 - remote_port=meta.remote_port,

881 - )

898 + if remote_engine_id not in self._remote_agents:

899 + # Being optimistic to assume engine is usually ready, apply

900 + # lock only when the optimistic check fails.

901 + with self._handshake_lock:

902 + if remote_engine_id not in self._remote_agents:

903 + fut = self._handshake_futures.get(remote_engine_id)

904 + if fut is None:

905 + fut = self._handshake_initiation_executor.submit(

906 + self._nixl_handshake, meta.remote_host,

907 + meta.remote_port)

908 + self._handshake_futures[remote_engine_id] = fut

909 +

910 + def done_callback(f: Future[dict[int, str]],

911 + eid=remote_engine_id):

912 + with self._handshake_lock:

913 + del self._handshake_futures[eid]

914 + try:

915 + self._remote_agents[eid] = f.result()

916 + except Exception:

917 + logger.exception(

918 + "Handshake with %s failed", eid)

919 +

920 + fut.add_done_callback(done_callback)

921 +

922 + # TODO: handle failure state of future in the

923 + # callback, we want to fail the request in this case.

924 + def request_ready(_f: Future[Any],

925 + entry=(req_id, meta)):

926 + self._ready_requests.put(entry)

927 +

928 + fut.add_done_callback(request_ready)

929 + continue

930 + self._read_blocks_for_req(req_id, meta)

931 +

932 + # Start transfers for requests whose handshakes have now finished.

933 + while not self._ready_requests.empty():

934 + self._read_blocks_for_req(*self._ready_requests.get_nowait())

935 +

936 + def _read_blocks_for_req(self, req_id: str, meta: ReqMeta):

937 + logger.debug(

938 + "Remote agent %s available, calling _read_blocks for req %s",

939 + meta.remote_engine_id, req_id)

940 + self._read_blocks(

941 + request_id=req_id,

942 + dst_engine_id=meta.remote_engine_id,

943 + local_block_ids=meta.local_block_ids,

944 + remote_block_ids=meta.remote_block_ids,

945 + )

946

947 def _read_blocks(

948 self,

949 local_block_ids: list[int],

950 remote_block_ids: list[int],

887 - remote_host: str,

888 - remote_port: int,

951 dst_engine_id: str,

952 request_id: str,

953 ):

892 - # NOTE(rob): this takes ~2s. We need to get this off the hotpath.

893 - if dst_engine_id not in self._remote_agents:

894 - self._nixl_handshake(remote_host, remote_port)

895 -

954 # NOTE(rob): having the staging blocks be on the READER side is

955 # not going to work well (since we will have to call rearrange tensors).

956 # after we detect the txn is complete (which means we cannot make the