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_device_sm,
|
||||||
get_int_env_var,
|
get_int_env_var,
|
||||||
get_quantization_config,
|
get_quantization_config,
|
||||||
|
human_readable_int,
|
||||||
is_blackwell_supported,
|
is_blackwell_supported,
|
||||||
is_cpu,
|
is_cpu,
|
||||||
is_cuda,
|
is_cuda,
|
||||||
@@ -3384,9 +3385,10 @@ class ServerArgs:
|
|||||||
)
|
)
|
||||||
parser.add_argument(
|
parser.add_argument(
|
||||||
"--context-length",
|
"--context-length",
|
||||||
type=int,
|
type=human_readable_int,
|
||||||
default=ServerArgs.context_length,
|
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(
|
parser.add_argument(
|
||||||
"--is-embedding",
|
"--is-embedding",
|
||||||
@@ -3613,10 +3615,11 @@ class ServerArgs:
|
|||||||
)
|
)
|
||||||
parser.add_argument(
|
parser.add_argument(
|
||||||
"--max-total-tokens",
|
"--max-total-tokens",
|
||||||
type=int,
|
type=human_readable_int,
|
||||||
default=ServerArgs.max_total_tokens,
|
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. "
|
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(
|
parser.add_argument(
|
||||||
"--chunked-prefill-size",
|
"--chunked-prefill-size",
|
||||||
@@ -3638,9 +3641,10 @@ class ServerArgs:
|
|||||||
)
|
)
|
||||||
parser.add_argument(
|
parser.add_argument(
|
||||||
"--max-prefill-tokens",
|
"--max-prefill-tokens",
|
||||||
type=int,
|
type=human_readable_int,
|
||||||
default=ServerArgs.max_prefill_tokens,
|
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(
|
parser.add_argument(
|
||||||
"--schedule-policy",
|
"--schedule-policy",
|
||||||
|
|||||||
@@ -47,6 +47,7 @@ import warnings
|
|||||||
from collections import OrderedDict, defaultdict
|
from collections import OrderedDict, defaultdict
|
||||||
from contextlib import contextmanager
|
from contextlib import contextmanager
|
||||||
from dataclasses import dataclass
|
from dataclasses import dataclass
|
||||||
|
from decimal import Decimal
|
||||||
from functools import lru_cache, partial
|
from functools import lru_cache, partial
|
||||||
from importlib.metadata import PackageNotFoundError, version
|
from importlib.metadata import PackageNotFoundError, version
|
||||||
from importlib.util import find_spec
|
from importlib.util import find_spec
|
||||||
@@ -2308,6 +2309,44 @@ def nullable_str(val: str):
|
|||||||
return val
|
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():
|
def pyspy_dump_schedulers():
|
||||||
"""py-spy dump on all scheduler in a local node."""
|
"""py-spy dump on all scheduler in a local node."""
|
||||||
try:
|
try:
|
||||||
|
|||||||
Reference in New Issue
Block a user