train_calc_agent.py

# Copyright (c) Microsoft. All rights reserved.

"""The training helper script for Calc-X agent with VERL algorithm.

Example usage:

```bash
python train_calc_agent.py --train-file data/train.parquet --val-file data/test.parquet --llm-proxy
```

To use an external store, run a store server first:

```bash
agl store --port 9999
```

Then run the training script with the external store address:

```bash
AGL_MANAGED_STORE=0 python train_calc_agent.py --external-store-address http://localhost:9999
```

Alternatively, you can also run algorithms and runners separately if needed:

```bash
AGL_MANAGED_STORE=0 AGL_CURRENT_ROLE=algorithm python train_calc_agent.py --external-store-address http://localhost:9999
AGL_MANAGED_STORE=0 AGL_CURRENT_ROLE=runner python train_calc_agent.py --external-store-address http://localhost:9999
```
"""

import argparse
import os
import uuid
from datetime import datetime
from typing import Any, Dict, Optional, cast

from calc_agent import MathProblem, calc_agent
from datasets import Dataset as HuggingFaceDataset

import agentlightning as agl
from agentlightning.env_var import LightningEnvVar, resolve_bool_env_var, resolve_str_env_var


def verl_default_config() -> Dict[str, Any]:
    config = {
        "algorithm": {
            "adv_estimator": "grpo",
            "use_kl_in_reward": False,
        },
        "data": {
            "train_batch_size": 32,
            "max_prompt_length": 4096,
            "max_response_length": 2048,
        },
        "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.6,
                "engine_kwargs": {
                    "vllm": {
                        "enable_auto_tool_choice": True,
                        "tool_call_parser": "hermes",
                    }
                },
            },
            "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-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": "calc_x",
            "nnodes": 1,
            "save_freq": 64,
            "test_freq": 32,
            "total_epochs": 2,
        },
    }
    return config


def train(
    *,
    train_file: str,
    val_file: str,
    model: Optional[str],
    llm_proxy: bool,
    ci: bool,
    ci_fast: bool,
    n_runners: int,
    external_store_address: str,
    lora: bool,
    lora_rank: int,
    lora_adapter_path: Optional[str],
    trajectory_level: bool = False,
    weave: bool,
    mongo_uri: Optional[str],
):
    """The training entrypoint function for Calc-X agent with VERL algorithm.

    Args:
        train_file: The path to the training parquet file.
        val_file: The path to the validation parquet file.
        model: The HF model id or path to override the default model.
        llm_proxy: Whether to enable LLM Proxy tracing/adapter.
        ci: Whether to run a minimal CI-style training loop.
        n_runners: The number of runners for the Trainer.
        ci_fast: Whether to cap the training loop at a single step (implies CI toggles).
        external_store_address: Connects to an external store instead of creating a new one in memory.
        lora: Whether to enable LoRA training.
        lora_rank: LoRA rank to use when LoRA is enabled.
        lora_adapter_path: Optional path to a pre-trained LoRA adapter to load.
        trajectory_level: Whether to enable trajectory level in trace aggregator.
        weave: Whether to enable Weave tracing.
        mongo_uri: MongoDB URI to use for the store.
    """
    # Load datasets (respect CLI file paths)
    train_dataset = cast(agl.Dataset[MathProblem], HuggingFaceDataset.from_parquet(train_file).to_list())  # type: ignore
    val_dataset = cast(agl.Dataset[MathProblem], HuggingFaceDataset.from_parquet(val_file).to_list())  # type: ignore

    print("First 5 rows of train dataset:")
    print(train_dataset[:5])  # type: ignore
    print("First 5 rows of val dataset:")
    print(val_dataset[:5])  # type: ignore

    config = verl_default_config()

    if model:
        config["actor_rollout_ref"]["model"]["path"] = model

    # Enable LoRA configuration if requested
    if lora:
        config["actor_rollout_ref"]["model"]["lora_rank"] = lora_rank
        print(f"LoRA enabled: lora_rank={lora_rank}")
        if lora_adapter_path:
            config["actor_rollout_ref"]["model"]["lora_adapter_path"] = lora_adapter_path
            print(f"Loading LoRA adapter from: {lora_adapter_path}")
        print("LoRA configuration will trigger verl to set ref_in_actor=True (LoRA mode)")

    if trajectory_level:
        config["agentlightning"] = {
            "trace_aggregator": {
                "level": "trajectory",
                "trajectory_max_prompt_length": 2048,
                "trajectory_max_response_length": 8192,
            }
        }
        print("Trajectory level enabled in trace aggregator.")

    # CI toggle keeps everything else the same but you can tweak the lightweight bits here if desired
    if ci or ci_fast:
        # Config the experiment name and project name so that they are available to CI
        timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
        random_suffix = uuid.uuid4().hex[:8]
        EXPERIMENT_NAME = f"calc_x_{timestamp}_{random_suffix}"

        PROJECT_NAME = "AgentLightningCI"

        # Skip this step if AGL_CURRENT_ROLE is runner
        agl_current_role = resolve_str_env_var(LightningEnvVar.AGL_CURRENT_ROLE)

        if agl_current_role != "runner":
            # 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}")

        # Keep it tiny/light without adding new knobs
        config["actor_rollout_ref"]["rollout"]["gpu_memory_utilization"] = 0.8
        config["trainer"]["total_epochs"] = 1
        config["trainer"]["total_training_steps"] = 20
        config["trainer"]["test_freq"] = 20
        config["trainer"]["experiment_name"] = EXPERIMENT_NAME
        config["trainer"]["project_name"] = PROJECT_NAME
        config["trainer"].pop("save_freq", None)

        if ci_fast:
            # Extra fast CI toggle for testing purposes.
            config["actor_rollout_ref"]["rollout"]["gpu_memory_utilization"] = 0.6
            config["trainer"]["total_training_steps"] = 1
            config["trainer"]["test_freq"] = 1

    algorithm = agl.VERL(config)

    if external_store_address:
        store: Optional[agl.LightningStore] = agl.LightningStoreClient(external_store_address)
    elif mongo_uri:
        from agentlightning.store.mongo import MongoLightningStore

        store = MongoLightningStore(mongo_uri=mongo_uri)
    else:
        store = None

    if llm_proxy:
        tracer = agl.OtelTracer()  # dummy tracer for LLM Proxy
        adapter = agl.LlmProxyTraceToTriplet()
        trainer = agl.Trainer(algorithm=algorithm, n_runners=n_runners, store=store, tracer=tracer, adapter=adapter)
    elif weave:
        # NOTE: Don't import WeaveTracer at the module level or in __init__.py files.
        # Always import it lazily/conditionally (behind a feature flag) to avoid interfering
        # with other libraries like LiteLLM/OpenTelemetry when weave is not explicitly enabled.
        from agentlightning.tracer.weave import WeaveTracer

        tracer = WeaveTracer()
        trainer = agl.Trainer(algorithm=algorithm, n_runners=n_runners, store=store, tracer=tracer)
    else:
        trainer = agl.Trainer(algorithm=algorithm, n_runners=n_runners, store=store)

    trainer.fit(calc_agent, train_dataset, val_dataset=val_dataset)


def main():
    parser = argparse.ArgumentParser(description="Train a math calc agent with Agent-lightning + VERL.")
    parser.add_argument("--train-file", type=str, default="data/train.parquet", help="Path to train parquet file")
    parser.add_argument("--val-file", type=str, default="data/test.parquet", help="Path to val parquet file")
    parser.add_argument("--model", type=str, default=None, help="HF model id or path (optional)")
    parser.add_argument("--llm-proxy", action="store_true", help="Enable LLM Proxy tracing/adapter")
    parser.add_argument("--weave", action="store_true", help="Enable Weave tracing")
    parser.add_argument("--ci", action="store_true", help="Run a minimal CI-style training loop")
    parser.add_argument(
        "--ci-fast", action="store_true", help="Limit the training loop to a single step (implies --ci)"
    )
    parser.add_argument("--n-runners", type=int, default=10, help="Number of runners for Trainer")
    parser.add_argument(
        "--external-store-address",
        type=str,
        default="",
        help="Connect to an external store instead of creating a new one in memory",
    )
    parser.add_argument("--debug", action="store_true", help="Enable debug logging")
    parser.add_argument(
        "--lora",
        action="store_true",
        help="Enable LoRA training. When enabled, the reference policy is computed by the actor rollout worker.",
    )
    parser.add_argument(
        "--lora-rank",
        type=int,
        default=32,
        help="LoRA rank to use when --lora is enabled (default: 32)",
    )
    parser.add_argument(
        "--lora-adapter-path",
        type=str,
        default=None,
        help="Optional path to a pre-trained LoRA adapter to load when --lora is enabled",
    )
    parser.add_argument(
        "--trajectory-level",
        action="store_true",
        help="Enable trajectory level in trace aggregator.",
    )
    parser.add_argument(
        "--mongo-uri",
        type=str,
        default=None,
        help="MongoDB URI to use for the store.",
    )

    args = parser.parse_args()

    if args.external_store_address:
        print(f"Connecting to external store at: {args.external_store_address}")
        if resolve_bool_env_var(LightningEnvVar.AGL_MANAGED_STORE, fallback=True):
            raise ValueError(
                "When using an external store, please set the environment variable AGL_MANAGED_STORE=0. "
                "Otherwise the trainer will still try to manage the store lifecycle for you!"
            )

    if args.ci_fast:
        args.ci = True

    agl.setup_logging("DEBUG" if args.debug else "INFO")

    train(
        train_file=args.train_file,
        val_file=args.val_file,
        model=args.model,
        llm_proxy=args.llm_proxy,
        ci=args.ci,
        ci_fast=args.ci_fast,
        n_runners=args.n_runners,
        external_store_address=args.external_store_address,
        lora=args.lora,
        lora_rank=args.lora_rank,
        lora_adapter_path=args.lora_adapter_path,
        trajectory_level=args.trajectory_level,
        weave=args.weave,
        mongo_uri=args.mongo_uri,
    )


if __name__ == "__main__":
    main()