Blame for train_lora.py

"""OpenCode-only async GRPO on LFM2.5-2.6B with a LoRA adapter.

All APIs verified against the released wheels (trl 1.16.0, openenv 0.10.0):
- HarnessRolloutOutcome fields: env_reward, tool_call_count, tool_failure_count, ...
- _AsyncRolloutLoop requires eos_token_ids; max_tool_calling_iterations defaults to
  max_turns=8 in loop-owning mode, so it is set explicitly.
- AsyncGRPOConfig is TrainingArguments-derived (no use_vllm/vllm_mode); the vLLM
  server is given as vllm_server_base_url; dtype defaults to float32, so set it.
- With an injected rollout_worker, the trainer never assigns a versioned lora_name,
  so adapter-only sync cannot address the policy through the capture proxy: the
  vLLM server runs WITHOUT --enable-lora and the trainer streams merged weights
  (its default when the server lacks --enable-lora). Checkpoints still contain
  only the LoRA adapter.

Reward: correctness (1/0) + tool-efficiency bonus of at most 0.1 on correct
answers, the recipe the article's LFM runs used. outcome.tool_call_count comes
from the harness's own ATIF trace.
"""

import os
import threading
from functools import partial

from datasets import Dataset
from huggingface_hub import HfApi
from peft import LoraConfig
from transformers import AutoTokenizer, TrainerCallback
from trl.experimental.async_grpo.async_grpo_config import AsyncGRPOConfig
from trl.experimental.async_grpo.async_grpo_trainer import AsyncGRPOTrainer
from trl.experimental.async_grpo.openenv_harness import (
    CaptureContractError,
    HarnessRolloutWorker,
    _HarnessRolloutLoop,
)

from harbor_env.harness import HarborSessionFactory  # via .pth to OpenEnv/envs

MODEL = "LiquidAI/LFM2.5-2.6B"
VLLM_URL = os.environ.get("VLLM_URL", "http://127.0.0.1:8001")   # serve_policy.sh
ENV_URL = os.environ.get("ENV_URL", "http://127.0.0.1:8000")     # serve_env.sh
TRAIN_SPLIT = os.environ.get("TRAIN_SPLIT", "/app/prepared/datasets/train")
SANDBOX = os.environ.get("SANDBOX", "hf-sandbox")                # Harbor backend
MAX_STEPS = int(os.environ.get("MAX_STEPS", "1000"))
ADAPTER_REPO = os.environ.get("ADAPTER_REPO", "")                # push checkpoints here

TOOL_CALL_BUDGET = 32   # reference for the efficiency bonus; tune to the task pool
MAX_BONUS = 0.1


class UploadCheckpointsCallback(TrainerCallback):
    """Push every checkpoint to the Hub, in the background.

    transformers v5 removed PushToHubCallback: periodic _save_checkpoint writes
    only to local disk (save_model is called with _internal_call=True, which
    skips the push), and `hub_strategy="every_save"` is vestigial. Without this
    callback a crashed job loses the whole run.
    """

    def __init__(self, repo_id):
        self.repo_id = repo_id
        self._threads = []

    def on_save(self, args, state, control, **kwargs):
        if not (self.repo_id and state.is_world_process_zero):
            return control
        ckpt = os.path.join(args.output_dir, f"checkpoint-{state.global_step}")
        t = threading.Thread(target=self._upload, args=(ckpt,), daemon=True)
        t.start()
        self._threads.append(t)
        return control

    def _upload(self, ckpt):
        try:
            HfApi().upload_folder(
                repo_id=self.repo_id,
                folder_path=ckpt,
                path_in_repo=os.path.basename(ckpt),
                # Keep adapter + trainer state, skip multi-GB optimizer shards.
                ignore_patterns=["optimizer.pt", "optimizer/**", "*.pth"],
            )
            print(f"[hub] uploaded {ckpt}", flush=True)
        except Exception as exc:  # never let an upload hiccup kill training
            print(f"[hub] checkpoint upload failed: {exc}", flush=True)

    def on_train_end(self, args, state, control, **kwargs):
        for t in self._threads:
            t.join(timeout=600)
        return control


class ResilientHarnessRolloutLoop(_HarnessRolloutLoop):
    """Score a contract-violating capture as unscorable instead of killing the worker.

    TRL makes CaptureContractError fatal: never train on a corrupted capture.
    Right in principle — but in practice the corruption source is transient infra
    (a dropped websocket between worker and env server), and one bad capture used
    to kill a 25h run at step 551. Dropping the rollout keeps bad data out of
    training just the same. To keep the safety property for SYSTEMATIC corruption
    (e.g. every capture broken, as in the 401 incident), re-raise after a full
    two groups' worth of consecutive contract errors.
    """

    MAX_CONSECUTIVE_CONTRACT_ERRORS = 16

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._contract_errors = 0

    def _run_session(self, prompt, group_id=0):
        import logging

        try:
            result = super()._run_session(prompt, group_id)
        except CaptureContractError as exc:
            self._contract_errors += 1
            if self._contract_errors >= self.MAX_CONSECUTIVE_CONTRACT_ERRORS:
                raise
            logging.warning("capture contract violation; scoring unscorable (%s)", exc)
            return self._EMPTY_ROLLOUT, None
        self._contract_errors = 0
        return result


class ResilientHarnessRolloutWorker(HarnessRolloutWorker):
    _loop_cls = ResilientHarnessRolloutLoop


class HeartbeatCallback(TrainerCallback):
    """Touch a file at every optimizer step; run_job.sh's watchdog kills the
    job if it goes stale. The first full run stalled silently for 13h when the
    rollout path died — steps just stopped, and nothing noticed."""

    PATH = "/tmp/rl-heartbeat"

    def on_step_end(self, args, state, control, **kwargs):
        with open(self.PATH, "w") as f:
            f.write(str(state.global_step))
        return control


def correctness_plus_efficiency(outcome):
    """1/0 correctness, plus <= 0.1 for a correct answer in fewer tool calls.

    A wrong answer stays 0, so the bonus can never be earned by giving up early.
    Returning None keeps an unverified rollout (dead sandbox) unscorable.
    """
    reward = outcome.env_reward
    if reward is None or reward <= 0:
        return reward
    bonus = MAX_BONUS * max(0.0, (TOOL_CALL_BUDGET - outcome.tool_call_count) / TOOL_CALL_BUDGET)
    return reward + min(bonus, MAX_BONUS)


def make_factory():
    return partial(
        HarborSessionFactory,
        server_url=ENV_URL,
        split=TRAIN_SPLIT,
        llm_url=VLLM_URL,
        model=MODEL,
        harness="opencode",       # OpenCode-only run
        sandbox=SANDBOX,
        reward_key="correctness,reward",
    )


def make_rollout_worker(dataset, tokenizer):
    return ResilientHarnessRolloutWorker(
        harness_session_factory=make_factory(),
        harness_adapter=None,     # let the harness own the loop (black-box capture)
        rollout_reward_fn=correctness_plus_efficiency,
        model_name=MODEL,
        dataset=dataset,
        processing_class=tokenizer,
        eos_token_ids=[tokenizer.eos_token_id],   # required loop arg
        reward_funcs=[],
        num_generations=8,        # one GRPO group = 8 rollouts of one task, one harness
        max_inflight_tasks=8,
        vllm_server_url=VLLM_URL,
        max_tokens=4096,          # same cap in training and eval — do not regress this
        max_tool_calling_iterations=32,   # loop default would cap the agent at 8 turns
        temperature=1.0,
        top_p=1.0,                # no truncation: keeps sampled dist. == policy dist.
        top_k=0,                  # 0 = disabled in current vLLM
        chat_template_kwargs={"enable_thinking": False},
    )


def main():
    tokenizer = AutoTokenizer.from_pretrained(MODEL)
    # LFM2.5-2.6B's chat_template.jinja doesn't byte-match any template TRL's
    # add_response_schema recognizes, and the rollout-loop ctor requires a
    # response parser. TRL ships one purpose-built for LFM2.5's Python-style
    # tool calls; set it directly. (Only the ctor needs it: in loop-owning mode
    # the harness, not TRL, parses responses.)
    from trl.chat_template_utils import lfm2_2_5_template

    tokenizer.response_template = lfm2_2_5_template
    # The trainer dataset is discovered from the env server, not loaded from the
    # Hub: HarborSessionFactory.create() only accepts instructions the server
    # itself serves, and prompt_rows() is its blessed export (one row per task,
    # shared by all num_generations of a group).
    discovery = make_factory()()
    dataset = Dataset.from_list(discovery.prompt_rows())
    print(f"discovered {len(dataset)} tasks from {ENV_URL} split={TRAIN_SPLIT}")
    worker = make_rollout_worker(dataset, tokenizer)

    config_kwargs = {
        "output_dir": "./checkpoints",
        "vllm_server_base_url": VLLM_URL,
        "dtype": "bfloat16",     # config default is float32
        "bf16": True,
        "learning_rate": 1e-6,
        "max_steps": MAX_STEPS,
        "logging_steps": 1,
        # Checkpoints ARE your LoRA adapters. Never set save_strategy="no".
        "save_strategy": "steps",
        "save_steps": 100,
        "save_total_limit": 12,
        "report_to": [],
        "seed": 42,
    }
    if ADAPTER_REPO:
        # push_to_hub covers the END-of-training push; periodic checkpoints are
        # pushed by UploadCheckpointsCallback (transformers v5 dropped that path).
        config_kwargs.update(push_to_hub=True, hub_model_id=ADAPTER_REPO)
    config = AsyncGRPOConfig(**config_kwargs)

    # LoRA on attention + MLP projections. Module names verified against the
    # LFM2.5-2.6B safetensors index (self_attn.{q,k,v,out}_proj, feed_forward.w{1,2,3}).
    # Short-conv projections (conv.in_proj/out_proj) are left out for portability.
    peft_config = LoraConfig(
        r=32,
        lora_alpha=64,
        lora_dropout=0.0,
        target_modules=["q_proj", "k_proj", "v_proj", "out_proj", "w1", "w2", "w3"],
        task_type="CAUSAL_LM",
    )

    trainer = AsyncGRPOTrainer(
        model=MODEL,
        args=config,
        train_dataset=dataset,   # required by the ctor even with an injected worker
        processing_class=tokenizer,
        rollout_worker=worker,
        peft_config=peft_config,
        callbacks=[UploadCheckpointsCallback(ADAPTER_REPO), HeartbeatCallback()],
    )
    resume = os.environ.get("RESUME_FROM", "") or None
    if resume:
        print(f"resuming from {resume}")
    trainer.train(resume_from_checkpoint=resume)
    trainer.save_model("./checkpoints/final")   # adapter_model.safetensors + adapter_config.json
    if ADAPTER_REPO:
        trainer.push_to_hub()


if __name__ == "__main__":
    main()