[Bugfix][NIXL] Fix Async Scheduler timeout issue

[NickLucche]

To keep P kv blocks from being stranded in the case a request is aborted during decode, we have introduced a TTL/timeout #20139.

With async scheduling + high concurrency requests, this log line (and consequent logic) can be observed

ERROR 09-27 17:06:12 [nixl_connector.py:1101] Potentially invalid KV blocks for unrecognized request cmpl-141e42e3-ad61-46ab-88cc-fb6aaa47e363-0 were retrieved by a decode worker. They may have expired.

The current workflow on main is something like:

____________________
iteration i (engine step)
____________________

(Scheduler) 
build_connector_meta(_)->meta.reqs_to_send = _reqs_need_send` # Requests to time-out are prepared for worker
      |
(Scheduler)  # Here P is done and we acknowledge it by setting timeout 
request_finished(_)->self._reqs_need_send[request.request_id] = some_time
      | 
worker runs . . .

____________________
iteration i+1
____________________
(Scheduler) 
build_connector_meta(_)->meta.reqs_to_send = _reqs_need_send` # Prepare the request for which we set timeout above
       |
     .  .  . 
       |
 (Worker) # Worker *receives* metadata about finished requests and starts tracking its timeout
start_load_kv(_) -> self._reqs_to_send.update(metadata.reqs_to_send)
. . .

Now if D is allowed in between iter_i and iter_i+1 to decode the requests and READ blocks from P (this happens on a separate async channel, not tied to scheduler step), the log.error can be observed as self._reqs_to_send has not yet been updated!

Basically async scheduling pushes the misalignment between the moment in which requests expiration is set (P side) and the moment in which blocks are read from D, by allowing the scheduler one step ahead and the runner to pick the next batch right up.

FIX #25777

[gemini-code-assist]

Code Review

This pull request addresses a race condition in the asynchronous scheduler that could lead to requests being incorrectly marked as expired. The fix introduces a new state to track requests within a processing batch, preventing premature expiration and erroneous error logging. The logic appears sound and effectively resolves the described issue. I have one suggestion to enhance the robustness of the implementation.

[gemini-code-assist]

Using discard() instead of remove() would be more robust. In a complex distributed system with asynchronous operations, it's safer to use discard() to avoid potential KeyError exceptions if the req_id is unexpectedly not in the set due to race conditions. This would make the worker more resilient.

A similar change would be beneficial on line 1126.

Suggested change

	self._reqs_to_process.remove(req_id)
	self._reqs_to_process.discard(req_id)

[tlrmchlsmth]

Tested in a WideEP setup -- fixes crashes I was seeing

[njhill]

Thanks for this @NickLucche. I'm wondering whether there's possibility for a "leak" of entries in self._reqs_to_process, specifically if the scheduler decides not to designate these as requests to send, for example if the prefill is aborted while in progress, or it's decided that it's too few tokens or whatever.

In this case the expiration won't be set and there will be no nixl notifications, so I can't see where the corresponding req_id would ever be removed from that set.