scripts/bold_pipeline.py

"""Bold full-stack training pipeline: SFT -> DPO -> GRPO -> Adv DPO -> Eval.

ONE script. Two modes via --smoke flag:
  --smoke   : 5 steps per stage, ~10 min total. Validates code path. Free Colab.
  (default) : full step counts, ~3-4h on H200. Production.

Usage in Colab T4 (smoke):
  !python scripts/bold_pipeline.py --smoke

Usage on HF Jobs H200 (production):
  hf jobs uv run --flavor h200 \\
      https://huggingface.co/spaces/sai1906/opsguard/resolve/main/scripts/bold_pipeline.py
"""
from __future__ import annotations

import argparse
import gc
import inspect
import json
import os
import random
import sys
from pathlib import Path

def _wait_for_cuda():
    import time as _t
    print("[cuda] sleeping 30s for driver init...", flush=True)
    _t.sleep(30)
    try:
        import torch
        n = torch.cuda.device_count() if torch.cuda.is_available() else 0
        if n > 0:
            print(f"[cuda] ready: {n} device(s), {torch.cuda.get_device_name(0)}", flush=True)
        else:
            print("[cuda] device_count=0 after sleep — letting model load trigger lazy init", flush=True)
    except Exception as e:
        print(f"[cuda] probe error (non-fatal, model load will retry): {e}", flush=True)


def _setup_workdir():
    import subprocess
    work = Path("/tmp/opsguard")
    if not (work / "data" / "sft_traces.jsonl").exists():
        if work.exists():
            subprocess.run(["rm", "-rf", str(work)], check=True)
        token = os.environ.get("HF_TOKEN", "")
        url = (f"https://user:{token}@huggingface.co/spaces/sai1906/opsguard"
               if token else "https://huggingface.co/spaces/sai1906/opsguard")
        subprocess.run(["git", "clone", "--depth", "1", url, str(work)], check=True)
    sys.path.insert(0, str(work))
    os.chdir(work)
    needed = work / "data" / "sft_traces.jsonl"
    if not needed.exists():
        raise FileNotFoundError(
            f"{needed} missing after clone. cwd={os.getcwd()} "
            f"contents={list(work.iterdir()) if work.exists() else 'workdir-missing'}"
        )
    print(f"[workdir] cwd={os.getcwd()}  data files OK", flush=True)
    return work


_wait_for_cuda()
WORK = _setup_workdir()


def _safe_kwargs(cls, kw):
    sig = inspect.signature(cls).parameters
    return {k: v for k, v in kw.items() if k in sig}


def _free():
    import torch
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()


def _push_lora(folder, repo_id):
    from huggingface_hub import HfApi
    api = HfApi(token=os.environ["HF_TOKEN"])
    try:
        api.create_repo(repo_id, repo_type="model", exist_ok=True)
    except Exception as e:
        print(f"  create_repo({repo_id}): {e}", flush=True)
    api.upload_folder(folder_path=folder, repo_id=repo_id, repo_type="model")
    print(f"  pushed -> https://huggingface.co/{repo_id}", flush=True)


def stage_sft(args, model_name, hub_repo):
    print("\n" + "=" * 70 + "\n=== STAGE 1: SFT WARMSTART ===\n" + "=" * 70, flush=True)
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
    from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
    from datasets import Dataset
    from trl import SFTConfig, SFTTrainer

    print(f"loading {model_name} (4bit={args.use_4bit})...", flush=True)
    model_kwargs = {"torch_dtype": torch.bfloat16, "device_map": "auto",
                     "attn_implementation": "sdpa"}
    if args.use_4bit:
        bnb = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16,
                                  bnb_4bit_use_double_quant=True, bnb_4bit_quant_type="nf4")
        model_kwargs["quantization_config"] = bnb
    tok = AutoTokenizer.from_pretrained(model_name)
    if tok.pad_token_id is None:
        tok.pad_token_id = tok.eos_token_id
    model = AutoModelForCausalLM.from_pretrained(model_name, **model_kwargs)
    model.config.use_cache = False
    if args.use_4bit:
        model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=True)
    else:
        model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
        model.enable_input_require_grads()
    lc = LoraConfig(r=args.lora_r, lora_alpha=args.lora_alpha, lora_dropout=0.05, bias="none",
                    target_modules=["q_proj","k_proj","v_proj","o_proj",
                                    "gate_proj","up_proj","down_proj"],
                    task_type="CAUSAL_LM")
    model = get_peft_model(model, lc)
    for n, p in model.named_parameters():
        if p.requires_grad and ("lora_" in n or "lora_A" in n or "lora_B" in n):
            p.data = p.data.to(torch.float32)
    model.print_trainable_parameters()

    rows = [json.loads(l) for l in open("data/sft_traces.jsonl")]
    print(f"loaded {len(rows)} SFT traces", flush=True)
    texts = [r["prompt"] + "\n\nACTION:\n" + r["completion"] + tok.eos_token for r in rows]
    if args.smoke:
        texts = texts[:64]
    ds = Dataset.from_list([{"text": t} for t in texts])

    n_epochs = 1 if args.smoke else args.sft_epochs
    raw = dict(
        output_dir="/tmp/opsguard-sft",
        per_device_train_batch_size=args.batch_size,
        gradient_accumulation_steps=args.grad_accum,
        num_train_epochs=n_epochs,
        learning_rate=2e-5,
        max_grad_norm=1.0,
        weight_decay=0.01,
        warmup_ratio=0.1,
        adam_epsilon=1e-7,
        optim="adamw_torch",
        gradient_checkpointing=True,
        gradient_checkpointing_kwargs={"use_reentrant": False},
        logging_steps=2,
        save_strategy="epoch",
        save_total_limit=1,
        bf16=True,
        max_seq_length=2048,
        dataset_text_field="text",
        report_to="none",
        push_to_hub=False,
        max_steps=5 if args.smoke else -1,
    )
    cfg = SFTConfig(**_safe_kwargs(SFTConfig, raw))
    trainer = SFTTrainer(model=model, train_dataset=ds, args=cfg, processing_class=tok)
    trainer.train()
    out_path = "/tmp/opsguard-sft-lora"
    model.save_pretrained(out_path)
    tok.save_pretrained(out_path)
    if hub_repo and not args.smoke:
        _push_lora(out_path, hub_repo + "-sft")
    del trainer, model
    _free()
    return out_path


def stage_dpo(args, model_name, sft_lora_path, hub_repo):
    print("\n" + "=" * 70 + "\n=== STAGE 2: DPO (R2Vul-style preference pairs) ===\n" + "=" * 70, flush=True)
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
    from peft import PeftModel
    from datasets import Dataset
    from trl import DPOConfig, DPOTrainer

    pref_path = "data/preference_pairs.jsonl"
    if not Path(pref_path).exists():
        print(f"WARN: {pref_path} missing — skipping DPO", flush=True)
        return sft_lora_path

    rows = [json.loads(l) for l in open(pref_path)]
    print(f"loaded {len(rows)} preference pairs", flush=True)
    if args.smoke:
        rows = rows[:32]
    ds = Dataset.from_list([
        {"prompt": r["prompt"], "chosen": r["chosen"], "rejected": r["rejected"]}
        for r in rows
    ])

    model_kwargs = {"torch_dtype": torch.bfloat16, "device_map": "auto",
                     "attn_implementation": "sdpa"}
    if args.use_4bit:
        bnb = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16,
                                  bnb_4bit_use_double_quant=True, bnb_4bit_quant_type="nf4")
        model_kwargs["quantization_config"] = bnb
    tok = AutoTokenizer.from_pretrained(sft_lora_path)
    base = AutoModelForCausalLM.from_pretrained(model_name, **model_kwargs)
    base.config.use_cache = False
    if not args.use_4bit:
        base.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
        base.enable_input_require_grads()
    model = PeftModel.from_pretrained(base, sft_lora_path, is_trainable=True)
    for n, p in model.named_parameters():
        if p.requires_grad and "lora_" in n:
            p.data = p.data.to(torch.float32)

    n_epochs = 1 if args.smoke else args.dpo_epochs
    raw = dict(
        output_dir="/tmp/opsguard-dpo",
        per_device_train_batch_size=1,
        gradient_accumulation_steps=8,
        num_train_epochs=n_epochs,
        learning_rate=5e-6,
        max_grad_norm=1.0,
        weight_decay=0.01,
        adam_epsilon=1e-7,
        optim="adamw_torch",
        gradient_checkpointing=True,
        gradient_checkpointing_kwargs={"use_reentrant": False},
        beta=0.1,
        max_length=2048,
        max_prompt_length=1500,
        logging_steps=2,
        save_strategy="epoch",
        save_total_limit=1,
        bf16=True,
        report_to="none",
        push_to_hub=False,
        remove_unused_columns=False,
        max_steps=5 if args.smoke else -1,
    )
    cfg = DPOConfig(**_safe_kwargs(DPOConfig, raw))
    try:
        trainer = DPOTrainer(model=model, args=cfg, train_dataset=ds, processing_class=tok)
    except TypeError:
        trainer = DPOTrainer(model=model, args=cfg, train_dataset=ds, tokenizer=tok)
    trainer.train()
    out_path = "/tmp/opsguard-dpo-lora"
    model.save_pretrained(out_path)
    tok.save_pretrained(out_path)
    if hub_repo and not args.smoke:
        _push_lora(out_path, hub_repo + "-dpo")
    del trainer, model, base
    _free()
    return out_path


def stage_grpo(args, model_name, dpo_lora_path, hub_repo):
    print("\n" + "=" * 70 + "\n=== STAGE 3: GRPO via custom reward fn (no TRL OpenEnv tools) ===\n" + "=" * 70, flush=True)
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
    from peft import PeftModel
    from datasets import Dataset
    from trl import GRPOConfig, GRPOTrainer

    from server.opsguard_environment import OpsguardEnvironment
    from models import OpsguardAction, ActionType
    from scripts.system_prompt import SYSTEM_PROMPT, format_observation, parse_action

    model_kwargs = {"torch_dtype": torch.bfloat16, "device_map": "auto",
                     "attn_implementation": "sdpa"}
    if args.use_4bit:
        bnb = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16,
                                  bnb_4bit_use_double_quant=True, bnb_4bit_quant_type="nf4")
        model_kwargs["quantization_config"] = bnb
    tok = AutoTokenizer.from_pretrained(dpo_lora_path)
    if tok.pad_token_id is None:
        tok.pad_token_id = tok.eos_token_id
    base = AutoModelForCausalLM.from_pretrained(model_name, **model_kwargs)
    base.config.use_cache = False
    if not args.use_4bit:
        base.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
        base.enable_input_require_grads()
    model = PeftModel.from_pretrained(base, dpo_lora_path, is_trainable=True)
    for n, p in model.named_parameters():
        if p.requires_grad and "lora_" in n:
            p.data = p.data.to(torch.float32)

    print("building prompt dataset (sampled env observations)...", flush=True)
    n_prompts = 24 if args.smoke else args.grpo_n_prompts
    scenarios = ["E2_social_eng_buildup", "E3_compromised_maintainer", "E4_multi_vector"]
    prompt_rows = []
    env = OpsguardEnvironment()
    for i in range(n_prompts):
        sid = scenarios[i % len(scenarios)]
        obs = env.reset(scenario_id=sid, seed=i)
        for _ in range(random.randint(0, 5)):
            obs = env.step(OpsguardAction(action_type=ActionType.WAIT))
            if obs.done:
                obs = env.reset(scenario_id=sid, seed=i)
                break
        prompt_text = SYSTEM_PROMPT + "\n\nOBSERVATION:\n" + format_observation(obs) + "\n\nACTION:\n"
        prompt_rows.append({"prompt": prompt_text, "scenario": sid, "seed": i})
    ds = Dataset.from_list(prompt_rows)

    def reward_fn(completions, prompts=None, scenario=None, seed=None, **kwargs):
        rewards = []
        scenarios_b = scenario if scenario is not None else ["E2_social_eng_buildup"] * len(completions)
        seeds_b = seed if seed is not None else [0] * len(completions)
        for comp, sid, sd in zip(completions, scenarios_b, seeds_b):
            try:
                action = parse_action(comp)
                env_local = OpsguardEnvironment()
                env_local.reset(scenario_id=sid, seed=int(sd))
                obs = env_local.step(action)
                rewards.append(float(obs.reward) if obs.reward is not None else 0.0)
            except Exception:
                rewards.append(-1.0)
        return rewards

    n_steps = 5 if args.smoke else args.grpo_steps
    raw = dict(
        output_dir="/tmp/opsguard-grpo",
        per_device_train_batch_size=1,
        gradient_accumulation_steps=4,
        num_generations=args.grpo_num_generations,
        max_steps=n_steps,
        max_completion_length=256,
        max_prompt_length=1500,
        beta=0.001,
        learning_rate=5e-6,
        max_grad_norm=1.0,
        weight_decay=0.01,
        adam_epsilon=1e-7,
        optim="adamw_torch",
        gradient_checkpointing=True,
        gradient_checkpointing_kwargs={"use_reentrant": False},
        logging_steps=1,
        save_strategy="no",
        bf16=True,
        report_to="none",
        push_to_hub=False,
        use_vllm=False,
    )
    cfg = GRPOConfig(**_safe_kwargs(GRPOConfig, raw))
    try:
        trainer = GRPOTrainer(model=model, args=cfg, reward_funcs=reward_fn,
                              train_dataset=ds, processing_class=tok)
    except TypeError:
        trainer = GRPOTrainer(model=model, args=cfg, reward_funcs=reward_fn,
                              train_dataset=ds, tokenizer=tok)
    trainer.train()
    out_path = "/tmp/opsguard-grpo-lora"
    model.save_pretrained(out_path)
    tok.save_pretrained(out_path)
    if hub_repo and not args.smoke:
        _push_lora(out_path, hub_repo + "-grpo")
    del trainer, model, base
    _free()
    return out_path


def stage_eval(args, model_name, final_lora_path):
    print("\n" + "=" * 70 + "\n=== STAGE 4: EVAL — trained vs base on Datadog held-out ===\n" + "=" * 70, flush=True)
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
    from peft import PeftModel
    from models import OpsguardAction, ActionType
    from scripts.system_prompt import SYSTEM_PROMPT, format_observation, parse_action

    cve_rows = [json.loads(l) for l in open("data/datadog_extracted.jsonl")]
    if args.smoke:
        cve_rows = cve_rows[:20]
    print(f"evaluating on {len(cve_rows)} real Datadog samples...", flush=True)

    model_kwargs = {"torch_dtype": torch.bfloat16, "device_map": "auto"}
    if args.use_4bit:
        bnb = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16,
                                  bnb_4bit_use_double_quant=True, bnb_4bit_quant_type="nf4")
        model_kwargs["quantization_config"] = bnb
    tok = AutoTokenizer.from_pretrained(final_lora_path)
    base = AutoModelForCausalLM.from_pretrained(model_name, **model_kwargs)
    model = PeftModel.from_pretrained(base, final_lora_path)
    model.eval()

    @torch.inference_mode()
    def predict(text_payload):
        prompt = SYSTEM_PROMPT + "\n\nOBSERVATION:\n" + text_payload + "\n\nACTION:\n"
        inp = tok(prompt, return_tensors="pt", truncation=True, max_length=1800).to(model.device)
        out = model.generate(**inp, max_new_tokens=200, do_sample=False,
                              pad_token_id=tok.eos_token_id)
        return tok.decode(out[0][inp["input_ids"].shape[1]:], skip_special_tokens=True)

    REJECT = {"reject_pr", "close_spam", "flag_security"}
    n_caught = 0
    for r in cve_rows:
        synth_obs = json.dumps({
            "current_issue": {
                "title": r.get("package", "") + " " + r.get("filename", ""),
                "body": r.get("diff_preview", "")[:1000],
                "is_pr": True,
                "author_login": "unknown",
                "pr_diff_preview": r.get("diff_preview", "")[:800],
            }
        })
        try:
            text = predict(synth_obs)
            action = parse_action(text)
            at = action.action_type.value if hasattr(action.action_type, "value") else str(action.action_type)
            if at in REJECT:
                n_caught += 1
        except Exception:
            pass
    catch_rate = n_caught / max(1, len(cve_rows))
    print(f"\nFINAL: trained LoRA caught {n_caught}/{len(cve_rows)} = {catch_rate:.3f} catch rate on Datadog held-out", flush=True)
    Path("/tmp/eval_post_train").mkdir(exist_ok=True)
    Path("/tmp/eval_post_train/cve_results.json").write_text(json.dumps({
        "n_samples": len(cve_rows),
        "n_caught": n_caught,
        "catch_rate": round(catch_rate, 3),
    }))
    return catch_rate


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--smoke", action="store_true",
                    help="5 steps per stage, ~10 min on Colab T4. Validates code path.")
    ap.add_argument("--model", type=str, default="Qwen/Qwen2.5-7B-Instruct")
    ap.add_argument("--use-4bit", action="store_true",
                    help="Use bnb 4-bit (default false = bf16 full). Toggle for VRAM-constrained env.")
    ap.add_argument("--lora-r", type=int, default=32,
                    help="LoRA rank (default 32 for stability). 128 for bold.")
    ap.add_argument("--lora-alpha", type=int, default=64,
                    help="LoRA alpha (default 64 = 2*r, mild scaling).")
    ap.add_argument("--sft-epochs", type=int, default=5)
    ap.add_argument("--dpo-epochs", type=int, default=3)
    ap.add_argument("--grpo-steps", type=int, default=200)
    ap.add_argument("--grpo-num-generations", type=int, default=4)
    ap.add_argument("--grpo-n-prompts", type=int, default=256)
    ap.add_argument("--batch-size", type=int, default=2)
    ap.add_argument("--grad-accum", type=int, default=4)
    ap.add_argument("--hub-repo", type=str, default="sai1906/opsguard",
                    help="Hub repo prefix; suffixes -sft, -dpo, -grpo appended.")
    ap.add_argument("--skip", type=str, default="",
                    help="Comma-separated stages to skip: sft,dpo,grpo,eval")
    args = ap.parse_args()

    if args.smoke and not args.use_4bit:
        print("[smoke] forcing --use-4bit on Colab/free GPUs", flush=True)
        args.use_4bit = True

    skipped = set(s.strip() for s in args.skip.split(",") if s.strip())

    print(f"=== BOLD PIPELINE  smoke={args.smoke}  4bit={args.use_4bit}  "
          f"r={args.lora_r}  alpha={args.lora_alpha} ===", flush=True)
    print(f"  SFT epochs={args.sft_epochs}  DPO epochs={args.dpo_epochs}  GRPO steps={args.grpo_steps}", flush=True)

    sft_lora = stage_sft(args, args.model, args.hub_repo) if "sft" not in skipped else None
    dpo_lora = stage_dpo(args, args.model, sft_lora or args.model, args.hub_repo) if "dpo" not in skipped else sft_lora
    grpo_lora = stage_grpo(args, args.model, dpo_lora or sft_lora or args.model, args.hub_repo) if "grpo" not in skipped else dpo_lora

    final = grpo_lora or dpo_lora or sft_lora
    if final and "eval" not in skipped:
        stage_eval(args, args.model, final)

    print("\n=== DONE ===", flush=True)


if __name__ == "__main__":
    main()