wehub-resource-sync
451 lines
17 KiB
Python
451 lines
17 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Tests for admin cache probe endpoint (POST /admin/api/cache/probe)."""
|
|
|
|
import asyncio
|
|
from types import SimpleNamespace
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
from fastapi import HTTPException
|
|
|
|
import omlx.admin.routes as admin_routes
|
|
import omlx.server # noqa: F401 — triggers set_admin_getters
|
|
from omlx.cache.paged_cache import compute_block_hash
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
BLOCK_SIZE = 4 # Small block size for readable tests
|
|
MODEL_ID = "test-model"
|
|
MODEL_NAME = "/models/test-model"
|
|
|
|
|
|
def _make_request(model_id=MODEL_ID, messages=None):
|
|
"""Build a CacheProbeRequest."""
|
|
return admin_routes.CacheProbeRequest(
|
|
model_id=model_id,
|
|
messages=messages or [{"role": "user", "content": "hello"}],
|
|
)
|
|
|
|
|
|
def _make_tokenizer(token_ids):
|
|
"""Return a minimal tokenizer mock that produces *token_ids*."""
|
|
tok = MagicMock(spec=[])
|
|
tok.apply_chat_template = MagicMock(return_value="rendered prompt")
|
|
tok.encode = MagicMock(return_value=token_ids)
|
|
return tok
|
|
|
|
|
|
def _compute_hashes(token_ids, block_size=BLOCK_SIZE, model_name=MODEL_NAME):
|
|
"""Compute the chain-hashed block sequence for *token_ids*."""
|
|
hashes = []
|
|
parent = b""
|
|
for start in range(0, len(token_ids), block_size):
|
|
block_tokens = token_ids[start : start + block_size]
|
|
h = compute_block_hash(parent, block_tokens, model_name=model_name)
|
|
hashes.append(h)
|
|
parent = h
|
|
return hashes
|
|
|
|
|
|
def _make_ssd_index(known_hashes):
|
|
"""Return a mock SSD index whose contains() recognises *known_hashes*."""
|
|
idx = MagicMock(spec=[])
|
|
idx.contains = MagicMock(side_effect=lambda h: h in known_hashes)
|
|
return idx
|
|
|
|
|
|
def _make_engine_entry(
|
|
tokenizer,
|
|
scheduler,
|
|
has_apply_chat_template=True,
|
|
):
|
|
"""Build the engine_pool._entries[model_id] namespace chain."""
|
|
engine_ns = SimpleNamespace(
|
|
_tokenizer=tokenizer,
|
|
_engine=SimpleNamespace(
|
|
engine=SimpleNamespace(scheduler=scheduler),
|
|
),
|
|
)
|
|
if has_apply_chat_template:
|
|
engine_ns._apply_chat_template = lambda msgs, tools, **kw: "rendered prompt"
|
|
return SimpleNamespace(engine=engine_ns)
|
|
|
|
|
|
def _make_scheduler(
|
|
ssd_hot=None,
|
|
ssd_index=None,
|
|
model_name=MODEL_NAME,
|
|
block_size=BLOCK_SIZE,
|
|
):
|
|
"""Build a scheduler SimpleNamespace with the attributes probe_cache reads."""
|
|
return SimpleNamespace(
|
|
block_aware_cache=SimpleNamespace(block_size=block_size),
|
|
paged_ssd_cache_manager=SimpleNamespace(
|
|
_hot_cache=ssd_hot or {},
|
|
_index=ssd_index or _make_ssd_index(set()),
|
|
),
|
|
paged_cache_manager=SimpleNamespace(model_name=model_name),
|
|
config=SimpleNamespace(paged_cache_block_size=block_size),
|
|
)
|
|
|
|
|
|
def _pool_with(entries):
|
|
"""Return a mock engine pool wrapping *entries*."""
|
|
pool = MagicMock(spec=[])
|
|
pool._entries = entries
|
|
return pool
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Error / edge-case tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestCacheProbeToolCallNormalization:
|
|
"""Verify cache probing mirrors chat-path tool argument normalization."""
|
|
|
|
@pytest.mark.parametrize(
|
|
("arguments", "expected"),
|
|
[
|
|
('{"city":"Seoul"}', {"city": "Seoul"}),
|
|
(" ", {}),
|
|
],
|
|
)
|
|
def test_normalizes_arguments_before_rendering(self, arguments, expected):
|
|
messages = [
|
|
{"role": "user", "content": "Check the weather."},
|
|
{
|
|
"role": "assistant",
|
|
"content": "",
|
|
"tool_calls": [
|
|
{
|
|
"id": "call_1",
|
|
"type": "function",
|
|
"function": {
|
|
"name": "weather",
|
|
"arguments": arguments,
|
|
},
|
|
}
|
|
],
|
|
},
|
|
{"role": "tool", "tool_call_id": "call_1", "content": "sunny"},
|
|
]
|
|
request = _make_request(messages=messages)
|
|
tokenizer = _make_tokenizer([1, 2, 3, 4])
|
|
scheduler = _make_scheduler()
|
|
entry = _make_engine_entry(tokenizer, scheduler)
|
|
rendered_messages = None
|
|
|
|
def render(normalized, tools, **kwargs):
|
|
nonlocal rendered_messages
|
|
rendered_messages = normalized
|
|
return "rendered prompt"
|
|
|
|
entry.engine._apply_chat_template = render
|
|
pool = _pool_with({MODEL_ID: entry})
|
|
|
|
with patch.object(admin_routes, "_get_engine_pool", return_value=pool):
|
|
asyncio.run(admin_routes.probe_cache(request, is_admin=True))
|
|
|
|
normalized_arguments = rendered_messages[1]["tool_calls"][0]["function"][
|
|
"arguments"
|
|
]
|
|
assert normalized_arguments == expected
|
|
assert (
|
|
request.messages[1]["tool_calls"][0]["function"]["arguments"] == arguments
|
|
)
|
|
|
|
|
|
class TestCacheProbeErrors:
|
|
"""Guard-clause and error-path coverage."""
|
|
|
|
def test_engine_pool_not_initialized(self):
|
|
with patch.object(admin_routes, "_get_engine_pool", return_value=None):
|
|
with pytest.raises(HTTPException) as exc_info:
|
|
asyncio.run(admin_routes.probe_cache(_make_request(), is_admin=True))
|
|
assert exc_info.value.status_code == 503
|
|
|
|
def test_model_not_found(self):
|
|
pool = _pool_with({})
|
|
with patch.object(admin_routes, "_get_engine_pool", return_value=pool):
|
|
with pytest.raises(HTTPException) as exc_info:
|
|
asyncio.run(admin_routes.probe_cache(_make_request(), is_admin=True))
|
|
assert exc_info.value.status_code == 404
|
|
|
|
def test_model_not_loaded(self):
|
|
pool = _pool_with({MODEL_ID: SimpleNamespace(engine=None)})
|
|
with patch.object(admin_routes, "_get_engine_pool", return_value=pool):
|
|
result = asyncio.run(
|
|
admin_routes.probe_cache(_make_request(), is_admin=True)
|
|
)
|
|
assert result["model_loaded"] is False
|
|
assert "reason" in result
|
|
|
|
def test_no_tokenizer(self):
|
|
entry = SimpleNamespace(
|
|
engine=SimpleNamespace(
|
|
_tokenizer=None,
|
|
_engine=SimpleNamespace(
|
|
engine=SimpleNamespace(scheduler=_make_scheduler()),
|
|
),
|
|
)
|
|
)
|
|
pool = _pool_with({MODEL_ID: entry})
|
|
with patch.object(admin_routes, "_get_engine_pool", return_value=pool):
|
|
with pytest.raises(HTTPException) as exc_info:
|
|
asyncio.run(admin_routes.probe_cache(_make_request(), is_admin=True))
|
|
assert exc_info.value.status_code == 400
|
|
|
|
def test_scheduler_unavailable(self):
|
|
tokenizer = _make_tokenizer([1, 2, 3])
|
|
entry = SimpleNamespace(
|
|
engine=SimpleNamespace(
|
|
_tokenizer=tokenizer,
|
|
_engine=None,
|
|
)
|
|
)
|
|
pool = _pool_with({MODEL_ID: entry})
|
|
with patch.object(admin_routes, "_get_engine_pool", return_value=pool):
|
|
with pytest.raises(HTTPException) as exc_info:
|
|
asyncio.run(admin_routes.probe_cache(_make_request(), is_admin=True))
|
|
assert exc_info.value.status_code == 500
|
|
|
|
def test_block_size_unavailable(self):
|
|
tokenizer = _make_tokenizer([1, 2, 3])
|
|
scheduler = SimpleNamespace(
|
|
block_aware_cache=None,
|
|
paged_ssd_cache_manager=None,
|
|
paged_cache_manager=None,
|
|
config=SimpleNamespace(paged_cache_block_size=0),
|
|
)
|
|
entry = _make_engine_entry(tokenizer, scheduler)
|
|
pool = _pool_with({MODEL_ID: entry})
|
|
with patch.object(admin_routes, "_get_engine_pool", return_value=pool):
|
|
with pytest.raises(HTTPException) as exc_info:
|
|
asyncio.run(admin_routes.probe_cache(_make_request(), is_admin=True))
|
|
assert exc_info.value.status_code == 500
|
|
|
|
def test_empty_tokens(self):
|
|
tokenizer = _make_tokenizer([])
|
|
scheduler = _make_scheduler()
|
|
entry = _make_engine_entry(tokenizer, scheduler)
|
|
pool = _pool_with({MODEL_ID: entry})
|
|
with patch.object(admin_routes, "_get_engine_pool", return_value=pool):
|
|
result = asyncio.run(
|
|
admin_routes.probe_cache(_make_request(), is_admin=True)
|
|
)
|
|
assert result["model_loaded"] is True
|
|
assert result["total_tokens"] == 0
|
|
assert result["total_blocks"] == 0
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Cache classification tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestCacheProbeClassification:
|
|
"""Verify hot / disk / cold block classification."""
|
|
|
|
def test_all_blocks_cold(self):
|
|
"""No SSD cache entries at all → everything cold."""
|
|
token_ids = list(range(12)) # 3 blocks of 4
|
|
tokenizer = _make_tokenizer(token_ids)
|
|
scheduler = _make_scheduler()
|
|
entry = _make_engine_entry(tokenizer, scheduler)
|
|
pool = _pool_with({MODEL_ID: entry})
|
|
|
|
with patch.object(admin_routes, "_get_engine_pool", return_value=pool):
|
|
result = asyncio.run(
|
|
admin_routes.probe_cache(_make_request(), is_admin=True)
|
|
)
|
|
|
|
assert result["total_tokens"] == 12
|
|
assert result["total_blocks"] == 3
|
|
assert result["blocks_ssd_hot"] == 0
|
|
assert result["blocks_ssd_disk"] == 0
|
|
assert result["blocks_cold"] == 3
|
|
assert result["ssd_hit_tokens"] == 0
|
|
assert result["cold_tokens"] == 12
|
|
|
|
def test_all_blocks_ssd_disk(self):
|
|
"""Every block found in SSD disk index."""
|
|
token_ids = list(range(8)) # 2 blocks of 4
|
|
hashes = _compute_hashes(token_ids)
|
|
ssd_index = _make_ssd_index(set(hashes))
|
|
tokenizer = _make_tokenizer(token_ids)
|
|
scheduler = _make_scheduler(ssd_index=ssd_index)
|
|
entry = _make_engine_entry(tokenizer, scheduler)
|
|
pool = _pool_with({MODEL_ID: entry})
|
|
|
|
with patch.object(admin_routes, "_get_engine_pool", return_value=pool):
|
|
result = asyncio.run(
|
|
admin_routes.probe_cache(_make_request(), is_admin=True)
|
|
)
|
|
|
|
assert result["blocks_ssd_hot"] == 0
|
|
assert result["blocks_ssd_disk"] == 2
|
|
assert result["blocks_cold"] == 0
|
|
assert result["ssd_hit_tokens"] == 8
|
|
assert result["cold_tokens"] == 0
|
|
|
|
def test_all_blocks_ssd_hot(self):
|
|
"""Every block found in SSD hot cache (RAM copy)."""
|
|
token_ids = list(range(8)) # 2 blocks of 4
|
|
hashes = _compute_hashes(token_ids)
|
|
ssd_hot = {h: {} for h in hashes}
|
|
tokenizer = _make_tokenizer(token_ids)
|
|
scheduler = _make_scheduler(ssd_hot=ssd_hot)
|
|
entry = _make_engine_entry(tokenizer, scheduler)
|
|
pool = _pool_with({MODEL_ID: entry})
|
|
|
|
with patch.object(admin_routes, "_get_engine_pool", return_value=pool):
|
|
result = asyncio.run(
|
|
admin_routes.probe_cache(_make_request(), is_admin=True)
|
|
)
|
|
|
|
assert result["blocks_ssd_hot"] == 2
|
|
assert result["blocks_ssd_disk"] == 0
|
|
assert result["blocks_cold"] == 0
|
|
|
|
def test_mixed_hot_and_disk(self):
|
|
"""First block hot, second block disk-only."""
|
|
token_ids = list(range(8)) # 2 blocks of 4
|
|
hashes = _compute_hashes(token_ids)
|
|
ssd_hot = {hashes[0]: {}}
|
|
ssd_index = _make_ssd_index({hashes[1]})
|
|
tokenizer = _make_tokenizer(token_ids)
|
|
scheduler = _make_scheduler(ssd_hot=ssd_hot, ssd_index=ssd_index)
|
|
entry = _make_engine_entry(tokenizer, scheduler)
|
|
pool = _pool_with({MODEL_ID: entry})
|
|
|
|
with patch.object(admin_routes, "_get_engine_pool", return_value=pool):
|
|
result = asyncio.run(
|
|
admin_routes.probe_cache(_make_request(), is_admin=True)
|
|
)
|
|
|
|
assert result["blocks_ssd_hot"] == 1
|
|
assert result["blocks_ssd_disk"] == 1
|
|
assert result["blocks_cold"] == 0
|
|
|
|
def test_partial_prefix_hit(self):
|
|
"""First 2 blocks cached, third block miss → third is cold."""
|
|
token_ids = list(range(12)) # 3 blocks of 4
|
|
hashes = _compute_hashes(token_ids)
|
|
ssd_index = _make_ssd_index({hashes[0], hashes[1]})
|
|
tokenizer = _make_tokenizer(token_ids)
|
|
scheduler = _make_scheduler(ssd_index=ssd_index)
|
|
entry = _make_engine_entry(tokenizer, scheduler)
|
|
pool = _pool_with({MODEL_ID: entry})
|
|
|
|
with patch.object(admin_routes, "_get_engine_pool", return_value=pool):
|
|
result = asyncio.run(
|
|
admin_routes.probe_cache(_make_request(), is_admin=True)
|
|
)
|
|
|
|
assert result["blocks_ssd_disk"] == 2
|
|
assert result["blocks_cold"] == 1
|
|
assert result["ssd_hit_tokens"] == 8
|
|
assert result["cold_tokens"] == 4
|
|
|
|
def test_gap_in_prefix_makes_rest_cold(self):
|
|
"""Block 0 cached, block 1 missing, block 2 cached → blocks 1+2 cold.
|
|
|
|
The walk stops at the first miss because of the contiguous prefix
|
|
assumption, so block 2 is never checked even though it exists in the
|
|
index.
|
|
"""
|
|
token_ids = list(range(12)) # 3 blocks of 4
|
|
hashes = _compute_hashes(token_ids)
|
|
# Only block 0 and block 2 in index (gap at block 1)
|
|
ssd_index = _make_ssd_index({hashes[0], hashes[2]})
|
|
tokenizer = _make_tokenizer(token_ids)
|
|
scheduler = _make_scheduler(ssd_index=ssd_index)
|
|
entry = _make_engine_entry(tokenizer, scheduler)
|
|
pool = _pool_with({MODEL_ID: entry})
|
|
|
|
with patch.object(admin_routes, "_get_engine_pool", return_value=pool):
|
|
result = asyncio.run(
|
|
admin_routes.probe_cache(_make_request(), is_admin=True)
|
|
)
|
|
|
|
assert result["blocks_ssd_disk"] == 1 # Only block 0
|
|
assert result["blocks_cold"] == 2 # Blocks 1 and 2
|
|
|
|
def test_partial_last_block(self):
|
|
"""Token count not a multiple of block_size → last block is smaller."""
|
|
token_ids = list(range(10)) # 2 full blocks + 1 partial (2 tokens)
|
|
hashes = _compute_hashes(token_ids)
|
|
ssd_index = _make_ssd_index(set(hashes))
|
|
tokenizer = _make_tokenizer(token_ids)
|
|
scheduler = _make_scheduler(ssd_index=ssd_index)
|
|
entry = _make_engine_entry(tokenizer, scheduler)
|
|
pool = _pool_with({MODEL_ID: entry})
|
|
|
|
with patch.object(admin_routes, "_get_engine_pool", return_value=pool):
|
|
result = asyncio.run(
|
|
admin_routes.probe_cache(_make_request(), is_admin=True)
|
|
)
|
|
|
|
assert result["total_tokens"] == 10
|
|
assert result["total_blocks"] == 3
|
|
assert result["blocks_ssd_disk"] == 3
|
|
assert result["blocks_cold"] == 0
|
|
assert result["ssd_hit_tokens"] == 10
|
|
assert result["cold_tokens"] == 0
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Response shape test
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestCacheProbeResponseShape:
|
|
"""Verify the response contains exactly the expected fields."""
|
|
|
|
EXPECTED_FIELDS = {
|
|
"model_id",
|
|
"model_loaded",
|
|
"total_tokens",
|
|
"block_size",
|
|
"total_blocks",
|
|
"blocks_ssd_hot",
|
|
"blocks_ssd_disk",
|
|
"blocks_cold",
|
|
"ssd_hit_tokens",
|
|
"cold_tokens",
|
|
}
|
|
|
|
def test_loaded_response_fields(self):
|
|
token_ids = list(range(4))
|
|
tokenizer = _make_tokenizer(token_ids)
|
|
scheduler = _make_scheduler()
|
|
entry = _make_engine_entry(tokenizer, scheduler)
|
|
pool = _pool_with({MODEL_ID: entry})
|
|
|
|
with patch.object(admin_routes, "_get_engine_pool", return_value=pool):
|
|
result = asyncio.run(
|
|
admin_routes.probe_cache(_make_request(), is_admin=True)
|
|
)
|
|
|
|
assert set(result.keys()) == self.EXPECTED_FIELDS
|
|
|
|
def test_no_dead_fields(self):
|
|
"""Ensure blocks_ram, ram_hit_tokens, prefix_index_hits are gone."""
|
|
token_ids = list(range(4))
|
|
tokenizer = _make_tokenizer(token_ids)
|
|
scheduler = _make_scheduler()
|
|
entry = _make_engine_entry(tokenizer, scheduler)
|
|
pool = _pool_with({MODEL_ID: entry})
|
|
|
|
with patch.object(admin_routes, "_get_engine_pool", return_value=pool):
|
|
result = asyncio.run(
|
|
admin_routes.probe_cache(_make_request(), is_admin=True)
|
|
)
|
|
|
|
assert "blocks_ram" not in result
|
|
assert "ram_hit_tokens" not in result
|
|
assert "prefix_index_hits" not in result
|