test: stand up the config tiers two unit tests read from (#33294)

This commit is contained in:
Cheng Wan
2026-08-02 15:46:34 -07:00
committed by GitHub
parent 8cc941a672
commit 88e5a0f635
3 changed files with 38 additions and 10 deletions
+2 -3
View File
@@ -552,10 +552,9 @@ def _acknowledge_deferred_cuda_ipc_cache_hits(
parallel = get_parallel()
if parallel.attn_tp_rank != 0:
return
server_args = get_server_args()
# The pool's recycler uses ServerArgs.tp_size, so its acknowledgement must
# The pool's recycler counts the whole TP group, so the acknowledgement must
# match that count even when an attention subgroup is smaller.
consumer_count = max(getattr(server_args, "tp_size", parallel.attn_tp_size), 1)
consumer_count = max(get_server_args().tp_size, 1)
for item in items:
item.acknowledge_deferred_cuda_ipc_feature(consumer_count)
@@ -14,10 +14,29 @@ import torch
from sglang.srt.managers import mm_utils
from sglang.srt.managers.schedule_batch import Modality, MultimodalDataItem
from sglang.srt.runtime_context import get_context, get_parallel
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=10, suite="base-b-test-cpu")
@pytest.fixture(autouse=True)
def publish_config_and_parallel_state():
"""Applied to every test in this module (``autouse``), named by none of them.
The embedding path reads the config namespaces and the attention-TP rank —
process state a served engine establishes at startup. Without this the
accessors raise instead of answering.
"""
override = get_context().override_server_args(tp_size=1)
override.install()
try:
with get_parallel().override(attn_tp_rank=0, attn_tp_size=1, tp_size=1):
yield
finally:
override.restore()
HIDDEN = 16
# Three items with text gaps between their placeholder runs; offsets are
@@ -13,6 +13,7 @@ import unittest
from types import SimpleNamespace
from unittest import mock
from sglang.srt.runtime_context import get_context
from sglang.test.ci.ci_register import register_cpu_ci, register_mlx_ci
from sglang.test.test_utils import CustomTestCase
@@ -40,6 +41,7 @@ def _arch(*, hybrid: bool):
def _stub_for_initialize(
test,
*,
dp_size: int,
attn_dp_size: int,
@@ -47,16 +49,22 @@ def _stub_for_initialize(
max_mamba_cache_size: int | None = None,
pool_size: int = 64,
):
stub = MlxModelRunnerStub.__new__(MlxModelRunnerStub)
stub._mlx_pool_size = pool_size
stub.device = "cpu"
stub.ps = ParallelState.trivial(dp_size=dp_size, attn_dp_size=attn_dp_size)
stub.server_args = SimpleNamespace(
# ``initialize`` reads the config namespaces, so the config has to be
# published rather than stubbed onto the runner.
override = get_context().override_server_args(
enable_memory_saver=False,
max_running_requests=max_running_requests,
max_mamba_cache_size=max_mamba_cache_size,
disable_radix_cache=False,
)
server_args = override.install()
test.addCleanup(override.restore)
stub = MlxModelRunnerStub.__new__(MlxModelRunnerStub)
stub._mlx_pool_size = pool_size
stub.device = "cpu"
stub.ps = ParallelState.trivial(dp_size=dp_size, attn_dp_size=attn_dp_size)
stub.server_args = server_args
stub.model_config = SimpleNamespace(
is_hybrid_swa=False,
sliding_window_size=None,
@@ -80,14 +88,14 @@ def _initialize_stub(stub, *, hybrid: bool = False):
class TestAttentionDpRequestCapacity(CustomTestCase):
def test_pure_dp_replica_retains_full_request_limit(self):
stub = _initialize_stub(
_stub_for_initialize(dp_size=4, attn_dp_size=1),
_stub_for_initialize(self, dp_size=4, attn_dp_size=1),
)
self.assertEqual(stub.max_running_requests, 8)
self.assertEqual(stub.req_to_token_pool.size, 8)
def test_attention_dp_partitions_request_limit(self):
stub = _initialize_stub(
_stub_for_initialize(dp_size=4, attn_dp_size=4),
_stub_for_initialize(self, dp_size=4, attn_dp_size=4),
)
self.assertEqual(stub.max_running_requests, 2)
self.assertEqual(stub.req_to_token_pool.size, 2)
@@ -95,6 +103,7 @@ class TestAttentionDpRequestCapacity(CustomTestCase):
def test_attention_dp_partitions_explicit_auxiliary_state_limit(self):
stub = _initialize_stub(
_stub_for_initialize(
self,
dp_size=4,
attn_dp_size=4,
max_running_requests=8,
@@ -109,6 +118,7 @@ class TestAttentionDpRequestCapacity(CustomTestCase):
def test_attention_dp_auxiliary_error_reports_global_cli_units(self):
stub = _stub_for_initialize(
self,
dp_size=4,
attn_dp_size=4,
max_running_requests=8,