Compare commits

...

35 Commits

Author SHA1 Message Date
Yuge Zhang e6b09e3fea update labels 2025-11-15 15:56:47 +08:00
Yuge Zhang e9317ba4a1 . 2025-11-15 15:19:46 +08:00
Yuge Zhang 8b1514a7cb . 2025-11-15 14:59:49 +08:00
Yuge Zhang 055638af1a . 2025-11-15 14:44:42 +08:00
Yuge Zhang ed62a8e3f6 fix tests 2025-11-15 14:42:10 +08:00
Yuge Zhang c08dc622ff fix agentops instrumentation tests 2025-11-15 14:13:00 +08:00
Yuge Zhang 908773d9ca fix port conflict 2025-11-15 13:31:16 +08:00
Yuge Zhang bd050860f9 fix tests 2025-11-15 13:28:21 +08:00
Yuge Zhang 2cf14121f0 add minimal examples tests 2025-11-15 13:04:16 +08:00
Yuge Zhang 2228aedf45 update documentation 2025-11-15 12:43:36 +08:00
Yuge Zhang 852d76f8bb update readme 2025-11-15 12:32:42 +08:00
Yuge Zhang d89d67c445 update minimal examples 2025-11-15 12:19:07 +08:00
Yuge Zhang 0f6f54b0fa add tests 2025-11-15 12:00:38 +08:00
Yuge Zhang 4f9bab87a4 . 2025-11-15 11:29:47 +08:00
Yuge Zhang 528020d372 add tests to llm proxy 2025-11-15 11:24:49 +08:00
Yuge Zhang 1cbd1c5f41 add tests 2025-11-15 10:43:47 +08:00
Yuge Zhang 87f6aed26b patch candidates 2025-11-15 10:22:41 +08:00
Yuge Zhang c9f0b90918 resolve comments and fix tests 2025-11-15 10:04:00 +08:00
Yuge Zhang 057d1f8d59 delete unused files 2025-11-15 01:27:56 +08:00
Yuge Zhang c3403c4b94 fix agentops and otel 2025-11-15 01:25:15 +08:00
Yuge Zhang 5173326d90 . 2025-11-15 00:29:08 +08:00
Yuge Zhang 51849445d5 move agentops logic to otel 2025-11-15 00:13:54 +08:00
Yuge Zhang b9831fea5c extract common util 2025-11-14 18:30:04 +08:00
Yuge Zhang 2e19bc6456 support llm proxy 2025-11-14 18:15:33 +08:00
Yuge Zhang bd301fd9ee agentops+ 2025-11-14 17:54:42 +08:00
Yuge Zhang cea5baf26d support agentops 2025-11-14 17:33:44 +08:00
Yuge Zhang 705b3ba98e add new examples 2025-11-14 15:52:12 +08:00
Yuge Zhang 5874453dd1 span proto to pan 2025-11-14 15:17:32 +08:00
Yuge Zhang 926d1ec7e7 Merge branch 'main' of github.com:microsoft/agent-lightning into feature/otlp 2025-11-14 14:43:47 +08:00
Yuge Zhang b3640786be . 2025-11-14 14:43:17 +08:00
Yuge Zhang e9b953e19d . 2025-11-14 01:02:20 +08:00
Yuge Zhang 2824fe7e12 . 2025-11-13 23:48:01 +08:00
Yuge Zhang 5106b73999 modify store 2025-11-13 23:45:53 +08:00
Yuge Zhang 85c581a41f update examples 2025-11-13 23:27:46 +08:00
Yuge Zhang 67e19143af first examples 2025-11-13 23:12:27 +08:00
30 changed files with 2350 additions and 366 deletions
+1 -1
View File
@@ -14,7 +14,7 @@ on:
run-name: >-
${{ github.event_name == 'repository_dispatch'
&& format(
'PR #{0} - Label {1} - {2}',
'APO - PR #{0} - {1} - {2}',
github.event.client_payload.pull_number,
github.event.client_payload.ci_label,
github.event.client_payload.correlation_id
+1 -1
View File
@@ -14,7 +14,7 @@ on:
run-name: >-
${{ github.event_name == 'repository_dispatch'
&& format(
'PR #{0} - Label {1} - {2}',
'Calc-X - PR #{0} - {1} - {2}',
github.event.client_payload.pull_number,
github.event.client_payload.ci_label,
github.event.client_payload.correlation_id
+1 -1
View File
@@ -14,7 +14,7 @@ on:
run-name: >-
${{ github.event_name == 'repository_dispatch'
&& format(
'PR #{0} - Label {1} - {2}',
'Backward Compatibility - PR #{0} - {1} - {2}',
github.event.client_payload.pull_number,
github.event.client_payload.ci_label,
github.event.client_payload.correlation_id
+1 -1
View File
@@ -14,7 +14,7 @@ on:
run-name: >-
${{ github.event_name == 'repository_dispatch'
&& format(
'PR #{0} - Label {1} - {2}',
'Spider - PR #{0} - {1} - {2}',
github.event.client_payload.pull_number,
github.event.client_payload.ci_label,
github.event.client_payload.correlation_id
+1 -1
View File
@@ -14,7 +14,7 @@ on:
run-name: >-
${{ github.event_name == 'repository_dispatch'
&& format(
'PR #{0} - Label {1} - {2}',
'Unsloth - PR #{0} - {1} - {2}',
github.event.client_payload.pull_number,
github.event.client_payload.ci_label,
github.event.client_payload.correlation_id
+176 -2
View File
@@ -14,7 +14,7 @@ on:
run-name: >-
${{ github.event_name == 'repository_dispatch'
&& format(
'PR #{0} - Label {1} - {2}',
'GPU Test - PR #{0} - {1} - {2}',
github.event.client_payload.pull_number,
github.event.client_payload.ci_label,
github.event.client_payload.correlation_id
@@ -69,7 +69,7 @@ jobs:
- name: Upload dependencies artifact
uses: actions/upload-artifact@v4
with:
name: dependencies-${{ matrix.python-version }}-${{ matrix.setup-script }}
name: dependencies-tests-full-${{ matrix.python-version }}-${{ matrix.setup-script }}
path: requirements-freeze.txt
compression-level: 0
@@ -95,3 +95,177 @@ jobs:
PYTEST_ADDOPTS: "--color=yes"
OPENAI_BASE_URL: http://localhost:12306/
OPENAI_API_KEY: dummy
minimal-examples:
if: >
github.event_name != 'repository_dispatch' ||
github.event.action == 'ci-gpu' ||
github.event.action == 'ci-all'
name: Minimal Examples with Python ${{ matrix.python-version }} (${{ matrix.setup-script }})
runs-on: [self-hosted, 1ES.Pool=agl-runner-gpu]
timeout-minutes: 30
strategy:
matrix:
include:
- python-version: '3.10'
setup-script: 'legacy'
- python-version: '3.12'
setup-script: 'stable'
- python-version: '3.13'
setup-script: 'latest'
fail-fast: false
steps:
- name: Check GPU status
run: nvidia-smi
- uses: actions/checkout@v4
with:
ref: ${{ github.event_name == 'repository_dispatch' && github.event.client_payload.pr_ref || (github.event.pull_request.number && format('refs/pull/{0}/merge', github.event.pull_request.number)) || github.ref }}
- uses: astral-sh/setup-uv@v7
with:
enable-cache: true
python-version: ${{ matrix.python-version }}
- name: Upgrade dependencies (latest)
run: uv lock --upgrade
if: matrix.setup-script == 'latest'
- name: Sync dependencies (latest)
run: uv sync --frozen --no-default-groups --extra apo --group dev --group agents --group torch-gpu-stable
if: matrix.setup-script == 'latest'
- name: Sync dependencies (stable & legacy)
run: uv sync --frozen --no-default-groups --extra apo --group dev --group agents --group torch-gpu-${{ matrix.setup-script }}
if: matrix.setup-script != 'latest'
- name: Freeze dependencies
run: |
set -ex
uv pip freeze | tee requirements-freeze.txt
echo "UV_LOCKED=1" >> $GITHUB_ENV
echo "UV_NO_SYNC=1" >> $GITHUB_ENV
- name: Upload dependencies artifact
uses: actions/upload-artifact@v4
with:
name: dependencies-minimal-examples-${{ matrix.python-version }}-${{ matrix.setup-script }}
path: requirements-freeze.txt
compression-level: 0
- name: Launch LiteLLM Proxy
run: |
./scripts/litellm_run.sh
env:
AZURE_API_BASE: ${{ secrets.AZURE_GROUP_SUBSCRIPTION_API_BASE }}
AZURE_API_KEY: ${{ secrets.AZURE_GROUP_SUBSCRIPTION_API_KEY }}
- name: Write Traces via Otel Tracer
run: |
set -euo pipefail
source .venv/bin/activate
cd examples/minimal
python write_traces.py otel
- name: Write Traces via AgentOps Tracer
env:
OPENAI_BASE_URL: http://localhost:12306/
OPENAI_API_KEY: dummy
run: |
set -euo pipefail
source .venv/bin/activate
cd examples/minimal
python write_traces.py agentops
- name: Write Traces via Otel Tracer with Client
run: |
set -euo pipefail
source .venv/bin/activate
cd examples/minimal
agl store --port 45993 --log-level DEBUG &
sleep 5
python write_traces.py otel --use-client
pkill -f agl && echo "SIGTERM sent to agl" || echo "No agl process found"
while pgrep -f agl; do
echo "Waiting for agl to finish..."
sleep 5
done
- name: Write Traces via AgentOps Tracer with Client
env:
OPENAI_BASE_URL: http://localhost:12306/
OPENAI_API_KEY: dummy
run: |
set -euo pipefail
source .venv/bin/activate
cd examples/minimal
agl store --port 45993 --log-level DEBUG &
sleep 5
python write_traces.py agentops --use-client
pkill -f agl && echo "SIGTERM sent to agl" || echo "No agl process found"
while pgrep -f agl; do
echo "Waiting for agl to finish..."
sleep 5
done
- name: vLLM Server
run: |
set -euo pipefail
source .venv/bin/activate
cd examples/minimal
python vllm_server.py Qwen/Qwen2.5-0.5B-Instruct
- name: LLM Proxy (OpenAI backend)
env:
OPENAI_API_BASE: http://localhost:12306/
OPENAI_API_KEY: dummy
run: |
set -euo pipefail
source .venv/bin/activate
cd examples/minimal
python llm_proxy.py openai gpt-4.1-mini &
LLM_PROXY_READY=0
for attempt in $(seq 1 30); do
if curl -sSf http://localhost:43886/health > /dev/null 2>&1; then
LLM_PROXY_READY=1
break
fi
sleep 2
done
if [[ "$LLM_PROXY_READY" != "1" ]]; then
echo "LLM proxy failed to become healthy" >&2
exit 1
fi
python llm_proxy.py test gpt-4.1-mini
pkill -f llm_proxy.py && echo "SIGTERM sent to llm_proxy.py" || echo "No llm_proxy.py process found"
while pgrep -f llm_proxy.py; do
echo "Waiting for llm_proxy.py to finish..."
sleep 5
done
- name: LLM Proxy (vLLM backend)
if: matrix.setup-script != 'legacy' # Skip if return_token_ids is not supported
run: |
set -euo pipefail
source .venv/bin/activate
cd examples/minimal
python llm_proxy.py vllm Qwen/Qwen2.5-0.5B-Instruct &
LLM_PROXY_READY=0
for attempt in $(seq 1 30); do
if curl -sSf http://localhost:43886/health > /dev/null 2>&1; then
LLM_PROXY_READY=1
break
fi
sleep 2
done
if [[ "$LLM_PROXY_READY" != "1" ]]; then
echo "LLM proxy failed to become healthy" >&2
exit 1
fi
python llm_proxy.py test Qwen/Qwen2.5-0.5B-Instruct
pkill -f llm_proxy.py && echo "SIGTERM sent to llm_proxy.py" || echo "No llm_proxy.py process found"
while pgrep -f llm_proxy.py; do
echo "Waiting for llm_proxy.py to finish..."
sleep 5
done
+7 -1
View File
@@ -25,9 +25,15 @@ def main(argv: Iterable[str] | None = None) -> int:
action="append",
help="Allowed CORS origin. Repeat for multiple origins. Use '*' to allow all origins.",
)
parser.add_argument(
"--log-level",
default="INFO",
choices=["DEBUG", "INFO", "WARNING", "ERROR"],
help="Configure the logging level for the store.",
)
args = parser.parse_args(list(argv) if argv is not None else None)
setup_logging()
setup_logging(args.log_level)
store = InMemoryLightningStore()
server = LightningStoreServer(
+25 -28
View File
@@ -13,7 +13,8 @@ from agentops.sdk.exporters import AuthenticatedOTLPExporter
from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter
from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
from opentelemetry.sdk.metrics.export import MetricExportResult
from opentelemetry.sdk.trace.export import SpanExportResult
from agentlightning.utils.otlp import LightningStoreOTLPExporter
logger = logging.getLogger(__name__)
@@ -32,25 +33,27 @@ def enable_agentops_service(enabled: bool = True) -> None:
"""
Enable or disable communication with the AgentOps service.
False (default): AgentOps exporters and clients will run in local mode
and will not attempt to communicate with the remote AgentOps service.
True: all exporters and clients will operate in normal mode and send data
to the AgentOps service as expected.
By default, AgentOps exporters and clients will run in local mode
and will NOT attempt to communicate with the remote AgentOps service.
Args:
enabled: If True, enable all AgentOps exporters and clients.
All exporters and clients will operate in normal mode and send data
to the [AgentOps service](https://www.agentops.ai).
"""
global _agentops_service_enabled
_agentops_service_enabled = enabled
logger.info(f"Switch set to {enabled} for exporters and clients.")
logger.info(f"AgentOps service enabled is set to {enabled}.")
def _patch_exporters():
import agentops.client.api
import agentops.sdk.core
import opentelemetry.exporter.otlp.proto.http.metric_exporter
import opentelemetry.exporter.otlp.proto.http.trace_exporter
agentops.sdk.core.AuthenticatedOTLPExporter = BypassableAuthenticatedOTLPExporter # type: ignore
opentelemetry.exporter.otlp.proto.http.metric_exporter.OTLPMetricExporter = BypassableOTLPMetricExporter
opentelemetry.exporter.otlp.proto.http.trace_exporter.OTLPSpanExporter = BypassableOTLPSpanExporter
agentops.sdk.core.OTLPMetricExporter = BypassableOTLPMetricExporter
if hasattr(agentops.sdk.core, "OTLPSpanExporter"):
agentops.sdk.core.OTLPSpanExporter = BypassableOTLPSpanExporter # type: ignore
agentops.client.api.V3Client = BypassableV3Client
agentops.client.api.V4Client = BypassableV4Client
@@ -58,12 +61,11 @@ def _patch_exporters():
def _unpatch_exporters():
import agentops.client.api
import agentops.sdk.core
import opentelemetry.exporter.otlp.proto.http.metric_exporter
import opentelemetry.exporter.otlp.proto.http.trace_exporter
agentops.sdk.core.AuthenticatedOTLPExporter = AuthenticatedOTLPExporter # type: ignore
opentelemetry.exporter.otlp.proto.http.metric_exporter.OTLPMetricExporter = OTLPMetricExporter
opentelemetry.exporter.otlp.proto.http.trace_exporter.OTLPSpanExporter = OTLPSpanExporter
agentops.sdk.core.OTLPMetricExporter = OTLPMetricExporter
if hasattr(agentops.sdk.core, "OTLPSpanExporter"):
agentops.sdk.core.OTLPSpanExporter = OTLPSpanExporter # type: ignore
agentops.client.api.V3Client = V3Client
agentops.client.api.V4Client = V4Client
@@ -243,18 +245,15 @@ def uninstrument_agentops():
pass
class BypassableAuthenticatedOTLPExporter(AuthenticatedOTLPExporter):
class BypassableAuthenticatedOTLPExporter(LightningStoreOTLPExporter, AuthenticatedOTLPExporter):
"""
AuthenticatedOTLPExporter with switchable service control.
When `_agentops_service_enabled` is False, skip export and return success.
"""
def export(self, *args: Any, **kwargs: Any) -> SpanExportResult:
if _agentops_service_enabled:
return super().export(*args, **kwargs)
else:
logger.debug("SwitchableAuthenticatedOTLPExporter is switched off, skipping export.")
return SpanExportResult.SUCCESS
def should_bypass(self) -> bool:
return not _agentops_service_enabled
class BypassableOTLPMetricExporter(OTLPMetricExporter):
@@ -271,18 +270,16 @@ class BypassableOTLPMetricExporter(OTLPMetricExporter):
return MetricExportResult.SUCCESS
class BypassableOTLPSpanExporter(OTLPSpanExporter):
class BypassableOTLPSpanExporter(LightningStoreOTLPExporter):
"""
OTLPSpanExporter with switchable service control.
When `_agentops_service_enabled` is False, skip export and return success.
This is used instead of BypassableAuthenticatedOTLPExporter on legacy AgentOps versions.
"""
def export(self, *args: Any, **kwargs: Any) -> SpanExportResult:
if _agentops_service_enabled:
return super().export(*args, **kwargs)
else:
logger.debug("SwitchableOTLPSpanExporter is switched off, skipping export.")
return SpanExportResult.SUCCESS
def should_bypass(self) -> bool:
return not _agentops_service_enabled
class BypassableV3Client(V3Client):
+58 -27
View File
@@ -40,12 +40,14 @@ from litellm.integrations.custom_logger import CustomLogger
from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig
from litellm.proxy.proxy_server import app, save_worker_config # pyright: ignore[reportUnknownVariableType]
from litellm.types.utils import CallTypes
from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
from opentelemetry.sdk.resources import Resource
from opentelemetry.sdk.trace import ReadableSpan
from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.types import Scope
from agentlightning.types import LLM, ProxyLLM
from agentlightning.types import LLM, ProxyLLM, SpanNames
from agentlightning.utils.server_launcher import (
LaunchMode,
PythonServerLauncher,
@@ -192,7 +194,7 @@ class LightningSpanExporter(SpanExporter):
def __init__(self, _store: Optional[LightningStore] = None):
self._store: Optional[LightningStore] = _store # this is only for testing purposes
self._buffer: List[ReadableSpan] = []
self._lock: Optional[threading.RLock] = None
self._lock: Optional[threading.Lock] = None
self._loop_lock_pid: Optional[int] = None
# Single dedicated event loop running in a daemon thread.
@@ -201,6 +203,8 @@ class LightningSpanExporter(SpanExporter):
self._loop: Optional[asyncio.AbstractEventLoop] = None
self._loop_thread: Optional[threading.Thread] = None
self._otlp_exporter = OTLPSpanExporter()
def _ensure_loop(self) -> asyncio.AbstractEventLoop:
"""Lazily initialize the event loop and thread on first use.
@@ -214,15 +218,15 @@ class LightningSpanExporter(SpanExporter):
self._loop_thread.start()
return self._loop
def _ensure_lock(self) -> threading.RLock:
def _ensure_lock(self) -> threading.Lock:
"""Lazily initialize the lock on first use.
Returns:
threading.RLock: The initialized lock.
threading.Lock: The initialized lock.
"""
self._clear_loop_and_lock()
if self._lock is None:
self._lock = threading.RLock()
self._lock = threading.Lock()
return self._lock
def _clear_loop_and_lock(self) -> None:
@@ -284,24 +288,18 @@ class LightningSpanExporter(SpanExporter):
with self._ensure_lock():
for span in spans:
self._buffer.append(span)
# Run the async flush on our private loop, synchronously from caller's POV.
async def _locked_flush():
# Take the lock inside the coroutine to serialize with other flushes.
with self._ensure_lock():
return await self._maybe_flush()
try:
loop = self._ensure_loop()
fut = asyncio.run_coroutine_threadsafe(_locked_flush(), loop)
fut.result() # Bubble up any exceptions from the coroutine.
except Exception as e:
logger.exception("Export flush failed: %s", e)
return SpanExportResult.FAILURE
default_endpoint = self._otlp_exporter._endpoint # pyright: ignore[reportPrivateUsage]
try:
self._maybe_flush()
except Exception as e:
logger.exception("Export flush failed: %s", e)
return SpanExportResult.FAILURE
finally:
self._otlp_exporter._endpoint = default_endpoint # pyright: ignore[reportPrivateUsage]
return SpanExportResult.SUCCESS
async def _maybe_flush(self):
def _maybe_flush(self):
"""Flush ready subtrees from the buffer.
Strategy:
@@ -323,11 +321,20 @@ class LightningSpanExporter(SpanExporter):
if not subtree_spans:
continue
# Store is initialized lazily here in most cases.
store = self._store or get_active_llm_proxy().get_store()
if store is None:
logger.warning("Store is not set in LLMProxy. Cannot log spans to store.")
continue
# If the store supports OTLP endpoint, use it.
if store.capabilities.get("otlp_traces", False):
otlp_traces_endpoint = store.otlp_traces_endpoint()
self._otlp_exporter._endpoint = otlp_traces_endpoint # pyright: ignore[reportPrivateUsage]
otlp_enabled = True
else:
otlp_enabled = False
# Merge all custom headers found in the subtree.
headers_merged: Dict[str, Any] = {}
@@ -383,10 +390,34 @@ class LightningSpanExporter(SpanExporter):
sequence_id_decimal = int(sequence_id)
# Persist each span in the subtree with the resolved identifiers.
for span in subtree_spans:
await store.add_otel_span(
rollout_id=rollout_id, attempt_id=attempt_id, sequence_id=sequence_id_decimal, readable_span=span
)
if otlp_enabled:
# If store has OTLP support, directly use OTLP exporter and export in batch
for span in subtree_spans:
span._resource = span._resource.merge( # pyright: ignore[reportPrivateUsage]
Resource.create(
{
SpanNames.ROLLOUT_ID: rollout_id,
SpanNames.ATTEMPT_ID: attempt_id,
SpanNames.SPAN_SEQUENCE_ID: sequence_id_decimal,
}
)
)
export_result = self._otlp_exporter.export(subtree_spans)
if export_result != SpanExportResult.SUCCESS:
raise RuntimeError(f"Failed to export spans via OTLP exporter. Result: {export_result}")
else:
# The old way: store does not support OTLP endpoint
for span in subtree_spans:
loop = self._ensure_loop()
add_otel_span_task = store.add_otel_span(
rollout_id=rollout_id,
attempt_id=attempt_id,
sequence_id=sequence_id_decimal,
readable_span=span,
)
fut = asyncio.run_coroutine_threadsafe(add_otel_span_task, loop)
fut.result() # Bubble up any exceptions from the coroutine.
def _get_root_span_ids(self) -> Iterable[int]:
"""Yield span_ids for root spans currently in the buffer.
@@ -1210,15 +1241,15 @@ class LLMProxy:
raise ValueError("Store is not set. Please set the store before starting the LLMProxy.")
store_capabilities = self.store.capabilities
if self.server_launcher.args.launch_mode == "mp" and not store_capabilities["zero_copy"]:
if self.server_launcher.args.launch_mode == "mp" and not store_capabilities.get("zero_copy", False):
raise RuntimeError(
"The store does not support zero-copy. Please use another store, or use asyncio or thread mode to launch the server."
)
elif self.server_launcher.args.launch_mode == "thread" and not store_capabilities["thread_safe"]:
elif self.server_launcher.args.launch_mode == "thread" and not store_capabilities.get("thread_safe", False):
raise RuntimeError(
"The store is not thread-safe. Please use another store, or use asyncio mode to launch the server."
)
elif self.server_launcher.args.launch_mode == "asyncio" and not store_capabilities["async_safe"]:
elif self.server_launcher.args.launch_mode == "asyncio" and not store_capabilities.get("async_safe", False):
raise RuntimeError("The store is not async-safe. Please use another store.")
logger.info(
+24 -2
View File
@@ -53,8 +53,11 @@ UNSET = _UnsetType()
Unset = _UnsetType # Alias for convenience
class LightningStoreCapabilities(TypedDict):
"""Capability of a LightningStore implementation."""
class LightningStoreCapabilities(TypedDict, total=False):
"""Capability of a LightningStore implementation.
All keys are optional and false by default.
"""
thread_safe: bool
"""Whether the store is thread-safe."""
@@ -62,6 +65,8 @@ class LightningStoreCapabilities(TypedDict):
"""Whether the store is async-safe."""
zero_copy: bool
"""Whether the store has only one copy across all threads/processes."""
otlp_traces: bool
"""Whether the store supports OTLP/HTTP traces."""
class LightningStore:
@@ -93,8 +98,25 @@ class LightningStore:
thread_safe=False,
async_safe=False,
zero_copy=False,
otlp_traces=False,
)
def otlp_traces_endpoint(self) -> str:
"""Return the OTLP/HTTP traces endpoint of the store.
The traces can have rollout ID and attempt ID (and optionally sequence ID)
saved in the "resource" of the spans.
The store, if it supports OTLP, should be able to receive the traces and save them
via [`add_span`][agentlightning.LightningStore.add_span] or
[`add_otel_span`][agentlightning.LightningStore.add_otel_span].
The endpoint should be compatible with [OTLP HTTP protocol](https://opentelemetry.io/docs/specs/otlp/).
It's not necessarily compatible with OTLP gRPC protocol.
The returned endpoint will usually ends with `/v1/traces`.
"""
raise NotImplementedError()
async def start_rollout(
self,
input: TaskInput,
+47 -12
View File
@@ -18,6 +18,12 @@ from fastapi import Request, Response
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import (
ExportTraceServiceRequest as PbExportTraceServiceRequest,
)
from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import (
ExportTraceServiceResponse as PbExportTraceServiceResponse,
)
from opentelemetry.sdk.trace import ReadableSpan
from pydantic import BaseModel, Field, TypeAdapter
@@ -35,6 +41,7 @@ from agentlightning.types import (
Worker,
WorkerStatus,
)
from agentlightning.utils.otlp import handle_otlp_export, spans_from_proto
from agentlightning.utils.server_launcher import LaunchMode, PythonServerLauncher, PythonServerLauncherArgs
from .base import UNSET, LightningStore, LightningStoreCapabilities, Unset
@@ -284,6 +291,8 @@ class LightningStoreServer(LightningStore):
and automatically delegate to an HTTP client instead of using the local store.
This ensures one single copy of the store will be shared across all processes.
This server exporting OTLP-compatible traces via the `/v1/traces` endpoint.
Args:
store: The underlying store to delegate operations to.
host: The hostname or IP address to bind the server to.
@@ -323,13 +332,13 @@ class LightningStoreServer(LightningStore):
)
store_capabilities = self.store.capabilities
if not store_capabilities["async_safe"]:
if not store_capabilities.get("async_safe", False):
raise ValueError("The store is not async-safe. Please use another store for the server.")
if self.launcher_args.launch_mode == "mp" and not store_capabilities["zero_copy"]:
if self.launcher_args.launch_mode == "mp" and not store_capabilities.get("zero_copy", False):
raise ValueError(
"The store does not support zero-copy. Please use another store, or use asyncio or thread mode to launch the server."
)
if self.launcher_args.launch_mode == "thread" and not store_capabilities["thread_safe"]:
if self.launcher_args.launch_mode == "thread" and not store_capabilities.get("thread_safe", False):
server_logger.warning(
"The store is not thread-safe. Please be careful when using the store server and the underlying store in different threads."
)
@@ -362,8 +371,13 @@ class LightningStoreServer(LightningStore):
async_safe=True,
thread_safe=True,
zero_copy=True,
otlp_traces=True,
)
def otlp_traces_endpoint(self) -> str:
"""Return the OTLP/HTTP traces endpoint of the store."""
return f"{self.endpoint}/v1/traces"
def __getstate__(self):
"""
Control pickling to prevent server state from being sent to subprocesses.
@@ -770,12 +784,33 @@ class LightningStoreServer(LightningStore):
async def wait_for_rollouts(request: WaitForRolloutsRequest): # pyright: ignore[reportUnusedFunction]
return await self.wait_for_rollouts(rollout_ids=request.rollout_ids, timeout=request.timeout)
# Setup OTLP endpoints
self._setup_otlp(api)
# Mount the API router of /v1/...
self.app.include_router(api)
# Finally, mount the dashboard assets
self._setup_dashboard()
def _setup_otlp(self, api: APIRouter):
"""Setup OTLP endpoints."""
async def _trace_handler(request: PbExportTraceServiceRequest) -> None:
spans = await spans_from_proto(request, self)
server_logger.debug(f"Received {len(spans)} OTLP spans: {', '.join([span.name for span in spans])}")
for span in spans:
await self.add_span(span)
# Reserved methods for OTEL traces
# https://opentelemetry.io/docs/specs/otlp/#otlphttp-request
@api.post("/traces")
async def otlp_traces(): # pyright: ignore[reportUnusedFunction]
return Response(status_code=501)
async def otlp_traces(request: Request): # pyright: ignore[reportUnusedFunction]
return await handle_otlp_export(
request, PbExportTraceServiceRequest, PbExportTraceServiceResponse, _trace_handler, "traces"
)
# Other API endpoints are not supported yet
@api.post("/metrics")
async def otlp_metrics(): # pyright: ignore[reportUnusedFunction]
return Response(status_code=501)
@@ -788,12 +823,6 @@ class LightningStoreServer(LightningStore):
async def otlp_development_profiles(): # pyright: ignore[reportUnusedFunction]
return Response(status_code=501)
# Mount the API router of /v1/...
self.app.include_router(api)
# Finally, mount the dashboard assets
self._setup_dashboard()
def _setup_dashboard(self):
"""Setup the dashboard static files and SPA."""
assert self.app is not None
@@ -1032,7 +1061,8 @@ class LightningStoreClient(LightningStore):
request_timeout: float = 30.0,
connection_timeout: float = 5.0,
):
self.server_address = server_address.rstrip("/") + API_V1_AGL_PREFIX
self.server_address_root = server_address.rstrip("/")
self.server_address = self.server_address_root + API_V1_AGL_PREFIX
self._sessions: Dict[int, aiohttp.ClientSession] = {} # id(loop) -> ClientSession
self._lock = threading.Lock()
@@ -1055,8 +1085,13 @@ class LightningStoreClient(LightningStore):
thread_safe=True,
async_safe=True,
zero_copy=True,
otlp_traces=True,
)
def otlp_traces_endpoint(self) -> str:
"""Return the OTLP/HTTP traces endpoint of the store."""
return f"{self.server_address_root}/v1/traces"
def __getstate__(self):
"""
When LightningStoreClient is pickled (e.g., passed to a subprocess), we only
+2 -1
View File
@@ -293,6 +293,7 @@ class InMemoryLightningStore(LightningStore):
thread_safe=False,
async_safe=True,
zero_copy=False,
otlp_traces=False,
)
@_healthcheck_wrapper
@@ -340,7 +341,7 @@ class InMemoryLightningStore(LightningStore):
self._attempts[rollout.rollout_id] = [attempt]
self._rollouts[rollout.rollout_id] = rollout
# Manully added rollout is not added to task queue. It's already preparing
# Manually added rollout is not added to task queue. It's already preparing
self._completion_events.setdefault(rollout.rollout_id, threading.Event())
return AttemptedRollout(**rollout.model_dump(), attempt=attempt)
+63 -180
View File
@@ -2,25 +2,23 @@
from __future__ import annotations
import asyncio
import logging
import os
import threading
from contextlib import asynccontextmanager, contextmanager
from typing import TYPE_CHECKING, Any, AsyncGenerator, Awaitable, Iterator, List, Optional
from typing import TYPE_CHECKING, Any, AsyncGenerator, Iterator, List, Optional
import agentops
import agentops.sdk.core
import opentelemetry.trace as trace_api
from agentops.sdk.core import TracingCore
from agentops.sdk.processors import SpanProcessor
from opentelemetry.instrumentation.utils import suppress_instrumentation
from opentelemetry.sdk.trace import ReadableSpan
from opentelemetry.sdk.trace import TracerProvider as TracerProviderImpl
from opentelemetry.trace import get_tracer_provider
from opentelemetry.trace.status import StatusCode
from agentlightning.instrumentation import instrument_all, uninstrument_all
from agentlightning.store.base import LightningStore
from .base import Tracer
from .otel import LightningSpanProcessor, OtelTracer
if TYPE_CHECKING:
from agentops.integration.callbacks.langchain import LangchainCallbackHandler
@@ -29,7 +27,7 @@ if TYPE_CHECKING:
logger = logging.getLogger(__name__)
class AgentOpsTracer(Tracer):
class AgentOpsTracer(OtelTracer):
"""Traces agent execution using AgentOps.
This tracer provides functionality to capture execution details using the
@@ -67,9 +65,8 @@ class AgentOpsTracer(Tracer):
def uninstrument(self, worker_id: int):
uninstrument_all()
def init_worker(self, worker_id: int):
super().init_worker(worker_id)
logger.info(f"[Worker {worker_id}] Setting up tracer...") # worker_id included in process name
def _initialize_tracer_provider(self, worker_id: int):
logger.info(f"[Worker {worker_id}] Setting up AgentOps tracer...") # worker_id included in process name
if self.instrument_managed:
self.instrument(worker_id)
@@ -85,16 +82,9 @@ class AgentOpsTracer(Tracer):
self._lightning_span_processor = LightningSpanProcessor()
try:
# new versions
instance = agentops.sdk.core.tracer
# TODO: The span processor cannot be deleted once added.
# This might be a problem if the tracer is entered and exited multiple times.
instance.provider.add_span_processor(self._lightning_span_processor) # type: ignore
except AttributeError:
# old versions
instance = TracingCore.get_instance() # type: ignore
instance._provider.add_span_processor(self._lightning_span_processor) # type: ignore
# TODO: The span processor cannot be deleted once added.
# This might be a problem if the tracer is entered and exited multiple times.
self._get_tracer_provider().add_span_processor(self._lightning_span_processor) # type: ignore
def teardown_worker(self, worker_id: int) -> None:
super().teardown_worker(worker_id)
@@ -111,7 +101,7 @@ class AgentOpsTracer(Tracer):
store: Optional[LightningStore] = None,
rollout_id: Optional[str] = None,
attempt_id: Optional[str] = None,
) -> AsyncGenerator[LightningSpanProcessor, None]:
) -> AsyncGenerator[trace_api.Tracer, None]:
"""
Starts a new tracing context. This should be used as a context manager.
@@ -122,12 +112,10 @@ class AgentOpsTracer(Tracer):
attempt_id: Optional attempt ID to add the spans to.
Yields:
The [`LightningSpanProcessor`][agentlightning.tracer.agentops.LightningSpanProcessor] instance to collect spans.
The OpenTelemetry tracer instance to collect spans.
"""
with self._trace_context_sync(
name=name, store=store, rollout_id=rollout_id, attempt_id=attempt_id
) as processor:
yield processor
with self._trace_context_sync(name=name, store=store, rollout_id=rollout_id, attempt_id=attempt_id) as tracer:
yield tracer
@contextmanager
def _trace_context_sync(
@@ -137,47 +125,51 @@ class AgentOpsTracer(Tracer):
store: Optional[LightningStore] = None,
rollout_id: Optional[str] = None,
attempt_id: Optional[str] = None,
) -> Iterator[LightningSpanProcessor]:
) -> Iterator[trace_api.Tracer]:
"""Implementation of `trace_context` for synchronous execution."""
if not self._lightning_span_processor:
raise RuntimeError("LightningSpanProcessor is not initialized. Call init_worker() first.")
tracer_provider = self._get_tracer_provider()
kwargs: dict[str, Any] = {}
if name is not None:
kwargs["trace_name"] = name
elif rollout_id is not None:
kwargs["trace_name"] = rollout_id
if store is not None and rollout_id is not None and attempt_id is not None:
if store.capabilities.get("otlp_traces", False) is True:
logger.debug(f"Tracing to LightningStore rollout_id={rollout_id}, attempt_id={attempt_id}")
self._enable_native_otlp_exporter(store, rollout_id, attempt_id)
else:
self._disable_native_otlp_exporter()
ctx = self._lightning_span_processor.with_context(store=store, rollout_id=rollout_id, attempt_id=attempt_id)
with ctx:
# AgentOps end_trace and start_trace must live inside the lightning span processor context.
# Otherwise some traces might not be recorded.
with self._agentops_trace_context(rollout_id, attempt_id, kwargs):
yield trace_api.get_tracer(__name__, tracer_provider=tracer_provider)
elif store is None and rollout_id is None and attempt_id is None:
# TODO: Add tests to cover both paths
self._disable_native_otlp_exporter()
with self._lightning_span_processor:
with self._agentops_trace_context(None, None, kwargs):
yield trace_api.get_tracer(__name__, tracer_provider=tracer_provider)
else:
raise ValueError("store, rollout_id, and attempt_id must be either all provided or all None")
@contextmanager
def _agentops_trace_context(self, rollout_id: Optional[str], attempt_id: Optional[str], kwargs: dict[str, Any]):
trace = agentops.start_trace(**kwargs)
status = StatusCode.OK # type: ignore
try:
if store is not None and rollout_id is not None and attempt_id is not None:
ctx = self._lightning_span_processor.with_context(
store=store, rollout_id=rollout_id, attempt_id=attempt_id
)
with ctx as processor:
yield processor
elif store is None and rollout_id is None and attempt_id is None:
with self._lightning_span_processor:
yield self._lightning_span_processor
else:
raise ValueError("store, rollout_id, and attempt_id must be either all provided or all None")
yield
except Exception as e:
# TODO: I'm not sure whether this will catch errors in user code.
status = StatusCode.ERROR # type: ignore
logger.error(f"Trace failed for rollout_id={rollout_id}, attempt_id={attempt_id}, error={e}")
logger.error(f"Trace failed for rollout_id={rollout_id}, attempt_id={attempt_id}: {e}")
finally:
agentops.end_trace(trace, end_state=status) # type: ignore
def get_last_trace(self) -> List[ReadableSpan]:
"""
Retrieves the raw list of captured spans from the most recent trace.
Returns:
A list of OpenTelemetry `ReadableSpan` objects.
"""
if not self._lightning_span_processor:
raise RuntimeError("LightningSpanProcessor is not initialized. Call init_worker() first.")
return self._lightning_span_processor.spans()
def get_langchain_handler(self, tags: List[str] | None = None) -> LangchainCallbackHandler:
"""
Get the Langchain callback handler for integrating with Langchain.
@@ -204,135 +196,26 @@ class AgentOpsTracer(Tracer):
get_langchain_callback_handler = get_langchain_handler # alias
def _get_tracer_provider(self) -> TracerProviderImpl:
try:
# new versions
instance = agentops.sdk.core.tracer
if instance.provider is None:
raise RuntimeError("AgentOps TracerProvider is not initialized.")
class LightningSpanProcessor(SpanProcessor):
"""Span processor that subclasses OpenTelemetry's `SpanProcessor` and adds support to dump traces
to a [`LightningStore`][agentlightning.LightningStore].
"""
if get_tracer_provider() is not instance.provider:
logger.error(
"Mismatch between global singleton TracerProvider and AgentOps TracerProvider. "
"AgentOps might not work properly."
)
def __init__(self):
self._spans: List[ReadableSpan] = []
if not isinstance(instance.provider, TracerProviderImpl): # type: ignore
raise RuntimeError("Unsupported TracerProvider type for AgentOps instrumentation.")
# Store related context and states
self._store: Optional[LightningStore] = None
self._rollout_id: Optional[str] = None
self._attempt_id: Optional[str] = None
self._lock = threading.Lock()
# private asyncio loop running in a daemon thread
self._loop_ready = threading.Event()
self._loop: Optional[asyncio.AbstractEventLoop] = None
self._loop_thread = threading.Thread(target=self._loop_runner, name="otel-loop", daemon=True)
self._loop_thread.start()
self._loop_ready.wait() # loop is ready
def _loop_runner(self):
loop = asyncio.new_event_loop()
self._loop = loop
asyncio.set_event_loop(loop)
self._loop_ready.set()
loop.run_forever()
loop.close()
def __enter__(self):
self._last_trace = None
self._spans = []
return self
def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any):
self._store = None
self._rollout_id = None
self._attempt_id = None
def _await_in_loop(self, coro: Awaitable[Any], timeout: Optional[float] = None) -> Any:
# submit to the dedicated loop and wait synchronously
if self._loop is None:
raise RuntimeError("Loop is not initialized. This should not happen.")
# If already on the exporter loop thread, schedule and return immediately.
# ---------------------------------------------------------------------------
# WHY THIS CONDITIONAL EXISTS:
# In rare cases, span.end() is triggered from a LangchainCallbackHandler.__del__
# (or another finalizer) while the Python garbage collector is running on the
# *same thread* that owns our exporter event loop ("otel-loop").
#
# When that happens, on_end() executes on the exporter loop thread itself.
# If we were to call `asyncio.run_coroutine_threadsafe(...).result()` here,
# it would deadlock immediately — because the loop cannot both wait on and run
# the same coroutine. The Future stays pending forever and the loop stops
# processing scheduled callbacks.
#
# To avoid that self-deadlock, we detect when on_end() runs on the exporter
# loop thread. If so, we *schedule* the coroutine on the loop (fire-and-forget)
# instead of blocking with .result().
#
# This situation can occur because Python calls __del__ in whatever thread
# releases the last reference, which can easily be our loop thread if the
# object is dereferenced during loop._run_once().
# ---------------------------------------------------------------------------
if threading.current_thread() is self._loop_thread:
self._loop.call_soon_threadsafe(asyncio.create_task, coro) # type: ignore
return None
fut = asyncio.run_coroutine_threadsafe(coro, self._loop) # type: ignore
return fut.result(timeout=timeout) # raises on error # type: ignore
def shutdown(self) -> None:
if self._loop:
self._loop.call_soon_threadsafe(self._loop.stop)
self._loop_thread.join(timeout=5)
self._loop = None
def force_flush(self, timeout_millis: int = 30000) -> bool:
return True
def spans(self) -> List[ReadableSpan]:
"""
Get the list of spans collected by this processor.
This is useful for debugging and testing purposes.
Returns:
List of ReadableSpan objects collected during tracing.
"""
return self._spans
def with_context(self, store: LightningStore, rollout_id: str, attempt_id: str):
# simple context manager without nesting into asyncio
class _Ctx:
def __enter__(_): # type: ignore
with self._lock:
self._store, self._rollout_id, self._attempt_id = store, rollout_id, attempt_id
self._last_trace = None
self._spans = []
return self
def __exit__(_, exc_type, exc, tb): # type: ignore
with self._lock:
self._store = self._rollout_id = self._attempt_id = None
return _Ctx()
def on_end(self, span: ReadableSpan) -> None:
"""
Process a span when it ends.
Args:
span: The span that has ended.
"""
# Skip if span is not sampled
if not span.context or not span.context.trace_flags.sampled:
return
if self._store and self._rollout_id and self._attempt_id:
try:
# Submit add_otel_span to the event loop and wait for it to complete
with suppress_instrumentation():
self._await_in_loop(
self._store.add_otel_span(self._rollout_id, self._attempt_id, span),
timeout=60.0,
)
except Exception:
# log; on_end MUST NOT raise
logger.exception(f"Error adding span to store: {span.name}")
self._spans.append(span)
self._tracer_provider = instance.provider
return self._tracer_provider
except AttributeError:
# old versions
instance = TracingCore.get_instance() # type: ignore
self._tracer_provider = instance._provider # type: ignore
return self._tracer_provider # type: ignore
+25
View File
@@ -3,6 +3,7 @@
from __future__ import annotations
import logging
from contextlib import contextmanager
from typing import TYPE_CHECKING, Any, AsyncContextManager, Awaitable, Callable, ContextManager, List, Optional
from opentelemetry.sdk.trace import ReadableSpan
@@ -138,3 +139,27 @@ class Tracer(ParallelWorkerBase):
"""
logger.warning(f"{self.__class__.__name__} does not provide a LangChain callback handler.")
return None
@contextmanager
def lifespan(self):
"""A context manager to manage the lifespan of the tracer.
This can be used to set up and tear down any necessary resources
for the tracer, useful for debugging purposes.
"""
has_init = False
has_init_worker = False
try:
self.init()
has_init = True
self.init_worker(0)
has_init_worker = True
yield
finally:
if has_init_worker:
self.teardown_worker(0)
if has_init:
self.teardown()
+255 -11
View File
@@ -2,16 +2,25 @@
from __future__ import annotations
import asyncio
import logging
import threading
from contextlib import asynccontextmanager
from typing import AsyncGenerator, List, Optional
from typing import Any, AsyncGenerator, Awaitable, List, Optional
import opentelemetry.trace as trace_api
from opentelemetry.sdk.trace import ReadableSpan, TracerProvider
from agentops.sdk.core import BatchSpanProcessor
from opentelemetry.instrumentation.utils import suppress_instrumentation
from opentelemetry.sdk.resources import Resource
from opentelemetry.sdk.trace import ReadableSpan, SpanProcessor
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace import TracerProvider as TracerProviderImpl
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from agentlightning.store.base import LightningStore
from agentlightning.types.tracer import SpanNames
from agentlightning.utils.otlp import LightningStoreOTLPExporter
from .agentops import LightningSpanProcessor # FIXME: This import should be from otel to agentops
from .base import Tracer
logger = logging.getLogger(__name__)
@@ -29,21 +38,31 @@ class OtelTracer(Tracer):
# This provider is only initialized when the worker is initialized.
self._tracer_provider: Optional[TracerProvider] = None
self._lightning_span_processor: Optional[LightningSpanProcessor] = None
self._simple_span_processor: Optional[SimpleSpanProcessor] = None
self._otlp_span_exporter: Optional[LightningStoreOTLPExporter] = None
self._initialized: bool = False
def init_worker(self, worker_id: int):
super().init_worker(worker_id)
self._initialize_tracer_provider(worker_id)
def _initialize_tracer_provider(self, worker_id: int):
logger.info(f"[Worker {worker_id}] Setting up OpenTelemetry tracer...")
if self._initialized:
logger.error("Tracer provider is already initialized. OpenTelemetry may not work as expected.")
tracer_provider = TracerProvider()
trace_api.set_tracer_provider(tracer_provider)
self._tracer_provider = TracerProvider()
trace_api.set_tracer_provider(self._tracer_provider)
self._lightning_span_processor = LightningSpanProcessor()
tracer_provider.add_span_processor(self._lightning_span_processor)
self._tracer_provider.add_span_processor(self._lightning_span_processor)
self._otlp_span_exporter = LightningStoreOTLPExporter()
self._simple_span_processor = SimpleSpanProcessor(self._otlp_span_exporter)
self._tracer_provider.add_span_processor(self._simple_span_processor)
self._initialized = True
logger.info(f"[Worker {worker_id}] OpenTelemetry tracer provider initialized.")
def teardown_worker(self, worker_id: int):
super().teardown_worker(worker_id)
logger.info(f"[Worker {worker_id}] Tearing down OpenTelemetry tracer...")
@@ -57,7 +76,7 @@ class OtelTracer(Tracer):
store: Optional[LightningStore] = None,
rollout_id: Optional[str] = None,
attempt_id: Optional[str] = None,
) -> AsyncGenerator[LightningSpanProcessor, None]:
) -> AsyncGenerator[trace_api.Tracer, None]:
"""
Starts a new tracing context. This should be used as a context manager.
@@ -68,18 +87,24 @@ class OtelTracer(Tracer):
attempt_id: Optional attempt ID to add the spans to.
Yields:
The LightningSpanProcessor instance to collect spans.
The OpenTelemetry tracer instance to collect spans.
"""
if not self._lightning_span_processor:
raise RuntimeError("LightningSpanProcessor is not initialized. Call init_worker() first.")
if store is not None and rollout_id is not None and attempt_id is not None:
if store.capabilities.get("otlp_traces", False) is True:
logger.debug(f"Tracing to LightningStore rollout_id={rollout_id}, attempt_id={attempt_id}")
self._enable_native_otlp_exporter(store, rollout_id, attempt_id)
else:
self._disable_native_otlp_exporter()
ctx = self._lightning_span_processor.with_context(store=store, rollout_id=rollout_id, attempt_id=attempt_id)
with ctx as processor:
yield processor
with ctx:
yield trace_api.get_tracer(__name__, tracer_provider=self._tracer_provider)
elif store is None and rollout_id is None and attempt_id is None:
self._disable_native_otlp_exporter()
with self._lightning_span_processor:
yield self._lightning_span_processor
yield trace_api.get_tracer(__name__, tracer_provider=self._tracer_provider)
else:
raise ValueError("store, rollout_id, and attempt_id must be either all provided or all None")
@@ -93,3 +118,222 @@ class OtelTracer(Tracer):
if not self._lightning_span_processor:
raise RuntimeError("LightningSpanProcessor is not initialized. Call init_worker() first.")
return self._lightning_span_processor.spans()
def _get_tracer_provider(self) -> TracerProviderImpl:
if self._tracer_provider is None:
raise RuntimeError("TracerProvider is not initialized. Call init_worker() first.")
return self._tracer_provider
def _enable_native_otlp_exporter(self, store: LightningStore, rollout_id: str, attempt_id: str):
tracer_provider = self._get_tracer_provider()
active_span_processor = tracer_provider._active_span_processor # pyright: ignore[reportPrivateUsage]
# Override the resources so that the server knows where the request comes from.
tracer_provider._resource = tracer_provider._resource.merge( # pyright: ignore[reportPrivateUsage]
Resource.create(
{
SpanNames.ROLLOUT_ID: rollout_id,
SpanNames.ATTEMPT_ID: attempt_id,
}
)
)
instrumented = False
candidates: List[str] = []
for processor in active_span_processor._span_processors: # pyright: ignore[reportPrivateUsage]
if isinstance(processor, LightningSpanProcessor):
# We don't need the LightningSpanProcessor any more.
logger.debug("LightningSpanProcessor already present in TracerProvider, disabling it.")
processor.disable_store_submission = True
elif isinstance(processor, (SimpleSpanProcessor, BatchSpanProcessor)):
# Instead, we rely on the OTLPSpanExporter to send spans to the store.
if isinstance(processor.span_exporter, LightningStoreOTLPExporter):
processor.span_exporter.enable_store_otlp(store.otlp_traces_endpoint(), rollout_id, attempt_id)
logger.debug(f"Set LightningStoreOTLPExporter endpoint to {store.otlp_traces_endpoint()}")
instrumented = True
else:
candidates.append(
f"{processor.__class__.__name__} with {processor.span_exporter.__class__.__name__}"
)
else:
candidates.append(f"{processor.__class__.__name__}")
if not instrumented:
raise RuntimeError(
"Failed to enable native OTLP exporter: no BatchSpanProcessor or SimpleSpanProcessor with "
"LightningStoreOTLPExporter found in TracerProvider. Please try using a non-OTLP store."
"Candidates are: " + ", ".join(candidates)
)
def _disable_native_otlp_exporter(self):
tracer_provider = self._get_tracer_provider()
active_span_processor = tracer_provider._active_span_processor # pyright: ignore[reportPrivateUsage]
tracer_provider._resource = tracer_provider._resource.merge( # pyright: ignore[reportPrivateUsage]
Resource.create(
{
SpanNames.ROLLOUT_ID: "",
SpanNames.ATTEMPT_ID: "",
}
)
) # reset resource
for processor in active_span_processor._span_processors: # pyright: ignore[reportPrivateUsage]
if isinstance(processor, LightningSpanProcessor):
# We will be in need of the LightningSpanProcessor again.
logger.debug("Enabling LightningSpanProcessor in TracerProvider.")
processor.disable_store_submission = False
class LightningSpanProcessor(SpanProcessor):
"""Span processor that subclasses OpenTelemetry's `SpanProcessor` and adds support to dump traces
to a [`LightningStore`][agentlightning.LightningStore].
It serves two purposes:
1. Records all the spans in a local buffer.
2. Submits the spans to the event loop to be added to the store.
"""
def __init__(self, disable_store_submission: bool = False):
self._disable_store_submission: bool = disable_store_submission
self._spans: List[ReadableSpan] = []
# Store related context and states
self._store: Optional[LightningStore] = None
self._rollout_id: Optional[str] = None
self._attempt_id: Optional[str] = None
self._lock = threading.Lock()
# private asyncio loop running in a daemon thread
self._loop_ready = threading.Event()
self._loop: Optional[asyncio.AbstractEventLoop] = None
self._loop_thread: Optional[threading.Thread] = None
@property
def disable_store_submission(self) -> bool:
"""Whether to disable submitting spans to the store."""
return self._disable_store_submission
@disable_store_submission.setter
def disable_store_submission(self, value: bool) -> None:
self._disable_store_submission = value
def _ensure_loop(self) -> None:
if self._loop_thread is None or self._loop is None:
self._loop_ready.clear()
self._loop_thread = threading.Thread(target=self._loop_runner, name="otel-loop", daemon=True)
self._loop_thread.start()
self._loop_ready.wait() # loop is ready
def _loop_runner(self):
loop = asyncio.new_event_loop()
self._loop = loop
asyncio.set_event_loop(loop)
self._loop_ready.set()
loop.run_forever()
loop.close()
def __enter__(self):
self._last_trace = None
self._spans = []
return self
def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any):
self._store = None
self._rollout_id = None
self._attempt_id = None
def _await_in_loop(self, coro: Awaitable[Any], timeout: Optional[float] = None) -> Any:
# submit to the dedicated loop and wait synchronously
self._ensure_loop()
if self._loop is None:
raise RuntimeError("Loop is not initialized. This should not happen.")
# If already on the exporter loop thread, schedule and return immediately.
# ---------------------------------------------------------------------------
# WHY THIS CONDITIONAL EXISTS:
# In rare cases, span.end() is triggered from a LangchainCallbackHandler.__del__
# (or another finalizer) while the Python garbage collector is running on the
# *same thread* that owns our exporter event loop ("otel-loop").
#
# When that happens, on_end() executes on the exporter loop thread itself.
# If we were to call `asyncio.run_coroutine_threadsafe(...).result()` here,
# it would deadlock immediately — because the loop cannot both wait on and run
# the same coroutine. The Future stays pending forever and the loop stops
# processing scheduled callbacks.
#
# To avoid that self-deadlock, we detect when on_end() runs on the exporter
# loop thread. If so, we *schedule* the coroutine on the loop (fire-and-forget)
# instead of blocking with .result().
#
# This situation can occur because Python calls __del__ in whatever thread
# releases the last reference, which can easily be our loop thread if the
# object is dereferenced during loop._run_once().
# ---------------------------------------------------------------------------
if threading.current_thread() is self._loop_thread:
self._loop.call_soon_threadsafe(asyncio.create_task, coro) # type: ignore
return None
fut = asyncio.run_coroutine_threadsafe(coro, self._loop) # type: ignore
return fut.result(timeout=timeout) # raises on error # type: ignore
def shutdown(self) -> None:
if self._loop:
self._loop.call_soon_threadsafe(self._loop.stop)
self._loop = None
if self._loop_thread:
self._loop_thread.join(timeout=5)
def force_flush(self, timeout_millis: int = 30000) -> bool:
return True
def spans(self) -> List[ReadableSpan]:
"""
Get the list of spans collected by this processor.
This is useful for debugging and testing purposes.
Returns:
List of ReadableSpan objects collected during tracing.
"""
return self._spans
def with_context(self, store: LightningStore, rollout_id: str, attempt_id: str):
# simple context manager without nesting into asyncio
class _Ctx:
def __enter__(_): # type: ignore
# Use _ instead of self to avoid shadowing the instance method.
with self._lock:
self._store, self._rollout_id, self._attempt_id = store, rollout_id, attempt_id
self._last_trace = None
self._spans = []
return self
def __exit__(_, exc_type, exc, tb): # type: ignore
with self._lock:
self._store = self._rollout_id = self._attempt_id = None
return _Ctx()
def on_end(self, span: ReadableSpan) -> None:
"""
Process a span when it ends.
Args:
span: The span that has ended.
"""
# Skip if span is not sampled
if not span.context or not span.context.trace_flags.sampled:
return
if not self._disable_store_submission and self._store and self._rollout_id and self._attempt_id:
try:
# Submit add_otel_span to the event loop and wait for it to complete
with suppress_instrumentation():
self._ensure_loop()
self._await_in_loop(
self._store.add_otel_span(self._rollout_id, self._attempt_id, span),
timeout=60.0,
)
except Exception:
# log; on_end MUST NOT raise
logger.exception(f"Error adding span to store: {span.name}")
self._spans.append(span)
+6
View File
@@ -411,6 +411,12 @@ class SpanNames(str, Enum):
"""The name of the exception span."""
VIRTUAL = "agentlightning.virtual"
"""The name of the virtual span. It represents derived spans without concrete operations."""
ROLLOUT_ID = "agentlightning.rollout_id"
"""The name of the rollout ID."""
ATTEMPT_ID = "agentlightning.attempt_id"
"""The name of the attempt ID."""
SPAN_SEQUENCE_ID = "agentlightning.span_sequence_id"
"""The name of the span sequence ID."""
class SpanAttributeNames(str, Enum):
+428
View File
@@ -0,0 +1,428 @@
# Copyright (c) Microsoft. All rights reserved.
import gzip
import logging
from typing import Any, Awaitable, Callable, Dict, List, Optional, Sequence, Tuple, Type, TypeVar
from fastapi import Request, Response
from google.protobuf import json_format
from google.rpc.status_pb2 import Status
from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
from opentelemetry.proto.collector.logs.v1.logs_service_pb2 import (
ExportLogsServiceRequest,
ExportLogsServiceResponse,
)
from opentelemetry.proto.collector.metrics.v1.metrics_service_pb2 import (
ExportMetricsServiceRequest,
ExportMetricsServiceResponse,
)
from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import (
ExportTraceServiceRequest,
ExportTraceServiceResponse,
)
from opentelemetry.proto.common.v1.common_pb2 import AnyValue, KeyValue
from opentelemetry.proto.resource.v1.resource_pb2 import Resource as ProtoResource
from opentelemetry.proto.trace.v1.trace_pb2 import Span as ProtoSpan
from opentelemetry.proto.trace.v1.trace_pb2 import Status as ProtoStatus
from opentelemetry.sdk.resources import Resource
from opentelemetry.sdk.trace import ReadableSpan
from opentelemetry.sdk.trace.export import SpanExportResult
from opentelemetry.util.types import AttributeValue
from agentlightning.store.base import LightningStore
from agentlightning.types.tracer import (
Attributes,
Event,
Link,
OtelResource,
Span,
SpanContext,
SpanNames,
TraceStatus,
convert_timestamp,
)
PROTOBUF_CT = "application/x-protobuf"
logger = logging.getLogger(__name__)
T_request = TypeVar("T_request", ExportLogsServiceRequest, ExportMetricsServiceRequest, ExportTraceServiceRequest)
T_response = TypeVar("T_response", ExportLogsServiceResponse, ExportMetricsServiceResponse, ExportTraceServiceResponse)
async def handle_otlp_export(
request: Request,
request_message_cls: Type[T_request],
response_message_cls: Type[T_response],
message_callback: Optional[Callable[[T_request], Awaitable[None]]],
signal_name: str,
) -> Response:
"""
Generic handler for /v1/traces, /v1/metrics, /v1/logs.
Convert the OTLP Protobuf request to a JSON-like object.
"""
content_type = request.headers.get("Content-Type", "").split(";")[0].strip()
if content_type != PROTOBUF_CT:
# For brevity we only support binary protobuf here.
return _bad_request_response(
request,
f"Unsupported Content-Type '{content_type}', expected '{PROTOBUF_CT}'",
content_type=PROTOBUF_CT,
)
raw_body = await request.body()
body = _read_body_maybe_gzip(request, raw_body)
# Empty request is allowed and should still succeed.
if not body:
req_msg = request_message_cls()
else:
req_msg = request_message_cls()
try:
req_msg.ParseFromString(body)
except Exception as exc:
return _bad_request_response(request, f"Unable to parse OTLP {signal_name} payload: {exc}")
if message_callback is not None:
await message_callback(req_msg)
# Build success response. Partial success field is left unset.
resp_msg = response_message_cls()
# Encode response in the same Content-Type as request.
if content_type == PROTOBUF_CT:
resp_bytes = resp_msg.SerializeToString()
else:
resp_bytes = json_format.MessageToJson(resp_msg).encode("utf-8")
resp_bytes, headers = _maybe_gzip_response(request, resp_bytes)
return Response(
content=resp_bytes,
media_type=content_type,
status_code=200,
headers=headers,
)
async def spans_from_proto(request: ExportTraceServiceRequest, store: LightningStore) -> List[Span]:
"""Parse an OTLP proto payload into List[Span].
A store is needed here for generating a sequence ID for each span.
"""
output_spans: List[Span] = []
for resource_spans in request.resource_spans:
# Resource-level attributes & IDs
resource_attrs = _kv_list_to_dict(resource_spans.resource.attributes)
# rollout_id, attempt_id from resource attributes when present.
rollout_id_resource = resource_attrs.get(SpanNames.ROLLOUT_ID)
attempt_id_resource = resource_attrs.get(SpanNames.ATTEMPT_ID)
# If sequence id is provided, all the spans will share the same sequence ID.
# unless otherwise overridden by span-level attributes.
sequence_id_resource = resource_attrs.get(SpanNames.SPAN_SEQUENCE_ID)
otel_resource = _resource_from_proto(resource_spans.resource, getattr(resource_spans, "schema_url", ""))
# Each ScopeSpans contains multiple spans
for scope_spans in resource_spans.scope_spans:
for proto_span in scope_spans.spans:
trace_id_hex = _bytes_to_trace_id_hex(proto_span.trace_id)
span_id_hex = _bytes_to_span_id_hex(proto_span.span_id)
parent_id_hex = _bytes_to_span_id_hex(proto_span.parent_span_id) if proto_span.parent_span_id else None
# Status
status_code_str = _STATUS_CODE_MAP.get(proto_span.status.code, "UNSET")
status = TraceStatus(
status_code=status_code_str,
description=proto_span.status.message or None,
)
# Attributes
span_attrs = _kv_list_to_dict(proto_span.attributes)
# Context
context = SpanContext(
trace_id=trace_id_hex,
span_id=span_id_hex,
is_remote=False,
trace_state={},
)
# Try to get if span attributes contain something like rollout_id or attempt_id
# Override the resource-level attributes with the span-level attributes if present.
rollout_id_span = span_attrs.get(SpanNames.ROLLOUT_ID)
attempt_id_span = span_attrs.get(SpanNames.ATTEMPT_ID)
sequence_id_span = span_attrs.get(SpanNames.SPAN_SEQUENCE_ID)
# Normalize to regular strings and ints
rollout_id_raw = rollout_id_span if rollout_id_span is not None else rollout_id_resource
attempt_id_raw = attempt_id_span if attempt_id_span is not None else attempt_id_resource
sequence_id_raw = sequence_id_span if sequence_id_span is not None else sequence_id_resource
rollout_id, attempt_id = _normalize_rollout_attempt_id(rollout_id_raw, attempt_id_raw)
sequence_id = _normalize_sequence_id(sequence_id_raw)
if rollout_id is None or attempt_id is None:
logger.warning(
"Both rollout_id and attempt_id must be present in resource attributes. "
"Spans will not be able to log to the store because of missing IDs: rollout_id=%s, attempt_id=%s, sequence_id=%s",
rollout_id,
attempt_id,
sequence_id,
)
continue
# Generate a new sequence ID if not provided
if sequence_id is None:
current_sequence_id = await store.get_next_span_sequence_id(
rollout_id=rollout_id, attempt_id=attempt_id
)
else:
current_sequence_id = sequence_id
# Build Span
span = Span(
rollout_id=rollout_id,
attempt_id=attempt_id,
sequence_id=current_sequence_id,
trace_id=trace_id_hex,
span_id=span_id_hex,
parent_id=parent_id_hex,
name=proto_span.name,
status=status,
attributes=span_attrs,
events=_events_from_proto(proto_span),
links=_links_from_proto(proto_span),
start_time=convert_timestamp(proto_span.start_time_unix_nano),
end_time=convert_timestamp(proto_span.end_time_unix_nano),
context=context,
parent=None, # OTLP only has parent_span_id; we don't have full SpanContext
resource=otel_resource,
)
output_spans.append(span)
return output_spans
class LightningStoreOTLPExporter(OTLPSpanExporter):
"""OTLP Exporter that write to a LightningStore-compatible backend.
The backend requires two special attributes on each span:
- `agentlightning.rollout_id`: The rollout ID to associate the span with.
- `agentlightning.attempt_id`: The attempt ID to associate the span with.
It can optionally use the following attribute to sequence spans:
- `agentlightning.span_sequence_id`: A decimal string representing the sequence ID of the span.
"""
_default_endpoint: Optional[str] = None
_rollout_id: Optional[str] = None
_attempt_id: Optional[str] = None
def enable_store_otlp(self, endpoint: str, rollout_id: str, attempt_id: str) -> None:
"""Enable storing OTLP data to a specific LightningStore rollout/attempt."""
self._rollout_id = rollout_id
self._attempt_id = attempt_id
self._default_endpoint = self._endpoint
self._endpoint = endpoint
def disable_store_otlp(self) -> None:
"""Disable storing OTLP data to LightningStore."""
self._rollout_id = None
self._attempt_id = None
if self._default_endpoint is not None:
self._endpoint = self._default_endpoint
def should_bypass(self) -> bool:
"""Check if the exporter should bypass the default export if rollout_id and attempt_id are not set."""
return True
def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
if self._rollout_id is not None and self._attempt_id is not None:
# rollout_id and attempt_id are present in resource attributes
# It means that the server supports OTLP endpoint.
for span in spans:
# Override the resources so that the server knows where the request comes from.
span._resource = span._resource.merge( # pyright: ignore[reportPrivateUsage]
Resource.create(
{
SpanNames.ROLLOUT_ID: self._rollout_id,
SpanNames.ATTEMPT_ID: self._attempt_id,
}
)
)
return super().export(spans)
elif not self.should_bypass():
logger.debug("Rollout ID and Attempt ID not set; using default OTLP exporter behavior.")
return super().export(spans)
else:
logger.debug("Rollout ID and Attempt ID not set; bypassing export.")
return SpanExportResult.SUCCESS
def _read_body_maybe_gzip(request: Request, raw_body: bytes) -> bytes:
"""
Decompress body if Content-Encoding: gzip; otherwise return as is.
"""
encoding = request.headers.get("Content-Encoding", "").lower()
if encoding == "gzip":
return gzip.decompress(raw_body)
return raw_body
def _maybe_gzip_response(request: Request, payload: bytes) -> Tuple[bytes, Dict[str, str]]:
"""
If Accept-Encoding includes gzip, gzip the payload and set Content-Encoding header.
"""
ae = request.headers.get("Accept-Encoding", "")
tokens = [token.split(";")[0].strip().lower() for token in ae.split(",") if token.strip()]
headers: Dict[str, str] = {}
if "gzip" in tokens:
payload = gzip.compress(payload)
headers["Content-Encoding"] = "gzip"
return payload, headers
def _bad_request_response(request: Request, message: str, content_type: str = PROTOBUF_CT) -> Response:
"""
Build a 400 response whose body is a protobuf Status message, encoded
in the same Content-Type as the request (OTLP/HTTP requirement).
"""
status_msg = Status(message=message)
if content_type == PROTOBUF_CT:
body = status_msg.SerializeToString()
else:
# Fallback: JSON representation of Status.
body = json_format.MessageToJson(status_msg).encode("utf-8")
body, headers = _maybe_gzip_response(request, body)
return Response(
content=body,
status_code=400,
media_type=content_type,
headers=headers,
)
def _normalize_rollout_attempt_id(
rollout_id: Optional[AttributeValue], attempt_id: Optional[AttributeValue]
) -> Tuple[Optional[str], Optional[str]]:
"""Normalize a rollout or attempt ID to a string."""
rollout_id_str = str(rollout_id) if rollout_id is not None else None
attempt_id_str = str(attempt_id) if attempt_id is not None else None
return rollout_id_str, attempt_id_str
def _normalize_sequence_id(sequence_id: Optional[AttributeValue]) -> Optional[int]:
"""Normalize a sequence ID to an integer."""
if sequence_id is None:
return None
try:
sequence_id_int = int(str(sequence_id))
except (ValueError, TypeError):
logger.warning(
"Invalid sequence_id value in resource attributes: %r. Must be an integer or string representing an integer. Assuming None.",
sequence_id,
)
sequence_id_int = None
return sequence_id_int
def _any_value_to_python(value: AnyValue) -> Any:
"""Convert OTLP AnyValue -> plain Python value."""
kind = value.WhichOneof("value")
if kind is None:
return None
if kind == "string_value":
return value.string_value
if kind == "bool_value":
return value.bool_value
if kind == "int_value":
return int(value.int_value)
if kind == "double_value":
return float(value.double_value)
if kind == "array_value":
return [_any_value_to_python(v) for v in value.array_value.values]
if kind == "kvlist_value":
# Map<string, AnyValue> -> dict
return {kv.key: _any_value_to_python(kv.value) for kv in value.kvlist_value.values}
if kind == "bytes_value":
# Serialize bytes as hex string to stay JSON-friendly
return value.bytes_value.hex()
return None
def _kv_list_to_dict(kvs: Sequence[KeyValue]) -> Attributes:
"""Convert repeated KeyValue -> Attributes dict."""
return {kv.key: _any_value_to_python(kv.value) for kv in kvs}
_STATUS_CODE_MAP = {
ProtoStatus.STATUS_CODE_UNSET: "UNSET",
ProtoStatus.STATUS_CODE_OK: "OK",
ProtoStatus.STATUS_CODE_ERROR: "ERROR",
}
def _bytes_to_trace_id_hex(b: bytes) -> str:
# OTLP uses 16-byte trace IDs; format as 32-char hex
if not b:
return "0" * 32
return b.hex().rjust(32, "0")
def _bytes_to_span_id_hex(b: bytes) -> str:
# OTLP uses 8-byte span IDs; format as 16-char hex
if not b:
return "0" * 16
return b.hex().rjust(16, "0")
def _events_from_proto(span: ProtoSpan) -> List[Event]:
"""Event converter from OTLP ProtoSpan to List[Event]."""
return [
Event(
name=e.name,
attributes=_kv_list_to_dict(e.attributes),
timestamp=convert_timestamp(e.time_unix_nano),
)
for e in span.events
]
def _links_from_proto(span: ProtoSpan) -> List[Link]:
"""Link converter from OTLP ProtoSpan to List[Link]."""
links: List[Link] = []
for link in span.links:
trace_id_hex = _bytes_to_trace_id_hex(link.trace_id)
span_id_hex = _bytes_to_span_id_hex(link.span_id)
ctx = SpanContext(
trace_id=trace_id_hex,
span_id=span_id_hex,
is_remote=False,
trace_state={}, # OTLP trace_state is currently a string; you can parse if needed
)
links.append(
Link(
context=ctx,
attributes=_kv_list_to_dict(link.attributes) or None,
)
)
return links
def _resource_from_proto(resource: ProtoResource, schema_url: str = "") -> OtelResource:
return OtelResource(
attributes=_kv_list_to_dict(resource.attributes),
schema_url=schema_url or "",
)
+9
View File
@@ -152,6 +152,13 @@ Programmatically this is encapsulated by [`Span.from_opentelemetry(readable_span
[`add_span`][agentlightning.LightningStore.add_span] or [`add_otel_span`][agentlightning.LightningStore.add_otel_span] both appends a span *and* acts as a heartbeat that can revive `unresponsive``running`.
## OTLP Compatibility
Some of the LightningStore implementations support exporting traces via the [OTLP/HTTP specification](https://opentelemetry.io/docs/specs/otlp/). For example, [`LightningStoreServer`][agentlightning.LightningStoreServer] exposes `/v1/traces` endpoint, it implements the binary Protobuf variant defined by the spec, including the required `Content-Type: application/x-protobuf`, optional `Content-Encoding: gzip`, and status responses encoded as `google.rpc.Status`. Agent-lightning helps parsing `ExportTraceServiceRequest` messages, validate identifiers, normalize resource metadata, and allocate sequence
numbers so store implementations only need to persist [`Span`][agentlightning.Span] objects in order.
Because the interface speaks standard OTLP, any OpenTelemetry-compatible SDK or collector can emit spans directly to a LightningStore OTLP endpoint without custom shims. The server responds according to the OTLP contract (status code, encoding, and error payloads), which keeps Agent-lightning interoperable with existing observability tooling. This compatibility serves as a strong complement to the OpenTelemetry conversion discussed above.
## Store Implementations
Currently, the only out-of-the-box implementation is [`InMemoryLightningStore`][agentlightning.InMemoryLightningStore]:
@@ -162,6 +169,8 @@ Currently, the only out-of-the-box implementation is [`InMemoryLightningStore`][
For production you will likely want persistence. Were actively building a SQLite-backed store that keeps the same API surface while adding durability, crash recovery, and better historical span queries. If you need something sooner, implement your own store by subclassing [`LightningStore`][agentlightning.LightningStore] and providing concrete storage for the small set of abstract methods (`enqueue_rollout`, `dequeue_rollout`, `update_attempt`, `add_span`, etc.). This document plus the tests in `tests/store/` illustrate the expected behavior.
Different store implementations may have different capabilities. For example, [`InMemoryLightningStore`][agentlightning.InMemoryLightningStore] does not support exporting traces via OTLP. Try to distinguish the capabilities of a store implementation by checking the [`capabilities`][agentlightning.LightningStore.capabilities] property.
## Thread Safety
**[`LightningStoreThreaded`][agentlightning.LightningStoreThreaded]** is a subclass of [`LightningStore`][agentlightning.LightningStore] that wraps another underlying store to make a store instance safe for multi-threaded callers. It wraps every state-mutating call in a mutex. Specifically:
+5 -1
View File
@@ -38,7 +38,7 @@
::: agentlightning.store.utils.propagate_status
::: agentlightning.tracer.agentops.LightningSpanProcessor
::: agentlightning.tracer.otel.LightningSpanProcessor
::: agentlightning.utils.server_launcher.PythonServerLauncher
@@ -46,6 +46,10 @@
::: agentlightning.utils.server_launcher.LaunchMode
::: agentlightning.utils.otlp.handle_otlp_export
::: agentlightning.utils.otlp.spans_from_proto
## Deprecated APIs
::: agentlightning.server.AgentLightningServer
+2
View File
@@ -2,6 +2,8 @@
::: agentlightning.LightningStore
::: agentlightning.LightningStoreCapabilities
## Store Implementations
::: agentlightning.InMemoryLightningStore
+2 -1
View File
@@ -7,10 +7,11 @@ This catalog highlights the examples shipped with Agent-lightning.
| [apo](./apo) | Automatic Prompt Optimization tutorials covering built-in, custom, and debugging workflows. | [![apo workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/badge-apo.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/badge-apo.yml) |
| [azure](./azure) | Supervised fine-tuning with Azure OpenAI. | **Unmaintained** — last verified with Agent-lightning v0.2.1 |
| [calc_x](./calc_x) | VERL-powered math reasoning agent training that uses AutoGen with an MCP calculator tool. | [![calc_x workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/badge-calc-x.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/badge-calc-x.yml) |
| [minimal](./minimal) | Bite-sized programs that demonstrate how individual Agent-lightning building blocks behave in isolation. | [![minimal workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/badge-unit.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/badge-unit.yml) |
| [rag](./rag) | Retrieval-Augmented Generation pipeline targeting the MuSiQue dataset with Wikipedia retrieval. | **Unmaintained** — last verified with Agent-lightning v0.1.1 |
| [search_r1](./search_r1) | Framework-free Search-R1 reinforcement learning training workflow with a retrieval backend. | **Unmaintained** — last verified with Agent-lightning v0.1.2 |
| [spider](./spider) | Text-to-SQL reinforcement learning training on the Spider dataset using LangGraph. | [![spider workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/badge-spider.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/badge-spider.yml) |
| [tinker](./tinker) | Reinforcement learning with Tinker as the backend training service. | **Unmaintained** — last verified with Agent-lightning v0.2.2 |
| [unsloth](./unsloth) | Supervised fine-tuning example powered by Unsloth with 4-bit quantization and LoRA. | [![unsloth workflow status](https://github.com/microsoft/agent-lightning/actions/workflows/badge-unsloth.yml/badge.svg)](https://github.com/microsoft/agent-lightning/actions/workflows/badge-unsloth.yml) |
| [tinker](./tinker) | Reinforcement learning with Tinker as the backend training service. | **Unmaintained** — last verified with Agent-lightning v0.2.1 |
*NOTE: CI status avoid taking any workflow running with latest dependencies into account. That's why we reference the corresponding `badge-*` workflows instead. Each example's own README also displays its `examples-*` workflow status whenever the project is maintained by CI.*
+17
View File
@@ -0,0 +1,17 @@
# Minimal Component Showcase
`examples/minimal` provides bite-sized programs that demonstrate how individual Agent-lightning building blocks behave in isolation.
Each module have been documented with its own CLI usage in the module-level docstring. Use this directory as a reference when wiring the same pieces into a larger system.
## Whats Included?
| Component | Demonstrated In | Highlights |
| --- | --- | --- |
| LightningStore + OTLP ingestion | `write_traces.py` | Shows how `OtelTracer` and `AgentOpsTracer` open rollouts, emit spans, and optionally forward them to a remote store client. |
| LLM proxying | `llm_proxy.py` | Guards either OpenAI or a local vLLM deployment with `LLMProxy`, proving how requests are routed through `/rollout/<id>/attempt/<id>` namespaces and captured in the store. |
| vLLM lifecycle | `vllm_server.py` | Minimal context manager that shells out to `vllm serve`, monitors readiness, and tears down the process safely. |
All runtime instructions (CLI arguments, required environment variables, etc.) are embedded directly in each scripts top-level docstring so the source stays self-documenting.
For full-fledged training workflows or multi-component experiments, browse the other subdirectories under `examples/`. This `minimal` folder deliberately keeps each demonstration focused on a single component so you can understand and test them independently.
+254
View File
@@ -0,0 +1,254 @@
# Copyright (c) Microsoft. All rights reserved.
"""Examples to serve an LLM proxy for a vLLM server or an OpenAI service.
Usage: run one of the following commands to start a server.
```bash
python llm_proxy.py vllm Qwen/Qwen2.5-0.5B-Instruct
```
Use the following command to test the LLM proxy.
```bash
python llm_proxy.py test Qwen/Qwen2.5-0.5B-Instruct
```
You can also test the OpenAI Proxy path (`OPENAI_API_KEY` environment variable is required).
```bash
dotenv run python llm_proxy.py openai gpt-4.1-mini
```
"""
import argparse
import asyncio
import os
from typing import List, no_type_check
import aiohttp
from portpicker import pick_unused_port
from rich.console import Console
from vllm_server import vllm_server
import agentlightning as agl
console = Console()
async def serve_llm_proxy_with_vllm(model_name: str, store_port: int = 43887):
"""Serve an LLM proxy for a vLLM server."""
# Create a store to store the traces
store = agl.InMemoryLightningStore()
store_server = agl.LightningStoreServer(store, "127.0.0.1", store_port)
await store_server.start()
# Create a vLLM server
vllm_port = pick_unused_port()
with vllm_server(model_name, vllm_port) as vllm_endpoint:
# Server is up.
# Create an LLM proxy to guard the vLLM server and catch the traces
llm_proxy = agl.LLMProxy(
port=43886,
model_list=[
{
"model_name": model_name,
"litellm_params": {
"model": f"hosted_vllm/{model_name}",
"api_base": vllm_endpoint,
},
}
],
store=store_server,
)
try:
await llm_proxy.start()
# Wait forever
await asyncio.sleep(float("inf"))
finally:
# Stop the LLM proxy and the store server
await llm_proxy.stop()
await store_server.stop()
async def serve_llm_proxy_with_openai(model_name: str, store_port: int = 43887):
"""Serve an LLM proxy for an OpenAI server."""
# Create a store to store the traces
store = agl.InMemoryLightningStore()
store_server = agl.LightningStoreServer(store, "127.0.0.1", store_port)
await store_server.start()
if not os.getenv("OPENAI_API_KEY"):
raise ValueError("OPENAI_API_KEY environment variable is not set")
# Create an LLM proxy to guard the OpenAI server and catch the traces
llm_proxy = agl.LLMProxy(
port=43886,
model_list=[
{
"model_name": model_name,
"litellm_params": {
"model": "openai/" + model_name,
# Must have OpenAI API key set in the environment variable
},
}
],
store=store_server,
callbacks=["opentelemetry"],
)
try:
await llm_proxy.start()
# Wait forever
await asyncio.sleep(float("inf"))
finally:
# Stop the LLM proxy and the store server
await llm_proxy.stop()
await store_server.stop()
async def test_llm_proxy(model_name: str, store_port: int = 43887):
"""Test the LLM proxy by sending a request to the proxy and checking the response.
We do it via aiohttp here. This can also be done with OpenAI client.
"""
# We first connect to the store server and start a rollout.
store = agl.LightningStoreClient(f"http://localhost:{store_port}")
rollout = await store.start_rollout(input={"origin": "test_llm_proxy"})
# The chat completion URL is simply /v1/chat/completions under the namespace of current rollout and attempt.
# This ensures the traces are properly put into the correct bucket.
chat_completion_url = (
f"http://localhost:43886/rollout/{rollout.rollout_id}/attempt/{rollout.attempt.attempt_id}/v1/chat/completions"
)
async with aiohttp.ClientSession() as session:
async with session.post(
chat_completion_url,
json={
"model": model_name,
"messages": [{"role": "user", "content": "Hello, what's your name?"}],
},
) as response:
response_body = await response.json()
console.print("Response body:", response_body)
_verify_response_body(response_body, model_name)
spans = await store.query_spans(rollout_id=rollout.rollout_id, attempt_id=rollout.attempt.attempt_id)
for span in spans:
console.print("Span:", span)
_verify_span(spans)
await store.close()
@no_type_check
def _verify_response_body(response_body: dict, model_name: str):
"""Expect Response body to be something like this:
```python
{
'id': 'chatcmpl-996a90a8678e4ed0a0d2724df2c0bba5',
'created': 1763178218,
'model': 'hosted_vllm/Qwen/Qwen2.5-0.5B-Instruct',
'object': 'chat.completion',
'choices': [
{
'finish_reason': 'stop',
'index': 0,
'message': {
'content': 'Hello! I am Qwen, an AI language model created by Alibaba Cloud. My name is Qwen, and I can assist you with
various tasks and provide information on a wide range of topics. How may I help you today?',
'role': 'assistant'
},
'provider_specific_fields': {
'stop_reason': None,
'token_ids': [9707, 0, ...],
}
}
],
'usage': {'completion_tokens': 48, 'prompt_tokens': 36, 'total_tokens': 84},
'prompt_token_ids': [151644, 8948, ...],
}
```
"""
if "qwen" in model_name.lower():
assert "qwen" in response_body["choices"][0]["message"]["content"].lower()
assert (
"provider_specific_fields" in response_body["choices"][0]
), "provider_specific_fields not found in response body"
assert (
"token_ids" in response_body["choices"][0]["provider_specific_fields"]
), "token_ids not found in response body"
assert "prompt_token_ids" in response_body, "prompt_token_ids not found in response body"
else:
assert "chatgpt" in response_body["choices"][0]["message"]["content"].lower()
def _verify_span(spans: List[agl.Span]):
"""Only a few spans are checked here.
`raw_gen_ai_request` span:
```python
Span(
rollout_id='ro-4c68a7e686a1',
attempt_id='at-308eb814',
sequence_id=1,
name='raw_gen_ai_request',
attributes={
'llm.hosted_vllm.messages': '[{\'role\': \'user\', \'content\': "Hello, what\'s your name?"}]',
'llm.hosted_vllm.extra_body': "{'return_token_ids': True}",
'llm.hosted_vllm.choices': '... \'token_ids\': [40, 1079, 1207, 16948, ...',
'llm.hosted_vllm.model': 'Qwen/Qwen2.5-0.5B-Instruct',
'llm.hosted_vllm.prompt_token_ids': '[151644, 8948, ...]',
},
resource=OtelResource(
attributes={
'agentlightning.rollout_id': 'ro-4c68a7e686a1',
'agentlightning.attempt_id': 'at-308eb814',
'agentlightning.span_sequence_id': 1
},
)
)
```
"""
assert len(spans) > 1
has_raw_gen_ai_request = False
for span in spans:
if span.name == "raw_gen_ai_request":
has_raw_gen_ai_request = True
if "llm.hosted_vllm.messages" in span.attributes:
assert "return_token_ids" in span.attributes["llm.hosted_vllm.extra_body"] # type: ignore
assert "token_ids" in span.attributes["llm.hosted_vllm.choices"] # type: ignore
assert span.attributes["llm.hosted_vllm.prompt_token_ids"]
assert "agentlightning.rollout_id" in span.resource.attributes
assert "agentlightning.attempt_id" in span.resource.attributes
assert "agentlightning.span_sequence_id" in span.resource.attributes
assert has_raw_gen_ai_request, "raw_gen_ai_request span not found"
if __name__ == "__main__":
agl.setup_logging()
parser = argparse.ArgumentParser(description="LLM Proxy runner")
parser.add_argument(
"mode",
choices=["vllm", "openai", "test"],
help="Which function to run",
)
parser.add_argument("model", type=str, help="Model name to serve.")
args = parser.parse_args()
if args.mode == "vllm":
asyncio.run(serve_llm_proxy_with_vllm(args.model))
elif args.mode == "openai":
asyncio.run(serve_llm_proxy_with_openai(args.model))
elif args.mode == "test":
asyncio.run(test_llm_proxy(args.model))
+91
View File
@@ -0,0 +1,91 @@
# Copyright (c) Microsoft. All rights reserved.
"""Programmatically launch and stop an vLLM server."""
import subprocess
import time
from contextlib import contextmanager
from typing import Optional
import httpx
from openai import OpenAI
from rich.console import Console
console = Console()
@contextmanager
def vllm_server(
model_path: str,
port: int,
startup_timeout: float = 300.0,
terminate_timeout: float = 10.0,
gpu_memory_utilization: float = 0.7,
auto_tool_choice: bool = True,
tool_call_parser: Optional[str] = "hermes",
):
"""Serves a vLLM model from command line.
Args:
model_path: The path to the vLLM model. It can be either a local path or a Hugging Face model ID.
port: The port to serve the model on.
startup_timeout: The timeout for the server to start.
terminate_timeout: The timeout for the server to terminate.
gpu_memory_utilization: The GPU memory utilization for the server. Set it lower to avoid OOM.
auto_tool_choice: Whether to enable auto tool choice.
tool_call_parser: The tool call parser to use.
"""
proc: Optional[subprocess.Popen[bytes]] = None
try:
vllm_serve_args = [
"--gpu-memory-utilization",
str(gpu_memory_utilization),
"--port",
str(port),
]
if auto_tool_choice:
vllm_serve_args.append("--enable-auto-tool-choice")
if tool_call_parser is not None:
vllm_serve_args.append("--tool-call-parser")
vllm_serve_args.append(tool_call_parser)
proc = subprocess.Popen(["vllm", "serve", model_path, *vllm_serve_args])
# Wait for the server to be ready
url = f"http://localhost:{port}/health"
start = time.time()
client = httpx.Client()
while True:
try:
if client.get(url).status_code == 200:
break
except Exception:
result = proc.poll()
if result is not None and result != 0:
raise RuntimeError("Server exited unexpectedly.") from None
time.sleep(0.5)
if time.time() - start > startup_timeout:
raise RuntimeError(f"Server failed to start in {startup_timeout} seconds.") from None
yield f"http://localhost:{port}/v1"
finally:
# Terminate the server
if proc is None:
return
proc.terminate()
try:
proc.wait(terminate_timeout)
except subprocess.TimeoutExpired:
proc.kill()
if __name__ == "__main__":
with vllm_server("Qwen/Qwen2.5-0.5B-Instruct", 8080) as endpoint:
client = OpenAI(base_url=endpoint, api_key="dummy")
response = client.chat.completions.create(
model="Qwen/Qwen2.5-0.5B-Instruct",
messages=[{"role": "user", "content": "Hello, what's your name?"}],
)
console.print(response)
assert "qwen" in response.choices[0].message.content.lower() # type: ignore
+187
View File
@@ -0,0 +1,187 @@
# Copyright (c) Microsoft. All rights reserved.
"""Example to write traces to a LightningStore via raw OpenTelemetry or AgentOpsTracer.
The example can be run with or without using a Lightning Store server.
When running this server, the traces will be written to the server via OTLP endpoint.
Prior to running this example with `--use-client` flag, please start a LightningStore server with OTLP enabled first:
```bash
agl store --port 45993 --log-level DEBUG
```
"""
import argparse
import asyncio
import time
from typing import List
from openai import AsyncOpenAI
from rich.console import Console
from agentlightning import AgentOpsTracer, LightningStoreClient, OtelTracer, Span, emit_reward, setup_logging
from agentlightning.store import InMemoryLightningStore
console = Console()
async def send_traces_via_otel(use_client: bool = False):
tracer = OtelTracer()
if not use_client:
store = InMemoryLightningStore()
else:
store = LightningStoreClient("http://localhost:45993")
rollout = await store.start_rollout(input={"origin": "write_traces_example"})
with tracer.lifespan():
# Initialize the capture of one single trace for one single rollout
async with tracer.trace_context(
"trace-manual", store=store, rollout_id=rollout.rollout_id, attempt_id=rollout.attempt.attempt_id
) as tracer:
with tracer.start_as_current_span("grpc-span-1"):
time.sleep(0.01)
# Nested Span
with tracer.start_as_current_span("grpc-span-2"):
time.sleep(0.01)
with tracer.start_as_current_span("grpc-span-3"):
time.sleep(0.01)
# This creates a reward span
emit_reward(1.0)
traces = await store.query_spans(rollout_id=rollout.rollout_id)
console.print(traces)
# Quickly validate the traces
assert len(traces) == 4
span_names = [span.name for span in traces]
assert "grpc-span-1" in span_names
assert "grpc-span-2" in span_names
assert "grpc-span-3" in span_names
assert "agentlightning.reward" in span_names
last_span = traces[-1]
assert last_span.name == "agentlightning.reward"
# NOTE: Try not to rely on this attribute. It may change in the future.
# Use utils from agentlightning.emitter to get the reward value.
assert last_span.attributes["reward"] == 1.0
if use_client:
# When using client, the resource should have rollout_id and attempt_id set
for span in traces:
assert "agentlightning.rollout_id" in span.resource.attributes
assert "agentlightning.attempt_id" in span.resource.attributes
if isinstance(store, LightningStoreClient):
await store.close()
async def send_traces_via_agentops(use_client: bool = False):
tracer = AgentOpsTracer()
if not use_client:
store = InMemoryLightningStore()
else:
store = LightningStoreClient("http://localhost:45993")
rollout = await store.start_rollout(input={"origin": "write_traces_example"})
# Initialize the tracer lifespan
# One lifespan can contain multiple traces
with tracer.lifespan():
# Initialize the capture of one single trace for one single rollout
async with tracer.trace_context(
"trace-1", store=store, rollout_id=rollout.rollout_id, attempt_id=rollout.attempt.attempt_id
):
openai_client = AsyncOpenAI()
response = await openai_client.chat.completions.create(
model="gpt-4.1-mini",
messages=[
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "Hello, what's your name?"},
],
)
assert response.choices[0].message.content is not None
assert "chatgpt" in response.choices[0].message.content.lower()
traces = await store.query_spans(rollout_id=rollout.rollout_id)
console.print(traces)
await _verify_agentops_traces(traces, use_client=use_client)
if isinstance(store, LightningStoreClient):
await store.close()
async def _verify_agentops_traces(spans: List[Span], use_client: bool = False):
"""Expected traces to something like:
```python
Span(
rollout_id='ro-ef9ff8a429d1',
attempt_id='at-37cc5f24',
sequence_id=1,
trace_id='b3a16b603f7805934215d467e717c9e7',
span_id='2782d5d750f49b2d',
parent_id='2fb97c818363bce3',
name='openai.chat.completion',
status=TraceStatus(status_code='OK', description=None),
attributes={
'gen_ai.request.type': 'chat',
'gen_ai.system': 'OpenAI',
'gen_ai.request.model': 'gpt-4.1-mini',
'gen_ai.request.streaming': False,
'gen_ai.prompt.0.role': 'system',
'gen_ai.prompt.0.content': 'You are a helpful assistant.',
'gen_ai.prompt.1.role': 'user',
'gen_ai.prompt.1.content': "Hello, what's your name?",
'gen_ai.response.id': 'chatcmpl-Cc1osPWiArOwCS8nUkp0kZuZPkpY4',
'gen_ai.response.model': 'gpt-4.1-mini-2025-04-14',
'gen_ai.completion.0.role': 'assistant',
'gen_ai.completion.0.content': "Hello! I'm ChatGPT, your AI assistant. How can I help you today?",
},
resource=OtelResource(
attributes={
'agentops.project.id': 'temporary',
'agentlightning.rollout_id': 'ro-ef9ff8a429d1',
'agentlightning.attempt_id': 'at-37cc5f24'
},
schema_url=''
)
)
```
"""
assert len(spans) == 2
for span in spans:
if span.name == "openai.chat.completion":
assert span.attributes["gen_ai.request.model"] == "gpt-4.1-mini"
assert span.attributes["gen_ai.request.streaming"] == False
assert span.attributes["gen_ai.prompt.0.role"] == "system"
assert span.attributes["gen_ai.prompt.0.content"] == "You are a helpful assistant."
assert span.attributes["gen_ai.prompt.1.role"] == "user"
assert span.attributes["gen_ai.prompt.1.content"] == "Hello, what's your name?"
assert "chatgpt" in span.attributes["gen_ai.completion.0.content"].lower() # type: ignore
if use_client:
assert "agentlightning.rollout_id" in span.resource.attributes
assert "agentlightning.attempt_id" in span.resource.attributes
else:
assert "trace-1" in span.name
assert span.attributes["agentops.span.kind"] == "session"
def main():
setup_logging("DEBUG")
parser = argparse.ArgumentParser()
parser.add_argument("mode", choices=["otel", "agentops"])
parser.add_argument("--use-client", action="store_true")
args = parser.parse_args()
if args.mode == "otel":
asyncio.run(send_traces_via_otel(use_client=args.use_client))
elif args.mode == "agentops":
asyncio.run(send_traces_via_agentops(use_client=args.use_client))
else:
raise ValueError(f"Invalid mode: {args.mode}")
if __name__ == "__main__":
main()
+6 -3
View File
@@ -17,7 +17,7 @@ def test_switchable_authenticated_exporter():
switchable_authenticated_exporter = BypassableAuthenticatedOTLPExporter(endpoint="http://dummy", jwt="dummy")
with patch.object(
switchable_authenticated_exporter.__class__.__bases__[0], "export", return_value=SpanExportResult.SUCCESS
switchable_authenticated_exporter.__class__.__bases__[-1], "export", return_value=SpanExportResult.SUCCESS
) as mock_export:
enable_agentops_service()
result = switchable_authenticated_exporter.export([])
@@ -34,7 +34,7 @@ def test_switchable_otlp_metric_exporter():
switchable_otlp_metric_exporter = BypassableOTLPMetricExporter()
with patch.object(
switchable_otlp_metric_exporter.__class__.__bases__[0], "export", return_value=MetricExportResult.SUCCESS
switchable_otlp_metric_exporter.__class__.__bases__[-1], "export", return_value=MetricExportResult.SUCCESS
) as mock_export:
enable_agentops_service()
result = switchable_otlp_metric_exporter.export(metrics_data=MagicMock())
@@ -51,7 +51,10 @@ def test_switchable_otlp_span_exporter():
switchable_otlp_span_exporter = BypassableOTLPSpanExporter()
with patch.object(
switchable_otlp_span_exporter.__class__.__bases__[0], "export", return_value=SpanExportResult.SUCCESS
# BypassableOTLPSpanExporter is a subclass of LightningStoreOTLPExporter, which is a subclass of OTLPSpanExporter
switchable_otlp_span_exporter.__class__.__bases__[-1].__bases__[0],
"export",
return_value=SpanExportResult.SUCCESS,
) as mock_export:
enable_agentops_service()
result = switchable_otlp_span_exporter.export([])
+89 -1
View File
@@ -2,18 +2,23 @@
import asyncio
import logging
import multiprocessing
import random
from typing import Any, List, cast
import litellm
import openai
import opentelemetry.trace as trace_api
import pytest
from agentops.sdk.core import BatchSpanProcessor
from litellm.llms.custom_llm import CustomLLM
from litellm.types.utils import ModelResponse
from litellm.utils import custom_llm_setup
from opentelemetry.sdk.trace import ReadableSpan
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from agentlightning.llm_proxy import LightningSpanExporter, LLMProxy
from agentlightning.store import LightningStoreServer
from agentlightning.store.memory import InMemoryLightningStore
from agentlightning.store.threading import LightningStoreThreaded
from agentlightning.types import Span
@@ -282,7 +287,6 @@ async def test_custom_llm_restarted_multiple_times(caplog: pytest.LogCaptureFixt
assert response.choices[0].message.content == f"Hi! {restart_idx}"
error_logs = [record.message for record in caplog.records if record.levelno >= logging.ERROR]
error_logs = [message for message in error_logs if "Task was destroyed but it is pending!" not in message]
assert not error_logs, f"Found error logs: {error_logs}"
assert not any("Cannot add callback" in record.message for record in caplog.records)
@@ -290,3 +294,87 @@ async def test_custom_llm_restarted_multiple_times(caplog: pytest.LogCaptureFixt
finally:
litellm.custom_provider_map = []
custom_llm_setup()
async def llm_proxy_span_exporter_loop(otlp_enabled: bool = False):
store = LightningStoreThreaded(InMemoryLightningStore())
if otlp_enabled:
store = LightningStoreServer(store, "127.0.0.1", get_free_port())
await store.start()
llm_instance = TestLLM(f"Hi! I'm a test LLM")
litellm.custom_provider_map = [{"provider": "test-llm", "custom_handler": llm_instance}]
custom_llm_setup()
proxy = LLMProxy(
launcher_args=PythonServerLauncherArgs(
launch_mode="thread",
healthcheck_url="/health",
port=get_free_port(),
),
store=store,
model_list=[
{
"model_name": "gpt-4o-arbitrary",
"litellm_params": {
"model": "test-llm/any-llm",
},
}
],
)
await proxy.start()
rollout = await store.start_rollout(None)
resource = proxy.as_resource(rollout.rollout_id, rollout.attempt.attempt_id)
client = openai.AsyncOpenAI(
base_url=resource.endpoint,
api_key="token-abc123",
timeout=5,
max_retries=0,
)
response = await client.chat.completions.create(
model="gpt-4o-arbitrary",
messages=[{"role": "user", "content": "Hello world"}],
stream=False,
)
assert response.choices[0].message.content == "Hi! I'm a test LLM"
spans = await store.query_spans(rollout.rollout_id, rollout.attempt.attempt_id)
assert len(spans) > 0, "Should have captured spans"
for span in spans:
assert span.rollout_id == rollout.rollout_id, f"Span {span.name} has incorrect rollout_id"
assert span.attempt_id == rollout.attempt.attempt_id, f"Span {span.name} has incorrect attempt_id"
assert span.sequence_id == 1, f"Span {span.name} has incorrect sequence_id"
tracer_provider = trace_api.get_tracer_provider()
have_asserted_loop = False
for span_processor in tracer_provider._active_span_processor._span_processors: # type: ignore
if isinstance(span_processor, (SimpleSpanProcessor, BatchSpanProcessor)):
if isinstance(span_processor.span_exporter, LightningSpanExporter):
if otlp_enabled:
assert span_processor.span_exporter._loop is None # type: ignore
else:
assert span_processor.span_exporter._loop is not None # type: ignore
have_asserted_loop = True
break
assert have_asserted_loop, f"LightningSpanExporter should be used with otlp_enabled={otlp_enabled}"
await proxy.stop()
if isinstance(store, LightningStoreServer):
await store.stop()
def llm_proxy_span_exporter_loop_sync(otlp_enabled: bool = False):
asyncio.run(llm_proxy_span_exporter_loop(otlp_enabled))
@pytest.mark.parametrize("otlp_enabled", [True, False])
def test_llm_proxy_span_exporter_loop(otlp_enabled: bool):
context = multiprocessing.get_context("spawn")
process = context.Process(target=llm_proxy_span_exporter_loop_sync, args=(otlp_enabled,))
process.start()
process.join(timeout=30.0)
assert process.exitcode == 0
+61 -36
View File
@@ -21,14 +21,14 @@ import anthropic
import openai
import pytest
from litellm.integrations.custom_logger import CustomLogger
from portpicker import pick_unused_port
from agentlightning import LlmProxyTraceToTriplet
from agentlightning.llm_proxy import LLMProxy, _reset_litellm_logging_worker # pyright: ignore[reportPrivateUsage]
from agentlightning.store import LightningStore, LightningStoreServer
from agentlightning.store import LightningStore, LightningStoreServer, LightningStoreThreaded
from agentlightning.store.memory import InMemoryLightningStore
from agentlightning.types import LLM, Span
from ..common.network import get_free_port
from ..common.tracer import clear_tracer_provider
from ..common.vllm import VLLM_VERSION, RemoteOpenAIServer
@@ -52,7 +52,7 @@ def qwen25_model():
"--tool-call-parser",
"hermes",
"--port",
str(get_free_port()),
str(pick_unused_port()),
],
) as server:
yield server
@@ -69,13 +69,17 @@ def test_qwen25_model_sanity(qwen25_model: RemoteOpenAIServer):
@pytest.mark.asyncio
async def test_basic_integration(qwen25_model: RemoteOpenAIServer):
@pytest.mark.parametrize("otlp_enabled", [True, False])
async def test_basic_integration(qwen25_model: RemoteOpenAIServer, otlp_enabled: bool):
clear_tracer_provider()
inmemory_store = InMemoryLightningStore()
store = LightningStoreServer(store=inmemory_store, host="127.0.0.1", port=get_free_port())
await store.start()
if otlp_enabled:
store = LightningStoreServer(store=inmemory_store, host="127.0.0.1", port=pick_unused_port())
await store.start()
else:
store = LightningStoreThreaded(inmemory_store)
proxy = LLMProxy(
port=get_free_port(),
port=pick_unused_port(),
model_list=[
{
"model_name": "gpt-4o-arbitrary",
@@ -86,6 +90,7 @@ async def test_basic_integration(qwen25_model: RemoteOpenAIServer):
}
],
store=store,
launch_mode="thread" if not otlp_enabled else "mp",
)
rollout = await store.start_rollout(None)
@@ -107,7 +112,8 @@ async def test_basic_integration(qwen25_model: RemoteOpenAIServer):
spans = await store.query_spans(rollout.rollout_id, rollout.attempt.attempt_id)
await store.stop()
if isinstance(store, LightningStoreServer):
await store.stop()
# Verify all spans have correct rollout_id, attempt_id, and sequence_id
assert len(spans) > 0, "Should have captured spans"
@@ -187,13 +193,18 @@ async def _make_proxy_and_store(
retries: int = 0,
gunicorn: bool = False,
callbacks: List[Union[Type[CustomLogger], str]] | None = None,
otlp_enabled: bool = False,
):
clear_tracer_provider()
_reset_litellm_logging_worker() # type: ignore
store = InMemoryLightningStore()
store_server = LightningStoreServer(store=store, host="127.0.0.1", port=get_free_port())
# When the server is forked into subprocess, it automatically becomes a client of the store
await store_server.start()
if otlp_enabled:
store = LightningStoreServer(store=store, host="127.0.0.1", port=pick_unused_port())
# When the server is forked into subprocess, it automatically becomes a client of the store
await store.start()
else:
# Backward compatibility with legacy thread + non-otlp mode
store = LightningStoreThreaded(store)
proxy = LLMProxy(
model_list=[
{
@@ -204,14 +215,15 @@ async def _make_proxy_and_store(
},
}
],
port=get_free_port(),
launch_mode="thread" if not otlp_enabled else "mp",
port=pick_unused_port(),
num_workers=4 if gunicorn else 1,
store=store_server,
store=store,
num_retries=retries,
callbacks=callbacks,
)
await proxy.start()
return proxy, store_server
return proxy, store
async def _new_resource(proxy: LLMProxy, store: LightningStore):
@@ -236,8 +248,9 @@ def _attr(s: Span, key: str, default: Any = None): # type: ignore
@pytest.mark.asyncio
async def test_multiple_requests_one_attempt(qwen25_model: RemoteOpenAIServer):
proxy, store = await _make_proxy_and_store(qwen25_model)
@pytest.mark.parametrize("otlp_enabled", [True, False])
async def test_multiple_requests_one_attempt(qwen25_model: RemoteOpenAIServer, otlp_enabled: bool):
proxy, store = await _make_proxy_and_store(qwen25_model, otlp_enabled=otlp_enabled)
try:
resource, rollout = await _new_resource(proxy, store)
client = _get_client_for_resource(resource)
@@ -259,13 +272,14 @@ async def test_multiple_requests_one_attempt(qwen25_model: RemoteOpenAIServer):
# TODO: Check response contents and token ids for the 3 requests respectively
finally:
await proxy.stop()
await store.stop()
if isinstance(store, LightningStoreServer):
await store.stop()
@pytest.mark.asyncio
@pytest.mark.parametrize("gunicorn", [False, True])
async def test_ten_concurrent_requests(qwen25_model: RemoteOpenAIServer, gunicorn: bool):
proxy, store = await _make_proxy_and_store(qwen25_model, gunicorn=gunicorn)
@pytest.mark.parametrize("mode", ["gunicorn", "thread", "uvicorn"])
async def test_ten_concurrent_requests(qwen25_model: RemoteOpenAIServer, mode: str):
proxy, store = await _make_proxy_and_store(qwen25_model, gunicorn=mode == "gunicorn", otlp_enabled=mode != "thread")
try:
resource, rollout = await _new_resource(proxy, store)
aclient = _get_async_client_for_resource(resource)
@@ -288,13 +302,15 @@ async def test_ten_concurrent_requests(qwen25_model: RemoteOpenAIServer, gunicor
# TODO: Check whether the sequence ids get mixed up or not
finally:
await proxy.stop()
await store.stop()
if isinstance(store, LightningStoreServer):
await store.stop()
@pytest.mark.asyncio
async def test_anthropic_client_compat(qwen25_model: RemoteOpenAIServer):
@pytest.mark.parametrize("otlp_enabled", [True, False])
async def test_anthropic_client_compat(qwen25_model: RemoteOpenAIServer, otlp_enabled: bool):
# litellm proxy accepts Anthropic schema and forwards to OpenAI backend
proxy, store = await _make_proxy_and_store(qwen25_model)
proxy, store = await _make_proxy_and_store(qwen25_model, otlp_enabled=otlp_enabled)
try:
resource, rollout = await _new_resource(proxy, store)
@@ -312,12 +328,14 @@ async def test_anthropic_client_compat(qwen25_model: RemoteOpenAIServer):
assert len(spans) > 0
finally:
await proxy.stop()
await store.stop()
if isinstance(store, LightningStoreServer):
await store.stop()
@pytest.mark.asyncio
async def test_tool_call_roundtrip(qwen25_model: RemoteOpenAIServer):
proxy, store = await _make_proxy_and_store(qwen25_model)
@pytest.mark.parametrize("otlp_enabled", [True, False])
async def test_tool_call_roundtrip(qwen25_model: RemoteOpenAIServer, otlp_enabled: bool):
proxy, store = await _make_proxy_and_store(qwen25_model, otlp_enabled=otlp_enabled)
try:
resource, rollout = await _new_resource(proxy, store)
client = _get_client_for_resource(resource)
@@ -381,12 +399,14 @@ async def test_tool_call_roundtrip(qwen25_model: RemoteOpenAIServer):
# TODO: Check response contents and token ids for the 2 requests respectively
finally:
await proxy.stop()
await store.stop()
if isinstance(store, LightningStoreServer):
await store.stop()
@pytest.mark.asyncio
async def test_streaming_chunks(qwen25_model: RemoteOpenAIServer):
proxy, store = await _make_proxy_and_store(qwen25_model)
@pytest.mark.parametrize("otlp_enabled", [True, False])
async def test_streaming_chunks(qwen25_model: RemoteOpenAIServer, otlp_enabled: bool):
proxy, store = await _make_proxy_and_store(qwen25_model, otlp_enabled=otlp_enabled)
try:
resource, rollout = await _new_resource(proxy, store)
client = _get_client_for_resource(resource)
@@ -417,12 +437,14 @@ async def test_streaming_chunks(qwen25_model: RemoteOpenAIServer):
assert "gen_ai.completion.0.content" in span.attributes
finally:
await proxy.stop()
await store.stop()
if isinstance(store, LightningStoreServer):
await store.stop()
@pytest.mark.asyncio
async def test_anthropic_token_ids(qwen25_model: RemoteOpenAIServer):
proxy, store = await _make_proxy_and_store(qwen25_model)
@pytest.mark.parametrize("otlp_enabled", [True, False])
async def test_anthropic_token_ids(qwen25_model: RemoteOpenAIServer, otlp_enabled: bool):
proxy, store = await _make_proxy_and_store(qwen25_model, otlp_enabled=otlp_enabled)
try:
resource, rollout = await _new_resource(proxy, store)
adapter = LlmProxyTraceToTriplet()
@@ -479,7 +501,8 @@ async def test_anthropic_token_ids(qwen25_model: RemoteOpenAIServer):
assert len(triplets) == 2
finally:
await proxy.stop()
await store.stop()
if isinstance(store, LightningStoreServer):
await store.stop()
class LogprobsCallback(CustomLogger):
@@ -489,9 +512,10 @@ class LogprobsCallback(CustomLogger):
@pytest.mark.asyncio
async def test_anthropic_logprobs(qwen25_model: RemoteOpenAIServer):
@pytest.mark.parametrize("otlp_enabled", [True, False])
async def test_anthropic_logprobs(qwen25_model: RemoteOpenAIServer, otlp_enabled: bool):
proxy, store = await _make_proxy_and_store(
qwen25_model, callbacks=[LogprobsCallback, "return_token_ids", "opentelemetry"]
qwen25_model, callbacks=[LogprobsCallback, "return_token_ids", "opentelemetry"], otlp_enabled=otlp_enabled
)
try:
resource, rollout = await _new_resource(proxy, store)
@@ -543,4 +567,5 @@ async def test_anthropic_logprobs(qwen25_model: RemoteOpenAIServer):
# TODO: Check logprobs
finally:
await proxy.stop()
await store.stop()
if isinstance(store, LightningStoreServer):
await store.stop()
+99 -55
View File
@@ -10,17 +10,25 @@ import pickle
import threading
import time
from multiprocessing.connection import Connection
from typing import Any
from typing import Any, Dict, List, Optional
from unittest.mock import AsyncMock, MagicMock
import pytest
import uvicorn
from fastapi import FastAPI, Request
from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import (
ExportTraceServiceRequest,
ExportTraceServiceResponse,
)
from opentelemetry.sdk.trace import ReadableSpan
from opentelemetry.trace import SpanContext, TraceFlags
from portpicker import pick_unused_port
from agentlightning.reward import emit_reward, find_reward_spans, get_reward_value, reward
from agentlightning.store.base import LightningStore
from agentlightning.tracer.agentops import LightningSpanProcessor
from agentlightning.tracer.otel import OtelTracer
from agentlightning.utils import otlp
from ..common.tracer import clear_agentops_init, clear_tracer_provider
@@ -41,13 +49,71 @@ def create_span(name: str, sampled: bool = True, with_context: bool = True) -> M
return span
def create_mock_store() -> MagicMock:
def create_mock_store(otlp_supported: bool = False) -> MagicMock:
"""Helper to create a mock LightningStore."""
store = MagicMock(spec=LightningStore)
store.add_otel_span = AsyncMock(return_value=None)
store.capabilities = {"otlp_traces": otlp_supported}
store.otlp_traces_endpoint.return_value = "http://store/v1/traces"
return store
@pytest.fixture(params=[False, True], ids=["store-no-otlp", "store-otlp"])
def store_supports_otlp(request: pytest.FixtureRequest) -> bool:
return bool(request.param)
@pytest.fixture
def otlp_server(store_supports_otlp: bool):
if not store_supports_otlp:
yield None
return
app = FastAPI()
received: List[ExportTraceServiceRequest] = []
@app.post("/v1/traces")
async def _export_traces(request: Request): # type: ignore
async def capture(message: ExportTraceServiceRequest) -> None:
received.append(message)
return await otlp.handle_otlp_export(
request,
ExportTraceServiceRequest,
ExportTraceServiceResponse,
capture,
signal_name="traces",
)
port = pick_unused_port()
config = uvicorn.Config(app, host="127.0.0.1", port=port, log_level="error")
server = uvicorn.Server(config)
thread = threading.Thread(target=server.run, daemon=True)
thread.start()
timeout = time.time() + 5
while not getattr(server, "started", False):
if time.time() > timeout:
raise RuntimeError("OTLP test server failed to start")
if not thread.is_alive():
raise RuntimeError("OTLP test server thread exited before startup")
time.sleep(0.01)
try:
yield {"url": f"http://127.0.0.1:{port}/v1/traces", "received": received}
finally:
server.should_exit = True
thread.join(timeout=5)
@pytest.fixture
def store(store_supports_otlp: bool, otlp_server: Optional[Dict[str, Any]]) -> MagicMock:
mock_store = create_mock_store(store_supports_otlp)
if store_supports_otlp:
assert otlp_server is not None
mock_store.otlp_traces_endpoint.return_value = otlp_server["url"]
return mock_store
def test_initialization_and_shutdown():
"""Test processor lifecycle: initialization, loop thread, and shutdown."""
processor = LightningSpanProcessor()
@@ -58,9 +124,16 @@ def test_initialization_and_shutdown():
assert processor._rollout_id is None
assert processor._attempt_id is None
assert processor._loop is None
assert processor._loop_thread is None
# Start the loop
processor._ensure_loop()
# Verify loop thread is running correctly
assert processor._loop is not None
assert processor._loop.is_running()
assert processor._loop_thread is not None
assert processor._loop_thread.is_alive()
assert processor._loop_thread.daemon is True
assert processor._loop_thread.name == "otel-loop"
@@ -99,10 +172,9 @@ def test_span_collection_with_filtering():
processor.shutdown()
def test_context_managers_clear_state():
def test_context_managers_clear_state(store: MagicMock):
"""Test that both context managers properly manage state."""
processor = LightningSpanProcessor()
store = create_mock_store()
# Add a span first
span1 = create_span("span1")
@@ -136,10 +208,9 @@ def test_context_managers_clear_state():
processor.shutdown()
def test_store_integration_complete():
def test_store_integration_complete(store: MagicMock):
"""Test all store integration scenarios: writes, errors, timeout, thread verification."""
processor = LightningSpanProcessor()
store = create_mock_store()
# Test 1: Successful store writes
with processor.with_context(store=store, rollout_id="r1", attempt_id="a1"):
@@ -231,10 +302,9 @@ def test_event_loop_operations():
processor.shutdown()
def test_concurrent_access():
def test_concurrent_access(store: MagicMock):
"""Test thread-safe concurrent span processing."""
processor = LightningSpanProcessor()
store = create_mock_store()
num_threads = 10
spans_per_thread = 5
@@ -261,34 +331,42 @@ def test_concurrent_access():
processor.shutdown()
def test_multiprocessing_behavior():
def test_multiprocessing_behavior(store_supports_otlp: bool):
"""Test processor behavior across process boundaries."""
# Test 1: Creating new processor in subprocess works
def subprocess_task(result_queue: "multiprocessing.Queue[tuple[str, Any]]") -> None:
def subprocess_task(result_queue: "multiprocessing.Queue[tuple[str, Any]]", supports_otlp: bool) -> None:
try:
processor = LightningSpanProcessor()
# Verify processor works in new process
assert processor._loop is not None
assert processor._loop_thread.is_alive()
span = create_span("subprocess_span")
processor.on_end(span)
# Loop is not used
assert processor._loop is None
assert processor._loop_thread is None
with processor.with_context(store=create_mock_store(supports_otlp), rollout_id="r1", attempt_id="a1"):
processor.on_end(span)
# Verify processor works in new process
assert processor._loop is not None
assert processor._loop_thread is not None
assert processor._loop_thread.is_alive()
result_queue.put(("success", len(processor.spans())))
processor.shutdown()
except Exception as e:
result_queue.put(("error", str(e)))
result_queue: multiprocessing.Queue[tuple[str, Any]] = multiprocessing.Queue()
process = multiprocessing.Process(target=subprocess_task, args=(result_queue,))
process = multiprocessing.Process(target=subprocess_task, args=(result_queue, store_supports_otlp))
process.start()
process.join(timeout=5)
assert not process.is_alive()
status, value = result_queue.get(timeout=1)
assert status == "success"
assert status == "success", f"Subprocess failed: {status}, {value}"
assert value == 1
# Test 2: Processor cannot be pickled (threads aren't picklable)
@@ -300,45 +378,9 @@ def test_multiprocessing_behavior():
processor.shutdown()
def test_edge_cases():
"""Test edge cases and error conditions."""
processor = LightningSpanProcessor()
# Test 1: force_flush always returns True
assert processor.force_flush() is True
assert processor.force_flush(timeout_millis=5000) is True
# Test 2: Calling _await_in_loop after shutdown raises
processor.shutdown()
async def dummy_coro():
return "test"
with pytest.raises(RuntimeError, match="Loop is not initialized"):
processor._await_in_loop(dummy_coro())
# Test 3: Verify shutdown thread join timeout is respected
processor2 = LightningSpanProcessor()
# Mock thread.join to verify timeout parameter
original_join = processor2._loop_thread.join
join_timeout: float | None = None
def mock_join(timeout: float | None = None) -> None:
nonlocal join_timeout
join_timeout = timeout
return original_join(timeout=timeout)
processor2._loop_thread.join = mock_join
processor2.shutdown()
assert join_timeout == 5 # Should pass 5 second timeout
def test_store_write_timeout():
def test_store_write_timeout(store: MagicMock):
"""Test that slow store writes respect timeout."""
processor = LightningSpanProcessor()
store = create_mock_store()
# Create a slow async function that exceeds timeout
async def slow_write(*args: Any, **kwargs: Any) -> None:
@@ -363,6 +405,9 @@ def test_multiple_processors_in_same_process():
processor1 = LightningSpanProcessor()
processor2 = LightningSpanProcessor()
processor1._ensure_loop()
processor2._ensure_loop()
# Both should have independent loops and threads
assert processor1._loop is not processor2._loop
assert processor1._loop_thread is not processor2._loop_thread
@@ -385,10 +430,9 @@ def test_multiple_processors_in_same_process():
processor2.shutdown()
def test_context_manager_reusability():
def test_context_manager_reusability(store: MagicMock):
"""Test that context managers can be entered and exited multiple times."""
processor = LightningSpanProcessor()
store = create_mock_store()
# First usage
with processor.with_context(store=store, rollout_id="r1", attempt_id="a1"):
+407
View File
@@ -0,0 +1,407 @@
# Copyright (c) Microsoft. All rights reserved.
# pyright: reportPrivateUsage=false
from __future__ import annotations
import gzip
from types import SimpleNamespace
from typing import Any, Dict, Iterable, List, Optional, cast
import pytest
from fastapi import Request
from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import (
ExportTraceServiceRequest,
ExportTraceServiceResponse,
)
from opentelemetry.proto.common.v1.common_pb2 import AnyValue, KeyValue
from opentelemetry.proto.resource.v1.resource_pb2 import Resource as ProtoResource
from opentelemetry.proto.trace.v1.trace_pb2 import Span as ProtoSpan
from opentelemetry.proto.trace.v1.trace_pb2 import Status as ProtoStatus
from opentelemetry.sdk.resources import Resource
from opentelemetry.sdk.trace import ReadableSpan
from opentelemetry.sdk.trace.export import SpanExportResult
from starlette.types import Message, Scope
from agentlightning.store import LightningStore
from agentlightning.types.tracer import SpanNames
from agentlightning.utils import otlp
BASE_TIME_NANOS = 1_700_000_000_000_000_000
EVENT_TIME_OFFSET = 3_000_000_000
EVENT_TIME_SECONDS = (BASE_TIME_NANOS + EVENT_TIME_OFFSET) / 1_000_000_000
EXTRA_EVENT_TIME_OFFSET = 4_000_000_000
EXTRA_EVENT_TIME_SECONDS = (BASE_TIME_NANOS + EXTRA_EVENT_TIME_OFFSET) / 1_000_000_000
class _StubStore(LightningStore):
def __init__(self) -> None:
self.sequence_calls: List[tuple[str, str]] = []
self.next_value = 1
async def get_next_span_sequence_id(self, rollout_id: str, attempt_id: str) -> int:
self.sequence_calls.append((rollout_id, attempt_id))
value = self.next_value
self.next_value += 1
return value
def _make_request(
body: bytes,
*,
content_type: str = otlp.PROTOBUF_CT,
content_encoding: Optional[str] = None,
accept_encoding: Optional[str] = None,
) -> Request:
headers: List[tuple[bytes, bytes]] = [(b"content-type", content_type.encode())]
if content_encoding:
headers.append((b"content-encoding", content_encoding.encode()))
if accept_encoding:
headers.append((b"accept-encoding", accept_encoding.encode()))
scope: Scope = {
"type": "http",
"asgi": {"version": "3.0", "spec_version": "2.3"},
"method": "POST",
"path": "/v1/test",
"headers": headers,
}
body_sent = False
async def receive() -> Message:
nonlocal body_sent
if body_sent:
return {"type": "http.request", "body": b"", "more_body": False}
body_sent = True
return {"type": "http.request", "body": body, "more_body": False}
return Request(scope, receive)
def _set_any_value(av: AnyValue, value: object) -> None:
if isinstance(value, bool):
av.bool_value = value
elif isinstance(value, int):
av.int_value = value
elif isinstance(value, float):
av.double_value = value
elif isinstance(value, bytes):
av.bytes_value = value
elif isinstance(value, list):
for item in cast(List[Any], value):
_set_any_value(av.array_value.values.add(), item)
elif isinstance(value, dict):
for key, item in cast(Dict[str, Any], value).items():
kv = av.kvlist_value.values.add()
kv.key = key
_set_any_value(kv.value, item)
else:
av.string_value = str(value)
def _add_attribute(attrs: Iterable[KeyValue], key: str, value: object) -> None:
kv = attrs.add() # type: ignore
kv.key = key
_set_any_value(kv.value, value) # type: ignore
def _build_span_request() -> ExportTraceServiceRequest:
request = ExportTraceServiceRequest()
resource_spans = request.resource_spans.add()
_add_attribute(resource_spans.resource.attributes, SpanNames.ROLLOUT_ID, "resource-rollout")
_add_attribute(resource_spans.resource.attributes, SpanNames.ATTEMPT_ID, "resource-attempt")
_add_attribute(resource_spans.resource.attributes, SpanNames.SPAN_SEQUENCE_ID, "5")
resource_spans.schema_url = "https://example/schema"
scope_spans = resource_spans.scope_spans.add()
span = scope_spans.spans.add()
span.trace_id = bytes.fromhex("01" * 16)
span.span_id = bytes.fromhex("02" * 8)
span.parent_span_id = bytes.fromhex("03" * 8)
span.name = "test-span"
span.start_time_unix_nano = BASE_TIME_NANOS
span.end_time_unix_nano = BASE_TIME_NANOS + 2_000_000_000
span.status.code = ProtoStatus.STATUS_CODE_ERROR
span.status.message = "boom"
_add_attribute(span.attributes, "foo", "bar")
_add_attribute(span.attributes, SpanNames.ROLLOUT_ID, "span-rollout")
_add_attribute(span.attributes, SpanNames.ATTEMPT_ID, "span-attempt")
_add_attribute(span.attributes, SpanNames.SPAN_SEQUENCE_ID, "7")
event = span.events.add()
event.name = "event"
event.time_unix_nano = BASE_TIME_NANOS + EVENT_TIME_OFFSET
_add_attribute(event.attributes, "event-attr", 9)
link = span.links.add()
link.trace_id = bytes.fromhex("04" * 16)
link.span_id = bytes.fromhex("05" * 8)
_add_attribute(link.attributes, "link-attr", True)
return request
@pytest.mark.asyncio
async def test_handle_otlp_export_success_with_gzip_response() -> None:
request_msg = _build_span_request()
body = request_msg.SerializeToString()
request = _make_request(
body,
accept_encoding="gzip;q=0.9,br",
)
received: List[ExportTraceServiceRequest] = []
async def callback(message: ExportTraceServiceRequest) -> None:
received.append(message)
response = await otlp.handle_otlp_export(
request,
ExportTraceServiceRequest,
ExportTraceServiceResponse,
callback,
signal_name="traces",
)
assert response.status_code == 200
assert received and received[0].SerializeToString() == body
assert response.headers["Content-Encoding"] == "gzip"
assert gzip.decompress(response.body) == ExportTraceServiceResponse().SerializeToString()
@pytest.mark.asyncio
async def test_handle_otlp_export_rejects_invalid_content_type() -> None:
request = _make_request(b"{}", content_type="application/json")
response = await otlp.handle_otlp_export(
request,
ExportTraceServiceRequest,
ExportTraceServiceResponse,
None,
signal_name="traces",
)
assert response.status_code == 400
status = otlp.Status() # type: ignore[attr-defined]
status.ParseFromString(response.body) # type: ignore
assert "Unsupported Content-Type" in status.message
@pytest.mark.asyncio
async def test_handle_otlp_export_rejects_bad_payload() -> None:
request = _make_request(b"not-a-proto")
response = await otlp.handle_otlp_export(
request,
ExportTraceServiceRequest,
ExportTraceServiceResponse,
None,
signal_name="traces",
)
assert response.status_code == 400
status = otlp.Status() # type: ignore[attr-defined]
status.ParseFromString(response.body) # type: ignore
assert "Unable to parse" in status.message
@pytest.mark.asyncio
async def test_handle_otlp_export_accepts_gzip_body() -> None:
request_msg = ExportTraceServiceRequest()
request_msg.resource_spans.add()
gz_body = gzip.compress(request_msg.SerializeToString())
request = _make_request(gz_body, content_encoding="gzip")
response = await otlp.handle_otlp_export(
request,
ExportTraceServiceRequest,
ExportTraceServiceResponse,
None,
signal_name="traces",
)
assert response.status_code == 200
@pytest.mark.asyncio
async def test_spans_from_proto_prefers_span_level_metadata() -> None:
store = _StubStore()
request = _build_span_request()
spans = await otlp.spans_from_proto(request, store)
assert len(spans) == 1
span = spans[0]
assert span.rollout_id == "span-rollout"
assert span.attempt_id == "span-attempt"
assert span.sequence_id == 7
assert span.status.status_code == "ERROR"
assert span.events[0].timestamp == pytest.approx(EVENT_TIME_SECONDS) # type: ignore
assert span.links[0].context.trace_id == "0404" * 8
assert span.links[0].attributes == {"link-attr": True}
assert span.resource.attributes[SpanNames.ROLLOUT_ID] == "resource-rollout"
assert span.resource.schema_url == "https://example/schema"
assert not store.sequence_calls
@pytest.mark.asyncio
async def test_spans_from_proto_requests_sequence_ids_when_missing() -> None:
store = _StubStore()
request = ExportTraceServiceRequest()
resource_spans = request.resource_spans.add()
_add_attribute(resource_spans.resource.attributes, SpanNames.ROLLOUT_ID, "r1")
_add_attribute(resource_spans.resource.attributes, SpanNames.ATTEMPT_ID, "a1")
scope_span = resource_spans.scope_spans.add()
span = scope_span.spans.add()
span.trace_id = b"" # exercise default ids
span.span_id = b""
span.name = "needs-seq"
spans = await otlp.spans_from_proto(request, store)
assert len(spans) == 1
assert spans[0].sequence_id == 1
assert store.sequence_calls == [("r1", "a1")]
@pytest.mark.asyncio
async def test_spans_from_proto_skips_spans_without_ids() -> None:
store = _StubStore()
request = ExportTraceServiceRequest()
request.resource_spans.add() # missing rollout and attempt
spans = await otlp.spans_from_proto(request, store)
assert spans == []
assert store.sequence_calls == []
def test_normalize_sequence_id_handles_bad_values(caplog: pytest.LogCaptureFixture) -> None:
caplog.set_level("WARNING")
assert otlp._normalize_sequence_id("not-int") is None
assert any("Invalid sequence_id" in record.message for record in caplog.records)
def test_any_value_to_python_full_roundtrip() -> None:
av = AnyValue()
_set_any_value(
av,
{
"s": "hello",
"b": True,
"i": 5,
"d": 1.5,
"arr": ["x", 2],
"nested": {"k": b"\x01"},
},
)
result = otlp._any_value_to_python(av)
assert result == {
"s": "hello",
"b": True,
"i": 5,
"d": 1.5,
"arr": ["x", 2],
"nested": {"k": "01"},
}
def test_kv_list_to_dict_converts_values() -> None:
resource = ProtoResource()
_add_attribute(resource.attributes, "num", 10)
_add_attribute(resource.attributes, "flag", False)
converted = otlp._kv_list_to_dict(resource.attributes)
assert converted == {"num": 10, "flag": False}
def test_bytes_to_hex_helpers() -> None:
assert otlp._bytes_to_trace_id_hex(b"") == "0" * 32
assert otlp._bytes_to_span_id_hex(b"") == "0" * 16
assert otlp._bytes_to_trace_id_hex(b"\xff") == "ff".rjust(32, "0")
assert otlp._bytes_to_span_id_hex(b"\xaa") == "aa".rjust(16, "0")
def test_events_and_links_from_proto() -> None:
span = ProtoSpan()
event = span.events.add()
event.name = "evt"
event.time_unix_nano = BASE_TIME_NANOS + EXTRA_EVENT_TIME_OFFSET
_add_attribute(event.attributes, "alpha", "beta")
link = span.links.add()
link.trace_id = bytes.fromhex("06" * 16)
link.span_id = bytes.fromhex("07" * 8)
_add_attribute(link.attributes, "delta", 1)
events = otlp._events_from_proto(span)
links = otlp._links_from_proto(span)
assert events[0].timestamp == pytest.approx(EXTRA_EVENT_TIME_SECONDS) # type: ignore
assert events[0].attributes == {"alpha": "beta"}
assert links[0].context.trace_id == "0606" * 8
assert links[0].attributes == {"delta": 1}
def test_resource_from_proto() -> None:
resource = ProtoResource()
_add_attribute(resource.attributes, "key", "value")
result = otlp._resource_from_proto(resource, schema_url="https://example/schema")
assert result.attributes == {"key": "value"}
assert result.schema_url == "https://example/schema"
def test_maybe_gzip_response_parses_quality_values() -> None:
request = SimpleNamespace(headers={"Accept-Encoding": "br, gzip;q=0.1"})
payload = b"payload"
compressed, headers = otlp._maybe_gzip_response(cast(Request, request), payload)
assert headers == {"Content-Encoding": "gzip"}
assert gzip.decompress(compressed) == payload
def test_bad_request_response_matches_request_encoding() -> None:
request = SimpleNamespace(headers={"Accept-Encoding": "gzip"})
response = otlp._bad_request_response(cast(Request, request), "error")
assert response.status_code == 400
assert response.media_type == otlp.PROTOBUF_CT
status = otlp.Status() # type: ignore[attr-defined]
status.ParseFromString(gzip.decompress(response.body))
assert status.message == "error"
class _DummyReadableSpan:
def __init__(self) -> None:
self._resource = Resource.create({"existing": "value"})
def test_lightning_store_otlp_exporter_overrides_resources(monkeypatch: pytest.MonkeyPatch) -> None:
exporter = otlp.LightningStoreOTLPExporter(endpoint="http://collector")
captured_spans: List[List[_DummyReadableSpan]] = []
def fake_export(self: otlp.LightningStoreOTLPExporter, spans: List[_DummyReadableSpan]) -> SpanExportResult:
captured_spans.append(spans)
return SpanExportResult.SUCCESS
monkeypatch.setattr(otlp.OTLPSpanExporter, "export", fake_export, raising=False)
exporter.enable_store_otlp("http://store", "rollout", "attempt")
span = _DummyReadableSpan()
result = exporter.export([cast(ReadableSpan, span)])
assert result == SpanExportResult.SUCCESS
assert captured_spans
attributes = captured_spans[0][0]._resource.attributes # type: ignore[attr-defined]
assert attributes[SpanNames.ROLLOUT_ID] == "rollout"
assert attributes[SpanNames.ATTEMPT_ID] == "attempt"
exporter.disable_store_otlp()
assert exporter._rollout_id is None
assert exporter._attempt_id is None