wehub-resource-sync
467 lines
17 KiB
Python
467 lines
17 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# Standard
|
|
from unittest.mock import MagicMock
|
|
|
|
# Third Party
|
|
from prometheus_client import REGISTRY
|
|
import pytest
|
|
|
|
# First Party
|
|
from lmcache.observability import (
|
|
LMCStatsMonitor,
|
|
PrometheusLogger,
|
|
)
|
|
|
|
|
|
@pytest.fixture(scope="function")
|
|
def stats_monitor():
|
|
LMCStatsMonitor.DestroyInstance()
|
|
return LMCStatsMonitor.GetOrCreate()
|
|
|
|
|
|
def test_on_retrieve_request(stats_monitor):
|
|
stats_monitor.on_retrieve_request(num_tokens=100)
|
|
stats = stats_monitor.get_stats_and_clear()
|
|
assert stats.interval_retrieve_requests == 1
|
|
assert stats.retrieve_hit_rate == 1.0
|
|
assert stats.local_cache_usage_bytes == 0
|
|
assert stats.remote_cache_usage_bytes == 0
|
|
assert len(stats.time_to_retrieve) == 0
|
|
|
|
|
|
def test_on_retrieve_finished(stats_monitor):
|
|
stats_obj = stats_monitor.on_retrieve_request(num_tokens=100)
|
|
stats_monitor.on_retrieve_finished(
|
|
retrieve_stats=stats_obj,
|
|
num_retrieved_tokens=100,
|
|
)
|
|
stats = stats_monitor.get_stats_and_clear()
|
|
assert stats.interval_retrieve_requests == 1
|
|
assert stats.retrieve_hit_rate == 1.0
|
|
assert len(stats.time_to_retrieve) == 1
|
|
|
|
|
|
def test_on_store_request_and_finished(stats_monitor):
|
|
stats_obj = stats_monitor.on_store_request(num_tokens=50)
|
|
stats_monitor.on_store_finished(store_stats=stats_obj)
|
|
stats = stats_monitor.get_stats_and_clear()
|
|
assert stats.interval_store_requests == 1
|
|
assert len(stats.time_to_store) == 1
|
|
|
|
|
|
def test_update_local_cache_usage(stats_monitor):
|
|
stats_monitor.update_local_cache_usage(usage=1024)
|
|
stats = stats_monitor.get_stats_and_clear()
|
|
assert stats.local_cache_usage_bytes == 1024
|
|
|
|
|
|
def test_update_remote_cache_usage(stats_monitor):
|
|
stats_monitor.update_remote_cache_usage(usage=2048)
|
|
stats = stats_monitor.get_stats_and_clear()
|
|
assert stats.remote_cache_usage_bytes == 2048
|
|
|
|
|
|
def test_update_local_storage_usage(stats_monitor):
|
|
stats_monitor.update_local_storage_usage(usage=4096)
|
|
stats = stats_monitor.get_stats_and_clear()
|
|
assert stats.local_storage_usage_bytes == 4096
|
|
|
|
|
|
def test_on_lookup_request(stats_monitor):
|
|
stats_monitor.on_lookup_request(num_tokens=50)
|
|
assert len(stats_monitor.lookup_requests) == 1
|
|
stats = stats_monitor.get_stats_and_clear()
|
|
assert stats.interval_lookup_requests == 1
|
|
assert stats.interval_lookup_tokens == 50
|
|
assert stats.lookup_hit_rate == 0
|
|
assert len(stats.interval_lookup_hit_rates) == 0
|
|
# on_lookup_finished is not called, lookup_requests is not clear
|
|
assert len(stats_monitor.lookup_requests) == 1
|
|
|
|
|
|
def test_on_lookup_finished(stats_monitor):
|
|
stats_obj = stats_monitor.on_lookup_request(num_tokens=100)
|
|
assert len(stats_monitor.lookup_requests) == 1
|
|
stats_monitor.on_lookup_finished(stats=stats_obj, num_hit_tokens=80)
|
|
stats = stats_monitor.get_stats_and_clear()
|
|
assert stats.interval_lookup_requests == 1
|
|
assert stats.interval_lookup_tokens == 100
|
|
assert stats.interval_lookup_hits == 80
|
|
assert stats.lookup_hit_rate == 0.8
|
|
assert len(stats.interval_lookup_hit_rates) == 1
|
|
assert stats.interval_lookup_hit_rates[0] == 0.8
|
|
assert len(stats_monitor.lookup_requests) == 0
|
|
|
|
|
|
def test_remote_read_metrics(stats_monitor):
|
|
stats_monitor.update_interval_remote_read_metrics(read_bytes=1024)
|
|
stats_monitor.update_interval_remote_read_metrics(read_bytes=2048)
|
|
stats = stats_monitor.get_stats_and_clear()
|
|
assert stats.interval_remote_read_requests == 2
|
|
assert stats.interval_remote_read_bytes == 3072
|
|
|
|
|
|
def test_remote_write_metrics(stats_monitor):
|
|
stats_monitor.update_interval_remote_write_metrics(write_bytes=512)
|
|
stats_monitor.update_interval_remote_write_metrics(write_bytes=1024)
|
|
stats = stats_monitor.get_stats_and_clear()
|
|
assert stats.interval_remote_write_requests == 2
|
|
assert stats.interval_remote_write_bytes == 1536
|
|
|
|
|
|
def test_remote_time_metrics(stats_monitor):
|
|
stats_monitor.update_interval_remote_time_to_get(get_time=10.5)
|
|
stats_monitor.update_interval_remote_time_to_put(put_time=15.2)
|
|
stats_monitor.update_interval_remote_time_to_get_sync(get_time_sync=12.3)
|
|
stats = stats_monitor.get_stats_and_clear()
|
|
assert stats.interval_remote_time_to_get == [10.5]
|
|
assert stats.interval_remote_time_to_put == [15.2]
|
|
assert stats.interval_remote_time_to_get_sync == [12.3]
|
|
|
|
|
|
def test_remote_ping_metrics(stats_monitor):
|
|
# Test successful ping
|
|
stats_monitor.update_remote_ping_latency(latency=25.5)
|
|
stats_monitor.update_remote_ping_error_code(error_code=0)
|
|
stats = stats_monitor.get_stats_and_clear()
|
|
assert stats.interval_remote_ping_latency == 25.5
|
|
assert stats.interval_remote_ping_success == 1
|
|
assert stats.interval_remote_ping_errors == 0
|
|
assert stats.interval_remote_ping_error_code == 0
|
|
|
|
|
|
def test_remote_ping_errors(stats_monitor):
|
|
# Test ping errors
|
|
stats_monitor.update_remote_ping_error_code(error_code=404)
|
|
stats_monitor.update_remote_ping_error_code(error_code=500)
|
|
stats = stats_monitor.get_stats_and_clear()
|
|
assert stats.interval_remote_ping_errors == 2
|
|
assert stats.interval_remote_ping_success == 0
|
|
assert stats.interval_remote_ping_error_code == 500
|
|
|
|
|
|
def test_retrieve_and_store_speed(stats_monitor):
|
|
# Test retrieve speed calculation
|
|
stats_obj_retrieve = stats_monitor.on_retrieve_request(num_tokens=1000)
|
|
stats_monitor.on_retrieve_finished(
|
|
retrieve_stats=stats_obj_retrieve, num_retrieved_tokens=1000
|
|
)
|
|
|
|
# Test store speed calculation
|
|
stats_obj_store = stats_monitor.on_store_request(num_tokens=500)
|
|
stats_monitor.on_store_finished(store_stats=stats_obj_store)
|
|
|
|
stats = stats_monitor.get_stats_and_clear()
|
|
assert len(stats.retrieve_speed) == 1
|
|
assert len(stats.store_speed) == 1
|
|
assert stats.retrieve_speed[0] > 0 # Should be tokens/second
|
|
assert stats.store_speed[0] > 0 # Should be tokens/second
|
|
|
|
|
|
def test_multiple_lookup_operations(stats_monitor):
|
|
# Test multiple lookup operations
|
|
stats_obj_1 = stats_monitor.on_lookup_request(num_tokens=100)
|
|
stats_monitor.on_lookup_finished(stats=stats_obj_1, num_hit_tokens=80)
|
|
stats_obj_2 = stats_monitor.on_lookup_request(num_tokens=200)
|
|
stats_monitor.on_lookup_finished(stats=stats_obj_2, num_hit_tokens=150)
|
|
assert len(stats_monitor.lookup_requests) == 2
|
|
assert stats_monitor.lookup_requests[stats_obj_1.request_id].hit_rate() == 0.8
|
|
assert stats_monitor.lookup_requests[stats_obj_2.request_id].hit_rate() == 0.75
|
|
|
|
stats = stats_monitor.get_stats_and_clear()
|
|
assert stats.interval_lookup_requests == 2
|
|
assert stats.interval_lookup_tokens == 300
|
|
assert stats.interval_lookup_hits == 230
|
|
assert stats.lookup_hit_rate == 230 / 300
|
|
assert len(stats.interval_lookup_hit_rates) == 2
|
|
assert len(stats_monitor.lookup_requests) == 0
|
|
|
|
|
|
def test_mixed_remote_operations(stats_monitor):
|
|
# Test a mix of remote operations
|
|
stats_monitor.update_interval_remote_read_metrics(read_bytes=1024)
|
|
stats_monitor.update_interval_remote_write_metrics(write_bytes=512)
|
|
stats_monitor.update_interval_remote_time_to_get(get_time=10.0)
|
|
stats_monitor.update_interval_remote_time_to_put(put_time=20.0)
|
|
stats_monitor.update_interval_remote_time_to_get_sync(get_time_sync=15.0)
|
|
stats_monitor.update_remote_ping_latency(latency=30.0)
|
|
stats_monitor.update_remote_ping_error_code(error_code=0)
|
|
|
|
stats = stats_monitor.get_stats_and_clear()
|
|
assert stats.interval_remote_read_requests == 1
|
|
assert stats.interval_remote_read_bytes == 1024
|
|
assert stats.interval_remote_write_requests == 1
|
|
assert stats.interval_remote_write_bytes == 512
|
|
assert stats.interval_remote_time_to_get == [10.0]
|
|
assert stats.interval_remote_time_to_put == [20.0]
|
|
assert stats.interval_remote_time_to_get_sync == [15.0]
|
|
assert stats.interval_remote_ping_latency == 30.0
|
|
assert stats.interval_remote_ping_success == 1
|
|
assert stats.interval_remote_ping_errors == 0
|
|
|
|
|
|
def test_combined_operations(stats_monitor):
|
|
stats_obj_retrieve = stats_monitor.on_retrieve_request(num_tokens=200)
|
|
stats_monitor.on_retrieve_finished(
|
|
retrieve_stats=stats_obj_retrieve,
|
|
num_retrieved_tokens=200,
|
|
)
|
|
stats_obj_store = stats_monitor.on_store_request(num_tokens=100)
|
|
stats_monitor.on_store_finished(store_stats=stats_obj_store)
|
|
stats_monitor.update_local_cache_usage(usage=512)
|
|
stats_monitor.update_remote_cache_usage(usage=1024)
|
|
stats_monitor.update_local_storage_usage(usage=2048)
|
|
|
|
stats_monitor2 = LMCStatsMonitor.GetOrCreate()
|
|
stats = stats_monitor2.get_stats_and_clear()
|
|
|
|
assert stats.interval_retrieve_requests == 1
|
|
assert stats.interval_store_requests == 1
|
|
assert stats.retrieve_hit_rate == 1.0
|
|
assert stats.local_cache_usage_bytes == 512
|
|
assert stats.remote_cache_usage_bytes == 1024
|
|
assert stats.local_storage_usage_bytes == 2048
|
|
assert len(stats.time_to_retrieve) == 1
|
|
assert len(stats.time_to_store) == 1
|
|
|
|
|
|
def test_stats_clearing(stats_monitor):
|
|
# Add some data
|
|
stats_obj = stats_monitor.on_lookup_request(num_tokens=100)
|
|
stats_monitor.update_interval_remote_read_metrics(read_bytes=1024)
|
|
stats_monitor.update_remote_ping_latency(latency=25.0)
|
|
|
|
assert len(stats_monitor.lookup_requests) == 1
|
|
|
|
# Get stats (which should clear them)
|
|
stats = stats_monitor.get_stats_and_clear()
|
|
assert stats.interval_lookup_requests == 1
|
|
assert stats.interval_lookup_tokens == 100
|
|
assert stats.interval_lookup_hits == 0
|
|
assert stats.interval_remote_read_requests == 1
|
|
assert stats.interval_remote_read_bytes == 1024
|
|
assert stats.interval_remote_ping_latency == 25.0
|
|
assert len(stats.interval_lookup_hit_rates) == 0
|
|
|
|
# Get stats again - should be cleared
|
|
stats2 = stats_monitor.get_stats_and_clear()
|
|
assert stats2.interval_lookup_requests == 0
|
|
assert stats2.interval_lookup_tokens == 0
|
|
assert stats2.interval_lookup_hits == 0
|
|
assert stats2.interval_remote_read_requests == 0
|
|
assert stats2.interval_remote_read_bytes == 0
|
|
assert stats2.interval_remote_ping_latency == 0
|
|
assert len(stats2.interval_lookup_hit_rates) == 0
|
|
|
|
assert len(stats_monitor.lookup_requests) == 1
|
|
|
|
# finish lookup request
|
|
stats_monitor.on_lookup_finished(stats=stats_obj, num_hit_tokens=80)
|
|
stats3 = stats_monitor.get_stats_and_clear()
|
|
assert stats3.interval_lookup_requests == 0
|
|
assert stats3.interval_lookup_tokens == 0
|
|
assert stats3.interval_lookup_hits == 80
|
|
assert stats3.interval_remote_read_requests == 0
|
|
assert stats3.interval_remote_read_bytes == 0
|
|
assert stats3.interval_remote_ping_latency == 0
|
|
assert len(stats3.interval_lookup_hit_rates) == 1
|
|
assert stats3.interval_lookup_hit_rates[0] == 0.8
|
|
|
|
assert len(stats_monitor.lookup_requests) == 0
|
|
|
|
|
|
def test_zero_division_protection(stats_monitor):
|
|
# Test that hit rates handle zero division gracefully
|
|
stats = stats_monitor.get_stats_and_clear()
|
|
assert stats.retrieve_hit_rate == 1.0
|
|
assert stats.lookup_hit_rate == 0
|
|
|
|
|
|
# ---- PrometheusLogger custom histogram bucket tests ----
|
|
|
|
|
|
@pytest.fixture(scope="function")
|
|
def _cleanup_prometheus_logger():
|
|
"""Reset PrometheusLogger instances and metrics between tests."""
|
|
LMCStatsMonitor.unregister_all_metrics()
|
|
PrometheusLogger._instances = {}
|
|
yield
|
|
LMCStatsMonitor.unregister_all_metrics()
|
|
PrometheusLogger._instances = {}
|
|
|
|
|
|
def _make_metadata():
|
|
"""Create a minimal mock LMCacheMetadata for testing."""
|
|
meta = MagicMock()
|
|
meta.model_name = "test_model"
|
|
meta.worker_id = 0
|
|
meta.role = "worker"
|
|
meta.served_model_name = None
|
|
return meta
|
|
|
|
|
|
def _make_config(extra_config=None):
|
|
"""Create a minimal mock config for testing."""
|
|
cfg = MagicMock()
|
|
cfg.extra_config = extra_config
|
|
cfg.get_extra_config_value = lambda key, default=None: (
|
|
extra_config.get(key, default) if extra_config is not None else default
|
|
)
|
|
return cfg
|
|
|
|
|
|
def _sample_labels(prom: PrometheusLogger) -> dict[str, str]:
|
|
"""Return labels as Prometheus exposes them in collected samples."""
|
|
return {name: str(value) for name, value in prom.labels.items()}
|
|
|
|
|
|
def test_prometheus_logger_default_buckets(_cleanup_prometheus_logger):
|
|
"""Histogram should use default buckets when no config."""
|
|
meta = _make_metadata()
|
|
prom = PrometheusLogger(meta)
|
|
# Verify histogram was created (basic sanity check)
|
|
assert prom.histogram_time_to_retrieve is not None
|
|
|
|
|
|
def test_prometheus_logger_custom_buckets(_cleanup_prometheus_logger):
|
|
"""Histogram should use custom buckets from config."""
|
|
meta = _make_metadata()
|
|
custom_buckets = [0.1, 0.5, 1.0, 5.0]
|
|
cfg = _make_config(
|
|
{
|
|
"histogram_bucket_time_to_retrieve": custom_buckets,
|
|
}
|
|
)
|
|
prom = PrometheusLogger(meta, config=cfg)
|
|
labels = prom.labels
|
|
|
|
# Verify the histogram's upper_bounds match custom buckets + Inf
|
|
hist = prom.histogram_time_to_retrieve.labels(**labels)
|
|
# prometheus_client Histogram stores buckets as _upper_bounds
|
|
upper_bounds = list(hist._upper_bounds)
|
|
for bucket_val in custom_buckets:
|
|
assert bucket_val in upper_bounds
|
|
|
|
|
|
def test_prometheus_logger_partial_custom_buckets(
|
|
_cleanup_prometheus_logger,
|
|
):
|
|
"""Only specified histograms should get custom buckets."""
|
|
meta = _make_metadata()
|
|
custom_buckets = [1, 10, 100]
|
|
cfg = _make_config(
|
|
{
|
|
"histogram_bucket_store_speed": custom_buckets,
|
|
}
|
|
)
|
|
prom = PrometheusLogger(meta, config=cfg)
|
|
labels = prom.labels
|
|
|
|
# store_speed should have custom buckets
|
|
hist = prom.histogram_store_speed.labels(**labels)
|
|
upper_bounds = list(hist._upper_bounds)
|
|
for bucket_val in custom_buckets:
|
|
assert bucket_val in upper_bounds
|
|
|
|
# time_to_retrieve should still have default buckets
|
|
hist_default = prom.histogram_time_to_retrieve.labels(**labels)
|
|
default_upper = list(hist_default._upper_bounds)
|
|
# default first bucket is 0.001
|
|
assert 0.001 in default_upper
|
|
|
|
|
|
def test_prometheus_logger_get_or_create_with_config(
|
|
_cleanup_prometheus_logger,
|
|
):
|
|
"""GetOrCreate should pass config to the constructor."""
|
|
meta = _make_metadata()
|
|
custom_buckets = [0.01, 0.1, 1.0]
|
|
cfg = _make_config(
|
|
{
|
|
"histogram_bucket_time_to_store": custom_buckets,
|
|
}
|
|
)
|
|
prom = PrometheusLogger.GetOrCreate(meta, config=cfg)
|
|
labels = prom.labels
|
|
|
|
hist = prom.histogram_time_to_store.labels(**labels)
|
|
upper_bounds = list(hist._upper_bounds)
|
|
for bucket_val in custom_buckets:
|
|
assert bucket_val in upper_bounds
|
|
|
|
|
|
def test_prometheus_logger_get_or_create_allows_multiple_roles(
|
|
_cleanup_prometheus_logger,
|
|
):
|
|
"""Different roles should not be treated as global metadata conflicts."""
|
|
worker_meta = _make_metadata()
|
|
scheduler_meta = _make_metadata()
|
|
scheduler_meta.role = "scheduler"
|
|
|
|
worker_prom = PrometheusLogger.GetOrCreate(worker_meta)
|
|
scheduler_prom = PrometheusLogger.GetOrCreate(scheduler_meta)
|
|
|
|
assert worker_prom.labels["role"] == "worker"
|
|
assert scheduler_prom.labels["role"] == "scheduler"
|
|
|
|
|
|
def test_prometheus_logger_label_view_rebinds_dynamic_metrics(
|
|
_cleanup_prometheus_logger,
|
|
):
|
|
"""Dynamic gauge callbacks should use each label view's labels."""
|
|
worker_meta = _make_metadata()
|
|
scheduler_meta = _make_metadata()
|
|
scheduler_meta.role = "scheduler"
|
|
|
|
worker_prom = PrometheusLogger.GetOrCreate(worker_meta)
|
|
scheduler_prom = PrometheusLogger.GetOrCreate(scheduler_meta)
|
|
|
|
worker_prom.lmcache_is_healthy.set_function(lambda: 1)
|
|
scheduler_prom.lmcache_is_healthy.set_function(lambda: 2)
|
|
|
|
assert (
|
|
REGISTRY.get_sample_value(
|
|
"lmcache:lmcache_is_healthy",
|
|
_sample_labels(worker_prom),
|
|
)
|
|
== 1
|
|
)
|
|
assert (
|
|
REGISTRY.get_sample_value(
|
|
"lmcache:lmcache_is_healthy",
|
|
_sample_labels(scheduler_prom),
|
|
)
|
|
== 2
|
|
)
|
|
|
|
|
|
def test_prometheus_logger_reset_reinitializes_all_label_views(
|
|
_cleanup_prometheus_logger,
|
|
):
|
|
"""Resetting one label view should keep metric children visible for all views."""
|
|
worker_meta = _make_metadata()
|
|
scheduler_meta = _make_metadata()
|
|
scheduler_meta.role = "scheduler"
|
|
|
|
worker_prom = PrometheusLogger.GetOrCreate(worker_meta)
|
|
scheduler_prom = PrometheusLogger.GetOrCreate(scheduler_meta)
|
|
|
|
worker_prom.counter_num_retrieve_requests.labels(**worker_prom.labels).inc()
|
|
scheduler_prom.counter_num_retrieve_requests.labels(**scheduler_prom.labels).inc()
|
|
scheduler_prom.reset_counters()
|
|
|
|
assert (
|
|
REGISTRY.get_sample_value(
|
|
"lmcache:num_retrieve_requests_total",
|
|
_sample_labels(worker_prom),
|
|
)
|
|
== 0
|
|
)
|
|
assert (
|
|
REGISTRY.get_sample_value(
|
|
"lmcache:num_retrieve_requests_total",
|
|
_sample_labels(scheduler_prom),
|
|
)
|
|
== 0
|
|
)
|