"""Modal PKPO pipeline: baseline eval, tiny warmup, PKPO LoRA, upload."""
from __future__ import annotations

import gc
import inspect
import json
import os
import shutil
import time
from pathlib import Path

import modal

from modal_common import BASE_MODEL, CACHE_DIR, HF_REPO, WORK_DIR, hf_cache_vol, image, work_vol

app = modal.App("pkpo-train")


def _copy_source_bundle(dst: Path):
    import agent_core
    import eval_lib
    import pkpo
    import shipped_tool
    import train_pkpo_modal

    for mod, name in [
        (agent_core, "agent_core.py"),
        (eval_lib, "eval_lib.py"),
        (pkpo, "pkpo.py"),
        (shipped_tool, "shipped_tool.py"),
        (train_pkpo_modal, "train_pkpo_modal.py"),
    ]:
        Path(dst / name).write_text(Path(inspect.getfile(mod)).read_text(encoding="utf-8"), encoding="utf-8")


def _load_tokenizer(model_ref: str):
    from transformers import AutoTokenizer
    from agent_core import CHAT_TEMPLATE

    tok = AutoTokenizer.from_pretrained(model_ref, trust_remote_code=True)
    tok.chat_template = CHAT_TEMPLATE
    if tok.pad_token is None:
        tok.pad_token = tok.eos_token
    return tok


def _write_generation_config(model, tok, dst: Path):
    from transformers import GenerationConfig
    from agent_core import GEN_TEMPERATURE, GEN_TOP_P, MAX_TURN_TOKENS, STOP_STRING

    gen = GenerationConfig.from_model_config(model.config)
    gen.do_sample = True
    gen.temperature = GEN_TEMPERATURE
    gen.top_p = GEN_TOP_P
    gen.max_new_tokens = MAX_TURN_TOKENS
    gen.pad_token_id = tok.pad_token_id
    gen.eos_token_id = tok.eos_token_id
    gen.stop_strings = [STOP_STRING]
    gen.save_pretrained(dst)


def _read_json(path: Path, default):
    if path.exists():
        return json.loads(path.read_text(encoding="utf-8"))
    return default


def _model_card(baseline, final, schedule, notes: str) -> str:
    b = baseline.get("pass_at_1", 0.0) if baseline else 0.0
    f = final.get("pass_at_1", 0.0) if final else 0.0
    return f"""---
license: other
base_model: Qwen/Qwen3-8B-Base
tags:
- code
- reinforcement-learning
- pkpo
- livecodebench
---

# Qwen3-8B Code PKPO

This repo contains a Qwen/Qwen3-8B-Base derivative trained for a small agentic
coding experiment using the shared tool path in `agent_core.py` and
`shipped_tool.py`.

## Method

- Base: `Qwen/Qwen3-8B-Base`.
- Prompt/template: custom `<think>...</think><answer>...</answer>` template saved
  in the tokenizer. The generation prompt ends with `Assistant: <think>`.
- Tool protocol: no system role; instructions are merged into the first user
  message; strict user/assistant alternation; plain-text `Tool type` and
  `Tool query` calls.
- Training data: `deepmind/code_contests` train split only, filtered to old
  stdin/stdout problems. The LiveCodeBench eval subset is not used for training.
- Reward: binary hidden-test pass/fail.
- PKPO: `sloo_minus_one` from the paper for `k >= 2`; centered `k=1` rewards for
  the first and final stages. No GRPO-style reward normalization is applied.
- Schedule actually run: `{schedule}`.

The run was intentionally small to fit the free-credit budget and deadline.
Results should be treated as a reproducible experiment, not a leaderboard model.

## Results

Evaluation uses `livecodebench/code_generation_lite` `v6`, a fixed subset saved
at `eval/eval_subset.json`, temperature 1.0, and the same one-turn tool path used
for training.

| model | pass@1 estimate |
|---|---:|
| base before training | {b:.4f} |
| final merged model | {f:.4f} |

Raw files:

- `eval/baseline_results.json`
- `eval/final_results.json`
- `eval/eval_subset.json`

## Usage

Serve with vLLM:

```bash
vllm serve bk1dr/qwen3-8b-code-pkpo --trust-remote-code --max-model-len 8192
```

Run the shipped tool:

```bash
python shipped_tool.py --base-url http://127.0.0.1:8000/v1 --model bk1dr/qwen3-8b-code-pkpo --max-turns 1 --cp < problem.txt
```

## Run Notes

{notes}
"""


@app.function(
    image=image,
    volumes={CACHE_DIR: hf_cache_vol, WORK_DIR: work_vol},
    secrets=[modal.Secret.from_name("hf-secret")],
    timeout=1800,
    cpu=4,
    memory=16384,
)
def push_scaffold():
    from huggingface_hub import HfApi, snapshot_download
    from transformers import AutoConfig
    from agent_core import CHAT_TEMPLATE

    base_path = snapshot_download(BASE_MODEL, cache_dir=CACHE_DIR)
    out = Path(WORK_DIR) / "scaffold_repo"
    if out.exists():
        shutil.rmtree(out)
    out.mkdir(parents=True)
    tok = _load_tokenizer(base_path)
    tok.save_pretrained(out)
    cfg = AutoConfig.from_pretrained(base_path, trust_remote_code=True)
    cfg.save_pretrained(out)
    (out / "chat_template.jinja").write_text(CHAT_TEMPLATE, encoding="utf-8")
    _copy_source_bundle(out)
    (out / "README.md").write_text(_model_card({}, {}, "pending", "Scaffold pushed before GPU training."), encoding="utf-8")
    HfApi().upload_folder(repo_id=HF_REPO, repo_type="model", folder_path=str(out), commit_message="Initial scaffold")
    work_vol.commit()
    return {"uploaded": str(out)}


@app.function(
    image=image,
    gpu="H100",
    volumes={CACHE_DIR: hf_cache_vol},
    timeout=15 * 60,
    cpu=4,
    memory=98304,
)
def smoke_batched_generation(max_new_tokens: int = 96):
    """Low-cost preflight for the exact batched stop-string generation API."""
    import torch
    from huggingface_hub import snapshot_download
    from transformers import AutoModelForCausalLM, AutoTokenizer

    from agent_core import CHAT_TEMPLATE, STOP_STRING, coding_task_message, render_messages

    base_path = snapshot_download(BASE_MODEL, cache_dir=CACHE_DIR)
    tok = AutoTokenizer.from_pretrained(base_path, trust_remote_code=True)
    tok.chat_template = CHAT_TEMPLATE
    if tok.pad_token is None:
        tok.pad_token = tok.eos_token
    tok.padding_side = "left"
    model = AutoModelForCausalLM.from_pretrained(
        base_path,
        torch_dtype=torch.bfloat16,
        device_map={"": 0},
        trust_remote_code=True,
        low_cpu_mem_usage=True,
    )
    prompt = render_messages([{
        "role": "user",
        "content": coding_task_message(
            "Read an integer N and print N plus one.", max_turns=1
        ),
    }], add_generation_prompt=True)
    inputs = tok([prompt, prompt], return_tensors="pt", padding=True).to(model.device)
    prompt_width = inputs.input_ids.shape[1]
    with torch.no_grad():
        out = model.generate(
            **inputs,
            do_sample=True,
            temperature=1.0,
            top_p=0.95,
            max_new_tokens=max_new_tokens,
            pad_token_id=tok.pad_token_id,
            eos_token_id=tok.eos_token_id,
            stop_strings=[STOP_STRING],
            tokenizer=tok,
        )
    completions = [tok.decode(row[prompt_width:], skip_special_tokens=False) for row in out]
    result = {
        "batch_size": len(completions),
        "completion_lengths": [len(c) for c in completions],
        "stopped": [STOP_STRING in c for c in completions],
        "samples": [c[:500] for c in completions],
    }
    print(json.dumps(result), flush=True)
    return result


@app.function(
    image=image,
    gpu="H100",
    volumes={CACHE_DIR: hf_cache_vol, WORK_DIR: work_vol},
    secrets=[modal.Secret.from_name("hf-secret")],
    timeout=125 * 60,
    cpu=10,
    memory=98304,
)
def train_eval_upload(
    eval_limit: int = 6,
    eval_samples: int = 6,
    train_groups_per_stage: int = 4,
    rollouts: int = 12,
    max_new_tokens: int = 1024,
    sft_examples: int = 96,
    finish_by_epoch: float = 0.0,
):
    """Train a checkpointed one-turn agentic coding LoRA with PKPO.

    A one-turn episode is intentional here: it uses the same ``Episode`` parser
    and final-action contract as the shipped tool while making the short budget
    practical.  Every generated completion is durably written to a private text
    trace and JSONL before it is judged.
    """
    import hashlib
    import numpy as np
    import torch
    from huggingface_hub import HfApi, snapshot_download
    from peft import LoraConfig, get_peft_model
    from torch.nn.utils import clip_grad_norm_
    from transformers import AutoModelForCausalLM, AutoTokenizer

    from agent_core import (CHAT_TEMPLATE, STOP_STRING, THINK_PREFIX, Episode,
                            extract_code, parse_action, render_messages)
    from eval_lib import (compact_json_dump, judge_final_answer, load_codecontest_train,
                          load_lcb_v6_subset, python3_verified_solutions,
                          short_completion_for_sft)
    from pkpo import transform_rewards

    torch.backends.cuda.matmul.allow_tf32 = True
    started = time.time()
    # Wall-clock guard: the final eval + merge + save + upload need a fixed
    # reserve that SFT/RL must never eat into. finish_by_epoch is when the
    # merged upload must be DONE (epoch seconds).
    finish_by = finish_by_epoch if finish_by_epoch > started else started + 100 * 60
    FINAL_RESERVE_S = 30 * 60
    rl_stop_at = finish_by - FINAL_RESERVE_S
    print("time_guard config", {"now": int(started), "finish_by": int(finish_by),
                                "rl_stop_at": int(rl_stop_at)}, flush=True)
    run_id = time.strftime("pkpo_%Y%m%dT%H%M%SZ", time.gmtime())
    run_dir = Path(WORK_DIR) / "results" / run_id
    checkpoint_root = Path(WORK_DIR) / "checkpoints" / run_id
    run_dir.mkdir(parents=True, exist_ok=False)
    checkpoint_root.mkdir(parents=True, exist_ok=True)

    def save_json(name: str, value):
        compact_json_dump(run_dir / name, value)

    def safe_verdict(verdict: dict) -> dict:
        """Keep useful diagnostics without saving hidden test values."""
        out = {
            "passed": bool(verdict.get("passed", False)),
            "error": verdict.get("error", "ok"),
            "passed_tests": int(verdict.get("passed_tests", 0)),
            "total_tests": int(verdict.get("total_tests", 0)),
        }
        for key in ("seconds", "exit_code"):
            if key in verdict:
                out[key] = verdict[key]
        if verdict.get("stderr"):
            out["stderr_tail"] = str(verdict["stderr"])[-500:]
        return out

    class RolloutTrace:
        """Private structured and plain-text traces, flushed per sampled output."""
        def __init__(self, root: Path):
            self.jsonl = (root / "rollouts.jsonl").open("a", encoding="utf-8")
            self.text = (root / "raw_rollouts.txt").open("a", encoding="utf-8")

        def _write(self, record: dict):
            self.jsonl.write(json.dumps(record, ensure_ascii=False, sort_keys=True) + "\n")
            self.jsonl.flush()

        def rollout(self, record: dict):
            self._write(record)
            header = {
                key: record.get(key)
                for key in ("run_id", "phase", "stage", "k", "group", "candidate_attempt",
                            "sample", "problem_id", "tool_type", "stop_seen")
            }
            self.text.write("\n===== ROLLOUT =====\n")
            self.text.write(json.dumps(header, sort_keys=True) + "\n")
            self.text.write("--- completion_raw ---\n")
            self.text.write(record.get("completion_raw", "") + "\n")
            self.text.write("--- final_answer ---\n")
            self.text.write(record.get("final_answer", "") + "\n")
            self.text.write("--- verdict ---\n")
            self.text.write(json.dumps(record.get("verdict", {}), ensure_ascii=False, sort_keys=True) + "\n")
            self.text.write("===== END ROLLOUT =====\n")
            self.text.flush()

        def summary(self, record: dict):
            self._write(record)
            self.text.write("\n===== GROUP SUMMARY =====\n")
            self.text.write(json.dumps(record, ensure_ascii=False, sort_keys=True) + "\n")
            self.text.write("===== END GROUP SUMMARY =====\n")
            self.text.flush()

        def close(self):
            self.jsonl.close()
            self.text.close()

    trace = RolloutTrace(run_dir)

    base_path = snapshot_download(BASE_MODEL, cache_dir=CACHE_DIR)
    tok = AutoTokenizer.from_pretrained(base_path, trust_remote_code=True)
    tok.chat_template = CHAT_TEMPLATE
    if tok.pad_token is None:
        tok.pad_token = tok.eos_token
    tok.padding_side = "left"

    def load_model():
        model = AutoModelForCausalLM.from_pretrained(
            base_path,
            torch_dtype=torch.bfloat16,
            device_map={"": 0},
            trust_remote_code=True,
            low_cpu_mem_usage=True,
        )
        model.config.use_cache = False
        model.gradient_checkpointing_enable()
        lora = LoraConfig(
            r=32,
            lora_alpha=64,
            lora_dropout=0.02,
            bias="none",
            task_type="CAUSAL_LM",
            target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
        )
        model = get_peft_model(model, lora)
        model.enable_input_require_grads()
        model.print_trainable_parameters()
        return model

    model = load_model()

    def generate_many(first_user_msg: str, count: int) -> list[str]:
        """Sample a group together; each row has an independent stop string."""
        if count <= 0:
            return []
        prompt = render_messages([{"role": "user", "content": first_user_msg}], add_generation_prompt=True)
        inputs = tok([prompt] * count, return_tensors="pt", padding=True).to(model.device)
        input_width = inputs.input_ids.shape[1]
        was_training = model.training
        old_cache = getattr(model.config, "use_cache", False)
        model.eval()
        model.config.use_cache = True
        try:
            with torch.no_grad():
                out = model.generate(
                    **inputs,
                    do_sample=True,
                    temperature=1.0,
                    top_p=0.95,
                    max_new_tokens=max_new_tokens,
                    pad_token_id=tok.pad_token_id,
                    eos_token_id=tok.eos_token_id,
                    use_cache=True,
                    stop_strings=[STOP_STRING],
                    tokenizer=tok,
                )
        finally:
            model.config.use_cache = old_cache
            if was_training:
                model.train()
        return [tok.decode(row[input_width:], skip_special_tokens=False) for row in out]

    eval_subset = load_lcb_v6_subset(limit=eval_limit, seed=7341)
    save_json("eval_subset.json", [p.to_public_dict() for p in eval_subset])

    def sample_episodes(problem, count: int, phase: str, stage=None, k=None, group=None,
                        candidate_attempt=None, max_tests: int = 18):
        """Drive the exact one-turn ``Episode`` state machine for every sample."""
        first_user_msg = problem.first_user_message(max_turns=1)
        prompt_hash = hashlib.sha256(
            render_messages([{"role": "user", "content": first_user_msg}], add_generation_prompt=True).encode()
        ).hexdigest()
        completions = generate_many(first_user_msg, count)
        samples = []
        for sample_idx, completion in enumerate(completions):
            episode = Episode(first_user_msg, max_turns=1)
            kind, _ = episode.step(completion)
            tool_type, _, parse_error = parse_action(completion)
            stop_seen = STOP_STRING in completion
            valid_final = bool(stop_seen and kind == "final" and episode.final_answer is not None)
            if valid_final:
                final_answer = episode.final_answer
                extracted = extract_code(final_answer)
                verdict = judge_final_answer(
                    final_answer, problem.hidden_tests, timeout_s=3, memory_mb=768, max_tests=max_tests
                )
            else:
                final_answer = episode.final_answer or ""
                extracted = ""
                reason = parse_error or episode.end_reason or "missing_final_action"
                verdict = {
                    "passed": False,
                    "passed_tests": 0,
                    "total_tests": min(max_tests, len(problem.hidden_tests)),
                    "error": "format_error: " + str(reason),
                }
            safe = safe_verdict(verdict)
            record = {
                "schema_version": 1,
                "record_type": "rollout",
                "run_id": run_id,
                "timestamp_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
                "phase": phase,
                "stage": stage,
                "k": k,
                "group": group,
                "candidate_attempt": candidate_attempt,
                "sample": sample_idx,
                "problem_id": problem.problem_id,
                "source": problem.source,
                "sampling": {"temperature": 1.0, "top_p": 0.95, "max_new_tokens": max_new_tokens},
                "prompt_sha256": prompt_hash,
                "prompt_chars": len(first_user_msg),
                "completion_raw": completion,
                "assistant_content": THINK_PREFIX + completion,
                "stop_seen": stop_seen,
                "tool_type": tool_type,
                "parse_error": parse_error,
                "final_answer": final_answer,
                "extracted_code": extracted,
                "verdict": safe,
            }
            trace.rollout(record)
            samples.append({
                "completion": completion,
                "valid_final": valid_final,
                "verdict": verdict,
                "safe_verdict": safe,
                "code": extracted,
                "tool_type": tool_type,
            })
            print(phase, problem.problem_id, sample_idx + 1,
                  "pass" if verdict.get("passed") else verdict.get("error", "fail"), flush=True)
        return samples

    def evaluate(tag: str) -> dict:
        details = []
        pass_count = 0
        total = 0
        for p in eval_subset:
            sample_results = []
            samples = sample_episodes(p, eval_samples, phase=tag, max_tests=80)
            for sample in samples:
                verdict = sample["verdict"]
                ok = bool(verdict["passed"])
                pass_count += int(ok)
                total += 1
                sample_results.append({
                    "passed": ok,
                    "end": sample["safe_verdict"].get("error", "ok"),
                    "completion_chars": len(sample["completion"]),
                    "code_chars": len(sample["code"]),
                    "stop_seen": STOP_STRING in sample["completion"],
                    "tool_type": sample["tool_type"],
                })
            details.append({"problem_id": p.problem_id, "difficulty": p.difficulty, "samples": sample_results})
            work_vol.commit()
        res = {
            "tag": tag,
            "pass_at_1": pass_count / max(total, 1),
            "passed_samples": pass_count,
            "total_samples": total,
            "eval_limit": eval_limit,
            "samples_per_problem": eval_samples,
            "temperature": 1.0,
            "top_p": 0.95,
            "max_new_tokens": max_new_tokens,
            "tool_turn_limit": 1,
            "details": details,
        }
        save_json(f"{tag}_results.json", res)
        work_vol.commit()
        print(tag, "pass@1", res["pass_at_1"], flush=True)
        return res

    def checkpoint(tag: str, state: dict):
        """Persist the small LoRA adapter after every recoverable unit of work."""
        path = checkpoint_root / tag
        if path.exists():
            shutil.rmtree(path)
        path.mkdir(parents=True)
        model.save_pretrained(path)
        tok.save_pretrained(path)
        compact_json_dump(path / "state.json", state)
        compact_json_dump(run_dir / "latest_state.json", {"checkpoint": str(path), **state})
        work_vol.commit()
        return str(path)

    def train_completion(first_user_msg: str, completion: str, weight: float):
        """Backprop only a complete assistant final turn; never truncate silently."""
        if abs(weight) < 1e-9:
            return None, 0.0, "zero_weight"
        stop_idx = completion.find(STOP_STRING)
        if stop_idx == -1:
            return None, 0.0, "missing_stop"
        prompt = render_messages([{"role": "user", "content": first_user_msg}], add_generation_prompt=True)
        prompt_ids = tok(prompt, return_tensors="pt", add_special_tokens=False).input_ids[0]
        # Batched generate right-pads finished rows with eos/pad until the longest
        # row stops; training on that tail drowns the real signal. The target is
        # exactly the turn content through </answer> plus ONE terminating EOS.
        target = completion[: stop_idx + len(STOP_STRING)]
        if tok.eos_token:
            target += tok.eos_token
        full = tok(prompt + target, return_tensors="pt", add_special_tokens=False)
        if full.input_ids.shape[1] > 8192:
            return None, 0.0, "overlength"
        input_ids = full.input_ids.to(model.device)
        labels = input_ids.clone()
        labels[:, : min(prompt_ids.numel(), labels.shape[1])] = -100
        if (labels != -100).sum() == 0:
            return None, 0.0, "empty_target"
        out = model(input_ids=input_ids, labels=labels)
        nll = float(out.loss.detach().cpu())
        weighted_loss = out.loss * float(weight)
        weighted_loss.backward()
        return nll, float(weighted_loss.detach().cpu()), None

    baseline = evaluate("baseline")
    opt = torch.optim.AdamW((p for p in model.parameters() if p.requires_grad), lr=4e-5, betas=(0.9, 0.95))

    print("verified Python-3 format warmup", flush=True)
    model.train()
    sft_log = []
    sft_pairs = python3_verified_solutions(limit=sft_examples)
    sft_time_floor = rl_stop_at - 25 * 60  # keep at least ~25 min of RL window
    for sft_idx, (prob, sol) in enumerate(sft_pairs, start=1):
        if time.time() > sft_time_floor:
            print("time_guard: stopping SFT early at", sft_idx - 1, "examples", flush=True)
            break
        opt.zero_grad(set_to_none=True)
        nll, loss, skipped = train_completion(
            prob.first_user_message(max_turns=1), short_completion_for_sft(sol), 1.0
        )
        entry = {"index": sft_idx, "problem_id": prob.problem_id, "nll": nll, "loss": loss,
                 "skipped": skipped}
        if skipped is None:
            entry["grad_norm"] = float(clip_grad_norm_(model.parameters(), 1.0))
            opt.step()
        sft_log.append(entry)
        print("sft", prob.problem_id, entry, flush=True)
        if sft_idx % 24 == 0:
            save_json("sft_log.json", sft_log)
            checkpoint(f"sft_{sft_idx:04d}", {"phase": "sft", "completed_examples": sft_idx})
    save_json("sft_log.json", sft_log)
    sft_checkpoint = checkpoint("after_sft", {
        "phase": "sft_complete",
        "completed_examples": len(sft_log),
        "verified_examples": len(sft_pairs),
    })

    print("loading disjoint RL curriculum", flush=True)
    schedule = [1, 8, 1]
    sft_ids = {p.problem_id for p, _ in sft_pairs}
    train_problems = load_codecontest_train(
        limit=max(72, train_groups_per_stage * len(schedule) * 8),
        seed=20260709,
        exclude_problem_ids=sft_ids,
    )
    print("loaded RL candidates", len(train_problems), flush=True)
    stage_log = []
    candidate_cursor = 0
    # Split the remaining RL window across stages so the mandatory final k=1
    # stage always gets its share even when earlier stages run long.
    rl_start = time.time()
    rl_window = max(rl_stop_at - rl_start, 0.0)
    stage_shares = [0.30, 0.40, 0.30]
    assert len(stage_shares) == len(schedule)
    stage_deadline = [rl_start + rl_window * sum(stage_shares[: i + 1]) for i in range(len(schedule))]
    print("rl window minutes", round(rl_window / 60, 1), flush=True)
    for stage_idx, k in enumerate(schedule):
        opt.param_groups[0]["lr"] = 2e-5 if k == 1 else 1e-5
        for group_idx in range(train_groups_per_stage):
            if time.time() > stage_deadline[stage_idx]:
                print("time_guard: closing stage", stage_idx, "k", k, "after", group_idx, "groups", flush=True)
                stage_log.append({"stage": stage_idx, "k": k, "group": group_idx,
                                  "skipped": "stage_time_guard"})
                save_json("training_stage_log.json", stage_log)
                break
            selected = None
            attempts = []
            for attempt_idx in range(1, 6):
                if candidate_cursor >= len(train_problems) or time.time() > stage_deadline[stage_idx]:
                    break
                p = train_problems[candidate_cursor]
                candidate_cursor += 1
                samples = sample_episodes(
                    p, rollouts, phase="rl_candidate", stage=stage_idx, k=k, group=group_idx,
                    candidate_attempt=attempt_idx,
                )
                rewards = [1.0 if sample["verdict"].get("passed") else 0.0 for sample in samples]
                advantages = transform_rewards(np.array(rewards, dtype=np.float64), k)
                eligible = bool(np.any(np.abs(advantages) > 1e-10))
                attempt = {
                    "problem_id": p.problem_id,
                    "passes": int(sum(rewards)),
                    "rewards": rewards,
                    "eligible": eligible,
                    "advantages": [float(x) for x in advantages],
                }
                attempts.append(attempt)
                trace.summary({
                    "schema_version": 1,
                    "record_type": "candidate_summary",
                    "run_id": run_id,
                    "stage": stage_idx,
                    "k": k,
                    "group": group_idx,
                    "candidate_attempt": attempt_idx,
                    **attempt,
                })
                if eligible:
                    selected = (p, samples, rewards, advantages)
                    break

            if selected is None:
                entry = {
                    "stage": stage_idx,
                    "k": k,
                    "group": group_idx,
                    "skipped": "no_pkpo_eligible_group",
                    "candidate_attempts": attempts,
                }
                stage_log.append(entry)
                save_json("training_stage_log.json", stage_log)
                checkpoint(f"stage{stage_idx}_group{group_idx}_skipped", {
                    "phase": "rl", "stage": stage_idx, "k": k, "group": group_idx,
                    "status": "skipped", "candidate_cursor": candidate_cursor,
                })
                print("stage skipped", entry, flush=True)
                continue

            p, samples, rewards, adv = selected
            opt.zero_grad(set_to_none=True)
            losses = []
            nlls = []
            skipped_samples = []
            first_user_msg = p.first_user_message(max_turns=1)
            for sample_idx, (sample, a) in enumerate(zip(samples, adv)):
                if not sample["valid_final"]:
                    skipped_samples.append({"sample": sample_idx, "reason": "invalid_final"})
                    losses.append(0.0)
                    nlls.append(None)
                    continue
                nll, loss, skipped = train_completion(first_user_msg, sample["completion"], float(a) / rollouts)
                nlls.append(nll)
                losses.append(loss)
                if skipped not in (None, "zero_weight"):
                    skipped_samples.append({"sample": sample_idx, "reason": skipped})
            if any(nll is not None for nll in nlls):
                grad_norm = float(clip_grad_norm_(model.parameters(), 1.0))
                opt.step()
            else:
                grad_norm = 0.0
            entry = {
                "stage": stage_idx,
                "k": k,
                "group": group_idx,
                "problem_id": p.problem_id,
                "rewards": rewards,
                "advantages": [float(x) for x in adv],
                "losses": losses,
                "nlls": nlls,
                "skipped_samples": skipped_samples,
                "candidate_attempts": attempts,
                "grad_norm": grad_norm,
            }
            stage_log.append(entry)
            save_json("training_stage_log.json", stage_log)
            ckpt = checkpoint(f"stage{stage_idx}_group{group_idx}", {
                "phase": "rl", "stage": stage_idx, "k": k, "group": group_idx,
                "problem_id": p.problem_id, "candidate_cursor": candidate_cursor,
                "after_sft_checkpoint": sft_checkpoint,
            })
            entry["checkpoint"] = ckpt
            save_json("training_stage_log.json", stage_log)
            trace.summary({
                "schema_version": 1,
                "record_type": "training_group_summary",
                "run_id": run_id,
                **entry,
            })
            print("stage", entry, flush=True)

    final = evaluate("final")

    final_dir = Path(WORK_DIR) / "final_model"
    if final_dir.exists():
        shutil.rmtree(final_dir)
    final_dir.mkdir(parents=True)
    model.eval()
    merged = model.merge_and_unload()
    merged.config.use_cache = True
    merged.save_pretrained(final_dir, safe_serialization=True, max_shard_size="4GB")
    tok.save_pretrained(final_dir)
    _write_generation_config(merged, tok, final_dir)
    _copy_source_bundle(final_dir)
    eval_dir = final_dir / "eval"
    eval_dir.mkdir()
    for name in ["baseline_results.json", "final_results.json", "eval_subset.json", "training_stage_log.json", "sft_log.json"]:
        src = run_dir / name
        if src.exists():
            shutil.copyfile(src, eval_dir / name)
    notes = (
        f"Modal H100 pipeline elapsed {(time.time() - started) / 60:.1f} minutes. "
        f"Run id {run_id}; eval subset size {eval_limit}, samples/problem {eval_samples}, "
        f"rollouts/group {rollouts}; verified Python-3 SFT examples {len(sft_pairs)}. "
        "Raw rollout traces remain in the private Modal work volume and are not published."
    )
    (final_dir / "README.md").write_text(_model_card(baseline, final, schedule, notes), encoding="utf-8")
    HfApi().upload_folder(repo_id=HF_REPO, repo_type="model", folder_path=str(final_dir), commit_message="Upload merged PKPO run")
    trace.close()
    work_vol.commit()
    return {
        "run_id": run_id,
        "baseline": baseline["pass_at_1"],
        "final": final["pass_at_1"],
        "model_dir": str(final_dir),
        "trace_text": str(run_dir / "raw_rollouts.txt"),
    }


@app.local_entrypoint()
def main(
    eval_limit: int = 6,
    eval_samples: int = 6,
    train_groups_per_stage: int = 4,
    rollouts: int = 12,
    max_new_tokens: int = 1024,
    sft_examples: int = 96,
    finish_by_epoch: float = 0.0,
):
    print("Spawning scaffold push (runs in parallel on CPU)...")
    scaffold_call = push_scaffold.spawn()
    print("Running GPU pipeline...")
    print(train_eval_upload.remote(
        eval_limit,
        eval_samples,
        train_groups_per_stage,
        rollouts,
        max_new_tokens,
        sft_examples,
        finish_by_epoch,
    ))
    print("scaffold:", scaffold_call.get())