training/train.py

"""
GRPO training script for the API Contract Validator environment.

Designed to be re-runnable by judges from a Colab notebook OR via HF
Jobs:

    # HF Jobs (T4 small, ~$0.50/hr — uses your $30 credit)
    hf jobs uv run \
        --with trl --with unsloth --with openenv-core --with wandb \
        --flavor t4-small \
        -s HF_TOKEN -s WANDB_API_KEY \
        -- python training/train.py

The script:

  1. Connects to a deployed HF Space (or local docker) running the env
  2. Loads a small base model with Unsloth 4-bit quantisation
  3. Applies LoRA adapters
  4. Rolls out episodes through the env, collecting (prompt, completion,
     reward) tuples
  5. Trains the LoRA adapters with GRPO from TRL
  6. Logs reward curves to WandB and writes results/reward_curve.png
  7. Pushes the trained adapter to the HuggingFace Hub

The reward function uses the env's grader directly — no synthetic
shaping. This is the key difference from a static-dataset SFT run:
the model learns from the env's verifiable signal, which is exactly
what the hackathon's "Improvement in Rewards" criterion rewards.
"""

from __future__ import annotations

import json
import os
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List

import matplotlib.pyplot as plt

# Ensure api_contract_validator is importable
ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))

# Load .env from api_contract_validator/ before reading os.getenv values
try:
    from dotenv import load_dotenv
    _ENV_FILE = ROOT / ".env"
    if _ENV_FILE.exists():
        load_dotenv(_ENV_FILE)
except ImportError:
    pass


# ── Configuration ────────────────────────────────────────────────────────


@dataclass
class TrainConfig:
    """Training configuration. All fields read from env vars at instantiation.

    Recommended HF Jobs configurations:

      Smoke test ($0.30, 5 min):
        BASE_MODEL=unsloth/Qwen2.5-1.5B-Instruct-bnb-4bit, MAX_STEPS=10,
        flavor t4-small

      Main run on L4 ($2.40, ~2 hr):
        BASE_MODEL=unsloth/Qwen2.5-7B-Instruct-bnb-4bit, MAX_STEPS=300,
        flavor l4x1

      Insurance run on T4 ($0.40, ~45 min):
        BASE_MODEL=unsloth/Qwen2.5-1.5B-Instruct-bnb-4bit, MAX_STEPS=200,
        flavor t4-small
    """

    base_model: str = os.getenv(
        "BASE_MODEL", "unsloth/Qwen2.5-7B-Instruct-bnb-4bit"
    )
    env_url: str = os.getenv("ENV_URL", "http://localhost:7860")
    push_to_hub_id: str | None = os.getenv("PUSH_TO_HUB", None)
    output_dir: str = os.getenv("OUTPUT_DIR", "checkpoints/grpo")
    results_dir: str = os.getenv(
        "RESULTS_DIR", str(ROOT / "results")
    )
    seed: int = int(os.getenv("SEED", "42"))

    # LoRA
    lora_r: int = int(os.getenv("LORA_R", "16"))
    lora_alpha: int = int(os.getenv("LORA_ALPHA", "32"))

    # GRPO
    max_seq_length: int = int(os.getenv("MAX_SEQ_LEN", "2048"))
    num_generations: int = int(os.getenv("NUM_GENERATIONS", "4"))
    max_steps: int = int(os.getenv("MAX_STEPS", "300"))
    learning_rate: float = float(os.getenv("LR", "5e-6"))
    per_device_batch_size: int = int(os.getenv("BATCH_SIZE", "1"))
    grad_accum: int = int(os.getenv("GRAD_ACCUM", "4"))

    # Tasks to train on (subset speeds up onsite training)
    train_tasks: List[str] | None = None

    # WandB
    wandb_project: str = os.getenv("WANDB_PROJECT", "openenv-contract-guardian")
    wandb_run: str = os.getenv("WANDB_RUN", "grpo-onsite")


# ── Reward function: rolls out one step against the live env ─────────────


def _list_value(values: Any, index: int, default: Any) -> Any:
    """Return ``values[index]`` for TRL batch kwargs, with a safe fallback."""
    if isinstance(values, list) and index < len(values):
        return values[index]
    return default


def make_reward_fn(env_url: str, task_pool: List[str]):
    """Return a TRL-compatible reward_fn that grades each completion via env.

    A fresh ``ValidatorEnv`` (WebSocket) is created per ``reward_fn``
    invocation and closed at the end. HF Spaces drops idle WebSockets
    after ~30 s, but GRPO's model-generation and backprop pauses are
    longer than that — sharing one WebSocket across batches caused
    "received 1011 keepalive ping timeout" on every batch after the
    first. A per-call client adds ~50 ms of TCP setup but eliminates
    the keepalive failures entirely.

    Within a single ``reward_fn`` call, all completions are graded
    through one client (calls are rapid so keepalive is not at risk).
    """
    from inference import _build_action, parse_llm_response  # noqa: WPS433
    import asyncio

    from client import ValidatorEnv  # noqa: WPS433

    def reward_fn(prompts, completions, **kwargs):  # noqa: ARG001
        rewards: List[float] = []
        try:
            loop = asyncio.get_event_loop()
            if loop.is_closed():
                raise RuntimeError("loop closed")
        except RuntimeError:
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)

        task_names = kwargs.get("task") or []
        seeds = kwargs.get("seed") or []

        env_client = ValidatorEnv(base_url=env_url)
        try:
            for idx, completion in enumerate(completions):
                text = (
                    completion
                    if isinstance(completion, str)
                    else completion[0]["content"]
                )
                task_name = _list_value(task_names, idx, task_pool[0])
                seed = _list_value(seeds, idx, 0)
                try:
                    loop.run_until_complete(
                        env_client.reset(task_name=task_name, seed=int(seed))
                    )
                    action_data = parse_llm_response(text)
                    action = _build_action(action_data)
                    step_result = loop.run_until_complete(
                        env_client.step(action)
                    )
                    rewards.append(float(step_result.reward or 0.0))
                except Exception as exc:  # noqa: BLE001
                    print(f"[WARN] reward_fn error: {exc}")
                    rewards.append(-0.5)
        finally:
            try:
                loop.run_until_complete(env_client.close())
            except Exception:  # noqa: BLE001
                pass

        return rewards

    return reward_fn


# ── Dataset: one prompt per env reset ────────────────────────────────────


def build_train_dataset(env_client, tasks: List[str], episodes_per_task: int = 50):
    """Roll out reset() to capture initial observations as training prompts.

    Each row is one episode start. During training, GRPO samples
    ``num_generations`` completions per prompt and uses the env to
    grade them.
    """
    import asyncio
    from datasets import Dataset  # type: ignore

    from inference import build_user_prompt, _system_prompt_for_phase  # noqa

    rows: List[Dict[str, Any]] = []
    loop = asyncio.get_event_loop()

    for task in tasks:
        for ep in range(episodes_per_task):
            seed = ep
            result = loop.run_until_complete(
                env_client.reset(task_name=task, seed=seed)
            )
            obs = result.observation.model_dump()
            phase = obs.get("phase", "detection")
            system = _system_prompt_for_phase(phase, task)
            user = build_user_prompt(obs, step=1, history=[])
            rows.append({
                "prompt": [
                    {"role": "system", "content": system},
                    {"role": "user", "content": user},
                ],
                "task": task,
                "seed": seed,
            })

    return Dataset.from_list(rows)


# ── Main entry point ─────────────────────────────────────────────────────


def main() -> None:
    cfg = TrainConfig()

    # ---- Imports happen inside main so the script can be inspected
    # ----  without the heavy deps installed.
    import asyncio
    from openenv.core.client_types import StepResult  # noqa: F401

    try:
        from unsloth import FastLanguageModel  # type: ignore
    except ImportError as exc:
        sys.exit(
            "unsloth not installed. Install with: pip install unsloth trl wandb. "
            f"({exc})"
        )

    from trl import GRPOConfig, GRPOTrainer  # type: ignore
    import wandb  # type: ignore

    from client import ValidatorEnv  # noqa: WPS433

    if os.getenv("WANDB_API_KEY"):
        wandb.init(
            project=cfg.wandb_project,
            name=cfg.wandb_run,
            config=cfg.__dict__,
        )

    # 1. Connect to the env
    env = ValidatorEnv(base_url=cfg.env_url)
    print(f"[INFO] connected to env at {cfg.env_url}")

    # 2. Choose tasks
    train_tasks = cfg.train_tasks or [
        "find_type_mismatches",
        "validate_nested_objects",
        "detect_breaking_changes",
        "validate_response_schema",
        "trace_downstream_blast_radius",
        "propose_backward_compat_fix",
    ]

    # 3. Build dataset
    print(f"[INFO] building dataset for tasks={train_tasks}")
    train_dataset = build_train_dataset(env, train_tasks)

    # 4. Mixed precision setup.
    #
    #    We force fp16 on every GPU rather than auto-selecting bf16 on
    #    Ampere+. Reason: unsloth's fast_lora kernel with bf16 autocast
    #    crashes inside its gradient-checkpointed LoRA forward pass with
    #    "self and mat2 must have the same dtype, but got Half and Float".
    #    fp16 avoids the autocast path that triggers the bug entirely
    #    and works on T4 (smoke test confirmed) and L4 alike.
    #
    #    Tradeoff: slightly less numerical range than bf16. Acceptable
    #    for LoRA training; bf16's main advantage is full-precision FT.
    import torch  # type: ignore
    use_bf16 = False
    torch_dtype = torch.float16
    print(f"[INFO] mixed precision: fp16 (bf16 disabled due to unsloth LoRA issue)")

    # 5. Load model + LoRA. Pass `dtype` explicitly so the model weights
    #    match the dtype the GRPO trainer will use. Without this Unsloth
    #    loads in fp16 by default; with bf16=True in GRPOConfig the LoRA
    #    forward pass crashes with "self and mat2 must have the same
    #    dtype, but got Half and Float".
    print(f"[INFO] loading model: {cfg.base_model}")
    model, tokenizer = FastLanguageModel.from_pretrained(
        model_name=cfg.base_model,
        max_seq_length=cfg.max_seq_length,
        load_in_4bit=True,
        dtype=torch_dtype,
    )
    model = FastLanguageModel.get_peft_model(
        model,
        r=cfg.lora_r,
        lora_alpha=cfg.lora_alpha,
        target_modules=[
            "q_proj", "k_proj", "v_proj", "o_proj",
            "gate_proj", "up_proj", "down_proj",
        ],
        random_state=cfg.seed,
    )

    grpo_cfg = GRPOConfig(
        output_dir=cfg.output_dir,
        learning_rate=cfg.learning_rate,
        per_device_train_batch_size=cfg.per_device_batch_size,
        gradient_accumulation_steps=cfg.grad_accum,
        num_generations=cfg.num_generations,
        max_steps=cfg.max_steps,
        max_prompt_length=cfg.max_seq_length // 2,
        max_completion_length=cfg.max_seq_length // 2,
        logging_steps=1,
        save_steps=50,
        report_to="wandb" if os.getenv("WANDB_API_KEY") else "none",
        bf16=use_bf16,
        fp16=not use_bf16,
    )

    reward_fn = make_reward_fn(cfg.env_url, train_tasks)

    trainer = GRPOTrainer(
        model=model,
        processing_class=tokenizer,
        reward_funcs=[reward_fn],
        args=grpo_cfg,
        train_dataset=train_dataset,
    )

    # 6. Train
    print("[INFO] starting GRPO training")
    trainer.train()

    # 7. Save reward curve
    results_dir = Path(cfg.results_dir)
    results_dir.mkdir(parents=True, exist_ok=True)

    history = [
        h for h in trainer.state.log_history if "reward" in h
    ]
    if history:
        steps = [h["step"] for h in history]
        rewards = [h["reward"] for h in history]
        plt.figure(figsize=(8, 5))
        plt.plot(steps, rewards, label="train reward", linewidth=2)
        plt.xlabel("Training step")
        plt.ylabel("Mean episode reward")
        plt.title("GRPO Training — Enterprise Contract Guardian")
        plt.grid(alpha=0.3)
        plt.legend()
        plt.tight_layout()
        out = results_dir / "reward_curve.png"
        plt.savefig(out, dpi=150)
        print(f"[INFO] wrote {out}")

    # 8. Save trainer state to JSON for plot.py to consume later
    state_path = results_dir / "training_state.json"
    state_path.write_text(json.dumps(trainer.state.log_history, indent=2))
    print(f"[INFO] wrote {state_path}")

    # 9. Push checkpoint AND training artefacts (reward_curve, state JSON)
    #    HF Jobs containers are ephemeral — anything written under
    #    results/ is lost when the job exits. To make the reward curve
    #    available after the job finishes, we upload it to the same HF
    #    Hub model repo where the LoRA adapter goes, under a
    #    "training_artifacts/" path.
    if cfg.push_to_hub_id:
        print(f"[INFO] pushing adapter to {cfg.push_to_hub_id}")
        model.push_to_hub(cfg.push_to_hub_id, token=os.getenv("HF_TOKEN"))

        # Upload training artefacts to the same repo
        try:
            from huggingface_hub import HfApi
            api = HfApi(token=os.getenv("HF_TOKEN"))
            for fname in ("reward_curve.png", "training_state.json"):
                local = results_dir / fname
                if local.exists():
                    print(f"[INFO] uploading {fname} -> {cfg.push_to_hub_id}/training_artifacts/{fname}")
                    api.upload_file(
                        path_or_fileobj=str(local),
                        path_in_repo=f"training_artifacts/{fname}",
                        repo_id=cfg.push_to_hub_id,
                        repo_type="model",
                        commit_message=f"Upload {fname} from GRPO run",
                    )
                else:
                    print(f"[WARN] {local} not found — skipping upload")
        except Exception as exc:  # noqa: BLE001
            print(f"[WARN] artefact upload failed: {exc}")

    # Clean up the dataset-build env client (reward_fn uses its own per-call clients)
    try:
        asyncio.get_event_loop().run_until_complete(env.close())
    except Exception:  # noqa: BLE001
        pass

    print("[INFO] done.")

    # Force a clean exit. The `websockets` library emits a non-zero exit
    # status from its __del__ finalizer when the event loop has been
    # closed, which makes HF Jobs mark the job ERROR even though every
    # artefact (adapter, reward_curve, training_state) was uploaded. Use
    # os._exit(0) to bypass interpreter shutdown finalizers entirely.
    import os as _os  # local alias to avoid shadowing module-level os
    _os._exit(0)


if __name__ == "__main__":
    main()