Upgrade SQL Agent Example to v0.2

.github/workflows/examples.yml

@@ -145,9 +145,8 @@ jobs:
           set -ex
           . .venv/bin/activate
           cd examples/spider
-          python sql_agent.py --trainer.n-workers 1 --trainer.dev true --trainer.max-tasks 2
+          python sql_agent.py
         env:
-          VERL_API_BASE: http://localhost:9999/
           OPENAI_API_BASE: http://localhost:12306/
           OPENAI_API_KEY: dummy
         if: success() || failure()
@@ -245,18 +244,10 @@ jobs:
           cd examples/spider
           ../../scripts/restart_ray.sh
           sleep 5
-          PYTHONUNBUFFERED=1 python sql_agent.py --trainer.n-workers 10 &
-          bash train_ci.sh
-          pkill -f sql_agent.py && echo "SIGTERM sent to sql_agent.py" || echo "No sql_agent.py process found"
-          while pgrep -f sql_agent.py; do
-            echo "Waiting for sql_agent.py to finish..."
-            sleep 5
-          done
-          echo "sql_agent.py has finished."
+          PYTHONUNBUFFERED=1 python train_sql_agent.py fast
           sleep 10
         shell: bash
         env:
-          VERL_API_BASE: http://localhost:9991/
           WANDB_BASE_URL: ${{ secrets.MSR_WANDB_BASE_URL }}
           WANDB_API_KEY: ${{ secrets.MSR_WANDB_API_KEY }}
         id: spider_train

agentlightning/algorithm/fast.py

@@ -5,10 +5,10 @@
 import asyncio
 import logging
 from datetime import datetime
-from typing import Any, List, Optional
+from typing import Any, List, Literal, Optional
 
 from agentlightning.llm_proxy import ModelConfig
-from agentlightning.types import Dataset, Rollout, RolloutStatus
+from agentlightning.types import Attempt, Dataset, Rollout, RolloutStatus, Span
 
 from .base import Algorithm
 
@@ -52,17 +52,38 @@ def __init__(
         train_split: float = 0.5,
         polling_interval: float = 5.0,
         max_queue_length: int = 4,
+        span_verbosity: Literal["keys", "key_values", "none"] = "keys",
     ) -> None:
         super().__init__()
         self.n_epochs = n_epochs
         self.train_split = train_split
         self.polling_interval = polling_interval
         self.max_queue_length = max_queue_length
+        self.span_verbosity = span_verbosity
         if not (0.0 < self.train_split < 1.0):
             raise ValueError("train_split must be between 0 and 1.")
 
         self._finished_rollout_count = 0
 
+    def _span_to_string(self, rollout_id: str, attempt: Attempt, span: Span) -> str:
+        if self.span_verbosity == "none":
+            return ""
+
+        prefix_msg = f"[Rollout {rollout_id} | Attempt {attempt.attempt_id} | Span {span.span_id}] #{span.sequence_id} ({span.name}) "
+        elapsed = f"{span.end_time - span.start_time:.2f}" if span.start_time and span.end_time else "unknown"
+
+        msg = (
+            prefix_msg
+            + f"From {_timestamp_to_iso_str(span.start_time) if span.start_time else 'unknown'}, "
+            + f"to {_timestamp_to_iso_str(span.end_time) if span.end_time else 'unknown'}, "
+            + f"{elapsed} seconds. "
+        )
+        if self.span_verbosity == "key_values":
+            msg += f"Attributes: {span.attributes}"
+        else:
+            msg += f"Attribute keys: {list(span.attributes.keys())}"
+        return msg
+
     async def _handle_rollout_finish(self, rollout: Rollout) -> None:
         store = self.get_store()
 
@@ -80,14 +101,8 @@ async def _handle_rollout_finish(self, rollout: Rollout) -> None:
             )
             spans = await store.query_spans(rollout_id=rollout_id)
             for span in spans:
-                prefix_msg = f"[Rollout {rollout_id} | Attempt {attempt.attempt_id} | Span {span.span_id}] #{span.sequence_id} ({span.name}) "
-                elapsed = f"{span.end_time - span.start_time:.2f}" if span.start_time and span.end_time else "unknown"
-                logger.info(
-                    prefix_msg
-                    + f"From {_timestamp_to_iso_str(span.start_time) if span.start_time else 'unknown'}, "
-                    + f"to {_timestamp_to_iso_str(span.end_time) if span.end_time else 'unknown'}, "
-                    + f"{elapsed} seconds. Attributes: {span.attributes}"
-                )
+                if self.span_verbosity != "none":
+                    logger.info(self._span_to_string(rollout.rollout_id, attempt, span))
 
         # Attempts to adapt the spans using the adapter if provided
         try:
@@ -158,6 +173,8 @@ async def run(
         ]
         train_indices = list(range(0, train_dataset_length))
         val_indices = list(range(train_dataset_length, train_dataset_length + val_dataset_length))
+        logger.debug(f"Train indices: {train_indices}")
+        logger.debug(f"Val indices: {val_indices}")
 
         store = self.get_store()
 
@@ -175,18 +192,24 @@ async def run(
             harvest_tasks: List[asyncio.Task[None]] = []
             logger.info(f"Proceeding epoch {epoch + 1}/{self.n_epochs}.")
             for index in train_indices + val_indices:
-                queuing_rollouts = await store.query_rollouts(status=["queuing", "requeuing"])
-                if len(queuing_rollouts) <= self.max_queue_length:
-                    # Only enqueue a new rollout when there is at most "max_queue_length" rollout in the queue.
-                    sample = concatenated_dataset[index]
-                    mode = "train" if index in train_indices else "val"
-                    rollout = await store.enqueue_rollout(input=sample, mode=mode, resources_id=resources_id)
-                    harvest_tasks.append(asyncio.create_task(self._harvest_rollout_spans(rollout.rollout_id)))
-                    logger.info(f"Enqueued rollout {rollout.rollout_id} in {mode} mode with sample: {sample}")
-                else:
-                    # Sleep a bit and try again later.
-                    await asyncio.sleep(self.polling_interval)
+                logger.info(
+                    f"Processing index {index}. {len(train_indices)} train indices and {len(val_indices)} val indices in total."
+                )
+                while True:
+                    queuing_rollouts = await store.query_rollouts(status=["queuing", "requeuing"])
+                    if len(queuing_rollouts) <= self.max_queue_length:
+                        # Only enqueue a new rollout when there is at most "max_queue_length" rollout in the queue.
+                        sample = concatenated_dataset[index]
+                        mode = "train" if index in train_indices else "val"
+                        rollout = await store.enqueue_rollout(input=sample, mode=mode, resources_id=resources_id)
+                        harvest_tasks.append(asyncio.create_task(self._harvest_rollout_spans(rollout.rollout_id)))
+                        logger.info(f"Enqueued rollout {rollout.rollout_id} in {mode} mode with sample: {sample}")
+                        break
+                    else:
+                        # Sleep a bit and try again later.
+                        await asyncio.sleep(self.polling_interval)
 
             # Wait for all harvest tasks to complete
+            logger.info(f"Waiting for {len(harvest_tasks)} harvest tasks to complete...")
             if len(harvest_tasks) > 0:
                 await asyncio.gather(*harvest_tasks)

agentlightning/config.py

@@ -309,6 +309,9 @@ def lightning_cli(cls1: Type[_C1], cls2: Type[_C2], cls3: Type[_C3], cls4: Type[
 def lightning_cli(*classes: Type[CliConfigurable]) -> Tuple[CliConfigurable, ...]: ...
 
 
+# FIXME: lightning_cli needs to be fixed to comply with the latest trainer implementation.
+
+
 def lightning_cli(*classes: Type[CliConfigurable]) -> CliConfigurable | Tuple[CliConfigurable, ...]:  # type: ignore
     """
     Parses command-line arguments to configure and instantiate provided CliConfigurable classes.

agentlightning/litagent/litagent.py

@@ -119,7 +119,10 @@ def get_tracer(self) -> Tracer:
         Returns:
             The Tracer instance associated with this agent.
         """
-        return self.trainer.tracer
+        if hasattr(self.runner, "tracer"):
+            return self.runner.tracer  # type: ignore
+        else:
+            return self.trainer.tracer
 
     @property
     def tracer(self) -> Tracer:

agentlightning/runner/agent.py

@@ -139,6 +139,15 @@ def teardown_worker(self, worker_id: int, *args: Any, **kwargs: Any) -> None:
 
         self._tracer.teardown_worker(worker_id)
 
+    @property
+    def tracer(self) -> Tracer:
+        """Get the tracer instance.
+
+        Returns:
+            The Tracer instance used by this runner.
+        """
+        return self._tracer
+
     def get_agent(self) -> LitAgent[T_task]:
         """Get the agent instance.
 

agentlightning/store/client_server.py

@@ -544,7 +544,8 @@ async def _get_session(self) -> aiohttp.ClientSession:
         with self._lock:
             sess = self._sessions.get(key)
             if sess is None or sess.closed:
-                sess = aiohttp.ClientSession()
+                timeout = aiohttp.ClientTimeout(total=30.0, connect=5.0, sock_connect=5.0, sock_read=30.0)
+                sess = aiohttp.ClientSession(timeout=timeout)
                 self._sessions[key] = sess
         return sess
 
@@ -608,7 +609,7 @@ async def _request_json(
             except aiohttp.ClientResponseError as cre:
                 # Respect app-level 4xx as final (server marks app faults as 400)
                 # 4xx => application issue; do not retry (except 408 which is transient)
-                logger.exception(f"ClientResponseError: {cre.status} {cre.message}")
+                logger.debug(f"ClientResponseError: {cre.status} {cre.message}", exc_info=True)
                 if 400 <= cre.status < 500 and cre.status != 408:
                     raise
                 # 5xx and others will be retried below if they raise
@@ -624,7 +625,7 @@ async def _request_json(
                 asyncio.TimeoutError,
             ) as net_exc:
                 # Network/session issue: probe health before retrying
-                logger.exception(f"Network/session issue: {net_exc}")
+                logger.debug(f"Network/session issue: {net_exc}", exc_info=True)
                 last_exc = net_exc
                 logger.info(f"Network/session issue will be retried. Retrying the request {method}: {path}")
                 if not await self._wait_until_healthy(session):

agentlightning/tracer/agentops.py

@@ -13,6 +13,7 @@
 import agentops.sdk.core
 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 agentlightning.instrumentation import instrument_all, uninstrument_all
@@ -197,7 +198,7 @@ def get_last_trace(self) -> List[ReadableSpan]:
             raise RuntimeError("LightningSpanProcessor is not initialized. Call init_worker() first.")
         return self._lightning_span_processor.spans()
 
-    def get_langchain_callback_handler(self, tags: List[str] | None = None) -> LangchainCallbackHandler:
+    def get_langchain_handler(self, tags: List[str] | None = None) -> LangchainCallbackHandler:
         """
         Get the Langchain callback handler for integrating with Langchain.
 
@@ -221,6 +222,8 @@ def get_langchain_callback_handler(self, tags: List[str] | None = None) -> Langc
             )
         return LangchainCallbackHandler(api_key=api_key, tags=tags)
 
+    get_langchain_callback_handler = get_langchain_handler  # alias
+
 
 class LightningSpanProcessor(SpanProcessor):
     def __init__(self):
@@ -261,6 +264,32 @@ def _await_in_loop(self, coro: Awaitable[Any], timeout: Optional[float] = None)
         # 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
 
@@ -313,10 +342,11 @@ def on_end(self, span: ReadableSpan) -> None:
         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
-                self._await_in_loop(
-                    self._store.add_otel_span(self._rollout_id, self._attempt_id, span),
-                    timeout=5.0,
-                )
+                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}")

agentlightning/tracer/base.py

@@ -1,14 +1,19 @@
 # Copyright (c) Microsoft. All rights reserved.
 
+from __future__ import annotations
+
 import logging
 from contextlib import contextmanager
-from typing import Any, Awaitable, Callable, Iterator, List, Optional
+from typing import TYPE_CHECKING, Any, Awaitable, Callable, Iterator, List, Optional
 
 from opentelemetry.sdk.trace import ReadableSpan
 
 from agentlightning.store.base import LightningStore
 from agentlightning.types import ParallelWorkerBase
 
+if TYPE_CHECKING:
+    from langchain.callbacks.base import BaseCallbackHandler  # type: ignore
+
 logger = logging.getLogger(__name__)
 
 
@@ -112,3 +117,11 @@ async def trace_run_async(self, func: Callable[..., Awaitable[Any]], *args: Any,
         """
         with self.trace_context(name=func.__name__):
             return await func(*args, **kwargs)
+
+    def get_langchain_handler(self) -> Optional[BaseCallbackHandler]:  # type: ignore
+        """Get a handler to install in langchain agent callback.
+
+        Agents are expected to use this handler in their agents to enable tracing.
+        """
+        logger.warning(f"{self.__class__.__name__} does not provide a LangChain callback handler.")
+        return None

examples/spider/spider_eval/__init__.py

@@ -0,0 +1 @@
+# Copyright (c) Microsoft. All rights reserved.

examples/spider/spider_eval/async_utils.py

@@ -0,0 +1,38 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+import queue
+import threading
+from typing import Any, Coroutine
+
+
+def run_sync_ephemeral(coro: Coroutine[Any, Any, Any]) -> Any:
+    """
+    Run an async coroutine from sync code.
+    - If no loop in this thread: use asyncio.run() directly.
+    - If already in an event loop: spawn a worker thread that calls asyncio.run()
+      (which creates and closes a brand-new event loop per call).
+    """
+    try:
+        asyncio.get_running_loop()
+    except RuntimeError:
+        # No running loop in this thread; safe to use asyncio.run
+        return asyncio.run(coro)
+
+    # Already in a running loop -> execute in a worker thread
+    q = queue.Queue[Any]()
+
+    def worker():
+        try:
+            result = asyncio.run(coro)  # creates & closes its own loop
+            q.put((True, result))
+        except BaseException as e:
+            q.put((False, e))
+
+    t = threading.Thread(target=worker, daemon=True)
+    t.start()
+    ok, payload = q.get()
+    t.join()
+    if ok:
+        return payload
+    raise payload

examples/spider/spider_eval/exec_eval.py

@@ -18,6 +18,7 @@
 
 import tqdm
 
+from .async_utils import run_sync_ephemeral
 from .parse import get_all_preds_for_execution, remove_distinct
 
 threadLock = threading.Lock()
@@ -225,8 +226,8 @@ def eval_exec_match(
             ranger = db_paths
 
         for db_path in ranger:
-            g_flag, g_denotation = asyncio.run(exec_on_db(db_path, g_str))
-            p_flag, p_denotation = asyncio.run(exec_on_db(db_path, pred))
+            g_flag, g_denotation = run_sync_ephemeral(exec_on_db(db_path, g_str))
+            p_flag, p_denotation = run_sync_ephemeral(exec_on_db(db_path, pred))
 
             # we should expect the gold to be succesfully executed on the database
             assert g_flag != "exception", "gold query %s has error on database file %s" % (g_str, db_path)