Commit 8ce5d319 authored by Robert Shaw's avatar Robert Shaw Committed by simon-mo
Browse files
parent 09c2cbc0
......@@ -474,8 +474,11 @@ class NixlConnectorWorker:
"backends", ["UCX"])
# Agent.
non_ucx_backends = [b for b in self.nixl_backends if b != "UCX"]
config = nixl_agent_config(backends=self.nixl_backends) if len(
non_ucx_backends) > 0 and nixl_agent_config is not None else None
if nixl_agent_config is None:
config = None
else:
config = nixl_agent_config(backends=self.nixl_backends) if len(
non_ucx_backends) > 0 else nixl_agent_config(num_threads=8)
self.nixl_wrapper = NixlWrapper(str(uuid.uuid4()), config)
# Map of engine_id -> {rank0: agent_name0, rank1: agent_name1..}.
......
......@@ -1288,4 +1288,9 @@ class Scheduler(SchedulerInterface):
self.finished_recving_kv_req_ids.add(req_id)
for req_id in (kv_connector_output.finished_sending or ()):
logger.debug("Finished sending KV transfer for request %s", req_id)
self._free_blocks(self.requests[req_id])
if req_id not in self.requests:
logger.warning(
"Got finished sending KV transfer for request %s,"
"but the request is already freed.", req_id)
else:
self._free_blocks(self.requests[req_id])
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment