[RDT]: In-flight transfers race condition

[BatshevaBlack]

What happened + What you expected to happen

fetch_and_get_rdt_objects() can call remove_remote_agent() while a sibling transfer on the same NIXL remote agent is still in-flight, causing crashes

Summary

RDTManager.fetch_and_get_rdt_objects() in rdt_manager.py posts multiple NIXL transfers before waiting on any of them, then waits for all of them to return.
The version of remote metadata can change in between two transfers posted.
This can lead to a 'remove_remote_agent' call while previous in progress transfers that were posted are still inflight, this causes these inflight transfers to crash (NIXL_ERR_REMOTE_DISCONNECT).

The mechanism

fetch_and_get_rdt_objects() rdt_manager.py pipelines fetches for a list of object IDs:

# Phase 1 - post every fetch, nothing waits between them
fetch_requests = {}
for object_id in object_ids:
    fetch_requests[object_id] = self._trigger_fetch(object_id, use_object_store)

# Phase 2 - only now does anything wait, after ALL are posted
while fetch_requests:
    result[object_id] = self._wait_fetch(object_id, fetch_request, ...)

_trigger_fetch() leads to tensor_transport_manager.fetch_multiple_tensors()

fetch_multiple_tensors() nixl_tensor_transport.py:

if remote_agent_meta_version != self._remote_agents[remote_name]:
    nixl_agent.remove_remote_agent(remote_name)   # no in-flight check

The race

In 'fetch_and_get_rdt_objects'
If there are two objects to transfer:

• Object A is posted
• The GC thread running in the background (from previous itterations) deregisters some memory and leads to a version change.
• Object B is posted, (A is still in-flight) object B gets to the fetch_multiple_tensors check- sees the version has changed, and triggers- remove_remote-agent

The agent was removed while A was inflight- causes a NIXL error and crashes the application.

Suggested fix:

Do not parallel transfers:

Don't proceed to next transfer tunill you get a completion from previous one in fetch_and_get_rdt_objects, this isn't the best solution for performance since we serialize transfers that can be parallelized.

Prevent tearing down connections while requests are in flight

This suggest adding a tracking on inflight requests and preventing removing the remote agent while requests are inflight.

Can be treated as a reader-writer lock:

The thread callingremove_remote_agent() will be treated as the writer.
Each transfer's lifetime - from initialize_xfer() through completion will be treated as a reader

The writer should not proceed while any reader holds the lock, readers should not be blocked by each other.
This maps cleanly onto the existing call structure: acquire read-lock at initialize_xfer(), release at transfer completion (wait_fetch_complete()/cleanup); acquire write-lock before remove_remote_agent(), release after add_remote_agent().

Versions / Dependencies

Environment

• NIXL v1.4.0 @ c0a1102b
• Ray 3.0.0.dev0, TransferQueue (RDT) 0.1.7.dev0
• PyTorch 2.10.0+cu129, CUDA 12.9, Python 3.13.12
• 2×8 NVIDIA H100 80GB HBM3 (driver 535.216.03)

Reproduction script

Repro script:

cpu_scatter_disconnect_race.py
python cpu_scatter_disconnect_race.py -N 16 --shape 3,2048,64000,128 --warmup 2 --iters 3

Issue Severity

No response