From b21db86e2f3ba76a51d8f14c711569d5fc31e5d9 Mon Sep 17 00:00:00 2001 From: Ke Bao Date: Fri, 3 Apr 2026 00:06:31 +0800 Subject: [PATCH] [CI] Fix gpu deps import in cpu test (#21950) --- .claude/skills/write-sglang-test/SKILL.md | 2 +- python/sglang/test/test_utils.py | 38 +++++++++++++++++++ test/registered/unit/README.md | 28 +++++++++++++- .../managers/test_scheduler_flush_cache.py | 8 +++- .../test_scheduler_pause_generation.py | 8 +++- 5 files changed, 78 insertions(+), 6 deletions(-) diff --git a/.claude/skills/write-sglang-test/SKILL.md b/.claude/skills/write-sglang-test/SKILL.md index af547ffa3..93bf5b786 100644 --- a/.claude/skills/write-sglang-test/SKILL.md +++ b/.claude/skills/write-sglang-test/SKILL.md @@ -172,7 +172,7 @@ if __name__ == "__main__": unittest.main() ``` -Use `unittest.mock.patch` / `MagicMock` to mock dependencies and isolate the logic under test. If the module fails to import on CPU CI (e.g., imports `torch` or CUDA ops at module level), use `sys.modules` stubs to make the import succeed. See existing tests in `test/registered/unit/` for examples. +Use `unittest.mock.patch` / `MagicMock` to mock dependencies and isolate the logic under test. If the module transitively imports GPU-only packages (e.g. `sgl_kernel`), they can be stubbed so the test runs on CPU CI. See `test/registered/unit/README.md` for details and examples. **Quality bar** — test real logic (validation boundaries, state transitions, error paths, branching, etc.). Skip tests that just verify Python itself works (e.g., "does calling an abstract method raise `NotImplementedError`?", "does a dataclass store the field I assigned?"). Consolidate repetitive patterns into parameterized tests. No production code changes in test PRs. diff --git a/python/sglang/test/test_utils.py b/python/sglang/test/test_utils.py index 6c6c29be7..a360358c7 100644 --- a/python/sglang/test/test_utils.py +++ b/python/sglang/test/test_utils.py @@ -2056,6 +2056,44 @@ def _distributed_worker(rank, world_size, backend, port, func, result_queue, kwa dist.destroy_process_group() +def maybe_stub_sgl_kernel(): + """Stub sgl_kernel if it cannot be imported (e.g. no GPU). + + Must be called before any import that transitively depends on sgl_kernel. + On machines with a working sgl_kernel this is a no-op. + """ + try: + import sgl_kernel # noqa: F401 + + return + except (ImportError, OSError): + pass + + import importlib.abc + import importlib.machinery + + class _SglKernelLoader(importlib.abc.Loader): + def create_module(self, spec): + return None + + def exec_module(self, module): + from unittest.mock import MagicMock + + module.__getattr__ = lambda name: MagicMock() + + class _SglKernelFinder(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path, target=None): + if fullname == "sgl_kernel" or fullname.startswith("sgl_kernel."): + return importlib.machinery.ModuleSpec( + fullname, + _SglKernelLoader(), + is_package=True, + ) + return None + + sys.meta_path.insert(0, _SglKernelFinder()) + + class CustomTestCase(unittest.TestCase): def __init_subclass__(cls, **kwargs): diff --git a/test/registered/unit/README.md b/test/registered/unit/README.md index 7d5b36fd0..f5023b431 100644 --- a/test/registered/unit/README.md +++ b/test/registered/unit/README.md @@ -32,7 +32,9 @@ Tests can use CPU or GPU — the key criterion is **no server process**. diff-cover coverage.xml --compare-branch=origin/main --fail-under=60 ``` -## Example +## Examples + +### Basic unit test ```python """Unit tests for — no server, no model loading.""" @@ -57,6 +59,30 @@ if __name__ == "__main__": unittest.main() ``` +### Stubbing GPU-only imports for CPU tests + +Some modules (e.g. `scheduler.py`, `io_struct.py`) transitively import packages like +`sgl_kernel` that require a GPU to initialize. To run pure-mock tests against these +modules on CPU-only CI, stub the problematic package **before** importing it. + +`maybe_stub_sgl_kernel()` in `test_utils.py` does this for `sgl_kernel`: it's a no-op +on GPU machines, and on CPU it installs a `sys.meta_path` finder that auto-creates empty +stub modules for all `sgl_kernel.*` submodules. + +```python +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import maybe_stub_sgl_kernel + +maybe_stub_sgl_kernel() # must precede any import that pulls in sgl_kernel + +from sglang.srt.managers.io_struct import FlushCacheReqInput +from sglang.srt.managers.scheduler import Scheduler + +register_cpu_ci(est_time=2, suite="stage-a-test-cpu") +``` + +The same pattern can be applied to other GPU-only packages: try importing the real package, and if it fails, register a `sys.meta_path` finder that stubs it. See `maybe_stub_sgl_kernel()` in `python/sglang/test/test_utils.py` for the implementation. + ## Rules - **No** `popen_launch_server()` or `Engine(...)`. diff --git a/test/registered/unit/managers/test_scheduler_flush_cache.py b/test/registered/unit/managers/test_scheduler_flush_cache.py index 76c610ea4..828854ed4 100644 --- a/test/registered/unit/managers/test_scheduler_flush_cache.py +++ b/test/registered/unit/managers/test_scheduler_flush_cache.py @@ -1,11 +1,15 @@ import unittest from unittest.mock import MagicMock, patch +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import maybe_stub_sgl_kernel + +maybe_stub_sgl_kernel() + from sglang.srt.managers.io_struct import FlushCacheReqInput from sglang.srt.managers.scheduler import Scheduler -from sglang.test.ci.ci_register import register_cpu_ci -register_cpu_ci(est_time=2, suite="stage-a-cpu-only") +register_cpu_ci(est_time=2, suite="stage-a-test-cpu") class TestSchedulerFlushCache(unittest.TestCase): diff --git a/test/registered/unit/managers/test_scheduler_pause_generation.py b/test/registered/unit/managers/test_scheduler_pause_generation.py index ee5813d30..210ba0aa6 100644 --- a/test/registered/unit/managers/test_scheduler_pause_generation.py +++ b/test/registered/unit/managers/test_scheduler_pause_generation.py @@ -2,11 +2,15 @@ import unittest from collections import deque from unittest.mock import MagicMock +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import maybe_stub_sgl_kernel + +maybe_stub_sgl_kernel() + from sglang.srt.managers.io_struct import PauseGenerationReqInput from sglang.srt.managers.scheduler import Scheduler -from sglang.test.ci.ci_register import register_cpu_ci -register_cpu_ci(est_time=2, suite="stage-a-cpu-only") +register_cpu_ci(est_time=2, suite="stage-a-test-cpu") class TestSchedulerPauseGeneration(unittest.TestCase):