update

a.py

@@ -0,0 +1,22 @@
+import asyncio
+
+
+async def a():
+    print("a")
+    b()
+    print("finish")
+
+
+def b():
+    print("b")
+    loop = asyncio.get_running_loop()
+    fut = asyncio.run_coroutine_threadsafe(c(), loop)
+    fut.result(timeout=5.0)
+
+
+async def c():
+    print("c")
+    await asyncio.sleep(0.1)
+
+
+asyncio.run(a())

agentlightning/algorithm/fast.py

@@ -5,10 +5,10 @@ from __future__ import annotations
 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 BaseAlgorithm
 
@@ -52,17 +52,38 @@ class Baseline(FastAlgorithm):
         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 @@ class Baseline(FastAlgorithm):
             )
             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:
@@ -188,5 +203,6 @@ class Baseline(FastAlgorithm):
                     await asyncio.sleep(self.polling_interval)
 
             # Wait for all harvest tasks to complete
+            print(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 @@ class LitAgent(Generic[T]):
         Returns:
             The BaseTracer 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) -> BaseTracer:

agentlightning/runner/agent.py

@@ -139,6 +139,15 @@ class LitAgentRunner(BaseRunner[T_task]):
 
         self._tracer.teardown_worker(worker_id)
 
+    @property
+    def tracer(self) -> BaseTracer:
+        """Get the tracer instance.
+
+        Returns:
+            The BaseTracer instance used by this runner.
+        """
+        return self._tracer
+
     def get_agent(self) -> LitAgent[T_task]:
         """Get the agent instance.
 

agentlightning/store/client_server.py

@@ -335,10 +335,12 @@ class LightningStoreServer(LightningStore):
 
         @self.app.post("/add_span", response_model=Span)
         async def add_span(span: Span):  # pyright: ignore[reportUnusedFunction]
+            print("!!!!! add_span received")
             return await self.store.add_span(span)
 
         @self.app.get("/get_next_span_sequence_id/{rollout_id}/{attempt_id}", response_model=int)
         async def get_next_span_sequence_id(rollout_id: str, attempt_id: str):  # pyright: ignore[reportUnusedFunction]
+            print("!!!!! get_next_span_sequence_id received")
             return await self.store.get_next_span_sequence_id(rollout_id, attempt_id)
 
         @self.app.post("/wait_for_rollouts", response_model=List[Rollout])
@@ -491,6 +493,18 @@ class LightningStoreServer(LightningStore):
         )
 
 
+# def _make_trace():
+#     tc = aiohttp.TraceConfig()
+#     # async def log_evt(session, context, params):
+#     #     print(f"[TRACE] {context}: {params}")
+#     tc.on_dns_resolvehost_start.append(lambda *a, **k: print("[TRACE] dns_start", a[-1].host))
+#     tc.on_connection_create_start.append(lambda *a, **k: print("[TRACE] conn_start"))
+#     tc.on_request_start.append(lambda *a, **k: print("[TRACE] req_start", a[-1].method, a[-1].url))
+#     tc.on_request_end.append(lambda *a, **k: print("[TRACE] req_end", a[-1].method, a[-1].url))
+#     tc.on_request_exception.append(lambda *a, **k: print("[TRACE] req_exc", a[-1].method, a[-1].url))
+#     return tc
+
+
 class LightningStoreClient(LightningStore):
     """HTTP client that talks to a remote LightningStoreServer.
 
@@ -541,11 +555,18 @@ class LightningStoreClient(LightningStore):
 
         loop = asyncio.get_running_loop()
         key = id(loop)
+        print("!!!!! _get_session received %s", key)
         with self._lock:
+            print("!!!!! _get_session with lock")
             sess = self._sessions.get(key)
             if sess is None or sess.closed:
+                # connector = aiohttp.TCPConnector(
+                #     limit=64, limit_per_host=16, ttl_dns_cache=300, enable_cleanup_closed=True
+                # )
+                # sess = aiohttp.ClientSession(trace_configs=[_make_trace()])
                 sess = aiohttp.ClientSession()
                 self._sessions[key] = sess
+        print(self._sessions)
         return sess
 
     async def _wait_until_healthy(self, session: aiohttp.ClientSession) -> bool:
@@ -591,6 +612,7 @@ class LightningStoreClient(LightningStore):
         """
         session = await self._get_session()
         url = f"{self.server_address}{path if path.startswith('/') else '/'+path}"
+        print("$$$$$$ session acquired", url)
 
         # attempt 0 is immediate, then follow retry schedule
         attempts = (0.0,) + self._retry_delays
@@ -602,13 +624,16 @@ class LightningStoreClient(LightningStore):
                 await asyncio.sleep(delay)
             try:
                 http_call = getattr(session, method)
-                async with http_call(url, json=json) as resp:
+                print("$$$$$$ http_call", http_call)
+                timeout = aiohttp.ClientTimeout(total=3.5, connect=1.0, sock_connect=1.0, sock_read=2.5)
+                async with http_call(url, json=json, timeout=timeout) as resp:
+                    print("$$$$$ resp", resp)
                     resp.raise_for_status()
                     return await resp.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 +649,7 @@ class LightningStoreClient(LightningStore):
                 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):
@@ -832,6 +857,7 @@ class LightningStoreClient(LightningStore):
             return None
 
     async def add_span(self, span: Span) -> Span:
+        print("$$$$$$ add_span received")
         data = await self._request_json("post", "/add_span", json=span.model_dump(mode="json"))
         return Span.model_validate(data)
 
@@ -848,6 +874,7 @@ class LightningStoreClient(LightningStore):
         sequence_id: int | None = None,
     ) -> Span:
         # unchanged logic, now benefits from retries inside add_span/get_next_span_sequence_id
+        print("$$$$$$ add_otel_span received")
         if sequence_id is None:
             sequence_id = await self.get_next_span_sequence_id(rollout_id, attempt_id)
         span = Span.from_opentelemetry(
@@ -856,6 +883,7 @@ class LightningStoreClient(LightningStore):
             attempt_id=attempt_id,
             sequence_id=sequence_id,
         )
+        print("$$$$$$ span created")
         await self.add_span(span)
         return span
 

agentlightning/tracer/agentops.py

@@ -197,7 +197,7 @@ class AgentOpsTracer(BaseTracer):
             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 +221,43 @@ class AgentOpsTracer(BaseTracer):
             )
         return LangchainCallbackHandler(api_key=api_key, tags=tags)
 
+    get_langchain_callback_handler = get_langchain_handler  # alias
+
+
+async def heartbeat(name="exporter-loop", period=0.5):
+    import asyncio
+    import time
+
+    last = time.perf_counter()
+    while True:
+        await asyncio.sleep(period)
+        now = time.perf_counter()
+        dt = now - last
+        last = now
+        if dt > period * 4:  # e.g., >2s if period=0.5s
+            print("!!!!!!! [%s] loop stall detected: slept %.3fs (expected %.3fs)" % (name, dt, period))
+
+
+import asyncio
+import logging
+
+# logging.basicConfig(level=logging.DEBUG)
+# asyncio.get_event_loop().set_debug(True)
+import time
+
+
+def debug_dump(loop):
+    while True:
+        try:
+            print("=== Pending tasks ===")
+            for t in asyncio.all_tasks(loop):
+                if not t.done():
+                    print(t, "awaiting", t.get_coro())
+                    t.print_stack()
+        except Exception:
+            pass
+        time.sleep(5)
+
 
 class LightningSpanProcessor(SpanProcessor):
     def __init__(self):
@@ -242,8 +279,13 @@ class LightningSpanProcessor(SpanProcessor):
     def _loop_runner(self):
         loop = asyncio.new_event_loop()
         self._loop = loop
+        self._loop.set_debug(True)
         asyncio.set_event_loop(loop)
         self._loop_ready.set()
+
+        thread = threading.Thread(target=debug_dump, args=(loop,), daemon=True)
+        thread.start()
+        # asyncio.create_task(heartbeat())
         loop.run_forever()
         loop.close()
 
@@ -261,6 +303,32 @@ class LightningSpanProcessor(SpanProcessor):
         # 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
 
@@ -306,6 +374,10 @@ class LightningSpanProcessor(SpanProcessor):
         Args:
             span: The span that has ended.
         """
+        import traceback
+
+        # print("ON_END")
+        # print(traceback.format_stack())
         # Skip if span is not sampled
         if not span.context or not span.context.trace_flags.sampled:
             return
@@ -313,10 +385,27 @@ class LightningSpanProcessor(SpanProcessor):
         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,
-                )
+                print("!!! before,")
+                print("Ready callbacks:", self._loop._ready)
+                print("Scheduled callbacks:", len(self._loop._scheduled))
+                if self._loop._scheduled:
+                    print("First in the queue:", self._loop._scheduled[0])
+                print("..... Current thread: ", threading.current_thread())
+                print("..... Loop thread: ", self._loop_thread)
+                if self._loop_thread.ident == threading.current_thread().ident:
+                    traceback.print_stack()
+                    print("Span content: ", span.attributes)
+                from opentelemetry.instrumentation.utils import suppress_instrumentation
+
+                with suppress_instrumentation():
+                    self._await_in_loop(
+                        self._store.add_otel_span(self._rollout_id, self._attempt_id, span),
+                        timeout=30.0,
+                    )
+                print("!!! after,")
+                print("All tasks")
+                print("Ready callbacks:", self._loop._ready)
+                print("Scheduled callbacks:", self._loop._scheduled)
             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
+
 logger = logging.getLogger(__name__)
 
 
@@ -112,3 +117,11 @@ class BaseTracer(ParallelWorkerBase):
         """
         with self.trace_context(name=func.__name__):
             return await func(*args, **kwargs)
+
+    def get_langchain_handler(self) -> Optional[BaseCallbackHandler]:
+        """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,37 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+import queue
+import threading
+
+
+def run_sync_ephemeral(coro):
+    """
+    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()
+
+    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 @@ from typing import Any, List, Set, Tuple
 
 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)

examples/spider/sql_agent.py

@@ -1,6 +1,8 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-"""
+"""Sample code that demonstrates an SQL agent using LangGraph and LangChain,
+trainable with Agent-lightning.
+
 Adapted from https://python.langchain.com/docs/tutorials/sql_qa/
 as well as https://langchain-ai.github.io/langgraph/tutorials/sql-agent/
 """
@@ -12,9 +14,9 @@ import re
 import shutil
 import tempfile
 import time
-from typing import Any, Dict, Literal, Optional, cast
+from typing import Any, Dict, List, Literal, Optional, cast
 
-import dotenv
+import pandas as pd
 import termcolor
 from langchain.chat_models import init_chat_model
 from langchain_community.tools.sql_database.tool import QuerySQLDatabaseTool
@@ -25,11 +27,11 @@ from langgraph.graph import END, START, MessagesState, StateGraph
 from langgraph.graph.state import CompiledStateGraph
 from spider_eval.exec_eval import eval_exec_match
 
-import agentlightning
+import agentlightning as agl
 
-agentlightning.configure_logger()
+agl.configure_logger()
 
-logger = agentlightning.configure_logger(name=__name__)
+logger = agl.configure_logger(name=__name__)
 
 
 WRITE_QUERY_PROMPT = ChatPromptTemplate(
@@ -411,7 +413,7 @@ def evaluate_query(query: str, ground_truth: str, database: str, raise_on_error:
             return 0.0
 
 
-class LitSQLAgent(agentlightning.LitAgent[Any]):
+class LitSQLAgent(agl.LitAgent[Dict[str, Any]]):
 
     def __init__(
         self,
@@ -428,20 +430,21 @@ class LitSQLAgent(agentlightning.LitAgent[Any]):
         self.table_info_truncate = table_info_truncate
         self.execution_truncate = execution_truncate
 
-    def _execute_rollout(
-        self, sample: dict[str, Any], *, resources: agentlightning.NamedResources, rollout_id: str, is_training: bool
+    def rollout(
+        self,
+        task: Dict[str, Any],
+        resources: agl.NamedResources,
+        rollout: agl.Rollout,
     ) -> float | None:
-        question = sample["question"]
+        question = task["question"]
         start_time = time.time()
-        llm: agentlightning.LLM = cast(agentlightning.LLM, resources["main_llm"])
+        llm: agl.LLM = cast(agl.LLM, resources["main_llm"])
 
-        if is_training:
-            original_db_path = os.path.join(self.spider_dir, "database", sample["db_id"], sample["db_id"] + ".sqlite")
+        if rollout.mode == "train":
+            original_db_path = os.path.join(self.spider_dir, "database", task["db_id"], task["db_id"] + ".sqlite")
         else:
-            original_db_path = os.path.join(
-                self.spider_dir, "test_database", sample["db_id"], sample["db_id"] + ".sqlite"
-            )
-        ground_truth = sample["query"]
+            original_db_path = os.path.join(self.spider_dir, "test_database", task["db_id"], task["db_id"] + ".sqlite")
+        ground_truth = task["query"]
 
         if not os.path.exists(original_db_path):
             logger.error(f"Database {original_db_path} does not exist. Skipping.")
@@ -455,6 +458,8 @@ class LitSQLAgent(agentlightning.LitAgent[Any]):
             logger.error("Schema file not found: %s", schema_path)
             schema = "No schema available."
 
+        rollout_id = rollout.rollout_id
+
         with tempfile.TemporaryDirectory() as temp_dir:
             db_path = os.path.join(temp_dir, os.path.basename(original_db_path))
             shutil.copyfile(original_db_path, db_path)
@@ -472,7 +477,7 @@ class LitSQLAgent(agentlightning.LitAgent[Any]):
                 endpoint=llm.endpoint,
                 verl_replacement=(
                     {"model": llm.model, **llm.sampling_parameters}
-                    if is_training
+                    if rollout.mode == "train"
                     else {
                         "model": llm.model,
                         "temperature": (
@@ -484,9 +489,11 @@ class LitSQLAgent(agentlightning.LitAgent[Any]):
                 ),
             ).graph()
             try:
+                # Required to make the langchain tracing work
+                handler = self.tracer.get_langchain_handler()
                 result = agent.invoke(  # type: ignore
                     {"question": question},  # type: ignore
-                    {"callbacks": [self.tracer.get_langchain_callback_handler()], "recursion_limit": 100},  # type: ignore
+                    {"callbacks": [handler] if handler else [], "recursion_limit": 100},
                 )
             except Exception as e:
                 logger.exception(f"[Rollout {rollout_id}] Error during agent invocation: {e}")
@@ -512,42 +519,27 @@ class LitSQLAgent(agentlightning.LitAgent[Any]):
 
         return reward
 
-    def training_rollout(self, task: Any, rollout_id: str, resources: agentlightning.NamedResources) -> Any:  # type: ignore
-        return self._execute_rollout(task, resources=resources, rollout_id=rollout_id, is_training=True)
-
-    def validation_rollout(self, task: Any, rollout_id: str, resources: agentlightning.NamedResources) -> Any:  # type: ignore
-        return self._execute_rollout(task, resources=resources, rollout_id=rollout_id, is_training=False)
-
-
-def spider_dev_data():
-    # Read from dev.parquet
-    import pandas as pd
 
+def debug_sql_agent():
     spider_dev_data_path = os.path.join(os.environ.get("VERL_SPIDER_DATA_DIR", "data"), "dev.parquet")
     if not os.path.exists(spider_dev_data_path):
         raise FileNotFoundError(f"Spider dev data file {spider_dev_data_path} does not exist.")
-    df = pd.read_parquet(spider_dev_data_path)  # type: ignore
-    if "OPENAI_API_BASE" not in os.environ:
-        logger.warning(
-            "Environment variable OPENAI_API_BASE is not set. Using default value 'https://api.openai.com/v1'."
-        )
-        openai_api_base = "https://api.openai.com/v1"
-    else:
-        openai_api_base = os.environ["OPENAI_API_BASE"]
+    df = pd.read_parquet(spider_dev_data_path).head(10)  # type: ignore
+    df = cast(List[Dict[str, Any]], df.to_dict(orient="records"))  # type: ignore
+    print("Debug data:", df)
 
-    resource = {
-        "main_llm": agentlightning.LLM(
-            model="gpt-4.1-nano",
-            endpoint=openai_api_base,
-            sampling_parameters={
-                "temperature": 0.0,
-            },
-        )
-    }
-    return agentlightning.DevTaskLoader(df.head(10).to_dict(orient="records"), resource)  # type: ignore
+    trainer = agl.Trainer(
+        n_workers=1,
+        initial_resources={
+            "main_llm": agl.LLM(
+                endpoint=os.environ["OPENAI_API_BASE"],
+                model="gpt-4.1-nano",
+                sampling_parameters={"temperature": 0.7},
+            )
+        },
+    )
+    trainer.dev(LitSQLAgent(), df)
 
 
 if __name__ == "__main__":
-    dotenv.load_dotenv()
-    agent, trainer = agentlightning.lightning_cli(LitSQLAgent, agentlightning.Trainer)
-    trainer.fit_v0(agent, os.environ["VERL_API_BASE"], dev_data=spider_dev_data())
+    debug_sql_agent()

examples/spider/train_sql_agent.py

@@ -0,0 +1,189 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Train an SQL agent on the Spider dataset using Agent-lightning.
+
+This module provides a training script for SQL agents using different model configurations.
+The script supports three different training configurations:
+
+1. 'fast' - A lightweight configuration optimized for CI testing with reduced epochs
+2. 'qwen' - Standard configuration using Qwen-2.5-Coder-1.5B-Instruct model
+3. 'llama' - Configuration using LLaMA-3.2-3B-Instruct model with JSON formatting
+
+Usage:
+    python train_sql_agent.py fast    # Fast training for CI/testing
+    python train_sql_agent.py qwen    # Standard Qwen model training
+    python train_sql_agent.py llama   # LLaMA model training
+
+The script uses reinforcement learning with VERL (Versatile Efficient RL) algorithm
+to train agents on the Spider dataset for text-to-SQL generation tasks.
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+from copy import deepcopy
+from datetime import datetime
+from typing import Any, Dict, Optional
+
+import pandas as pd
+from sql_agent import LitSQLAgent
+
+import agentlightning as agl
+
+RL_TRAINING_CONFIG: Dict[str, Any] = {
+    "algorithm": {
+        "adv_estimator": "grpo",
+        "use_kl_in_reward": False,
+    },
+    "data": {
+        "train_files": "data/train_spider.parquet",
+        "val_files": "data/test_dev_500.parquet",
+        "train_batch_size": 32,
+        "max_prompt_length": 4096,
+        "max_response_length": 2048,
+        "truncation": "error",
+    },
+    "actor_rollout_ref": {
+        "rollout": {
+            "tensor_model_parallel_size": 1,
+            "n": 4,
+            "log_prob_micro_batch_size_per_gpu": 4,
+            "multi_turn": {"format": "hermes"},
+            "name": "vllm",
+            "gpu_memory_utilization": 0.8,
+        },
+        "actor": {
+            "ppo_mini_batch_size": 32,
+            "ppo_micro_batch_size_per_gpu": 4,
+            "optim": {"lr": 1e-6},
+            "use_kl_loss": False,
+            "kl_loss_coef": 0.0,
+            "entropy_coeff": 0,
+            "clip_ratio_low": 0.2,
+            "clip_ratio_high": 0.3,
+            "fsdp_config": {
+                "param_offload": True,
+                "optimizer_offload": True,
+            },
+        },
+        "ref": {
+            "log_prob_micro_batch_size_per_gpu": 8,
+            "fsdp_config": {"param_offload": True},
+        },
+        "model": {
+            "path": "Qwen/Qwen2.5-Coder-1.5B-Instruct",
+            "use_remove_padding": True,
+            "enable_gradient_checkpointing": True,
+        },
+    },
+    "trainer": {
+        "n_gpus_per_node": 1,
+        "val_before_train": True,
+        "critic_warmup": 0,
+        "logger": ["console", "wandb"],
+        "project_name": "AgentLightning",
+        "experiment_name": "spider",
+        "nnodes": 1,
+        "test_freq": 32,
+        "total_epochs": 2,
+    },
+}
+
+
+def config_train_fast() -> Dict[str, Any]:
+    """A fast training run for CI testing purposes."""
+
+    # `EXPERIMENT_NAME="spider_$(date +%Y%m%d%H%M%S)"`
+    timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
+    EXPERIMENT_NAME = f"spider_{timestamp}"
+
+    # `PROJECT_NAME=AgentLightningCI`
+    PROJECT_NAME = "AgentLightningCI"
+
+    # Simulate writing to $GITHUB_OUTPUT if it’s set
+    github_output = os.getenv("GITHUB_OUTPUT")
+    if github_output:
+        with open(github_output, "a") as f:
+            f.write(f"project_name={PROJECT_NAME}\n")
+            f.write(f"run_name={EXPERIMENT_NAME}\n")
+
+    print("Set environment variables:")
+    print(f"PROJECT_NAME={PROJECT_NAME}")
+    print(f"EXPERIMENT_NAME={EXPERIMENT_NAME}")
+
+    config = deepcopy(RL_TRAINING_CONFIG)
+    config["actor_rollout_ref"]["rollout"]["gpu_memory_utilization"] = 0.6
+    config["model"]["path"] = "Qwen/Qwen2.5-Coder-0.5B-Instruct"
+    config["data"]["val_files"] = "data/test_dev.parquet"
+    config["trainer"]["total_epochs"] = 1
+    config["trainer"]["total_training_steps"] = 1
+    config["trainer"]["experiment_name"] = EXPERIMENT_NAME
+    config["trainer"]["project_name"] = PROJECT_NAME
+    config["trainer"]["test_freq"] = 1
+    return config
+
+
+def config_train_qwen() -> Dict[str, Any]:
+    """A configuration for training with Qwen-2.5B."""
+
+    config = deepcopy(RL_TRAINING_CONFIG)
+    return config
+
+
+def config_train_llama() -> Dict[str, Any]:
+    """A configuration for training with LLaMA-3.2-3B-Instruct."""
+
+    config = deepcopy(RL_TRAINING_CONFIG)
+    config["actor_rollout_ref"]["rollout"]["multi_turn"]["format"] = "llama3_json"
+    config["actor_rollout_ref"]["model"]["path"] = "meta-llama/Llama-3.2-3B-Instruct"
+    return config
+
+
+def train(config: Dict[str, Any], active_agent: Optional[str]) -> None:
+    """Train the SQL agent with the given configuration."""
+
+    agent = LitSQLAgent()
+    algorithm = agl.VERL(config)
+    trainer = agl.Trainer(n_workers=10, algorithm=algorithm, adapter={"agent_match": active_agent})
+    print("Adapter agent match acknowledged:", trainer.adapter.agent_match)  # type: ignore
+
+    train_data = pd.read_parquet(config["data"]["train_files"]).to_dict(orient="records")  # type: ignore
+    val_data = pd.read_parquet(config["data"]["val_files"]).to_dict(orient="records")  # type: ignore
+    trainer.fit(agent, train_dataset=train_data, val_dataset=val_data)  # type: ignore
+
+
+def main() -> None:
+    """Main function to parse arguments and run training."""
+    parser = argparse.ArgumentParser(
+        description="Train an SQL agent on the Spider dataset using different model configurations"
+    )
+
+    parser.add_argument(
+        "config",
+        choices=["fast", "qwen", "llama"],
+        help="Training configuration: 'fast' (CI testing), 'qwen' (Qwen-2.5-Coder-1.5B), 'llama' (LLaMA-3.2-3B)",
+    )
+
+    parser.add_argument(
+        "--active-agent", type=str, help="Override the active agent name (default: auto-generated based on config)"
+    )
+
+    args = parser.parse_args()
+
+    # Get the appropriate configuration
+    config_functions = {"fast": config_train_fast, "qwen": config_train_qwen, "llama": config_train_llama}
+
+    config = config_functions[args.config]()
+
+    # Set active agent - use provided value or default based on config choice
+    active_agent = args.active_agent
+
+    print(f"Starting training with '{args.config}' configuration...")
+    print(f"Active agent: {active_agent}")
+
+    train(config, active_agent)
+
+
+if __name__ == "__main__":
+    main()

pyproject.toml

@@ -48,6 +48,7 @@ dev = [
 ]
 experiment = [
   "random-word",
+  "gdown",
 ]
 agent = [
   "autogen-agentchat",