PLUGINS/UCX: Single UCX connection per handle

[iyastreb]

What?

Keep just a single connection pointer per handle, not set of them

• Simplify the code and reasoning about handle connection
• Improve performance (no set nodes allocation on heap)
• Improve debuggability

Why?

A transfer handle targets exactly one remote agent, since createXferReq/prepXferReq take a single remote_agent and resolve the remote descriptors only against that agent's section. The engine keeps at most one connection per agent in remoteConnMap, and each nixlUcxPublicMetadata::conn is a const shared_ptr copied from that entry, so one agent means one nixlUcxConnection. The handle is also pinned to one worker, and a connection owns one endpoint per worker, so the whole descriptor list resolves to a single nixlUcxEp.

Testing

BW and post latency improves by 1% on low dimensions

UCX_TLS=^cuda_ipc nixlbench --initiator_seg_type=VRAM --target_seg_type=VRAM --start_block_size=4096 --max_block_size=4096 --num_iter=100000

Block Size (B)      Batch Size     B/W (GB/Sec)   Avg Lat. (us)  Avg Prep (us)  P99 Prep (us)  Avg Post (us)  P99 Post (us)  Avg Tx (us)    P99 Tx (us)
----------------------------------------------------------------------------------------------------------------------------------------------------------------
4096                1              0.943396       4.3            12.0           12.0           1.1           2.0            3.2            4.0
4096                1              0.952939       4.3            12.0           12.0           1.0           2.0            3.2            4.0

This also improves sglang performance by 1-2%

Summary by CodeRabbit

• Bug Fixes
  • Improved UCX connection handling to prevent conflicting connection assignments.
  • Increased transfer reliability by consistently using and validating a single endpoint for grouped requests.
  • Improved request cleanup and completion tracking after transfers.
  • Improved error handling by releasing pending transfer requests when operations fail.
  • Simplified transfer processing to improve consistency and reliability during range-based transfers.

[github-actions]

👋 Hi iyastreb! Thank you for contributing to ai-dynamo/nixl.

Your PR reviewers will review your contribution then trigger the CI to test your changes.

🚀

[coderabbitai]

📝 Walkthrough

Walkthrough

The UCX backend now tracks one connection per request handle. Standard range transfers use one endpoint, release superseded or failed requests, and flush the selected endpoint once.

Changes

UCX connection and transfer handling

Layer / File(s)	Summary
Single connection request state src/plugins/ucx/ucx_backend.cpp	Request handles track one UCX connection. Assignment, validation, append, release, completion, and cleanup use that connection.
Single-endpoint transfer posting src/plugins/ucx/ucx_backend.cpp, src/plugins/ucx/ucx_backend.h	Standard transfers use the endpoint from the first remote descriptor, validate all descriptors, release superseded or failed requests, post the range, and flush once. The batching helper and its declaration were removed.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Suggested reviewers: brminich, colinnv

Sequence Diagram(s)

sequenceDiagram
  participant nixlUcxEngine
  participant RemoteDescriptor
  participant UCXEndpoint
  nixlUcxEngine->>RemoteDescriptor: Read endpoint from first descriptor
  nixlUcxEngine->>RemoteDescriptor: Validate endpoint for each descriptor
  nixlUcxEngine->>UCXEndpoint: Post requests through selected endpoint
  nixlUcxEngine->>UCXEndpoint: Flush selected endpoint once

Loading

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title clearly and concisely describes the primary UCX change to use a single connection per transfer handle.
Description check	✅ Passed	The description explains what changed, why it changed, and provides benchmark testing results; the optional How section is not required.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@src/plugins/ucx/ucx_backend.cpp`:
- Around line 1240-1251: Update the request handling around pending_req so every
NIXL_IN_PROG UCX request is retained rather than freeing or overwriting the
previous request. Before the failure path exits, cancel or drain all tracked UCX
requests and only then release them, ensuring UCX no longer accesses transfer
buffers before returning the error.

🪄 Autofix

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Enterprise

Run ID: a1146478-0bdd-4d10-b9fc-99ed1641dfe4

📥 Commits

Reviewing files that changed from the base of the PR and between 32794e9 and d85016c.

📒 Files selected for processing (2)

• src/plugins/ucx/ucx_backend.cpp
• src/plugins/ucx/ucx_backend.h

💤 Files with no reviewable changes (1)

• src/plugins/ucx/ucx_backend.h

[iyastreb]

/build

[svc-nixl]

🤖 CI Triage Agent — nixl-ci-build-container-pr · commit 8c099327

TL;DR: The container builds failed because apt-get could not reach archive.ubuntu.com/security.ubuntu.com (network unreachable / connection timed out) during the Docker image build — an infrastructure/network flake, not a code defect. Re-run the build once network is restored (and consider a mirror/retry hardening).

Full analysis

Summary: Multiple "Build image" stages (nixl/x86_64 node 156, nixlbench/x86_64 node 233, and other arches) failed with apt-get ... exit status 100.

Root cause: Loss of network connectivity to the Ubuntu package mirrors during the image build. In node 233, apt-get update could not connect to archive.ubuntu.com:80 or security.ubuntu.com:80 at all (connect (111: Connection refused) / connect (101: Network is unreachable)), so every package "Unable to locate package". In node 156, the environment was already degraded (package list fetch took 10m37s at 42.4 kB/s), and the build ultimately died fetching linux-libc-dev_6.8.0-137.137_amd64.deb with Cannot initiate the connection to archive.ubuntu.com:80 ... Network is unreachable ... connection timed out. Both are the same transient network failure surfaced as apt exit code 100. This is unrelated to PR #2040 / commit 8c09932.

Implicated commit: unknown — not code-related; environmental network failure on the Jenkins build agents.

File: contrib/Dockerfile STEP 22–23 (apt-get install / DOCA install) for the nixl image; nixlbench Dockerfile STEP 4 (apt-get install) — these are where apt hits the network, but the files are not the cause.

Suggested fix:

1. Re-run the failed build (retry) — the network outage is transient.
2. To harden against these flakes: point apt at an internal/caching mirror (e.g. apt-cacher-ng or a corporate mirror) instead of archive.ubuntu.com, and/or force IPv4 (Acquire::ForceIPv4 "true"; — note the many 2620:.../2a06:... IPv6 attempts failing with "Network is unreachable"). Adding apt-get -o Acquire::Retries=5 and wrapping the apt steps in a retry loop would also reduce sensitivity to brief outages.

Related: none (no matching code issue/PR found).