wehub-resource-sync
380 lines
13 KiB
Python
380 lines
13 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Tests for the POST /cache/checksums endpoint (and other MP HTTP routes)."""
|
|
|
|
# Standard
|
|
from unittest.mock import MagicMock, PropertyMock
|
|
|
|
# Third Party
|
|
from fastapi.testclient import TestClient
|
|
import pytest
|
|
import torch
|
|
|
|
# First Party
|
|
from lmcache.v1.multiprocess.http_apis.dependencies import build_context
|
|
from lmcache.v1.multiprocess.http_server import app
|
|
import lmcache.c_ops as lmc_ops
|
|
|
|
|
|
def _make_kv_tensors(
|
|
num_layers: int = 2,
|
|
num_blocks: int = 4,
|
|
block_size: int = 4,
|
|
num_heads: int = 2,
|
|
head_size: int = 8,
|
|
dtype: torch.dtype = torch.float32,
|
|
) -> list[torch.Tensor]:
|
|
"""Create deterministic CPU KV tensors for testing."""
|
|
torch.manual_seed(42)
|
|
return [
|
|
torch.randn(
|
|
2,
|
|
num_blocks,
|
|
block_size,
|
|
num_heads,
|
|
head_size,
|
|
dtype=dtype,
|
|
)
|
|
for _ in range(num_layers)
|
|
]
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_gpu_ctx():
|
|
"""Create a mock GPUCacheContext with kv_tensors."""
|
|
ctx = MagicMock()
|
|
tensors = _make_kv_tensors()
|
|
type(ctx).kv_tensors = PropertyMock(return_value=tensors)
|
|
type(ctx).block_size = PropertyMock(return_value=4)
|
|
# KV tensors are built as [2, NB, BS, NH, HS] -> NL_X_TWO_NB_BS_NH_HS;
|
|
# one homogeneous kernel group, so every layer reports the same format.
|
|
fmt = lmc_ops.EngineKVFormat.NL_X_TWO_NB_BS_NH_HS
|
|
ctx.engine_kv_formats.return_value = [fmt]
|
|
ctx.engine_kv_format_per_layer.return_value = [fmt] * len(tensors)
|
|
return ctx
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_mixed_engine():
|
|
"""Engine whose context mixes two KV formats (a key+value layer and a
|
|
key-only layer), so the endpoint gathers each along its own axis."""
|
|
ctx = MagicMock()
|
|
kv_kv = torch.randn(2, 4, 4, 2, 8)
|
|
kv_idx = torch.randn(4, 4, 8)
|
|
type(ctx).kv_tensors = PropertyMock(return_value=[kv_kv, kv_idx])
|
|
ctx.engine_kv_format_per_layer.return_value = [
|
|
lmc_ops.EngineKVFormat.NL_X_TWO_NB_BS_NH_HS,
|
|
lmc_ops.EngineKVFormat.NL_X_NB_BS_HS,
|
|
]
|
|
engine = MagicMock()
|
|
engine.cache_contexts = {0: ctx}
|
|
return engine
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_engine(mock_gpu_ctx):
|
|
"""Create a mock engine with cache_contexts."""
|
|
engine = MagicMock()
|
|
engine.cache_contexts = {0: mock_gpu_ctx}
|
|
return engine
|
|
|
|
|
|
@pytest.fixture
|
|
def client_with_engine(mock_engine):
|
|
"""Create a test client with a mocked engine wired onto the per-app context.
|
|
|
|
The cache handlers resolve the engine through ``get_context`` (i.e.
|
|
``app.state.context``), so the context must be built, not just
|
|
``app.state.engine``.
|
|
"""
|
|
app.state.engine = mock_engine
|
|
app.state.context = build_context(mock_engine)
|
|
client = TestClient(app)
|
|
yield client
|
|
client.close()
|
|
app.state.engine = None
|
|
app.state.context = None
|
|
|
|
|
|
@pytest.fixture
|
|
def client_no_engine():
|
|
"""Create a test client without an engine/context."""
|
|
app.state.engine = None
|
|
app.state.context = None
|
|
client = TestClient(app)
|
|
yield client
|
|
client.close()
|
|
|
|
|
|
class TestCacheChecksumsEndpoint:
|
|
"""Tests for POST /cache/checksums."""
|
|
|
|
def test_success_basic(self, client_with_engine):
|
|
"""Basic successful checksum request."""
|
|
# chunk_size is in blocks; 1 block + chunk_size=1 -> 1 chunk.
|
|
resp = client_with_engine.post(
|
|
"/cache/checksums", json={"block_ids": [0], "chunk_size": 1}
|
|
)
|
|
assert resp.status_code == 200
|
|
data = resp.json()
|
|
assert data["status"] == "success"
|
|
assert data["chunk_size"] == 1
|
|
assert data["num_chunks"] == 1
|
|
assert len(data["chunk_checksums"]) == 1
|
|
assert "block_id_ranges" in data
|
|
|
|
def test_success_layerwise(self, client_with_engine):
|
|
"""Layerwise mode returns per-layer checksums."""
|
|
resp = client_with_engine.post(
|
|
"/cache/checksums",
|
|
json={"block_ids": [0], "chunk_size": 1, "layerwise": True},
|
|
)
|
|
assert resp.status_code == 200
|
|
data = resp.json()
|
|
assert data["layerwise"] is True
|
|
cks = data["chunk_checksums"]
|
|
assert "layer_0" in cks
|
|
assert "layer_1" in cks
|
|
|
|
def test_deterministic(self, client_with_engine):
|
|
"""Same request produces identical checksums."""
|
|
body = {"block_ids": [0], "chunk_size": 1}
|
|
d1 = client_with_engine.post("/cache/checksums", json=body).json()
|
|
d2 = client_with_engine.post("/cache/checksums", json=body).json()
|
|
assert d1["chunk_checksums"] == d2["chunk_checksums"]
|
|
|
|
def test_mixed_format_supported(self, mock_mixed_engine):
|
|
"""Two different KV formats are gathered per layer, not rejected."""
|
|
app.state.engine = mock_mixed_engine
|
|
app.state.context = build_context(mock_mixed_engine)
|
|
client = TestClient(app)
|
|
try:
|
|
resp = client.post(
|
|
"/cache/checksums",
|
|
json={"block_ids": [0], "chunk_size": 1, "layerwise": True},
|
|
)
|
|
finally:
|
|
client.close()
|
|
app.state.engine = None
|
|
app.state.context = None
|
|
assert resp.status_code == 200
|
|
data = resp.json()
|
|
assert data["status"] == "success"
|
|
# Both layers checksummed despite different formats / block axes.
|
|
assert set(data["chunk_checksums"]) == {"layer_0", "layer_1"}
|
|
|
|
def test_range_block_ids(self, client_with_engine):
|
|
"""Two block ids -> two single-block chunks."""
|
|
resp = client_with_engine.post(
|
|
"/cache/checksums", json={"block_ids": [0, 1], "chunk_size": 1}
|
|
)
|
|
assert resp.status_code == 200
|
|
data = resp.json()
|
|
# 2 blocks, chunk_size=1 block -> 2 chunks
|
|
assert data["num_chunks"] == 2
|
|
|
|
def test_multiple_block_ids(self, client_with_engine):
|
|
"""Three block ids -> three single-block chunks."""
|
|
resp = client_with_engine.post(
|
|
"/cache/checksums", json={"block_ids": [0, 1, 2], "chunk_size": 1}
|
|
)
|
|
assert resp.status_code == 200
|
|
data = resp.json()
|
|
# 3 blocks, chunk_size=1 block -> 3 chunks
|
|
assert data["num_chunks"] == 3
|
|
|
|
# ------------------------------------------------------------------
|
|
# Error cases
|
|
# ------------------------------------------------------------------
|
|
|
|
def test_no_engine(self, client_no_engine):
|
|
"""503 when the engine/context is not initialized."""
|
|
resp = client_no_engine.post(
|
|
"/cache/checksums", json={"block_ids": [0], "chunk_size": 1}
|
|
)
|
|
assert resp.status_code == 503
|
|
|
|
def test_no_cache_contexts(self, client_with_engine, mock_engine):
|
|
"""501 when engine has no cache_contexts attribute."""
|
|
mock_engine.cache_contexts = None
|
|
resp = client_with_engine.post(
|
|
"/cache/checksums", json={"block_ids": [0], "chunk_size": 1}
|
|
)
|
|
assert resp.status_code == 501
|
|
|
|
def test_unknown_instance_id(self, client_with_engine):
|
|
"""404 when instance_id is not registered."""
|
|
resp = client_with_engine.post(
|
|
"/cache/checksums",
|
|
json={"block_ids": [0], "chunk_size": 1, "instance_id": 99},
|
|
)
|
|
assert resp.status_code == 404
|
|
|
|
def test_missing_block_ids(self, client_with_engine):
|
|
"""400 when block_ids is empty."""
|
|
resp = client_with_engine.post("/cache/checksums", json={"chunk_size": 1})
|
|
assert resp.status_code == 400
|
|
assert "block_ids" in resp.json()["detail"]
|
|
|
|
def test_missing_chunk_size(self, client_with_engine):
|
|
"""400 when chunk_size is missing (defaults to 0)."""
|
|
resp = client_with_engine.post("/cache/checksums", json={"block_ids": [0]})
|
|
assert resp.status_code == 400
|
|
assert "chunk_size" in resp.json()["detail"]
|
|
|
|
def test_zero_chunk_size(self, client_with_engine):
|
|
"""400 when chunk_size is zero."""
|
|
resp = client_with_engine.post(
|
|
"/cache/checksums", json={"block_ids": [0], "chunk_size": 0}
|
|
)
|
|
assert resp.status_code == 400
|
|
|
|
def test_negative_chunk_size(self, client_with_engine):
|
|
"""400 when chunk_size is negative."""
|
|
resp = client_with_engine.post(
|
|
"/cache/checksums", json={"block_ids": [0], "chunk_size": -1}
|
|
)
|
|
assert resp.status_code == 400
|
|
|
|
def test_empty_kv_caches(self, client_with_engine, mock_gpu_ctx):
|
|
"""404 when kv_tensors is empty."""
|
|
type(mock_gpu_ctx).kv_tensors = PropertyMock(
|
|
return_value=[],
|
|
)
|
|
resp = client_with_engine.post(
|
|
"/cache/checksums", json={"block_ids": [0], "chunk_size": 1}
|
|
)
|
|
assert resp.status_code == 404
|
|
|
|
# ------------------------------------------------------------------
|
|
# Chunk boundary edge cases
|
|
# ------------------------------------------------------------------
|
|
|
|
def test_partial_last_chunk(self, client_with_engine):
|
|
"""3 blocks with chunk_size=2 blocks -> 2 chunks (2+1)."""
|
|
resp = client_with_engine.post(
|
|
"/cache/checksums", json={"block_ids": [0, 1, 2], "chunk_size": 2}
|
|
)
|
|
data = resp.json()
|
|
assert data["num_chunks"] == 2
|
|
assert len(data["chunk_checksums"]) == 2
|
|
|
|
def test_single_block_single_chunk(self, client_with_engine):
|
|
"""Single block with chunk_size=1 block produces one chunk."""
|
|
resp = client_with_engine.post(
|
|
"/cache/checksums", json={"block_ids": [0], "chunk_size": 1}
|
|
)
|
|
data = resp.json()
|
|
assert data["num_chunks"] == 1
|
|
|
|
# ------------------------------------------------------------------
|
|
# Checksum validity
|
|
# ------------------------------------------------------------------
|
|
|
|
def test_checksums_are_valid_md5(self, client_with_engine):
|
|
"""All checksums are 32-char hex strings."""
|
|
# 2 blocks, chunk_size=1 block -> 2 md5 digests.
|
|
resp = client_with_engine.post(
|
|
"/cache/checksums", json={"block_ids": [0, 1], "chunk_size": 1}
|
|
)
|
|
data = resp.json()
|
|
for cksum in data["chunk_checksums"]:
|
|
assert len(cksum) == 32
|
|
int(cksum, 16) # must be valid hex
|
|
|
|
def test_block_id_ranges_in_response(
|
|
self,
|
|
client_with_engine,
|
|
):
|
|
"""Response includes compressed block_id_ranges."""
|
|
resp = client_with_engine.post(
|
|
"/cache/checksums", json={"block_ids": [0], "chunk_size": 1}
|
|
)
|
|
data = resp.json()
|
|
assert "block_id_ranges" in data
|
|
# Single block id is kept as-is (no range compression).
|
|
assert data["block_id_ranges"] == [0]
|
|
|
|
# ------------------------------------------------------------------
|
|
# block_ids addressing
|
|
# ------------------------------------------------------------------
|
|
|
|
def test_block_ids_basic(self, client_with_engine):
|
|
"""block_ids with a block-level chunk_size."""
|
|
# 2 blocks, chunk_size=1 block -> 2 chunks.
|
|
resp = client_with_engine.post(
|
|
"/cache/checksums", json={"block_ids": [0, 1], "chunk_size": 1}
|
|
)
|
|
assert resp.status_code == 200
|
|
data = resp.json()
|
|
assert data["status"] == "success"
|
|
assert data["num_chunks"] == 2
|
|
assert data["block_id_ranges"] == [0, 1]
|
|
|
|
def test_block_ids_invalid_type(self, client_with_engine):
|
|
"""422 when block_ids is not a list of ints (FastAPI validation)."""
|
|
resp = client_with_engine.post(
|
|
"/cache/checksums", json={"block_ids": "abc", "chunk_size": 1}
|
|
)
|
|
assert resp.status_code == 422
|
|
|
|
|
|
class TestHealthAndMiscEndpoints:
|
|
"""Smoke tests for other endpoints on the MP HTTP server."""
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _reset_state(self):
|
|
yield
|
|
app.state.engine = None
|
|
app.state.context = None
|
|
|
|
def test_root(self, client_no_engine):
|
|
"""GET / returns ok."""
|
|
resp = client_no_engine.get("/")
|
|
assert resp.status_code == 200
|
|
assert resp.json()["status"] == "ok"
|
|
|
|
def test_healthcheck_no_engine(self, client_no_engine):
|
|
"""503 when engine is not set."""
|
|
resp = client_no_engine.get("/healthcheck")
|
|
assert resp.status_code == 503
|
|
|
|
def test_healthcheck_with_engine(
|
|
self,
|
|
client_with_engine,
|
|
):
|
|
"""200 when engine is available."""
|
|
resp = client_with_engine.get("/healthcheck")
|
|
assert resp.status_code == 200
|
|
assert resp.json()["status"] == "healthy"
|
|
|
|
def test_clear_cache_no_engine(self, client_no_engine):
|
|
"""503 when engine/context is not set."""
|
|
resp = client_no_engine.post("/cache/clear", json={})
|
|
assert resp.status_code == 503
|
|
|
|
def test_clear_cache_success(
|
|
self,
|
|
client_with_engine,
|
|
mock_engine,
|
|
):
|
|
"""200 and engine.clear() called."""
|
|
resp = client_with_engine.post("/cache/clear", json={"tier": "l1"})
|
|
assert resp.status_code == 200
|
|
mock_engine.clear.assert_called_once()
|
|
|
|
def test_status_no_engine(self, client_no_engine):
|
|
"""503 when engine is not set."""
|
|
resp = client_no_engine.get("/status")
|
|
assert resp.status_code == 503
|
|
|
|
def test_status_success(
|
|
self,
|
|
client_with_engine,
|
|
mock_engine,
|
|
):
|
|
"""200 and engine.report_status() called."""
|
|
mock_engine.report_status.return_value = {"ok": True}
|
|
resp = client_with_engine.get("/status")
|
|
assert resp.status_code == 200
|
|
assert resp.json()["ok"] is True
|