examples/sdpo_with_real_traces_production/run.py

"""Production-grade SDPO end-to-end on real Claude Code traces (CPU, ~2min).

This is the FIFTH example in the SDPO progression — the production-grade
sibling to `examples/sdpo_with_real_traces/`:

  examples/gsm8k_grpo/                       -- plain GRPO baseline
  examples/gsm8k_grpo_with_sdpo/             -- SDPO on hand-crafted prompts
  examples/sdpo_with_real_traces/            -- SDPO WIRING smoke (misaligned)
  examples/sdpo_with_real_traces_production/ -- SDPO PRODUCTION-GRADE (this)

Where `sdpo_with_real_traces` was a wiring-only smoke (HINT appended to
messages → student/teacher right-edge tokens diverge → JSD measured on
different content), THIS example uses the production path:

  ClaudeCodeIngester
    → claude_states_to_trace_examples()  [Wave 19 NEW adapter]
    → ComposerDataCollator(hint_generator=...)
    → batch with PROPERLY-ALIGNED ctx_teacher_input_ids + sdpo_loss_mask
    → compose_loss

The data collator's `_build_hint_injected_trace` walks the turns,
detects error sites via `tool_error` markers, injects the hint as a
system turn BEFORE the assistant recovery turn, and builds an
`sdpo_loss_mask` that's 1 only at the post-hint assistant tokens
(positions where student and teacher are predicting the SAME content).

This example demonstrates:
  ✅ The full production data path: ingester → adapter → collator
  ✅ SDPO column firing on PROPERLY-ALIGNED student/teacher contexts
  ✅ Real tool error detection via the [TOOL_RESULT (ERROR)] tag flow
  ✅ A deterministic hint generator wired into CollatorConfig
  ✅ Gradient flow through Qwen2.5-0.5B-Instruct's params

Closes the V5 gap end-to-end (the path is production-grade and
content-honest, with a detailed hint at the actual error site of the
trace), within the constraint that the trace fixture is hand-authored
(PII reasons; users can point at their own JSONL).

Usage:
    pip install -e ".[train]"
    python examples/sdpo_with_real_traces_production/run.py

Cross-references:
  - composer_replication.ingestion.trace_examples.claude_states_to_trace_examples
  - composer_replication.trainer.data_collator.ComposerDataCollator
  - composer_replication.trainer.data_collator._build_hint_injected_trace
  - examples/sdpo_with_real_traces/ (the wiring-only sibling for comparison)
"""
from __future__ import annotations

import logging
import math
import sys
import time
from pathlib import Path

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

from composer_replication import compose_loss
from composer_replication.ingestion import (
    ClaudeCodeIngester,
    claude_states_to_trace_examples,
)
from composer_replication.trainer.data_collator import (
    CollatorConfig,
    ComposerDataCollator,
)

# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------

MODEL_REPO = "Qwen/Qwen2.5-0.5B-Instruct"
N_STEPS = 5
LR = 1e-5
ALPHA_SDPO = 0.5
BETA_REPLAY = 0.0
MAX_SEQ_LEN = 1024  # generous; the with-error fixture is short

OUTPUT_DIR = Path(__file__).resolve().parent / "output"
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

# This fixture is the WITH-ERROR variant — it has an `is_error: true`
# tool_result that the adapter detects and the collator injects a hint
# before. The clean Spike 007 fixture has no errors and would produce
# a no-op SDPO batch.
FIXTURE_PATH = (
    Path(__file__).resolve().parents[2]
    / "spikes" / "007-real-trace-ingestion" / "fixtures" / "synthetic_session_with_error.jsonl"
)


# ---------------------------------------------------------------------------
# Hint generator — deterministic, error-kind-aware
# ---------------------------------------------------------------------------


def hint_for_error(error_kind: str, error_meta: dict) -> str | None:
    """Return a hint text given the classified error kind.

    A real production hint generator would pull from a curated hint
    library or an LLM-as-teacher; here we use a small static map for
    determinism. Returning None for an error kind tells the collator
    to skip the SDPO injection for that turn.
    """
    library = {
        "file_not_found": (
            "Hint: when reading a file fails with 'does not exist', "
            "first verify the path with `ls` on the parent directory "
            "or use a glob to find similar names before retrying."
        ),
        "permission_denied": (
            "Hint: when 'permission denied', check ownership with `ls -l` "
            "before retrying. Don't blindly add `sudo`; read the situation."
        ),
        "command_not_found": (
            "Hint: when a command isn't found, check `which <command>` "
            "and `echo $PATH`; the binary may need to be installed first."
        ),
        "tool_error": (
            "Hint: this tool call failed. Read the error carefully and "
            "consider whether to retry, change inputs, or pivot to a "
            "different tool before continuing."
        ),
    }
    return library.get(error_kind, library["tool_error"])


# ---------------------------------------------------------------------------
# Build batch via production path
# ---------------------------------------------------------------------------


def build_production_batch(
    tokenizer, fixture_path: Path,
) -> tuple[dict[str, torch.Tensor], int, int]:
    """Run the full production pipeline.

    Returns:
        (batch, n_states, n_error_sites)
    """
    ingester = ClaudeCodeIngester(skip_sidechain=True, strip_thinking=True)
    states = list(ingester.ingest(fixture_path))
    if not states:
        raise RuntimeError(f"No TraceState yielded from {fixture_path}")

    examples = claude_states_to_trace_examples(states)
    n_error_sites = sum(
        1 for ex in examples for t in ex["turns"] if t.get("tool_error")
    )

    config = CollatorConfig(
        hint_generator=hint_for_error,
        enable_replay_dpo=False,  # this example focuses on SDPO
        pad_token_id=tokenizer.pad_token_id or 0,
        max_seq_len=MAX_SEQ_LEN,
    )
    collator = ComposerDataCollator(tokenizer=tokenizer, config=config)
    batch = collator(examples)
    return batch, len(states), n_error_sites


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------


def main() -> int:
    torch.manual_seed(42)

    log_path = OUTPUT_DIR.parent / "run.log"
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
        handlers=[logging.StreamHandler(sys.stdout), logging.FileHandler(log_path, mode="w")],
    )
    log = logging.getLogger("sdpo_production")

    log.info("=" * 64)
    log.info("PRODUCTION-GRADE SDPO + ClaudeCodeIngester + ComposerDataCollator")
    log.info("Model: %s (CPU)", MODEL_REPO)
    log.info("=" * 64)

    if not FIXTURE_PATH.is_file():
        log.error("Fixture not found at %s", FIXTURE_PATH)
        return 1
    log.info("[1/5] Fixture: %s (size=%d bytes)",
             FIXTURE_PATH.name, FIXTURE_PATH.stat().st_size)

    log.info("[2/5] Loading model + tokenizer ...")
    t0 = time.time()
    tokenizer = AutoTokenizer.from_pretrained(MODEL_REPO)
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token = tokenizer.eos_token
    model = AutoModelForCausalLM.from_pretrained(MODEL_REPO, torch_dtype=torch.float32)
    model.to("cpu")
    n_params = sum(p.numel() for p in model.parameters())
    log.info("  loaded in %.1fs (%.3fB params)", time.time() - t0, n_params / 1e9)

    log.info("[3/5] Building batch via production pipeline ...")
    log.info("  ClaudeCodeIngester → claude_states_to_trace_examples → ComposerDataCollator")
    batch, n_states, n_error_sites = build_production_batch(tokenizer, FIXTURE_PATH)
    log.info("  ingested %d states; adapter detected %d error site(s)",
             n_states, n_error_sites)
    if n_error_sites == 0:
        log.error("  No error sites detected — SDPO will be a no-op. "
                  "Use the with-error fixture or extend the adapter.")
        return 1
    for k, v in batch.items():
        log.info("    %s: shape=%s dtype=%s", k, tuple(v.shape), v.dtype)
    if "ctx_teacher_input_ids" not in batch:
        log.error("  Collator did not produce ctx_teacher_input_ids — "
                  "no error sites survived hint generator. Aborting.")
        return 1
    sdpo_in_loss = (batch["sdpo_loss_mask"] == 1).sum().item()
    log.info("  sdpo_loss_mask: %d positions in loss (per-row: %s)",
             sdpo_in_loss, (batch["sdpo_loss_mask"] == 1).sum(dim=-1).tolist())

    s_shape = batch["input_ids"].shape
    t_shape = batch["ctx_teacher_input_ids"].shape
    log.info("  shape reconciliation: student %s vs teacher %s — %s",
             tuple(s_shape), tuple(t_shape),
             "ALIGNED" if s_shape == t_shape else "MISMATCH (collator bug?)")
    assert s_shape == t_shape, (
        f"Shape mismatch after collator: student {s_shape} vs teacher {t_shape}. "
        f"compose_loss requires student_logits.shape == teacher_logits.shape; "
        f"the collator's __call__ must reconcile them."
    )

    log.info("[4/5] Running %d SGD steps with alpha_sdpo=%.2f ...", N_STEPS, ALPHA_SDPO)
    optim = torch.optim.SGD(model.parameters(), lr=LR)
    history: list[dict[str, float]] = []

    model.train()
    t0 = time.time()
    for step in range(N_STEPS):
        optim.zero_grad()
        out = compose_loss(
            model, batch,
            alpha_sdpo=ALPHA_SDPO, beta_replay=BETA_REPLAY,
        )
        out.total.backward()
        gnorm = sum(
            p.grad.abs().sum().item() for p in model.parameters() if p.grad is not None
        )
        optim.step()

        components = out.detached()
        components["grad_norm"] = gnorm
        history.append(components)
        log.info(
            "  step %d/%d: total=%.4f  lm_ce=%.4f  sdpo_jsd=%.4f  trace_replay_dpo=%.4f  |grad|=%.2e",
            step + 1, N_STEPS,
            components["total"], components["lm_ce"],
            components["sdpo_jsd"], components["trace_replay_dpo"],
            gnorm,
        )
    dt = time.time() - t0
    log.info("Training complete in %.1fs (avg %.1fs/step)", dt, dt / N_STEPS)

    log.info("[5/5] Verifying production-grade SDPO behavior ...")
    sdpo_values = [h["sdpo_jsd"] for h in history]

    # Production-grade SDPO MUST produce a non-zero JSD signal because
    # the collator put the hint in a position where it actually changes
    # the teacher's prediction at the masked positions.
    assert all(abs(s) > 1e-7 for s in sdpo_values), (
        f"Production-grade SDPO column produced negligible JSD: {sdpo_values}. "
        f"The hint isn't perturbing teacher logits at masked positions — "
        f"check the collator's hint injection or the loss mask."
    )
    log.info("  ✓ sdpo_jsd > 1e-7 at every step (min=%.6f max=%.6f)",
             min(sdpo_values), max(sdpo_values))

    # The composed total must differ from lm_ce alone — confirms SDPO contributes
    diffs = [abs(h["total"] - h["lm_ce"]) for h in history]
    assert all(d > 1e-6 for d in diffs), (
        f"total ≈ lm_ce — SDPO contribution negligible. abs(total-lm_ce)={diffs}"
    )
    log.info("  ✓ total != lm_ce at every step (min |diff|=%.4f)", min(diffs))

    gnorms = [h["grad_norm"] for h in history]
    assert all(g > 0.0 and math.isfinite(g) for g in gnorms), (
        f"Some grads non-finite or zero: {gnorms}"
    )
    log.info("  ✓ |grad| finite at every step (min=%.2e max=%.2e)",
             min(gnorms), max(gnorms))

    # ----------------------------------------------------------------
    # Alignment audit (Wave 19 honesty: documents the residual drift)
    # ----------------------------------------------------------------
    s_in = batch["input_ids"]
    t_in = batch["ctx_teacher_input_ids"]
    m_in = batch["sdpo_loss_mask"]
    n_aligned = 0
    n_total_in_loss = 0
    for row in range(s_in.shape[0]):
        in_loss = (m_in[row] == 1)
        n_pos = in_loss.sum().item()
        if n_pos == 0:
            continue
        s_at = s_in[row][in_loss]
        t_at = t_in[row][in_loss]
        n_aligned += int((s_at == t_at).sum().item())
        n_total_in_loss += n_pos
    if n_total_in_loss:
        ratio = n_aligned / n_total_in_loss
        log.info("  alignment audit: %d / %d in-loss positions match student==teacher (%.1f%%)",
                 n_aligned, n_total_in_loss, 100 * ratio)
        if ratio < 0.95:
            log.warning(
                "  NOTE: %d positions (%.1f%%) of the SDPO mask cover non-aligned "
                "tokens. As of Wave 20 the chat-template drift was fixed via "
                "ComposerDataCollator._build_chat_aligned_mask (per-message "
                "apply_chat_template prefix deltas). A ratio below ~100%% now "
                "indicates a NEW regression — investigate the collator, not a "
                "known-residual bug.",
                n_total_in_loss - n_aligned,
                100 * (1 - ratio),
            )
        else:
            log.info(
                "  ✓ Wave 20 chat-template alignment holding (%.1f%% — was ~67%% "
                "before the _build_chat_aligned_mask fix).", 100 * ratio,
            )

    log.info("=" * 64)
    log.info("Summary")
    log.info("=" * 64)
    log.info("  trace fixture:   %s", FIXTURE_PATH.name)
    log.info("  states:          %d", n_states)
    log.info("  error sites:     %d", n_error_sites)
    log.info("  sdpo_loss_mask:  %d positions in loss", sdpo_in_loss)
    log.info("  alpha_sdpo:      %.2f", ALPHA_SDPO)
    log.info("  total step 1:    %.4f", history[0]["total"])
    log.info("  total step %d:    %.4f", N_STEPS, history[-1]["total"])
    log.info("  wall-clock:      %.1fs", dt)
    log.info("=" * 64)
    log.info("✅ Production-grade SDPO verified end-to-end via ComposerDataCollator.")
    return 0


if __name__ == "__main__":
    sys.exit(main())