mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-12-03 09:53:45 +00:00
a8bc99408c
The integration-tests.sh script already sets LLAMA_STACK_TEST_STACK_CONFIG_TYPE based on the stack config. Our custom detection logic was unnecessary and potentially interfering. Revert to relying on the environment variable set by the test script. The LLAMA_STACK_DISABLE_GUNICORN environment variable is still set correctly when stack_mode == 'server', which happens for both server: and docker: configs.
82 lines
2.9 KiB
Python
82 lines
2.9 KiB
Python
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
|
# All rights reserved.
|
|
#
|
|
# This source code is licensed under the terms described in the LICENSE file in
|
|
# the root directory of this source tree.
|
|
|
|
"""Telemetry test configuration supporting both library and server test modes."""
|
|
|
|
import os
|
|
|
|
import pytest
|
|
|
|
import llama_stack.core.telemetry.telemetry as telemetry_module
|
|
from llama_stack.testing.api_recorder import patch_httpx_for_test_id
|
|
from tests.integration.fixtures.common import instantiate_llama_stack_client
|
|
from tests.integration.telemetry.collectors import InMemoryTelemetryManager, OtlpHttpTestCollector
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def telemetry_test_collector():
|
|
# Stack mode is set by integration-tests.sh based on STACK_CONFIG
|
|
stack_mode = os.environ.get("LLAMA_STACK_TEST_STACK_CONFIG_TYPE", "library_client")
|
|
|
|
if stack_mode == "server":
|
|
try:
|
|
collector = OtlpHttpTestCollector()
|
|
except RuntimeError as exc:
|
|
pytest.skip(str(exc))
|
|
env_overrides = {
|
|
"OTEL_EXPORTER_OTLP_ENDPOINT": collector.endpoint,
|
|
"OTEL_EXPORTER_OTLP_PROTOCOL": "http/protobuf",
|
|
"OTEL_BSP_SCHEDULE_DELAY": "200",
|
|
"OTEL_BSP_EXPORT_TIMEOUT": "2000",
|
|
"LLAMA_STACK_DISABLE_GUNICORN": "true", # Disable multi-process for telemetry collection
|
|
}
|
|
|
|
previous_env = {key: os.environ.get(key) for key in env_overrides}
|
|
previous_force_restart = os.environ.get("LLAMA_STACK_TEST_FORCE_SERVER_RESTART")
|
|
|
|
for key, value in env_overrides.items():
|
|
os.environ[key] = value
|
|
|
|
os.environ["LLAMA_STACK_TEST_FORCE_SERVER_RESTART"] = "1"
|
|
telemetry_module._TRACER_PROVIDER = None
|
|
|
|
try:
|
|
yield collector
|
|
finally:
|
|
collector.shutdown()
|
|
for key, prior in previous_env.items():
|
|
if prior is None:
|
|
os.environ.pop(key, None)
|
|
else:
|
|
os.environ[key] = prior
|
|
if previous_force_restart is None:
|
|
os.environ.pop("LLAMA_STACK_TEST_FORCE_SERVER_RESTART", None)
|
|
else:
|
|
os.environ["LLAMA_STACK_TEST_FORCE_SERVER_RESTART"] = previous_force_restart
|
|
else:
|
|
manager = InMemoryTelemetryManager()
|
|
try:
|
|
yield manager.collector
|
|
finally:
|
|
manager.shutdown()
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def llama_stack_client(telemetry_test_collector, request):
|
|
"""Ensure telemetry collector is ready before initializing the stack client."""
|
|
patch_httpx_for_test_id()
|
|
client = instantiate_llama_stack_client(request.session)
|
|
return client
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_otlp_collector(telemetry_test_collector):
|
|
"""Provides access to telemetry data and clears between tests."""
|
|
telemetry_test_collector.clear()
|
|
try:
|
|
yield telemetry_test_collector
|
|
finally:
|
|
telemetry_test_collector.clear()
|