diff --git a/.github/workflows/examples.yml b/.github/workflows/examples.yml index 6e8d39738..8d12235c2 100644 --- a/.github/workflows/examples.yml +++ b/.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 diff --git a/agentlightning/algorithm/fast.py b/agentlightning/algorithm/fast.py index f9fb699cd..17f23a30a 100644 --- a/agentlightning/algorithm/fast.py +++ b/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) diff --git a/agentlightning/config.py b/agentlightning/config.py index 7e625bf95..98081d9f5 100644 --- a/agentlightning/config.py +++ b/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. diff --git a/agentlightning/litagent/litagent.py b/agentlightning/litagent/litagent.py index 8a4eb156b..27bd0bdd0 100644 --- a/agentlightning/litagent/litagent.py +++ b/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: diff --git a/agentlightning/runner/agent.py b/agentlightning/runner/agent.py index bd9644d65..892c3fc23 100644 --- a/agentlightning/runner/agent.py +++ b/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. diff --git a/agentlightning/store/client_server.py b/agentlightning/store/client_server.py index 9a1ea41f1..fc86bda77 100644 --- a/agentlightning/store/client_server.py +++ b/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): diff --git a/agentlightning/tracer/agentops.py b/agentlightning/tracer/agentops.py index ddae131fb..e597c2436 100644 --- a/agentlightning/tracer/agentops.py +++ b/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}") diff --git a/agentlightning/tracer/base.py b/agentlightning/tracer/base.py index b17e7b564..7f91c55ab 100644 --- a/agentlightning/tracer/base.py +++ b/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 diff --git a/examples/spider/spider_eval/__init__.py b/examples/spider/spider_eval/__init__.py new file mode 100644 index 000000000..2a50eae89 --- /dev/null +++ b/examples/spider/spider_eval/__init__.py @@ -0,0 +1 @@ +# Copyright (c) Microsoft. All rights reserved. diff --git a/examples/spider/spider_eval/async_utils.py b/examples/spider/spider_eval/async_utils.py new file mode 100644 index 000000000..fffd54c2e --- /dev/null +++ b/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 diff --git a/examples/spider/spider_eval/exec_eval.py b/examples/spider/spider_eval/exec_eval.py index f1b1ec483..69e2081cf 100644 --- a/examples/spider/spider_eval/exec_eval.py +++ b/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) diff --git a/examples/spider/sql_agent.py b/examples/spider/sql_agent.py index 63f3b0b48..216e70874 100644 --- a/examples/spider/sql_agent.py +++ b/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 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.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 @@ def __init__( 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 @@ def _execute_rollout( 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) @@ -469,10 +474,10 @@ def _execute_rollout( execution_truncate=self.execution_truncate, debug=False, db_schema=schema, - endpoint=llm.endpoint, + endpoint=llm.get_base_url(rollout.rollout_id, rollout.attempt.attempt_id), # type: ignore verl_replacement=( {"model": llm.model, **llm.sampling_parameters} - if is_training + if rollout.mode == "train" else { "model": llm.model, "temperature": ( @@ -484,9 +489,11 @@ def _execute_rollout( ), ).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 @@ def _execute_rollout( 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"] - - 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 + 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) + + 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() diff --git a/examples/spider/train_sql_agent.py b/examples/spider/train_sql_agent.py new file mode 100644 index 000000000..cfceac277 --- /dev/null +++ b/examples/spider/train_sql_agent.py @@ -0,0 +1,192 @@ +# 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 framework +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["actor_rollout_ref"]["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-1B-Instruct. + + You will need a `HF_TOKEN` set to run with this config. + """ + + 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-1B-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() diff --git a/pyproject.toml b/pyproject.toml index 4b2916b43..bfac5e1ad 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -48,6 +48,7 @@ dev = [ ] experiment = [ "random-word", + "gdown", ] agent = [ "autogen-agentchat", diff --git a/scripts/setup_latest_gpu.sh b/scripts/setup_latest_gpu.sh index b25827d55..ab7ed08d0 100755 --- a/scripts/setup_latest_gpu.sh +++ b/scripts/setup_latest_gpu.sh @@ -9,7 +9,8 @@ pip install --no-cache-dir flash-attn --no-build-isolation # This must match pytorch version. pip install --no-cache-dir vllm==0.10.2 # Latest VERL release version. -pip install --no-cache-dir verl +# FIXME: Make VERL 0.5.0 work +pip install --no-cache-dir "verl<0.6.0" pip install --no-cache-dir -e .[dev,agent,trl,apo] # Upgrade agentops to the latest version diff --git a/tests/algorithm/test_baseline.py b/tests/algorithm/test_baseline.py index 9b80a7e26..1496c3574 100644 --- a/tests/algorithm/test_baseline.py +++ b/tests/algorithm/test_baseline.py @@ -112,7 +112,7 @@ async def _mock_runner( async def test_mock_algorithm_collects_rollout_logs(caplog: pytest.LogCaptureFixture) -> None: store = InMemoryLightningStore() await store.update_resources("default", _make_resources()) - algorithm = Baseline(polling_interval=0.01) + algorithm = Baseline(polling_interval=0.01, span_verbosity="key_values") algorithm.set_store(store) adapter = _AdapterStub() algorithm.set_adapter(adapter) @@ -158,3 +158,92 @@ async def test_mock_algorithm_collects_rollout_logs(caplog: pytest.LogCaptureFix and entry.attempt_id in msg for msg in log_messages ) + + +@pytest.mark.asyncio +async def test_baseline_does_not_skip_samples_when_queue_full() -> None: + """Test that Baseline waits and retries when queue is full instead of skipping samples. + + This is a regression test for a bug where samples would be skipped when the queue + exceeded max_queue_length. The fix wraps the queue check in a while loop to ensure + all samples are eventually processed. + """ + store = InMemoryLightningStore() + await store.update_resources("default", _make_resources()) + + # Use a small max_queue_length and fast polling to test queue full behavior + algorithm = Baseline(polling_interval=0.01, max_queue_length=1) + algorithm.set_store(store) + + # Create a dataset with 5 samples + train_dataset = [f"sample-{i}" for i in range(5)] + expected_rollouts = len(train_dataset) + + # Track which samples were enqueued + enqueued_samples: List[Any] = [] + artifacts: List[_RolloutArtifacts] = [] + + async def _slow_runner() -> None: + """A slow runner that creates backpressure by processing rollouts with delays.""" + processed = 0 + while processed < expected_rollouts: + attempted = await store.dequeue_rollout() + if attempted is None: + await asyncio.sleep(0.01) + continue + + attempt = attempted.attempt + rollout_id = attempted.rollout_id + rollout = await store.get_rollout_by_id(rollout_id) + + # Track the sample that was enqueued + if rollout: + enqueued_samples.append(rollout.input) + + await store.update_attempt( + rollout_id, + attempt.attempt_id, + status="running", + worker_id="slow-runner", + ) + + # Add a delay to create backpressure and cause queue to fill up + await asyncio.sleep(0.05) + + span = _build_span(rollout_id, attempt.attempt_id, sequence_id=1, index=processed + 1) + await store.add_span(span) + await store.update_attempt(rollout_id, attempt.attempt_id, status="succeeded") + await store.update_rollout(rollout_id, status="succeeded") + + artifacts.append( + _RolloutArtifacts( + rollout_id=rollout_id, + attempt_id=attempt.attempt_id, + attempt_sequence=attempt.sequence_id, + span=span, + ) + ) + processed += 1 + + runner_task = asyncio.create_task(_slow_runner()) + try: + await algorithm.run(train_dataset=train_dataset) + await asyncio.wait_for(runner_task, timeout=5) + finally: + if not runner_task.done(): + runner_task.cancel() + with suppress(asyncio.CancelledError): + await runner_task + + # Verify that ALL samples were enqueued and processed (no samples skipped) + assert ( + len(enqueued_samples) == expected_rollouts + ), f"Expected {expected_rollouts} samples to be enqueued, but got {len(enqueued_samples)}" + assert ( + len(artifacts) == expected_rollouts + ), f"Expected {expected_rollouts} rollouts to be processed, but got {len(artifacts)}" + + # Verify that the enqueued samples match the dataset (in order) + assert ( + enqueued_samples == train_dataset + ), f"Enqueued samples {enqueued_samples} do not match dataset {train_dataset}"