diff --git a/agentlightning/runner/agent.py b/agentlightning/runner/agent.py index 892c3fc23..0f25fd3bd 100644 --- a/agentlightning/runner/agent.py +++ b/agentlightning/runner/agent.py @@ -364,7 +364,7 @@ async def _step_impl(self, next_rollout: AttemptedRollout, raise_on_exception: b await self._trigger_hooks(hook_type="on_rollout_start", agent=agent, runner=self, rollout=next_rollout) start_time = time.time() - with self._tracer.trace_context( + async with self._tracer.trace_context( name=rollout_id, store=store, rollout_id=rollout_id, attempt_id=next_rollout.attempt.attempt_id ): await self._trigger_hooks( diff --git a/agentlightning/runner/legacy.py b/agentlightning/runner/legacy.py index 1e308b8b1..a6f177c7e 100644 --- a/agentlightning/runner/legacy.py +++ b/agentlightning/runner/legacy.py @@ -180,7 +180,7 @@ def run(self) -> bool: # type: ignore except Exception: logger.exception(f"{self._log_prefix(rollout_id)} Exception during on_rollout_start hook.") - with self.tracer.trace_context(name=f"rollout_{rollout_id}"): + with self.tracer._trace_context_sync(name=f"rollout_{rollout_id}"): # pyright: ignore[reportPrivateUsage] start_time = time.time() rollout_method = self.agent.training_rollout if task.mode == "train" else self.agent.validation_rollout # Pass the task input, not the whole task object @@ -257,7 +257,7 @@ async def run_async(self) -> bool: except Exception: logger.exception(f"{self._log_prefix(rollout_id)} Exception during on_rollout_start hook.") - with self.tracer.trace_context(name=f"rollout_{rollout_id}"): + async with self.tracer.trace_context(name=f"rollout_{rollout_id}"): start_time = time.time() rollout_method = ( self.agent.training_rollout_async if task.mode == "train" else self.agent.validation_rollout_async diff --git a/agentlightning/tracer/agentops.py b/agentlightning/tracer/agentops.py index e597c2436..0e03ee862 100644 --- a/agentlightning/tracer/agentops.py +++ b/agentlightning/tracer/agentops.py @@ -6,8 +6,8 @@ import logging import os import threading -from contextlib import contextmanager -from typing import TYPE_CHECKING, Any, Awaitable, Iterator, List, Optional +from contextlib import asynccontextmanager, contextmanager +from typing import TYPE_CHECKING, Any, AsyncGenerator, Awaitable, Iterator, List, Optional import agentops import agentops.sdk.core @@ -153,15 +153,15 @@ def teardown_worker(self, worker_id: int) -> None: self.uninstrument(worker_id) logger.info(f"[Worker {worker_id}] Instrumentation removed.") - @contextmanager - def trace_context( + @asynccontextmanager + async def trace_context( self, name: Optional[str] = None, *, store: Optional[LightningStore] = None, rollout_id: Optional[str] = None, attempt_id: Optional[str] = None, - ) -> Iterator[LightningSpanProcessor]: + ) -> AsyncGenerator[LightningSpanProcessor, None]: """ Starts a new tracing context. This should be used as a context manager. @@ -174,6 +174,21 @@ def trace_context( Yields: The LightningSpanProcessor 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 + + @contextmanager + def _trace_context_sync( + self, + name: Optional[str] = None, + *, + store: Optional[LightningStore] = None, + rollout_id: Optional[str] = None, + attempt_id: Optional[str] = None, + ) -> Iterator[LightningSpanProcessor]: + """Implementation of `trace_context` for synchronous execution.""" if not self._lightning_span_processor: raise RuntimeError("LightningSpanProcessor is not initialized. Call init_worker() first.") diff --git a/agentlightning/tracer/base.py b/agentlightning/tracer/base.py index 7f91c55ab..0f4a6afa0 100644 --- a/agentlightning/tracer/base.py +++ b/agentlightning/tracer/base.py @@ -3,8 +3,7 @@ from __future__ import annotations import logging -from contextlib import contextmanager -from typing import TYPE_CHECKING, Any, Awaitable, Callable, Iterator, List, Optional +from typing import TYPE_CHECKING, Any, AsyncContextManager, Awaitable, Callable, ContextManager, List, Optional from opentelemetry.sdk.trace import ReadableSpan @@ -36,9 +35,9 @@ class Tracer(ParallelWorkerBase): tracer = YourTracerImplementation() try: - with tracer.trace_context(name="my_traced_task"): + async with tracer.trace_context(name="my_traced_task"): # ... code to be traced ... - run_my_agent_logic() + await run_my_agent_logic() except Exception as e: print(f"An error occurred: {e}") @@ -52,7 +51,6 @@ class Tracer(ParallelWorkerBase): ``` """ - @contextmanager def trace_context( self, name: Optional[str] = None, @@ -60,7 +58,7 @@ def trace_context( store: Optional[LightningStore] = None, rollout_id: Optional[str] = None, attempt_id: Optional[str] = None, - ) -> Iterator[Any]: + ) -> AsyncContextManager[Any]: """ Starts a new tracing context. This should be used as a context manager. @@ -79,6 +77,17 @@ def trace_context( """ raise NotImplementedError() + def _trace_context_sync( + self, + name: Optional[str] = None, + *, + store: Optional[LightningStore] = None, + rollout_id: Optional[str] = None, + attempt_id: Optional[str] = None, + ) -> ContextManager[Any]: + """Internal API for CI backward compatibility.""" + raise NotImplementedError() + def get_last_trace(self) -> List[ReadableSpan]: """ Retrieves the raw list of captured spans from the most recent trace. @@ -92,6 +101,8 @@ def trace_run(self, func: Callable[..., Any], *args: Any, **kwargs: Any) -> Any: """ A convenience wrapper to trace the execution of a single synchronous function. + Deprecated in favor of customizing Runners. + Args: func: The synchronous function to execute and trace. *args: Positional arguments to pass to the function. @@ -100,13 +111,15 @@ def trace_run(self, func: Callable[..., Any], *args: Any, **kwargs: Any) -> Any: Returns: The return value of the function. """ - with self.trace_context(name=func.__name__): + with self._trace_context_sync(name=func.__name__): return func(*args, **kwargs) async def trace_run_async(self, func: Callable[..., Awaitable[Any]], *args: Any, **kwargs: Any) -> Any: """ A convenience wrapper to trace the execution of a single asynchronous function. + Deprecated in favor of customizing Runners. + Args: func: The asynchronous function to execute and trace. *args: Positional arguments to pass to the function. @@ -115,7 +128,7 @@ async def trace_run_async(self, func: Callable[..., Awaitable[Any]], *args: Any, Returns: The return value of the function. """ - with self.trace_context(name=func.__name__): + async with self.trace_context(name=func.__name__): return await func(*args, **kwargs) def get_langchain_handler(self) -> Optional[BaseCallbackHandler]: # type: ignore diff --git a/agentlightning/tracer/http.py b/agentlightning/tracer/http.py index 3da12e4a0..6fb28fdfa 100644 --- a/agentlightning/tracer/http.py +++ b/agentlightning/tracer/http.py @@ -5,8 +5,8 @@ import multiprocessing import queue import uuid -from contextlib import contextmanager -from typing import Any, Awaitable, Callable, Dict, Iterator, List, Optional, Tuple +from contextlib import asynccontextmanager, contextmanager +from typing import Any, AsyncGenerator, Awaitable, Callable, Dict, Iterator, List, Optional, Tuple from urllib.parse import urlparse from httpdbg.hooks.all import httprecord @@ -78,8 +78,19 @@ def init_worker(self, worker_id: int) -> None: super().init_worker(worker_id) logger.info(f"[Worker {worker_id}] HttpTracer initialized.") + @asynccontextmanager + async def trace_context(self, name: Optional[str] = None, **kwargs: Any) -> AsyncGenerator[HTTPRecords, None]: + """ + Starts a new HTTP tracing context. This should be used as a context manager. + + Args: + name: Optional name for the tracing context. + """ + with self._trace_context_sync(name=name, **kwargs) as records: + yield records + @contextmanager - def trace_context(self, name: Optional[str] = None, **kwargs: Any) -> Iterator[HTTPRecords]: + def _trace_context_sync(self, name: Optional[str] = None, **kwargs: Any) -> Iterator[HTTPRecords]: """ Starts a new HTTP tracing context. This should be used as a context manager. diff --git a/agentlightning/tracer/otel.py b/agentlightning/tracer/otel.py index b786b6b2a..91a5c606c 100644 --- a/agentlightning/tracer/otel.py +++ b/agentlightning/tracer/otel.py @@ -3,8 +3,8 @@ from __future__ import annotations import logging -from contextlib import contextmanager -from typing import Iterator, List, Optional +from contextlib import asynccontextmanager +from typing import AsyncGenerator, List, Optional import opentelemetry.trace as trace_api from opentelemetry.sdk.trace import ReadableSpan, TracerProvider @@ -49,15 +49,15 @@ def teardown_worker(self, worker_id: int): logger.info(f"[Worker {worker_id}] Tearing down OpenTelemetry tracer...") self._tracer_provider = None - @contextmanager - def trace_context( + @asynccontextmanager + async def trace_context( self, name: Optional[str] = None, *, store: Optional[LightningStore] = None, rollout_id: Optional[str] = None, attempt_id: Optional[str] = None, - ) -> Iterator[LightningSpanProcessor]: + ) -> AsyncGenerator[LightningSpanProcessor, None]: """ Starts a new tracing context. This should be used as a context manager. diff --git a/tests/runner/test_agent_runner.py b/tests/runner/test_agent_runner.py index fc12c7bb1..62ae72ad3 100644 --- a/tests/runner/test_agent_runner.py +++ b/tests/runner/test_agent_runner.py @@ -2,8 +2,8 @@ import asyncio import random -from contextlib import contextmanager -from typing import Any, Dict, Iterator, List, Optional, Sequence, cast +from contextlib import asynccontextmanager +from typing import Any, AsyncGenerator, Dict, List, Optional, Sequence, cast import pytest from opentelemetry import trace as trace_api @@ -79,15 +79,15 @@ def teardown(self, *args: Any, **kwargs: Any) -> None: def get_last_trace(self) -> List[ReadableSpan]: return list(self._last_trace) - @contextmanager - def trace_context( + @asynccontextmanager + async def trace_context( self, name: Optional[str] = None, *, store: Optional[LightningStore] = None, rollout_id: Optional[str] = None, attempt_id: Optional[str] = None, - ) -> Iterator[List[ReadableSpan]]: + ) -> AsyncGenerator[List[ReadableSpan], None]: previous = self._contexts[-1] if self._contexts else None current = { "name": name, diff --git a/tests/runner/test_runner_context.py b/tests/runner/test_runner_context.py index b0a195705..97de82ef5 100644 --- a/tests/runner/test_runner_context.py +++ b/tests/runner/test_runner_context.py @@ -1,7 +1,7 @@ # Copyright (c) Microsoft. All rights reserved. -from contextlib import contextmanager -from typing import Any, Dict, Iterator, List, Optional +from contextlib import asynccontextmanager +from typing import Any, AsyncGenerator, Dict, List, Optional import pytest from opentelemetry import trace as trace_api @@ -54,15 +54,15 @@ def teardown_worker(self, worker_id: int, *args: Any, **kwargs: Any) -> None: def get_last_trace(self) -> List[ReadableSpan]: return list(self._last_trace) - @contextmanager - def trace_context( + @asynccontextmanager + async def trace_context( self, name: Optional[str] = None, *, store: Optional[LightningStore] = None, rollout_id: Optional[str] = None, attempt_id: Optional[str] = None, - ) -> Iterator[List[ReadableSpan]]: + ) -> AsyncGenerator[List[ReadableSpan], None]: self._last_trace = [] try: yield self._last_trace diff --git a/tests/tracer/test_integration.py b/tests/tracer/test_integration.py index 1d88c3522..6223c964f 100644 --- a/tests/tracer/test_integration.py +++ b/tests/tracer/test_integration.py @@ -78,7 +78,7 @@ OPENAI_MODEL = "gpt-4.1-mini" OPENAI_API_KEY = os.environ["OPENAI_API_KEY"] else: - OPENAI_BASE_URL = "http://127.0.0.1:8000/v1" + OPENAI_BASE_URL = "http://127.0.0.1:58000/v1" OPENAI_MODEL = "gpt-4.1-mini" OPENAI_API_KEY = "token-abc123" @@ -101,7 +101,7 @@ class MockOpenAICompatibleServer: Now supports replaying from prompt caches. """ - def __init__(self, host: str = "127.0.0.1", port: int = 8000) -> None: + def __init__(self, host: str = "127.0.0.1", port: int = 58000) -> None: self.host = host self.port = port self.app = FastAPI() @@ -767,7 +767,7 @@ def create_prompt_caches() -> None: if USE_OPENAI: tracer = HttpTracer() - with tracer.trace_context(): + with tracer._trace_context_sync(): run_all() with open(os.path.join(os.path.dirname(__file__), "../assets/prompt_caches.jsonl"), "w") as f: diff --git a/tests/tracer/test_otel.py b/tests/tracer/test_otel.py index 211080519..bdb13e407 100644 --- a/tests/tracer/test_otel.py +++ b/tests/tracer/test_otel.py @@ -412,6 +412,10 @@ def test_context_manager_reusability(): def _otel_reward_subprocess(mode: str, conn: Connection[tuple[str, Any]]) -> None: + asyncio.run(_otel_reward_subprocess_async(mode, conn)) + + +async def _otel_reward_subprocess_async(mode: str, conn: Connection[tuple[str, Any]]) -> None: tracer: OtelTracer | None = None try: try: @@ -431,14 +435,14 @@ def _otel_reward_subprocess(mode: str, conn: Connection[tuple[str, Any]]) -> Non def compute_reward() -> float: return expected_reward - with tracer.trace_context(name="reward-decorator"): + async with tracer.trace_context(name="reward-decorator"): returned = compute_reward() if returned != expected_reward: raise AssertionError(f"Expected reward {expected_reward}, got {returned}") elif mode == "emit": expected_reward = 4.5 - with tracer.trace_context(name="reward-emit"): + async with tracer.trace_context(name="reward-emit"): emit_reward(expected_reward) else: raise ValueError(f"Unsupported mode: {mode}")