Skip to content

Commit

Permalink
[Core] Optimize SequenceStatus.is_finished by switching to IntEnum (v…
Browse files Browse the repository at this point in the history
…llm-project#5974)

Signed-off-by: Alvant <alvasian@yandex.ru>
  • Loading branch information
Yard1 authored and Alvant committed Oct 26, 2024
1 parent cf1d8b7 commit c349e81
Showing 1 changed file with 11 additions and 14 deletions.
25 changes: 11 additions & 14 deletions vllm/sequence.py
Original file line number Diff line number Diff line change
Expand Up @@ -39,24 +39,21 @@ class Logprob:
SampleLogprobs = List[Dict[int, Logprob]]


class SequenceStatus(enum.Enum):
class SequenceStatus(enum.IntEnum):
"""Status of a sequence."""
WAITING = enum.auto()
RUNNING = enum.auto()
SWAPPED = enum.auto()
FINISHED_STOPPED = enum.auto()
FINISHED_LENGTH_CAPPED = enum.auto()
FINISHED_ABORTED = enum.auto()
FINISHED_IGNORED = enum.auto()
WAITING = 0
RUNNING = 1
SWAPPED = 2
# Note: anything after SWAPPED (2) will be considered
# as a finished status.
FINISHED_STOPPED = 3
FINISHED_LENGTH_CAPPED = 4
FINISHED_ABORTED = 5
FINISHED_IGNORED = 6

@staticmethod
def is_finished(status: "SequenceStatus") -> bool:
return status in [
SequenceStatus.FINISHED_STOPPED,
SequenceStatus.FINISHED_LENGTH_CAPPED,
SequenceStatus.FINISHED_ABORTED,
SequenceStatus.FINISHED_IGNORED,
]
return status > SequenceStatus.SWAPPED

@staticmethod
def get_finished_reason(status: "SequenceStatus") -> Union[str, None]:
Expand Down

0 comments on commit c349e81

Please sign in to comment.