wehub-resource-sync
260 lines
8.4 KiB
Python
260 lines
8.4 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Tests for the SpecPrefill draft-scoring workflow."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Callable
|
|
from contextlib import nullcontext
|
|
from types import SimpleNamespace
|
|
from typing import Any
|
|
from unittest.mock import patch
|
|
|
|
import mlx.core as mx
|
|
|
|
import omlx.specprefill.draft as draft_workflow
|
|
from omlx.request import Request, SamplingParams
|
|
from omlx.specprefill.policy import plan_specprefill_scoring
|
|
|
|
|
|
class _Logger:
|
|
def __init__(self) -> None:
|
|
self.debug_messages: list[str] = []
|
|
self.info_messages: list[str] = []
|
|
self.error_messages: list[str] = []
|
|
|
|
def debug(self, message: str, *args: Any, **kwargs: Any) -> None:
|
|
self.debug_messages.append(message)
|
|
|
|
def info(self, message: str, *args: Any, **kwargs: Any) -> None:
|
|
self.info_messages.append(message)
|
|
|
|
def error(self, message: str, *args: Any, **kwargs: Any) -> None:
|
|
self.error_messages.append(message)
|
|
|
|
|
|
class _Tracker:
|
|
def __init__(self) -> None:
|
|
self.updates: list[dict[str, Any]] = []
|
|
self.removed: list[str] = []
|
|
|
|
def update(
|
|
self,
|
|
request_id: str,
|
|
processed: int,
|
|
total: int,
|
|
model_id: str,
|
|
phase: str = "prefill",
|
|
detail: str | None = None,
|
|
extra: dict[str, Any] | None = None,
|
|
) -> None:
|
|
self.updates.append(
|
|
{
|
|
"request_id": request_id,
|
|
"processed": processed,
|
|
"total": total,
|
|
"model_id": model_id,
|
|
"phase": phase,
|
|
"detail": detail,
|
|
"extra": extra,
|
|
}
|
|
)
|
|
|
|
def remove(self, request_id: str) -> None:
|
|
self.removed.append(request_id)
|
|
|
|
|
|
class _DraftCache:
|
|
def __init__(
|
|
self,
|
|
block_table: Any = None,
|
|
reconstructed_cache: Any = None,
|
|
fetch_error: Exception | None = None,
|
|
) -> None:
|
|
self.block_table = block_table
|
|
self.reconstructed_cache = reconstructed_cache
|
|
self.fetch_error = fetch_error
|
|
self.fetches: list[tuple[str, list[int]]] = []
|
|
self.preloads: list[Any] = []
|
|
self.reconstructions: list[Any] = []
|
|
self.stores: list[tuple[str, list[int], list[Any], Any]] = []
|
|
|
|
def fetch_cache(self, request_id: str, tokens: list[int]) -> tuple[Any, list[int]]:
|
|
self.fetches.append((request_id, list(tokens)))
|
|
if self.fetch_error is not None:
|
|
raise self.fetch_error
|
|
return self.block_table, []
|
|
|
|
def preload_blocks(self, block_table: Any) -> int:
|
|
self.preloads.append(block_table)
|
|
return block_table.num_tokens
|
|
|
|
def reconstruct_cache(self, block_table: Any) -> Any:
|
|
self.reconstructions.append(block_table)
|
|
return self.reconstructed_cache
|
|
|
|
def store_cache(
|
|
self,
|
|
request_id: str,
|
|
tokens: list[int],
|
|
cache_data: list[Any],
|
|
model_cache_config: Any = None,
|
|
) -> None:
|
|
self.stores.append(
|
|
(request_id, list(tokens), cache_data, model_cache_config)
|
|
)
|
|
|
|
|
|
def _request_and_plan() -> tuple[Request, Any]:
|
|
request = Request(
|
|
request_id="request-1",
|
|
prompt=list(range(20)),
|
|
sampling_params=SamplingParams(),
|
|
)
|
|
request.prompt_token_ids = list(range(20))
|
|
request.num_prompt_tokens = 20
|
|
request.remaining_tokens = request.prompt_token_ids
|
|
request.specprefill_system_end = 4
|
|
request.cached_tokens = 0
|
|
plan = plan_specprefill_scoring(
|
|
remaining_tokens=request.remaining_tokens,
|
|
system_prompt_end=request.specprefill_system_end,
|
|
cached_tokens=request.cached_tokens,
|
|
requested_threshold=None,
|
|
requested_keep_pct=None,
|
|
default_threshold=8,
|
|
default_keep_pct=0.2,
|
|
)
|
|
assert plan is not None
|
|
return request, plan
|
|
|
|
|
|
def _run(
|
|
request: Request,
|
|
plan: Any,
|
|
*,
|
|
draft_cache: _DraftCache | None = None,
|
|
score_tokens: Callable[..., Any] | None = None,
|
|
extract_cache_states: Callable[[list[Any]], tuple[list[dict[str, Any]], Any]] | None = None,
|
|
) -> tuple[_Tracker, _Logger, dict[str, Any]]:
|
|
tracker = _Tracker()
|
|
logger = _Logger()
|
|
selected_indices = mx.arange(3)
|
|
stream = object()
|
|
trace: dict[str, Any] = {"streams": [], "syncs": [], "score_calls": []}
|
|
|
|
def default_score_tokens(
|
|
model: Any, tokens: list[int], **kwargs: Any
|
|
) -> tuple[Any, list[str]]:
|
|
trace["score_calls"].append(kwargs)
|
|
return mx.zeros(plan.n_to_score), ["draft-cache"]
|
|
|
|
def select_chunks(importance: Any, keep_pct: float) -> Any:
|
|
return selected_indices
|
|
|
|
def use_stream(selected_stream: Any):
|
|
trace["streams"].append(selected_stream)
|
|
return nullcontext()
|
|
|
|
with (
|
|
patch.object(draft_workflow, "get_prefill_tracker", return_value=tracker),
|
|
patch(
|
|
"omlx.patches.specprefill.score_tokens",
|
|
side_effect=score_tokens or default_score_tokens,
|
|
),
|
|
patch("omlx.patches.specprefill.select_chunks", side_effect=select_chunks),
|
|
patch.object(draft_workflow.mx, "stream", side_effect=use_stream),
|
|
):
|
|
draft_workflow.run_specprefill_draft_scoring(
|
|
request=request,
|
|
plan=plan,
|
|
draft_model=object(),
|
|
draft_prefix_cache=draft_cache,
|
|
model_id="model-id",
|
|
prefill_step_size=4,
|
|
stream=stream,
|
|
extract_cache_states=extract_cache_states or (lambda cache: ([], None)),
|
|
sync_and_clear_cache=lambda: trace["syncs"].append(stream),
|
|
log=logger,
|
|
)
|
|
trace["selected_indices"] = selected_indices
|
|
trace["stream"] = stream
|
|
return tracker, logger, trace
|
|
|
|
|
|
def test_success_updates_request_tracker_logger_and_stream():
|
|
request, plan = _request_and_plan()
|
|
|
|
tracker, logger, trace = _run(request, plan)
|
|
|
|
assert request.specprefill_indices is trace["selected_indices"]
|
|
assert request.specprefill_total_tokens == plan.n_to_score
|
|
assert request.specprefill_position_offset == plan.effective_system
|
|
assert request._specprefill_system_tokens == plan.effective_system
|
|
assert [update["phase"] for update in tracker.updates] == [
|
|
"specprefill_scoring",
|
|
"specprefill_selected",
|
|
"prefill",
|
|
]
|
|
assert tracker.updates[-1]["processed"] == plan.n_to_score
|
|
assert tracker.removed == []
|
|
assert trace["streams"] == [trace["stream"]]
|
|
assert trace["syncs"] == [trace["stream"]]
|
|
assert logger.info_messages[0].startswith("SpecPrefill: scored")
|
|
|
|
|
|
def test_reconstructed_cache_is_scored_and_stored():
|
|
request, plan = _request_and_plan()
|
|
block_table = SimpleNamespace(num_tokens=3)
|
|
reconstructed_cache = ["reconstructed"]
|
|
draft_cache = _DraftCache(block_table, reconstructed_cache)
|
|
model_cache_config = object()
|
|
|
|
def extract_cache_states(cache: list[Any]) -> tuple[list[dict[str, Any]], Any]:
|
|
assert cache == ["draft-cache"]
|
|
return [{"state": "value"}], model_cache_config
|
|
|
|
_, _, trace = _run(
|
|
request,
|
|
plan,
|
|
draft_cache=draft_cache,
|
|
extract_cache_states=extract_cache_states,
|
|
)
|
|
|
|
assert trace["score_calls"][0]["existing_cache"] is reconstructed_cache
|
|
assert draft_cache.fetches == [(request.request_id, list(plan.tokens_to_score))]
|
|
assert draft_cache.preloads == [block_table]
|
|
assert draft_cache.reconstructions == [block_table]
|
|
assert draft_cache.stores == [
|
|
(
|
|
request.request_id,
|
|
list(plan.tokens_to_score),
|
|
[{"state": "value"}],
|
|
model_cache_config,
|
|
)
|
|
]
|
|
|
|
|
|
def test_cache_fetch_error_falls_back_to_uncached_scoring():
|
|
request, plan = _request_and_plan()
|
|
draft_cache = _DraftCache(fetch_error=RuntimeError("disk gone"))
|
|
|
|
_, logger, trace = _run(request, plan, draft_cache=draft_cache)
|
|
|
|
assert trace["score_calls"][0]["existing_cache"] is None
|
|
assert any("draft cache fetch failed: disk gone" in message for message in logger.debug_messages)
|
|
|
|
|
|
def test_scoring_error_clears_request_and_tracker():
|
|
request, plan = _request_and_plan()
|
|
|
|
def fail_scoring(*args: Any, **kwargs: Any) -> None:
|
|
raise RuntimeError("boom")
|
|
|
|
tracker, logger, _ = _run(request, plan, score_tokens=fail_scoring)
|
|
|
|
assert request.specprefill_indices is None
|
|
assert tracker.removed == [request.request_id]
|
|
assert logger.error_messages == [
|
|
"SpecPrefill scoring failed, falling back to normal path: boom"
|
|
]
|