fix: unbatch removals of requests from input_batch (#511)

# Description

Calls to `input_batch.remove_request` adjust the internal state of the
batch and change the mapping from request ids to indices. If multiple
removals are batched, they can cause the `batch_update` to have repeat
indices, which leads to bugs. The change here forces input batch
metadata updates to happen for each removal one-at-a-time; completely
unbatching the batch_updates.

## Related Issues

Fixes #508
See also
#492 (comment)

---------

Signed-off-by: Travis Johnson <[email protected]>

Author: tjohnson31415

tests/v1/worker/test_spyre_input_batch.py

@@ -36,6 +36,14 @@ def _remove_requests(input_batch: SamplingInputBatch, batch_size: int,
     for index in req_indices_to_remove:
         input_batch.remove_request(reqs[index].req_id)
         req_ids_to_remove.add(reqs[index].req_id)
+
+    # FIXME: it is a bug in the current implementation that removed indices may
+    # be duplicated, which can break logitsprocs tracking. Once fixed we should
+    # add this assert.
+    # see also: https://github.com/vllm-project/vllm-spyre/issues/508
+    # removed = input_batch.batch_update_builder.removed
+    # assert len(set(removed)) == len(removed), "Duplicate removed indices"
+
     return req_ids_to_remove
 
 

vllm_spyre/v1/worker/spyre_model_runner.py

@@ -394,7 +394,10 @@ def update_states(self, scheduler_output: SchedulerOutput):
             for req_id in scheduler_output.finished_req_ids:
                 self.input_batch.remove_request(req_id)
                 self.requests.pop(req_id, None)
-            self.input_batch.refresh_metadata()
+                # TODO: Processing multiple removals at once can break alignment
+                # of logitprocs. Refactor so that we can batch removals to the
+                # `input_batch`
+                self.input_batch.refresh_metadata()
 
     def _get_prompt_logprobs_dict(
         self,