wehub-resource-sync
120 lines
3.9 KiB
Python
120 lines
3.9 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""End-to-end test: the MP cache server emits usage telemetry at startup.
|
|
|
|
Spawns a real ``run_cache_server`` in a child process with the stats-server
|
|
URL pointed at a local HTTP sink, and asserts the startup report arrives
|
|
over the wire.
|
|
"""
|
|
|
|
# Standard
|
|
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
|
import json
|
|
import multiprocessing
|
|
import threading
|
|
import time
|
|
|
|
# Third Party
|
|
import pytest
|
|
import torch
|
|
|
|
# First Party
|
|
from lmcache.v1.distributed.config import (
|
|
EvictionConfig,
|
|
L1ManagerConfig,
|
|
L1MemoryManagerConfig,
|
|
StorageManagerConfig,
|
|
)
|
|
from lmcache.v1.mp_observability.config import DEFAULT_OBSERVABILITY_CONFIG
|
|
from lmcache.v1.multiprocess.config import MPServerConfig
|
|
from lmcache.v1.multiprocess.server import run_cache_server
|
|
|
|
if not torch.cuda.is_available():
|
|
pytest.skip(
|
|
"MP cache server requires an accelerator device",
|
|
allow_module_level=True,
|
|
)
|
|
|
|
SERVER_PORT = 15599
|
|
STARTUP_TIMEOUT_SECONDS = 90.0
|
|
|
|
|
|
class _SinkHandler(BaseHTTPRequestHandler):
|
|
def do_POST(self) -> None:
|
|
length = int(self.headers.get("Content-Length", 0))
|
|
body = json.loads(self.rfile.read(length))
|
|
self.server.received.append((self.path, body)) # type: ignore[attr-defined]
|
|
self.send_response(200)
|
|
self.end_headers()
|
|
|
|
def log_message(self, *args: object) -> None:
|
|
pass
|
|
|
|
|
|
class _UsageSink(ThreadingHTTPServer):
|
|
def __init__(self) -> None:
|
|
super().__init__(("127.0.0.1", 0), _SinkHandler)
|
|
self.received: list[tuple[str, dict[str, object]]] = []
|
|
|
|
|
|
def _usage_server_runner(port: int) -> None:
|
|
"""Child-process entry point running a minimal cache server."""
|
|
mp_config = MPServerConfig(host="localhost", port=port)
|
|
storage_manager_config = StorageManagerConfig(
|
|
l1_manager_config=L1ManagerConfig(
|
|
memory_config=L1MemoryManagerConfig(
|
|
size_in_bytes=1 << 30,
|
|
use_lazy=True,
|
|
),
|
|
),
|
|
eviction_config=EvictionConfig(eviction_policy="LRU"),
|
|
)
|
|
run_cache_server(
|
|
mp_config=mp_config,
|
|
storage_manager_config=storage_manager_config,
|
|
obs_config=DEFAULT_OBSERVABILITY_CONFIG,
|
|
)
|
|
|
|
|
|
def test_server_startup_emits_usage_report(monkeypatch, tmp_path):
|
|
sink = _UsageSink()
|
|
sink_thread = threading.Thread(target=sink.serve_forever, daemon=True)
|
|
sink_thread.start()
|
|
|
|
# The child inherits this environment at spawn: tracking enabled,
|
|
# stats URL pointed at the sink, HOME isolated so no real state files
|
|
# (machine_id) are touched.
|
|
monkeypatch.setenv("LMCACHE_TRACK_USAGE", "true")
|
|
monkeypatch.setenv(
|
|
"LMCACHE_USAGE_TRACK_URL", f"http://127.0.0.1:{sink.server_address[1]}"
|
|
)
|
|
monkeypatch.setenv("HOME", str(tmp_path))
|
|
|
|
context = multiprocessing.get_context("spawn")
|
|
process = context.Process(target=_usage_server_runner, args=(SERVER_PORT,))
|
|
process.start()
|
|
try:
|
|
deadline = time.monotonic() + STARTUP_TIMEOUT_SECONDS
|
|
while time.monotonic() < deadline and len(sink.received) < 2:
|
|
time.sleep(0.1)
|
|
received = list(sink.received)
|
|
assert len(received) == 2, f"expected 2 usage messages, got {received}"
|
|
|
|
assert {path for path, _ in received} == {"/context"}
|
|
payloads = {payload["message_type"]: payload for _, payload in received}
|
|
assert set(payloads) == {"EnvMessage", "MPServerMessage"}
|
|
|
|
mp_payload = payloads["MPServerMessage"]
|
|
assert mp_payload["deployment_mode"] == "mp_server"
|
|
assert mp_payload["chunk_size"] == 256
|
|
assert mp_payload["l1_size_bytes"] == 1 << 30
|
|
assert "instance_id" not in mp_payload
|
|
assert payloads["EnvMessage"]["session_id"] == mp_payload["session_id"]
|
|
finally:
|
|
process.terminate()
|
|
process.join(timeout=10)
|
|
if process.is_alive():
|
|
process.kill()
|
|
process.join()
|
|
sink.shutdown()
|
|
sink.server_close()
|