Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
146 changes: 117 additions & 29 deletions python/sglang/test/ci/ci_register.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,20 @@
import warnings
from dataclasses import dataclass
from enum import Enum, auto
from typing import List
from typing import List, Optional

__all__ = [
"HWBackend",
"CIRegistry",
"collect_tests",
"register_cpu_ci",
"register_cuda_ci",
"register_amd_ci",
"ut_parse_one_file",
]

_PARAM_ORDER = ("est_time", "suite", "nightly", "disabled")
_UNSET = object()


class HWBackend(Enum):
Expand All @@ -17,18 +30,32 @@ class CIRegistry:
filename: str
est_time: float
suite: str
nightly: bool = False
disabled: Optional[str] = None # None = enabled, string = disabled with reason


def register_cpu_ci(est_time: float, suite: str):
pass
def register_cpu_ci(
est_time: float, suite: str, nightly: bool = False, disabled: Optional[str] = None
):
"""Marker for CPU CI registration (parsed via AST; runtime no-op)."""
return None


def register_cuda_ci(est_time: float, suite: str):
pass
def register_cuda_ci(
est_time: float, suite: str, nightly: bool = False, disabled: Optional[str] = None
):
"""Marker for CUDA CI registration (parsed via AST; runtime no-op)."""
return None


def register_amd_ci(est_time: float, suite: str):
pass
def register_amd_ci(
est_time: float,
suite: str,
nightly: bool = False,
disabled: Optional[str] = None,
):
"""Marker for AMD CI registration (parsed via AST; runtime no-op)."""
return None


REGISTER_MAPPING = {
Expand All @@ -43,35 +70,96 @@ def __init__(self, filename: str):
self.filename = filename
self.registries: list[CIRegistry] = []

def _constant_value(self, node: ast.AST) -> object:
if isinstance(node, ast.Constant):
return node.value
return _UNSET

def _parse_call_args(
self, func_call: ast.Call
) -> tuple[float, str, bool, Optional[str]]:
args = {name: _UNSET for name in _PARAM_ORDER}
seen = set()

if any(isinstance(arg, ast.Starred) for arg in func_call.args):
raise ValueError(
f"{self.filename}: starred arguments are not supported in {func_call.func.id}()"
)
if len(func_call.args) > len(_PARAM_ORDER):
raise ValueError(
f"{self.filename}: too many positional arguments in {func_call.func.id}()"
)

for name, arg in zip(_PARAM_ORDER, func_call.args):
seen.add(name)
args[name] = self._constant_value(arg)

for kw in func_call.keywords:
if kw.arg is None:
raise ValueError(
f"{self.filename}: **kwargs are not supported in {func_call.func.id}()"
)
if kw.arg not in args:
raise ValueError(
f"{self.filename}: unknown argument '{kw.arg}' in {func_call.func.id}()"
)
if kw.arg in seen:
raise ValueError(
f"{self.filename}: duplicated argument '{kw.arg}' in {func_call.func.id}()"
)
seen.add(kw.arg)
args[kw.arg] = self._constant_value(kw.value)

if args["est_time"] is _UNSET or args["suite"] is _UNSET:
raise ValueError(
f"{self.filename}: est_time and suite are required constants in {func_call.func.id}()"
)

est_time, suite = args["est_time"], args["suite"]
nightly_value = args["nightly"]

if not isinstance(est_time, (int, float)):
raise ValueError(
f"{self.filename}: est_time must be a number in {func_call.func.id}()"
)
if not isinstance(suite, str):
raise ValueError(
f"{self.filename}: suite must be a string in {func_call.func.id}()"
)

if nightly_value is _UNSET:
nightly = False
elif isinstance(nightly_value, bool):
nightly = nightly_value
else:
raise ValueError(
f"{self.filename}: nightly must be a boolean in {func_call.func.id}()"
)

disabled = args["disabled"] if args["disabled"] is not _UNSET else None
if disabled is not None and not isinstance(disabled, str):
raise ValueError(
f"{self.filename}: disabled must be a string in {func_call.func.id}()"
)

return float(est_time), suite, nightly, disabled

def _collect_ci_registry(self, func_call: ast.Call):
if not isinstance(func_call.func, ast.Name):
return None

if func_call.func.id not in REGISTER_MAPPING:
backend = REGISTER_MAPPING.get(func_call.func.id)
if backend is None:
return None

hw = REGISTER_MAPPING[func_call.func.id]
est_time, suite = None, None
for kw in func_call.keywords:
if kw.arg == "est_time":
if isinstance(kw.value, ast.Constant):
est_time = kw.value.value
elif kw.arg == "suite":
if isinstance(kw.value, ast.Constant):
suite = kw.value.value

for i, arg in enumerate(func_call.args):
if isinstance(arg, ast.Constant):
if i == 0:
est_time = arg.value
elif i == 1:
suite = arg.value
assert (
est_time is not None
), "esimation_time is required and should be a constant"
assert suite is not None, "suite is required and should be a constant"
est_time, suite, nightly, disabled = self._parse_call_args(func_call)
return CIRegistry(
backend=hw, filename=self.filename, est_time=est_time, suite=suite
backend=backend,
filename=self.filename,
est_time=est_time,
suite=suite,
nightly=nightly,
disabled=disabled,
)

def visit_Module(self, node):
Expand Down
10 changes: 10 additions & 0 deletions test/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,16 @@ python3 test_choices.py
- Ensure you added `unittest.main()` for unittest and `pytest.main([__file__])` for pytest in the scripts. The CI run them via `python3 test_file.py`.
- The CI will run some suites such as `per-commit-1-gpu`, `per-commit-2-gpu`, and `nightly-1-gpu` automatically. If you need special setup or custom test groups, you may modify the workflows in [`.github/workflows/`](https://github.com/sgl-project/sglang/tree/main/.github/workflows).

## CI Registry Quick Peek

Tests in `test/registered/` declare CI metadata via lightweight markers:

```python
from sglang.test.ci.ci_register import register_cuda_ci

register_cuda_ci(est_time=80, suite="stage-a-test-1")
```

## Writing Elegant Test Cases

- Learn from existing examples in [sglang/test/srt](https://github.com/sgl-project/sglang/tree/main/test/srt).
Expand Down
Loading