Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>
This commit is contained in:
rshaw@neuralmagic.com 2025-06-30 01:59:02 +00:00
parent cda2f2c453
commit 5d8b665366

View File

@ -340,7 +340,7 @@ class NixlConnectorWorker:
self.nixl_wrapper = NixlWrapper(str(uuid.uuid4()), self.nixl_wrapper = NixlWrapper(str(uuid.uuid4()),
None, None,
num_workers=num_workers, num_workers=num_workers,
num_shared_workers=None) num_shared_workers=None)
# Map of engine_id -> {rank0: agent_name0, rank1: agent_name1..}. # Map of engine_id -> {rank0: agent_name0, rank1: agent_name1..}.
self._remote_agents: dict[str, dict[int, str]] = defaultdict(dict) self._remote_agents: dict[str, dict[int, str]] = defaultdict(dict)
@ -999,9 +999,12 @@ class NixlConnectorWorker:
# Begin async xfer. # Begin async xfer.
start = time.perf_counter() start = time.perf_counter()
for handle in handles: # IT WORKS WITH THIS:
self.nixl_wrapper.transfer(handle) # for handle in handles:
# self.nixl_wrapper.transfer_batched(handles) # self.nixl_wrapper.transfer(handle)
# IT FAILS WITH THIS:
self.nixl_wrapper.transfer_batched(handles)
end = time.perf_counter() end = time.perf_counter()
logger.info("======== LAUNCH TIME: %s ========", end - start) logger.info("======== LAUNCH TIME: %s ========", end - start)