[Core] Optimize SequenceStatus.is_finished by switching to IntEnum (#5974)

This commit is contained in:
Antoni Baum 2024-06-29 05:47:53 -07:00 committed by GitHub
parent 51e971d39e
commit 7c01f70641
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -39,24 +39,21 @@ PromptLogprobs = List[Optional[Dict[int, 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]: