Early-stop on cancelled RPCs

PiperOrigin-RevId: 680872920
Change-Id: I1171fab0cce5add718b6bc2fa9c3787283947877

saxml/server/model_service_base.py

@@ -2450,9 +2450,15 @@ def _run_generation_loop(
 
           state.scores += scores * prev_mask
           state.steps += prev_mask
-          done = done * prev_mask
-
           done = np.logical_or(done, state.steps >= state.max_decode_steps)
+          done = np.logical_or(
+              done,
+              [
+                  t.rpc.should_cancel() if t and t.rpc else False
+                  for t in state.rpc_tasks
+              ],
+          )
+          done = np.logical_and(done, prev_mask)
 
           if np.any(done):
             # Detokenize and send RPC in another thread for done slots
@@ -2500,6 +2506,11 @@ def _postprocess():
         if rpc_task.aux['slot_count'] > len(rpc_task.aux['finished_results']):
           assert not state.method.streamable_output
           continue
+        if rpc_task.rpc and rpc_task.rpc.should_cancel():
+          logging.info('request cancelled.')
+          rpc_task.done(utils.cancelled())
+          state.rpc_tasks[slot] = None
+          continue
         # [num_samples, ...]
         seqs = np.stack(
             [x for x, _ in rpc_task.aux['finished_results']], axis=0