mirror of
https://github.com/meta-llama/llama-stack.git
synced 2025-12-03 09:53:45 +00:00
c8f82cad6a
The telemetry fixture was only checking LLAMA_STACK_TEST_STACK_CONFIG_TYPE environment variable, which defaults to 'library_client'. In CI, tests run with --stack-config=docker:ci-tests, which wasn't being detected as server mode. This commit checks the --stack-config argument and treats both 'server:' and 'docker:' prefixes as server mode, ensuring LLAMA_STACK_DISABLE_GUNICORN is set when needed for telemetry span collection.
90 lines
3.3 KiB
Python
90 lines
3.3 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(request):
|
|
# Determine stack mode from --stack-config argument
|
|
stack_config = request.session.config.getoption("--stack-config", default=None)
|
|
if not stack_config:
|
|
stack_config = os.environ.get("LLAMA_STACK_CONFIG", "")
|
|
|
|
# Check if running in server or docker mode (both need server-side telemetry)
|
|
if stack_config.startswith("server:") or stack_config.startswith("docker:"):
|
|
stack_mode = "server"
|
|
else:
|
|
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()
|