feat: support human-readable suffixes (25.6k, 1M, 1Mi) for token CLI (#20577)
This commit is contained in:
@@ -43,6 +43,7 @@ from sglang.srt.utils.common import (
|
||||
get_device_sm,
|
||||
get_int_env_var,
|
||||
get_quantization_config,
|
||||
human_readable_int,
|
||||
is_blackwell_supported,
|
||||
is_cpu,
|
||||
is_cuda,
|
||||
@@ -3384,9 +3385,10 @@ class ServerArgs:
|
||||
)
|
||||
parser.add_argument(
|
||||
"--context-length",
|
||||
type=int,
|
||||
type=human_readable_int,
|
||||
default=ServerArgs.context_length,
|
||||
help="The model's maximum context length. Defaults to None (will use the value from the model's config.json instead).",
|
||||
help="The model's maximum context length. Defaults to None (will use the value from the model's config.json instead)."
|
||||
+ f"\n\n{human_readable_int.__doc__}",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--is-embedding",
|
||||
@@ -3613,10 +3615,11 @@ class ServerArgs:
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-total-tokens",
|
||||
type=int,
|
||||
type=human_readable_int,
|
||||
default=ServerArgs.max_total_tokens,
|
||||
help="The maximum number of tokens in the memory pool. If not specified, it will be automatically calculated based on the memory usage fraction. "
|
||||
"This option is typically used for development and debugging purposes.",
|
||||
"This option is typically used for development and debugging purposes."
|
||||
+ f"\n\n{human_readable_int.__doc__}",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--chunked-prefill-size",
|
||||
@@ -3638,9 +3641,10 @@ class ServerArgs:
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-prefill-tokens",
|
||||
type=int,
|
||||
type=human_readable_int,
|
||||
default=ServerArgs.max_prefill_tokens,
|
||||
help="The maximum number of tokens in a prefill batch. The real bound will be the maximum of this value and the model's maximum context length.",
|
||||
help="The maximum number of tokens in a prefill batch. The real bound will be the maximum of this value and the model's maximum context length."
|
||||
+ f"\n\n{human_readable_int.__doc__}",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--schedule-policy",
|
||||
|
||||
@@ -47,6 +47,7 @@ import warnings
|
||||
from collections import OrderedDict, defaultdict
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
from decimal import Decimal
|
||||
from functools import lru_cache, partial
|
||||
from importlib.metadata import PackageNotFoundError, version
|
||||
from importlib.util import find_spec
|
||||
@@ -2308,6 +2309,44 @@ def nullable_str(val: str):
|
||||
return val
|
||||
|
||||
|
||||
def human_readable_int(value: str) -> int:
|
||||
"""Supports standard SI suffixes (k, M, G, T) and IEC suffixes
|
||||
(Ki, Mi, Gi, Ti). Suffixes are case-sensitive.
|
||||
|
||||
Decimals are allowed for SI suffixes only.
|
||||
|
||||
Examples:
|
||||
'1k' -> 1000 '1M' -> 1000000 '25.6k' -> 25600
|
||||
'1Ki' -> 1024 '1Mi' -> 1048576
|
||||
"""
|
||||
value = value.strip()
|
||||
|
||||
si_multiplier = {"k": 10**3, "M": 10**6, "G": 10**9, "T": 10**12}
|
||||
iec_multiplier = {"Ki": 2**10, "Mi": 2**20, "Gi": 2**30, "Ti": 2**40}
|
||||
|
||||
match = re.fullmatch(r"(\d+(?:\.\d+)?)(Ki|Mi|Gi|Ti|k|M|G|T)", value)
|
||||
if match:
|
||||
number, suffix = match.groups()
|
||||
if suffix in iec_multiplier:
|
||||
if "." in number:
|
||||
raise argparse.ArgumentTypeError(
|
||||
f"Decimals are not allowed with IEC suffixes like '{suffix}'. "
|
||||
f"Use an integer IEC value such as '{int(Decimal(number))}{suffix}', "
|
||||
f"or an SI value such as '{number}{suffix[0]}'."
|
||||
)
|
||||
return int(number) * iec_multiplier[suffix]
|
||||
return int(Decimal(number) * si_multiplier[suffix])
|
||||
|
||||
try:
|
||||
return int(value)
|
||||
except ValueError:
|
||||
raise argparse.ArgumentTypeError(
|
||||
f"Invalid integer value: '{value}'. "
|
||||
"Use a plain integer, SI suffixes (1k, 1M), or IEC suffixes (1Ki, 1Mi). "
|
||||
"Suffixes are case-sensitive."
|
||||
)
|
||||
|
||||
|
||||
def pyspy_dump_schedulers():
|
||||
"""py-spy dump on all scheduler in a local node."""
|
||||
try:
|
||||
|
||||
Reference in New Issue
Block a user