wehub-resource-sync
1148 lines
43 KiB
Python
1148 lines
43 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""
|
|
Tests for interleaved chunked prefill + decode (SchedulerConfig.chunked_prefill).
|
|
|
|
Strategy: keep tests fast by mocking MLX model calls and cache operations.
|
|
_begin_prefill() and _step_prefill_chunk() are tested by patching
|
|
make_prompt_cache and mx.eval; the scheduler-level flow is tested by
|
|
patching _step_prefill_chunk directly.
|
|
"""
|
|
|
|
from collections import deque
|
|
from types import SimpleNamespace
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
from omlx.exceptions import PrefillMemoryExceededError
|
|
from omlx.request import Request, RequestStatus, SamplingParams
|
|
from omlx.scheduler import (
|
|
PrefillEvictionRequest,
|
|
Scheduler,
|
|
SchedulerConfig,
|
|
_PrefillAbortedError,
|
|
_PrefillEvictionNeeded,
|
|
_PrefillState,
|
|
)
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _make_scheduler(chunked_prefill: bool = True, step_size: int = 4) -> Scheduler:
|
|
"""Return a Scheduler with a mock model/tokenizer and chunked_prefill config."""
|
|
model = MagicMock()
|
|
model.layers = [] # No attention layers — keeps _build_state_machine simple
|
|
|
|
tokenizer = MagicMock()
|
|
tokenizer.eos_token_id = 2
|
|
|
|
config = SchedulerConfig(
|
|
max_num_seqs=8,
|
|
prefill_step_size=step_size,
|
|
chunked_prefill=chunked_prefill,
|
|
paged_cache_block_size=0, # Disable boundary snapshots
|
|
)
|
|
|
|
scheduler = Scheduler(model=model, tokenizer=tokenizer, config=config)
|
|
|
|
# Replace the real batch_generator factory so insert() returns a uid.
|
|
mock_bg = MagicMock()
|
|
mock_bg.insert.return_value = [42]
|
|
mock_bg.next_generated.return_value = iter([])
|
|
scheduler.batch_generator = mock_bg
|
|
scheduler._current_sampler_params = ()
|
|
|
|
return scheduler
|
|
|
|
|
|
def _make_request(request_id: str = "req-1", n_tokens: int = 10) -> Request:
|
|
"""Return a pre-tokenized request with *n_tokens* prompt tokens."""
|
|
req = Request(
|
|
request_id=request_id,
|
|
prompt=list(range(n_tokens)),
|
|
sampling_params=SamplingParams(max_tokens=32),
|
|
)
|
|
req.prompt_token_ids = list(range(n_tokens))
|
|
req.num_prompt_tokens = n_tokens
|
|
req.remaining_tokens = list(range(n_tokens))
|
|
return req
|
|
|
|
|
|
def _make_prefill_state(
|
|
scheduler: Scheduler, request: Request, n_remaining: int = 20
|
|
) -> _PrefillState:
|
|
"""Build a minimal _PrefillState for direct testing."""
|
|
import mlx.core as mx
|
|
|
|
tokens_remaining = mx.zeros((1, n_remaining), dtype=mx.int32)
|
|
state = _PrefillState(
|
|
request=request,
|
|
cache=[],
|
|
tokens_remaining=tokens_remaining,
|
|
last_token=[99],
|
|
tokens_processed=0,
|
|
base_size=0,
|
|
emitted_boundaries={},
|
|
boundary_enabled=False,
|
|
block_size=0,
|
|
total_length=n_remaining + 1,
|
|
sampler=MagicMock(),
|
|
sm=MagicMock(),
|
|
per_row_lps=[],
|
|
)
|
|
return state
|
|
|
|
|
|
class _RecordingModel:
|
|
def __init__(self, model_type: str):
|
|
self.model_type = model_type
|
|
self.layers = []
|
|
self.chunk_lengths: list[int] = []
|
|
|
|
def __call__(self, tokens, cache=None):
|
|
self.chunk_lengths.append(int(tokens.shape[1]))
|
|
|
|
|
|
def _make_recording_scheduler(
|
|
model_type: str,
|
|
*,
|
|
uses_minimax_m3_positions: bool = False,
|
|
nested_vlm_model_type: str | None = None,
|
|
model_name: str = "",
|
|
) -> tuple[Scheduler, _RecordingModel]:
|
|
model = _RecordingModel(model_type)
|
|
if uses_minimax_m3_positions:
|
|
model._uses_minimax_m3_positions = True
|
|
if nested_vlm_model_type is not None:
|
|
model._vlm_model = SimpleNamespace(
|
|
config=SimpleNamespace(model_type=nested_vlm_model_type)
|
|
)
|
|
tokenizer = MagicMock()
|
|
tokenizer.eos_token_id = 2
|
|
scheduler = Scheduler(
|
|
model=model,
|
|
tokenizer=tokenizer,
|
|
config=SchedulerConfig(
|
|
prefill_step_size=2048,
|
|
chunked_prefill=True,
|
|
paged_cache_block_size=0,
|
|
model_name=model_name,
|
|
),
|
|
)
|
|
return scheduler, model
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# SchedulerConfig
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestSchedulerConfigChunkedPrefill:
|
|
def test_default_is_false(self):
|
|
config = SchedulerConfig()
|
|
assert config.chunked_prefill is False
|
|
|
|
def test_can_be_enabled(self):
|
|
config = SchedulerConfig(chunked_prefill=True)
|
|
assert config.chunked_prefill is True
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _PrefillState
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestPrefillState:
|
|
def test_fields_accessible(self):
|
|
import mlx.core as mx
|
|
|
|
state = _PrefillState(
|
|
request=MagicMock(),
|
|
cache=[],
|
|
tokens_remaining=mx.zeros((1, 5), dtype=mx.int32),
|
|
last_token=[7],
|
|
tokens_processed=0,
|
|
base_size=0,
|
|
emitted_boundaries={},
|
|
boundary_enabled=False,
|
|
block_size=256,
|
|
total_length=6,
|
|
)
|
|
assert state.tokens_processed == 0
|
|
assert state.sampler is None
|
|
assert state.per_row_lps is None
|
|
|
|
def test_insert_params_settable(self):
|
|
import mlx.core as mx
|
|
|
|
state = _PrefillState(
|
|
request=MagicMock(),
|
|
cache=[],
|
|
tokens_remaining=mx.zeros((1, 3), dtype=mx.int32),
|
|
last_token=[1],
|
|
tokens_processed=0,
|
|
base_size=0,
|
|
emitted_boundaries={},
|
|
boundary_enabled=False,
|
|
block_size=256,
|
|
total_length=4,
|
|
)
|
|
state.sampler = "s"
|
|
state.sm = "sm"
|
|
state.per_row_lps = []
|
|
assert state.sampler == "s"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Scheduler queues initialised
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestSchedulerQueues:
|
|
def test_prefilling_queue_exists(self):
|
|
sched = _make_scheduler()
|
|
assert hasattr(sched, "prefilling")
|
|
assert isinstance(sched.prefilling, deque)
|
|
assert len(sched.prefilling) == 0
|
|
|
|
def test_prefill_states_dict_exists(self):
|
|
sched = _make_scheduler()
|
|
assert hasattr(sched, "_prefill_states")
|
|
assert isinstance(sched._prefill_states, dict)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# has_requests includes prefilling
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestHasRequests:
|
|
def test_false_when_all_empty(self):
|
|
sched = _make_scheduler()
|
|
assert not sched.has_requests()
|
|
|
|
def test_true_when_prefilling(self):
|
|
sched = _make_scheduler()
|
|
req = _make_request()
|
|
sched.prefilling.append(req)
|
|
assert sched.has_requests()
|
|
|
|
def test_still_true_with_waiting_only(self):
|
|
sched = _make_scheduler()
|
|
req = _make_request()
|
|
sched.waiting.append(req)
|
|
assert sched.has_requests()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# get_stats includes num_prefilling
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestGetStats:
|
|
def test_num_prefilling_in_stats(self):
|
|
sched = _make_scheduler()
|
|
stats = sched.get_stats()
|
|
assert "num_prefilling" in stats
|
|
assert stats["num_prefilling"] == 0
|
|
|
|
def test_num_prefilling_counts_correctly(self):
|
|
sched = _make_scheduler()
|
|
sched.prefilling.append(_make_request("r1"))
|
|
sched.prefilling.append(_make_request("r2"))
|
|
assert sched.get_stats()["num_prefilling"] == 2
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# GLM adaptive chunked prefill
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestGLMAdaptiveChunkedPrefill:
|
|
def test_glm_uses_adaptive_prefill_chunk_size(self, monkeypatch):
|
|
monkeypatch.delenv("MLX_LM_GLM_DSA_ADAPTIVE_PREFILL_STEP", raising=False)
|
|
monkeypatch.delenv("MLX_LM_GLM_DSA_ADAPTIVE_PREFILL_STEP_SIZE", raising=False)
|
|
monkeypatch.delenv("MLX_LM_GLM_DSA_ADAPTIVE_PREFILL_AFTER", raising=False)
|
|
monkeypatch.delenv(
|
|
"MLX_LM_GLM_DSA_ADAPTIVE_PREFILL_MIN_REMAINING", raising=False
|
|
)
|
|
|
|
sched, model = _make_recording_scheduler("glm_moe_dsa")
|
|
req = _make_request("glm", n_tokens=8194)
|
|
state = _make_prefill_state(sched, req, n_remaining=8193)
|
|
|
|
with patch("omlx.scheduler._sync_and_clear_cache"):
|
|
done = sched._step_prefill_chunk(state)
|
|
|
|
assert not done
|
|
assert model.chunk_lengths == [8192]
|
|
assert state.tokens_processed == 8192
|
|
|
|
def test_non_glm_keeps_configured_prefill_chunk_size(self, monkeypatch):
|
|
monkeypatch.delenv("MLX_LM_GLM_DSA_ADAPTIVE_PREFILL_STEP", raising=False)
|
|
|
|
sched, model = _make_recording_scheduler("deepseek_v32")
|
|
req = _make_request("deepseek", n_tokens=8193)
|
|
state = _make_prefill_state(sched, req, n_remaining=8192)
|
|
|
|
with patch("omlx.scheduler._sync_and_clear_cache"):
|
|
done = sched._step_prefill_chunk(state)
|
|
|
|
assert not done
|
|
assert model.chunk_lengths == [2048]
|
|
assert state.tokens_processed == 2048
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# MiniMax M3 adaptive chunked prefill
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestMiniMaxM3AdaptiveChunkedPrefill:
|
|
def test_minimax_m3_uses_4096_for_long_prefill(self, monkeypatch):
|
|
monkeypatch.delenv("MLX_MINIMAX_M3_ADAPTIVE_PREFILL_STEP", raising=False)
|
|
monkeypatch.delenv("MLX_MINIMAX_M3_ADAPTIVE_PREFILL_STEP_SIZE", raising=False)
|
|
monkeypatch.delenv("MLX_MINIMAX_M3_ADAPTIVE_PREFILL_AFTER", raising=False)
|
|
monkeypatch.delenv(
|
|
"MLX_MINIMAX_M3_ADAPTIVE_PREFILL_MIN_REMAINING", raising=False
|
|
)
|
|
|
|
sched, model = _make_recording_scheduler("minimax_m3")
|
|
req = _make_request("minimax", n_tokens=4098)
|
|
state = _make_prefill_state(sched, req, n_remaining=4097)
|
|
|
|
with patch("omlx.scheduler._sync_and_clear_cache"):
|
|
done = sched._step_prefill_chunk(state)
|
|
|
|
assert not done
|
|
assert model.chunk_lengths == [4096]
|
|
assert state.tokens_processed == 4096
|
|
|
|
def test_minimax_m3_keeps_2048_for_short_prefill(self, monkeypatch):
|
|
monkeypatch.delenv("MLX_MINIMAX_M3_ADAPTIVE_PREFILL_STEP", raising=False)
|
|
|
|
sched, model = _make_recording_scheduler("minimax_m3_vl")
|
|
req = _make_request("minimax-short", n_tokens=4096)
|
|
state = _make_prefill_state(sched, req, n_remaining=4095)
|
|
|
|
with patch("omlx.scheduler._sync_and_clear_cache"):
|
|
done = sched._step_prefill_chunk(state)
|
|
|
|
assert not done
|
|
assert model.chunk_lengths == [2048]
|
|
assert state.tokens_processed == 2048
|
|
|
|
def test_minimax_m3_env_can_disable_adaptive_prefill(self, monkeypatch):
|
|
monkeypatch.setenv("MLX_MINIMAX_M3_ADAPTIVE_PREFILL_STEP", "0")
|
|
|
|
sched, model = _make_recording_scheduler("minimax_m3")
|
|
req = _make_request("minimax-disabled", n_tokens=4098)
|
|
state = _make_prefill_state(sched, req, n_remaining=4097)
|
|
|
|
with patch("omlx.scheduler._sync_and_clear_cache"):
|
|
done = sched._step_prefill_chunk(state)
|
|
|
|
assert not done
|
|
assert model.chunk_lengths == [2048]
|
|
assert state.tokens_processed == 2048
|
|
|
|
def test_minimax_m3_vlm_adapter_flag_enables_adaptive_prefill(self, monkeypatch):
|
|
monkeypatch.delenv("MLX_MINIMAX_M3_ADAPTIVE_PREFILL_STEP", raising=False)
|
|
|
|
sched, model = _make_recording_scheduler(
|
|
"vlm",
|
|
uses_minimax_m3_positions=True,
|
|
)
|
|
req = _make_request("minimax-adapter", n_tokens=4098)
|
|
state = _make_prefill_state(sched, req, n_remaining=4097)
|
|
|
|
with patch("omlx.scheduler._sync_and_clear_cache"):
|
|
done = sched._step_prefill_chunk(state)
|
|
|
|
assert not done
|
|
assert model.chunk_lengths == [4096]
|
|
assert state.tokens_processed == 4096
|
|
|
|
def test_minimax_m3_nested_vlm_model_enables_adaptive_prefill(self, monkeypatch):
|
|
monkeypatch.delenv("MLX_MINIMAX_M3_ADAPTIVE_PREFILL_STEP", raising=False)
|
|
|
|
sched, model = _make_recording_scheduler(
|
|
"vlm",
|
|
nested_vlm_model_type="minimax_m3_vl",
|
|
)
|
|
req = _make_request("minimax-nested-vlm", n_tokens=4098)
|
|
state = _make_prefill_state(sched, req, n_remaining=4097)
|
|
|
|
with patch("omlx.scheduler._sync_and_clear_cache"):
|
|
done = sched._step_prefill_chunk(state)
|
|
|
|
assert not done
|
|
assert model.chunk_lengths == [4096]
|
|
assert state.tokens_processed == 4096
|
|
|
|
def test_minimax_m3_model_path_enables_adaptive_prefill(
|
|
self, tmp_path, monkeypatch
|
|
):
|
|
monkeypatch.delenv("MLX_MINIMAX_M3_ADAPTIVE_PREFILL_STEP", raising=False)
|
|
(tmp_path / "config.json").write_text(
|
|
'{"model_type": "minimax_m3_vl"}',
|
|
encoding="utf-8",
|
|
)
|
|
|
|
sched, model = _make_recording_scheduler(
|
|
"vlm",
|
|
model_name=str(tmp_path),
|
|
)
|
|
req = _make_request("minimax-model-path", n_tokens=4098)
|
|
state = _make_prefill_state(sched, req, n_remaining=4097)
|
|
|
|
with patch("omlx.scheduler._sync_and_clear_cache"):
|
|
done = sched._step_prefill_chunk(state)
|
|
|
|
assert not done
|
|
assert model.chunk_lengths == [4096]
|
|
assert state.tokens_processed == 4096
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# reset() clears prefilling
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestReset:
|
|
def test_reset_clears_prefilling(self):
|
|
sched = _make_scheduler()
|
|
req = _make_request()
|
|
sched.prefilling.append(req)
|
|
sched._prefill_states[req.request_id] = MagicMock()
|
|
sched.requests[req.request_id] = req
|
|
|
|
sched.reset()
|
|
|
|
assert len(sched.prefilling) == 0
|
|
assert len(sched._prefill_states) == 0
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# fail_all_requests() includes prefilling
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestFailAllRequests:
|
|
def test_fail_all_includes_prefilling(self):
|
|
sched = _make_scheduler()
|
|
req = _make_request("pf-req")
|
|
sched.prefilling.append(req)
|
|
sched._prefill_states[req.request_id] = MagicMock()
|
|
sched.requests[req.request_id] = req
|
|
|
|
failed = sched.fail_all_requests()
|
|
|
|
assert "pf-req" in failed
|
|
assert len(sched.prefilling) == 0
|
|
assert len(sched._prefill_states) == 0
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _do_abort_request() cleans up prefilling
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestAbortPrefilling:
|
|
def test_abort_removes_from_prefilling(self):
|
|
sched = _make_scheduler()
|
|
req = _make_request("abort-me")
|
|
req.status = RequestStatus.WAITING
|
|
sched.prefilling.append(req)
|
|
sched._prefill_states[req.request_id] = MagicMock()
|
|
sched.requests[req.request_id] = req
|
|
|
|
sched._do_abort_request(req.request_id)
|
|
|
|
assert req.request_id not in sched._prefill_states
|
|
assert all(r.request_id != req.request_id for r in sched.prefilling)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _advance_chunked_prefills(): core logic
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestAdvanceChunkedPrefills:
|
|
def test_no_op_when_queue_empty(self):
|
|
sched = _make_scheduler()
|
|
scheduled = []
|
|
rejected = []
|
|
# Should not raise
|
|
sched._advance_chunked_prefills(scheduled, rejected)
|
|
assert scheduled == []
|
|
assert rejected == []
|
|
|
|
def test_advances_chunk_when_not_done(self):
|
|
"""Requests that still have tokens stay in prefilling queue."""
|
|
sched = _make_scheduler()
|
|
req = _make_request("r1")
|
|
sched.requests[req.request_id] = req
|
|
state = _make_prefill_state(sched, req, n_remaining=20)
|
|
sched.prefilling.append(req)
|
|
sched._prefill_states[req.request_id] = state
|
|
|
|
with patch.object(
|
|
sched, "_step_prefill_chunk", return_value=False
|
|
) as mock_step:
|
|
scheduled = []
|
|
rejected = []
|
|
sched._advance_chunked_prefills(scheduled, rejected)
|
|
|
|
mock_step.assert_called_once_with(state)
|
|
# Not done → stays in prefilling, not moved to running
|
|
assert req in sched.prefilling
|
|
assert scheduled == []
|
|
assert rejected == []
|
|
assert req.request_id not in sched.running
|
|
|
|
def test_inserts_when_done(self):
|
|
"""Completed prefill is inserted into BatchGenerator and moved to running."""
|
|
sched = _make_scheduler()
|
|
req = _make_request("r1")
|
|
sched.requests[req.request_id] = req
|
|
state = _make_prefill_state(sched, req, n_remaining=1)
|
|
state.sampler = MagicMock()
|
|
state.sm = MagicMock()
|
|
state.per_row_lps = []
|
|
sched.prefilling.append(req)
|
|
sched._prefill_states[req.request_id] = state
|
|
|
|
with patch.object(sched, "_step_prefill_chunk", return_value=True):
|
|
with patch.object(sched, "_emit_final_boundary_if_needed"):
|
|
scheduled = []
|
|
rejected = []
|
|
sched._advance_chunked_prefills(scheduled, rejected)
|
|
|
|
# Moved to running, removed from prefilling
|
|
assert req not in sched.prefilling
|
|
assert req.request_id not in sched._prefill_states
|
|
assert req.request_id in sched.running
|
|
assert req in scheduled
|
|
assert rejected == []
|
|
assert req.status == RequestStatus.RUNNING
|
|
|
|
def test_skips_aborted_request(self):
|
|
"""Request whose state was cleared by abort is silently skipped."""
|
|
sched = _make_scheduler()
|
|
req = _make_request("gone")
|
|
# State NOT added to _prefill_states (simulates post-abort cleanup)
|
|
sched.prefilling.append(req)
|
|
|
|
scheduled = []
|
|
rejected = []
|
|
sched._advance_chunked_prefills(scheduled, rejected) # Must not raise
|
|
|
|
assert scheduled == []
|
|
assert rejected == []
|
|
assert len(sched.prefilling) == 0
|
|
|
|
def test_abort_during_chunk_discards_state(self):
|
|
"""_PrefillAbortedError from _step_prefill_chunk is swallowed cleanly."""
|
|
sched = _make_scheduler()
|
|
req = _make_request("r1")
|
|
sched.requests[req.request_id] = req
|
|
state = _make_prefill_state(sched, req)
|
|
sched.prefilling.append(req)
|
|
sched._prefill_states[req.request_id] = state
|
|
|
|
with patch.object(
|
|
sched, "_step_prefill_chunk", side_effect=_PrefillAbortedError([], 4)
|
|
):
|
|
scheduled = []
|
|
rejected = []
|
|
sched._advance_chunked_prefills(scheduled, rejected) # Must not raise
|
|
|
|
assert req.request_id not in sched._prefill_states
|
|
assert req not in sched.prefilling
|
|
assert scheduled == []
|
|
assert rejected == []
|
|
|
|
def test_runtime_error_surfaces_as_request_error(self):
|
|
"""A non-memory RuntimeError mid-chunk yields a finish_reason="error"
|
|
RequestOutput immediately (only memory-pressure errors are requeued)."""
|
|
sched = _make_scheduler()
|
|
req = _make_request("oom")
|
|
sched.requests[req.request_id] = req
|
|
state = _make_prefill_state(sched, req)
|
|
sched.prefilling.append(req)
|
|
sched._prefill_states[req.request_id] = state
|
|
|
|
with patch.object(
|
|
sched, "_step_prefill_chunk", side_effect=RuntimeError("kernel panic")
|
|
):
|
|
scheduled = []
|
|
rejected = []
|
|
sched._advance_chunked_prefills(scheduled, rejected)
|
|
|
|
assert req.request_id not in sched._prefill_states
|
|
assert req not in sched.prefilling
|
|
assert req.request_id not in sched.requests
|
|
assert scheduled == []
|
|
assert len(rejected) == 1
|
|
out = rejected[0]
|
|
assert out.request_id == "oom"
|
|
assert out.finished is True
|
|
assert out.finish_reason == "error"
|
|
assert "kernel panic" in out.error
|
|
|
|
def test_memory_error_requeues_instead_of_surfacing(self):
|
|
"""A memory-pressure RuntimeError mid-chunk requeues the request for a
|
|
fresh attempt instead of immediately surfacing an error to the client."""
|
|
sched = _make_scheduler()
|
|
req = _make_request("oom-mem")
|
|
sched.requests[req.request_id] = req
|
|
state = _make_prefill_state(sched, req)
|
|
sched.prefilling.append(req)
|
|
sched._prefill_states[req.request_id] = state
|
|
|
|
with patch.object(
|
|
sched,
|
|
"_step_prefill_chunk",
|
|
side_effect=RuntimeError("Memory limit exceeded during chunked prefill"),
|
|
):
|
|
scheduled = []
|
|
rejected = []
|
|
sched._advance_chunked_prefills(scheduled, rejected)
|
|
|
|
# No client-facing error; the request is reset and back on the queue.
|
|
assert rejected == []
|
|
assert req.request_id not in sched._prefill_states
|
|
assert req not in sched.prefilling
|
|
assert sched.requests.get(req.request_id) is req
|
|
assert req in sched.waiting
|
|
assert req.prefill_oom_retries == 1
|
|
|
|
def test_capacity_error_surfaces_as_typed_request_error(self):
|
|
"""A deterministic capacity rejection is not retried as transient OOM."""
|
|
sched = _make_scheduler()
|
|
req = _make_request("capacity")
|
|
sched.requests[req.request_id] = req
|
|
state = _make_prefill_state(sched, req)
|
|
sched.prefilling.append(req)
|
|
sched._prefill_states[req.request_id] = state
|
|
|
|
err = PrefillMemoryExceededError(
|
|
message="Prefill context too large for available memory",
|
|
request_id=req.request_id,
|
|
estimated_bytes=123,
|
|
limit_bytes=100,
|
|
)
|
|
with patch.object(sched, "_step_prefill_chunk", side_effect=err):
|
|
scheduled = []
|
|
rejected = []
|
|
sched._advance_chunked_prefills(scheduled, rejected)
|
|
|
|
assert scheduled == []
|
|
assert len(rejected) == 1
|
|
out = rejected[0]
|
|
assert out.error == str(err)
|
|
assert out.error_code == "prefill_memory_exceeded"
|
|
assert out.error_metadata == {
|
|
"request_id": req.request_id,
|
|
"estimated_bytes": 123,
|
|
"limit_bytes": 100,
|
|
}
|
|
assert req.prefill_oom_retries == 0
|
|
|
|
def test_multiple_requests_all_advanced(self):
|
|
"""All requests in prefilling get one chunk advanced per call."""
|
|
sched = _make_scheduler()
|
|
reqs = [_make_request(f"r{i}") for i in range(3)]
|
|
for req in reqs:
|
|
sched.requests[req.request_id] = req
|
|
state = _make_prefill_state(sched, req, n_remaining=20)
|
|
state.sampler = MagicMock()
|
|
state.sm = MagicMock()
|
|
state.per_row_lps = []
|
|
sched.prefilling.append(req)
|
|
sched._prefill_states[req.request_id] = state
|
|
|
|
call_count = 0
|
|
|
|
def fake_step(state):
|
|
nonlocal call_count
|
|
call_count += 1
|
|
return False # All still in-progress
|
|
|
|
with patch.object(sched, "_step_prefill_chunk", side_effect=fake_step):
|
|
sched._advance_chunked_prefills([], [])
|
|
|
|
assert call_count == 3 # One chunk per request
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# _schedule_waiting(): chunked fork is taken for long prompts
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestScheduleWaitingChunkedFork:
|
|
def _setup(self, n_tokens: int, chunked: bool = True, step_size: int = 4):
|
|
sched = _make_scheduler(chunked_prefill=chunked, step_size=step_size)
|
|
req = _make_request("r1", n_tokens=n_tokens)
|
|
sched.add_request(req)
|
|
return sched, req
|
|
|
|
def test_short_prompt_stays_on_normal_path(self):
|
|
"""Prompts that fit in one chunk use the normal prefill path."""
|
|
# step_size=4, prompt=3 tokens → not long enough to trigger chunked fork
|
|
sched, req = self._setup(n_tokens=3, step_size=4)
|
|
|
|
with patch.object(
|
|
sched, "_do_external_prefill", return_value=([], [0])
|
|
) as mock_ep:
|
|
with patch.object(sched, "_begin_prefill") as mock_bp:
|
|
sched._schedule_waiting()
|
|
|
|
mock_ep.assert_called_once()
|
|
mock_bp.assert_not_called()
|
|
|
|
def test_long_prompt_enters_prefilling_queue(self):
|
|
"""Prompts longer than step_size+1 enter the chunked prefill queue."""
|
|
# step_size=4, 10 tokens → triggers chunked path
|
|
sched, req = self._setup(n_tokens=10, step_size=4)
|
|
|
|
with patch.object(
|
|
sched, "_begin_prefill", return_value=_make_prefill_state(sched, req)
|
|
) as mock_bp:
|
|
with patch.object(sched, "_step_prefill_chunk", return_value=False):
|
|
sched._schedule_waiting()
|
|
|
|
mock_bp.assert_called_once()
|
|
assert req.request_id in sched._prefill_states
|
|
assert req in sched.prefilling
|
|
assert req.request_id not in sched.running
|
|
|
|
def test_prefilling_request_counts_against_concurrency_cap(self):
|
|
"""A chunked prefill already in flight consumes a scheduler slot."""
|
|
sched = _make_scheduler(chunked_prefill=True, step_size=4)
|
|
sched.config.max_num_seqs = 1
|
|
|
|
inflight = _make_request("inflight", n_tokens=10)
|
|
sched.requests[inflight.request_id] = inflight
|
|
sched.prefilling.append(inflight)
|
|
sched._prefill_states[inflight.request_id] = _make_prefill_state(
|
|
sched,
|
|
inflight,
|
|
)
|
|
|
|
queued = _make_request("queued", n_tokens=10)
|
|
sched.add_request(queued)
|
|
|
|
with patch.object(sched, "_begin_prefill") as mock_begin:
|
|
scheduled, rejected = sched._schedule_waiting()
|
|
|
|
mock_begin.assert_not_called()
|
|
assert scheduled == []
|
|
assert rejected == []
|
|
assert queued in sched.waiting
|
|
assert inflight in sched.prefilling
|
|
|
|
def test_long_prompt_completes_in_first_chunk_goes_to_running(self):
|
|
"""If the first chunk happens to finish the prefill, request goes to running."""
|
|
sched, req = self._setup(n_tokens=10, step_size=4)
|
|
fake_state = _make_prefill_state(sched, req, n_remaining=1)
|
|
|
|
with patch.object(sched, "_begin_prefill", return_value=fake_state):
|
|
with patch.object(sched, "_step_prefill_chunk", return_value=True):
|
|
with patch.object(sched, "_emit_final_boundary_if_needed"):
|
|
with patch("omlx.scheduler._sync_and_clear_cache"):
|
|
sched._schedule_waiting()
|
|
|
|
assert req.request_id not in sched._prefill_states
|
|
assert req not in sched.prefilling
|
|
assert req.request_id in sched.running
|
|
|
|
def test_chunked_disabled_uses_normal_path(self):
|
|
"""chunked_prefill=False always uses the full-prefill path."""
|
|
sched, req = self._setup(n_tokens=100, chunked=False, step_size=4)
|
|
|
|
with patch.object(
|
|
sched, "_do_external_prefill", return_value=([], [0])
|
|
) as mock_ep:
|
|
with patch.object(sched, "_begin_prefill") as mock_bp:
|
|
sched._schedule_waiting()
|
|
|
|
mock_ep.assert_called_once()
|
|
mock_bp.assert_not_called()
|
|
|
|
def test_non_chunked_path_runtime_error_cleans_up_and_rejects(self):
|
|
"""RuntimeError from _do_external_prefill in the non-chunked path
|
|
must pop self.requests, drop the temp uid mappings, remove the
|
|
PrefillProgressTracker entry, and emit a finish_reason=\"error\"
|
|
RequestOutput so the client sees the failure (#1405)."""
|
|
from omlx.prefill_progress import get_prefill_tracker
|
|
|
|
sched, req = self._setup(n_tokens=3, step_size=4)
|
|
rid = req.request_id
|
|
tracker = get_prefill_tracker()
|
|
tracker.clear()
|
|
tracker.update(rid, processed=1, total=3, model_id="test")
|
|
assert tracker.get_model_progress("test"), "tracker entry not set up"
|
|
|
|
try:
|
|
with patch.object(
|
|
sched,
|
|
"_do_external_prefill",
|
|
side_effect=RuntimeError("Memory limit exceeded during prefill"),
|
|
):
|
|
scheduled, rejected = sched._schedule_waiting()
|
|
|
|
assert rid not in sched.requests
|
|
assert rid not in sched.request_id_to_uid
|
|
assert not any(v == rid for v in sched.uid_to_request_id.values())
|
|
assert tracker.get_model_progress("test") == []
|
|
assert scheduled == []
|
|
assert len(rejected) == 1
|
|
out = rejected[0]
|
|
assert out.request_id == rid
|
|
assert out.finished is True
|
|
assert out.finish_reason == "error"
|
|
assert "Memory limit" in out.error
|
|
finally:
|
|
tracker.clear()
|
|
|
|
def _setup_throttle(self, max_bytes_gb=10, hard_cap_gb=12):
|
|
"""Build a scheduler with watermark fields set for throttle tests."""
|
|
sched = _make_scheduler()
|
|
sched._memory_limit_bytes = max_bytes_gb * 1024**3
|
|
sched._memory_hard_limit_bytes = hard_cap_gb * 1024**3
|
|
sched._prefill_safe_zone_ratio = 0.80
|
|
sched._prefill_min_chunk_tokens = 32
|
|
return sched
|
|
|
|
def _mock_current(self, sched, current_gb):
|
|
"""Context manager-ish — patch both memory probes to current_gb."""
|
|
target = int(current_gb * 1024**3)
|
|
return patch("omlx.scheduler.mx.get_active_memory", return_value=target), patch(
|
|
"omlx.scheduler.get_phys_footprint", return_value=target
|
|
)
|
|
|
|
def test_adaptive_throttle_below_soft_watermark_passthrough(self):
|
|
"""current < soft watermark → no throttle, full chunk."""
|
|
sched = self._setup_throttle(max_bytes_gb=10, hard_cap_gb=12)
|
|
# soft_watermark = 10 * 0.80 = 8 GB; current 5 GB is below
|
|
a, b = self._mock_current(sched, 5)
|
|
with a, b:
|
|
result = sched._adaptive_chunk_size(
|
|
2048, request_id="r1", loop_label="external"
|
|
)
|
|
assert result == 2048
|
|
|
|
def test_adaptive_throttle_tier_1024(self):
|
|
"""First quarter of the soft-to-hard band → 1024."""
|
|
sched = self._setup_throttle(max_bytes_gb=10, hard_cap_gb=12)
|
|
# soft_wm = 8 GB, band = 12 - 8 = 4 GB. 10% into band = 8.4 GB.
|
|
a, b = self._mock_current(sched, 8.4)
|
|
with a, b:
|
|
result = sched._adaptive_chunk_size(
|
|
2048, request_id="r1", loop_label="external"
|
|
)
|
|
assert result == 1024
|
|
|
|
def test_adaptive_throttle_tier_512(self):
|
|
"""50%+ of band → 512."""
|
|
sched = self._setup_throttle(max_bytes_gb=10, hard_cap_gb=12)
|
|
# 60% of band: 8 + 4*0.60 = 10.4 GB
|
|
a, b = self._mock_current(sched, 10.4)
|
|
with a, b:
|
|
result = sched._adaptive_chunk_size(
|
|
2048, request_id="r1", loop_label="external"
|
|
)
|
|
assert result == 512
|
|
|
|
def test_adaptive_throttle_requested_smaller_than_tier(self):
|
|
"""Requested chunk already smaller than the tier target → pass through."""
|
|
sched = self._setup_throttle(max_bytes_gb=10, hard_cap_gb=12)
|
|
# 60% of band → tier 512. But requested=256 < 512.
|
|
a, b = self._mock_current(sched, 10.4)
|
|
with a, b:
|
|
result = sched._adaptive_chunk_size(
|
|
256, request_id="r1", loop_label="external"
|
|
)
|
|
assert result == 256
|
|
|
|
def test_adaptive_throttle_no_cap_passthrough(self):
|
|
"""When hard limit or soft base is unset (=0), no throttle."""
|
|
sched = self._setup_throttle()
|
|
sched._memory_hard_limit_bytes = 0
|
|
result = sched._adaptive_chunk_size(
|
|
2048, request_id="r1", loop_label="external"
|
|
)
|
|
assert result == 2048
|
|
|
|
sched._memory_hard_limit_bytes = 10 * 1024**3
|
|
sched._memory_limit_bytes = 0
|
|
result = sched._adaptive_chunk_size(
|
|
2048, request_id="r1", loop_label="external"
|
|
)
|
|
assert result == 2048
|
|
|
|
def test_chunked_first_chunk_runtime_error_cleans_up_and_rejects(self):
|
|
"""RuntimeError on the chunked first chunk must pop self.requests,
|
|
remove the PrefillProgressTracker entry, and emit an error
|
|
RequestOutput. _step_prefill_chunk updates the tracker before the
|
|
hard-limit check, so without this catch the entry would leak
|
|
(#1405)."""
|
|
from omlx.prefill_progress import get_prefill_tracker
|
|
|
|
sched, req = self._setup(n_tokens=10, step_size=4)
|
|
rid = req.request_id
|
|
tracker = get_prefill_tracker()
|
|
tracker.clear()
|
|
tracker.update(rid, processed=2, total=10, model_id="test")
|
|
assert tracker.get_model_progress("test"), "tracker entry not set up"
|
|
|
|
try:
|
|
with patch.object(
|
|
sched,
|
|
"_begin_prefill",
|
|
return_value=_make_prefill_state(sched, req),
|
|
):
|
|
with patch.object(
|
|
sched,
|
|
"_step_prefill_chunk",
|
|
side_effect=RuntimeError(
|
|
"Memory limit exceeded during chunked prefill"
|
|
),
|
|
):
|
|
scheduled, rejected = sched._schedule_waiting()
|
|
|
|
assert rid not in sched.requests
|
|
assert rid not in sched._prefill_states
|
|
assert req not in sched.prefilling
|
|
assert tracker.get_model_progress("test") == []
|
|
assert scheduled == []
|
|
assert len(rejected) == 1
|
|
out = rejected[0]
|
|
assert out.request_id == rid
|
|
assert out.finished is True
|
|
assert out.finish_reason == "error"
|
|
assert "Memory limit" in out.error
|
|
finally:
|
|
tracker.clear()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Prefill-rejection paged-cache cleanup
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestPrefillRejectionReleasesPagedCache:
|
|
"""Rejection paths must release block_aware_cache refs / paged_cache
|
|
block_table entries that ``add_request`` populated via ``fetch_cache``.
|
|
|
|
Without this, every rejected request leaks an entry in
|
|
``BlockAwarePrefixCache._request_tables`` plus the ref counts on its
|
|
prefix-matched blocks — pinning the paged cache and compounding the
|
|
very memory pressure that triggered the rejection. The existing
|
|
``self.requests.pop(...)`` and ``get_prefill_tracker().remove(...)``
|
|
cleanups handle scheduler-side state but never reach into the
|
|
paged-cache layer.
|
|
"""
|
|
|
|
def test_helper_calls_block_aware_cache_release(self):
|
|
"""The helper delegates to block_aware_cache.release_cache when one
|
|
is attached — the normal production wiring."""
|
|
sched = _make_scheduler()
|
|
sched.block_aware_cache = MagicMock()
|
|
sched.paged_cache_manager = MagicMock()
|
|
|
|
sched._release_paged_cache_for_request("rid-1")
|
|
|
|
sched.block_aware_cache.release_cache.assert_called_once_with("rid-1")
|
|
# release_cache delegates to delete_block_table internally; the
|
|
# helper must NOT also call it directly (double-delete).
|
|
sched.paged_cache_manager.delete_block_table.assert_not_called()
|
|
|
|
def test_helper_falls_back_to_paged_cache_manager(self):
|
|
"""Without a BlockAwarePrefixCache, fall back to deleting the block
|
|
table directly on the paged cache manager."""
|
|
sched = _make_scheduler()
|
|
sched.block_aware_cache = None
|
|
sched.paged_cache_manager = MagicMock()
|
|
|
|
sched._release_paged_cache_for_request("rid-2")
|
|
|
|
sched.paged_cache_manager.delete_block_table.assert_called_once_with("rid-2")
|
|
|
|
def test_helper_is_noop_without_any_paged_cache(self):
|
|
"""No paged-cache layer attached → silent no-op."""
|
|
sched = _make_scheduler()
|
|
sched.block_aware_cache = None
|
|
sched.paged_cache_manager = None
|
|
|
|
# Should not raise.
|
|
sched._release_paged_cache_for_request("rid-3")
|
|
|
|
def test_advance_chunked_prefills_releases_on_runtime_error(self):
|
|
"""_advance_chunked_prefills' RuntimeError handler must call
|
|
release_cache so the paged-cache block refs from the request's
|
|
prefix-cache lookup don't leak."""
|
|
sched = _make_scheduler()
|
|
sched.block_aware_cache = MagicMock()
|
|
req = _make_request("oom-chunked")
|
|
sched.requests[req.request_id] = req
|
|
state = _make_prefill_state(sched, req)
|
|
sched.prefilling.append(req)
|
|
sched._prefill_states[req.request_id] = state
|
|
|
|
with patch.object(
|
|
sched,
|
|
"_step_prefill_chunk",
|
|
side_effect=RuntimeError("Memory limit exceeded"),
|
|
):
|
|
sched._advance_chunked_prefills([], [])
|
|
|
|
sched.block_aware_cache.release_cache.assert_called_once_with("oom-chunked")
|
|
|
|
def test_schedule_waiting_non_chunked_releases_on_runtime_error(self):
|
|
"""The non-chunked _do_external_prefill rejection path must release
|
|
the paged-cache footprint before popping self.requests."""
|
|
sched = _make_scheduler(step_size=4)
|
|
sched.block_aware_cache = MagicMock()
|
|
# No prefix-cache hit: fetch_cache returns (None, prompt_tokens) so
|
|
# add_request falls through to the waiting queue without trying to
|
|
# preload/reconstruct.
|
|
sched.block_aware_cache.fetch_cache.return_value = (None, [0, 1, 2])
|
|
req = _make_request("oom-direct", n_tokens=3)
|
|
sched.add_request(req)
|
|
sched.block_aware_cache.reset_mock()
|
|
|
|
with patch.object(
|
|
sched,
|
|
"_do_external_prefill",
|
|
side_effect=RuntimeError("kernel panic"),
|
|
):
|
|
sched._schedule_waiting()
|
|
|
|
sched.block_aware_cache.release_cache.assert_called_once_with("oom-direct")
|
|
|
|
def test_schedule_waiting_chunked_first_chunk_releases_on_runtime_error(self):
|
|
"""The chunked first-chunk rejection path must release the
|
|
paged-cache footprint before popping self.requests."""
|
|
sched = _make_scheduler(step_size=4)
|
|
sched.block_aware_cache = MagicMock()
|
|
sched.block_aware_cache.fetch_cache.return_value = (None, list(range(10)))
|
|
req = _make_request("oom-first-chunk", n_tokens=10)
|
|
sched.add_request(req)
|
|
sched.block_aware_cache.reset_mock()
|
|
|
|
with patch.object(
|
|
sched,
|
|
"_begin_prefill",
|
|
return_value=_make_prefill_state(sched, req),
|
|
):
|
|
with patch.object(
|
|
sched,
|
|
"_step_prefill_chunk",
|
|
side_effect=RuntimeError("kernel panic"),
|
|
):
|
|
sched._schedule_waiting()
|
|
|
|
sched.block_aware_cache.release_cache.assert_called_once_with("oom-first-chunk")
|
|
|
|
def test_schedule_waiting_preflight_rejection_releases(self):
|
|
"""_preflight_memory_check rejection (the non-RuntimeError path
|
|
inside _schedule_waiting) must also release the paged-cache
|
|
footprint. Same leak shape as the RuntimeError rejections — the
|
|
request reached this point via add_request → fetch_cache so
|
|
_request_tables is populated and prefix block refs are held."""
|
|
sched = _make_scheduler(step_size=4)
|
|
sched.block_aware_cache = MagicMock()
|
|
sched.block_aware_cache.fetch_cache.return_value = (None, list(range(5)))
|
|
req = _make_request("oom-preflight", n_tokens=5)
|
|
sched.add_request(req)
|
|
sched.block_aware_cache.reset_mock()
|
|
|
|
from omlx.scheduler import _PreflightRejection
|
|
|
|
with patch.object(
|
|
sched,
|
|
"_preflight_memory_check",
|
|
return_value=_PreflightRejection(
|
|
message="Memory limit exceeded by preflight estimate",
|
|
estimated_bytes=1,
|
|
limit_bytes=1,
|
|
),
|
|
):
|
|
scheduled, rejected = sched._schedule_waiting()
|
|
|
|
assert scheduled == []
|
|
assert len(rejected) == 1
|
|
assert rejected[0].request_id == "oom-preflight"
|
|
assert rejected[0].finish_reason == "error"
|
|
sched.block_aware_cache.release_cache.assert_called_once_with("oom-preflight")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# First-chunk eviction pause must preserve a reconstructed prefix (#2180)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestFirstChunkEvictionPreservesPrefix:
|
|
def test_first_chunk_eviction_pause_keeps_reconstructed_prefix(self):
|
|
"""_PrefillEvictionNeeded raised before the first chunk's forward
|
|
pass must not discard a reconstructed SSD prefix. The eviction pause
|
|
keeps prompt_cache / block_table / cached_tokens / remaining_tokens
|
|
attached, so when no idle model can be evicted the retry prefills
|
|
only the uncached suffix instead of recomputing the whole prompt
|
|
cold (#2180)."""
|
|
sched = _make_scheduler(step_size=4)
|
|
sched.block_aware_cache = MagicMock()
|
|
sched.block_aware_cache.fetch_cache.return_value = (None, list(range(100)))
|
|
req = _make_request("evict-first-chunk", n_tokens=100)
|
|
sched.add_request(req)
|
|
sched.block_aware_cache.reset_mock()
|
|
|
|
# Simulate the state _prepare_prefix_cache_for_request leaves after a
|
|
# successful paged/SSD cache hit + reconstruction: 90 cached tokens,
|
|
# a 10-token uncached suffix, and a live block table.
|
|
prompt_cache = [MagicMock()]
|
|
block_table = MagicMock()
|
|
sched._prefix_cache_prepared.add(req.request_id)
|
|
req.prompt_cache = prompt_cache
|
|
req.cached_tokens = 90
|
|
req.remaining_tokens = req.prompt_token_ids[90:]
|
|
req.block_table = block_table
|
|
req.shared_prefix_blocks = 3
|
|
|
|
eviction = PrefillEvictionRequest(
|
|
request_id=req.request_id,
|
|
model_id="test",
|
|
current_bytes=1,
|
|
target_cap_bytes=1,
|
|
predicted_transient_bytes=1,
|
|
requested_tokens=4,
|
|
reason="adaptive_prefill_throttle",
|
|
)
|
|
with patch.object(
|
|
sched,
|
|
"_begin_prefill",
|
|
return_value=_make_prefill_state(sched, req),
|
|
):
|
|
with patch.object(
|
|
sched,
|
|
"_step_prefill_chunk",
|
|
side_effect=_PrefillEvictionNeeded(eviction),
|
|
):
|
|
scheduled, rejected = sched._schedule_waiting()
|
|
|
|
assert scheduled == []
|
|
assert rejected == []
|
|
# Paused back into the waiting queue with the eviction request pending.
|
|
assert req in sched.waiting
|
|
assert sched._pending_prefill_eviction_request is eviction
|
|
# The reconstructed prefix must survive the pause untouched.
|
|
assert req.prompt_cache is prompt_cache
|
|
assert req.cached_tokens == 90
|
|
assert req.remaining_tokens == req.prompt_token_ids[90:]
|
|
assert req.block_table is block_table
|
|
assert req.shared_prefix_blocks == 3
|
|
sched.block_aware_cache.release_cache.assert_not_called()
|