Cheng Wan
b99175dc7d
Last of four; stacked on #38048. The record is the operator's input; the bags are what is in effect. A reader that takes the record and reads a field off it gets the input, which is the wrong one of the two whenever resolution decided something -- and the mistake is silent, because for most fields and most launches the two agree. Several of these files already read both ways, sometimes in the same expression: ```python get_tokenizer( get_serving().tokenizer_path, tokenizer_mode=server_args.tokenizer_mode, # the input, not the decision ... ) ``` Sixty-odd files convert. Record field reads in runtime code go from 199 to 11. Nine parameters that the conversion emptied are dropped along with the argument at every call site -- the dead-parameter ratchet is what names them. ### "Runs after its process publishes" is a per-entry-point claim Most converted reads sit in the serving and model-executor layers, which only exist after publication, or in the two subprocess entry points, which publish first thing. Three places are not like that, and they keep reading the record they were handed: - **`HttpServerEngineAdapter`** launches the server as a *child*. The parent resolves the record and never publishes, so the adapter's own reads -- the launch banner, the API key in its readiness loop, the TP width in `update_weights_from_tensor` -- are of `self.server_args`. A bag read here fails closed in a bare process, or answers for an unrelated engine in one that happens to have published. - **`serve_grpc`** reads its sidecar port before the integrated servicer builds the `Engine` that publishes. The comment above that line already said so and already bound `cfg = resolving_view(server_args)` for it; the sidecar port and the port it derives from read `cfg`. - **`initialize_dp_attention`** runs from callers whose publish is not guaranteed, so its one predicate stays on the resolution view. `ROLE_NAMESPACE_SETS["dp_controller"]` gains `observability` and `serving`, because the controller's metrics gate, tracing setup and worker-port broadcast now read those namespaces. Under `SGLANG_ROLE_NAMESPACES=enforce` that set is what the process may read, so a conversion that reaches a new namespace has to widen it in the same change. ## Three things worth a reviewer's attention **Eleven reads were `getattr(record, "field", default)`.** An AST scan for attribute access does not see those, so the census that said "43 readers" was counting the shape it could match rather than the thing it was after. `incremental_streaming_output` was read that way twice, and the transcription tests were the only reason it surfaced. **Not every record read is a bag read waiting to happen.** A multimodal processor's `base_gpu_id` is the instance's, not the process's: two engines in one process keep different ones, and `test_publishing_another_config_does_not_move_the_device` exists to say so. It stays on the record while `rl_on_policy_target` beside it moves. `RequestMetricsExporter` is the same shape -- it is handed the directory it writes to, and a test builds several with different ones. `configure_logger` is a third: 17 call sites, one of which passes an `argparse.Namespace`, so it is not a global-context reader at all. Those eleven remaining reads are the ones with a reason. **The fixtures move with the code.** Tests that hung config off a mock manager now publish a record, which is what the serving layer reads; where a test states a value it says so with `override_server_args` instead of assigning through the mock. `test_hisparse_unit` is the last of them: it stubbed a `server_args` onto a fake scheduler to say the decode radix cache was off, and the value it was standing in for is the published default, so the stub goes and the class publishes. ## Two things CI caught that a local sweep could not **`unittest.TestCase.enterContext` is Python 3.11+.** The converted fixtures used it at 18 sites; `requires-python` is `>=3.10` and CI runs 3.10, so every one of them raised `AttributeError` there while passing on a newer local interpreter. They call `enter_override(self, ...)` now -- a four-line helper in `sglang/test/test_utils.py` over the override's own `install()` / `restore()`. **A batched sweep cannot see a missing publish.** Three fixtures needed a published config and did not have one; each *passed* inside a shard where some other file had published, and failed when run alone. The affected cases are `test_serving_completions` (which set `incremental_streaming_output` on the mock manager's record, where nothing reads it now), `test_qwen3_vl_feature_materialization` (same shape for `mm_enable_dp_encoder`), and the two Qwen Rust tests -- whose fixture already carried the comment `# Non-auto: get_resolved_model_impl would choke on a SimpleNamespace` next to the `model_impl` it sets, which is exactly what happened once `get_mm_processor_cls` started reading that value from the bag. Its `publish` mirrors `model_impl` now, like the four fields it already mirrored. ## Verification A full registered-unit sweep (648 files) against this stack's merge-base: 19 failures on both sides, the same 19, none of them config. That sweep is what caught 23 failures the file-scoped runs missed -- and, later, that the narrower 139-file list did not even contain the files this change reaches. It is also what caught the `test_hisparse_unit` fixture above: the file passes inside a shard where something else published, and fails when it is run on its own, which is why every failing file is re-run alone before it is counted.
436 lines
17 KiB
Python
436 lines
17 KiB
Python
import unittest
|
|
from concurrent.futures import Future
|
|
from types import SimpleNamespace
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
from sglang.srt.disaggregation.base import KVPoll
|
|
from sglang.srt.disaggregation.decode import (
|
|
DecodePreallocQueue,
|
|
DecodeTransferQueue,
|
|
HiCacheRestoreResult,
|
|
)
|
|
from sglang.srt.disaggregation.fake.conn import FakeKVManager, FakeKVReceiver
|
|
from sglang.srt.disaggregation.utils import DisaggregationMode
|
|
from sglang.srt.distributed.parallel_state_wrapper import ParallelState
|
|
from sglang.srt.managers.schedule_batch import FINISH_ABORT
|
|
from sglang.srt.managers.scheduler import Scheduler
|
|
from sglang.srt.runtime_context import get_context, publish, reset_context
|
|
from sglang.srt.server_args import ServerArgs
|
|
from sglang.test.ci.ci_register import register_cpu_ci
|
|
from sglang.test.test_utils import CustomTestCase
|
|
|
|
register_cpu_ci(est_time=11, suite="base-a-test-cpu")
|
|
|
|
|
|
class FakeReceiver:
|
|
def __init__(self):
|
|
self.clear_called = False
|
|
self.conclude_state = None
|
|
|
|
def clear(self):
|
|
self.clear_called = True
|
|
|
|
def failure_exception(self):
|
|
return None
|
|
|
|
|
|
class TestDecodeQueueCleanup(CustomTestCase):
|
|
def setUp(self):
|
|
# The code under test reads its config from the bags.
|
|
reset_context()
|
|
self.addCleanup(reset_context)
|
|
publish(ServerArgs(model_path="dummy"), role="tokenizer")
|
|
|
|
def test_paged_swa_retraction_resume_uses_physical_page_budget(self):
|
|
# resume_retracted_reqs reads the retraction backend off the disagg
|
|
# bag, so the case publishes a config instead of injecting one.
|
|
override = get_context().override_server_args(
|
|
disaggregation_decode_retraction_backup="cpu_tensor"
|
|
)
|
|
override.install()
|
|
self.addCleanup(override.restore)
|
|
|
|
page_size = 128
|
|
fill_len = 574
|
|
physical_tokens_per_req = 5 * page_size
|
|
physical_available = 18 * page_size
|
|
|
|
reqs = [
|
|
SimpleNamespace(
|
|
rid=f"req-{i}",
|
|
origin_input_ids=[0] * fill_len,
|
|
output_ids=[],
|
|
is_retracted=True,
|
|
retraction_backup=None,
|
|
load_kv_cache=MagicMock(),
|
|
)
|
|
for i in range(4)
|
|
]
|
|
|
|
queue = DecodePreallocQueue.__new__(DecodePreallocQueue)
|
|
queue.retracted_queue = reqs.copy()
|
|
queue.num_reserved_decode_tokens = 0
|
|
queue.req_to_token_pool = SimpleNamespace(available_size=lambda: len(reqs))
|
|
queue.token_to_kv_pool_allocator = SimpleNamespace(page_size=page_size)
|
|
queue.tree_cache = MagicMock()
|
|
queue.scheduler = SimpleNamespace(
|
|
sliding_window_size=2047,
|
|
server_args=SimpleNamespace(disable_radix_cache=True),
|
|
)
|
|
queue._uses_swa_tail_prealloc = MagicMock(return_value=True)
|
|
queue._swa_aware_allocatable_token_budgets = MagicMock(
|
|
return_value=(physical_available, physical_available)
|
|
)
|
|
queue._swa_tail_allocatable_token_budget = MagicMock(
|
|
side_effect=lambda **_: physical_available
|
|
)
|
|
|
|
def pre_alloc(_req):
|
|
nonlocal physical_available
|
|
self.assertGreaterEqual(physical_available, physical_tokens_per_req)
|
|
physical_available -= physical_tokens_per_req
|
|
|
|
queue._pre_alloc = MagicMock(side_effect=pre_alloc)
|
|
|
|
resumed = queue.resume_retracted_reqs()
|
|
|
|
self.assertEqual(resumed, reqs[:3])
|
|
self.assertEqual(queue.retracted_queue, reqs[3:])
|
|
self.assertEqual(physical_available, 3 * page_size)
|
|
self.assertEqual(queue._pre_alloc.call_count, 3)
|
|
|
|
def test_prealloc_abort_clears_receiver_before_removing_request(self):
|
|
receiver = FakeReceiver()
|
|
req = SimpleNamespace(
|
|
rid="abort-prealloc",
|
|
bootstrap_room=42,
|
|
finished_reason=None,
|
|
return_logprob=False,
|
|
)
|
|
decode_req = SimpleNamespace(
|
|
req=req, kv_receiver=receiver, waiting_for_input=True
|
|
)
|
|
|
|
queue = DecodePreallocQueue.__new__(DecodePreallocQueue)
|
|
queue.pp_size = 1
|
|
queue.tp_rank = 0
|
|
queue.gloo_group = object()
|
|
queue.queue = [decode_req]
|
|
queue.pending_reqs = []
|
|
queue.retracted_queue = []
|
|
queue._resolve_pending_reqs = MagicMock()
|
|
queue._uses_swa_tail_prealloc = MagicMock(return_value=False)
|
|
queue._allocatable_token_budgets = MagicMock(return_value=0)
|
|
queue._hicache_pending_restore_tokens = MagicMock(return_value=0)
|
|
|
|
scheduler = MagicMock()
|
|
scheduler.running_batch.reqs = []
|
|
scheduler.enable_priority_scheduling = False
|
|
scheduler.enable_hisparse = False
|
|
scheduler.metrics_reporter.enable_metrics = False
|
|
scheduler.output_streamer = MagicMock()
|
|
queue.scheduler = scheduler
|
|
|
|
with patch(
|
|
"sglang.srt.disaggregation.decode.poll_and_all_reduce",
|
|
return_value=[KVPoll.Failed],
|
|
) as poll:
|
|
preallocated, failed = queue.pop_preallocated()
|
|
|
|
poll.assert_called_once_with([receiver], queue.gloo_group)
|
|
self.assertEqual(preallocated, [])
|
|
self.assertEqual(failed, [decode_req])
|
|
self.assertEqual(queue.queue, [])
|
|
self.assertTrue(receiver.clear_called)
|
|
self.assertIsNone(decode_req.kv_receiver)
|
|
self.assertIsInstance(req.finished_reason, FINISH_ABORT)
|
|
scheduler.output_streamer.stream_output.assert_called_once_with(
|
|
[req], req.return_logprob
|
|
)
|
|
|
|
def test_prealloc_abort_also_drops_from_pending_reqs(self):
|
|
# Same DecodeRequest lives in both queue and pending_reqs (add() slow
|
|
# path). Aborting must drop it from both, and compare by identity since
|
|
# DecodeRequest's dataclass __eq__ would compare the tensor receiver.
|
|
class BadEqReceiver(FakeReceiver):
|
|
def __eq__(self, other):
|
|
raise TypeError("use identity comparison, not value equality")
|
|
|
|
__hash__ = object.__hash__
|
|
|
|
receiver = BadEqReceiver()
|
|
req = SimpleNamespace(
|
|
rid="abort-shared",
|
|
finished_reason=FINISH_ABORT("aborted"),
|
|
return_logprob=False,
|
|
)
|
|
decode_req = SimpleNamespace(req=req, kv_receiver=receiver)
|
|
|
|
queue = DecodePreallocQueue.__new__(DecodePreallocQueue)
|
|
queue.pp_size = 1
|
|
queue.queue = [decode_req]
|
|
queue.pending_reqs = [decode_req] # same object, dual ownership
|
|
queue.retracted_queue = []
|
|
queue._resolve_pending_reqs = MagicMock()
|
|
queue._update_handshake_waiters = MagicMock()
|
|
queue._uses_swa_tail_prealloc = MagicMock(return_value=False)
|
|
queue._allocatable_token_budgets = MagicMock(return_value=0)
|
|
queue._hicache_pending_restore_tokens = MagicMock(return_value=0)
|
|
|
|
scheduler = MagicMock()
|
|
scheduler.running_batch.reqs = []
|
|
scheduler.enable_priority_scheduling = False
|
|
scheduler.enable_hisparse = False
|
|
scheduler.output_streamer = MagicMock()
|
|
queue.scheduler = scheduler
|
|
|
|
# Must not raise on the receiver __eq__ above.
|
|
preallocated, failed = queue.pop_preallocated()
|
|
|
|
self.assertEqual(preallocated, [])
|
|
self.assertEqual(failed, [decode_req])
|
|
self.assertEqual(queue.queue, [])
|
|
self.assertTrue(all(r is not decode_req for r in queue.pending_reqs))
|
|
self.assertIsNone(decode_req.kv_receiver)
|
|
|
|
def test_swa_reclaim_failure_rejects_only_request(self):
|
|
receiver = FakeReceiver()
|
|
req = SimpleNamespace(
|
|
rid="swa-reclaim-failed",
|
|
origin_input_ids=[1, 2, 3],
|
|
output_ids=[],
|
|
finished_reason=None,
|
|
return_logprob=False,
|
|
sampling_params=SimpleNamespace(max_new_tokens=1),
|
|
)
|
|
decode_req = SimpleNamespace(
|
|
req=req,
|
|
kv_receiver=receiver,
|
|
waiting_for_input=True,
|
|
is_rebootstrap=False,
|
|
)
|
|
|
|
queue = DecodePreallocQueue.__new__(DecodePreallocQueue)
|
|
queue.pp_size = 1
|
|
queue.queue = [decode_req]
|
|
queue.pending_reqs = [decode_req]
|
|
queue.retracted_queue = []
|
|
queue.num_reserved_decode_tokens = 0
|
|
queue._resolve_pending_reqs = MagicMock()
|
|
queue._update_handshake_waiters = MagicMock()
|
|
queue._uses_swa_tail_prealloc = MagicMock(return_value=True)
|
|
queue._swa_aware_allocatable_token_budgets = MagicMock(
|
|
return_value=(1024, 1024)
|
|
)
|
|
queue._prealloc_required_tokens = MagicMock(return_value=(3, 3))
|
|
queue._prealloc_kv_lens = MagicMock(return_value=(3, 3))
|
|
queue._reclaim_swa_tail_capacity = MagicMock(
|
|
return_value=(
|
|
"SWA eviction insufficient: needed=64, available=0, "
|
|
"req=swa-reclaim-failed"
|
|
)
|
|
)
|
|
queue._hicache_pending_restore_tokens = MagicMock(return_value=0)
|
|
queue._pre_alloc = MagicMock()
|
|
queue.req_to_token_pool = MagicMock()
|
|
queue.req_to_token_pool.available_size.return_value = 1
|
|
queue.req_to_metadata_buffer_idx_allocator = MagicMock()
|
|
queue.req_to_metadata_buffer_idx_allocator.available_size.return_value = 1
|
|
|
|
scheduler = MagicMock()
|
|
scheduler.running_batch.reqs = []
|
|
scheduler.enable_priority_scheduling = False
|
|
scheduler.enable_hisparse = False
|
|
scheduler.server_args.disaggregation_decode_enable_radix_cache = False
|
|
scheduler.output_streamer = MagicMock()
|
|
queue.scheduler = scheduler
|
|
|
|
preallocated, failed = queue.pop_preallocated()
|
|
|
|
self.assertEqual(preallocated, [])
|
|
self.assertEqual(failed, [decode_req])
|
|
self.assertEqual(queue.queue, [])
|
|
self.assertEqual(queue.pending_reqs, [])
|
|
self.assertTrue(receiver.clear_called)
|
|
self.assertIsNone(decode_req.kv_receiver)
|
|
self.assertIsInstance(req.finished_reason, FINISH_ABORT)
|
|
queue._pre_alloc.assert_not_called()
|
|
scheduler.output_streamer.stream_output.assert_called_once_with(
|
|
[req], req.return_logprob
|
|
)
|
|
|
|
def test_ensure_prefill_info_tolerates_cleared_receiver(self):
|
|
# A req whose kv_receiver was already cleared must not crash on .abort().
|
|
queue = DecodePreallocQueue.__new__(DecodePreallocQueue)
|
|
queue._max_ensure_retries = 1
|
|
queue._ensure_retry_interval = 0
|
|
queue._ensure_retry_count = {"127.0.0.1:11500": 0}
|
|
queue._ensure_last_attempt_time = {}
|
|
queue.kv_manager = MagicMock()
|
|
queue.kv_manager.try_ensure_parallel_info.return_value = False
|
|
|
|
cleared_req = SimpleNamespace(
|
|
req=SimpleNamespace(rid="cleared"), kv_receiver=None
|
|
)
|
|
addr_to_reqs = {"127.0.0.1:11500": [cleared_req]}
|
|
|
|
ready, remaining = queue._ensure_prefill_info(addr_to_reqs)
|
|
|
|
self.assertEqual(ready, {})
|
|
self.assertEqual(remaining, [])
|
|
|
|
def test_prefetches_prefill_dp_rank_query(self):
|
|
addr = "127.0.0.1:11500"
|
|
executor = MagicMock()
|
|
future = Future()
|
|
future.set_result({"7": 1})
|
|
executor.submit.return_value = future
|
|
|
|
def decode_req(room):
|
|
return SimpleNamespace(
|
|
req=SimpleNamespace(
|
|
bootstrap_host="127.0.0.1",
|
|
bootstrap_port=11500,
|
|
bootstrap_room=room,
|
|
),
|
|
kv_receiver=MagicMock(),
|
|
)
|
|
|
|
first = decode_req(7)
|
|
queue = DecodePreallocQueue.__new__(DecodePreallocQueue)
|
|
queue.pending_reqs = [first]
|
|
queue._prefill_dp_rank_queries = {}
|
|
queue.kv_manager = SimpleNamespace(
|
|
prefill_info_table={addr: object()},
|
|
_ensure_prefill_recompute_executor=lambda: executor,
|
|
)
|
|
queue._resolve_prefill_dp_rank = MagicMock(return_value=None)
|
|
queue._ensure_prefill_info = lambda groups: (groups, [])
|
|
|
|
queue.prefetch_prefill_dp_rank_queries()
|
|
tail = decode_req(8)
|
|
queue.pending_reqs.append(tail)
|
|
with patch(
|
|
"sglang.srt.disaggregation.decode.CommonKVReceiver.query_prefill_dp_ranks",
|
|
return_value={"8": 2},
|
|
) as query:
|
|
queue._resolve_pending_reqs()
|
|
|
|
_, called_addr, called_rooms = executor.submit.call_args.args
|
|
self.assertEqual((called_addr, called_rooms), (addr, [7]))
|
|
query.assert_called_once_with(addr, [8])
|
|
first.kv_receiver.init.assert_called_once_with(1)
|
|
tail.kv_receiver.init.assert_called_once_with(2)
|
|
self.assertEqual(queue.pending_reqs, [])
|
|
|
|
@patch("sglang.srt.disaggregation.decode.release_kv_cache")
|
|
@patch("sglang.srt.disaggregation.decode.prepare_abort")
|
|
@patch("sglang.srt.disaggregation.decode.poll_and_all_reduce")
|
|
def test_transfer_failure_cleanup_respects_deferred_release_gates(
|
|
self, mock_poll, mock_prepare_abort, mock_release_kv_cache
|
|
):
|
|
receiver = FakeReceiver()
|
|
req = SimpleNamespace(
|
|
rid="failed-transfer",
|
|
bootstrap_room=7,
|
|
return_logprob=False,
|
|
)
|
|
decode_req = SimpleNamespace(
|
|
req=req,
|
|
kv_receiver=receiver,
|
|
metadata_buffer_index=3,
|
|
hicache_restore_status=HiCacheRestoreResult.READY,
|
|
)
|
|
|
|
queue = DecodeTransferQueue.__new__(DecodeTransferQueue)
|
|
queue.queue = [decode_req]
|
|
queue.enable_staging = False
|
|
queue.enable_deferred_kv_release = False
|
|
queue.gloo_group = MagicMock()
|
|
queue.req_to_metadata_buffer_idx_allocator = MagicMock()
|
|
queue.tp_rank = 0
|
|
queue.tree_cache = MagicMock()
|
|
queue.metadata_buffers = SimpleNamespace(bootstrap_room=[None] * 4)
|
|
queue.spec_algorithm = MagicMock()
|
|
queue.spec_algorithm.is_none.return_value = True
|
|
queue._clean_hicache_prefetch_resources = MagicMock()
|
|
|
|
scheduler = MagicMock()
|
|
scheduler.enable_decode_hicache = False
|
|
scheduler.enable_hisparse = False
|
|
scheduler.output_streamer = MagicMock()
|
|
scheduler.metrics_reporter.enable_metrics = False
|
|
queue.scheduler = scheduler
|
|
|
|
mock_poll.return_value = [KVPoll.Failed]
|
|
|
|
transferred = queue.pop_transferred()
|
|
|
|
self.assertEqual(transferred, [])
|
|
self.assertEqual(queue.queue, [])
|
|
self.assertTrue(receiver.clear_called)
|
|
self.assertIsNone(decode_req.kv_receiver)
|
|
queue.req_to_metadata_buffer_idx_allocator.free.assert_called_once_with(3)
|
|
scheduler.output_streamer.stream_output.assert_called_once_with(
|
|
[req], req.return_logprob
|
|
)
|
|
mock_prepare_abort.assert_called_once()
|
|
mock_release_kv_cache.assert_called_once_with(
|
|
req, queue.tree_cache, is_insert=False
|
|
)
|
|
|
|
receiver = FakeReceiver()
|
|
receiver.kv_mgr = FakeKVManager.__new__(FakeKVManager)
|
|
decode_req.kv_receiver = receiver
|
|
queue.queue = [decode_req]
|
|
queue.enable_deferred_kv_release = True
|
|
queue.req_to_metadata_buffer_idx_allocator.reset_mock()
|
|
mock_release_kv_cache.reset_mock()
|
|
|
|
transferred = queue.pop_transferred()
|
|
|
|
self.assertEqual(transferred, [])
|
|
self.assertEqual(queue.queue, [])
|
|
self.assertTrue(receiver.clear_called)
|
|
self.assertIsNone(decode_req.kv_receiver)
|
|
queue.req_to_metadata_buffer_idx_allocator.free.assert_called_once_with(3)
|
|
mock_release_kv_cache.assert_called_once_with(
|
|
req, queue.tree_cache, is_insert=False
|
|
)
|
|
|
|
def test_fake_receiver_initializes_deferred_release_state(self):
|
|
manager = MagicMock()
|
|
receiver = FakeKVReceiver(manager, "")
|
|
|
|
self.assertIs(receiver.kv_mgr, manager)
|
|
self.assertFalse(receiver.abort_notified)
|
|
|
|
def test_retracted_decode_requests_keep_scheduler_non_idle(self):
|
|
scheduler = Scheduler.__new__(Scheduler)
|
|
scheduler.running_batch = MagicMock()
|
|
scheduler.running_batch.is_empty.return_value = True
|
|
scheduler.chunked_req = None
|
|
scheduler.dllm_manager = MagicMock()
|
|
scheduler.dllm_manager.any_staging_reqs.return_value = False
|
|
scheduler.last_batch = None
|
|
scheduler.cur_batch_for_debug = None
|
|
scheduler.enable_overlap = False
|
|
scheduler.ps = ParallelState.trivial()
|
|
scheduler.running_mbs = []
|
|
scheduler.waiting_queue = []
|
|
scheduler.grammar_manager = SimpleNamespace(grammar_queue=[])
|
|
scheduler.disaggregation_mode = DisaggregationMode.DECODE
|
|
scheduler.disagg_decode_prealloc_queue = SimpleNamespace(
|
|
queue=[], retracted_queue=[object()]
|
|
)
|
|
scheduler.disagg_decode_transfer_queue = SimpleNamespace(queue=[])
|
|
scheduler.decode_offload_manager = None
|
|
scheduler.enable_hisparse = False
|
|
scheduler.enable_hierarchical_cache = False
|
|
|
|
self.assertFalse(scheduler.is_fully_idle())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|