fix request finish condition

octoml · Dec 11, 2023 · 159ca14 · 159ca14
1 parent 7ed88cd
commit 159ca14
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/serve/mlc_serve/engine/staging_engine.py b/serve/mlc_serve/engine/staging_engine.py
@@ -242,7 +242,7 @@ def step(self) -> InferenceStepResult:
                 prompt_len[request_id] = state.prompt_len
 
                 if seq_output.finish_reason is not None:
-                    del self.requests[request_id]
+                    gen_seq.is_finished = True
 
             for request_id, out_seqs in seq_outputs.items():
                 outputs.append(
@@ -252,6 +252,8 @@ def step(self) -> InferenceStepResult:
                         num_prompt_tokens=prompt_len[request_id],
                     )
                 )
+                if self.requests[request_id].is_finished:
+                    del self.requests[request_id]
 
         return InferenceStepResult(outputs=outputs)