[Bug]: RPC client can discard a reply received before send completes and hang the caller #22150

System Info

Upstream main 48acf026 (the affected RPCClient code is also unchanged from 672e049c). macOS arm64, Python 3.11.18, PyZMQ 27.3.0. CPU-only; no CUDA build, GPU execution, or performance measurements.

Reproduction

RPCClient._call_async awaits self._client_socket.put_async(request) before creating the entry in _pending_futures. If the response reader processes a reply while that await is still pending, _handle_regular_response finds no receiver, drops the reply, and the caller subsequently waits on a future that can never complete. With the default timeout of None, it hangs indefinitely. Both successful results and remote errors can be lost.

This deterministic reproduction exercises the real response handler:

import asyncio
import threading
from types import SimpleNamespace
from tensorrt_llm.executor.rpc.rpc_client import RPCClient
from tensorrt_llm.executor.rpc.rpc_common import RPCResponse

async def main():
    client = RPCClient.__new__(RPCClient)
    client._closed = client._server_stopped = False
    client._timeout = None
    client._call_lock = threading.Lock()
    client._pending_futures = {}

    async def send(request):
        client._handle_regular_response(
            RPCResponse(request.request_id, result=42))
        await asyncio.sleep(0)

    client._client_socket = SimpleNamespace(put_async=send)
    print(await asyncio.wait_for(client._call_async("answer"), 1))

asyncio.run(main())

Actual: TimeoutError. Expected: 42.

I also reproduced the ordering with the real RPCServer, RPCClient, HMAC serialization and ZeroMQ IPC transport: a wrapper around the real send awaits confirmation that the real response reader has handled the reply before returning. Upstream loses the response; registering before send returns 42. This forces an adverse scheduling order; it is not a claim about its production frequency.

Proposed fix and validation

Register the response future under the existing close/admission lock before sending. Put the send inside the existing try/finally so send errors and cancellation remove the registration. Fire-and-forget calls should send without registering a response future.

I have a minimal patch and seven CPU regression cases. On upstream: five failures and two passes; patched: seven passes. Twenty focused existing RPC tests also pass, including the 10,000-call synchronous and asynchronous cases.

Local execution uses a source-loading harness because this machine cannot import the full NVIDIA native package. It executes the actual RPC/IPC source and the actual AsyncQueue/_SyncQueue utilities with native tracing/bootstrap imports bypassed and logging replaced. It does not validate the complete TensorRT-LLM installation. The tests will be registered in l0_cpu.yml for matched-build upstream CPU CI.

I checked related RPC searches and changed-file lists of matching open PRs. #21997 changes error logging in this file, not registration ordering; #21221 changes executor submission registration in rpc_proxy_mixin.py; #20313 concerns close/shutdown settlement already present on main. I did not find a competing fix for this send/response ordering.

AI assistance (Codex) was used for source inspection, implementation and testing.