[Test] Consolidate test cleanup and CI taxonomy (net -11.4K lines) (#37436)

Co-authored-by: Mick Qian <mickqian@users.noreply.github.com>
This commit is contained in:
Xiaoyu Zhang
2026-09-07 15:13:59 +08:00
committed by GitHub
co-authored by Mick Qian
parent 6a1ff90f2d
commit 4d23a4fa6d
199 changed files with 1185 additions and 11812 deletions
+9
View File
@@ -1,6 +1,7 @@
---
paths:
- "test/**/*.py"
- "python/sglang/multimodal_gen/test/**/*.py"
---
# Unit Test Admission Criteria
@@ -68,5 +69,13 @@ One strong case beats several weak ones: each additional case must guard a
distinct failure mode. Ask "which bug escapes if I delete this case?" -- no
answer means delete it.
New cases join an existing file in the same subsystem by default. Create a new
file only when it needs a different fixture, dependency, owner, or CI contract;
every file pays a separate interpreter-import cost in the CPU gate.
Suite cadence is part of admission: if a failing run cannot be attributed to a
single PR's diff, the test belongs in a nightly or weekly suite rather than a
per-commit lane.
Test mechanics (placement, CI registration, fixtures) live in
[`write-sglang-test`](../skills/write-sglang-test/SKILL.md).
+29 -25
View File
@@ -11,14 +11,14 @@ This skill covers **how to write and register tests**. For CI pipeline internals
1. **Always use `CustomTestCase`** — never raw `unittest.TestCase`. It ensures `tearDownClass` runs even when `setUpClass` fails, preventing resource leaks in CI.
2. **`tearDownClass` must be defensive** — use `hasattr`/null checks before accessing resources (e.g. `cls.process`) that `setUpClass` may not have finished allocating.
3. **Place tests in `test/registered/<category>/`** — including JIT kernel tests and benchmarks, which live in `test/registered/jit/` and `test/registered/jit/benchmark/` (nested subfolders are allowed)
3. **Place tests in `test/registered/<kind>/<subsystem>/`**`<kind>` is `unit`, `kernel`, `e2e`, `accuracy`, `perf`, or `stress`; hardware belongs in registrations, not directory names
4. **Reuse server fixtures** — inherit from `DefaultServerBase` or write `setUpClass`/`tearDownClass` with `popen_launch_server`
5. **Prefer mock over real server** — when testing logic that doesn't need a server / engine launch (middleware, request routing, config validation, argument parsing), use `unittest.mock.patch` / `MagicMock` and place tests in `test/registered/unit/`. Only launch a real server when the test genuinely needs inference results or server lifecycle behavior.
5. **Mock boundaries, not SGLang behavior** — mock slow or external dependencies only when the assertion still checks an observable result, state transition, or error. A test whose evidence is only `assert_called*` mirrors its mock and is not admissible. Launch a real server only when inference results or lifecycle behavior are the contract under test.
JIT kernel notes:
- If the task is adding or updating code under `python/sglang/kernels/jit/`, prefer the `add-jit-kernel` skill first.
- JIT kernel correctness tests use `test/registered/jit/**/test_*.py`.
- JIT kernel benchmarks use `test/registered/jit/benchmark/**/bench_*.py`.
- New JIT kernel correctness tests use `test/registered/kernel/jit/**/test_*.py`.
- New JIT kernel benchmarks use `test/registered/kernel/jit/benchmark/**/bench_*.py`.
- Those files are executed by `test/run_suite.py` through dedicated kernel suites (`base-b-kernel-*`); a `register_*_ci(...)` call placed under `python/sglang/` is rejected by the `check-no-registered-tests-in-package` pre-commit hook.
---
@@ -27,7 +27,7 @@ JIT kernel notes:
| Scenario | Model | CI Registration | Suite |
|----------|-------|-----------------|-------|
| **Unit tests** (no server / engine launch) | None | `register_cpu_ci` (prefer) or `register_cuda_ci` | `base-a-test-cpu` or `base-b-test-1-gpu-small` |
| **Unit tests** (no server / engine launch) | None | `register_cpu_ci` | `base-a-test-cpu` |
| **Common / backend-independent** (middleware, abort, routing, config, arg parsing) | `DEFAULT_SMALL_MODEL_NAME_FOR_TEST` (1B) | `register_cuda_ci` only | `base-b-test-1-gpu-small` |
| **Model-agnostic functionality** (sampling, session, OpenAI API features) | `DEFAULT_SMALL_MODEL_NAME_FOR_TEST` (1B) | `register_cuda_ci` (+ AMD if relevant) | `base-b-test-1-gpu-small` |
| **General performance** (single node, no spec/DP/parallelism) | `DEFAULT_MODEL_NAME_FOR_TEST` (8B) | `register_cuda_ci` | `base-b-test-1-gpu-large` |
@@ -70,10 +70,10 @@ A per-commit suite name is **generated** from registration metadata as `{stage}-
| `base-b-test-1-gpu-large` | `1-gpu-h100` | Tests that need H100-class memory or kernels (e.g. FA3) |
| `base-b-test-2-gpu-large` | `2-gpu-h100` | Two-GPU correctness and parallelism (TP/PP) on H100 |
| `base-b-test-4-gpu-b200` | `4-gpu-b200` | Early Blackwell coverage (SM100+ paths) on four GPUs |
| `base-b-kernel-unit-test-1-gpu-large` | `1-gpu-h100` | JIT kernel correctness tests under `test/registered/jit/` |
| `base-b-kernel-unit-test-1-gpu-large` | `1-gpu-h100` | JIT kernel correctness tests under `test/registered/kernel/jit/` |
| `base-b-kernel-unit-test-4-gpu-b200` | `4-gpu-b200` | JIT kernel correctness tests for Blackwell / SM100-specific paths |
| `base-b-kernel-unit-test-8-gpu-h200` | `8-gpu-h200` | Multi-GPU JIT kernel correctness tests under `test/registered/jit/` |
| `base-b-kernel-benchmark-test-1-gpu-large` | `1-gpu-h100` | JIT kernel benchmark files under `test/registered/jit/benchmark/` |
| `base-b-kernel-unit-test-8-gpu-h200` | `8-gpu-h200` | Multi-GPU JIT kernel correctness tests under `test/registered/kernel/jit/` |
| `base-b-kernel-benchmark-test-1-gpu-large` | `1-gpu-h100` | JIT kernel benchmark files under `test/registered/kernel/jit/benchmark/` |
| `base-c-test-4-gpu-h100` | `4-gpu-h100` | Large 4-GPU H100 integration and scaling tests |
| `base-c-test-8-gpu-h200` | `8-gpu-h200` | Large 8-GPU H200 runs for big models and parallelism |
| `base-c-test-8-gpu-h20` | `8-gpu-h20` | Large 8-GPU H20 runs for big models |
@@ -162,7 +162,7 @@ from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
# Prefer CPU. Only use register_cuda_ci when the test truly needs a GPU.
# Unit tests are CPU-only. GPU operator tests belong under `kernel/`.
class TestTargetClass(CustomTestCase):
def test_basic_behavior(self):
@@ -180,7 +180,13 @@ if __name__ == "__main__":
unittest.main()
```
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. Do not modify `sys.modules` at module level — use `patch.dict` (as a class decorator or with `start`/`stop`) to ensure cleanup and avoid cross-test pollution. See `test/registered/unit/README.md` for details and examples.
Use `unittest.mock.patch` / `MagicMock` only at dependency boundaries. Assert the
resulting value, state, protocol output, or error—not merely that the mock was
called. If the module transitively imports GPU-only packages (e.g. `sgl_kernel`),
they can be stubbed so the test runs on CPU CI. Do not modify `sys.modules` at
module level—use `patch.dict` (as a class decorator or with `start`/`stop`) to
ensure cleanup and avoid cross-test pollution. 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.
@@ -380,15 +386,12 @@ Every call generates a suite named `{stage}-test-{runner_config}`, e.g. `base-b-
```
test/
├── registered/ # CI tests (auto-discovered by run_suite.py)
│ ├── unit/ # No server / engine launch (see test/registered/unit/README.md)
│ ├── kernels/ # CUDA kernel correctness (no server, GPU required)
│ ├── sampling/ # test_penalty.py, test_sampling_params.py ...
│ ├── sessions/ # test_session_control.py ...
│ ├── openai_server/ # basic/, features/, validation/ ...
── spec/ # eagle/, utils/ ...
│ ├── models/ # model-specific accuracy tests
│ ├── perf/ # performance benchmarks
│ └── <category>/ # create new category if needed
│ ├── unit/<subsystem>/ # CPU-only; no server or model weights
│ ├── kernel/<group>/ # accelerator operator correctness/benchmarks
│ ├── e2e/<subsystem>/ # engine/server integration
│ ├── accuracy/<family>/ # scheduled eval floors
│ ├── perf/<family>/ # scheduled latency/throughput contracts
── stress/<subsystem>/ # stress/weekly coverage
├── manual/ # Non-CI: debugging, one-off, manual verification
└── run_suite.py # CI runner (scans registered/ plus jit_kernel test/benchmark files)
@@ -398,10 +401,11 @@ python/sglang/kernels/jit/
```
**Decision rule** (see also `test/registered/README.md`):
- Component logic, no server → `registered/unit/`
- JIT kernel correctness / benchmarks → `test/registered/jit/` or `test/registered/jit/benchmark/`
- Other kernel correctness → `registered/kernels/`
- Server needed → `registered/<category>/`
- CPU component logic, no server → `registered/unit/<subsystem>/`
- JIT kernel correctness / benchmarks → `registered/kernel/jit/`
- Other accelerator operator correctness → `registered/kernel/<group>/`
- Server needed → `registered/e2e/<subsystem>/`
- Eval floor / performance contract → `registered/{accuracy,perf}/<family>/`
- Local debugging → `manual/`
---
@@ -443,8 +447,8 @@ Before submitting a test:
- [ ] Inherits from `CustomTestCase` (not `unittest.TestCase`)
- [ ] Has `register_*_ci(...)` call at module level
- [ ] Placed in `test/registered/<category>/` (JIT kernel test/benchmark → `test/registered/jit/` or `test/registered/jit/benchmark/`)
- [ ] JIT kernel work: test files live in `test/registered/jit/`; only test-only helpers stay under `python/sglang/kernels/jit/`
- [ ] Placed in `test/registered/<kind>/<subsystem>/`
- [ ] JIT kernel work: test files live in `test/registered/kernel/jit/`; only test-only helpers stay under `python/sglang/kernels/jit/`
- [ ] Backend-independent tests: `register_cuda_ci` only + smallest model
- [ ] Logic that doesn't need a server / engine launch → unit test in `registered/unit/` (see Unit Tests section)
- [ ] `setUpClass` launches server, `tearDownClass` kills it (if server-based)
+40 -40
View File
@@ -119,17 +119,16 @@ jobs:
timeout-minutes: 240
env:
RUNAI_STREAMER_MEMORY_LIMIT: 0
CONTINUE_ON_ERROR_FLAG: ${{ inputs.continue_on_error == 'true' && '--continue-on-error' || '' }}
PARTITION_PLAN_JSON: ${{ needs.compute-diffusion-partitions.outputs.plan-1gpu }}
DIFFUSION_CONTINUE_ON_ERROR: ${{ inputs.continue_on_error }}
DIFFUSION_PARTITION_ID: ${{ matrix.part }}
DIFFUSION_TOTAL_PARTITIONS: ${{ needs.compute-diffusion-partitions.outputs['partition-count-1gpu'] }}
DIFFUSION_PARTITION_PLAN_JSON: ${{ needs.compute-diffusion-partitions.outputs.plan-1gpu }}
SGLANG_DIFFUSION_ARTIFACT_DIR: ${{ github.workspace }}/diffusion-failures
run: |
cd python
python3 sglang/multimodal_gen/test/run_suite.py \
--suite 1-gpu \
--partition-id ${{ matrix.part }} \
--total-partitions ${{ needs.compute-diffusion-partitions.outputs['partition-count-1gpu'] }} \
--partition-plan-json "$PARTITION_PLAN_JSON" \
$CONTINUE_ON_ERROR_FLAG
python3 test/run_suite.py \
--hw cuda \
--suite base-b-test-diffusion-1-gpu-h100 \
--timeout-per-file 14400
- name: Upload execution report
if: always()
@@ -192,13 +191,13 @@ jobs:
timeout-minutes: 120
env:
RUNAI_STREAMER_MEMORY_LIMIT: 0
CONTINUE_ON_ERROR_FLAG: ${{ inputs.continue_on_error == 'true' && '--continue-on-error' || '' }}
DIFFUSION_CONTINUE_ON_ERROR: ${{ inputs.continue_on_error }}
SGLANG_DIFFUSION_ARTIFACT_DIR: ${{ github.workspace }}/diffusion-failures
run: |
cd python
python3 sglang/multimodal_gen/test/run_suite.py \
--suite 1-gpu-5090 \
$CONTINUE_ON_ERROR_FLAG
python3 test/run_suite.py \
--hw cuda \
--suite base-b-test-diffusion-1-gpu-5090 \
--timeout-per-file 7200
- name: Upload execution report
if: always()
@@ -261,13 +260,13 @@ jobs:
timeout-minutes: 60
env:
RUNAI_STREAMER_MEMORY_LIMIT: 0
CONTINUE_ON_ERROR_FLAG: ${{ inputs.continue_on_error == 'true' && '--continue-on-error' || '' }}
DIFFUSION_CONTINUE_ON_ERROR: ${{ inputs.continue_on_error }}
SGLANG_DIFFUSION_ARTIFACT_DIR: ${{ github.workspace }}/diffusion-bcg-artifacts
run: |
cd python
python3 sglang/multimodal_gen/test/run_suite.py \
--suite bcg-diffusion \
$CONTINUE_ON_ERROR_FLAG
python3 test/run_suite.py \
--hw cuda \
--suite base-b-test-diffusion-bcg-1-gpu-h100 \
--timeout-per-file 3600
- name: Upload BCG diffusion artifacts
if: always()
@@ -329,17 +328,16 @@ jobs:
HF_TOKEN: ${{ secrets.SGLANG_DIFFUSION_CI_HF_TOKEN || secrets.HF_TOKEN }}
HUGGING_FACE_HUB_TOKEN: ${{ secrets.SGLANG_DIFFUSION_CI_HF_TOKEN || secrets.HF_TOKEN }}
RUNAI_STREAMER_MEMORY_LIMIT: 0
CONTINUE_ON_ERROR_FLAG: ${{ inputs.continue_on_error == 'true' && '--continue-on-error' || '' }}
PARTITION_PLAN_JSON: ${{ needs.compute-diffusion-partitions.outputs.plan-2gpu }}
DIFFUSION_CONTINUE_ON_ERROR: ${{ inputs.continue_on_error }}
DIFFUSION_PARTITION_ID: ${{ matrix.part }}
DIFFUSION_TOTAL_PARTITIONS: ${{ needs.compute-diffusion-partitions.outputs['partition-count-2gpu'] }}
DIFFUSION_PARTITION_PLAN_JSON: ${{ needs.compute-diffusion-partitions.outputs.plan-2gpu }}
SGLANG_DIFFUSION_ARTIFACT_DIR: ${{ github.workspace }}/diffusion-failures
run: |
cd python
python3 sglang/multimodal_gen/test/run_suite.py \
--suite 2-gpu \
--partition-id ${{ matrix.part }} \
--total-partitions ${{ needs.compute-diffusion-partitions.outputs['partition-count-2gpu'] }} \
--partition-plan-json "$PARTITION_PLAN_JSON" \
$CONTINUE_ON_ERROR_FLAG
python3 test/run_suite.py \
--hw cuda \
--suite base-b-test-diffusion-2-gpu-h100 \
--timeout-per-file 14400
- name: Upload execution report
if: always()
@@ -402,12 +400,12 @@ jobs:
timeout-minutes: 240
env:
RUNAI_STREAMER_MEMORY_LIMIT: 0
CONTINUE_ON_ERROR_FLAG: ${{ inputs.continue_on_error == 'true' && '--continue-on-error' || '' }}
DIFFUSION_CONTINUE_ON_ERROR: ${{ inputs.continue_on_error }}
run: |
cd python
python3 sglang/multimodal_gen/test/run_suite.py \
--suite component-accuracy \
$CONTINUE_ON_ERROR_FLAG
python3 test/run_suite.py \
--hw cuda \
--suite base-b-test-diffusion-component-2-gpu-h100 \
--timeout-per-file 14400
- uses: ./.github/actions/upload-cuda-coredumps
if: always()
@@ -453,13 +451,13 @@ jobs:
timeout-minutes: 240
env:
RUNAI_STREAMER_MEMORY_LIMIT: 0
CONTINUE_ON_ERROR_FLAG: ${{ inputs.continue_on_error == 'true' && '--continue-on-error' || '' }}
DIFFUSION_CONTINUE_ON_ERROR: ${{ inputs.continue_on_error }}
SGLANG_DIFFUSION_ARTIFACT_DIR: ${{ github.workspace }}/diffusion-failures
run: |
cd python
python3 sglang/multimodal_gen/test/run_suite.py \
--suite 1-gpu-b200 \
$CONTINUE_ON_ERROR_FLAG
python3 test/run_suite.py \
--hw cuda \
--suite base-b-test-diffusion-1-gpu-b200 \
--timeout-per-file 14400
- name: Upload diffusion failure artifacts
if: always()
@@ -519,8 +517,10 @@ jobs:
- name: Run diffusion unit tests
timeout-minutes: 60
run: |
cd python
python3 sglang/multimodal_gen/test/run_suite.py --suite unit
python3 test/run_suite.py \
--hw cuda \
--suite base-b-test-diffusion-unit-1-gpu-h100 \
--timeout-per-file 3600
diffusion-coverage-check:
needs: [multimodal-gen-test-1-gpu, multimodal-gen-test-2-gpu]
@@ -17,7 +17,7 @@ from sglang.kernels.ops.diffusion import fused_rmsnorm_scale_shift_bitexact
```
**Import from the package, never from a submodule.** The internal layout is
free to move; the facade is not. `test_import_surface.py` enforces this, with
free to move; the facade is not. Callers should use the facade, with
a small allowlist for tests that deliberately exercise one backend.
Resolution is lazy (PEP 562): the backends have disjoint heavy dependencies
@@ -189,7 +189,7 @@ inspecting model modules is its whole job.
generated by the KDA workflow in `sglang.kernels.kda_kernels`, together
with its source revision and any JIT CUDA source files.
2. Export it from `__init__.py` (`_EXPORTS`) and register a `KernelSpec`
(`_SPECS`)`test_import_surface.py` checks both resolve.
(`_SPECS`).
3. Give it a `can_use_*` predicate; raise, don't return `None`.
4. State the numerical contract in the module docstring, including which
shapes it was verified on.
@@ -5,8 +5,8 @@ This module is the **only** supported import surface for these kernels::
from sglang.kernels.ops.diffusion import fused_rmsnorm_scale_shift_bitexact
Importing a submodule directly (``...diffusion.norm.norm_triton``) couples the
caller to the file layout; ``test_import_surface.py`` guards against it. The
one exception is a test that deliberately exercises a single backend.
caller to the file layout. The one exception is a test that deliberately
exercises a single backend.
Layout -- ordinary implementations use one subpackage per **operator domain**
(``norm``, ``modulate``, ``rope``, ``activation``, ``attention``, ``routing``,
+2 -687
View File
@@ -1,691 +1,6 @@
"""
Test runner for multimodal_gen that manages test suites and parallel execution.
For diffusion 1-gpu/2-gpu suites, cases are partitioned by estimated runtime
using LPT so each CI shard has a similar total runtime.
"""
import argparse
import copy
import json
import os
import random
import subprocess
import sys
import time
from dataclasses import dataclass
from pathlib import Path
import tabulate
from sglang.multimodal_gen.runtime.platforms import current_platform
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
from sglang.multimodal_gen.test.partitioning import PartitionItem, assign_partition
from sglang.multimodal_gen.test.runner.pytest_runner import (
partition_items_by_index,
run_pytest,
)
from sglang.multimodal_gen.test.server.testcase_configs import (
BASELINE_CONFIG,
DiffusionTestCase,
)
# TODO: remove duplicated code
if current_platform.is_npu():
from sglang.multimodal_gen.test.server.ascend.testcase_configs_npu import (
_UPDATE_WEIGHTS_FROM_DISK_TEST_FILE,
COMPONENT_ACCURACY_SUITES,
DEFAULT_EST_TIME_SECONDS,
DEFAULT_STANDALONE_EST_TIME_SECONDS,
FILE_SUITES,
PARAMETRIZED_CASE_GROUPS,
STANDALONE_FILE_EST_TIMES,
STANDALONE_FILES,
STARTUP_OVERHEAD_SECONDS,
SUITES,
)
else:
from sglang.multimodal_gen.test.server.gpu_cases import ( # noqa: F401 It is used by ci scripts
_UPDATE_WEIGHTS_FROM_DISK_TEST_FILE,
_UPDATE_WEIGHTS_MODEL_PAIR_ENV,
_UPDATE_WEIGHTS_MODEL_PAIR_IDS,
COMPONENT_ACCURACY_FILE_NUM_GPUS,
COMPONENT_ACCURACY_SUITES,
DEFAULT_EST_TIME_SECONDS,
DEFAULT_STANDALONE_EST_TIME_SECONDS,
FILE_SUITES,
ONE_GPU_CASES,
PARAMETRIZED_CASE_GROUPS,
STANDALONE_FILE_EST_TIMES,
STANDALONE_FILES,
STARTUP_OVERHEAD_SECONDS,
STRICT_SUITES,
SUITES,
TWO_GPU_CASES,
)
logger = init_logger(__name__)
@dataclass(frozen=True)
class PartitionAssignment:
case_ids: list[str]
standalone_files: list[str]
estimated_time: float | None = None
missing_standalone_estimates: list[str] | None = None
def get_case_est_time(case_id: str) -> float:
scenario = BASELINE_CONFIG.scenarios.get(case_id)
if scenario is None:
return DEFAULT_EST_TIME_SECONDS
if scenario.estimated_full_test_time_s is not None:
return scenario.estimated_full_test_time_s
return scenario.expected_e2e_ms / 1000.0 + STARTUP_OVERHEAD_SECONDS
def get_standalone_file_est_time(
suite: str, standalone_file: str
) -> tuple[float, bool]:
suite_est_times = STANDALONE_FILE_EST_TIMES.get(suite, {})
if standalone_file not in suite_est_times:
return DEFAULT_STANDALONE_EST_TIME_SECONDS, True
return suite_est_times[standalone_file], False
def get_all_standalone_file_est_times() -> dict[str, dict[str, float]]:
return copy.deepcopy(STANDALONE_FILE_EST_TIMES)
def validate_standalone_file_est_times() -> dict[str, list[str]]:
missing_by_suite: dict[str, list[str]] = {}
for suite, standalone_files in STANDALONE_FILES.items():
suite_est_times = STANDALONE_FILE_EST_TIMES.get(suite, {})
missing = [
standalone_file
for standalone_file in standalone_files
if standalone_file not in suite_est_times
]
if missing:
missing_by_suite[suite] = missing
return missing_by_suite
def get_suite_files_rel(suite: str, parametrized_only: bool = False) -> list[str]:
if parametrized_only and suite in PARAMETRIZED_CASE_GROUPS:
return [filename for filename, _ in PARAMETRIZED_CASE_GROUPS[suite]]
return SUITES[suite]
def _normalize_standalone_key(standalone_file: str) -> str:
return f"standalone:{standalone_file}"
def parse_partition_plan(
suite: str,
partition_id: int,
total_partitions: int,
plan_json: str,
) -> PartitionAssignment:
plan = json.loads(plan_json)
if plan.get("suite") != suite:
raise ValueError(
f"Partition plan suite mismatch: expected {suite!r}, "
f"got {plan.get('suite')!r}"
)
partition_count = plan.get("partition_count")
if partition_count != total_partitions:
raise ValueError(
f"Partition count mismatch for suite {suite!r}: "
f"plan={partition_count}, matrix={total_partitions}"
)
partitions = plan.get("partitions", [])
selected_partition = None
for partition in partitions:
if partition.get("part") == partition_id:
selected_partition = partition
break
if selected_partition is None:
raise ValueError(
f"Partition {partition_id} not found in plan for suite {suite!r}"
)
return PartitionAssignment(
case_ids=list(selected_partition.get("case_ids", [])),
standalone_files=list(selected_partition.get("standalone_files", [])),
estimated_time=selected_partition.get("estimated_time"),
missing_standalone_estimates=list(
selected_partition.get("missing_standalone_estimates", [])
),
)
def build_local_partition_assignment(
suite: str,
partition_id: int,
total_partitions: int,
) -> PartitionAssignment:
"""Assign this shard's work when CI did not precompute a partition plan.
Lanes with a hardcoded ``--total-partitions`` (the AMD ones) cannot give
every standalone file a shard of its own, so standalone files are LPT
balanced together with the parametrized cases instead.
"""
items = [
PartitionItem(kind="case", item_id=case.id, est_time=get_case_est_time(case.id))
for case in _get_dynamic_suite_cases(suite)
]
for standalone_file in STANDALONE_FILES.get(suite, []):
items.append(
PartitionItem(
kind="standalone",
item_id=standalone_file,
est_time=get_standalone_file_est_time(suite, standalone_file)[0],
)
)
my_items = assign_partition(items, partition_id, total_partitions)
return PartitionAssignment(
case_ids=[item.item_id for item in my_items if item.kind == "case"],
standalone_files=[
item.item_id for item in my_items if item.kind == "standalone"
],
)
def _merge_execution_results(
executed_cases: list[str],
case_results: dict[str, str],
new_executed_cases: list[str],
new_case_results: dict[str, str],
) -> None:
executed_cases.extend(
case_id for case_id in new_executed_cases if case_id not in executed_cases
)
case_results.update(new_case_results)
def _format_standalone_estimate_snippet(
suite: str, standalone_file: str, measured_full_test_time_s: float
) -> str:
return (
f'"{suite}": {{\n "{standalone_file}": {measured_full_test_time_s:.1f},\n}}'
)
def _print_missing_standalone_estimate_message(
suite: str,
standalone_file: str,
measured_full_test_time_s: float,
) -> None:
snippet = _format_standalone_estimate_snippet(
suite, standalone_file, measured_full_test_time_s
)
logger.error(
f"\n{'=' * 60}\n"
f'Add standalone estimate for suite "{suite}" and file "{standalone_file}":\n\n'
f"File: python/sglang/multimodal_gen/test/run_suite.py\n\n"
f"Current partition used fallback estimate: "
f"{DEFAULT_STANDALONE_EST_TIME_SECONDS:.1f}s\n\n"
f"{snippet}\n"
f"{'=' * 60}\n"
)
def _run_standalone_file(
suite: str,
standalone_rel: str,
target_dir: Path,
extra_filter: str | None = None,
) -> tuple[int, list[str], dict[str, str], dict]:
if standalone_rel == _UPDATE_WEIGHTS_FROM_DISK_TEST_FILE:
_maybe_pin_update_weights_model_pair([standalone_rel])
est_time, used_fallback_estimate = get_standalone_file_est_time(
suite, standalone_rel
)
standalone_file = _resolve_suite_files(target_dir, [standalone_rel], strict=True)[0]
junit_xml_path = str(
target_dir / f"junit_results_{suite}_{Path(standalone_rel).stem}.xml"
)
start_time = time.perf_counter()
exit_code, _, _ = run_pytest(
[standalone_file],
filter_expr=extra_filter,
junit_xml_path=junit_xml_path,
)
measured_full_test_time_s = round(time.perf_counter() - start_time, 1)
standalone_key = _normalize_standalone_key(standalone_rel)
measurement = {
"suite": suite,
"standalone_file": standalone_rel,
"measured_full_test_time_s": measured_full_test_time_s,
"used_fallback_estimate": used_fallback_estimate,
"fallback_estimate_s": DEFAULT_STANDALONE_EST_TIME_SECONDS,
"had_configured_estimate": not used_fallback_estimate,
"configured_or_fallback_estimate_s": est_time,
}
if used_fallback_estimate:
_print_missing_standalone_estimate_message(
suite, standalone_rel, measured_full_test_time_s
)
return (
exit_code,
[standalone_key],
{standalone_key: "pass" if exit_code == 0 else "fail"},
measurement,
)
def parse_args():
suite_choices = sorted(set(FILE_SUITES) | set(PARAMETRIZED_CASE_GROUPS))
parser = argparse.ArgumentParser(description="Run multimodal_gen test suite")
parser.add_argument(
"--suite",
type=str,
required=True,
choices=suite_choices,
help="The test suite to run.",
)
parser.add_argument(
"--partition-id",
type=int,
default=0,
help="Index of the current partition (for parallel execution)",
)
parser.add_argument(
"--total-partitions",
type=int,
default=1,
help="Total number of partitions",
)
parser.add_argument(
"--base-dir",
type=str,
default="server",
help="Base directory for tests relative to this script's parent",
)
parser.add_argument(
"-k",
"--filter",
type=str,
default=None,
help="Pytest filter expression (passed to pytest -k)",
)
parser.add_argument(
"--continue-on-error",
action="store_true",
default=False,
help="Continue running remaining tests even if one fails.",
)
parser.add_argument(
"--partition-plan-json",
type=str,
default=None,
help="Full partition plan JSON for the current suite.",
)
return parser.parse_args()
def write_execution_report(
suite: str,
partition_id: int,
total_partitions: int,
executed_cases: list[str],
is_standalone: bool = False,
standalone_file: str | None = None,
case_results: dict[str, str] | None = None,
missing_standalone_estimates: list[str] | None = None,
standalone_measurements: list[dict] | None = None,
) -> str:
report = {
"suite": suite,
"partition_id": partition_id,
"total_partitions": total_partitions,
"is_standalone": is_standalone,
"standalone_file": standalone_file,
"executed_cases": executed_cases,
"case_results": case_results or {},
"missing_standalone_estimates": missing_standalone_estimates or [],
"standalone_measurements": standalone_measurements or [],
}
report_filename = f"execution_report_{suite}_{partition_id}.json"
report_path = Path(__file__).parent / report_filename
with open(report_path, "w", encoding="utf-8") as f:
json.dump(report, f, indent=2)
logger.info("Execution report written to: %s", report_path)
return str(report_path)
def run_component_accuracy_files(files, filter_expr=None, continue_on_error=False):
exit_code = 0
for file_path in files:
file_name = Path(file_path).name
num_gpus = COMPONENT_ACCURACY_FILE_NUM_GPUS.get(file_name, 1)
if num_gpus > 1:
cmd = [
sys.executable,
"-m",
"torch.distributed.run",
f"--nproc_per_node={num_gpus}",
"-m",
"pytest",
"-s",
"-v",
]
else:
cmd = [sys.executable, "-m", "pytest", "-s", "-v"]
if filter_expr:
cmd.extend(["-k", filter_expr])
cmd.append(file_path)
print(f"Running command: {' '.join(cmd)}")
file_exit_code = subprocess.call(cmd)
if file_exit_code == 5:
print(
"No tests collected (exit code 5). This is expected when filters "
"deselect all tests in a file. Treating as success."
)
file_exit_code = 0
if file_exit_code != 0 and exit_code == 0:
exit_code = file_exit_code
if file_exit_code != 0 and not continue_on_error:
return file_exit_code
return exit_code
def _is_in_ci() -> bool:
return os.environ.get("SGLANG_IS_IN_CI", "").lower() in ("1", "true", "yes", "on")
def _maybe_pin_update_weights_model_pair(suite_files_rel: list[str]) -> None:
if not _is_in_ci():
return
if _UPDATE_WEIGHTS_FROM_DISK_TEST_FILE not in suite_files_rel:
return
if os.environ.get(_UPDATE_WEIGHTS_MODEL_PAIR_ENV):
print(
f"Using preset {_UPDATE_WEIGHTS_MODEL_PAIR_ENV}="
f"{os.environ[_UPDATE_WEIGHTS_MODEL_PAIR_ENV]}"
)
return
selected_pair = random.choice(_UPDATE_WEIGHTS_MODEL_PAIR_IDS)
os.environ[_UPDATE_WEIGHTS_MODEL_PAIR_ENV] = selected_pair
print(f"Selected {_UPDATE_WEIGHTS_MODEL_PAIR_ENV}={selected_pair} for this CI run")
def _resolve_suite_files(
target_dir: Path, suite_files_rel: list[str], strict: bool
) -> list[str]:
suite_files_abs = []
for f_rel in suite_files_rel:
f_abs = target_dir / f_rel
if not f_abs.exists():
msg = f"Test file {f_rel} not found in {target_dir}."
if strict:
print(f"Error: {msg}")
sys.exit(1)
print(f"Warning: {msg} Skipping.")
continue
suite_files_abs.append(str(f_abs))
return suite_files_abs
def _run_file_suite(args, target_dir: Path) -> int:
suite_files_rel = FILE_SUITES[args.suite]
_maybe_pin_update_weights_model_pair(suite_files_rel)
suite_files_abs = _resolve_suite_files(
target_dir, suite_files_rel, args.suite in STRICT_SUITES
)
if not suite_files_abs:
print(f"No valid test files found for suite '{args.suite}'.")
return 1 if args.suite in STRICT_SUITES else 0
exit_code, _, _ = run_pytest(
suite_files_abs,
filter_expr=args.filter,
junit_xml_path=None,
)
return exit_code
def _get_dynamic_suite_cases(suite: str) -> list[DiffusionTestCase]:
cases = []
for _, case_group in PARAMETRIZED_CASE_GROUPS[suite]:
cases.extend(case_group)
return cases
def _get_parametrized_files_for_case_ids(
suite: str, case_ids: set[str], target_dir: Path
) -> list[str]:
files = []
for filename, case_group in PARAMETRIZED_CASE_GROUPS[suite]:
if any(case.id in case_ids for case in case_group):
file_path = target_dir / filename
if file_path.exists():
files.append(str(file_path))
else:
logger.warning("Test file %s not found in %s", filename, target_dir)
return files
def _run_dynamic_suite(args, target_dir: Path) -> int:
if args.partition_plan_json:
assignment = parse_partition_plan(
suite=args.suite,
partition_id=args.partition_id,
total_partitions=args.total_partitions,
plan_json=args.partition_plan_json,
)
else:
assignment = build_local_partition_assignment(
suite=args.suite,
partition_id=args.partition_id,
total_partitions=args.total_partitions,
)
return _run_partition_assignment(args, target_dir, assignment)
def _run_partition_assignment(
args, target_dir: Path, assignment: PartitionAssignment
) -> int:
rows = [[args.suite, f"{args.partition_id + 1}/{args.total_partitions}"]]
print(tabulate.tabulate(rows, headers=["Suite", "Partition"], tablefmt="psql"))
total_est_time = 0.0
executed_cases: list[str] = []
case_results: dict[str, str] = {}
missing_standalone_estimates: list[str] = []
standalone_measurements: list[dict] = []
overall_exit_code = 0
if assignment.case_ids:
case_id_set = set(assignment.case_ids)
total_est_time += sum(
get_case_est_time(case_id) for case_id in assignment.case_ids
)
suite_files = _get_parametrized_files_for_case_ids(
args.suite, case_id_set, target_dir
)
if not suite_files:
print(f"No valid parametrized test files found for suite '{args.suite}'.")
return 0
partition_filter = " or ".join(
f"[{case_id}]" for case_id in assignment.case_ids
)
filter_expr = (
f"({partition_filter}) and ({args.filter})"
if args.filter
else partition_filter
)
print(
f"Running {len(assignment.case_ids)} parametrized cases with estimated total "
f"{sum(get_case_est_time(case_id) for case_id in assignment.case_ids):.1f}s:"
)
for case_id in assignment.case_ids:
print(f" - case: {case_id} ({get_case_est_time(case_id):.1f}s)")
print(f"Test files: {[Path(f).name for f in suite_files]}")
print(f"Filter expression: {filter_expr}")
junit_xml_path = str(
target_dir / f"junit_results_{args.suite}_{args.partition_id}.xml"
)
exit_code, new_executed_cases, new_case_results = run_pytest(
suite_files,
filter_expr=filter_expr,
junit_xml_path=junit_xml_path,
)
_merge_execution_results(
executed_cases, case_results, new_executed_cases, new_case_results
)
# A failing case must not swallow this shard's standalone files: they
# are separate pytest runs, and they only share a shard because the
# shard count is fixed. --continue-on-error still decides whether a
# failing standalone file stops the ones queued behind it.
if exit_code != 0 and overall_exit_code == 0:
overall_exit_code = exit_code
if assignment.standalone_files:
standalone_estimate = sum(
get_standalone_file_est_time(args.suite, standalone_file)[0]
for standalone_file in assignment.standalone_files
)
total_est_time += standalone_estimate
print(
f"Running {len(assignment.standalone_files)} standalone file(s) with estimated total "
f"{standalone_estimate:.1f}s:"
)
for standalone_file in assignment.standalone_files:
est_time, used_fallback_estimate = get_standalone_file_est_time(
args.suite, standalone_file
)
fallback_suffix = (
f", fallback estimate {DEFAULT_STANDALONE_EST_TIME_SECONDS:.1f}s"
if used_fallback_estimate
else ""
)
print(
f" - standalone: {standalone_file} ({est_time:.1f}s{fallback_suffix})"
)
for standalone_file in assignment.standalone_files:
exit_code, new_executed_cases, new_case_results, measurement = (
_run_standalone_file(
args.suite,
standalone_file,
target_dir,
extra_filter=args.filter,
)
)
if measurement["used_fallback_estimate"]:
missing_standalone_estimates.append(standalone_file)
standalone_measurements.append(measurement)
_merge_execution_results(
executed_cases,
case_results,
new_executed_cases,
new_case_results,
)
if exit_code != 0 and overall_exit_code == 0:
overall_exit_code = exit_code
if exit_code != 0 and not args.continue_on_error:
break
if not assignment.case_ids and not assignment.standalone_files:
print(f"No work assigned to partition {args.partition_id}. Exiting success.")
print(f"Partition estimated total time: {total_est_time:.1f}s")
write_execution_report(
suite=args.suite,
partition_id=args.partition_id,
total_partitions=args.total_partitions,
executed_cases=executed_cases,
is_standalone=False,
standalone_file=None,
case_results=case_results,
missing_standalone_estimates=missing_standalone_estimates,
standalone_measurements=standalone_measurements,
)
return overall_exit_code
def main():
args = parse_args()
validate_standalone_file_est_times()
test_root_dir = Path(__file__).resolve().parent
target_dir = test_root_dir / args.base_dir
if not target_dir.exists():
print(f"Error: Target directory {target_dir} does not exist.")
sys.exit(1)
if args.suite in COMPONENT_ACCURACY_SUITES:
suite_files_rel = FILE_SUITES[args.suite]
suite_files_abs = _resolve_suite_files(
target_dir, suite_files_rel, args.suite in STRICT_SUITES
)
if not suite_files_abs:
print(f"No valid test files found for suite '{args.suite}'.")
sys.exit(1 if args.suite in STRICT_SUITES else 0)
my_files = partition_items_by_index(
suite_files_abs, args.partition_id, args.total_partitions
)
partition_info = (
f"{args.partition_id + 1}/{args.total_partitions} "
f"(0-based id={args.partition_id})"
)
headers = ["Suite", "Partition"]
rows = [[args.suite, partition_info]]
msg = tabulate.tabulate(rows, headers=headers, tablefmt="psql") + "\n"
msg += f"Enabled {len(my_files)} file(s):\n"
for file_path in my_files:
msg += f" - {file_path}\n"
print(msg, flush=True)
print(
f"Suite: {args.suite} | Partition: {args.partition_id}/{args.total_partitions}"
)
print(f"Selected {len(suite_files_abs)} files:")
for f in suite_files_abs:
print(f" - {os.path.basename(f)}")
if not my_files:
print("No files assigned to this partition. Exiting success.")
sys.exit(0)
print(f"Running {len(my_files)} files in this shard: {', '.join(my_files)}")
exit_code = run_component_accuracy_files(
my_files,
filter_expr=args.filter,
continue_on_error=args.continue_on_error,
)
msg = "\n" + tabulate.tabulate(rows, headers=headers, tablefmt="psql") + "\n"
msg += f"Executed {len(my_files)} file(s):\n"
for file_path in my_files:
msg += f" - {file_path}\n"
print(msg, flush=True)
elif args.suite in PARAMETRIZED_CASE_GROUPS:
exit_code = _run_dynamic_suite(args, target_dir)
else:
exit_code = _run_file_suite(args, target_dir)
sys.exit(exit_code)
"""Compatibility entry point; CI dispatches diffusion via ``test/run_suite.py``."""
from sglang.multimodal_gen.test.runner.diffusion_suite_runner import main
if __name__ == "__main__":
main()
@@ -0,0 +1,690 @@
"""Internal diffusion-suite adapter used by ``test/run_suite.py`` bridges.
For diffusion 1-gpu/2-gpu suites, cases are partitioned by estimated runtime
using LPT so each CI shard has a similar total runtime.
"""
import argparse
import copy
import json
import os
import random
import subprocess
import sys
import time
from dataclasses import dataclass
from pathlib import Path
import tabulate
from sglang.multimodal_gen.runtime.platforms import current_platform
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
from sglang.multimodal_gen.test.partitioning import PartitionItem, assign_partition
from sglang.multimodal_gen.test.runner.pytest_runner import (
partition_items_by_index,
run_pytest,
)
from sglang.multimodal_gen.test.server.testcase_configs import (
BASELINE_CONFIG,
DiffusionTestCase,
)
# TODO: remove duplicated code
if current_platform.is_npu():
from sglang.multimodal_gen.test.server.ascend.testcase_configs_npu import (
_UPDATE_WEIGHTS_FROM_DISK_TEST_FILE,
COMPONENT_ACCURACY_SUITES,
DEFAULT_EST_TIME_SECONDS,
DEFAULT_STANDALONE_EST_TIME_SECONDS,
FILE_SUITES,
PARAMETRIZED_CASE_GROUPS,
STANDALONE_FILE_EST_TIMES,
STANDALONE_FILES,
STARTUP_OVERHEAD_SECONDS,
SUITES,
)
else:
from sglang.multimodal_gen.test.server.gpu_cases import ( # noqa: F401 It is used by ci scripts
_UPDATE_WEIGHTS_FROM_DISK_TEST_FILE,
_UPDATE_WEIGHTS_MODEL_PAIR_ENV,
_UPDATE_WEIGHTS_MODEL_PAIR_IDS,
COMPONENT_ACCURACY_FILE_NUM_GPUS,
COMPONENT_ACCURACY_SUITES,
DEFAULT_EST_TIME_SECONDS,
DEFAULT_STANDALONE_EST_TIME_SECONDS,
FILE_SUITES,
ONE_GPU_CASES,
PARAMETRIZED_CASE_GROUPS,
STANDALONE_FILE_EST_TIMES,
STANDALONE_FILES,
STARTUP_OVERHEAD_SECONDS,
STRICT_SUITES,
SUITES,
TWO_GPU_CASES,
)
logger = init_logger(__name__)
@dataclass(frozen=True)
class PartitionAssignment:
case_ids: list[str]
standalone_files: list[str]
estimated_time: float | None = None
missing_standalone_estimates: list[str] | None = None
def get_case_est_time(case_id: str) -> float:
scenario = BASELINE_CONFIG.scenarios.get(case_id)
if scenario is None:
return DEFAULT_EST_TIME_SECONDS
if scenario.estimated_full_test_time_s is not None:
return scenario.estimated_full_test_time_s
return scenario.expected_e2e_ms / 1000.0 + STARTUP_OVERHEAD_SECONDS
def get_standalone_file_est_time(
suite: str, standalone_file: str
) -> tuple[float, bool]:
suite_est_times = STANDALONE_FILE_EST_TIMES.get(suite, {})
if standalone_file not in suite_est_times:
return DEFAULT_STANDALONE_EST_TIME_SECONDS, True
return suite_est_times[standalone_file], False
def get_all_standalone_file_est_times() -> dict[str, dict[str, float]]:
return copy.deepcopy(STANDALONE_FILE_EST_TIMES)
def validate_standalone_file_est_times() -> dict[str, list[str]]:
missing_by_suite: dict[str, list[str]] = {}
for suite, standalone_files in STANDALONE_FILES.items():
suite_est_times = STANDALONE_FILE_EST_TIMES.get(suite, {})
missing = [
standalone_file
for standalone_file in standalone_files
if standalone_file not in suite_est_times
]
if missing:
missing_by_suite[suite] = missing
return missing_by_suite
def get_suite_files_rel(suite: str, parametrized_only: bool = False) -> list[str]:
if parametrized_only and suite in PARAMETRIZED_CASE_GROUPS:
return [filename for filename, _ in PARAMETRIZED_CASE_GROUPS[suite]]
return SUITES[suite]
def _normalize_standalone_key(standalone_file: str) -> str:
return f"standalone:{standalone_file}"
def parse_partition_plan(
suite: str,
partition_id: int,
total_partitions: int,
plan_json: str,
) -> PartitionAssignment:
plan = json.loads(plan_json)
if plan.get("suite") != suite:
raise ValueError(
f"Partition plan suite mismatch: expected {suite!r}, "
f"got {plan.get('suite')!r}"
)
partition_count = plan.get("partition_count")
if partition_count != total_partitions:
raise ValueError(
f"Partition count mismatch for suite {suite!r}: "
f"plan={partition_count}, matrix={total_partitions}"
)
partitions = plan.get("partitions", [])
selected_partition = None
for partition in partitions:
if partition.get("part") == partition_id:
selected_partition = partition
break
if selected_partition is None:
raise ValueError(
f"Partition {partition_id} not found in plan for suite {suite!r}"
)
return PartitionAssignment(
case_ids=list(selected_partition.get("case_ids", [])),
standalone_files=list(selected_partition.get("standalone_files", [])),
estimated_time=selected_partition.get("estimated_time"),
missing_standalone_estimates=list(
selected_partition.get("missing_standalone_estimates", [])
),
)
def build_local_partition_assignment(
suite: str,
partition_id: int,
total_partitions: int,
) -> PartitionAssignment:
"""Assign this shard's work when CI did not precompute a partition plan.
Lanes with a hardcoded ``--total-partitions`` (the AMD ones) cannot give
every standalone file a shard of its own, so standalone files are LPT
balanced together with the parametrized cases instead.
"""
items = [
PartitionItem(kind="case", item_id=case.id, est_time=get_case_est_time(case.id))
for case in _get_dynamic_suite_cases(suite)
]
for standalone_file in STANDALONE_FILES.get(suite, []):
items.append(
PartitionItem(
kind="standalone",
item_id=standalone_file,
est_time=get_standalone_file_est_time(suite, standalone_file)[0],
)
)
my_items = assign_partition(items, partition_id, total_partitions)
return PartitionAssignment(
case_ids=[item.item_id for item in my_items if item.kind == "case"],
standalone_files=[
item.item_id for item in my_items if item.kind == "standalone"
],
)
def _merge_execution_results(
executed_cases: list[str],
case_results: dict[str, str],
new_executed_cases: list[str],
new_case_results: dict[str, str],
) -> None:
executed_cases.extend(
case_id for case_id in new_executed_cases if case_id not in executed_cases
)
case_results.update(new_case_results)
def _format_standalone_estimate_snippet(
suite: str, standalone_file: str, measured_full_test_time_s: float
) -> str:
return (
f'"{suite}": {{\n "{standalone_file}": {measured_full_test_time_s:.1f},\n}}'
)
def _print_missing_standalone_estimate_message(
suite: str,
standalone_file: str,
measured_full_test_time_s: float,
) -> None:
snippet = _format_standalone_estimate_snippet(
suite, standalone_file, measured_full_test_time_s
)
logger.error(
f"\n{'=' * 60}\n"
f'Add standalone estimate for suite "{suite}" and file "{standalone_file}":\n\n'
"File: python/sglang/multimodal_gen/test/server/gpu_cases.py\n\n"
f"Current partition used fallback estimate: "
f"{DEFAULT_STANDALONE_EST_TIME_SECONDS:.1f}s\n\n"
f"{snippet}\n"
f"{'=' * 60}\n"
)
def _run_standalone_file(
suite: str,
standalone_rel: str,
target_dir: Path,
extra_filter: str | None = None,
) -> tuple[int, list[str], dict[str, str], dict]:
if standalone_rel == _UPDATE_WEIGHTS_FROM_DISK_TEST_FILE:
_maybe_pin_update_weights_model_pair([standalone_rel])
est_time, used_fallback_estimate = get_standalone_file_est_time(
suite, standalone_rel
)
standalone_file = _resolve_suite_files(target_dir, [standalone_rel], strict=True)[0]
junit_xml_path = str(
target_dir / f"junit_results_{suite}_{Path(standalone_rel).stem}.xml"
)
start_time = time.perf_counter()
exit_code, _, _ = run_pytest(
[standalone_file],
filter_expr=extra_filter,
junit_xml_path=junit_xml_path,
)
measured_full_test_time_s = round(time.perf_counter() - start_time, 1)
standalone_key = _normalize_standalone_key(standalone_rel)
measurement = {
"suite": suite,
"standalone_file": standalone_rel,
"measured_full_test_time_s": measured_full_test_time_s,
"used_fallback_estimate": used_fallback_estimate,
"fallback_estimate_s": DEFAULT_STANDALONE_EST_TIME_SECONDS,
"had_configured_estimate": not used_fallback_estimate,
"configured_or_fallback_estimate_s": est_time,
}
if used_fallback_estimate:
_print_missing_standalone_estimate_message(
suite, standalone_rel, measured_full_test_time_s
)
return (
exit_code,
[standalone_key],
{standalone_key: "pass" if exit_code == 0 else "fail"},
measurement,
)
def parse_args():
suite_choices = sorted(set(FILE_SUITES) | set(PARAMETRIZED_CASE_GROUPS))
parser = argparse.ArgumentParser(description="Run multimodal_gen test suite")
parser.add_argument(
"--suite",
type=str,
required=True,
choices=suite_choices,
help="The test suite to run.",
)
parser.add_argument(
"--partition-id",
type=int,
default=0,
help="Index of the current partition (for parallel execution)",
)
parser.add_argument(
"--total-partitions",
type=int,
default=1,
help="Total number of partitions",
)
parser.add_argument(
"--base-dir",
type=str,
default="server",
help="Base directory for tests relative to this script's parent",
)
parser.add_argument(
"-k",
"--filter",
type=str,
default=None,
help="Pytest filter expression (passed to pytest -k)",
)
parser.add_argument(
"--continue-on-error",
action="store_true",
default=False,
help="Continue running remaining tests even if one fails.",
)
parser.add_argument(
"--partition-plan-json",
type=str,
default=None,
help="Full partition plan JSON for the current suite.",
)
return parser.parse_args()
def write_execution_report(
suite: str,
partition_id: int,
total_partitions: int,
executed_cases: list[str],
is_standalone: bool = False,
standalone_file: str | None = None,
case_results: dict[str, str] | None = None,
missing_standalone_estimates: list[str] | None = None,
standalone_measurements: list[dict] | None = None,
) -> str:
report = {
"suite": suite,
"partition_id": partition_id,
"total_partitions": total_partitions,
"is_standalone": is_standalone,
"standalone_file": standalone_file,
"executed_cases": executed_cases,
"case_results": case_results or {},
"missing_standalone_estimates": missing_standalone_estimates or [],
"standalone_measurements": standalone_measurements or [],
}
report_filename = f"execution_report_{suite}_{partition_id}.json"
report_path = Path(__file__).resolve().parents[1] / report_filename
with open(report_path, "w", encoding="utf-8") as f:
json.dump(report, f, indent=2)
logger.info("Execution report written to: %s", report_path)
return str(report_path)
def run_component_accuracy_files(files, filter_expr=None, continue_on_error=False):
exit_code = 0
for file_path in files:
file_name = Path(file_path).name
num_gpus = COMPONENT_ACCURACY_FILE_NUM_GPUS.get(file_name, 1)
if num_gpus > 1:
cmd = [
sys.executable,
"-m",
"torch.distributed.run",
f"--nproc_per_node={num_gpus}",
"-m",
"pytest",
"-s",
"-v",
]
else:
cmd = [sys.executable, "-m", "pytest", "-s", "-v"]
if filter_expr:
cmd.extend(["-k", filter_expr])
cmd.append(file_path)
print(f"Running command: {' '.join(cmd)}")
file_exit_code = subprocess.call(cmd)
if file_exit_code == 5:
print(
"No tests collected (exit code 5). This is expected when filters "
"deselect all tests in a file. Treating as success."
)
file_exit_code = 0
if file_exit_code != 0 and exit_code == 0:
exit_code = file_exit_code
if file_exit_code != 0 and not continue_on_error:
return file_exit_code
return exit_code
def _is_in_ci() -> bool:
return os.environ.get("SGLANG_IS_IN_CI", "").lower() in ("1", "true", "yes", "on")
def _maybe_pin_update_weights_model_pair(suite_files_rel: list[str]) -> None:
if not _is_in_ci():
return
if _UPDATE_WEIGHTS_FROM_DISK_TEST_FILE not in suite_files_rel:
return
if os.environ.get(_UPDATE_WEIGHTS_MODEL_PAIR_ENV):
print(
f"Using preset {_UPDATE_WEIGHTS_MODEL_PAIR_ENV}="
f"{os.environ[_UPDATE_WEIGHTS_MODEL_PAIR_ENV]}"
)
return
selected_pair = random.choice(_UPDATE_WEIGHTS_MODEL_PAIR_IDS)
os.environ[_UPDATE_WEIGHTS_MODEL_PAIR_ENV] = selected_pair
print(f"Selected {_UPDATE_WEIGHTS_MODEL_PAIR_ENV}={selected_pair} for this CI run")
def _resolve_suite_files(
target_dir: Path, suite_files_rel: list[str], strict: bool
) -> list[str]:
suite_files_abs = []
for f_rel in suite_files_rel:
f_abs = target_dir / f_rel
if not f_abs.exists():
msg = f"Test file {f_rel} not found in {target_dir}."
if strict:
print(f"Error: {msg}")
sys.exit(1)
print(f"Warning: {msg} Skipping.")
continue
suite_files_abs.append(str(f_abs))
return suite_files_abs
def _run_file_suite(args, target_dir: Path) -> int:
suite_files_rel = FILE_SUITES[args.suite]
_maybe_pin_update_weights_model_pair(suite_files_rel)
suite_files_abs = _resolve_suite_files(
target_dir, suite_files_rel, args.suite in STRICT_SUITES
)
if not suite_files_abs:
print(f"No valid test files found for suite '{args.suite}'.")
return 1 if args.suite in STRICT_SUITES else 0
exit_code, _, _ = run_pytest(
suite_files_abs,
filter_expr=args.filter,
junit_xml_path=None,
)
return exit_code
def _get_dynamic_suite_cases(suite: str) -> list[DiffusionTestCase]:
cases = []
for _, case_group in PARAMETRIZED_CASE_GROUPS[suite]:
cases.extend(case_group)
return cases
def _get_parametrized_files_for_case_ids(
suite: str, case_ids: set[str], target_dir: Path
) -> list[str]:
files = []
for filename, case_group in PARAMETRIZED_CASE_GROUPS[suite]:
if any(case.id in case_ids for case in case_group):
file_path = target_dir / filename
if file_path.exists():
files.append(str(file_path))
else:
logger.warning("Test file %s not found in %s", filename, target_dir)
return files
def _run_dynamic_suite(args, target_dir: Path) -> int:
if args.partition_plan_json:
assignment = parse_partition_plan(
suite=args.suite,
partition_id=args.partition_id,
total_partitions=args.total_partitions,
plan_json=args.partition_plan_json,
)
else:
assignment = build_local_partition_assignment(
suite=args.suite,
partition_id=args.partition_id,
total_partitions=args.total_partitions,
)
return _run_partition_assignment(args, target_dir, assignment)
def _run_partition_assignment(
args, target_dir: Path, assignment: PartitionAssignment
) -> int:
rows = [[args.suite, f"{args.partition_id + 1}/{args.total_partitions}"]]
print(tabulate.tabulate(rows, headers=["Suite", "Partition"], tablefmt="psql"))
total_est_time = 0.0
executed_cases: list[str] = []
case_results: dict[str, str] = {}
missing_standalone_estimates: list[str] = []
standalone_measurements: list[dict] = []
overall_exit_code = 0
if assignment.case_ids:
case_id_set = set(assignment.case_ids)
total_est_time += sum(
get_case_est_time(case_id) for case_id in assignment.case_ids
)
suite_files = _get_parametrized_files_for_case_ids(
args.suite, case_id_set, target_dir
)
if not suite_files:
print(f"No valid parametrized test files found for suite '{args.suite}'.")
return 0
partition_filter = " or ".join(
f"[{case_id}]" for case_id in assignment.case_ids
)
filter_expr = (
f"({partition_filter}) and ({args.filter})"
if args.filter
else partition_filter
)
print(
f"Running {len(assignment.case_ids)} parametrized cases with estimated total "
f"{sum(get_case_est_time(case_id) for case_id in assignment.case_ids):.1f}s:"
)
for case_id in assignment.case_ids:
print(f" - case: {case_id} ({get_case_est_time(case_id):.1f}s)")
print(f"Test files: {[Path(f).name for f in suite_files]}")
print(f"Filter expression: {filter_expr}")
junit_xml_path = str(
target_dir / f"junit_results_{args.suite}_{args.partition_id}.xml"
)
exit_code, new_executed_cases, new_case_results = run_pytest(
suite_files,
filter_expr=filter_expr,
junit_xml_path=junit_xml_path,
)
_merge_execution_results(
executed_cases, case_results, new_executed_cases, new_case_results
)
# A failing case must not swallow this shard's standalone files: they
# are separate pytest runs, and they only share a shard because the
# shard count is fixed. --continue-on-error still decides whether a
# failing standalone file stops the ones queued behind it.
if exit_code != 0 and overall_exit_code == 0:
overall_exit_code = exit_code
if assignment.standalone_files:
standalone_estimate = sum(
get_standalone_file_est_time(args.suite, standalone_file)[0]
for standalone_file in assignment.standalone_files
)
total_est_time += standalone_estimate
print(
f"Running {len(assignment.standalone_files)} standalone file(s) with estimated total "
f"{standalone_estimate:.1f}s:"
)
for standalone_file in assignment.standalone_files:
est_time, used_fallback_estimate = get_standalone_file_est_time(
args.suite, standalone_file
)
fallback_suffix = (
f", fallback estimate {DEFAULT_STANDALONE_EST_TIME_SECONDS:.1f}s"
if used_fallback_estimate
else ""
)
print(
f" - standalone: {standalone_file} ({est_time:.1f}s{fallback_suffix})"
)
for standalone_file in assignment.standalone_files:
exit_code, new_executed_cases, new_case_results, measurement = (
_run_standalone_file(
args.suite,
standalone_file,
target_dir,
extra_filter=args.filter,
)
)
if measurement["used_fallback_estimate"]:
missing_standalone_estimates.append(standalone_file)
standalone_measurements.append(measurement)
_merge_execution_results(
executed_cases,
case_results,
new_executed_cases,
new_case_results,
)
if exit_code != 0 and overall_exit_code == 0:
overall_exit_code = exit_code
if exit_code != 0 and not args.continue_on_error:
break
if not assignment.case_ids and not assignment.standalone_files:
print(f"No work assigned to partition {args.partition_id}. Exiting success.")
print(f"Partition estimated total time: {total_est_time:.1f}s")
write_execution_report(
suite=args.suite,
partition_id=args.partition_id,
total_partitions=args.total_partitions,
executed_cases=executed_cases,
is_standalone=False,
standalone_file=None,
case_results=case_results,
missing_standalone_estimates=missing_standalone_estimates,
standalone_measurements=standalone_measurements,
)
return overall_exit_code
def main():
args = parse_args()
validate_standalone_file_est_times()
test_root_dir = Path(__file__).resolve().parents[1]
target_dir = test_root_dir / args.base_dir
if not target_dir.exists():
print(f"Error: Target directory {target_dir} does not exist.")
sys.exit(1)
if args.suite in COMPONENT_ACCURACY_SUITES:
suite_files_rel = FILE_SUITES[args.suite]
suite_files_abs = _resolve_suite_files(
target_dir, suite_files_rel, args.suite in STRICT_SUITES
)
if not suite_files_abs:
print(f"No valid test files found for suite '{args.suite}'.")
sys.exit(1 if args.suite in STRICT_SUITES else 0)
my_files = partition_items_by_index(
suite_files_abs, args.partition_id, args.total_partitions
)
partition_info = (
f"{args.partition_id + 1}/{args.total_partitions} "
f"(0-based id={args.partition_id})"
)
headers = ["Suite", "Partition"]
rows = [[args.suite, partition_info]]
msg = tabulate.tabulate(rows, headers=headers, tablefmt="psql") + "\n"
msg += f"Enabled {len(my_files)} file(s):\n"
for file_path in my_files:
msg += f" - {file_path}\n"
print(msg, flush=True)
print(
f"Suite: {args.suite} | Partition: {args.partition_id}/{args.total_partitions}"
)
print(f"Selected {len(suite_files_abs)} files:")
for f in suite_files_abs:
print(f" - {os.path.basename(f)}")
if not my_files:
print("No files assigned to this partition. Exiting success.")
sys.exit(0)
print(f"Running {len(my_files)} files in this shard: {', '.join(my_files)}")
exit_code = run_component_accuracy_files(
my_files,
filter_expr=args.filter,
continue_on_error=args.continue_on_error,
)
msg = "\n" + tabulate.tabulate(rows, headers=headers, tablefmt="psql") + "\n"
msg += f"Executed {len(my_files)} file(s):\n"
for file_path in my_files:
msg += f" - {file_path}\n"
print(msg, flush=True)
elif args.suite in PARAMETRIZED_CASE_GROUPS:
exit_code = _run_dynamic_suite(args, target_dir)
else:
exit_code = _run_file_suite(args, target_dir)
sys.exit(exit_code)
if __name__ == "__main__":
main()
@@ -2,7 +2,7 @@
"""
Generate diffusion CI outputs for consistency testing.
This script reuses the CI test code by calling run_suite.py with SGLANG_GEN_GT=1,
This script reuses the CI test adapter with SGLANG_GEN_GT=1,
ensuring that GT generation uses exactly the same code path as CI tests.
Usage:
@@ -21,7 +21,7 @@ from sglang.multimodal_gen.test.partitioning import (
PartitionItem,
partition_items_by_lpt,
)
from sglang.multimodal_gen.test.run_suite import (
from sglang.multimodal_gen.test.runner.diffusion_suite_runner import (
SUITES,
_maybe_pin_update_weights_model_pair,
get_case_est_time,
@@ -1,112 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
from pathlib import Path
def test_realtime_webui_presets_do_not_emit_camera_scripts():
repo_root = Path(__file__).resolve().parents[6]
app_js = (
repo_root / "python/sglang/multimodal_gen/apps/realtime_webui/app.js"
).read_text()
index_html = (
repo_root / "python/sglang/multimodal_gen/apps/realtime_webui/index.html"
).read_text()
styles_css = (
repo_root / "python/sglang/multimodal_gen/apps/realtime_webui/styles.css"
).read_text()
playback_js = (
repo_root
/ "python/sglang/multimodal_gen/apps/realtime_webui/playback_controller.js"
).read_text()
assert "preset.actions" not in app_js
assert "repeatActions" not in app_js
assert 'id="eventFrames"' not in index_html
assert "ControlStateController" in app_js
assert 'const DEFAULT_PREVIEW_OUTPUT_FORMAT = "webp";' in app_js
assert 'id="transportFormat"' in index_html
assert 'id="fps" type="number" value="25"' in index_html
assert 'id="superResolution" type="checkbox"' in index_html
assert 'id="upscalingScale"' in index_html
assert 'class="workspace"' in index_html
assert 'class="preview-frame"' in index_html
assert 'id="previewOverlay" class="preview-overlay"' in index_html
assert 'id="previewScale" type="range" min="80" max="170" value="120"' in index_html
assert 'id="previewScaleText"' in index_html
assert 'id="outputSizeText"' in index_html
assert 'id="frameInterpolation" type="checkbox" />' in index_html
assert (
'id="serverUrl" value="ws://127.0.0.1:30000/v1/realtime_video/generate"'
in index_html
)
assert '<option value="webp" selected>WebP preview</option>' in index_html
assert 'id="serverSendText"' in index_html
assert 'id="theoreticalFpsText"' in index_html
assert 'id="renderFps"' in index_html
assert 'id="stageRenderFps"' not in index_html
assert "sglang-diffusion Realtime Studio" in index_html
assert "SGLD" not in index_html
assert 'class="tabs"' not in index_html
assert "Recordings" not in index_html
assert "API" not in index_html
assert "Info" not in index_html
assert 'id="steps" type="number" value="4"' in index_html
assert 'id="guidance" type="number" value="1"' in index_html
assert "styles.css?v=realtime-record-v49" in index_html
assert "app.js?v=realtime-record-v75" in index_html
assert (
'const DECODER_WORKER_URL = "./decoder_worker.js?v=rgb-worker-v10";' in app_js
)
assert "const DEFAULT_TARGET_FPS = 25;" in app_js
assert "const DEFAULT_FRAME_INTERPOLATION_EXP = 1;" in app_js
assert "const DEFAULT_FRAME_INTERPOLATION_SCALE = 1.0;" in app_js
assert "const DEFAULT_UPSCALING_SCALE = 2;" in app_js
assert "const DEFAULT_PREVIEW_SCALE = 120;" in app_js
assert 'setPreviewState("waiting")' in app_js
assert "stage.dataset.previewState = state" in app_js
assert "previewProgressSpin" in styles_css
assert "previewDotPulse" not in styles_css
assert 'document.querySelector(".preview-frame")' in app_js
assert 'previewFrame.style.setProperty("--preview-scale"' in app_js
assert "cancelAnimationFrame(previewScaleFrame)" in app_js
assert "enable_frame_interpolation: true" in app_js
assert "frame_interpolation_exp: DEFAULT_FRAME_INTERPOLATION_EXP" in app_js
assert "frame_interpolation_scale: DEFAULT_FRAME_INTERPOLATION_SCALE" in app_js
assert "readSuperResolutionParams()" in app_js
assert "enable_upscaling: true" in app_js
assert "upscaling_scale: readUpscalingScale()" in app_js
assert "updateOutputSizeFromHeader(header)" in app_js
assert "setPreviewScale(DEFAULT_PREVIEW_SCALE)" in app_js
assert "preview_scale" in app_js
assert "sr_scale" in app_js
assert "elapsedMs < targetMs" in playback_js
assert "queuedDecodeFrames > maxQueuedFrames" in app_js
assert (
'const REACTOR_PRESET_BASE_URL = "https://www.reactor.inc/lingbot-world-fast-v1";'
in app_js
)
assert "Dragon Dolly" in app_js
assert "no creature morphing" in app_js
assert "the Plastic Beach island stays centered" in app_js
assert "no camera descent, no push-in, no orbit" in app_js
assert "Ziggy Stardust" in app_js
assert "blue K. West sign" in app_js
assert "wet pavement reflecting a yellow streetlamp" in app_js
assert "ZiggyStardust.jpg" in app_js
assert "A slow aerial orbit around a pastel floating island hotel" not in app_js
assert app_js.index("Dragon Ride") < app_js.index("Dragon Dolly")
assert app_js.index("Ziggy Stardust") < app_js.index("Plastic Beach")
assert app_js.index("Dragon Dolly") < app_js.index("Kid A")
assert "dragon-ride.jpg" in app_js
assert "stageRenderFps" not in app_js
assert 'setStatus("Receiving", "live")' in app_js
assert "decodeQueue.push(" in app_js
assert "receiveChain" not in app_js
assert 'message.type === "chunk_stats"' in app_js
assert "chunkTotal > 0 ? numFrames / chunkTotal" in app_js
assert ".stage-stat" in styles_css
assert ".workspace" in styles_css
assert ".preview-frame" in styles_css
assert ".preview-overlay" in styles_css
assert ".preview-scale-control" in styles_css
assert "--preview-scale" in styles_css
@@ -47,18 +47,6 @@ def test_consistency_gt_urls_are_pinned_to_ci_data_revision():
assert pinned_revision_path in test_utils.SGL_TEST_FILES_SGLANG_CONSISTENCY_GT_BASE
def test_remote_file_exists_returns_false_for_definitive_404(monkeypatch):
class Response:
status_code = 404
def close(self):
pass
monkeypatch.setattr(test_utils.requests, "head", lambda *args, **kwargs: Response())
assert test_utils._remote_file_exists("https://example.com/missing.png") is False
def test_remote_video_gt_candidates_survive_inconclusive_probe(monkeypatch):
monkeypatch.setenv(test_utils.CONSISTENCY_PLATFORM_ENV, "h100")
monkeypatch.setattr(test_utils, "_remote_file_exists", lambda url: None)
@@ -1907,18 +1907,6 @@ class TestCosmos3ModalitySamplingParams(unittest.TestCase):
self.assertEqual(sp.condition_frame_indexes, [0, 1])
self.assertEqual(sp.condition_video_keep, "first")
def test_action_fields_default_none(self):
sp = Cosmos3SamplingParams(prompt="t")
for field in (
"action_mode",
"domain_id",
"domain_name",
"raw_action_dim",
"action_fps",
"action",
):
self.assertIsNone(getattr(sp, field))
class TestCosmos3CaptionMetadata(unittest.TestCase):
"""Structured captions get generation metadata; prose prompts opt out."""
@@ -1,658 +0,0 @@
import importlib.util
import json
import sys
import tempfile
import types
import unittest
from pathlib import Path
from unittest.mock import patch
def _load_benchmark_module(temp_root: Path):
multimodal_gen_root = Path(__file__).resolve().parents[2]
script_path = (
multimodal_gen_root
/ ".claude"
/ "skills"
/ "sglang-diffusion-benchmark-profile"
/ "scripts"
/ "bench_diffusion_denoise.py"
)
fake_env = types.ModuleType("diffusion_skill_env")
fake_env.ensure_dir = lambda path: (
Path(path).mkdir(parents=True, exist_ok=True) or Path(path)
)
fake_env.get_assets_dir = lambda _root: temp_root / "assets"
fake_env.get_output_dir = lambda _kind, _root: temp_root / "outputs"
fake_env.get_repo_root = lambda: temp_root / "repo"
fake_env.pick_idle_gpus = lambda count: list(range(count))
spec = importlib.util.spec_from_file_location(
"test_bench_diffusion_denoise", script_path
)
assert spec is not None and spec.loader is not None
module = importlib.util.module_from_spec(spec)
with patch.dict(sys.modules, {"diffusion_skill_env": fake_env}):
spec.loader.exec_module(module)
return module
def _load_skill_env_module():
script_path = (
Path(__file__).resolve().parents[2]
/ ".claude"
/ "skills"
/ "sglang-diffusion-benchmark-profile"
/ "scripts"
/ "diffusion_skill_env.py"
)
spec = importlib.util.spec_from_file_location(
"test_diffusion_skill_env", script_path
)
assert spec is not None and spec.loader is not None
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
return module
class TestDiffusionBenchmarkSkill(unittest.TestCase):
def test_skill_env_prefers_own_worktree_over_installed_package(self):
module = _load_skill_env_module()
installed = types.ModuleType("sglang")
installed.__file__ = "/sgl-workspace/sglang/python/sglang/__init__.py"
with patch.dict(sys.modules, {"sglang": installed}):
repo_root = module.get_repo_root()
self.assertEqual(repo_root, Path(__file__).resolve().parents[5])
def test_nightly_presets_remain_aligned(self):
with tempfile.TemporaryDirectory() as tmpdir:
module = _load_benchmark_module(Path(tmpdir))
repo_root = Path(__file__).resolve().parents[5]
module.NIGHTLY_CONFIG_PATH = (
repo_root
/ "scripts"
/ "ci"
/ "utils"
/ "diffusion"
/ "comparison_configs.json"
)
self.assertEqual(module.validate_nightly_alignment(), 0)
def test_recent_model_presets_are_eager_by_default(self):
with tempfile.TemporaryDirectory() as tmpdir:
module = _load_benchmark_module(Path(tmpdir))
expected = {
"longcat-image",
"longcat-image-edit",
"longcat-image-edit-turbo",
"qwen-edit-base",
"qwen-image-layered",
"stable-diffusion-3.5-medium",
"sana-video",
"sana-wm-bidirectional",
"sana-wm-streaming",
"lingbot-video-moe",
"lingbot-world",
"lingbot-world-v2",
"fastwan21-t2v-1.3b",
"fasth3-t2va-vsa",
"wan22-t2v-nvfp4",
"krea2-turbo",
"krea2-raw",
"ideogram4-fast",
"ideogram4-instant",
"longlive2-t2v",
"longlive2-i2v",
"fast-hunyuan",
"turbowan21-t2v-1.3b",
"helios-mid",
"helios-distilled",
"joy-echo",
"cosmos3-edge-t2i",
"cosmos3-super-t2v-cfg2tp2",
"cosmos3-super-i2v",
"cosmos3-super-t2i-distilled",
"ltx25",
"ltx25-diffusion-decoder",
}
self.assertTrue(expected.issubset(module.MODELS))
eager_cmd = module.build_sglang_cmd("longcat-image")
self.assertNotIn("--enable-torch-compile", eager_cmd)
self.assertIn("--enable-prompt-rewrite=false", eager_cmd)
self.assertIn("--quality=lossless", eager_cmd)
compiled_cmd = module.build_sglang_cmd("longcat-image", torch_compile=True)
self.assertIn("--enable-torch-compile", compiled_cmd)
longcat_edit_cmd = module.build_sglang_cmd("longcat-image-edit")
self.assertIn(
"--model-path=meituan-longcat/LongCat-Image-Edit",
longcat_edit_cmd,
)
self.assertTrue(
any(arg.startswith("--image-path=") for arg in longcat_edit_cmd)
)
self.assertIn("--enable-prompt-rewrite=false", longcat_edit_cmd)
longcat_edit_bcg_cmd = module.build_sglang_cmd(
"longcat-image-edit", breakable_cuda_graph=True
)
resolution_index = longcat_edit_bcg_cmd.index("--warmup-resolutions")
self.assertEqual(longcat_edit_bcg_cmd[resolution_index + 1], "1264x848")
longcat_edit_turbo_cmd = module.build_sglang_cmd("longcat-image-edit-turbo")
self.assertIn(
"--model-path=meituan-longcat/LongCat-Image-Edit-Turbo",
longcat_edit_turbo_cmd,
)
layered_cmd = module.build_sglang_cmd("qwen-image-layered")
self.assertIn("--model-path=Qwen/Qwen-Image-Layered", layered_cmd)
self.assertIn("--num-frames=4", layered_cmd)
sd35_cmd = module.build_sglang_cmd("stable-diffusion-3.5-medium")
self.assertIn(
"--model-path=stabilityai/stable-diffusion-3.5-medium-diffusers",
sd35_cmd,
)
self.assertIn("stable-diffusion-3.5-medium", module.GATED_MODELS)
h3_cmd = module.build_sglang_cmd("minimax-h3-t2va", torch_compile=True)
self.assertNotIn("--enable-torch-compile", h3_cmd)
fastwan_cmd = module.build_sglang_cmd("fastwan21-t2v-1.3b")
self.assertIn("--num-frames=61", fastwan_cmd)
self.assertIn("--num-inference-steps=3", fastwan_cmd)
self.assertIn("--dit-layerwise-offload=false", fastwan_cmd)
wan_nvfp4_cmd = module.build_sglang_cmd("wan22-t2v-nvfp4")
self.assertIn(
"--model-path=nvidia/Wan2.2-T2V-A14B-Diffusers-NVFP4",
wan_nvfp4_cmd,
)
self.assertIn("--num-frames=81", wan_nvfp4_cmd)
self.assertIn("--dit-layerwise-offload=false", wan_nvfp4_cmd)
self.assertEqual(module.required_gpus_for_model("wan22-t2v-nvfp4"), 1)
krea_raw_cmd = module.build_sglang_cmd("krea2-raw")
self.assertIn("--num-inference-steps=50", krea_raw_cmd)
self.assertIn("--guidance-scale=4.5", krea_raw_cmd)
cosmos_i2v_cmd = module.build_sglang_cmd("cosmos3-super-i2v")
self.assertIn(
"--model-path=nvidia/Cosmos3-Super-Image2Video", cosmos_i2v_cmd
)
self.assertIn("--num-gpus=2", cosmos_i2v_cmd)
self.assertIn("--tp-size=2", cosmos_i2v_cmd)
self.assertIn("--num-frames=81", cosmos_i2v_cmd)
cosmos_cfg_cmd = module.build_sglang_cmd("cosmos3-super-t2v-cfg2tp2")
self.assertIn("--model-path=nvidia/Cosmos3-Super", cosmos_cfg_cmd)
self.assertIn("--num-gpus=4", cosmos_cfg_cmd)
self.assertIn("--tp-size=2", cosmos_cfg_cmd)
sana_wm_dense_cmd = module.build_sglang_cmd("sana-wm-bidirectional")
self.assertIn(
"--model-path=Efficient-Large-Model/SANA-WM_bidirectional",
sana_wm_dense_cmd,
)
self.assertIn("--num-inference-steps=20", sana_wm_dense_cmd)
self.assertNotIn("--streaming", sana_wm_dense_cmd)
sana_wm_streaming_cmd = module.build_sglang_cmd("sana-wm-streaming")
self.assertIn("--streaming", sana_wm_streaming_cmd)
self.assertIn("--refiner-chunked", sana_wm_streaming_cmd)
self.assertIn("--action=w-16,wl-16,l-16", sana_wm_streaming_cmd)
lingbot_world_cmd = module.build_sglang_cmd("lingbot-world")
self.assertIn(
"--model-path=robbyant/lingbot-world-fast-diffusers",
lingbot_world_cmd,
)
self.assertIn("--num-frames=9", lingbot_world_cmd)
self.assertIn("--warmup-mode=off", lingbot_world_cmd)
self.assertIn("--config=", " ".join(lingbot_world_cmd))
self.assertEqual(
module.MODELS["lingbot-world"]["config_overrides"]["actions"],
[["w"] for _ in range(9)],
)
lingbot_world_v2_cmd = module.build_sglang_cmd("lingbot-world-v2")
self.assertIn(
"--model-path=robbyant/lingbot-world-v2-14b-causal-fast-diffusers",
lingbot_world_v2_cmd,
)
self.assertIn("--num-frames=9", lingbot_world_v2_cmd)
self.assertIn("--num-inference-steps=4", lingbot_world_v2_cmd)
ideogram_cmd = module.build_sglang_cmd("ideogram4-instant")
self.assertFalse(
any(arg.startswith("--num-inference-steps") for arg in ideogram_cmd)
)
longlive_i2v_cmd = module.build_sglang_cmd("longlive2-i2v")
self.assertIn("--num-frames=61", longlive_i2v_cmd)
self.assertTrue(
any(arg.startswith("--image-path=") for arg in longlive_i2v_cmd)
)
joy_echo_cmd = module.build_sglang_cmd("joy-echo")
self.assertIn("--num-gpus=2", joy_echo_cmd)
self.assertIn("--ulysses-degree=2", joy_echo_cmd)
config_arg = next(
arg for arg in joy_echo_cmd if arg.startswith("--config=")
)
config = json.loads(Path(config_arg.removeprefix("--config=")).read_text())
self.assertFalse(config["enable_memory_bank"])
def test_quality_and_bcg_comparators_are_explicit_and_exclusive(self):
with tempfile.TemporaryDirectory() as tmpdir:
module = _load_benchmark_module(Path(tmpdir))
high_cmd = module.build_sglang_cmd("longcat-image", quality="high")
self.assertIn("--quality=high", high_cmd)
self.assertNotIn("--enable-breakable-cuda-graph", high_cmd)
extra_high_cmd = module.build_sglang_cmd(
"longcat-image", quality="extra-high"
)
self.assertIn("--quality=extra-high", extra_high_cmd)
self.assertNotIn("--enable-breakable-cuda-graph", extra_high_cmd)
bcg_cmd = module.build_sglang_cmd(
"longcat-image",
breakable_cuda_graph=True,
bcg_text_buckets=[256, 512],
)
self.assertIn("--enable-breakable-cuda-graph", bcg_cmd)
self.assertEqual(
bcg_cmd[bcg_cmd.index("--warmup-resolutions") + 1], "1024x1024"
)
bucket_index = bcg_cmd.index("--bcg-text-buckets")
self.assertEqual(
bcg_cmd[bucket_index + 1 : bucket_index + 3], ["256", "512"]
)
sana_video_bcg_cmd = module.build_sglang_cmd(
"sana-video", breakable_cuda_graph=True
)
self.assertEqual(
sana_video_bcg_cmd[sana_video_bcg_cmd.index("--warmup-num-frames") + 1],
"17",
)
for _, quality, breakable_cuda_graph in module.QUALITY_BCG_ABBA_MATRIX:
module.build_sglang_cmd(
"longcat-image",
quality=quality,
breakable_cuda_graph=breakable_cuda_graph,
bcg_text_buckets=[256, 512] if breakable_cuda_graph else None,
)
with self.assertRaisesRegex(ValueError, "comparators"):
module.build_sglang_cmd(
"longcat-image",
torch_compile=True,
breakable_cuda_graph=True,
)
with self.assertRaisesRegex(ValueError, "requires"):
module.build_sglang_cmd("longcat-image", bcg_text_buckets=[256])
def test_isolated_cache_cleanup_writes_zero_residual_ledger(self):
with tempfile.TemporaryDirectory() as tmpdir:
temp_root = Path(tmpdir)
module = _load_benchmark_module(temp_root)
cache_root = temp_root / "model-caches"
cache_dir = module._prepare_model_cache(
cache_root, "longcat-image", "baseline"
)
weight_path = cache_dir / "huggingface" / "hub" / "model.safetensors"
weight_path.parent.mkdir(parents=True)
weight_path.write_bytes(b"weights")
env = module._model_cache_env(cache_dir)
self.assertTrue(env["HF_HOME"].startswith(str(cache_dir)))
self.assertTrue(env["HF_XET_CACHE"].startswith(str(cache_dir)))
self.assertTrue(env["TRANSFORMERS_CACHE"].startswith(str(cache_dir)))
self.assertTrue(env["MODELSCOPE_CACHE"].startswith(str(cache_dir)))
ledger_path = temp_root / "artifacts" / "cleanup.jsonl"
record = module._cleanup_model_cache(
cache_root,
cache_dir,
ledger_path,
"longcat-image",
"baseline",
"success",
)
self.assertFalse(cache_dir.exists())
self.assertEqual(record["before"]["weight_file_count"], 1)
self.assertEqual(record["after"]["file_count"], 0)
ledger = json.loads(ledger_path.read_text(encoding="utf-8"))
self.assertEqual(ledger["exit_reason"], "success")
self.assertEqual(ledger["after"]["weight_file_count"], 0)
def test_isolated_cache_refuses_to_reuse_existing_run_directory(self):
with tempfile.TemporaryDirectory() as tmpdir:
temp_root = Path(tmpdir)
module = _load_benchmark_module(temp_root)
cache_root = temp_root / "model-caches"
module._prepare_model_cache(cache_root, "sana-video", "baseline")
with self.assertRaises(FileExistsError):
module._prepare_model_cache(cache_root, "sana-video", "baseline")
def test_isolated_cache_seeds_read_only_hf_cache_with_writable_overlay(self):
with tempfile.TemporaryDirectory() as tmpdir:
temp_root = Path(tmpdir)
module = _load_benchmark_module(temp_root)
seed_root = temp_root / "shared-hf"
source_model = seed_root / "hub" / "models--org--model"
source_weight = source_model / "snapshots" / "abc" / "model.safetensors"
source_weight.parent.mkdir(parents=True)
source_weight.write_bytes(b"shared weights")
source_ref = source_model / "refs" / "main"
source_ref.parent.mkdir()
source_ref.write_text("abc")
cache_root = temp_root / "model-caches"
cache_dir = module._prepare_model_cache(
cache_root,
"sana-video",
"baseline",
seed_model_cache_roots=[seed_root],
)
seeded_model = cache_dir / "huggingface" / "hub" / "models--org--model"
self.assertTrue(seeded_model.is_dir())
self.assertFalse(seeded_model.is_symlink())
seeded_weight = seeded_model / "snapshots" / "abc" / "model.safetensors"
self.assertTrue(seeded_weight.is_symlink())
self.assertEqual(
seeded_weight.read_bytes(),
b"shared weights",
)
new_blob = seeded_model / "blobs" / "downloaded"
new_blob.parent.mkdir(exist_ok=True)
new_blob.write_bytes(b"new download")
(seeded_model / "refs" / "main").write_text("new-revision")
self.assertEqual(new_blob.read_bytes(), b"new download")
self.assertEqual(source_ref.read_text(), "abc")
module._cleanup_model_cache(
cache_root,
cache_dir,
temp_root / "cleanup.jsonl",
"sana-video",
"baseline",
"success",
)
self.assertFalse(cache_dir.exists())
self.assertEqual(source_weight.read_bytes(), b"shared weights")
def test_interrupted_run_cleans_isolated_cache_in_finally(self):
with tempfile.TemporaryDirectory() as tmpdir:
temp_root = Path(tmpdir)
module = _load_benchmark_module(temp_root)
cache_root = temp_root / "model-caches"
output_dir = temp_root / "outputs"
output_dir.mkdir()
with (
patch.object(
module, "_run_benchmark_once_impl", side_effect=KeyboardInterrupt
),
self.assertRaises(KeyboardInterrupt),
):
module.run_benchmark_once(
"sana-video",
"baseline",
output_dir,
model_cache_root=cache_root,
cleanup_model_cache=True,
)
self.assertFalse((cache_root / "sana-video-baseline").exists())
ledger = json.loads(
(output_dir / "cleanup.jsonl").read_text(encoding="utf-8")
)
self.assertEqual(ledger["exit_reason"], "interrupted")
self.assertEqual(ledger["after"]["weight_file_count"], 0)
def test_failed_run_is_recorded_as_error_and_cleaned(self):
with tempfile.TemporaryDirectory() as tmpdir:
temp_root = Path(tmpdir)
module = _load_benchmark_module(temp_root)
cache_root = temp_root / "model-caches"
output_dir = temp_root / "outputs"
output_dir.mkdir()
with (
patch.object(
module, "_run_benchmark_once_impl", side_effect=RuntimeError("boom")
),
self.assertRaisesRegex(RuntimeError, "boom"),
):
module.run_benchmark_once(
"sana-video",
"baseline",
output_dir,
model_cache_root=cache_root,
cleanup_model_cache=True,
)
self.assertFalse((cache_root / "sana-video-baseline").exists())
ledger = json.loads(
(output_dir / "cleanup.jsonl").read_text(encoding="utf-8")
)
self.assertEqual(ledger["exit_reason"], "error")
def test_zero_exit_without_artifacts_is_invalid(self):
with tempfile.TemporaryDirectory() as tmpdir:
temp_root = Path(tmpdir)
module = _load_benchmark_module(temp_root)
output_dir = temp_root / "outputs"
output_dir.mkdir()
with patch.object(module.subprocess, "Popen") as popen:
popen.return_value.stdout = iter(())
popen.return_value.wait.return_value = 0
result = module._run_benchmark_once_impl(
"sana-video",
"missing-artifacts",
output_dir,
warmup=False,
cuda_visible_devices="0",
)
command = popen.call_args.args[0]
self.assertIn("--output-path", command)
self.assertIn("--output-file-name", command)
self.assertTrue(result["error"])
self.assertEqual(
result["missing_artifacts"], ["perf dump", "generated output"]
)
def test_mesh_artifacts_are_accepted_and_hashed(self):
with tempfile.TemporaryDirectory() as tmpdir:
temp_root = Path(tmpdir)
module = _load_benchmark_module(temp_root)
output_dir = temp_root / "outputs"
output_dir.mkdir()
def finish_run():
(output_dir / "hunyuan3d-shape_mesh-output.json").write_text(
json.dumps({"total_duration_ms": 1000, "steps": []}),
encoding="utf-8",
)
(output_dir / "hunyuan3d-shape-mesh-output.obj").write_bytes(
b"v 0 0 0\n"
)
return 0
with patch.object(module.subprocess, "Popen") as popen:
popen.return_value.stdout = iter(())
popen.return_value.wait.side_effect = finish_run
result = module._run_benchmark_once_impl(
"hunyuan3d-shape",
"mesh-output",
output_dir,
warmup=False,
cuda_visible_devices="0",
)
self.assertFalse(result["error"])
self.assertEqual(
result["output_artifacts"],
[str(output_dir / "hunyuan3d-shape-mesh-output.obj")],
)
self.assertEqual(len(result["output_sha256"]), 1)
def test_high_bcg_rejects_quality_fusion_mounted_after_capture(self):
with tempfile.TemporaryDirectory() as tmpdir:
temp_root = Path(tmpdir)
module = _load_benchmark_module(temp_root)
output_dir = temp_root / "outputs"
output_dir.mkdir()
with patch.object(module.subprocess, "Popen") as popen:
popen.return_value.stdout = iter(
(
"[Diffusion BCG] captured 3 segment(s)\n",
"Mounted LTX-2 fused RMSNorm+modulate for quality=high\n",
)
)
popen.return_value.wait.return_value = 0
result = module._run_benchmark_once_impl(
"longcat-image",
"bcg-high",
output_dir,
warmup=False,
quality="high",
breakable_cuda_graph=True,
cuda_visible_devices="0",
)
self.assertTrue(result["error"])
self.assertEqual(
result["bcg_invalid_signals"],
[module.BCG_LATE_QUALITY_FUSION_SIGNAL],
)
def test_extra_high_bcg_rejects_quality_fusion_mounted_after_capture(self):
with tempfile.TemporaryDirectory() as tmpdir:
temp_root = Path(tmpdir)
module = _load_benchmark_module(temp_root)
output_dir = temp_root / "outputs"
output_dir.mkdir()
with patch.object(module.subprocess, "Popen") as popen:
popen.return_value.stdout = iter(
(
"[Diffusion BCG] captured 3 segment(s)\n",
"Mounted Qwen fused added-QKV for quality=extra-high\n",
)
)
popen.return_value.wait.return_value = 0
result = module._run_benchmark_once_impl(
"longcat-image",
"bcg-extra-high",
output_dir,
warmup=False,
quality="extra-high",
breakable_cuda_graph=True,
cuda_visible_devices="0",
)
self.assertTrue(result["error"])
self.assertEqual(
result["bcg_invalid_signals"],
[module.BCG_LATE_QUALITY_FUSION_SIGNAL],
)
def test_quality_bcg_matrix_reuses_one_gpu_set_and_cleans_once(self):
with tempfile.TemporaryDirectory() as tmpdir:
temp_root = Path(tmpdir)
module = _load_benchmark_module(temp_root)
cache_root = temp_root / "model-caches"
output_dir = temp_root / "outputs"
output_dir.mkdir()
calls = []
def fake_run(model_key, label, _output_dir, **kwargs):
calls.append((model_key, label, kwargs))
cache_dir = kwargs["model_cache_dir"]
weight_path = cache_dir / "hub" / "model.safetensors"
weight_path.parent.mkdir(parents=True, exist_ok=True)
weight_path.write_bytes(b"weights")
return {"model": model_key, "label": label, "error": False}
with patch.object(module, "_run_benchmark_once_impl", side_effect=fake_run):
results = module.run_quality_bcg_matrix(
"sana-video",
"h200",
output_dir,
model_cache_root=cache_root,
cleanup_model_cache=True,
)
self.assertEqual(len(results), 12)
self.assertEqual(
[
(call[2]["quality"], call[2]["breakable_cuda_graph"])
for call in calls
],
[
(quality, breakable_cuda_graph)
for _, quality, breakable_cuda_graph in module.QUALITY_BCG_ABBA_MATRIX
],
)
self.assertEqual({call[2]["cuda_visible_devices"] for call in calls}, {"0"})
self.assertEqual(
{call[2]["model_cache_dir"] for call in calls},
{calls[0][2]["model_cache_dir"]},
)
self.assertFalse(calls[0][2]["model_cache_dir"].exists())
ledger = json.loads(
(output_dir / "cleanup.jsonl").read_text(encoding="utf-8")
)
self.assertEqual(ledger["exit_reason"], "success")
self.assertEqual(ledger["before"]["weight_file_count"], 1)
self.assertEqual(ledger["after"]["weight_file_count"], 0)
def test_quality_bcg_matrix_rejects_output_hash_mismatch(self):
with tempfile.TemporaryDirectory() as tmpdir:
module = _load_benchmark_module(Path(tmpdir))
results = [
{
"quality": "lossless",
"breakable_cuda_graph": False,
"output_sha256": ["eager"],
"error": False,
},
{
"quality": "lossless",
"breakable_cuda_graph": True,
"output_sha256": ["bcg"],
"error": False,
},
]
module._validate_quality_bcg_output_hashes(results)
self.assertFalse(results[0]["error"])
self.assertTrue(results[1]["error"])
self.assertEqual(
results[1]["output_hash_error"],
"BCG lossless output hash differs from eager",
)
if __name__ == "__main__":
unittest.main()
@@ -617,13 +617,6 @@ class TestStageAffinityAndValidation(_GlobalStageArgsMixin, unittest.TestCase):
pipeline.create_pipeline_stages(pipeline.server_args)
self.assertEqual(list(pipeline._stage_name_mapping.keys()), stage_names)
def test_hunyuan3d_shape_stage_no_longer_stores_model_dtype(self):
pipeline = self._make_hunyuan_pipeline(RoleType.ENCODER, paint_enable=False)
pipeline.create_pipeline_stages(pipeline.server_args)
stage = pipeline._stage_name_mapping["shape_before_denoising"]
self.assertIsInstance(stage, Hunyuan3DShapeBeforeDenoisingStage)
self.assertFalse(hasattr(stage, "model_dtype"))
def test_ltx2_refinement_stage_keeps_class_name_stage_key(self):
stage = object.__new__(LTX2RefinementStage)
self.assertEqual(
@@ -29,17 +29,6 @@ from sglang.multimodal_gen.runtime.utils.nvtx_pytorch_hooks import (
class TestMaybeNvtxRange(unittest.TestCase):
def test_disabled_returns_noop_context_manager(self) -> None:
ran = False
with maybe_nvtx_range("never", enabled=False):
ran = True
self.assertTrue(ran)
def test_disabled_propagates_exception(self) -> None:
with self.assertRaises(RuntimeError):
with maybe_nvtx_range("never", enabled=False):
raise RuntimeError("boom")
def test_disabled_does_not_call_nvtx(self) -> None:
with (
patch.object(nvtx_pytorch_hooks.nvtx, "range_push") as push,
@@ -15,8 +15,6 @@ from sglang.multimodal_gen.runtime.models.encoders.qwen2_5vl import (
Qwen2_5_VLAttention,
Qwen2_5_VLForConditionalGeneration,
_apply_repetition_penalty,
_make_column_linear,
_make_row_linear,
_select_next_token,
)
from sglang.multimodal_gen.runtime.models.encoders.qwen2_5vl_vision import (
@@ -29,14 +27,7 @@ from sglang.multimodal_gen.runtime.pipelines.longcat_image import LongCatImagePi
from sglang.srt.layers.linear import (
ColumnParallelLinear,
ReplicatedLinear,
RowParallelLinear,
)
from sglang.srt.models.qwen2_5_vl import (
Qwen2_5_VisionPatchEmbed,
Qwen2_5_VisionPatchMerger,
Qwen2_5_VLMLP,
)
from sglang.srt.runtime_context import get_parallel
class _StubQwen2_5VL(Qwen2_5_VLForConditionalGeneration):
@@ -79,49 +70,6 @@ class _AttentionRecorder(nn.Module):
return query
def test_native_vision_reuses_srt_modules():
config = SimpleNamespace(
hidden_size=16,
intermediate_size=24,
hidden_act="silu",
num_heads=2,
depth=0,
patch_size=2,
temporal_patch_size=1,
in_channels=3,
spatial_merge_size=2,
out_hidden_size=12,
fullatt_block_indexes=[],
window_size=8,
)
with get_parallel().override(tp_size=1, tp_rank=0):
model = Qwen2_5VLVisionTransformer(config)
mlp = Qwen2_5_VLMLP(
16,
24,
fuse_gate_up=False,
)
fused_mlp = Qwen2_5_VLMLP(16, 24)
assert isinstance(model.patch_embed, Qwen2_5_VisionPatchEmbed)
assert isinstance(model.merger, Qwen2_5_VisionPatchMerger)
assert not mlp.fuse_gate_up
assert isinstance(mlp.gate_proj, ColumnParallelLinear)
assert isinstance(mlp.up_proj, ColumnParallelLinear)
assert mlp.gate_proj.tp_size == mlp.up_proj.tp_size == 1
assert isinstance(mlp.down_proj, ReplicatedLinear)
assert isinstance(fused_mlp.down_proj, RowParallelLinear)
assert mlp.act is not None
assert isinstance(
_make_column_linear(16, 24, bias=False, use_tensor_parallel=False),
ReplicatedLinear,
)
assert isinstance(
_make_row_linear(24, 16, bias=False, use_tensor_parallel=False),
ReplicatedLinear,
)
def test_text_mlp_uses_single_rank_when_intermediate_size_is_not_tp_divisible(
monkeypatch,
):
@@ -12,7 +12,6 @@ from sglang.multimodal_gen.runtime.models.encoders.minimax_h3_qwen3vl import (
)
from sglang.multimodal_gen.runtime.models.encoders.qwen3vl import (
Qwen3VLForConditionalGeneration,
_make_text_rms_norm,
)
from sglang.multimodal_gen.runtime.models.encoders.qwen3vl_vision import (
Qwen3VLVisionRotaryEmbedding,
@@ -20,7 +19,6 @@ from sglang.multimodal_gen.runtime.models.encoders.qwen3vl_vision import (
_vision_cu_seqlens,
_vision_position_ids,
)
from sglang.srt.layers.layernorm import RMSNorm
from sglang.srt.models.qwen3_vl import (
Qwen3VLMoeVisionPatchMerger,
Qwen3VLVisionPatchEmbed,
@@ -48,13 +46,6 @@ def test_native_vision_layout_matches_qwen3_merge_order():
assert cu_seqlens.tolist() == [0, 24, 32, 40]
def test_qwen3vl_text_reuses_srt_rms_norm():
norm = _make_text_rms_norm(16, 1e-6)
assert isinstance(norm, RMSNorm)
assert norm.cast_x_before_out_mul
def test_native_vision_keeps_checkpoint_parameter_names():
config = SimpleNamespace(
hidden_size=16,
@@ -770,11 +770,6 @@ class TestWarmupModeNormalization(unittest.TestCase):
sa = self._resolve(warmup_mode="server")
self.assertEqual(sa.warmup_mode, "server")
def test_defaulted_mode_applies_without_legacy_flags(self):
# Bare `sglang serve` defaults to server-based warmup.
sa = self._resolve(warmup_mode="server")
self.assertEqual(sa.warmup_mode, "server")
def test_resolutions_force_warmup_on(self):
sa = self._resolve(
warmup_mode="off",
@@ -1,36 +0,0 @@
"""Unit tests for server warmup progress reporting."""
import unittest
from unittest.mock import MagicMock, patch
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch
from sglang.multimodal_gen.runtime.server_warmup import SchedulerWarmupMixin
class TestServerWarmupProgress(unittest.TestCase):
def test_ci_progress_uses_scheduler_counter_when_tqdm_is_disabled(self):
scheduler = SchedulerWarmupMixin()
scheduler._show_warmup_progress = True
scheduler._warmup_total = 1
scheduler._warmup_processed = 1
progress_bar = MagicMock(total=1, n=0)
scheduler._warmup_progress_bar = progress_bar
with (
patch(
"sglang.multimodal_gen.runtime.server_warmup._is_ci_log_env",
return_value=True,
),
patch("sglang.multimodal_gen.runtime.server_warmup.logger") as logger,
):
scheduler._advance_warmup_progress_bar(object(), OutputBatch())
logger.info.assert_called_once_with(
"Warmup requests: %s/%s %s", 1, 1, "warmup req"
)
progress_bar.close.assert_called_once_with()
self.assertIsNone(scheduler._warmup_progress_bar)
if __name__ == "__main__":
unittest.main()
@@ -25,11 +25,6 @@ class TestSingleRankDeviceGroup(unittest.TestCase):
new_device_group(ranks, requested)
new_group.assert_called_once_with(ranks, backend=requested)
def test_backend_defaults_to_none_for_multi_rank(self):
with patch(NEW_GROUP_PATH) as new_group:
new_device_group([0, 1])
new_group.assert_called_once_with([0, 1], backend=None)
if __name__ == "__main__":
unittest.main()
@@ -1,46 +1,14 @@
import importlib.util
import unittest
from unittest.mock import MagicMock, patch
import torch
from sglang.multimodal_gen.runtime.layers.attention.backends.sol_attn import (
SolAttnBackend,
SolAttnImpl,
_get_sol_attn_runtime_config,
_parse_layer_ranges,
)
from sglang.multimodal_gen.runtime.platforms.cuda import CudaPlatformBase
from sglang.multimodal_gen.runtime.platforms.interface import AttentionBackendEnum
class FakeCudaPlatform(CudaPlatformBase):
is_sm120_device = False
is_blackwell_device = False
supports_flash_attention = True
@classmethod
def is_sm120(cls):
return cls.is_sm120_device
@classmethod
def is_blackwell(cls):
return cls.is_blackwell_device
@classmethod
def has_device_capability(
cls,
capability: tuple[int, int] | int,
device_id: int = 0,
) -> bool:
return cls.supports_flash_attention
class TestSolAttnBackend(unittest.TestCase):
def test_enum_name(self):
self.assertEqual(str(AttentionBackendEnum.SOL_ATTN), "sol_attn")
self.assertTrue(AttentionBackendEnum.SOL_ATTN.is_sparse)
def test_parse_layer_ranges(self):
self.assertEqual(_parse_layer_ranges("0,1,3-5"), frozenset({0, 1, 3, 4, 5}))
@@ -70,9 +38,6 @@ class TestSolAttnBackend(unittest.TestCase):
):
_get_sol_attn_runtime_config()
def test_backend_head_size(self):
self.assertEqual(SolAttnBackend.get_supported_head_sizes(), [128])
def test_dense_guard_uses_early_steps(self):
impl = SolAttnImpl(
num_heads=8,
@@ -127,19 +92,6 @@ class TestSolAttnBackend(unittest.TestCase):
):
self.assertFalse(impl._should_use_dense())
def test_cuda_resolver(self):
if importlib.util.find_spec("sol_attn") is None:
self.skipTest("sol_attn package is not available")
cls_str = FakeCudaPlatform.get_attn_backend_cls_str(
selected_backend=AttentionBackendEnum.SOL_ATTN,
head_size=128,
dtype=torch.bfloat16,
)
self.assertTrue(cls_str.endswith("SolAttnBackend"))
def test_supports_packed_varlen(self):
self.assertTrue(SolAttnBackend.supports_packed_varlen())
if __name__ == "__main__":
unittest.main()
@@ -156,12 +156,6 @@ def test_strategy_sp1_replicates(monkeypatch):
assert sps.plan_text_strategy(100) == "replicate"
def test_strategy_shard_when_legal(monkeypatch):
_fake_sp(monkeypatch, 2)
assert sps.plan_text_strategy(15) == "shard"
assert sps.plan_text_strategy(16) == "shard"
def test_strategy_replicates_when_padding_spans_multiple_shards(monkeypatch):
_fake_sp(monkeypatch, 8)
assert sps.plan_text_strategy(1) == "replicate"
@@ -34,12 +34,6 @@ class _FakeProjection(nn.Module):
return hidden_states, None
def test_mmgen_clip_reuses_srt_components():
assert mmgen_clip.CLIPEncoder is srt_clip.CLIPEncoder
assert mmgen_clip.CLIPTextEmbeddings is srt_clip.CLIPTextEmbeddings
assert mmgen_clip.CLIPVisionEmbeddings is srt_clip.CLIPVisionEmbeddings
def test_clip_encoder_propagates_causal_semantics():
with (
patch.object(srt_clip, "CLIPAttention", return_value=nn.Identity()) as attn,
@@ -24,7 +24,6 @@ from sglang.multimodal_gen.runtime.layers.attention.backends.subblock_sparse.rou
_snap_up_to_8,
)
from sglang.multimodal_gen.runtime.layers.attention.backends.subblock_sparse_attn import (
SubBlockSparseAttentionBackend,
SubBlockSparseAttentionImpl,
SubBlockSparseSchedule,
_dit_layer_index,
@@ -187,17 +186,6 @@ class TestBudgetGranularity(unittest.TestCase):
class TestSubBlockSparseBackend(unittest.TestCase):
def test_the_advertised_builder_can_be_built(self):
"""`AttentionMetadataBuilder.__init__` is abstract; a builder that does
not override it makes `get_builder_cls()()` a TypeError."""
builder = SubBlockSparseAttentionBackend.get_builder_cls()()
builder.prepare()
metadata = builder.build(current_timestep=7)
self.assertIsInstance(
metadata, SubBlockSparseAttentionBackend.get_metadata_cls()
)
self.assertEqual(metadata.current_timestep, 7)
def test_sm90_adapter_uses_presorted_indices_and_64x64_blocks(self):
captured = {}
@@ -11,9 +11,9 @@ from types import SimpleNamespace
import pytest
from sglang.multimodal_gen.test import run_suite
from sglang.multimodal_gen.test.partitioning import PartitionItem, assign_partition
from sglang.multimodal_gen.test.run_suite import (
from sglang.multimodal_gen.test.runner import diffusion_suite_runner as run_suite
from sglang.multimodal_gen.test.runner.diffusion_suite_runner import (
PartitionAssignment,
build_local_partition_assignment,
)
@@ -71,9 +71,6 @@ class _FakeServerArgs:
def should_start_component_on_cpu(self, _component_name):
return False
def should_use_fsdp_for_component(self, _component_name):
return False
def should_configure_layerwise_offload_for_lazy_component(self, component_name):
return component_name in self.layerwise_components
@@ -615,12 +612,6 @@ class TestVAELoader(unittest.TestCase):
native_load.assert_not_called()
def test_pipeline_config_declares_an_empty_native_only_default(self):
loader = vae_loader.VAELoader()
server_args = _FakeServerArgs(QwenImagePipelineConfig())
self.assertFalse(loader.should_raise_customized_load_error(server_args, "vae"))
def test_backfill_ltx2_audio_vae_latent_stats_maps_official_keys(self):
loaded = {
"per_channel_statistics.mean-of-means": torch.tensor([1.0, 2.0]),
@@ -77,12 +77,6 @@ class _DispatchProbeVAE(ParallelTiledVAE):
class TestVAESpatialParallelDecode(unittest.TestCase):
def test_base_vae_config_defaults_to_auto_parallel_decode(self):
config = VAEConfig()
self.assertTrue(config.use_parallel_decode)
self.assertEqual(config.parallel_decode_mode, "auto")
def test_image_video_vae_configs_default_to_auto_parallel_decode(self):
configs = (
ErnieImageVAEConfig(),
@@ -1,36 +0,0 @@
import unittest
from unittest.mock import patch
from torch import nn
from sglang.multimodal_gen.runtime.models.dits.wanvideo import WanSelfAttention
from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
_WAN = "sglang.multimodal_gen.runtime.models.dits.wanvideo"
class TestWanAttentionBackendRole(unittest.TestCase):
def test_cross_attention_role_is_forwarded_to_usp(self):
with (
patch(f"{_WAN}.ColumnParallelLinear", return_value=nn.Identity()),
patch(f"{_WAN}.RowParallelLinear", return_value=nn.Identity()),
patch(f"{_WAN}.get_tp_world_size", return_value=1),
patch(f"{_WAN}.USPAttention") as usp_attention,
):
WanSelfAttention(
dim=128,
num_heads=1,
qk_norm=False,
is_cross_attention=True,
supported_attention_backends={
AttentionBackendEnum.FA,
AttentionBackendEnum.TORCH_SDPA,
},
)
self.assertTrue(usp_attention.call_args.kwargs["is_cross_attention"])
self.assertTrue(usp_attention.call_args.kwargs["skip_sequence_parallel"])
if __name__ == "__main__":
unittest.main()
@@ -5,8 +5,7 @@ Importing this package is what registers them. An architecture may be claimed
by more than one module here -- one supplies its attention shape, another its
MoE runner -- but two of them must never declare the *same* field for it:
nobody would own that value, and which module supplied it would come down to
the order of the imports below. ``test_model_override_split.py`` forbids the
overlap, which is why this list needs no particular order.
the order of the imports below. Keep each field owned by one family module.
"""
from sglang.srt.arg_groups.model_overrides import cohere2_moe # noqa: F401
@@ -33,8 +33,8 @@ def _minicpm_sala_overrides(server_args: Any, hf_config: Any) -> dict:
"minicpm_flashattn": ("fa4" if get_platform().is_blackwell else "fa3"),
"minicpm_flashinfer": "flashinfer",
}
# Literal keys keep the written-field set statically derivable; a loop
# variable hides it from the census in test_chain_read_ratchet.py.
# Keep the three backend decisions explicit so each resolved field is
# easy to review independently.
dense_attention = dense_backends.get(cfg.attention_backend)
if dense_attention is not None:
overrides["attention_backend"] = dense_attention
+1 -2
View File
@@ -797,8 +797,7 @@ def m3_fp8_attn_gemm_enabled(args) -> bool:
# NOTE: The process-wide ServerArgs is owned by the runtime context
# (sglang.srt.runtime_context). The two functions below are LEGACY shims kept
# for the existing call-sites; they publish/read the same live object by
# reference. Do not add new call-sites — the counts are ratcheted
# (decrease-only) by test/registered/unit/test_legacy_global_ratchet.py.
# reference. Do not add new call-sites.
# Imports are in-function so the two modules stay cycle-free at import time.
@functools.lru_cache(maxsize=1)
def _underscore_field_names() -> frozenset:
-13
View File
@@ -98,19 +98,6 @@ def register_amd_ci(
return None
def register_musa_ci(
est_time: float,
suite: Optional[str] = None,
nightly: bool = False,
disabled: Optional[str] = None,
*,
stage: Optional[str] = None,
runner_config: Optional[str] = None,
):
"""Marker for MUSA CI registration (parsed via AST; runtime no-op)."""
return None
def register_npu_ci(
est_time: float,
suite: Optional[str] = None,
@@ -0,0 +1,35 @@
"""Bridge registered diffusion suites to their case-aware pytest adapter."""
import os
import sys
from pathlib import Path
def _enabled(name: str) -> bool:
return os.environ.get(name, "").lower() in {"1", "true", "yes", "on"}
def run_diffusion_suite(suite: str) -> None:
"""Run one legacy-named diffusion suite without exposing a second CI CLI."""
from sglang.multimodal_gen.test.runner.diffusion_suite_runner import main
args = [sys.argv[0], "--suite", suite]
optional_values = (
("DIFFUSION_PARTITION_ID", "--partition-id"),
("DIFFUSION_TOTAL_PARTITIONS", "--total-partitions"),
("DIFFUSION_PARTITION_PLAN_JSON", "--partition-plan-json"),
("DIFFUSION_PYTEST_FILTER", "--filter"),
)
for environment_name, option in optional_values:
value = os.environ.get(environment_name)
if value:
args.extend([option, value])
if _enabled("DIFFUSION_CONTINUE_ON_ERROR"):
args.append("--continue-on-error")
# Preserve the historical cwd: several diffusion fixtures emit artifacts
# relative to ``python/`` rather than to their source file.
os.chdir(Path(__file__).resolve().parents[4] / "python")
sys.argv = args
main()
@@ -44,7 +44,9 @@ STARTUP_OVERHEAD_SECONDS = 120.0
# Paths relative to repository root
BASELINE_REL_PATH = "python/sglang/multimodal_gen/test/server/perf_baselines"
BASELINE_PLATFORM_ORDER = ("h100", "b200", "5090")
RUN_SUITE_REL_PATH = "python/sglang/multimodal_gen/test/run_suite.py"
RUN_SUITE_REL_PATH = (
"python/sglang/multimodal_gen/test/runner/diffusion_suite_runner.py"
)
USE_NPU_CONFIGS = os.getenv("USE_NPU_CONFIGS", "0").lower() in ("1", "true")
@@ -155,7 +155,8 @@ def print_missing_standalone_estimates_summary(
print("\n" + "=" * 60)
print(
"Add standalone estimate(s) to python/sglang/multimodal_gen/test/run_suite.py"
"Add standalone estimate(s) to "
"python/sglang/multimodal_gen/test/runner/diffusion_suite_runner.py"
)
print("=" * 60)
print("The following standalone file(s) used fallback estimate 300.0s.")
+106
View File
@@ -26,6 +26,7 @@ import glob
import importlib.util
import os
import re
import subprocess
import sys
# Suite names of the form `{stage}-test-{runner_config}` are exactly what the
@@ -38,6 +39,8 @@ _MODERN_SHAPE = re.compile(r"^(.+)-test-(.+)$")
# no suite any workflow invokes and the test silently never runs.
_LEGACY_CUDA_PREFIXES = ("stress",)
_TEST_KINDS = {"unit", "kernel", "e2e", "accuracy", "perf", "stress"}
def _defines_testcase(tree: ast.AST) -> bool:
"""True if the file defines unittest classes, statically or via type()."""
@@ -70,6 +73,99 @@ def _main_runs_tests(tree: ast.Module) -> bool:
return False
def _git_lines(*args: str) -> list[str] | None:
result = subprocess.run(["git", *args], capture_output=True, text=True, check=False)
if result.returncode != 0:
return None
return [line for line in result.stdout.splitlines() if line]
def _changed_registered_files() -> set[str]:
"""Return added, copied, or renamed registered-test destinations."""
lines = _git_lines("diff", "--cached", "--name-status", "--diff-filter=ACR")
if not lines:
base_ref = os.environ.get("GITHUB_BASE_REF", "main")
for candidate in (f"origin/{base_ref}", base_ref):
if _git_lines("rev-parse", "--verify", candidate) is None:
continue
merge_base = _git_lines("merge-base", candidate, "HEAD")
if not merge_base:
continue
lines = _git_lines(
"diff",
"--name-status",
"--diff-filter=ACR",
merge_base[0],
"HEAD",
)
break
selected = set()
for line in lines or []:
fields = line.split("\t")
destination = fields[-1]
if destination.startswith("test/registered/") and destination.endswith(".py"):
selected.add(destination)
return selected
def _contains_call(tree: ast.AST, name: str) -> bool:
return any(
isinstance(node, ast.Call)
and (
(isinstance(node.func, ast.Name) and node.func.id == name)
or (isinstance(node.func, ast.Attribute) and node.func.attr == name)
)
for node in ast.walk(tree)
)
def taxonomy_errors(path: str, registries: list, tree: ast.AST) -> list[str]:
"""Validate the kind/subsystem contract for a newly admitted path."""
parts = path.split("/")
relative_parts = parts[2:] if parts[:2] == ["test", "registered"] else []
if len(relative_parts) < 3 or relative_parts[0] not in _TEST_KINDS:
return [
f"{path}: registered tests must live under "
"test/registered/<kind>/<subsystem>/; kind must be one of "
+ ", ".join(sorted(_TEST_KINDS))
]
kind = relative_parts[0]
errors = []
if kind == "unit":
non_cpu = [r for r in registries if r.backend.name != "CPU"]
if non_cpu:
errors.append(f"{path}: unit tests may register only CPU suites")
if any(r.est_time > 60 for r in registries):
errors.append(f"{path}: unit test est_time must be <= 60 seconds")
if _contains_call(tree, "popen_launch_server"):
errors.append(f"{path}: unit tests may not launch a server")
elif kind == "kernel":
if any("-kernel-" not in (r.effective_suite or "") for r in registries):
errors.append(f"{path}: kernel tests must use a *-kernel-* suite")
elif kind in {"accuracy", "perf"}:
invalid = [
r
for r in registries
if not (r.effective_suite or "").startswith(("nightly-", "weekly-"))
]
if invalid:
errors.append(f"{path}: {kind} tests must use nightly/weekly suites")
elif kind == "stress":
invalid = [
r
for r in registries
if (r.effective_suite or "") != "stress"
and not (r.effective_suite or "").startswith("weekly-")
]
if invalid:
errors.append(f"{path}: stress tests must use stress/weekly suites")
return errors
def main() -> int:
# Import ci_register directly to avoid pulling in all of sglang
spec = importlib.util.spec_from_file_location(
@@ -93,6 +189,8 @@ def main() -> int:
legacy_shape = [] # (file, suite, stage, runner_config) -- has a -test- split
non_dispatchable = [] # (file, suite) -- legacy CUDA suite no workflow invokes
dead_tests = [] # (file) -- TestCase classes that `python3 file.py` never runs
taxonomy_violations = []
changed_files = _changed_registered_files()
for f in files:
try:
registries, _has_main_entry = ci_register.ut_parse_one_file(f)
@@ -106,6 +204,8 @@ def main() -> int:
# `python3 file.py`); the ERROR text below explains the fix.
with open(f, "r", encoding="utf-8") as fh:
tree = ast.parse(fh.read(), filename=f)
if f in changed_files:
taxonomy_violations.extend(taxonomy_errors(f, registries, tree))
if _defines_testcase(tree) and not _main_runs_tests(tree):
dead_tests.append(f)
for r in registries:
@@ -172,6 +272,12 @@ def main() -> int:
print(f" {f}")
print()
exit_code = 1
if taxonomy_violations:
print("ERROR: Registered-test taxonomy violations:")
for error in taxonomy_violations:
print(f" {error}")
print()
exit_code = 1
return exit_code
+1 -2
View File
@@ -18,9 +18,8 @@ bindings/python/
│ └── mini_lb.py
├── tests/ # Python unit tests
│ ├── conftest.py
│ ├── test_validation.py
│ ├── test_arg_parser.py
│ ├── test_router_config.py
│ ├── test_pyo3_binding.py
│ └── test_startup_sequence.py
├── Cargo.toml # Rust package configuration for bindings
├── pyproject.toml # Python package configuration
@@ -1,423 +0,0 @@
"""
Unit tests for router configuration validation and setup.
These tests focus on testing the router configuration logic in isolation,
including validation of configuration parameters and their interactions.
"""
from unittest.mock import MagicMock, patch
import pytest
from sglang_router.launch_router import RouterArgs, launch_router
from sglang_router.router import policy_from_str
from sglang_router.sglang_router_rs import PolicyType
class TestRouterConfigValidation:
"""Test router configuration validation logic."""
def test_valid_basic_config(self):
"""Test that a valid basic configuration passes validation."""
args = RouterArgs(
host="127.0.0.1",
port=30000,
worker_urls=["http://worker1:8000", "http://worker2:8000"],
policy="cache_aware",
)
# Should not raise any exceptions
assert args.host == "127.0.0.1"
assert args.port == 30000
assert args.worker_urls == ["http://worker1:8000", "http://worker2:8000"]
assert args.policy == "cache_aware"
def test_valid_pd_config(self):
"""Test that a valid PD configuration passes validation."""
args = RouterArgs(
host="127.0.0.1",
port=30000,
pd_disaggregation=True,
prefill_urls=[
("http://prefill1:8000", 9000),
("http://prefill2:8000", None),
],
decode_urls=["http://decode1:8001", "http://decode2:8001"],
policy="cache_aware",
)
assert args.pd_disaggregation is True
assert args.prefill_urls == [
("http://prefill1:8000", 9000),
("http://prefill2:8000", None),
]
assert args.decode_urls == ["http://decode1:8001", "http://decode2:8001"]
assert args.policy == "cache_aware"
def test_pd_config_without_urls_allowed(self):
"""Test that PD mode without URLs is now allowed (URLs are optional)."""
args = RouterArgs(
pd_disaggregation=True,
prefill_urls=[],
decode_urls=[],
service_discovery=False,
)
# Should not raise validation error - URLs are now optional
with patch("sglang_router.launch_router.Router") as router_mod:
mock_router_instance = MagicMock()
router_mod.from_args = MagicMock(return_value=mock_router_instance)
# This should succeed without raising an error
launch_router(args)
router_mod.from_args.assert_called_once()
def test_pd_config_with_service_discovery_allows_empty_urls(self):
"""Test that PD mode with service discovery allows empty URLs."""
args = RouterArgs(
pd_disaggregation=True,
prefill_urls=[],
decode_urls=[],
service_discovery=True,
)
# Should not raise validation error when service discovery is enabled
with patch("sglang_router.launch_router.Router") as router_mod:
mock_router_instance = MagicMock()
router_mod.from_args = MagicMock(return_value=mock_router_instance)
launch_router(args)
# Should create router instance via from_args
router_mod.from_args.assert_called_once()
def test_regular_mode_without_workers_allows_empty_urls(self):
"""Test that regular mode allows empty worker URLs."""
args = RouterArgs(worker_urls=[], service_discovery=False)
# Should not raise validation error
with patch("sglang_router.launch_router.Router") as router_mod:
mock_router_instance = MagicMock()
router_mod.from_args = MagicMock(return_value=mock_router_instance)
launch_router(args)
# Should create router instance via from_args
router_mod.from_args.assert_called_once()
def test_cache_threshold_validation(self):
"""Test cache threshold validation."""
# Valid cache threshold
args = RouterArgs(cache_threshold=0.5)
assert args.cache_threshold == 0.5
# Edge cases
args = RouterArgs(cache_threshold=0.0)
assert args.cache_threshold == 0.0
args = RouterArgs(cache_threshold=1.0)
assert args.cache_threshold == 1.0
def test_balance_threshold_validation(self):
"""Test load balancing threshold validation."""
# Valid thresholds
args = RouterArgs(balance_abs_threshold=64, balance_rel_threshold=1.5)
assert args.balance_abs_threshold == 64
assert args.balance_rel_threshold == 1.5
# Edge cases
args = RouterArgs(balance_abs_threshold=0, balance_rel_threshold=1.0)
assert args.balance_abs_threshold == 0
assert args.balance_rel_threshold == 1.0
def test_timeout_validation(self):
"""Test timeout parameter validation."""
# Valid timeouts
args = RouterArgs(
worker_startup_timeout_secs=600,
worker_startup_check_interval=30,
request_timeout_secs=1800,
queue_timeout_secs=60,
)
assert args.worker_startup_timeout_secs == 600
assert args.worker_startup_check_interval == 30
assert args.request_timeout_secs == 1800
assert args.queue_timeout_secs == 60
def test_retry_config_validation(self):
"""Test retry configuration validation."""
# Valid retry config
args = RouterArgs(
retry_max_retries=5,
retry_initial_backoff_ms=50,
retry_max_backoff_ms=30000,
retry_backoff_multiplier=1.5,
retry_jitter_factor=0.2,
disable_retries=False,
)
assert args.retry_max_retries == 5
assert args.retry_initial_backoff_ms == 50
assert args.retry_max_backoff_ms == 30000
assert args.retry_backoff_multiplier == 1.5
assert args.retry_jitter_factor == 0.2
assert args.disable_retries is False
def test_circuit_breaker_config_validation(self):
"""Test circuit breaker configuration validation."""
# Valid circuit breaker config
args = RouterArgs(
cb_failure_threshold=10,
cb_success_threshold=3,
cb_timeout_duration_secs=60,
cb_window_duration_secs=120,
disable_circuit_breaker=False,
)
assert args.cb_failure_threshold == 10
assert args.cb_success_threshold == 3
assert args.cb_timeout_duration_secs == 60
assert args.cb_window_duration_secs == 120
assert args.disable_circuit_breaker is False
def test_health_check_config_validation(self):
"""Test health check configuration validation."""
# Valid health check config
args = RouterArgs(
health_failure_threshold=3,
health_success_threshold=2,
health_check_timeout_secs=5,
health_check_interval_secs=60,
health_check_endpoint="/health",
)
assert args.health_failure_threshold == 3
assert args.health_success_threshold == 2
assert args.health_check_timeout_secs == 5
assert args.health_check_interval_secs == 60
assert args.health_check_endpoint == "/health"
def test_rate_limiting_config_validation(self):
"""Test rate limiting configuration validation."""
# Valid rate limiting config
args = RouterArgs(
max_concurrent_requests=256,
queue_size=100,
queue_timeout_secs=60,
rate_limit_tokens_per_second=100,
)
assert args.max_concurrent_requests == 256
assert args.queue_size == 100
assert args.queue_timeout_secs == 60
assert args.rate_limit_tokens_per_second == 100
def test_service_discovery_config_validation(self):
"""Test service discovery configuration validation."""
# Valid service discovery config
args = RouterArgs(
service_discovery=True,
selector={"app": "worker", "env": "prod"},
service_discovery_port=8080,
service_discovery_namespace="default",
)
assert args.service_discovery is True
assert args.selector == {"app": "worker", "env": "prod"}
assert args.service_discovery_port == 8080
assert args.service_discovery_namespace == "default"
def test_pd_service_discovery_config_validation(self):
"""Test PD service discovery configuration validation."""
# Valid PD service discovery config
args = RouterArgs(
pd_disaggregation=True,
service_discovery=True,
prefill_selector={"app": "prefill"},
decode_selector={"app": "decode"},
bootstrap_port_annotation="sglang.ai/bootstrap-port",
)
assert args.pd_disaggregation is True
assert args.service_discovery is True
assert args.prefill_selector == {"app": "prefill"}
assert args.decode_selector == {"app": "decode"}
assert args.bootstrap_port_annotation == "sglang.ai/bootstrap-port"
def test_prometheus_config_validation(self):
"""Test Prometheus configuration validation."""
# Valid Prometheus config
args = RouterArgs(prometheus_port=29000, prometheus_host="127.0.0.1")
assert args.prometheus_port == 29000
assert args.prometheus_host == "127.0.0.1"
def test_cors_config_validation(self):
"""Test CORS configuration validation."""
# Valid CORS config
args = RouterArgs(
cors_allowed_origins=["http://localhost:3000", "https://example.com"]
)
assert args.cors_allowed_origins == [
"http://localhost:3000",
"https://example.com",
]
def test_tokenizer_config_validation(self):
"""Test tokenizer configuration validation."""
# Note: model_path and tokenizer_path are not available in current RouterArgs
pytest.skip("Tokenizer configuration not available in current implementation")
def test_dp_aware_config_validation(self):
"""Test data parallelism aware configuration validation."""
# Valid DP aware config
args = RouterArgs(dp_aware=True, api_key="test-api-key")
assert args.dp_aware is True
assert args.api_key == "test-api-key"
def test_request_id_headers_validation(self):
"""Test request ID headers configuration validation."""
# Valid request ID headers config
args = RouterArgs(
request_id_headers=["x-request-id", "x-trace-id", "x-correlation-id"]
)
assert args.request_id_headers == [
"x-request-id",
"x-trace-id",
"x-correlation-id",
]
def test_policy_consistency_validation(self):
"""Test policy consistency validation in PD mode."""
# Test with both prefill and decode policies specified
args = RouterArgs(
pd_disaggregation=True,
prefill_urls=[("http://prefill1:8000", None)],
decode_urls=["http://decode1:8001"],
policy="cache_aware",
prefill_policy="power_of_two",
decode_policy="round_robin",
)
# Should not raise validation error
with patch("sglang_router.launch_router.Router") as router_mod:
mock_router_instance = MagicMock()
router_mod.from_args = MagicMock(return_value=mock_router_instance)
launch_router(args)
# Should create router instance via from_args
router_mod.from_args.assert_called_once()
def test_policy_fallback_validation(self):
"""Test policy fallback validation in PD mode."""
# Test with only prefill policy specified
args = RouterArgs(
pd_disaggregation=True,
prefill_urls=[("http://prefill1:8000", None)],
decode_urls=["http://decode1:8001"],
policy="cache_aware",
prefill_policy="power_of_two",
decode_policy=None,
)
# Should not raise validation error
with patch("sglang_router.launch_router.Router") as router_mod:
mock_router_instance = MagicMock()
router_mod.from_args = MagicMock(return_value=mock_router_instance)
launch_router(args)
# Should create router instance via from_args
router_mod.from_args.assert_called_once()
def test_policy_enum_conversion(self):
"""Test policy string to enum conversion."""
# Test all valid policy conversions
assert policy_from_str("random") == PolicyType.Random
assert policy_from_str("round_robin") == PolicyType.RoundRobin
assert policy_from_str("cache_aware") == PolicyType.CacheAware
assert policy_from_str("power_of_two") == PolicyType.PowerOfTwo
def test_invalid_policy_enum_conversion(self):
"""Test invalid policy string to enum conversion."""
with pytest.raises(KeyError):
policy_from_str("invalid_policy")
def test_config_immutability(self):
"""Test that configuration objects are properly immutable."""
args = RouterArgs(
host="127.0.0.1", port=30000, worker_urls=["http://worker1:8000"]
)
# Test that we can't modify the configuration after creation
# (This is more of a design test - dataclasses are mutable by default)
original_host = args.host
args.host = "0.0.0.0"
assert args.host == "0.0.0.0" # Dataclasses are mutable
assert args.host != original_host
def test_config_defaults_consistency(self):
"""Test that configuration defaults are consistent."""
args1 = RouterArgs()
args2 = RouterArgs()
# Both instances should have the same defaults
assert args1.host == args2.host
assert args1.port == args2.port
assert args1.policy == args2.policy
assert args1.worker_urls == args2.worker_urls
assert args1.pd_disaggregation == args2.pd_disaggregation
def test_config_serialization(self):
"""Test that configuration can be serialized/deserialized."""
args = RouterArgs(
host="127.0.0.1",
port=30000,
worker_urls=["http://worker1:8000"],
policy="cache_aware",
cache_threshold=0.5,
)
# Test that we can access all attributes
assert hasattr(args, "host")
assert hasattr(args, "port")
assert hasattr(args, "worker_urls")
assert hasattr(args, "policy")
assert hasattr(args, "cache_threshold")
def test_config_with_none_values(self):
"""Test configuration with None values."""
args = RouterArgs(
api_key=None,
log_dir=None,
log_level=None,
prometheus_port=None,
prometheus_host=None,
request_id_headers=None,
rate_limit_tokens_per_second=None,
service_discovery_namespace=None,
)
# All None values should be preserved
assert args.api_key is None
assert args.log_dir is None
assert args.log_level is None
assert args.prometheus_port is None
assert args.prometheus_host is None
assert args.request_id_headers is None
assert args.rate_limit_tokens_per_second is None
assert args.service_discovery_namespace is None
def test_config_with_empty_lists(self):
"""Test configuration with empty lists."""
args = RouterArgs(
worker_urls=[], prefill_urls=[], decode_urls=[], cors_allowed_origins=[]
)
# All empty lists should be preserved
assert args.worker_urls == []
assert args.prefill_urls == []
assert args.decode_urls == []
assert args.cors_allowed_origins == []
def test_config_with_empty_dicts(self):
"""Test configuration with empty dictionaries."""
args = RouterArgs(selector={}, prefill_selector={}, decode_selector={})
# All empty dictionaries should be preserved
assert args.selector == {}
assert args.prefill_selector == {}
assert args.decode_selector == {}
@@ -1,642 +1,4 @@
"""
Unit tests for startup sequence logic in sglang_router.
These tests focus on testing the startup sequence logic in isolation,
including router initialization, configuration validation, and startup flow.
"""
import logging
from unittest.mock import MagicMock, patch
import pytest
from sglang_router.launch_router import RouterArgs, launch_router
from sglang_router.router import policy_from_str
# Local helper mirroring the router logger setup used in production
def setup_logger():
logger = logging.getLogger("router")
logger.setLevel(logging.INFO)
if not logger.handlers:
formatter = logging.Formatter(
"[Router (Python)] %(asctime)s - %(levelname)s - %(message)s",
datefmt="%Y-%m-%d %H:%M:%S",
)
handler = logging.StreamHandler()
handler.setFormatter(formatter)
logger.addHandler(handler)
return logger
from sglang_router.sglang_router_rs import PolicyType
class TestSetupLogger:
"""Test logger setup functionality."""
def test_setup_logger_returns_logger(self):
"""Test that setup_logger returns a logger instance."""
logger = setup_logger()
assert isinstance(logger, logging.Logger)
assert logger.name == "router"
assert logger.level == logging.INFO
def test_setup_logger_has_handler(self):
"""Test that setup_logger configures a handler."""
logger = setup_logger()
assert len(logger.handlers) > 0
handler = logger.handlers[0]
assert isinstance(handler, logging.StreamHandler)
def test_setup_logger_has_formatter(self):
"""Test that setup_logger configures a formatter."""
logger = setup_logger()
handler = logger.handlers[0]
formatter = handler.formatter
assert formatter is not None
assert "[Router (Python)]" in formatter._fmt
def test_setup_logger_multiple_calls(self):
"""Test that multiple calls to setup_logger work correctly."""
logger1 = setup_logger()
logger2 = setup_logger()
# Should return the same logger instance
assert logger1 is logger2
class TestPolicyFromStr:
"""Test policy string to enum conversion in startup context."""
def test_policy_conversion_in_startup(self):
"""Test policy conversion during startup sequence."""
# Test all valid policies
policies = ["random", "round_robin", "cache_aware", "power_of_two"]
expected_enums = [
PolicyType.Random,
PolicyType.RoundRobin,
PolicyType.CacheAware,
PolicyType.PowerOfTwo,
]
for policy_str, expected_enum in zip(policies, expected_enums):
result = policy_from_str(policy_str)
assert result == expected_enum
def test_invalid_policy_in_startup(self):
"""Test handling of invalid policy during startup."""
with pytest.raises(KeyError):
policy_from_str("invalid_policy")
class TestRouterInitialization:
"""Test router initialization logic."""
def test_router_initialization_basic(self):
"""Test basic router initialization."""
args = RouterArgs(
host="127.0.0.1",
port=30000,
worker_urls=["http://worker1:8000"],
policy="cache_aware",
)
with patch("sglang_router.launch_router.Router") as router_mod:
captured_args = {}
mock_router_instance = MagicMock()
def fake_from_args(router_args):
# capture needed fields from RouterArgs
captured_args.update(
dict(
host=router_args.host,
port=router_args.port,
worker_urls=router_args.worker_urls,
policy=policy_from_str(router_args.policy),
)
)
return mock_router_instance
router_mod.from_args = MagicMock(side_effect=fake_from_args)
result = launch_router(args)
# Verify Router.from_args was called and captured fields match
router_mod.from_args.assert_called_once()
assert captured_args["host"] == "127.0.0.1"
assert captured_args["port"] == 30000
assert captured_args["worker_urls"] == ["http://worker1:8000"]
assert captured_args["policy"] == PolicyType.CacheAware
# Verify router.start() was called
mock_router_instance.start.assert_called_once()
# Function returns None; ensure start was invoked
def test_router_initialization_pd_mode(self):
"""Test router initialization in PD mode."""
args = RouterArgs(
pd_disaggregation=True,
prefill_urls=[("http://prefill1:8000", 9000)],
decode_urls=["http://decode1:8001"],
policy="power_of_two",
)
with patch("sglang_router.launch_router.Router") as router_mod:
captured_args = {}
mock_router_instance = MagicMock()
def fake_from_args(router_args):
captured_args.update(
dict(
pd_disaggregation=router_args.pd_disaggregation,
prefill_urls=router_args.prefill_urls,
decode_urls=router_args.decode_urls,
policy=policy_from_str(router_args.policy),
)
)
return mock_router_instance
router_mod.from_args = MagicMock(side_effect=fake_from_args)
result = launch_router(args)
# Verify Router.from_args was called with PD parameters
router_mod.from_args.assert_called_once()
assert captured_args["pd_disaggregation"] is True
assert captured_args["prefill_urls"] == [("http://prefill1:8000", 9000)]
assert captured_args["decode_urls"] == ["http://decode1:8001"]
assert captured_args["policy"] == PolicyType.PowerOfTwo
# Verify router.start() was called
mock_router_instance.start.assert_called_once()
# Function returns None; ensure start was invoked
def test_router_initialization_with_service_discovery(self):
"""Test router initialization with service discovery."""
args = RouterArgs(
service_discovery=True,
selector={"app": "worker", "env": "prod"},
service_discovery_port=8080,
service_discovery_namespace="default",
)
with patch("sglang_router.launch_router.Router") as router_mod:
captured_args = {}
mock_router_instance = MagicMock()
def fake_from_args(router_args):
captured_args.update(
dict(
service_discovery=router_args.service_discovery,
selector=router_args.selector,
service_discovery_port=router_args.service_discovery_port,
service_discovery_namespace=router_args.service_discovery_namespace,
)
)
return mock_router_instance
router_mod.from_args = MagicMock(side_effect=fake_from_args)
result = launch_router(args)
# Verify Router.from_args was called with service discovery parameters
router_mod.from_args.assert_called_once()
assert captured_args["service_discovery"] is True
assert captured_args["selector"] == {"app": "worker", "env": "prod"}
assert captured_args["service_discovery_port"] == 8080
assert captured_args["service_discovery_namespace"] == "default"
# Verify router.start() was called
mock_router_instance.start.assert_called_once()
# Function returns None; ensure start was invoked
def test_router_initialization_with_retry_config(self):
"""Test router initialization with retry configuration."""
args = RouterArgs(
retry_max_retries=3,
retry_initial_backoff_ms=100,
retry_max_backoff_ms=10000,
retry_backoff_multiplier=2.0,
retry_jitter_factor=0.1,
disable_retries=False,
)
with patch("sglang_router.launch_router.Router") as router_mod:
captured_args = {}
mock_router_instance = MagicMock()
def fake_from_args(router_args):
captured_args.update(
dict(
retry_max_retries=router_args.retry_max_retries,
retry_initial_backoff_ms=router_args.retry_initial_backoff_ms,
retry_max_backoff_ms=router_args.retry_max_backoff_ms,
retry_backoff_multiplier=router_args.retry_backoff_multiplier,
retry_jitter_factor=router_args.retry_jitter_factor,
disable_retries=router_args.disable_retries,
)
)
return mock_router_instance
router_mod.from_args = MagicMock(side_effect=fake_from_args)
result = launch_router(args)
# Verify router was created with retry parameters
router_mod.from_args.assert_called_once()
assert captured_args["retry_max_retries"] == 3
assert captured_args["retry_initial_backoff_ms"] == 100
assert captured_args["retry_max_backoff_ms"] == 10000
assert captured_args["retry_backoff_multiplier"] == 2.0
assert captured_args["retry_jitter_factor"] == 0.1
assert captured_args["disable_retries"] is False
# Verify router.start() was called
mock_router_instance.start.assert_called_once()
# Function returns None; ensure start was invoked
def test_router_initialization_with_circuit_breaker_config(self):
"""Test router initialization with circuit breaker configuration."""
args = RouterArgs(
cb_failure_threshold=5,
cb_success_threshold=2,
cb_timeout_duration_secs=30,
cb_window_duration_secs=60,
disable_circuit_breaker=False,
)
with patch("sglang_router.launch_router.Router") as router_mod:
captured_args = {}
mock_router_instance = MagicMock()
def fake_from_args(router_args):
captured_args.update(
dict(
cb_failure_threshold=router_args.cb_failure_threshold,
cb_success_threshold=router_args.cb_success_threshold,
cb_timeout_duration_secs=router_args.cb_timeout_duration_secs,
cb_window_duration_secs=router_args.cb_window_duration_secs,
disable_circuit_breaker=router_args.disable_circuit_breaker,
)
)
return mock_router_instance
router_mod.from_args = MagicMock(side_effect=fake_from_args)
result = launch_router(args)
# Verify router was created with circuit breaker parameters
router_mod.from_args.assert_called_once()
assert captured_args["cb_failure_threshold"] == 5
assert captured_args["cb_success_threshold"] == 2
assert captured_args["cb_timeout_duration_secs"] == 30
assert captured_args["cb_window_duration_secs"] == 60
assert captured_args["disable_circuit_breaker"] is False
# Verify router.start() was called
mock_router_instance.start.assert_called_once()
# Function returns None; ensure start was invoked
def test_router_initialization_with_rate_limiting_config(self):
"""Test router initialization with rate limiting configuration."""
args = RouterArgs(
max_concurrent_requests=512,
queue_size=200,
queue_timeout_secs=120,
rate_limit_tokens_per_second=100,
)
with patch("sglang_router.launch_router.Router") as router_mod:
captured_args = {}
mock_router_instance = MagicMock()
def fake_from_args(router_args):
captured_args.update(
dict(
max_concurrent_requests=router_args.max_concurrent_requests,
queue_size=router_args.queue_size,
queue_timeout_secs=router_args.queue_timeout_secs,
rate_limit_tokens_per_second=router_args.rate_limit_tokens_per_second,
)
)
return mock_router_instance
router_mod.from_args = MagicMock(side_effect=fake_from_args)
result = launch_router(args)
# Verify router was created with rate limiting parameters
router_mod.from_args.assert_called_once()
assert captured_args["max_concurrent_requests"] == 512
assert captured_args["queue_size"] == 200
assert captured_args["queue_timeout_secs"] == 120
assert captured_args["rate_limit_tokens_per_second"] == 100
# Verify router.start() was called
mock_router_instance.start.assert_called_once()
# Function returns None; ensure start was invoked
def test_router_initialization_with_health_check_config(self):
"""Test router initialization with health check configuration."""
args = RouterArgs(
health_failure_threshold=2,
health_success_threshold=1,
health_check_timeout_secs=3,
health_check_interval_secs=30,
health_check_endpoint="/healthz",
)
with patch("sglang_router.launch_router.Router") as router_mod:
captured_args = {}
mock_router_instance = MagicMock()
def fake_from_args(router_args):
captured_args.update(
dict(
health_failure_threshold=router_args.health_failure_threshold,
health_success_threshold=router_args.health_success_threshold,
health_check_timeout_secs=router_args.health_check_timeout_secs,
health_check_interval_secs=router_args.health_check_interval_secs,
health_check_endpoint=router_args.health_check_endpoint,
)
)
return mock_router_instance
router_mod.from_args = MagicMock(side_effect=fake_from_args)
result = launch_router(args)
# Verify router was created with health check parameters
router_mod.from_args.assert_called_once()
assert captured_args["health_failure_threshold"] == 2
assert captured_args["health_success_threshold"] == 1
assert captured_args["health_check_timeout_secs"] == 3
assert captured_args["health_check_interval_secs"] == 30
assert captured_args["health_check_endpoint"] == "/healthz"
# Verify router.start() was called
mock_router_instance.start.assert_called_once()
# Function returns None; ensure start was invoked
def test_router_initialization_with_prometheus_config(self):
"""Test router initialization with Prometheus configuration."""
args = RouterArgs(prometheus_port=29000, prometheus_host="127.0.0.1")
with patch("sglang_router.launch_router.Router") as router_mod:
captured_args = {}
mock_router_instance = MagicMock()
def fake_from_args(router_args):
captured_args.update(
dict(
prometheus_port=router_args.prometheus_port,
prometheus_host=router_args.prometheus_host,
)
)
return mock_router_instance
router_mod.from_args = MagicMock(side_effect=fake_from_args)
result = launch_router(args)
# Verify router was created with Prometheus parameters
router_mod.from_args.assert_called_once()
assert captured_args["prometheus_port"] == 29000
assert captured_args["prometheus_host"] == "127.0.0.1"
# Verify router.start() was called
mock_router_instance.start.assert_called_once()
# Function returns None; ensure start was invoked
def test_router_initialization_with_cors_config(self):
"""Test router initialization with CORS configuration."""
args = RouterArgs(
cors_allowed_origins=["http://localhost:3000", "https://example.com"]
)
with patch("sglang_router.launch_router.Router") as router_mod:
captured_args = {}
mock_router_instance = MagicMock()
def fake_from_args(router_args):
captured_args.update(
dict(cors_allowed_origins=router_args.cors_allowed_origins)
)
return mock_router_instance
router_mod.from_args = MagicMock(side_effect=fake_from_args)
result = launch_router(args)
# Verify router was created with CORS parameters
router_mod.from_args.assert_called_once()
assert captured_args["cors_allowed_origins"] == [
"http://localhost:3000",
"https://example.com",
]
# Verify router.start() was called
mock_router_instance.start.assert_called_once()
# Function returns None; ensure start was invoked
def test_router_initialization_with_tokenizer_config(self):
"""Test router initialization with tokenizer configuration."""
# Note: model_path and tokenizer_path are not available in current RouterArgs
pytest.skip("Tokenizer configuration not available in current implementation")
class TestStartupValidation:
"""Test startup validation logic."""
def test_pd_mode_validation_during_startup(self):
"""Test PD mode validation during startup."""
# PD mode without URLs is now allowed (URLs are optional)
args = RouterArgs(
pd_disaggregation=True,
prefill_urls=[],
decode_urls=[],
service_discovery=False,
)
# Should not raise validation error - URLs are now optional
with patch("sglang_router.launch_router.Router") as router_mod:
mock_router_instance = MagicMock()
router_mod.from_args = MagicMock(return_value=mock_router_instance)
# This should succeed without raising an error
launch_router(args)
router_mod.from_args.assert_called_once()
def test_pd_mode_with_service_discovery_validation(self):
"""Test PD mode with service discovery validation during startup."""
args = RouterArgs(
pd_disaggregation=True,
prefill_urls=[],
decode_urls=[],
service_discovery=True,
)
# Should not raise validation error
with patch("sglang_router.launch_router.Router") as router_mod:
mock_router_instance = MagicMock()
router_mod.from_args = MagicMock(return_value=mock_router_instance)
result = launch_router(args)
# Should create router instance
router_mod.from_args.assert_called_once()
def test_policy_warning_during_startup(self):
"""Test policy warning during startup in PD mode."""
args = RouterArgs(
pd_disaggregation=True,
prefill_urls=[("http://prefill1:8000", None)],
decode_urls=["http://decode1:8001"],
policy="cache_aware",
prefill_policy="power_of_two",
decode_policy="round_robin",
)
with patch("sglang_router.launch_router.Router") as router_mod:
mock_router_instance = MagicMock()
router_mod.from_args = MagicMock(return_value=mock_router_instance)
# The policy messages are emitted by router_args logger
with patch("sglang_router.router_args.logger") as mock_logger:
result = launch_router(args)
# Should log warning about policy usage
mock_logger.warning.assert_called_once()
warning_call = mock_logger.warning.call_args[0][0]
assert (
"Both --prefill-policy and --decode-policy are specified"
in warning_call
)
# Should create router instance
router_mod.from_args.assert_called_once()
def test_policy_info_during_startup(self):
"""Test policy info logging during startup in PD mode."""
# Test with only prefill policy specified
args = RouterArgs(
pd_disaggregation=True,
prefill_urls=[("http://prefill1:8000", None)],
decode_urls=["http://decode1:8001"],
policy="cache_aware",
prefill_policy="power_of_two",
decode_policy=None,
)
with patch("sglang_router.launch_router.Router") as router_mod:
mock_router_instance = MagicMock()
router_mod.from_args = MagicMock(return_value=mock_router_instance)
# The policy messages are emitted by router_args logger
with patch("sglang_router.router_args.logger") as mock_logger:
result = launch_router(args)
# Should log info about policy usage
mock_logger.info.assert_called_once()
info_call = mock_logger.info.call_args[0][0]
assert "Using --prefill-policy 'power_of_two'" in info_call
assert "and --policy 'cache_aware'" in info_call
# Should create router instance
router_mod.from_args.assert_called_once()
def test_policy_info_decode_only_during_startup(self):
"""Test policy info logging during startup with only decode policy specified."""
args = RouterArgs(
pd_disaggregation=True,
prefill_urls=[("http://prefill1:8000", None)],
decode_urls=["http://decode1:8001"],
policy="cache_aware",
prefill_policy=None,
decode_policy="round_robin",
)
with patch("sglang_router.launch_router.Router") as router_mod:
mock_router_instance = MagicMock()
router_mod.from_args = MagicMock(return_value=mock_router_instance)
# The policy messages are emitted by router_args logger
with patch("sglang_router.router_args.logger") as mock_logger:
result = launch_router(args)
# Should log info about policy usage
mock_logger.info.assert_called_once()
info_call = mock_logger.info.call_args[0][0]
assert "Using --policy 'cache_aware'" in info_call
assert "and --decode-policy 'round_robin'" in info_call
# Should create router instance
router_mod.from_args.assert_called_once()
class TestStartupErrorHandling:
"""Test startup error handling logic."""
def test_router_creation_error_handling(self):
"""Test error handling when router creation fails."""
args = RouterArgs(
host="127.0.0.1", port=30000, worker_urls=["http://worker1:8000"]
)
with patch("sglang_router.launch_router.Router") as router_mod:
# Simulate router creation failure in from_args
router_mod.from_args = MagicMock(
side_effect=Exception("Router creation failed")
)
with patch("sglang_router.launch_router.logger") as mock_logger:
with pytest.raises(Exception, match="Router creation failed"):
launch_router(args)
# Should log error
mock_logger.error.assert_called_once()
error_call = mock_logger.error.call_args[0][0]
assert "Error starting router: Router creation failed" in error_call
def test_router_start_error_handling(self):
"""Test error handling when router start fails."""
args = RouterArgs(
host="127.0.0.1", port=30000, worker_urls=["http://worker1:8000"]
)
with patch("sglang_router.launch_router.Router") as router_mod:
mock_router_instance = MagicMock()
router_mod.from_args = MagicMock(return_value=mock_router_instance)
# Simulate router start failure
mock_router_instance.start.side_effect = Exception("Router start failed")
with patch("sglang_router.launch_router.logger") as mock_logger:
with pytest.raises(Exception, match="Router start failed"):
launch_router(args)
# Should log error
mock_logger.error.assert_called_once()
error_call = mock_logger.error.call_args[0][0]
assert "Error starting router: Router start failed" in error_call
# --- Added unit tests for Router wrapper and launch_server helpers ---
"""Focused runtime tests for the Python router and server-launch helpers."""
def _install_sglang_stubs(monkeypatch):
@@ -827,7 +189,7 @@ def test_launch_server_process_and_cleanup(monkeypatch):
sa = SA()
sa.tp_size = 2
proc = ls.launch_server_process(sa, worker_port=31001, dp_id=3)
ls.launch_server_process(sa, worker_port=31001, dp_id=3)
assert created.get("started") is True
targ, targ_args = created["target"], created["args"]
assert targ is ls.run_server
@@ -911,206 +273,3 @@ def test_launch_server_process_declares_on_a_resolved_record(monkeypatch):
worker = created["args"][0]
assert (worker.port, worker.base_gpu_id, worker.dp_size) == (31002, 6, 1)
assert (parent.port, parent.base_gpu_id, parent.dp_size) == (30000, 0, 4)
def test_validation_error_handling(self):
"""Test error handling when validation fails."""
args = RouterArgs(
pd_disaggregation=True,
prefill_urls=[],
decode_urls=[],
service_discovery=False,
)
with patch("sglang_router.launch_router.logger") as mock_logger:
with pytest.raises(
ValueError, match="PD disaggregation mode requires --prefill"
):
launch_router(args)
# Should log error for validation failures
mock_logger.error.assert_called_once()
class TestStartupFlow:
"""Test complete startup flow."""
def test_complete_startup_flow_basic(self):
"""Test complete startup flow for basic configuration."""
args = RouterArgs(
host="127.0.0.1",
port=30000,
worker_urls=["http://worker1:8000", "http://worker2:8000"],
policy="cache_aware",
cache_threshold=0.5,
balance_abs_threshold=32,
balance_rel_threshold=1.5,
)
with patch("sglang_router.launch_router.Router") as router_mod:
mock_router_instance = MagicMock()
router_mod.from_args = MagicMock(return_value=mock_router_instance)
result = launch_router(args)
# Verify complete flow
router_mod.from_args.assert_called_once()
mock_router_instance.start.assert_called_once()
def test_complete_startup_flow_pd_mode(self):
"""Test complete startup flow for PD mode configuration."""
args = RouterArgs(
pd_disaggregation=True,
prefill_urls=[
("http://prefill1:8000", 9000),
("http://prefill2:8000", None),
],
decode_urls=["http://decode1:8001", "http://decode2:8001"],
policy="power_of_two",
prefill_policy="cache_aware",
decode_policy="round_robin",
)
with patch("sglang_router.launch_router.Router") as router_mod:
mock_router_instance = MagicMock()
router_mod.from_args = MagicMock(return_value=mock_router_instance)
with patch("sglang_router.router_args.logger") as mock_logger:
result = launch_router(args)
# Verify complete flow
router_mod.from_args.assert_called_once()
mock_router_instance.start.assert_called_once()
# Verify policy warning was logged
mock_logger.warning.assert_called_once()
def test_complete_startup_flow_with_all_features(self):
"""Test complete startup flow with all features enabled."""
args = RouterArgs(
host="0.0.0.0",
port=30001,
worker_urls=["http://worker1:8000"],
policy="round_robin",
service_discovery=True,
selector={"app": "worker"},
service_discovery_port=8080,
service_discovery_namespace="default",
dp_aware=True,
api_key="test-key",
log_dir="/tmp/logs",
log_level="debug",
prometheus_port=29000,
prometheus_host="0.0.0.0",
request_id_headers=["x-request-id", "x-trace-id"],
request_timeout_secs=1200,
max_concurrent_requests=512,
queue_size=200,
queue_timeout_secs=120,
rate_limit_tokens_per_second=100,
cors_allowed_origins=["http://localhost:3000"],
retry_max_retries=3,
retry_initial_backoff_ms=100,
retry_max_backoff_ms=10000,
retry_backoff_multiplier=2.0,
retry_jitter_factor=0.1,
cb_failure_threshold=5,
cb_success_threshold=2,
cb_timeout_duration_secs=30,
cb_window_duration_secs=60,
health_failure_threshold=2,
health_success_threshold=1,
health_check_timeout_secs=3,
health_check_interval_secs=30,
health_check_endpoint="/healthz",
)
with patch("sglang_router.launch_router.Router") as router_mod:
captured_args = {}
mock_router_instance = MagicMock()
def fake_from_args(router_args):
captured_args.update(
dict(
host=router_args.host,
port=router_args.port,
worker_urls=router_args.worker_urls,
policy=policy_from_str(router_args.policy),
service_discovery=router_args.service_discovery,
selector=router_args.selector,
service_discovery_port=router_args.service_discovery_port,
service_discovery_namespace=router_args.service_discovery_namespace,
dp_aware=router_args.dp_aware,
api_key=router_args.api_key,
log_dir=router_args.log_dir,
log_level=router_args.log_level,
prometheus_port=router_args.prometheus_port,
prometheus_host=router_args.prometheus_host,
request_id_headers=router_args.request_id_headers,
request_timeout_secs=router_args.request_timeout_secs,
max_concurrent_requests=router_args.max_concurrent_requests,
queue_size=router_args.queue_size,
queue_timeout_secs=router_args.queue_timeout_secs,
rate_limit_tokens_per_second=router_args.rate_limit_tokens_per_second,
cors_allowed_origins=router_args.cors_allowed_origins,
retry_max_retries=router_args.retry_max_retries,
retry_initial_backoff_ms=router_args.retry_initial_backoff_ms,
retry_max_backoff_ms=router_args.retry_max_backoff_ms,
retry_backoff_multiplier=router_args.retry_backoff_multiplier,
retry_jitter_factor=router_args.retry_jitter_factor,
cb_failure_threshold=router_args.cb_failure_threshold,
cb_success_threshold=router_args.cb_success_threshold,
cb_timeout_duration_secs=router_args.cb_timeout_duration_secs,
cb_window_duration_secs=router_args.cb_window_duration_secs,
health_failure_threshold=router_args.health_failure_threshold,
health_success_threshold=router_args.health_success_threshold,
health_check_timeout_secs=router_args.health_check_timeout_secs,
health_check_interval_secs=router_args.health_check_interval_secs,
health_check_endpoint=router_args.health_check_endpoint,
)
)
return mock_router_instance
router_mod.from_args = MagicMock(side_effect=fake_from_args)
result = launch_router(args)
# Verify complete flow
router_mod.from_args.assert_called_once()
mock_router_instance.start.assert_called_once()
# Verify key parameters were propagated into RouterArgs
assert captured_args["host"] == "0.0.0.0"
assert captured_args["port"] == 30001
assert captured_args["worker_urls"] == ["http://worker1:8000"]
assert captured_args["policy"] == PolicyType.RoundRobin
assert captured_args["service_discovery"] is True
assert captured_args["selector"] == {"app": "worker"}
assert captured_args["service_discovery_port"] == 8080
assert captured_args["service_discovery_namespace"] == "default"
assert captured_args["dp_aware"] is True
assert captured_args["api_key"] == "test-key"
assert captured_args["log_dir"] == "/tmp/logs"
assert captured_args["log_level"] == "debug"
assert captured_args["prometheus_port"] == 29000
assert captured_args["prometheus_host"] == "0.0.0.0"
assert captured_args["request_id_headers"] == ["x-request-id", "x-trace-id"]
assert captured_args["request_timeout_secs"] == 1200
assert captured_args["max_concurrent_requests"] == 512
assert captured_args["queue_size"] == 200
assert captured_args["queue_timeout_secs"] == 120
assert captured_args["rate_limit_tokens_per_second"] == 100
assert captured_args["cors_allowed_origins"] == ["http://localhost:3000"]
assert captured_args["retry_max_retries"] == 3
assert captured_args["retry_initial_backoff_ms"] == 100
assert captured_args["retry_max_backoff_ms"] == 10000
assert captured_args["retry_backoff_multiplier"] == 2.0
assert captured_args["retry_jitter_factor"] == 0.1
assert captured_args["cb_failure_threshold"] == 5
assert captured_args["cb_success_threshold"] == 2
assert captured_args["cb_timeout_duration_secs"] == 30
assert captured_args["cb_window_duration_secs"] == 60
assert captured_args["health_failure_threshold"] == 2
assert captured_args["health_success_threshold"] == 1
assert captured_args["health_check_timeout_secs"] == 3
assert captured_args["health_check_interval_secs"] == 30
assert captured_args["health_check_endpoint"] == "/healthz"
@@ -1,509 +0,0 @@
"""
Unit tests for validation logic in sglang_router.
These tests focus on testing the validation logic in isolation,
including parameter validation, URL validation, and configuration validation.
"""
from unittest.mock import MagicMock, patch
import pytest
from sglang_router.launch_router import RouterArgs, launch_router
class TestURLValidation:
"""Test URL validation logic."""
def test_valid_worker_urls(self):
"""Test validation of valid worker URLs."""
valid_urls = [
"http://worker1:8000",
"https://worker2:8000",
"http://localhost:8000",
"http://127.0.0.1:8000",
"http://192.168.1.100:8000",
"http://worker.example.com:8000",
]
for url in valid_urls:
args = RouterArgs(worker_urls=[url])
# Should not raise any validation errors
assert url in args.worker_urls
def test_valid_prefill_urls(self):
"""Test validation of valid prefill URLs."""
valid_prefill_urls = [
("http://prefill1:8000", 9000),
("https://prefill2:8000", None),
("http://localhost:8000", 9000),
("http://127.0.0.1:8000", None),
]
for url, bootstrap_port in valid_prefill_urls:
args = RouterArgs(prefill_urls=[(url, bootstrap_port)])
# Should not raise any validation errors
assert (url, bootstrap_port) in args.prefill_urls
def test_valid_decode_urls(self):
"""Test validation of valid decode URLs."""
valid_decode_urls = [
"http://decode1:8001",
"https://decode2:8001",
"http://localhost:8001",
"http://127.0.0.1:8001",
]
for url in valid_decode_urls:
args = RouterArgs(decode_urls=[url])
# Should not raise any validation errors
assert url in args.decode_urls
def test_malformed_urls(self):
"""Test handling of malformed URLs."""
# Note: The current implementation doesn't validate URL format
# This test documents the current behavior
malformed_urls = [
"not-a-url",
"ftp://worker1:8000", # Wrong protocol
"http://", # Missing host
":8000", # Missing protocol and host
"http://worker1", # Missing port
]
for url in malformed_urls:
args = RouterArgs(worker_urls=[url])
# Currently, malformed URLs are accepted
# This might be something to improve in the future
assert url in args.worker_urls
class TestPortValidation:
"""Test port validation logic."""
def test_valid_ports(self):
"""Test validation of valid port numbers."""
valid_ports = [1, 80, 8000, 30000, 65535]
for port in valid_ports:
args = RouterArgs(port=port)
assert args.port == port
def test_invalid_ports(self):
"""Test handling of invalid port numbers."""
# Note: The current implementation doesn't validate port ranges
# This test documents the current behavior
invalid_ports = [0, -1, 65536, 70000]
for port in invalid_ports:
args = RouterArgs(port=port)
# Currently, invalid ports are accepted
# This might be something to improve in the future
assert args.port == port
def test_bootstrap_port_validation(self):
"""Test validation of bootstrap ports in PD mode."""
valid_bootstrap_ports = [1, 80, 9000, 30000, 65535, None]
for bootstrap_port in valid_bootstrap_ports:
args = RouterArgs(prefill_urls=[("http://prefill1:8000", bootstrap_port)])
assert args.prefill_urls[0][1] == bootstrap_port
class TestParameterValidation:
"""Test parameter validation logic."""
def test_cache_threshold_validation(self):
"""Test cache threshold parameter validation."""
# Valid cache thresholds
valid_thresholds = [0.0, 0.1, 0.5, 0.9, 1.0]
for threshold in valid_thresholds:
args = RouterArgs(cache_threshold=threshold)
assert args.cache_threshold == threshold
def test_balance_threshold_validation(self):
"""Test load balancing threshold parameter validation."""
# Valid absolute thresholds
valid_abs_thresholds = [0, 1, 32, 64, 128, 1000]
for threshold in valid_abs_thresholds:
args = RouterArgs(balance_abs_threshold=threshold)
assert args.balance_abs_threshold == threshold
# Valid relative thresholds
valid_rel_thresholds = [1.0, 1.1, 1.5, 2.0, 10.0]
for threshold in valid_rel_thresholds:
args = RouterArgs(balance_rel_threshold=threshold)
assert args.balance_rel_threshold == threshold
def test_timeout_validation(self):
"""Test timeout parameter validation."""
# Valid timeouts
valid_timeouts = [1, 30, 60, 300, 600, 1800, 3600]
for timeout in valid_timeouts:
args = RouterArgs(
worker_startup_timeout_secs=timeout,
worker_startup_check_interval=timeout,
request_timeout_secs=timeout,
queue_timeout_secs=timeout,
)
assert args.worker_startup_timeout_secs == timeout
assert args.worker_startup_check_interval == timeout
assert args.request_timeout_secs == timeout
assert args.queue_timeout_secs == timeout
def test_retry_parameter_validation(self):
"""Test retry parameter validation."""
# Valid retry parameters
valid_retry_counts = [0, 1, 3, 5, 10]
for count in valid_retry_counts:
args = RouterArgs(retry_max_retries=count)
assert args.retry_max_retries == count
# Valid backoff parameters
valid_backoff_ms = [1, 50, 100, 1000, 30000]
for backoff in valid_backoff_ms:
args = RouterArgs(
retry_initial_backoff_ms=backoff, retry_max_backoff_ms=backoff
)
assert args.retry_initial_backoff_ms == backoff
assert args.retry_max_backoff_ms == backoff
# Valid multiplier parameters
valid_multipliers = [1.0, 1.5, 2.0, 3.0]
for multiplier in valid_multipliers:
args = RouterArgs(retry_backoff_multiplier=multiplier)
assert args.retry_backoff_multiplier == multiplier
# Valid jitter parameters
valid_jitter = [0.0, 0.1, 0.2, 0.5]
for jitter in valid_jitter:
args = RouterArgs(retry_jitter_factor=jitter)
assert args.retry_jitter_factor == jitter
def test_circuit_breaker_parameter_validation(self):
"""Test circuit breaker parameter validation."""
# Valid failure thresholds
valid_failure_thresholds = [1, 3, 5, 10, 20]
for threshold in valid_failure_thresholds:
args = RouterArgs(cb_failure_threshold=threshold)
assert args.cb_failure_threshold == threshold
# Valid success thresholds
valid_success_thresholds = [1, 2, 3, 5]
for threshold in valid_success_thresholds:
args = RouterArgs(cb_success_threshold=threshold)
assert args.cb_success_threshold == threshold
# Valid timeout durations
valid_timeouts = [10, 30, 60, 120, 300]
for timeout in valid_timeouts:
args = RouterArgs(
cb_timeout_duration_secs=timeout, cb_window_duration_secs=timeout
)
assert args.cb_timeout_duration_secs == timeout
assert args.cb_window_duration_secs == timeout
def test_health_check_parameter_validation(self):
"""Test health check parameter validation."""
# Valid failure thresholds
valid_failure_thresholds = [1, 2, 3, 5, 10]
for threshold in valid_failure_thresholds:
args = RouterArgs(health_failure_threshold=threshold)
assert args.health_failure_threshold == threshold
# Valid success thresholds
valid_success_thresholds = [1, 2, 3, 5]
for threshold in valid_success_thresholds:
args = RouterArgs(health_success_threshold=threshold)
assert args.health_success_threshold == threshold
# Valid timeouts and intervals
valid_times = [1, 5, 10, 30, 60, 120]
for time_val in valid_times:
args = RouterArgs(
health_check_timeout_secs=time_val, health_check_interval_secs=time_val
)
assert args.health_check_timeout_secs == time_val
assert args.health_check_interval_secs == time_val
def test_rate_limiting_parameter_validation(self):
"""Test rate limiting parameter validation."""
# Valid concurrent request limits
valid_limits = [1, 10, 64, 256, 512, 1000]
for limit in valid_limits:
args = RouterArgs(max_concurrent_requests=limit)
assert args.max_concurrent_requests == limit
# Valid queue sizes
valid_queue_sizes = [0, 10, 50, 100, 500, 1000]
for size in valid_queue_sizes:
args = RouterArgs(queue_size=size)
assert args.queue_size == size
# Valid token rates
valid_rates = [1, 10, 50, 100, 500, 1000]
for rate in valid_rates:
args = RouterArgs(rate_limit_tokens_per_second=rate)
assert args.rate_limit_tokens_per_second == rate
def test_tree_size_validation(self):
"""Test tree size parameter validation."""
# Valid tree sizes (powers of 2)
valid_sizes = [2**10, 2**20, 2**24, 2**26, 2**28, 2**30]
for size in valid_sizes:
args = RouterArgs(max_tree_size=size)
assert args.max_tree_size == size
def test_payload_size_validation(self):
"""Test payload size parameter validation."""
# Valid payload sizes
valid_sizes = [
1024, # 1KB
1024 * 1024, # 1MB
10 * 1024 * 1024, # 10MB
100 * 1024 * 1024, # 100MB
512 * 1024 * 1024, # 512MB
1024 * 1024 * 1024, # 1GB
]
for size in valid_sizes:
args = RouterArgs(max_payload_size=size)
assert args.max_payload_size == size
class TestConfigurationValidation:
"""Test configuration validation logic."""
def test_pd_mode_validation(self):
"""Test PD mode configuration validation."""
# Valid PD configuration
args = RouterArgs(
pd_disaggregation=True,
prefill_urls=[("http://prefill1:8000", 9000)],
decode_urls=["http://decode1:8001"],
)
assert args.pd_disaggregation is True
assert len(args.prefill_urls) > 0
assert len(args.decode_urls) > 0
def test_service_discovery_validation(self):
"""Test service discovery configuration validation."""
# Valid service discovery configuration
args = RouterArgs(
service_discovery=True,
selector={"app": "worker", "env": "prod"},
service_discovery_port=8080,
service_discovery_namespace="default",
)
assert args.service_discovery is True
assert args.selector == {"app": "worker", "env": "prod"}
assert args.service_discovery_port == 8080
assert args.service_discovery_namespace == "default"
def test_pd_service_discovery_validation(self):
"""Test PD service discovery configuration validation."""
# Valid PD service discovery configuration
args = RouterArgs(
pd_disaggregation=True,
service_discovery=True,
prefill_selector={"app": "prefill"},
decode_selector={"app": "decode"},
)
assert args.pd_disaggregation is True
assert args.service_discovery is True
assert args.prefill_selector == {"app": "prefill"}
assert args.decode_selector == {"app": "decode"}
def test_policy_validation(self):
"""Test policy configuration validation."""
# Valid policies
valid_policies = ["random", "round_robin", "cache_aware", "power_of_two"]
for policy in valid_policies:
args = RouterArgs(policy=policy)
assert args.policy == policy
def test_pd_policy_validation(self):
"""Test PD policy configuration validation."""
# Valid PD policies
valid_policies = ["random", "round_robin", "cache_aware", "power_of_two"]
for prefill_policy in valid_policies:
for decode_policy in valid_policies:
args = RouterArgs(
pd_disaggregation=True,
prefill_urls=[("http://prefill1:8000", None)],
decode_urls=["http://decode1:8001"],
prefill_policy=prefill_policy,
decode_policy=decode_policy,
)
assert args.prefill_policy == prefill_policy
assert args.decode_policy == decode_policy
def test_cors_validation(self):
"""Test CORS configuration validation."""
# Valid CORS origins
valid_origins = [
[],
["http://localhost:3000"],
["https://example.com"],
["http://localhost:3000", "https://example.com"],
["*"], # Wildcard (if supported)
]
for origins in valid_origins:
args = RouterArgs(cors_allowed_origins=origins)
assert args.cors_allowed_origins == origins
def test_logging_validation(self):
"""Test logging configuration validation."""
# Valid log levels
valid_log_levels = ["debug", "info", "warning", "error", "critical"]
for level in valid_log_levels:
args = RouterArgs(log_level=level)
assert args.log_level == level
def test_prometheus_validation(self):
"""Test Prometheus configuration validation."""
# Valid Prometheus configuration
args = RouterArgs(prometheus_port=29000, prometheus_host="127.0.0.1")
assert args.prometheus_port == 29000
assert args.prometheus_host == "127.0.0.1"
def test_tokenizer_validation(self):
"""Test tokenizer configuration validation."""
# Note: model_path and tokenizer_path are not available in current RouterArgs
pytest.skip("Tokenizer configuration not available in current implementation")
def test_request_id_headers_validation(self):
"""Test request ID headers configuration validation."""
# Valid request ID headers
valid_headers = [
["x-request-id"],
["x-request-id", "x-trace-id"],
["x-request-id", "x-trace-id", "x-correlation-id"],
["custom-header"],
]
for headers in valid_headers:
args = RouterArgs(request_id_headers=headers)
assert args.request_id_headers == headers
class TestLaunchValidation:
"""Test launch-time validation logic."""
def test_pd_mode_allows_empty_urls(self):
"""Test that PD mode now allows empty URLs (URLs are optional)."""
# PD mode without URLs is now allowed
args = RouterArgs(
pd_disaggregation=True,
prefill_urls=[],
decode_urls=[],
service_discovery=False,
)
# Should not raise validation error - URLs are now optional
with patch("sglang_router.launch_router.Router") as router_mod:
mock_router_instance = MagicMock()
router_mod.from_args = MagicMock(return_value=mock_router_instance)
# This should succeed without raising an error
launch_router(args)
router_mod.from_args.assert_called_once()
def test_pd_mode_with_service_discovery_allows_empty_urls(self):
"""Test that PD mode with service discovery allows empty URLs."""
args = RouterArgs(
pd_disaggregation=True,
prefill_urls=[],
decode_urls=[],
service_discovery=True,
)
# Should not raise validation error
with patch("sglang_router.launch_router.Router") as router_mod:
mock_router_instance = MagicMock()
router_mod.from_args = MagicMock(return_value=mock_router_instance)
launch_router(args)
# Should create router instance via from_args
router_mod.from_args.assert_called_once()
def test_regular_mode_allows_empty_worker_urls(self):
"""Test that regular mode allows empty worker URLs."""
args = RouterArgs(worker_urls=[], service_discovery=False)
# Should not raise validation error
with patch("sglang_router.launch_router.Router") as router_mod:
mock_router_instance = MagicMock()
router_mod.from_args = MagicMock(return_value=mock_router_instance)
launch_router(args)
# Should create router instance via from_args
router_mod.from_args.assert_called_once()
def test_launch_with_valid_config(self):
"""Test launching with valid configuration."""
args = RouterArgs(
host="127.0.0.1",
port=30000,
worker_urls=["http://worker1:8000"],
policy="cache_aware",
)
# Should not raise validation error
with patch("sglang_router.launch_router.Router") as router_mod:
mock_router_instance = MagicMock()
router_mod.from_args = MagicMock(return_value=mock_router_instance)
launch_router(args)
# Should create router instance via from_args
router_mod.from_args.assert_called_once()
def test_launch_with_pd_config(self):
"""Test launching with valid PD configuration."""
args = RouterArgs(
pd_disaggregation=True,
prefill_urls=[("http://prefill1:8000", 9000)],
decode_urls=["http://decode1:8001"],
policy="cache_aware",
)
# Should not raise validation error
with patch("sglang_router.launch_router.Router") as router_mod:
mock_router_instance = MagicMock()
router_mod.from_args = MagicMock(return_value=mock_router_instance)
launch_router(args)
# Should create router instance via from_args
router_mod.from_args.assert_called_once()
def test_launch_with_service_discovery_config(self):
"""Test launching with valid service discovery configuration."""
args = RouterArgs(
service_discovery=True,
selector={"app": "worker"},
service_discovery_port=8080,
)
# Should not raise validation error
with patch("sglang_router.launch_router.Router") as router_mod:
mock_router_instance = MagicMock()
router_mod.from_args = MagicMock(return_value=mock_router_instance)
launch_router(args)
# Should create router instance via from_args
router_mod.from_args.assert_called_once()
+37 -4
View File
@@ -72,9 +72,28 @@ Parameters: `est_time` (seconds), `stage` + `runner_config` (target stage and ru
Keep `est_time`, `stage`, `runner_config` as **literal values**`run_suite.py` collects them by AST parsing.
JIT kernel correctness tests and benchmarks live under `test/registered/jit/`, same as other registered tests (their helpers stay alongside the kernel source under `python/sglang/kernels/jit/` and are imported by absolute path):
- Correctness tests: `test/registered/jit/test_*.py``base-b-kernel-unit-test-1-gpu-large`
- Benchmarks: `test/registered/jit/benchmark/bench_*.py``base-b-kernel-benchmark-test-1-gpu-large`
New and renamed tests use this layout:
```text
test/registered/<kind>/<subsystem>/test_*.py
```
`<kind>` is one of `unit`, `kernel`, `e2e`, `accuracy`, `perf`, or `stress`.
Hardware is expressed by one or more `register_*_ci` calls, never by creating a
new top-level hardware directory. The admission checker applies the layout and
kind/suite contract incrementally while legacy paths are migrated.
Diffusion workflows also enter through `test/run_suite.py`; registered bridge
files preserve their case-level pytest partitioning until the remaining
diffusion cases are moved out of the package test-support tree.
New JIT kernel correctness tests and benchmarks live under
`test/registered/kernel/jit/`; legacy `test/registered/jit/` files are migrated
incrementally. Helpers stay alongside the kernel source under
`python/sglang/kernels/jit/` and are imported by absolute path:
- Correctness tests: `test/registered/kernel/jit/test_*.py``base-b-kernel-unit-test-1-gpu-large`
- Benchmarks: `test/registered/kernel/jit/benchmark/bench_*.py``base-b-kernel-benchmark-test-1-gpu-large`
## Choosing a Suite
@@ -94,6 +113,20 @@ Use the lightest suite that meets your test's needs. Full suite tables are in th
See the [write-sglang-test skill](../.claude/skills/write-sglang-test/SKILL.md) for templates, fixtures, model selection, and a complete checklist.
Before adding a registered test, identify the production change that would make
it fail. Prefer extending an existing fixture/server launch over adding another
file. The incremental admission check applies these ratchets to new or modified
registered tests:
- Temporary `disabled=` registrations and unconditional skips must reference an
issue and include `until YYYY-MM-DD`; expired entries fail lint.
- A file registered on CUDA plus another accelerator must place a nearby
`backend-specific:` comment above the extra registration and name the path or
failure mode that only that backend can catch.
- Default PR registrations are limited to 1,200 estimated weighted accelerator-seconds
per backend (`est_time * GPU count`). Move larger matrices to extra/nightly,
or document a nearby `ci-cost-override:` rationale.
## Multi-Hardware Backends
This README mostly describes the NVIDIA GPU CI pipeline. Other hardware backends (AMD, NPU) follow the same practices and use the multi-backend registry system. A scheduled job summarizes test coverage across all backends; [here is an example run](https://github.com/sgl-project/sglang/actions/runs/23424304300).
@@ -111,4 +144,4 @@ This README mostly describes the NVIDIA GPU CI pipeline. Other hardware backends
### Adding New Models to Nightly CI
- **Text models**: Extend the [global model list variables](https://github.com/sgl-project/sglang/blob/85c1f7937781199203b38bb46325a2840f353a04/python/sglang/test/test_utils.py#L104) in `test_utils.py`.
- **VLMs**: Extend the `MODEL_THRESHOLDS` dictionary in `test/registered/eval/test_vlms_mmmu_eval.py`.
- **VLMs**: Extend the `MODEL_THRESHOLDS` dictionary in `test/registered/accuracy/models/test_vlms_mmmu_eval.py`.
@@ -1,290 +0,0 @@
"""AMD GROK GSM8K Completion Evaluation Test (8-GPU)
Tests GROK models (Grok-1 FP8, Grok-1 INT4, Grok-2) using
few-shot completion benchmark on MI300X.
Registry: nightly-amd-8-gpu-grok suite
"""
import ast
import os
import re
import time
import unittest
from dataclasses import dataclass
from typing import List, Optional, Tuple
import numpy as np
from sglang.srt.utils import kill_process_tree
from sglang.test.ci.ci_register import register_amd_ci
from sglang.test.test_utils import (
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
is_in_ci,
popen_launch_server,
write_github_step_summary,
)
from sglang.utils import download_and_cache_file, read_jsonl
# DISABLED: Split into individual files for each model variant
# See: test_grok1_fp8_eval_amd.py, test_grok1_int4_eval_amd.py, test_grok2_eval_amd.py
register_amd_ci(
est_time=2700,
suite="nightly-amd-8-gpu-grok",
nightly=True,
disabled="Split into test_grok1_fp8_eval_amd.py, test_grok1_int4_eval_amd.py, test_grok2_eval_amd.py",
)
INVALID = -9999999
@dataclass
class ModelConfig:
"""Configuration for a model to test."""
model_path: str
tp_size: int = 8
accuracy_threshold: float = 0.50
other_args: Optional[List[str]] = None
env_vars: Optional[dict] = None
tokenizer_path: Optional[str] = None
timeout: Optional[int] = None
def __post_init__(self):
if self.other_args is None:
self.other_args = []
if self.env_vars is None:
self.env_vars = {}
# GROK models for MI300X
GROK_MODELS = [
# GROK1-FP8
ModelConfig(
model_path="lmzheng/grok-1",
tp_size=8,
accuracy_threshold=0.80,
timeout=3600,
tokenizer_path="Xenova/grok-1-tokenizer",
other_args=[
"--quantization",
"fp8",
"--attention-backend",
"aiter",
"--mem-fraction-static",
"0.85",
"--trust-remote-code",
],
env_vars={
"RCCL_MSCCL_ENABLE": "0",
"SGLANG_USE_AITER": "1",
"SGLANG_INT4_WEIGHT": "0",
},
),
# GROK1-INT4
ModelConfig(
model_path="amd/grok-1-W4A8KV8",
tp_size=8,
accuracy_threshold=0.80,
timeout=3600,
tokenizer_path="Xenova/grok-1-tokenizer",
other_args=[
"--quantization",
"fp8",
"--attention-backend",
"aiter",
"--mem-fraction-static",
"0.85",
"--trust-remote-code",
],
env_vars={
"RCCL_MSCCL_ENABLE": "0",
"SGLANG_USE_AITER": "1",
"SGLANG_INT4_WEIGHT": "1",
},
),
# GROK2
ModelConfig(
model_path="xai-org/grok-2",
tp_size=8,
accuracy_threshold=0.915,
timeout=3600,
tokenizer_path="alvarobartt/grok-2-tokenizer",
other_args=[
"--quantization",
"fp8",
"--attention-backend",
"aiter",
"--mem-fraction-static",
"0.85",
"--trust-remote-code",
],
env_vars={
"RCCL_MSCCL_ENABLE": "0",
"SGLANG_USE_AITER": "1",
"SGLANG_INT4_WEIGHT": "0",
},
),
]
def get_one_example(lines, i, include_answer):
"""Format a single GSM8K example."""
ret = "Question: " + lines[i]["question"] + "\nAnswer:"
if include_answer:
ret += " " + lines[i]["answer"]
return ret
def get_few_shot_examples(lines, k):
"""Get k few-shot examples for prompting."""
ret = ""
for i in range(k):
ret += get_one_example(lines, i, True) + "\n\n"
return ret
def get_answer_value(answer_str):
"""Extract numerical answer from response."""
answer_str = answer_str.replace(",", "")
numbers = re.findall(r"\d+", answer_str)
if len(numbers) < 1:
return INVALID
try:
return ast.literal_eval(numbers[-1])
except SyntaxError:
return INVALID
def run_gsm8k_benchmark(
base_url: str,
num_questions: int = 200,
num_shots: int = 5,
parallel: int = 64,
) -> Tuple[float, float, float]:
"""Run GSM8K few-shot completion benchmark."""
import sglang as sgl
from sglang.lang.backend.runtime_endpoint import RuntimeEndpoint
url = "https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl"
data_path = download_and_cache_file(url)
lines = list(read_jsonl(data_path))
few_shot_examples = get_few_shot_examples(lines, num_shots)
questions = []
labels = []
for i in range(len(lines[:num_questions])):
questions.append(get_one_example(lines, i, False))
labels.append(get_answer_value(lines[i]["answer"]))
assert all(l != INVALID for l in labels)
arguments = [{"question": q} for q in questions]
@sgl.function
def few_shot_gsm8k(s, question):
s += few_shot_examples + question
s += sgl.gen(
"answer", max_tokens=512, stop=["Question", "Assistant:", "<|separator|>"]
)
backend = RuntimeEndpoint(base_url)
sgl.set_default_backend(backend)
tic = time.perf_counter()
states = few_shot_gsm8k.run_batch(
arguments, temperature=0, num_threads=parallel, progress_bar=True
)
latency = time.perf_counter() - tic
preds = [get_answer_value(states[i]["answer"]) for i in range(len(states))]
acc = np.mean(np.array(preds) == np.array(labels))
invalid = np.mean(np.array(preds) == INVALID)
return float(acc), float(invalid), float(latency)
class TestGrokEvalAMD(unittest.TestCase):
"""GROK GSM8K Completion Evaluation Test for AMD MI300X."""
@classmethod
def setUpClass(cls):
cls.models = GROK_MODELS
cls.base_url = DEFAULT_URL_FOR_TEST
cls.num_questions = int(os.environ.get("GSM8K_NUM_QUESTIONS", "200"))
def test_grok_accuracy(self):
"""Test GROK models with GSM8K completion benchmark."""
all_results = []
summary = "### GROK Models (MI300X)\n\n"
summary += "| Model | TP | Accuracy | Threshold | Status |\n"
summary += "| ----- | -- | -------- | --------- | ------ |\n"
for config in self.models:
with self.subTest(model=config.model_path):
print(f"\n{'=' * 60}")
print(f"Testing: {config.model_path}")
print(f"{'=' * 60}")
env = os.environ.copy()
for key, value in config.env_vars.items():
env[key] = value
other_args = list(config.other_args)
other_args.extend(["--tp", str(config.tp_size)])
if config.tokenizer_path:
other_args.extend(["--tokenizer-path", config.tokenizer_path])
timeout = config.timeout or DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH
try:
process = popen_launch_server(
model=config.model_path,
base_url=self.base_url,
timeout=timeout,
other_args=other_args,
env=env,
)
try:
acc, invalid, latency = run_gsm8k_benchmark(
self.base_url, num_questions=self.num_questions
)
passed = acc >= config.accuracy_threshold
status = "✅ PASS" if passed else "❌ FAIL"
print(
f" accuracy={acc:.3f} threshold={config.accuracy_threshold} {status}"
)
all_results.append(
{
"model": config.model_path,
"accuracy": acc,
"passed": passed,
}
)
summary += f"| {config.model_path} | {config.tp_size} | {acc:.3f} | {config.accuracy_threshold} | {status} |\n"
finally:
kill_process_tree(process.pid)
except Exception as e:
summary += f"| {config.model_path} | {config.tp_size} | N/A | {config.accuracy_threshold} | ❌ ERROR |\n"
all_results.append(
{
"model": config.model_path,
"accuracy": None,
"passed": False,
"error": str(e),
}
)
if is_in_ci():
write_github_step_summary(summary)
failed = [r for r in all_results if not r["passed"]]
if failed:
raise AssertionError(f"Failed models: {[r['model'] for r in failed]}")
if __name__ == "__main__":
unittest.main()
@@ -23,7 +23,7 @@ which is below the >=3.5.0 the aiter gluon DSA kernels need, so it logs
what loses the accuracy. Re-add a 7.0 job once its image ships Triton >=3.5.0,
or once the legacy DSA fallback is fixed on gfx950.
The eval matches the CUDA GLM-5.2-FP8 nightly (`test/registered/8-gpu-models/
The eval matches the CUDA GLM-5.2-FP8 nightly (`test/registered/e2e/models_large/
test_glm52_fp8.py`): same dataset and same 0.92 baseline, so a red run here
means AMD diverged from CUDA rather than the harness diverging.
@@ -14,14 +14,13 @@ import unittest
import psutil
import sglang as sgl
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import (
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
CustomTestCase,
)
register_cuda_ci(est_time=38, stage="base-b", runner_config="1-gpu-small")
register_amd_ci(est_time=77, suite="stage-b-test-1-gpu-small-amd")
register_cuda_ci(est_time=77, stage="base-b", runner_config="1-gpu-small")
class TestEngineChildPids(CustomTestCase):
@@ -4,7 +4,7 @@ import re
import unittest
from sglang.srt.utils import kill_process_tree
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import (
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
@@ -17,8 +17,7 @@ from sglang.test.test_utils import (
send_generate_requests,
)
register_cuda_ci(est_time=65, stage="base-b", runner_config="1-gpu-small")
register_amd_ci(est_time=70, suite="stage-b-test-1-gpu-small-amd")
register_cuda_ci(est_time=53, stage="base-b", runner_config="1-gpu-small")
class TestMaxQueuedRequests(CustomTestCase):
@@ -2,7 +2,7 @@
Balanced recipe (TP=4, DeepEP, EAGLE) plus --attn-cp-size=4 with the
DSA prefill-CP interleave strategy. Split out of
models_e2e/test_deepseek_v4_flash_fp4_b200.py so the `cp` group covers
e2e/models/test_deepseek_v4_flash_fp4_b200.py so the `cp` group covers
all context-parallel tests.
Registry: extra-b-test-4-gpu-b200 (label-gated extra CI, 4x B200)
@@ -1,71 +0,0 @@
"""EAGLE spec-decoding core on CPU: the standard config (topk=1, page_size=1)
on the synchronous (non-overlap) path. topk > 1 tree drafting is covered in
test_spec_eagle_topk_cpu.py (split to stay under the per-file CI timeout).
"""
import unittest
from sglang.srt.environ import envs
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.kits.matched_stop_kit import MatchedStopMixin
from sglang.test.kits.spec_server_kits import (
SpecAccuracyKit,
SpecCorrectnessKit,
SpecFeatureKit,
SpecLogprobKit,
SpecPenaltyKit,
)
from sglang.test.server_fixtures.spec_eagle_fixture import EagleLlama2Base
# Measured 780s all-green on a 40-core GNR socket (1 launch + 18 methods).
register_cpu_ci(
est_time=800,
suite="stage-a-test-cpu-intel",
disabled="EagleLlama2Base needs gated meta-llama/Llama-2-7b-chat-hf",
)
_KITS = (
SpecCorrectnessKit,
SpecAccuracyKit,
SpecLogprobKit,
SpecPenaltyKit,
SpecFeatureKit,
MatchedStopMixin,
)
class _Core(EagleLlama2Base):
"""EAGLE (Llama-2) preset on CPU."""
attention_backend = "intel_amx"
disable_overlap = True
mem_fraction_static = 0.3
# CPU decode is compute-bound; a wider batch buys nothing here.
max_running_requests = 8
gsm8k_num_examples = 64
env_overrides = ((envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY, 1),)
class TestEagleLlama2NoOverlap(_Core, *_KITS):
"""Spec v1 (overlap scheduler off) -- the only mode reachable on CPU."""
# Standard chain config (topk=1, page_size=1), same shape as the CUDA core.
spec_steps = 5
spec_topk = 1
spec_tokens = 6
# EAGLE/Llama-2 topk=1 accepts modestly; tune against CI if needed.
acc_length_thres = 1.6
batch_accept_len_thres = 1.3
gsm8k_accept_len_thres = 1.3
@unittest.skip(
"constrained decoding on CPU needs a vocab-mask CPU branch in the "
"xgrammar backend (upstream gap, not spec-specific); the other grammar "
"backends lack the rollback spec verification requires"
)
def test_constrained_decoding(self):
pass
if __name__ == "__main__":
unittest.main()
@@ -1,29 +0,0 @@
import unittest
from sglang.srt.environ import envs
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.kits.spec_server_kits import SpecParityKit
from sglang.test.server_fixtures.spec_eagle_fixture import Eagle3Base
# Estimated: 2 sequential 8B server launches + one 4-prompt greedy method
# (CUDA sibling: 360); tune from CI TIMINGS once it has run there.
register_cpu_ci(
est_time=480,
suite="stage-a-test-cpu-intel",
disabled="EAGLE3 numerical parity mismatches on CPU intel_amx",
)
class TestEagle3ParityCPU(SpecParityKit, Eagle3Base):
"""EAGLE3 spec (intel_amx) greedy output == non-spec reference."""
attention_backend = "intel_amx"
disable_overlap = True
mem_fraction_static = 0.3
# CPU decode is compute-bound; a wider batch buys nothing here.
max_running_requests = 8
env_overrides = ((envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY, 1),)
if __name__ == "__main__":
unittest.main()
@@ -1,68 +0,0 @@
"""EAGLE topk > 1 tree drafting on CPU (Llama-2 topk=4, synchronous path).
Split from test_spec_eagle_cpu.py, mirroring the CUDA test_spec_eagle.py /
test_spec_eagle_topk.py layout, so each file stays under the per-file CI
timeout.
"""
import unittest
from sglang.srt.environ import envs
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.kits.spec_server_kits import (
SpecAccuracyKit,
SpecCorrectnessKit,
SpecFeatureKit,
SpecLogprobKit,
SpecPenaltyKit,
)
from sglang.test.server_fixtures.spec_eagle_fixture import EagleLlama2Base
# Measured 830s all-green on a 40-core GNR socket (1 launch + 14 methods).
register_cpu_ci(
est_time=850,
suite="stage-a-test-cpu-intel",
disabled="EagleLlama2Base needs gated meta-llama/Llama-2-7b-chat-hf",
)
class _Core(EagleLlama2Base):
"""EAGLE (Llama-2) preset on CPU."""
attention_backend = "intel_amx"
disable_overlap = True
mem_fraction_static = 0.3
# CPU decode is compute-bound; a wider batch buys nothing here.
max_running_requests = 8
gsm8k_num_examples = 64
env_overrides = ((envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY, 1),)
class TestEagleLlama2Topk4(
_Core,
SpecCorrectnessKit,
SpecAccuracyKit,
SpecLogprobKit,
SpecPenaltyKit,
SpecFeatureKit,
):
"""EAGLE/Llama-2 topk=4 tree coverage (kits listed in bases)."""
spec_steps = 3
spec_topk = 4
spec_tokens = 8
acc_length_thres = 2.4
batch_accept_len_thres = 1.6
gsm8k_accept_len_thres = 2.0
@unittest.skip(
"constrained decoding on CPU needs a vocab-mask CPU branch in the "
"xgrammar backend (upstream gap, not spec-specific); the other grammar "
"backends lack the rollback spec verification requires"
)
def test_constrained_decoding(self):
pass
if __name__ == "__main__":
unittest.main()
@@ -802,269 +802,6 @@ class TestReduceSum:
without_dim_names(unsharder_result.tensors[0]), without_dim_names(expected)
)
def test_recompute_pseudo_mismatch(self) -> None:
"""_verify_replicated_group returns failed check for RECOMPUTE_PSEUDO axis mismatch."""
tensor_a = torch.ones(4)
tensor_b = torch.ones(4) + 0.1
checks: list[ReplicatedCheckResult] = _verify_replicated_group(
[tensor_a, tensor_b],
axis=ParallelAxis.RECOMPUTE_PSEUDO,
group_index=0,
)
assert len(checks) == 1
assert checks[0].axis == "recompute_pseudo"
assert checks[0].group_index == 0
assert checks[0].compared_index == 1
assert checks[0].baseline_index == 0
assert not checks[0].passed
assert checks[0].diff.max_abs_diff == pytest.approx(0.1, abs=1e-5)
class TestThdCpConcat:
def test_single_seq(self) -> None:
"""Single seq THD unshard: 2 ranks → per-seq concat."""
rank0 = apply_dim_names(torch.tensor([1, 2, 3]), ["t"])
rank1 = apply_dim_names(torch.tensor([4, 5, 6]), ["t"])
plan = UnsharderPlan(
axis=ParallelAxis.CP,
params=CpThdConcatParams(dim_name="t", seq_lens_per_rank=[3]),
groups=[[0, 1]],
)
unsharder_result: UnsharderResult = execute_unsharder_plan(plan, [rank0, rank1])
assert len(unsharder_result.tensors) == 1
expected = torch.tensor([1, 2, 3, 4, 5, 6])
assert torch.equal(without_dim_names(unsharder_result.tensors[0]), expected)
def test_multi_seq(self) -> None:
"""Multi-seq THD unshard: 2 ranks, seq_lens=[50, 32, 46]."""
# rank0: [seqA_r0(50) | seqB_r0(32) | pad_r0(46)]
# rank1: [seqA_r1(50) | seqB_r1(32) | pad_r1(46)]
seq_a_r0 = torch.arange(0, 50)
seq_b_r0 = torch.arange(100, 132)
pad_r0 = torch.full((46,), -1)
rank0 = apply_dim_names(torch.cat([seq_a_r0, seq_b_r0, pad_r0]), ["t"])
seq_a_r1 = torch.arange(50, 100)
seq_b_r1 = torch.arange(132, 164)
pad_r1 = torch.full((46,), -2)
rank1 = apply_dim_names(torch.cat([seq_a_r1, seq_b_r1, pad_r1]), ["t"])
plan = UnsharderPlan(
axis=ParallelAxis.CP,
params=CpThdConcatParams(dim_name="t", seq_lens_per_rank=[50, 32, 46]),
groups=[[0, 1]],
)
unsharder_result: UnsharderResult = execute_unsharder_plan(plan, [rank0, rank1])
assert len(unsharder_result.tensors) == 1
unsharded: torch.Tensor = without_dim_names(unsharder_result.tensors[0])
# seqA: r0(50) + r1(50) = 100 tokens, values 0..99
assert torch.equal(unsharded[:100], torch.cat([seq_a_r0, seq_a_r1]))
# seqB: r0(32) + r1(32) = 64 tokens
assert torch.equal(unsharded[100:164], torch.cat([seq_b_r0, seq_b_r1]))
# pad: r0(46) + r1(46) = 92 tokens
assert torch.equal(unsharded[164:256], torch.cat([pad_r0, pad_r1]))
def test_with_hidden_dim(self) -> None:
"""THD unshard with trailing hidden dim: shape [T, H]."""
torch.manual_seed(42)
hidden: int = 4
# rank0: [seqA_r0(3, 4) | seqB_r0(2, 4)]
# rank1: [seqA_r1(3, 4) | seqB_r1(2, 4)]
seq_a_r0 = torch.randn(3, hidden)
seq_b_r0 = torch.randn(2, hidden)
rank0 = apply_dim_names(torch.cat([seq_a_r0, seq_b_r0]), ["t", "h"])
seq_a_r1 = torch.randn(3, hidden)
seq_b_r1 = torch.randn(2, hidden)
rank1 = apply_dim_names(torch.cat([seq_a_r1, seq_b_r1]), ["t", "h"])
plan = UnsharderPlan(
axis=ParallelAxis.CP,
params=CpThdConcatParams(dim_name="t", seq_lens_per_rank=[3, 2]),
groups=[[0, 1]],
)
unsharder_result: UnsharderResult = execute_unsharder_plan(plan, [rank0, rank1])
assert len(unsharder_result.tensors) == 1
unsharded: torch.Tensor = without_dim_names(unsharder_result.tensors[0])
assert unsharded.shape == (10, hidden)
assert torch.equal(unsharded[:6], torch.cat([seq_a_r0, seq_a_r1]))
assert torch.equal(unsharded[6:10], torch.cat([seq_b_r0, seq_b_r1]))
def test_with_leading_batch_dim(self) -> None:
"""THD unshard with leading batch dim: shape [B, T, H], t is dim=1."""
torch.manual_seed(42)
batch: int = 2
hidden: int = 4
# rank0: [seqA_r0(3) | seqB_r0(2)] per batch item
# rank1: [seqA_r1(3) | seqB_r1(2)] per batch item
seq_a_r0 = torch.randn(batch, 3, hidden)
seq_b_r0 = torch.randn(batch, 2, hidden)
rank0 = apply_dim_names(torch.cat([seq_a_r0, seq_b_r0], dim=1), ["b", "t", "h"])
seq_a_r1 = torch.randn(batch, 3, hidden)
seq_b_r1 = torch.randn(batch, 2, hidden)
rank1 = apply_dim_names(torch.cat([seq_a_r1, seq_b_r1], dim=1), ["b", "t", "h"])
plan = UnsharderPlan(
axis=ParallelAxis.CP,
params=CpThdConcatParams(dim_name="t", seq_lens_per_rank=[3, 2]),
groups=[[0, 1]],
)
unsharder_result: UnsharderResult = execute_unsharder_plan(plan, [rank0, rank1])
assert len(unsharder_result.tensors) == 1
unsharded: torch.Tensor = without_dim_names(unsharder_result.tensors[0])
assert unsharded.shape == (batch, 10, hidden)
# seqA: r0(3) + r1(3) = 6 tokens per batch
assert torch.equal(unsharded[:, :6, :], torch.cat([seq_a_r0, seq_a_r1], dim=1))
# seqB: r0(2) + r1(2) = 4 tokens per batch
assert torch.equal(
unsharded[:, 6:10, :], torch.cat([seq_b_r0, seq_b_r1], dim=1)
)
class TestReduceSum:
def test_basic_tp2_reduce(self) -> None:
"""2 partial tensors sum to full tensor."""
torch.manual_seed(42)
full_tensor = torch.randn(4, 8)
part_a = full_tensor * 0.6
part_b = full_tensor * 0.4
dim_specs = parse_dims("h[tp:partial] d").dims
parallel_infos = [
{ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=2)} for i in range(2)
]
plans = compute_unsharder_plan(dim_specs, parallel_infos)
assert len(plans) == 1
assert isinstance(plans[0].params, ReduceSumParams)
named_parts: list[torch.Tensor] = _name_tensors([part_a, part_b], dim_specs)
unsharder_result: UnsharderResult = execute_unsharder_plan(
plans[0], named_parts
)
assert len(unsharder_result.tensors) == 1
assert torch.allclose(
without_dim_names(unsharder_result.tensors[0]), full_tensor
)
def test_tp4_reduce(self) -> None:
"""4 partial tensors sum to full tensor."""
torch.manual_seed(42)
full_tensor = torch.randn(4, 8)
parts: list[torch.Tensor] = [full_tensor * 0.25 for _ in range(4)]
dim_specs = parse_dims("h[tp:partial] d").dims
parallel_infos = [
{ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=4)} for i in range(4)
]
plans = compute_unsharder_plan(dim_specs, parallel_infos)
assert len(plans) == 1
named_parts: list[torch.Tensor] = _name_tensors(parts, dim_specs)
unsharder_result: UnsharderResult = execute_unsharder_plan(
plans[0], named_parts
)
assert len(unsharder_result.tensors) == 1
assert torch.allclose(
without_dim_names(unsharder_result.tensors[0]), full_tensor
)
def test_multi_axis_concat_then_reduce(self) -> None:
"""CP concat + TP reduce end-to-end."""
torch.manual_seed(42)
full_tensor = torch.randn(4, 8, 16)
cp_chunks = list(full_tensor.chunk(2, dim=1))
# Each CP chunk is held as partial sums across TP ranks
tensors: list[torch.Tensor] = []
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = []
for cp_rank in range(2):
for tp_rank in range(2):
tensors.append(cp_chunks[cp_rank] * 0.5)
parallel_infos.append(
{
ParallelAxis.CP: AxisInfo(axis_rank=cp_rank, axis_size=2),
ParallelAxis.TP: AxisInfo(axis_rank=tp_rank, axis_size=2),
}
)
dim_specs = parse_dims("b s[cp] h[tp:partial]").dims
plans = compute_unsharder_plan(dim_specs, parallel_infos)
assert len(plans) == 2
current: list[torch.Tensor] = _name_tensors(tensors, dim_specs)
for plan in plans:
unsharder_result: UnsharderResult = execute_unsharder_plan(plan, current)
current = unsharder_result.tensors
assert len(current) == 1
assert torch.allclose(without_dim_names(current[0]), full_tensor)
def test_reduce_scrambled_ranks(self) -> None:
"""Scrambled rank order — sum is commutative so result is the same."""
torch.manual_seed(42)
full_tensor = torch.randn(4, 8)
parts: list[torch.Tensor] = [
full_tensor * 0.1,
full_tensor * 0.2,
full_tensor * 0.3,
full_tensor * 0.4,
]
parallel_infos = [
{ParallelAxis.TP: AxisInfo(axis_rank=2, axis_size=4)},
{ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=4)},
{ParallelAxis.TP: AxisInfo(axis_rank=3, axis_size=4)},
{ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=4)},
]
dim_specs = parse_dims("h[tp:partial] d").dims
plans = compute_unsharder_plan(dim_specs, parallel_infos)
named_parts: list[torch.Tensor] = _name_tensors(parts, dim_specs)
unsharder_result: UnsharderResult = execute_unsharder_plan(
plans[0], named_parts
)
assert len(unsharder_result.tensors) == 1
assert torch.allclose(
without_dim_names(unsharder_result.tensors[0]), full_tensor
)
def test_reduce_preserves_named_dims(self) -> None:
"""Named tensor dimensions are preserved through reduce_sum."""
dim_specs = parse_dims("h[tp:partial] d").dims
part_a = apply_dim_names(torch.randn(4, 8), ["h", "d"])
part_b = apply_dim_names(torch.randn(4, 8), ["h", "d"])
plan = UnsharderPlan(
axis=ParallelAxis.TP,
params=ReduceSumParams(),
groups=[[0, 1]],
)
unsharder_result: UnsharderResult = execute_unsharder_plan(
plan, [part_a, part_b]
)
assert len(unsharder_result.tensors) == 1
assert get_dim_names(unsharder_result.tensors[0]) == ("h", "d")
expected = apply_dim_names(
without_dim_names(part_a) + without_dim_names(part_b), ["h", "d"]
)
assert torch.allclose(
without_dim_names(unsharder_result.tensors[0]), without_dim_names(expected)
)
class TestFusedDimExecutor:
def test_fused_tp2_concat(self) -> None:
@@ -19,79 +19,6 @@ from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=20, stage="weekly", runner_config="cpu")
class TestComputeTensorInfo:
def test_basic_tensor_returns_correct_shape_and_dtype(self) -> None:
tensor = torch.randn(2, 3)
info = compute_tensor_info(tensor)
assert info.shape == [2, 3]
assert info.dtype == "torch.float32"
assert info.stats.mean == pytest.approx(tensor.float().mean().item(), abs=1e-4)
def test_include_sample_false_returns_none_sample(self) -> None:
tensor = torch.randn(2, 3)
info = compute_tensor_info(tensor, include_sample=False)
assert info.sample is None
def test_include_sample_true_returns_string_sample(self) -> None:
tensor = torch.randn(2, 3)
info = compute_tensor_info(tensor, include_sample=True)
assert info.sample is not None
assert isinstance(info.sample, str)
def test_empty_tensor_stats_are_zero(self) -> None:
tensor = torch.tensor([])
info = compute_tensor_info(tensor)
assert info.stats.mean == 0.0
assert info.stats.std == 0.0
assert info.shape == [0]
def test_integer_tensor_converted_to_float_for_stats(self) -> None:
"""Integer tensors should be cast to float internally for stats computation."""
tensor = torch.tensor([1, 2, 3, 4], dtype=torch.int32)
info = compute_tensor_info(tensor)
assert info.dtype == "torch.int32"
assert info.stats.mean == pytest.approx(2.5, abs=1e-4)
assert info.stats.min == pytest.approx(1.0, abs=1e-4)
assert info.stats.max == pytest.approx(4.0, abs=1e-4)
def test_bfloat16_tensor_shape_and_stats(self) -> None:
"""bfloat16 tensors produce correct shape and dtype string."""
tensor = torch.ones(3, 4, dtype=torch.bfloat16)
info = compute_tensor_info(tensor)
assert info.shape == [3, 4]
assert info.dtype == "torch.bfloat16"
assert info.stats.mean == pytest.approx(1.0, abs=1e-2)
def test_multidimensional_shape(self) -> None:
"""Shape is preserved for high-rank tensors."""
tensor = torch.randn(2, 3, 4, 5)
info = compute_tensor_info(tensor)
assert info.shape == [2, 3, 4, 5]
def test_scalar_tensor(self) -> None:
"""Scalar (0-dim) tensor produces empty shape list."""
tensor = torch.tensor(3.14)
info = compute_tensor_info(tensor)
assert info.shape == []
assert info.stats.mean == pytest.approx(3.14, abs=1e-4)
assert info.stats.min == pytest.approx(3.14, abs=1e-4)
assert info.stats.max == pytest.approx(3.14, abs=1e-4)
def test_include_sample_true_contains_tensor_representation(self) -> None:
"""Sample string should contain some recognizable tensor content."""
tensor = torch.tensor([1.0, 2.0])
info = compute_tensor_info(tensor, include_sample=True)
assert info.sample is not None
assert "1." in info.sample or "2." in info.sample
def test_percentiles_present_for_small_tensor(self) -> None:
"""Small tensors (< threshold) should have percentile data."""
tensor = torch.randn(100)
info = compute_tensor_info(tensor)
assert len(info.stats.percentiles) > 0
assert 50 in info.stats.percentiles
class TestComputeTensorInfo:
def test_basic_tensor_returns_correct_shape_and_dtype(self) -> None:
tensor = torch.randn(2, 3)
@@ -389,47 +389,6 @@ class TestTorchSave:
assert "skip the tensor" in captured.out
class TestLog:
def test_log_format(self):
with _capture_stdout() as captured:
_log("hello")
out = captured.getvalue()
assert "hello" in out, out
assert "[Dumper, rank=" in out, out
assert ", t=" in out, out
class TestCompareTensorsQuick:
def test_identical(self):
a = torch.tensor([1.0, 2.0, 3.0])
s = _compare_tensors_quick(a, a.clone())
assert "rel_diff=0" in s, s
assert "max_abs=0" in s, s
def test_diverged(self):
a = torch.tensor([1.0, 2.0, 3.0])
b = torch.tensor([1.0, 2.0, 4.0]) # last element differs by 1
s = _compare_tensors_quick(a, b)
assert "max_abs=1" in s, s
assert "rel_diff=" in s, s
def test_shape_mismatch(self):
s = _compare_tensors_quick(torch.zeros(3), torch.zeros(4))
assert "shape mismatch" in s, s
def test_dtype_unified(self):
s = _compare_tensors_quick(
torch.zeros(3, dtype=torch.float32),
torch.zeros(3, dtype=torch.float64),
)
assert "rel_diff=" in s, s
assert "max_abs=" in s, s
def test_empty(self):
s = _compare_tensors_quick(torch.zeros(0), torch.zeros(0))
assert s == "empty"
class TestCollectiveTimeout:
def test_watchdog_fires_on_timeout(self):
block_event = threading.Event()
@@ -1,4 +1,6 @@
import tempfile
import unittest
from pathlib import Path
import torch
from torch import nn
@@ -12,20 +14,11 @@ from sglang.srt.layers.linear import LinearBase
from sglang.srt.models.qwen2 import Qwen2MLP
from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
from sglang.srt.utils import add_prefix, get_device
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.layer_ut_utils import init_single_process_dist
from sglang.test.test_utils import CustomTestCase
register_cuda_ci(
est_time=9,
stage="base-b",
runner_config="1-gpu-small",
disabled="Test uses pytest-style function without TestCase class - see #17145",
)
register_amd_ci(
est_time=15,
suite="stage-b-test-1-gpu-small-amd",
disabled="Test uses pytest-style function without TestCase class - see #17145",
)
register_cuda_ci(est_time=9, stage="base-b", runner_config="1-gpu-small")
TEST_HIDDEN_SIZE = 32
@@ -73,26 +66,29 @@ def init_weights(module):
torch.nn.init.ones_(module.weight)
def test_model_forward_dump(tmp_path):
set_global_server_args_for_scheduler(ServerArgs(model_path="dummy"))
device = get_device()
init_single_process_dist(backend=get_default_distributed_backend(device))
model = MockCausalLM()
model.apply(init_weights)
model = model.to(device=device, dtype=torch.bfloat16)
dumper = register_forward_hook_for_model(
model, tmp_path / "sglang_dump", [0], 0, 0, 0
)
class TestTensorDumpForwardHook(CustomTestCase):
def test_model_forward_dump(self):
set_global_server_args_for_scheduler(ServerArgs(model_path="dummy"))
device = get_device()
init_single_process_dist(backend=get_default_distributed_backend(device))
model = MockCausalLM()
model.apply(init_weights)
model = model.to(device=device, dtype=torch.bfloat16)
dir_path = dumper.get_dump_dir()
inp = torch.randn(4, TEST_HIDDEN_SIZE, dtype=torch.bfloat16) * 0.01
result = model(inp.to(device))
data = torch.load(f"{dir_path}/Pass00000.pt")
assert "model.layernorm" in data
assert "model.mlp.down_proj" in data
assert torch.allclose(
data["model.mlp.down_proj"], result.cpu(), rtol=1e-5, atol=1e-5
)
with tempfile.TemporaryDirectory() as temp_dir:
dumper = register_forward_hook_for_model(
model, Path(temp_dir) / "sglang_dump", [0], 0, 0, 0
)
dir_path = dumper.get_dump_dir()
inp = torch.randn(4, TEST_HIDDEN_SIZE, dtype=torch.bfloat16) * 0.01
result = model(inp.to(device))
data = torch.load(f"{dir_path}/Pass00000.pt")
self.assertIn("model.layernorm", data)
self.assertIn("model.mlp.down_proj", data)
torch.testing.assert_close(
data["model.mlp.down_proj"], result.cpu(), rtol=1e-5, atol=1e-5
)
if __name__ == "__main__":
@@ -0,0 +1,13 @@
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.ci.diffusion_suite_bridge import run_diffusion_suite
# ci-cost-override: compatibility bridge preserves the existing diffusion lane.
register_cuda_ci(
est_time=14400,
stage="base-b",
runner_config="diffusion-1-gpu-h100",
)
if __name__ == "__main__":
run_diffusion_suite("1-gpu")
@@ -0,0 +1,13 @@
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.ci.diffusion_suite_bridge import run_diffusion_suite
# ci-cost-override: compatibility bridge preserves the existing diffusion lane.
register_cuda_ci(
est_time=7200,
stage="base-b",
runner_config="diffusion-1-gpu-5090",
)
if __name__ == "__main__":
run_diffusion_suite("1-gpu-5090")
@@ -0,0 +1,13 @@
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.ci.diffusion_suite_bridge import run_diffusion_suite
# ci-cost-override: compatibility bridge preserves the existing diffusion lane.
register_cuda_ci(
est_time=14400,
stage="base-b",
runner_config="diffusion-1-gpu-b200",
)
if __name__ == "__main__":
run_diffusion_suite("1-gpu-b200")
@@ -0,0 +1,13 @@
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.ci.diffusion_suite_bridge import run_diffusion_suite
# ci-cost-override: compatibility bridge preserves the existing diffusion lane.
register_cuda_ci(
est_time=14400,
stage="base-b",
runner_config="diffusion-2-gpu-h100",
)
if __name__ == "__main__":
run_diffusion_suite("2-gpu")
@@ -0,0 +1,13 @@
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.ci.diffusion_suite_bridge import run_diffusion_suite
# ci-cost-override: compatibility bridge preserves the existing diffusion lane.
register_cuda_ci(
est_time=3600,
stage="base-b",
runner_config="diffusion-bcg-1-gpu-h100",
)
if __name__ == "__main__":
run_diffusion_suite("bcg-diffusion")
@@ -0,0 +1,13 @@
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.ci.diffusion_suite_bridge import run_diffusion_suite
# ci-cost-override: compatibility bridge preserves the existing diffusion lane.
register_cuda_ci(
est_time=14400,
stage="base-b",
runner_config="diffusion-component-2-gpu-h100",
)
if __name__ == "__main__":
run_diffusion_suite("component-accuracy")
@@ -0,0 +1,13 @@
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.ci.diffusion_suite_bridge import run_diffusion_suite
# ci-cost-override: compatibility bridge preserves the existing diffusion lane.
register_cuda_ci(
est_time=3600,
stage="base-b",
runner_config="diffusion-unit-1-gpu-h100",
)
if __name__ == "__main__":
run_diffusion_suite("unit")

Some files were not shown because too many files have changed in this diff Show More