Skip to content

Cut fixed per-request costs out of the reshard control plane - #808

Open
yinlin09 wants to merge 1 commit into
mainfrom
yinlin/reshard-rpc-floor
Open

Cut fixed per-request costs out of the reshard control plane#808
yinlin09 wants to merge 1 commit into
mainfrom
yinlin/reshard-rpc-floor

Conversation

@yinlin09

@yinlin09 yinlin09 commented Aug 28, 2026

Copy link
Copy Markdown
Collaborator

Every Stage-3 coordination pays fixed costs on its framed RPCs (coordinate, GET_METADATA, receiver arm):

  1. The 4-byte length prefix and the body go out as two separate send() calls with Nagle enabled, on requests and responses alike, exposing every hop to the delayed-ACK stall.
  2. The destination controller is asked for every registered unit's full pool manifest on every request, although work units register once per engine lifetime.
  3. Client sockets carry no keepalive, so a black-holed peer is only detected at the full receive timeout (300 s on the coordinate hop).
  4. The framed server never reaps its per-connection threads: one std::thread handle and stack per request, held until shutdown.

Changes: single-buffer framing both directions plus TCP_NODELAY on client and accepted sockets; SO_KEEPALIVE and TCP_USER_TIMEOUT bounded by the call's I/O timeout; the accept loop joins finished connection threads; and the destination metadata is cached per controller address. Staleness (engine replacement) surfaces as a plan-build or receiver-arm failure. A failed attempt that used the cache drops the entry; while no receiver has acknowledged its arm the attempt is side-effect-free beyond the abandoned claim and is replayed once on fresh metadata. Once any receiver has acknowledged (multi-destination plans arm every receiver concurrently), the failure is returned as is and the next request re-queries.

Validation

reshard_service_test passes in the OSS bazel build (ml-build container, clang-18), including the two new tests: RemoteMetadataCachedAndRefreshedOnStaleFailure (cache hit on the second request; exactly one refetch after a fingerprint-mismatch replay) and PartialReceiverArmFailureDropsCacheWithoutReplay (one of two receivers refuses its arm: no replay, the cache entry is dropped, the next request re-queries).

Benchmarked on tpu7x 1P1D (Qwen3.5-397B, prefill PCP8 → decode DP8, vllm-torchtpu with background Stage-3 submission, 888 transfers per leg, zero transfer failures either leg): plan build including the peer-metadata path drops from p50 2.1 ms / max 5.1 ms to p50 1.5 ms / max 3.0 ms; end-to-end serving metrics are unchanged. The coordinate call's floor is dominated by sender dispatch (p50 ~26 ms under 4-way concurrent coordination — the reply currently awaits all 8 sender-worker acks), which this PR deliberately does not touch: returning the coordinate reply at receiver-arm and letting sender dispatch complete asynchronously is the follow-up that would remove most of the remaining floor.

copybara-service Bot pushed a commit that referenced this pull request Aug 29, 2026
Every Stage-3 coordination paid four avoidable fixed costs on its framed
RPCs (coordinate, GET_METADATA, receiver arm):

1. The 4-byte length prefix and the body went out as two separate send()
   calls with Nagle enabled, on requests and responses alike, exposing
   every hop to the delayed-ACK stall (tens of ms on small RPCs).
2. The destination controller was asked for every registered unit's full
   pool manifest on every request, although work units register once per
   engine lifetime.
3. Client sockets carried no keepalive, so a black-holed peer was only
   detected at the full receive timeout.
4. The framed server never reaped its per-connection threads: one
   std::thread handle and stack per request, held until shutdown.

Changes:

- framed_rpc: single-buffer framing on both directions; TCP_NODELAY on
  client and accepted sockets; SO_KEEPALIVE plus TCP_USER_TIMEOUT bounded
  by the call's I/O timeout on client sockets; the accept loop joins
  finished connection threads.
- reshard_coordinator: destination metadata is cached per controller
  address. Staleness (engine replacement) surfaces as a plan-build or
  receiver-arm failure, both side-effect-free beyond the abandoned claim,
  and is repaired by invalidate + fresh query + one replay; a failure
  after the receiver ack is never replayed.

Validation: reshard package builds and reshard_service_test passes in the
glibc-2.36 container, including the new
RemoteMetadataCachedAndRefreshedOnStaleFailure test (cache hit on the
second request, exactly one refetch after a fingerprint-mismatch replay).

GitHub: #808
PiperOrigin-RevId: 973205820
Every Stage-3 coordination paid four avoidable fixed costs on its framed
RPCs (coordinate, GET_METADATA, receiver arm):

1. The 4-byte length prefix and the body went out as two separate send()
   calls with Nagle enabled, on requests and responses alike, exposing
   every hop to the delayed-ACK stall (tens of ms on small RPCs).
2. The destination controller was asked for every registered unit's full
   pool manifest on every request, although work units register once per
   engine lifetime.
3. Client sockets carried no keepalive, so a black-holed peer was only
   detected at the full receive timeout.
4. The framed server never reaped its per-connection threads: one
   std::thread handle and stack per request, held until shutdown.

Changes:

- framed_rpc: single-buffer framing on both directions; TCP_NODELAY on
  client and accepted sockets; SO_KEEPALIVE plus TCP_USER_TIMEOUT bounded
  by the call's I/O timeout on client sockets; the accept loop joins
  finished connection threads.
- reshard_coordinator: destination metadata is cached per controller
  address. Staleness (engine replacement) surfaces as a plan-build or
  receiver-arm failure. A failed attempt that used the cache drops the
  entry; while no receiver has acknowledged its arm the attempt is
  side-effect-free beyond the abandoned claim and is replayed once on
  fresh metadata. Once any receiver has acknowledged, the failure is
  returned as is and the next request re-queries.

Validation: reshard package builds and reshard_service_test passes in the
OSS bazel build (ml-build container, clang-18), including the new
RemoteMetadataCachedAndRefreshedOnStaleFailure test (cache hit on the
second request, exactly one refetch after a fingerprint-mismatch replay)
and PartialReceiverArmFailureDropsCacheWithoutReplay test (one of two
receivers refuses its arm: no replay, the cache entry is dropped, and the
next request re-queries).
@yinlin09
yinlin09 force-pushed the yinlin/reshard-rpc-floor branch from 89a3bf1 to 45e4d88 Compare September 3, 2026 10:59
copybara-service Bot pushed a commit that referenced this pull request Sep 3, 2026
Every Stage-3 coordination paid four avoidable fixed costs on its framed
RPCs (coordinate, GET_METADATA, receiver arm):

1. The 4-byte length prefix and the body went out as two separate send()
   calls with Nagle enabled, on requests and responses alike, exposing
   every hop to the delayed-ACK stall (tens of ms on small RPCs).
2. The destination controller was asked for every registered unit's full
   pool manifest on every request, although work units register once per
   engine lifetime.
3. Client sockets carried no keepalive, so a black-holed peer was only
   detected at the full receive timeout.
4. The framed server never reaped its per-connection threads: one
   std::thread handle and stack per request, held until shutdown.

Changes:

- framed_rpc: single-buffer framing on both directions; TCP_NODELAY on
  client and accepted sockets; SO_KEEPALIVE plus TCP_USER_TIMEOUT bounded
  by the call's I/O timeout on client sockets; the accept loop joins
  finished connection threads.
- reshard_coordinator: destination metadata is cached per controller
  address. Staleness (engine replacement) surfaces as a plan-build or
  receiver-arm failure. A failed attempt that used the cache drops the
  entry; while no receiver has acknowledged its arm the attempt is
  side-effect-free beyond the abandoned claim and is replayed once on
  fresh metadata. Once any receiver has acknowledged, the failure is
  returned as is and the next request re-queries.

Validation: reshard package builds and reshard_service_test passes in the
OSS bazel build (ml-build container, clang-18), including the new
RemoteMetadataCachedAndRefreshedOnStaleFailure test (cache hit on the
second request, exactly one refetch after a fingerprint-mismatch replay)
and PartialReceiverArmFailureDropsCacheWithoutReplay test (one of two
receivers refuses its arm: no replay, the cache entry is dropped, and the
next request re-queries).

GitHub: #808
PiperOrigin-RevId: 973205820
copybara-service Bot pushed a commit that referenced this pull request Sep 3, 2026
Every Stage-3 coordination paid four avoidable fixed costs on its framed
RPCs (coordinate, GET_METADATA, receiver arm):

1. The 4-byte length prefix and the body went out as two separate send()
   calls with Nagle enabled, on requests and responses alike, exposing
   every hop to the delayed-ACK stall (tens of ms on small RPCs).
2. The destination controller was asked for every registered unit's full
   pool manifest on every request, although work units register once per
   engine lifetime.
3. Client sockets carried no keepalive, so a black-holed peer was only
   detected at the full receive timeout.
4. The framed server never reaped its per-connection threads: one
   std::thread handle and stack per request, held until shutdown.

Changes:

- framed_rpc: single-buffer framing on both directions; TCP_NODELAY on
  client and accepted sockets; SO_KEEPALIVE plus TCP_USER_TIMEOUT bounded
  by the call's I/O timeout on client sockets; the accept loop joins
  finished connection threads.
- reshard_coordinator: destination metadata is cached per controller
  address. Staleness (engine replacement) surfaces as a plan-build or
  receiver-arm failure. A failed attempt that used the cache drops the
  entry; while no receiver has acknowledged its arm the attempt is
  side-effect-free beyond the abandoned claim and is replayed once on
  fresh metadata. Once any receiver has acknowledged, the failure is
  returned as is and the next request re-queries.

Validation: reshard package builds and reshard_service_test passes in the
OSS bazel build (ml-build container, clang-18), including the new
RemoteMetadataCachedAndRefreshedOnStaleFailure test (cache hit on the
second request, exactly one refetch after a fingerprint-mismatch replay)
and PartialReceiverArmFailureDropsCacheWithoutReplay test (one of two
receivers refuses its arm: no replay, the cache entry is dropped, and the
next request re-queries).

GitHub: #808
PiperOrigin-RevId: 973205820
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants