refactor: try a different fix

Signed-off-by: Travis Johnson <tsjohnso@us.ibm.com>

Author: tjohnson31415

vllm/outputs.py

@@ -178,9 +178,9 @@ def from_seq_group(
         if seq_group.request_id in seq_id_to_seq_group:
             group: SequenceGroupBase = seq_id_to_seq_group[
                 seq_group.request_id]
-            assembled_seq_group = group.maybe_assemble_group(seq_group)
             if finished:
                 group.finish_seq(seq_group)
+            assembled_seq_group = group.maybe_assemble_group(seq_group)
             if assembled_seq_group is None:
                 return None
 

vllm/sequence.py

@@ -1471,6 +1471,9 @@ def add_request(request_id: str, engine, params, *args, **kwargs):
     def finish_seq(self, seq: SequenceGroup):
         """The sequence `seq` finishes, we should record the information.
         """
+        # idempotent
+        if seq.request_id not in self.to_be_finished:
+            return
         del self.to_be_finished[seq.request_id]
         self.finished_reqs[seq.request_id] = seq
 
@@ -1529,34 +1532,30 @@ def add_request(request_id: str, engine, params, **kwargs):
     def maybe_assemble_group(
             self, seq_group: SequenceGroup) -> Optional[SequenceGroup]:
 
-        # in the streaming mode, we will return the assembled sequence for the
-        # last remaining sequence, and return None for the rest of sequences
-        if self.streaming:
-            last_remaining_id = list(self.to_be_finished)[-1]
-            if seq_group.request_id == last_remaining_id:
+        # in the streaming mode, we must return the assembled sequence
+        # group while sequences are still processing, but only for one of
+        # the remaining sequences
+        if self.streaming and not seq_group.is_finished():
+            first_remaining_id = next(iter(self.to_be_finished))
+            if seq_group.request_id == first_remaining_id:
                 return self.assembled_seq_group
             return None
 
-        # in the non-streaming mode, we will return the assembled sequence
-        # when the last sequences finishes, and then return None for the
-        # rest of the time
-        if (len(self.to_be_finished) == 1
-                and seq_group.request_id in self.to_be_finished
-                and seq_group.is_finished()):
-            assert self.assembled_seq_group is not None
-            params = self.assembled_seq_group.sampling_params
-            assert isinstance(params, SamplingParams)
-            if not self.output_produced:
-                self.output_produced = True
-                if params._real_n is not None:
-                    # Get the top-n sequences.
-                    n = params._real_n or params.n
-                    seqs = self.assembled_seq_group.seqs
-                    sorting_key = lambda seq: seq.get_cumulative_logprob()
-                    sorted_seqs = sorted(seqs, key=sorting_key, reverse=True)
-                    top_n_seqs = sorted_seqs[:n]
-                    self.assembled_seq_group.seqs = top_n_seqs
-                return self.assembled_seq_group
-            if self.output_produced:
-                return None
-        return None
+        # for non-streaming and when all streamed sequences are finished,
+        # we will return the assembled sequence for the last finished sequence
+        if len(self.to_be_finished) > 0 or self.output_produced:
+            return None
+
+        assert self.assembled_seq_group is not None
+        params = self.assembled_seq_group.sampling_params
+        assert isinstance(params, SamplingParams)
+        self.output_produced = True
+        if params._real_n is not None:
+            # Get the top-n sequences.
+            n = params._real_n or params.n
+            seqs = self.assembled_seq_group.seqs
+            sorting_key = lambda seq: seq.get_cumulative_logprob()
+            sorted_seqs = sorted(seqs, key=sorting_key, reverse=True)
+            top_n_seqs = sorted_seqs[:n]
+            self.assembled_seq_group.seqs = top_n_seqs
+        return self.assembled_seq_group