[CI] CI registry update (#13927)

Co-authored-by: alisonshao <54658187+alisonshao@users.noreply.github.com>
This commit is contained in:
Liangsheng Yin
2025-11-25 22:43:19 +08:00
committed by GitHub
parent 64225a8ae9
commit d3d404d3d7
2 changed files with 127 additions and 29 deletions

View File

@@ -2,7 +2,20 @@ import ast
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):
@@ -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 = {
@@ -43,35 +70,96 @@ class RegistryVisitor(ast.NodeVisitor):
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):