Make tracer.trace_context async

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(

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

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.")
 

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,15 +51,14 @@ class Tracer(ParallelWorkerBase):
     ```
     """
 
-    @contextmanager
     def trace_context(
         self,
         name: Optional[str] = None,
         *,
         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

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.
 

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.
 

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,

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

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:

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}")