Put store init in the right place of tracer (#321)

Author: ultmaster

agentlightning/runner/agent.py

@@ -138,7 +138,7 @@ def init_worker(self, worker_id: int, store: LightningStore, **kwargs: Any) -> N
         self._store = store
         self.worker_id = worker_id
 
-        self._tracer.init_worker(worker_id)
+        self._tracer.init_worker(worker_id, store)
 
     def teardown(self, *args: Any, **kwargs: Any) -> None:
         """Teardown the runner and clean up all resources.
@@ -469,7 +469,7 @@ async def _step_impl(self, next_rollout: AttemptedRollout, raise_on_exception: b
 
             start_time = time.time()
             async with self._tracer.trace_context(
-                name=rollout_id, store=store, rollout_id=rollout_id, attempt_id=next_rollout.attempt.attempt_id
+                name=rollout_id, rollout_id=rollout_id, attempt_id=next_rollout.attempt.attempt_id
             ):
                 await self._trigger_hooks(
                     hook_type="on_trace_start", agent=agent, runner=self, tracer=self._tracer, rollout=next_rollout

agentlightning/tracer/agentops.py

@@ -4,6 +4,7 @@
 
 import logging
 import os
+import warnings
 from contextlib import asynccontextmanager, contextmanager
 from typing import TYPE_CHECKING, Any, AsyncGenerator, Iterator, List, Optional
 
@@ -114,6 +115,14 @@ async def trace_context(
         Yields:
             The OpenTelemetry tracer instance to collect spans.
         """
+        if store is not None:
+            warnings.warn(
+                "store is deprecated in favor of init_worker(). It will be removed in the future.",
+                DeprecationWarning,
+                stacklevel=3,
+            )
+        else:
+            store = self._store
         with self._trace_context_sync(name=name, store=store, rollout_id=rollout_id, attempt_id=attempt_id) as tracer:
             yield tracer
 
@@ -164,9 +173,10 @@ def _agentops_trace_context(self, rollout_id: Optional[str], attempt_id: Optiona
         try:
             yield
         except Exception as e:
-            # TODO: I'm not sure whether this will catch errors in user code.
+            # 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}: {e}")
+            raise  # should reraise the error here so that runner can handle it
         finally:
             agentops.end_trace(trace, end_state=status)  # type: ignore
 

agentlightning/tracer/base.py

@@ -52,6 +52,18 @@ class Tracer(ParallelWorkerBase):
     ```
     """
 
+    _store: Optional[LightningStore] = None
+
+    def init_worker(self, worker_id: int, store: Optional[LightningStore] = None) -> None:
+        """Initialize the tracer for a worker.
+
+        Args:
+            worker_id: The ID of the worker.
+            store: The store to add the spans to. If it's provided, traces will be added to the store when tracing.
+        """
+        super().init_worker(worker_id)
+        self._store = store
+
     def trace_context(
         self,
         name: Optional[str] = None,
@@ -68,11 +80,9 @@ def trace_context(
         within the `with` block are collected and made available via
         [`get_last_trace`][agentlightning.Tracer.get_last_trace].
 
-        If a store is provided, the spans will be added to the store when tracing.
-
         Args:
             name: The name for the root span of this trace context.
-            store: The store to add the spans to.
+            store: The store to add the spans to. Deprecated in favor of passing store to init_worker().
             rollout_id: The rollout ID to add the spans to.
             attempt_id: The attempt ID to add the spans to.
         """
@@ -82,7 +92,6 @@ 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]:
@@ -141,19 +150,22 @@ def get_langchain_handler(self) -> Optional[BaseCallbackHandler]:  # type: ignor
         return None
 
     @contextmanager
-    def lifespan(self):
+    def lifespan(self, store: Optional[LightningStore] = None):
         """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.
+
+        Args:
+            store: The store to add the spans to. If it's provided, traces will be added to the store when tracing.
         """
         has_init = False
         has_init_worker = False
         try:
             self.init()
             has_init = True
 
-            self.init_worker(0)
+            self.init_worker(0, store)
             has_init_worker = True
 
             yield

agentlightning/tracer/http.py

@@ -19,6 +19,8 @@
     TraceState,
 )
 
+from agentlightning.store import LightningStore
+
 from .base import Tracer
 
 logger = logging.getLogger(__name__)
@@ -68,14 +70,15 @@ def __init__(
         self.subprocess_mode = subprocess_mode
         self.subprocess_timeout = subprocess_timeout
 
-    def init_worker(self, worker_id: int) -> None:
+    def init_worker(self, worker_id: int, store: Optional[LightningStore] = None) -> None:
         """
         Initialize the tracer in a worker process.
 
         Args:
             worker_id: The ID of the worker process.
+            store: The store to add the spans to.
         """
-        super().init_worker(worker_id)
+        super().init_worker(worker_id, store)
         logger.info(f"[Worker {worker_id}] HttpTracer initialized.")
 
     @asynccontextmanager

agentlightning/tracer/otel.py

@@ -5,6 +5,7 @@
 import asyncio
 import logging
 import threading
+import warnings
 from contextlib import asynccontextmanager
 from typing import Any, AsyncGenerator, Awaitable, List, Optional
 
@@ -42,8 +43,8 @@ def __init__(self):
         self._otlp_span_exporter: Optional[LightningStoreOTLPExporter] = None
         self._initialized: bool = False
 
-    def init_worker(self, worker_id: int):
-        super().init_worker(worker_id)
+    def init_worker(self, worker_id: int, store: Optional[LightningStore] = None):
+        super().init_worker(worker_id, store)
         self._initialize_tracer_provider(worker_id)
 
     def _initialize_tracer_provider(self, worker_id: int):
@@ -92,7 +93,18 @@ async def trace_context(
         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 is not None:
+            warnings.warn(
+                "store is deprecated in favor of init_worker(). It will be removed in the future.",
+                DeprecationWarning,
+                stacklevel=3,
+            )
+        else:
+            store = self._store
+
+        if rollout_id is not None and attempt_id is not None:
+            if store is None:
+                raise ValueError("store is required to be initialized when rollout_id and attempt_id are provided")
             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)
@@ -101,12 +113,12 @@ async def trace_context(
             ctx = self._lightning_span_processor.with_context(store=store, rollout_id=rollout_id, attempt_id=attempt_id)
             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:
+        elif rollout_id is None and attempt_id is None:
             self._disable_native_otlp_exporter()
             with 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")
+            raise ValueError("rollout_id and attempt_id must be either all provided or all None")
 
     def get_last_trace(self) -> List[ReadableSpan]:
         """

examples/minimal/write_traces.py

@@ -34,7 +34,7 @@ async def send_traces_via_otel(use_client: bool = False):
         store = LightningStoreClient("http://localhost:45993")
     rollout = await store.start_rollout(input={"origin": "write_traces_example"})
 
-    with tracer.lifespan():
+    with tracer.lifespan(store):
         # 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
@@ -89,10 +89,10 @@ async def send_traces_via_agentops(use_client: bool = False):
 
     # Initialize the tracer lifespan
     # One lifespan can contain multiple traces
-    with tracer.lifespan():
+    with tracer.lifespan(store):
         # 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
+            "trace-1", rollout_id=rollout.rollout_id, attempt_id=rollout.attempt.attempt_id
         ):
             openai_client = AsyncOpenAI()
             response = await openai_client.chat.completions.create(

examples/tinker/tests/test_tinker_llm.py

@@ -56,7 +56,7 @@ async def test_tracer():
     try:
         tracer = AgentOpsTracer()
         tracer.init()
-        tracer.init_worker(0)
+        tracer.init_worker(worker_id=0, store=store)
 
         # init tracer before llm_proxy to avoid tracer provider being not active.
         console.print("Starting LLM proxy...")
@@ -70,7 +70,7 @@ async def test_tracer():
         client = openai.OpenAI(base_url="http://localhost:4000/v1", api_key="dummy")
 
         async with tracer.trace_context(
-            name="test_llm", store=store, rollout_id=rollout.rollout_id, attempt_id=rollout.attempt.attempt_id
+            name="test_llm", rollout_id=rollout.rollout_id, attempt_id=rollout.attempt.attempt_id
         ):
             response = client.chat.completions.create(
                 model=model_name,

tests/tracer/test_agentops.py

@@ -1,9 +1,11 @@
 # Copyright (c) Microsoft. All rights reserved.
 
 import multiprocessing
+import sys
 from typing import Any, Optional, Union
 
 import agentops
+import pytest
 from agentops.sdk.core import TraceContext
 from opentelemetry.trace.status import StatusCode
 
@@ -20,7 +22,8 @@ def _func_without_exception():
     pass
 
 
-def test_trace_error_status_from_instance():
+@pytest.mark.parametrize("with_exception", [True, False])
+def test_trace_error_status_from_instance(with_exception: bool):
     """
     Test that AgentOpsTracer correctly sets trace end state based on execution result.
 
@@ -30,7 +33,7 @@ def test_trace_error_status_from_instance():
     """
 
     ctx = multiprocessing.get_context("spawn")
-    proc = ctx.Process(target=_test_trace_error_status_from_instance_imp)
+    proc = ctx.Process(target=_test_trace_error_status_from_instance_imp, args=(with_exception,))
     proc.start()
     proc.join(30.0)  # On GPU server, the time is around 10 seconds.
 
@@ -42,13 +45,19 @@ def test_trace_error_status_from_instance():
 
         assert False, "Child process hung. Check test output for details."
 
-    assert proc.exitcode == 0, (
-        f"Child process for test_trace_error_status_from_instance failed with exit code {proc.exitcode}. "
-        "Check child traceback in test output."
-    )
+    if with_exception:
+        assert proc.exitcode != 0, (
+            f"Child process for test_trace_error_status_from_instance with exception exited with exit code {proc.exitcode}. "
+            "Check child traceback in test output."
+        )
+    else:
+        assert proc.exitcode == 0, (
+            f"Child process for test_trace_error_status_from_instance without exception failed with exit code {proc.exitcode}. "
+            "Check child traceback in test output."
+        )
 
 
-def _test_trace_error_status_from_instance_imp():
+def _test_trace_error_status_from_instance_imp(with_exception: bool):
     captured_state = {}
     old_end_trace = agentops.end_trace
 
@@ -65,12 +74,14 @@ def custom_end_trace(
     tracer.init_worker(0)
 
     try:
-        tracer.trace_run(_func_with_exception)
-        assert captured_state["state"] == StatusCode.ERROR
-
-        tracer.trace_run(_func_without_exception)
-        assert captured_state["state"] == StatusCode.OK
-
+        if with_exception:
+            tracer.trace_run(_func_with_exception)
+            if captured_state["state"] != StatusCode.ERROR:
+                sys.exit(-1)
+        else:
+            tracer.trace_run(_func_without_exception)
+            if captured_state["state"] != StatusCode.OK:
+                sys.exit(-1)
     finally:
         agentops.end_trace = old_end_trace
         tracer.teardown_worker(0)