From 34d8392422134ec6bf6877ffda9607f3790182d4 Mon Sep 17 00:00:00 2001 From: Alex Bozarth Date: Wed, 11 Mar 2026 12:48:16 -0500 Subject: [PATCH 1/2] fix: hf metrics tests run out of memory Signed-off-by: Alex Bozarth --- test/telemetry/test_metrics_backend.py | 41 ++++++++++++++++++-------- 1 file changed, 29 insertions(+), 12 deletions(-) diff --git a/test/telemetry/test_metrics_backend.py b/test/telemetry/test_metrics_backend.py index 9bba4503a..a1376b4aa 100644 --- a/test/telemetry/test_metrics_backend.py +++ b/test/telemetry/test_metrics_backend.py @@ -51,6 +51,33 @@ def enable_metrics(monkeypatch): importlib.reload(mellea.telemetry.metrics) +@pytest.fixture(scope="module") +def hf_metrics_backend(gh_run): + """Shared HuggingFace backend for telemetry metrics tests. + + Uses module scope to load the model once and reuse it across all tests, + preventing memory exhaustion from loading multiple model instances. + """ + if gh_run: + pytest.skip("Skipping HuggingFace backend creation in CI") + + from mellea.backends.cache import SimpleLRUCache + from mellea.backends.huggingface import LocalHFBackend + + backend = LocalHFBackend( + model_id=IBM_GRANITE_4_HYBRID_MICRO.hf_model_name, # type: ignore + cache=SimpleLRUCache(5), + ) + + yield backend + + # Cleanup + import gc + + del backend + gc.collect() + + def get_metric_value(metrics_data, metric_name, attributes=None): """Helper to extract metric value from metrics data. @@ -299,13 +326,9 @@ async def test_litellm_token_metrics_integration( @pytest.mark.huggingface @pytest.mark.parametrize("stream", [False, True], ids=["non-streaming", "streaming"]) async def test_huggingface_token_metrics_integration( - enable_metrics, metric_reader, stream, gh_run + enable_metrics, metric_reader, stream, hf_metrics_backend ): """Test that HuggingFace backend records token metrics correctly.""" - if gh_run: - pytest.skip("Skipping in CI - requires model download") - - from mellea.backends.huggingface import LocalHFBackend from mellea.backends.model_options import ModelOption from mellea.telemetry import metrics as metrics_module @@ -315,17 +338,11 @@ async def test_huggingface_token_metrics_integration( metrics_module._input_token_counter = None metrics_module._output_token_counter = None - from mellea.backends.cache import SimpleLRUCache - - backend = LocalHFBackend( - model_id=IBM_GRANITE_4_HYBRID_MICRO.hf_model_name, # type: ignore - cache=SimpleLRUCache(5), - ) ctx = SimpleContext() ctx = ctx.add(Message(role="user", content="Say 'hello' and nothing else")) model_options = {ModelOption.STREAM: True} if stream else {} - mot, _ = await backend.generate_from_context( + mot, _ = await hf_metrics_backend.generate_from_context( Message(role="assistant", content=""), ctx, model_options=model_options ) From 7a966e439e4182fd12d474c31c2cf65ca41bac1e Mon Sep 17 00:00:00 2001 From: Alex Bozarth Date: Thu, 12 Mar 2026 10:29:59 -0500 Subject: [PATCH 2/2] test: add requires_heavy_ram marker to HuggingFace backend tests Add @pytest.mark.requires_heavy_ram to tests that instantiate LocalHFBackend to address memory leak issues when running these tests in pytest. This ensures tests are skipped on systems without sufficient RAM. Changes: - test/telemetry/test_metrics_backend.py: Added marker to test_huggingface_token_metrics_integration - test/stdlib/test_spans.py: Added marker to module-level pytestmark Signed-off-by: Alex Bozarth --- test/stdlib/test_spans.py | 7 ++++++- test/telemetry/test_metrics_backend.py | 1 + 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/test/stdlib/test_spans.py b/test/stdlib/test_spans.py index 06c9f8a2f..82e63c547 100644 --- a/test/stdlib/test_spans.py +++ b/test/stdlib/test_spans.py @@ -8,7 +8,12 @@ from mellea.stdlib.session import MelleaSession, start_session # Module-level markers for all tests using Granite 4 hybrid micro (3B model) -pytestmark = [pytest.mark.huggingface, pytest.mark.requires_gpu, pytest.mark.llm] +pytestmark = [ + pytest.mark.huggingface, + pytest.mark.requires_gpu, + pytest.mark.requires_heavy_ram, + pytest.mark.llm, +] # We edit the context type in the async tests below. Don't change the scope here. diff --git a/test/telemetry/test_metrics_backend.py b/test/telemetry/test_metrics_backend.py index a1376b4aa..5b2702bd1 100644 --- a/test/telemetry/test_metrics_backend.py +++ b/test/telemetry/test_metrics_backend.py @@ -324,6 +324,7 @@ async def test_litellm_token_metrics_integration( @pytest.mark.asyncio @pytest.mark.llm @pytest.mark.huggingface +@pytest.mark.requires_heavy_ram @pytest.mark.parametrize("stream", [False, True], ids=["non-streaming", "streaming"]) async def test_huggingface_token_metrics_integration( enable_metrics, metric_reader, stream, hf_metrics_backend