From 5ffdb02d0c6818f26e08408f79c668b7965e76c0 Mon Sep 17 00:00:00 2001 From: Alex Nails Date: Mon, 24 Aug 2026 18:22:43 -0700 Subject: [PATCH] [CI] Cut repeated tokenizer loads, serial subprocesses and a double scan (#36241) Co-authored-by: Claude Opus 5 (1M context) --- .../bench_fn/test_benchmark_datasets_api.py | 122 ++++++++---------- .../ops/diffusion/test_import_surface.py | 98 +++++++------- .../test_function_call_parser.py | 18 ++- 3 files changed, 118 insertions(+), 120 deletions(-) diff --git a/test/registered/bench_fn/test_benchmark_datasets_api.py b/test/registered/bench_fn/test_benchmark_datasets_api.py index ac7671f74..9cc5d23b4 100644 --- a/test/registered/bench_fn/test_benchmark_datasets_api.py +++ b/test/registered/bench_fn/test_benchmark_datasets_api.py @@ -1,6 +1,7 @@ import argparse import asyncio import base64 +import functools import io import json import pickle @@ -59,10 +60,58 @@ from sglang.benchmark.serving import ( from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase -register_cpu_ci(est_time=40, suite="base-a-test-cpu") +register_cpu_ci(est_time=30, suite="base-a-test-cpu") register_cpu_ci(est_time=46, suite="base-c-test-cpu") +_BENCH_SERVING_CLI_CASES = { + "help": ["--help"], + "invalid_distribution": [ + "--dataset-name", + "generated-shared-prefix", + "--gsp-group-distribution", + "invalid_name", + ], + "flush_cache_timeout": ["--flush-cache-timeout", "inf"], + "zipf_without_alpha": [ + "--dataset-name", + "generated-shared-prefix", + "--gsp-group-distribution", + "zipf", + "--ready-check-timeout-sec", + "0", + ], + "uniform_with_alpha": [ + "--dataset-name", + "generated-shared-prefix", + "--gsp-group-distribution", + "uniform", + "--gsp-zipf-alpha", + "1.0", + "--ready-check-timeout-sec", + "0", + ], +} + + +@functools.lru_cache(maxsize=1) +def _bench_serving_cli_results(): + def run(args): + return subprocess.run( + [sys.executable, "-m", "sglang.benchmark.serving", *args], + capture_output=True, + text=True, + timeout=180, + ) + + with ThreadPoolExecutor(max_workers=len(_BENCH_SERVING_CLI_CASES)) as pool: + futures = { + name: pool.submit(run, args) + for name, args in _BENCH_SERVING_CLI_CASES.items() + } + return {name: future.result() for name, future in futures.items()} + + class _DummyTokenTensor: def __init__(self, value: int): self.value = value @@ -1363,12 +1412,7 @@ class TestBenchmarkDatasetsAPI(CustomTestCase): # Subprocess-driven coverage of the live CLI: --help advertises both # flags with the rank-based Zipf formula and the alpha constraint, # and argparse rejects an unknown distribution choice. - help_res = subprocess.run( - [sys.executable, "-m", "sglang.benchmark.serving", "--help"], - capture_output=True, - text=True, - timeout=90, - ) + help_res = _bench_serving_cli_results()["help"] self.assertEqual(help_res.returncode, 0, help_res.stderr) out = help_res.stdout # Both new flags appear. @@ -1380,37 +1424,13 @@ class TestBenchmarkDatasetsAPI(CustomTestCase): self.assertIn("finite float", out) # Argparse rejects unknown distribution choice. - bad_choice_res = subprocess.run( - [ - sys.executable, - "-m", - "sglang.benchmark.serving", - "--dataset-name", - "generated-shared-prefix", - "--gsp-group-distribution", - "invalid_name", - ], - capture_output=True, - text=True, - timeout=90, - ) + bad_choice_res = _bench_serving_cli_results()["invalid_distribution"] self.assertNotEqual(bad_choice_res.returncode, 0) self.assertIn("invalid choice", (bad_choice_res.stderr + bad_choice_res.stdout)) def test_serving_benchmark_cli_rejects_invalid_flush_cache_timeout(self): """Invalid timeouts fail before the benchmark contacts a server or hangs.""" - res = subprocess.run( - [ - sys.executable, - "-m", - "sglang.benchmark.serving", - "--flush-cache-timeout", - "inf", - ], - capture_output=True, - text=True, - timeout=90, - ) + res = _bench_serving_cli_results()["flush_cache_timeout"] self.assertEqual(res.returncode, 2, res.stderr) self.assertIn("expected a finite float > 0", res.stderr) @@ -1423,22 +1443,7 @@ class TestBenchmarkDatasetsAPI(CustomTestCase): # Malformed CLI combinations (zipf with no alpha) must fail at # argparse time so users see the GSP-flag error directly, not a # downstream connection or model-fetch failure. - res = subprocess.run( - [ - sys.executable, - "-m", - "sglang.benchmark.serving", - "--dataset-name", - "generated-shared-prefix", - "--gsp-group-distribution", - "zipf", - "--ready-check-timeout-sec", - "0", - ], - capture_output=True, - text=True, - timeout=90, - ) + res = _bench_serving_cli_results()["zipf_without_alpha"] # parser.error() exits with code 2 (argparse convention). self.assertEqual(res.returncode, 2, res.stderr) stderr = res.stderr + res.stdout @@ -1458,24 +1463,7 @@ class TestBenchmarkDatasetsAPI(CustomTestCase): def test_bench_serving_cli_rejects_uniform_with_alpha_before_server(self): # The complementary malformation: uniform distribution with an # explicit alpha value. Must also fail at argparse time. - res = subprocess.run( - [ - sys.executable, - "-m", - "sglang.benchmark.serving", - "--dataset-name", - "generated-shared-prefix", - "--gsp-group-distribution", - "uniform", - "--gsp-zipf-alpha", - "1.0", - "--ready-check-timeout-sec", - "0", - ], - capture_output=True, - text=True, - timeout=90, - ) + res = _bench_serving_cli_results()["uniform_with_alpha"] self.assertEqual(res.returncode, 2, res.stderr) stderr = res.stderr + res.stdout self.assertIn("--gsp-group-distribution", stderr) diff --git a/test/registered/kernels/ops/diffusion/test_import_surface.py b/test/registered/kernels/ops/diffusion/test_import_surface.py index 9eef11a19..67af3116a 100644 --- a/test/registered/kernels/ops/diffusion/test_import_surface.py +++ b/test/registered/kernels/ops/diffusion/test_import_surface.py @@ -14,6 +14,7 @@ resolve them with ``importlib``/``ast`` without importing torch backends. """ import ast +import functools import importlib import pathlib import subprocess @@ -25,7 +26,7 @@ from sglang.kernels.ops.diffusion import _EXPORTS, _SPECS from sglang.kernels.registry import registry from sglang.test.ci.ci_register import register_cpu_ci -register_cpu_ci(est_time=4, suite="base-a-test-cpu") +register_cpu_ci(est_time=25, suite="base-a-test-cpu") PACKAGE = "sglang.kernels.ops.diffusion" _PACKAGE_DIR = pathlib.Path(importlib.import_module(PACKAGE).__file__ or "").parent @@ -74,6 +75,52 @@ def _module_defines(module_path: str) -> set[str]: return names +@functools.lru_cache(maxsize=None) +def _scan_root(root: str) -> tuple[frozenset[str], tuple[str, ...]]: + unexported: set[str] = set() + offenders: list[str] = [] + root_dir = _REPO_ROOT / root + if not root_dir.exists(): + return frozenset(), () + + for path in root_dir.rglob("*.py"): + rel = path.relative_to(_REPO_ROOT).as_posix() + if rel.startswith("python/sglang/kernels/ops/diffusion/"): + continue + try: + source = path.read_text(encoding="utf-8") + except UnicodeDecodeError: + continue + if PACKAGE not in source: + continue + try: + tree = ast.parse(source) + except SyntaxError: + continue + allowlisted = rel in _DEEP_IMPORT_ALLOWLIST + for node in ast.walk(tree): + if isinstance(node, ast.ImportFrom): + if node.module == PACKAGE: + unexported.update( + a.name + for a in node.names + if a.name not in _EXPORTS and not a.name.startswith("_") + ) + elif ( + not allowlisted + and node.module + and node.module.startswith(f"{PACKAGE}.") + ): + offenders.append(f"{rel}:{node.lineno} imports {node.module}") + elif isinstance(node, ast.Import) and not allowlisted: + offenders.extend( + f"{rel}:{node.lineno} imports {a.name}" + for a in node.names + if a.name.startswith(f"{PACKAGE}.") + ) + return frozenset(unexported), tuple(offenders) + + def test_every_export_resolves_to_a_real_symbol(): missing = [ f"{symbol} -> {module}" @@ -92,26 +139,9 @@ def test_every_symbol_imported_from_the_facade_is_exported(): or code path runs, on the platform that has the backend. Enumerating the call sites catches it here instead. """ - unexported = set() + unexported: set[str] = set() for root in ("python/sglang", "test", "benchmark"): - root_dir = _REPO_ROOT / root - if not root_dir.exists(): - continue - for path in root_dir.rglob("*.py"): - rel = path.relative_to(_REPO_ROOT).as_posix() - if rel.startswith("python/sglang/kernels/ops/diffusion/"): - continue - try: - tree = ast.parse(path.read_text(encoding="utf-8")) - except (SyntaxError, UnicodeDecodeError): - continue - for node in ast.walk(tree): - if isinstance(node, ast.ImportFrom) and node.module == PACKAGE: - unexported.update( - a.name - for a in node.names - if a.name not in _EXPORTS and not a.name.startswith("_") - ) + unexported.update(_scan_root(root)[0]) assert not unexported, f"imported but not in _EXPORTS: {sorted(unexported)}" @@ -174,34 +204,10 @@ def test_importing_the_package_does_not_import_any_leaf_module(): @pytest.mark.parametrize("root", ["python/sglang", "test", "benchmark"]) def test_runtime_code_imports_only_through_the_facade(root): - root_dir = _REPO_ROOT / root - if not root_dir.exists(): # source checkouts only + if not (_REPO_ROOT / root).exists(): # source checkouts only pytest.skip(f"{root} not present in this install") - offenders = [] - for path in root_dir.rglob("*.py"): - rel = path.relative_to(_REPO_ROOT).as_posix() - if rel.startswith("python/sglang/kernels/ops/diffusion/"): - continue # intra-package imports are the point of the subpackages - if rel in _DEEP_IMPORT_ALLOWLIST: - continue - try: - tree = ast.parse(path.read_text(encoding="utf-8")) - except (SyntaxError, UnicodeDecodeError): - continue - for node in ast.walk(tree): - if ( - isinstance(node, ast.ImportFrom) - and node.module - and node.module.startswith(f"{PACKAGE}.") - ): - offenders.append(f"{rel}:{node.lineno} imports {node.module}") - elif isinstance(node, ast.Import): - offenders.extend( - f"{rel}:{node.lineno} imports {a.name}" - for a in node.names - if a.name.startswith(f"{PACKAGE}.") - ) + offenders = _scan_root(root)[1] assert not offenders, ( "import from sglang.kernels.ops.diffusion instead of a submodule:\n " + "\n ".join(offenders) diff --git a/test/registered/unit/function_call/test_function_call_parser.py b/test/registered/unit/function_call/test_function_call_parser.py index 8fcd0eb59..9c683d25e 100644 --- a/test/registered/unit/function_call/test_function_call_parser.py +++ b/test/registered/unit/function_call/test_function_call_parser.py @@ -1,3 +1,4 @@ +import functools import json import unittest import warnings @@ -33,10 +34,17 @@ from sglang.srt.function_call.pythonic_detector import PythonicDetector from sglang.srt.function_call.qwen3_coder_detector import Qwen3CoderDetector from sglang.test.ci.ci_register import register_cpu_ci -register_cpu_ci(est_time=15, suite="base-a-test-cpu") +register_cpu_ci(est_time=20, suite="base-a-test-cpu") register_cpu_ci(est_time=70, suite="base-c-test-cpu") +@functools.lru_cache(maxsize=None) +def _shared_tokenizer(path: str): + from sglang.srt.utils.hf_transformers_utils import get_tokenizer + + return get_tokenizer(path) + + class TestInklingDetector(unittest.TestCase): def setUp(self): self.tools = [ @@ -1599,9 +1607,7 @@ class TestDeepSeekV32Detector(unittest.TestCase): ), ] self.detector = DeepSeekV32Detector() - from sglang.srt.utils.hf_transformers_utils import get_tokenizer - - self.tokenizer = get_tokenizer("deepseek-ai/DeepSeek-V3.2") + self.tokenizer = _shared_tokenizer("deepseek-ai/DeepSeek-V3.2") self.interval = 1 def test_detect_and_parse_xml_format(self): @@ -2049,9 +2055,7 @@ class TestDeepSeekV4Detector(unittest.TestCase): ), ] self.detector = DeepSeekV4Detector() - from sglang.srt.utils.hf_transformers_utils import get_tokenizer - - self.tokenizer = get_tokenizer("deepseek-ai/DeepSeek-V3.2") + self.tokenizer = _shared_tokenizer("deepseek-ai/DeepSeek-V3.2") self.interval = 1 def test_detect_and_parse_xml_format(self):