feat: bandwidth-aware load_aware P2P source selection

[yixinh-nv]

Summary

Adds a third P2P source-selection policy, load_aware, that biases weight-pull
source selection away from sources whose RDMA path is already busy, so scaling up
or rolling a live inference fleet does not steal bandwidth from sources that are
actively serving. It builds on the random / rendezvous_hash selector interface
from #461: same ScoredSelector seam, one new scalar on the wire, no new RPCs, and
it collapses to rendezvous_hash whenever the load signal is absent — never worse
than the deterministic baseline.

The MX servers stay stateless: the load signal is published by the source about
itself and passed through by the server — the server never counts or accumulates
anything — so replicas can sit behind a load balancer and scale arbitrarily, and any
two replicas rank an identical ListSources identically.

What changed

source_load: a source-published load scalar (proto + server passthrough)

• New float source_load on SourceInstanceRef (=6) and UpdateStatusRequest (=5)
  in p2p.proto (Python + Go bindings regenerated).
• The server plumbs it as an opaque passthrough: WorkerRecord.source_load, refreshed
  on update_status and surfaced on list_sources; memory / Redis (#[serde(default)])
  / Kubernetes-CRD (sourceLoad) backends all carry it. The server never derives it.

load_aware selector (client)

• score = unit_hash(target, candidate) − w_load · source_load, source_load clamped
  to [0,1], w_load = MX_P2P_LOAD_WEIGHT (default 1.0, clamped ≥ 0).
• Absent/zero load ⇒ every penalty is 0 ⇒ exact rendezvous_hash ordering (safe
  default; also covers old servers that predate the field).

Where the load number comes from (two providers, blended)

source_load is produced on the source by make_source_load_provider(device_id),
reporting max() of two independent signals so selection reacts to a NIC that is
physically hot and to an imminent serving spike a counter has not seen yet:

1. RDMA NIC utilization (nic_metrics.py, runtime-agnostic, default). Reads the
   GPU-affine RDMA NIC's InfiniBand port byte counters
   (/sys/class/infiniband/<dev>/ports/<p>/counters/{port_xmit_data,port_rcv_data},
   ×4 for the 4-octet spec unit) and divides the delta since the last sample by the
   elapsed time and the link capacity (parsed from .../rate). The NIC is the exact
   rail transfers pin to — resolved via the same ucx_utils.probe_nic_pin_for_device
   PCIe-affinity assignment the NIXL transfer path uses — so it reflects the contention
   a puller would actually hit. Degrades safely to 0.0 (no device / unreadable sysfs
   / unparseable rate), so IB/RoCE report a real number and other fabrics fall back to
   rendezvous_hash. Sampled off the existing publisher heartbeat (no extra sleeps).
2. Inference-runtime serving load (runtime_load.py, opt-in via
   MX_P2P_RUNTIME_METRICS_URL). Scrapes the co-located vLLM/SGLang Prometheus
   /metrics for a KV-cache-utilization gauge (vllm:kv_cache_usage_perc, or the
   legacy vllm:gpu_cache_usage_perc; sglang:token_usage) as a [0,1] serving-load
   proxy — predictive, since a full KV cache foretells RDMA the NIC counter has not
   measured yet, and it works on fabrics whose sysfs counters cannot be read.

Neither provider touches the server, proto, or selector; the only wire contract is the
normalized source_load scalar, and any provider error degrades to 0.0.

Observability

• mx_p2p_source_load{scheme, source_worker_id} gauge recorded per candidate at
  selection (opt-in MX_METRICS_ENABLED=1), alongside the existing
  mx_p2p_source_selections_total / mx_p2p_transfer_seconds.

Compatibility

• Additive proto field; old clients/servers ignore it and load_aware degrades to
  rendezvous_hash. Default policy remains random. No new RPCs, no publish-cadence
  change, no server-side state.

Validation

Unit tests

• Python: load_aware collapse-to-rendezvous at zero load, steering away from the
  busiest source, w_load monotonicity, negative-weight clamp, [0,1] clamping,
  missing-field tolerance; NIC sampler (injected reader/clock/link-rate); runtime
  provider (vLLM new/legacy gauge names, SGLang, unreachable ⇒ 0); publisher forwards a
  non-zero provider value and falls back to 0 on error. Rust: source_load passthrough
  across memory/Redis/k8s backends with non-zero round-trip assertions.

B200 E2E

Ran on B200, InfiniBand RDMA with a vLLM disaggregated
prefill/decode fleet, TP=1, kubernetes metadata backend, MX_METRICS_ENABLED=1, and
source_load = max(NIC util, vLLM KV-cache usage). Fleet driven by AIPerf (closed-loop
concurrency 48, 512-in / 128-out, streaming chat) against the dynamo frontend, capturing
TTFT / request-latency / MX metrics at baseline, during a Case-1 scale-up
(decode 4→8 under load, new replicas pull weights over MX P2P), and during a Case-2
rolling update (decode replicas replaced one at a time under load).

• Mechanism, end to end. New replicas load weights via --load-format modelexpress
  and register as P2P sources publishing source_load; the NIC+KV blend, the selector,
  and the mx_p2p_source_load telemetry all work on real hardware.
• The source_load signal is real and heterogeneous under load. On a memory-tight
  Qwen2.5-72B fleet the KV-cache signal saturated — prefill workers reported
  99.7% KV usage vs. 0% on the idle decode workers — and during the scale-up
  load_aware routed the new pullers to the idle sources, away from the 99.7%-loaded
  prefill sources (the intended steering).
• No regression at typical model sizes. On Qwen2.5-14B, a scale-up under load
  held TTFT flat (p99 ≈ 380 ms baseline → ≈ 375 ms during scale-up) with throughput up
  ~5%; here the B200's KV headroom keeps source_load low, so load_aware rides its
  rendezvous_hash fallback — i.e. it matches the deterministic baseline exactly, the
  designed safety property.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[yixinh-nv]

/ok to test 29a8d4e

[coderabbitai]

Walkthrough

Changes

Load-aware P2P source selection

Layer / File(s)	Summary
Protocol and server load propagation modelexpress_common/proto/p2p.proto, modelexpress_server/src/p2p/*, modelexpress_client/python/modelexpress/p2p_pb2.py	Adds source_load to status, discovery, server records, and memory, Redis, and Kubernetes metadata paths.
Client load sampling and heartbeat publication modelexpress_client/python/modelexpress/{nic_metrics.py,runtime_load.py}, modelexpress_client/python/modelexpress/metadata/*, modelexpress_client/python/modelexpress/client.py	Samples NIC and optional runtime metrics, combines them into source_load, and publishes the value during heartbeats with fallback behavior.
Load-aware selection and observability modelexpress_client/python/modelexpress/source_selection.py, modelexpress_client/python/modelexpress/load_strategy/*, modelexpress_client/python/modelexpress/metrics.py, docs/*	Adds the load_aware policy, load-penalized rendezvous scoring, candidate load metrics, tests, and configuration documentation.

Estimated code review effort: 4 (Complex) | ~45 minutes

Poem

I’m a rabbit with metrics, hopping light,
Busy sources fade from sight.
NICs and runtimes whisper, “Load!”
Heartbeats carry the woodland code.
Hashes rank the paths just right—
Carrots saved for retries tonight.

🚥 Pre-merge checks | ✅ 3 | ❌ 2

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Linked Issues check	⚠️ Warning	The PR adds load-aware source_load telemetry, but #477 also required active_transfers metadata and server-side TTL-based load tracking, which are not present.	Add the missing active_transfers/TTL load-tracking path and expose it through ListSources, or revise the issue scope to match the source_load design.
Docstring Coverage	⚠️ Warning	Docstring coverage is 33.33% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Out of Scope Changes check	✅ Passed	The changes stay focused on the load_aware P2P selection feature, related telemetry, docs, and tests, with no clear unrelated additions.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title is concise and accurately summarizes the main change: adding load-aware P2P source selection.

---

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

[coderabbitai]

Actionable comments posted: 4

🧹 Nitpick comments (2)

modelexpress_client/python/tests/test_publisher.py (1)

98-119: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Cover the injected provider path.

These tests only assert the default 0.0 value. Add direct _tick() tests for a non-zero source_load_provider result and a raising provider falling back to 0.0.

🤖 Prompt for 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.

In `@modelexpress_client/python/tests/test_publisher.py` around lines 98 - 119,
Extend the PublisherThread _tick tests to cover source_load_provider: verify a
non-zero provider result is passed as source_load in mx_client.update_status,
and verify a provider that raises still publishes with source_load=0.0. Reuse
the existing publisher setup and assert both outcomes through
update_status.call_args_list.

modelexpress_server/src/p2p/backend/redis.rs (1)

740-740: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win

Assert nonzero source_load propagation. Current fixtures use 0.0 and wildcard predicates, so they validate field arity but not the new value flow.

• modelexpress_server/src/p2p/backend/redis.rs#L740-L740: use a nonzero value and assert it survives JSON round-trip.
• modelexpress_server/src/p2p/service.rs#L637-L646: pass a nonzero request value and assert the backend receives it.
• modelexpress_server/src/p2p/service.rs#L688-L698: return distinct nonzero loads and assert they appear in ListSourcesResponse.
• modelexpress_server/src/p2p/state.rs#L525-L544: replace the final wildcard matcher with the expected nonzero load.

🤖 Prompt for 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.

In `@modelexpress_server/src/p2p/backend/redis.rs` at line 740, Update the
source_load test fixtures to verify nonzero value propagation: in
modelexpress_server/src/p2p/backend/redis.rs:740, use a nonzero source_load and
assert it survives the JSON round-trip; in
modelexpress_server/src/p2p/service.rs:637-646, send a nonzero request value and
assert the backend receives it; in
modelexpress_server/src/p2p/service.rs:688-698, return distinct nonzero loads
and assert they appear in ListSourcesResponse; in
modelexpress_server/src/p2p/state.rs:525-544, replace the final wildcard matcher
with the expected nonzero load.

🤖 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 `@docs/ARCHITECTURE.md`:
- Around line 617-625: Update the stale architecture references around the
source-selection documentation: revise the old method signature at Line 559 to
include the current load-aware/source-load behavior, and change the Line 631
wording so load_aware is documented as shipped rather than deferred. Keep the
surrounding selector and topology-aware policy descriptions consistent with the
implemented three-policy registry.

In `@modelexpress_client/python/modelexpress/envs.py`:
- Around line 202-203: Update the MX_P2P_LOAD_WEIGHT configuration parser in
_env_float usage to reject or clamp negative values, ensuring LoadAwareSelector
receives a nonnegative w_load. Add a test covering negative MX_P2P_LOAD_WEIGHT
input and verify the invalid configuration follows the chosen behavior.

In `@modelexpress_client/python/modelexpress/runtime_load.py`:
- Around line 31-35: Add “vllm:kv_cache_usage_perc” first in _RATIO_METRICS,
retaining “vllm:gpu_cache_usage_perc” as the legacy fallback and preserving the
existing sglang metrics. Add or update a test covering runtime load selection
when the new vLLM metric is present.

In `@modelexpress_client/python/tests/test_runtime_load.py`:
- Around line 70-79: Update test_make_source_load_provider_blends_when_url_set
to avoid real socket I/O by patching the runtime metrics fetch/provider seam,
and the NIC sampler if necessary, to return deterministic failure values.
Preserve the test’s assertions that the blended provider does not crash and
returns a float within [0, 1].

---

Nitpick comments:
In `@modelexpress_client/python/tests/test_publisher.py`:
- Around line 98-119: Extend the PublisherThread _tick tests to cover
source_load_provider: verify a non-zero provider result is passed as source_load
in mx_client.update_status, and verify a provider that raises still publishes
with source_load=0.0. Reuse the existing publisher setup and assert both
outcomes through update_status.call_args_list.

In `@modelexpress_server/src/p2p/backend/redis.rs`:
- Line 740: Update the source_load test fixtures to verify nonzero value
propagation: in modelexpress_server/src/p2p/backend/redis.rs:740, use a nonzero
source_load and assert it survives the JSON round-trip; in
modelexpress_server/src/p2p/service.rs:637-646, send a nonzero request value and
assert the backend receives it; in
modelexpress_server/src/p2p/service.rs:688-698, return distinct nonzero loads
and assert they appear in ListSourcesResponse; in
modelexpress_server/src/p2p/state.rs:525-544, replace the final wildcard matcher
with the expected nonzero load.

🪄 Autofix (Beta)

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: CHILL

Plan: Enterprise

Run ID: b0e3d3b3-d05d-448f-8ce7-cb03da69d17e

📥 Commits

Reviewing files that changed from the base of the PR and between 656133f and a2b16a6.

⛔ Files ignored due to path filters (1)

• modelexpress_client/go/gen/modelexpress/p2p/p2p.pb.go is excluded by !**/*.pb.go, !**/gen/**

📒 Files selected for processing (26)

• docs/ARCHITECTURE.md
• docs/DEPLOYMENT.md
• modelexpress_client/python/modelexpress/client.py
• modelexpress_client/python/modelexpress/envs.py
• modelexpress_client/python/modelexpress/load_strategy/rdma_strategy.py
• modelexpress_client/python/modelexpress/metadata/k8s_service_client.py
• modelexpress_client/python/modelexpress/metadata/publish.py
• modelexpress_client/python/modelexpress/metadata/publisher.py
• modelexpress_client/python/modelexpress/metrics.py
• modelexpress_client/python/modelexpress/nic_metrics.py
• modelexpress_client/python/modelexpress/p2p_pb2.py
• modelexpress_client/python/modelexpress/runtime_load.py
• modelexpress_client/python/modelexpress/source_selection.py
• modelexpress_client/python/tests/test_nic_metrics.py
• modelexpress_client/python/tests/test_publisher.py
• modelexpress_client/python/tests/test_runtime_load.py
• modelexpress_client/python/tests/test_source_selection.py
• modelexpress_common/proto/p2p.proto
• modelexpress_server/src/p2p/backend.rs
• modelexpress_server/src/p2p/backend/kubernetes.rs
• modelexpress_server/src/p2p/backend/memory.rs
• modelexpress_server/src/p2p/backend/redis.rs
• modelexpress_server/src/p2p/k8s_types.rs
• modelexpress_server/src/p2p/reaper.rs
• modelexpress_server/src/p2p/service.rs
• modelexpress_server/src/p2p/state.rs