fix(tp): acknowledge shm reads so rank0 cannot corrupt in-flight RPC (#246)

[Anai-Guo]

Summary

Fixes #246.

In tensor-parallel mode (tensor_parallel_size > 1), rank0 dispatches method calls to worker ranks through a shared-memory buffer plus a per-worker multiprocessing.Event. The protocol only had a forward (rank0 → worker) notification: there was no acknowledgement that a worker had finished reading the previous command.

As a result rank0 could return from one call() and immediately write_shm() the next RPC, overwriting the shared buffer while a worker was still inside pickle.loads(self.shm.buf[...]) for the previous command. That corrupts the payload (_pickle.UnpicklingError); once one worker rank dies, the surviving ranks hang in the next NCCL collective until timeout.

Fix

Add a per-worker acknowledgement event that closes the handshake:

• read_shm (worker): after fully reading and unpickling the payload, self.ack_event.set().
• write_shm (rank0): before overwriting the buffer, wait() + clear() every worker's ack event, so the buffer is only reused once all workers have consumed the previous command.
• The ack events are pre-set in LLMEngine so the very first command isn't blocked waiting on a read that hasn't happened yet.

This is the ping-pong the buffer needs: rank0 clears each ack before setting the forward event, and a worker sets its ack only after the read completes, so no signal is lost and no read races an overwrite. No change to the public LLM/LLMEngine API; only ModelRunner's internal constructor gains an ack_event argument, and its sole caller (LLMEngine) is updated.

Testing

• python -m py_compile on both changed modules.
• Handshake reasoning verified against the exit path (call("exit") still drains cleanly) and the bootstrap case (pre-set acks let the first write_shm proceed).

🤖 Generated with Claude Code