From 52d50d67fe5db71d67925278f408d2faa6f52f18 Mon Sep 17 00:00:00 2001 From: claude-bot-go Date: Tue, 28 Jul 2026 18:22:22 +0800 Subject: [PATCH] =?UTF-8?q?Fix=20#54:=20[milestone=20Milestone=205=20]=20D?= =?UTF-8?q?PO/PPO/SFT=20training=20pipeline=20integration=20=E2=80=94=20wi?= =?UTF-8?q?re=20the=20validated=20trace=20exporters...?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- evomerge/training_pipeline.py | 545 ++++++++++++++++++++++++++++++++ evomerge/trust_score.py | 39 +++ scripts/train_dpo.py | 36 +++ scripts/train_sft.py | 34 ++ tests/test_training_pipeline.py | 323 +++++++++++++++++++ 5 files changed, 977 insertions(+) create mode 100644 evomerge/training_pipeline.py create mode 100644 tests/test_training_pipeline.py diff --git a/evomerge/training_pipeline.py b/evomerge/training_pipeline.py new file mode 100644 index 0000000..7f79449 --- /dev/null +++ b/evomerge/training_pipeline.py @@ -0,0 +1,545 @@ +"""Training-pipeline integration — wire trace exporters to training-job orchestration. + +This module is the integration layer between the validated trace exporters +(``evomerge.pipeline.to_sft_records`` / ``to_dpo_records`` / ``to_ppo_records``) +and an external training-job orchestrator (Ray / Lightning / the local TRL shims +in ``scripts/train_*.py``). + +Per the repo boundary (CLAUDE.md), trace-pipeline produces training *data* and +thin orchestration *contracts* — it is **not** a training framework. This +module therefore provides four data-side primitives that together "wire the +validated trace exporters to actual training job orchestration": + + 1. Dataset versioning — content-addressed dataset manifests (``version_dataset``) + 2. Checkpoint management — append-only lineage registry (``CheckpointRegistry``) + 3. Loss telemetry — training-loss → trust-score feedback (``LossTelemetry``) + 4. Job manifest — wires exporters → versioned dataset → checkpoint plan + → telemetry sink → orchestrator (``build_training_job_manifest``) + +The actual Ray/Lightning cluster execution lives outside this repo; the +``TrainingJobManifest`` is the contract an external runner consumes. +""" +from __future__ import annotations + +import hashlib +import json +import time +from collections.abc import Iterable, Sequence +from dataclasses import asdict, dataclass, field +from pathlib import Path +from typing import Any + +from evomerge.io import write_dicts_jsonl + +TRAINING_MODES: tuple[str, ...] = ("sft", "dpo", "ppo") +ORCHESTRATORS: tuple[str, ...] = ("local", "ray", "lightning") + +#: Default convergence threshold for ``LossTelemetry`` (mirrors train_sft.py). +_DEFAULT_CONVERGED_BELOW = 0.1 + + +# =========================================================================== +# Helpers +# =========================================================================== + +def _record_to_dict(rec: Any) -> dict: + """Normalise a training record (pydantic model or plain dict) to a dict.""" + if isinstance(rec, dict): + return rec + if hasattr(rec, "model_dump"): # pydantic v2 + return rec.model_dump(mode="json") + if hasattr(rec, "dict"): # pydantic v1 fallback + return rec.dict() + return dict(rec) + + +def _canonical_json(obj: Any) -> str: + return json.dumps(obj, sort_keys=True, separators=(",", ":"), ensure_ascii=False) + + +def _short_version(digest: str) -> str: + return "v-" + digest[:6] + + +def _now_ms() -> int: + return int(time.time() * 1000) + + +# =========================================================================== +# 1. Dataset versioning +# =========================================================================== + +def compute_dataset_digest(records: Iterable[Any]) -> str: + """SHA-256 over the canonical serialisation of every record. + + Deterministic: identical record content yields an identical digest + regardless of dict key ordering, so re-versioning unchanged data produces + the same version tag (idempotent). + """ + h = hashlib.sha256() + for rec in records: + h.update(_canonical_json(_record_to_dict(rec)).encode("utf-8")) + h.update(b"\n") + return h.hexdigest() + + +@dataclass +class DatasetVersion: + """Content-addressed version of a training-record dataset.""" + + name: str # 'sft' | 'dpo' | 'ppo' + schema_version: str # e.g. 'sft/v1' + content_digest: str # full sha256 + version: str # short tag, e.g. 'v-a1b2c3' + record_count: int + path: str # versioned jsonl path + manifest_path: str # sidecar manifest json path + created_at_ms: int + sources: list[str] = field(default_factory=list) + + def to_dict(self) -> dict: + return asdict(self) + + +def version_dataset( + records: Sequence[Any], + *, + name: str, + out_dir: str | Path, + sources: list[str] | None = None, +) -> DatasetVersion: + """Write a content-addressed, versioned copy of ``records`` plus a manifest. + + The output filename embeds the short content digest, so the same record + content always maps to the same versioned path (idempotent). A sidecar + ``-.manifest.json`` records the schema version, record count, + full digest, source lineage, and creation time. + + Args: + records: training records (pydantic models or plain dicts). + name: dataset name — one of ``TRAINING_MODES``. + out_dir: directory to write into (created if absent). + sources: optional lineage strings (e.g. source rollout JSONL paths). + """ + if name not in TRAINING_MODES: + raise ValueError(f"unknown training mode {name!r}; expected one of {TRAINING_MODES}") + + out = Path(out_dir) + out.mkdir(parents=True, exist_ok=True) + + dicts = [_record_to_dict(r) for r in records] + digest = compute_dataset_digest(dicts) + version = _short_version(digest) + schema_version = (dicts[0].get("schema_version") if dicts else None) or f"{name}/v1" + + data_path = out / f"{name}-{version}.jsonl" + manifest_path = out / f"{name}-{version}.manifest.json" + + write_dicts_jsonl(dicts, data_path) + dv = DatasetVersion( + name=name, + schema_version=schema_version, + content_digest=digest, + version=version, + record_count=len(dicts), + path=str(data_path), + manifest_path=str(manifest_path), + created_at_ms=_now_ms(), + sources=list(sources or []), + ) + manifest_path.write_text(json.dumps(dv.to_dict(), indent=2, ensure_ascii=False)) + return dv + + +def load_dataset_manifest(path: str | Path) -> DatasetVersion: + """Load a ``DatasetVersion`` previously written by :func:`version_dataset`.""" + data = json.loads(Path(path).read_text(encoding="utf-8")) + return DatasetVersion(**data) + + +# =========================================================================== +# 2. Checkpoint management +# =========================================================================== + +@dataclass +class CheckpointEntry: + """One checkpoint in the lineage registry.""" + + checkpoint_id: str # stable id, e.g. '-' + mode: str # 'sft' | 'dpo' | 'ppo' + dataset_version: str # content digest of the training dataset + base_ref: str # base model name OR parent checkpoint_id + status: str = "ready" # 'training' | 'ready' | 'failed' + path: str = "" + metrics: dict[str, Any] = field(default_factory=dict) + created_at_ms: int = 0 + + def to_dict(self) -> dict: + return asdict(self) + + +class CheckpointRegistry: + """Append-only JSONL registry of training checkpoints with lineage lookup. + + The registry is a simple append-only log; each ``checkpoint_id`` is recorded + at most once. Lineage is reconstructed by walking ``base_ref`` pointers back + from a checkpoint to its base model. + """ + + def __init__(self, path: str | Path) -> None: + self.path = Path(path) + self.path.parent.mkdir(parents=True, exist_ok=True) + self.path.touch(exist_ok=True) + + def _read(self) -> list[CheckpointEntry]: + entries: list[CheckpointEntry] = [] + for line in self.path.read_text(encoding="utf-8").splitlines(): + line = line.strip() + if line: + entries.append(CheckpointEntry(**json.loads(line))) + return entries + + def all(self) -> list[CheckpointEntry]: + return self._read() + + def find(self, checkpoint_id: str) -> CheckpointEntry | None: + for entry in self._read(): + if entry.checkpoint_id == checkpoint_id: + return entry + return None + + def latest(self, mode: str | None = None) -> CheckpointEntry | None: + entries = [e for e in self._read() if mode is None or e.mode == mode] + if not entries: + return None + return max(entries, key=lambda e: (e.created_at_ms, e.checkpoint_id)) + + def register(self, entry: CheckpointEntry) -> CheckpointEntry: + """Append ``entry`` unless its checkpoint_id is already registered.""" + if entry.mode not in TRAINING_MODES: + raise ValueError(f"unknown mode {entry.mode!r}; expected one of {TRAINING_MODES}") + if not entry.created_at_ms: + entry.created_at_ms = _now_ms() + if self.find(entry.checkpoint_id) is None: + with self.path.open("a", encoding="utf-8") as fh: + fh.write(json.dumps(entry.to_dict(), ensure_ascii=False) + "\n") + return entry + + def lineage(self, checkpoint_id: str) -> list[CheckpointEntry]: + """Return the chain ``[ckpt, parent, ...]`` following ``base_ref`` pointers. + + Stops when ``base_ref`` no longer names a registered checkpoint (i.e. it + is a base model name). Cycle-safe. + """ + by_id = {e.checkpoint_id: e for e in self._read()} + chain: list[CheckpointEntry] = [] + seen: set[str] = set() + cur = by_id.get(checkpoint_id) + while cur is not None and cur.checkpoint_id not in seen: + seen.add(cur.checkpoint_id) + chain.append(cur) + cur = by_id.get(cur.base_ref) + return chain + + +# =========================================================================== +# 3. Training-loss telemetry → trust +# =========================================================================== + +@dataclass +class LossTelemetry: + """Training-loss telemetry extracted from a training run.""" + + mode: str + steps: int = 0 + initial_loss: float | None = None + final_loss: float | None = None + loss_history: list[float] = field(default_factory=list) + converged: bool = False + converged_below: float = _DEFAULT_CONVERGED_BELOW + + def to_dict(self) -> dict: + return asdict(self) + + +def _extract_losses_from_log(lines: Iterable[str]) -> list[float]: + """Pull ``loss`` values from HF Trainer-style log lines. + + Handles two shapes: + - JSON-object lines: ``{"loss": 0.42, "step": 10}`` + - key=value lines: ``'loss' = 0.42`` / ``loss=0.42`` / ``loss: 0.42`` + """ + losses: list[float] = [] + for line in lines: + s = line.strip() + if not s: + continue + loss: float | None = None + if s.startswith("{"): + try: + obj = json.loads(s) + if isinstance(obj, dict) and obj.get("loss") is not None: + loss = float(obj["loss"]) + except (json.JSONDecodeError, TypeError, ValueError): + loss = None + if loss is None: + low = s.lower() + idx = low.find("loss") + if idx != -1: + rest = s[idx + 4:].lstrip("_=: \t'\"") + num = "" + for ch in rest: + if ch in "-.0123456789eE": + num += ch + else: + break + if num and num not in ("-", "."): + try: + loss = float(num) + except ValueError: + loss = None + if loss is not None: + losses.append(loss) + return losses + + +def telemetry_from_loss_history( + loss_history: Sequence[float], + *, + mode: str, + steps: int | None = None, + converged_below: float = _DEFAULT_CONVERGED_BELOW, +) -> LossTelemetry: + """Build telemetry from an explicit loss history.""" + hist = [float(x) for x in loss_history] + final = hist[-1] if hist else None + initial = hist[0] if hist else None + return LossTelemetry( + mode=mode, + steps=steps if steps is not None else len(hist), + initial_loss=initial, + final_loss=final, + loss_history=hist, + converged=bool(hist) and final is not None and final < converged_below, + converged_below=converged_below, + ) + + +def parse_trainer_log( + lines: Iterable[str], + *, + mode: str, + converged_below: float = _DEFAULT_CONVERGED_BELOW, +) -> LossTelemetry: + """Build telemetry from HF Trainer log lines (see ``_extract_losses_from_log``).""" + return telemetry_from_loss_history( + _extract_losses_from_log(lines), mode=mode, converged_below=converged_below + ) + + +def telemetry_from_summary( + summary: dict, + *, + mode: str, + converged_below: float = _DEFAULT_CONVERGED_BELOW, +) -> LossTelemetry: + """Build telemetry from a ``training_summary.json``-style dict. + + Recognises ``final_loss`` / ``loss`` / ``train_loss`` and optional + ``loss_history`` / ``initial_loss`` / ``steps`` / ``max_steps``. + """ + hist = [float(x) for x in (summary.get("loss_history") or [])] + final = None + for key in ("final_loss", "loss", "train_loss"): + v = summary.get(key) + if v is not None: + try: + final = float(v) + except (TypeError, ValueError): + pass + break + initial = summary.get("initial_loss") + if initial is not None: + try: + initial = float(initial) + except (TypeError, ValueError): + initial = None + if not hist: + if final is not None and initial is not None: + hist = [initial, final] + elif final is not None: + hist = [final] + steps_raw = summary.get("steps") or summary.get("max_steps") + try: + steps = int(steps_raw) if steps_raw is not None else (len(hist) or 0) + except (TypeError, ValueError): + steps = len(hist) or 0 + return LossTelemetry( + mode=mode, + steps=steps, + initial_loss=initial, + final_loss=final, + loss_history=hist, + converged=final is not None and final < converged_below, + converged_below=converged_below, + ) + + +def telemetry_to_training_health(telemetry: LossTelemetry) -> float: + """Map a :class:`LossTelemetry` to a ``[0, 1]`` training-health score. + + Combines three signals: + - convergence (final_loss at/below the threshold → 1.0; decays to 0 at 2x) + - improvement (relative drop from initial → final; 0 if it got worse) + - stability (low relative variance in the tail of the history) + + Returns 0.0 when there is no loss evidence at all. + """ + if not telemetry.loss_history or telemetry.final_loss is None: + return 0.0 + + final = telemetry.final_loss + thr = telemetry.converged_below if telemetry.converged_below > 0 else _DEFAULT_CONVERGED_BELOW + if final <= thr: + convergence = 1.0 + elif final >= 2 * thr: + convergence = 0.0 + else: + convergence = (2 * thr - final) / thr + + if telemetry.initial_loss and telemetry.initial_loss > 0: + improvement = max(0.0, min(1.0, (telemetry.initial_loss - final) / telemetry.initial_loss)) + else: + improvement = 0.5 # no initial → neutral + + tail = telemetry.loss_history[max(1, len(telemetry.loss_history) // 2):] + if len(tail) >= 2 and final > 0: + mean = sum(tail) / len(tail) + var = sum((x - mean) ** 2 for x in tail) / len(tail) + stability = max(0.0, 1.0 - (var ** 0.5) / max(mean, 1e-9)) + else: + stability = 0.8 + + return max(0.0, min(1.0, convergence * 0.5 + improvement * 0.3 + stability * 0.2)) + + +# =========================================================================== +# 4. Job manifest — wire exporters → orchestration +# =========================================================================== + +@dataclass +class TrainingJobManifest: + """Contract consumed by an external training orchestrator (Ray/Lightning/local). + + trace-pipeline produces the manifest; the actual cluster execution is out of + scope (CLAUDE.md: not a training framework). The manifest pins the versioned + dataset, checkpoint plan, and telemetry sink so a run is reproducible and its + loss telemetry can flow back into trust scores. + """ + + job_id: str + mode: str + orchestrator: str + dataset: dict # DatasetVersion.to_dict() + base_ref: str + checkpoint_plan: dict # {save_steps, save_total_limit, resume_from} + telemetry_sink: str # path where LossTelemetry json is written + hyperparameters: dict + created_at_ms: int + + def to_dict(self) -> dict: + return asdict(self) + + +def records_for_mode(rollouts: Sequence[Any], mode: str) -> list[Any]: + """Dispatch to the validated trace exporter for ``mode``. + + This is the explicit wire from the ``evomerge.pipeline`` converters into the + orchestration manifest. ``mode='dpo'`` preserves the converter's default + attestation gate (branches without verifier_results are implicitly trusted). + """ + if mode == "sft": + from evomerge.pipeline import to_sft_records + return to_sft_records(rollouts) + if mode == "dpo": + from evomerge.pipeline import to_dpo_records + return to_dpo_records(rollouts) + if mode == "ppo": + from evomerge.pipeline import to_ppo_records + return to_ppo_records(rollouts) + raise ValueError(f"unknown mode {mode!r}; expected one of {TRAINING_MODES}") + + +def build_training_job_manifest( + *, + records: Sequence[Any], + mode: str, + out_dir: str | Path, + base_ref: str = "", + orchestrator: str = "local", + hyperparameters: dict | None = None, + save_steps: int = 50, + save_total_limit: int = 3, + resume_from: str | None = None, + sources: list[str] | None = None, +) -> tuple[TrainingJobManifest, DatasetVersion]: + """Version a dataset and emit a training-job manifest for an orchestrator. + + Writes: + - ``/-.jsonl`` (versioned dataset) + - ``/-.manifest.json`` (dataset manifest) + - ``/training-job.manifest.json`` (this job manifest) + + Returns ``(manifest, dataset_version)``. + """ + if mode not in TRAINING_MODES: + raise ValueError(f"unknown mode {mode!r}; expected one of {TRAINING_MODES}") + if orchestrator not in ORCHESTRATORS: + raise ValueError( + f"unknown orchestrator {orchestrator!r}; expected one of {ORCHESTRATORS}" + ) + + out = Path(out_dir) + out.mkdir(parents=True, exist_ok=True) + + dataset_version = version_dataset(records, name=mode, out_dir=out, sources=sources) + + manifest = TrainingJobManifest( + job_id=f"{mode}-{dataset_version.version}-{_now_ms()}", + mode=mode, + orchestrator=orchestrator, + dataset=dataset_version.to_dict(), + base_ref=base_ref, + checkpoint_plan={ + "save_steps": save_steps, + "save_total_limit": save_total_limit, + "resume_from": resume_from, + }, + telemetry_sink=str(out / f"{mode}-{dataset_version.version}.telemetry.json"), + hyperparameters=dict(hyperparameters or {}), + created_at_ms=_now_ms(), + ) + (out / "training-job.manifest.json").write_text( + json.dumps(manifest.to_dict(), indent=2, ensure_ascii=False) + ) + return manifest, dataset_version + + +__all__ = [ + "CheckpointEntry", + "CheckpointRegistry", + "DatasetVersion", + "LossTelemetry", + "ORCHESTRATORS", + "TRAINING_MODES", + "TrainingJobManifest", + "build_training_job_manifest", + "compute_dataset_digest", + "load_dataset_manifest", + "parse_trainer_log", + "records_for_mode", + "telemetry_from_loss_history", + "telemetry_from_summary", + "telemetry_to_training_health", + "version_dataset", +] diff --git a/evomerge/trust_score.py b/evomerge/trust_score.py index 67840b4..51cead2 100644 --- a/evomerge/trust_score.py +++ b/evomerge/trust_score.py @@ -339,6 +339,45 @@ def add_dimension(self, name: str, score: float) -> AgentTrustScoreBuilder: self._dims[name] = max(0.0, min(1.0, score)) return self + def add_training_telemetry(self, telemetry: Any) -> AgentTrustScoreBuilder: + """Fold training-loss telemetry back into the score as ``training_health``. + + Closes the trace→training→trust feedback loop (Milestone 5): the loss + telemetry from a training job is mapped to a [0, 1] health dimension and + enters the geometric mean, so a model that converged well lifts the trust + score and a divergent/untrained run does not. + + Accepts: + - a plain ``int``/``float`` in [0, 1] (used directly as the health score), or + - a ``evomerge.training_pipeline.LossTelemetry`` (or any duck-typed object + exposing ``loss_history`` / ``final_loss`` / ``initial_loss`` / + ``converged_below``), scored via ``telemetry_to_training_health``. + """ + if isinstance(telemetry, bool): + # bool is an int subclass — treat as a 0/1 health score explicitly. + health = 1.0 if telemetry else 0.0 + elif isinstance(telemetry, (int, float)): + health = float(telemetry) + else: + # Lazy import avoids a hard dependency / import cycle. + from evomerge.training_pipeline import LossTelemetry, telemetry_to_training_health + + if isinstance(telemetry, LossTelemetry): + health = telemetry_to_training_health(telemetry) + else: # duck-typed object with the same attributes + proxy = LossTelemetry( + mode=getattr(telemetry, "mode", "sft"), + steps=getattr(telemetry, "steps", 0), + initial_loss=getattr(telemetry, "initial_loss", None), + final_loss=getattr(telemetry, "final_loss", None), + loss_history=list(getattr(telemetry, "loss_history", []) or []), + converged=bool(getattr(telemetry, "converged", False)), + converged_below=float(getattr(telemetry, "converged_below", 0.1) or 0.1), + ) + health = telemetry_to_training_health(proxy) + self._dims["training_health"] = max(0.0, min(1.0, health)) + return self + def build(self) -> AgentTrustScore: # Only include non-None dimensions in the geometric mean known_values = [v for v in self._dims.values() if v is not None] diff --git a/scripts/train_dpo.py b/scripts/train_dpo.py index 7e0d666..b5dc579 100644 --- a/scripts/train_dpo.py +++ b/scripts/train_dpo.py @@ -31,6 +31,16 @@ REPO_ROOT = Path(__file__).parent.parent sys.path.insert(0, str(REPO_ROOT)) +# Training-pipeline integration (issue #54): dataset versioning, checkpoint +# registry, and loss-telemetry feedback. Torch-agnostic, so importing here does +# not affect the torch-optional flow below. +from evomerge.training_pipeline import ( # noqa: E402 + CheckpointEntry, + CheckpointRegistry, + telemetry_from_summary, + version_dataset, +) + # Default DPO data sources relative to REPO_ROOT _DEFAULT_DPO_SOURCES = [ "data/training/ifeval/compliance_dpo.jsonl", @@ -195,6 +205,13 @@ def main() -> int: out = Path(args.out_dir) out.mkdir(parents=True, exist_ok=True) + # Version the training dataset (content-addressed) so this run is reproducible. + # `paths` is only defined for --dpo-data; --merge-all has no explicit source list. + _sources = None if args.merge_all else paths + dataset_version = version_dataset(records, name="dpo", out_dir=out, sources=_sources) + print(f" dataset version : {dataset_version.version} " + f"({dataset_version.record_count} records)") + training_args = DPOConfig( output_dir=str(out), num_train_epochs=args.epochs if args.max_steps <= 0 else 1, @@ -232,6 +249,10 @@ def main() -> int: trainer.save_model(str(out / "final")) tokenizer.save_pretrained(str(out / "final")) + # Capture training-loss history for telemetry (issue #54). + log_history = getattr(getattr(trainer, "state", None), "log_history", []) or [] + loss_history = [float(e["loss"]) for e in log_history if isinstance(e, dict) and "loss" in e] + summary = { "sft_checkpoint": str(sft_ckpt), "n_dpo_pairs": len(records), @@ -239,8 +260,23 @@ def main() -> int: "num_train_epochs": args.epochs if args.max_steps <= 0 else None, "max_steps": args.max_steps if args.max_steps > 0 else None, "lr": args.lr, + "loss_history": loss_history, + "final_loss": loss_history[-1] if loss_history else None, } (out / "training_summary.json").write_text(json.dumps(summary, indent=2)) + + # Emit loss telemetry and register the checkpoint in the lineage registry. + telemetry = telemetry_from_summary(summary, mode="dpo") + (out / "training_telemetry.json").write_text(json.dumps(telemetry.to_dict(), indent=2)) + CheckpointRegistry(out / "checkpoints.jsonl").register(CheckpointEntry( + checkpoint_id=f"dpo-{dataset_version.version}", + mode="dpo", + dataset_version=dataset_version.content_digest, + base_ref=str(sft_ckpt), + status="ready", + path=str(out / "final"), + metrics={"final_loss": telemetry.final_loss, "converged": telemetry.converged}, + )) print(f"\n✓ DPO training complete → {out}/final") return 0 diff --git a/scripts/train_sft.py b/scripts/train_sft.py index 30f201a..0d90644 100644 --- a/scripts/train_sft.py +++ b/scripts/train_sft.py @@ -41,6 +41,16 @@ REPO_ROOT = Path(__file__).parent.parent +# Training-pipeline integration (issue #54): dataset versioning, checkpoint +# registry, and loss-telemetry feedback. These helpers are torch-agnostic, so +# importing them here does not affect the torch-optional flow below. +from evomerge.training_pipeline import ( # noqa: E402 + CheckpointEntry, + CheckpointRegistry, + telemetry_from_summary, + version_dataset, +) + def _load_sft_records(paths: list[str]) -> list[dict]: """Load and merge multiple SFT JSONL files.""" @@ -219,6 +229,11 @@ def format_fn(example): out = Path(args.out_dir) out.mkdir(parents=True, exist_ok=True) + # Version the training dataset (content-addressed) so this run is reproducible. + dataset_version = version_dataset(records, name="sft", out_dir=out, sources=paths) + print(f" dataset version : {dataset_version.version} " + f"({dataset_version.record_count} records)") + # EarlyStoppingCallback based on loss threshold from transformers import TrainerCallback @@ -279,6 +294,10 @@ def on_log(self, args, state, control, logs=None, **kwargs): trainer.save_model(str(out / "final")) tokenizer.save_pretrained(str(out / "final")) + # Capture training-loss history for telemetry (issue #54). + log_history = getattr(getattr(trainer, "state", None), "log_history", []) or [] + loss_history = [float(e["loss"]) for e in log_history if isinstance(e, dict) and "loss" in e] + # save training summary summary = { "base_model": args.base_model, @@ -289,8 +308,23 @@ def on_log(self, args, state, control, logs=None, **kwargs): "dtype": dtype_str, "max_steps": args.max_steps, "lr": args.lr, + "loss_history": loss_history, + "final_loss": loss_history[-1] if loss_history else None, } (out / "training_summary.json").write_text(json.dumps(summary, indent=2)) + + # Emit loss telemetry and register the checkpoint in the lineage registry. + telemetry = telemetry_from_summary(summary, mode="sft") + (out / "training_telemetry.json").write_text(json.dumps(telemetry.to_dict(), indent=2)) + CheckpointRegistry(out / "checkpoints.jsonl").register(CheckpointEntry( + checkpoint_id=f"sft-{dataset_version.version}", + mode="sft", + dataset_version=dataset_version.content_digest, + base_ref=args.base_model, + status="ready", + path=str(out / "final"), + metrics={"final_loss": telemetry.final_loss, "converged": telemetry.converged}, + )) print(f"\n✓ training complete → {out}/final") return 0 diff --git a/tests/test_training_pipeline.py b/tests/test_training_pipeline.py new file mode 100644 index 0000000..c84880a --- /dev/null +++ b/tests/test_training_pipeline.py @@ -0,0 +1,323 @@ +"""Tests for evomerge.training_pipeline — Milestone 5 training-pipeline integration. + +Covers the four sub-requirements of issue #54 (DPO/PPO/SFT training pipeline +integration) each in its own test class: + + - TestDatasetVersioning — dataset versioning (content-addressed) + - TestCheckpointRegistry — checkpoint management (lineage) + - TestLossTelemetry — training-loss telemetry → trust-score feedback + - TestTrainingJobManifest — wire exporters → orchestration manifest +""" +from __future__ import annotations + +import json + +import pytest + +from evomerge.schemas.rollout import RolloutBranchRecord +from evomerge.schemas.training import Message, Provenance, SftTrainingRecord +from evomerge.training_pipeline import ( + ORCHESTRATORS, + TRAINING_MODES, + CheckpointEntry, + CheckpointRegistry, + DatasetVersion, + LossTelemetry, + TrainingJobManifest, + build_training_job_manifest, + compute_dataset_digest, + load_dataset_manifest, + parse_trainer_log, + records_for_mode, + telemetry_from_loss_history, + telemetry_from_summary, + telemetry_to_training_health, + version_dataset, +) +from evomerge.trust_score import AgentTrustScoreBuilder + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + +def _sft_dicts(n: int = 2) -> list[dict]: + return [ + { + "schema_version": "sft/v1", + "messages": [ + {"role": "user", "content": f"question {i}"}, + {"role": "assistant", "content": f"answer {i}"}, + ], + "output_type": "final_answer", + "provenance": {"source": "test"}, + } + for i in range(n) + ] + + +def _sft_model() -> SftTrainingRecord: + return SftTrainingRecord( + messages=[Message(role="user", content="hi"), Message(role="assistant", content="yo")], + output_type="final_answer", + provenance=Provenance(source="test"), + ) + + +def _branch(rollout_id="r1", branch_index=0, score=1, status="pass", answer="Good."): + return RolloutBranchRecord( + rollout_id=rollout_id, + task="Do the thing.", + branch_index=branch_index, + temperature=0.7, + session_id="s1", + final_answer=answer, + objective_score=score, + objective_status=status, + rank=branch_index, + total_score=float(score), + ) + + +# =========================================================================== +# 1. Dataset versioning +# =========================================================================== + +class TestDatasetVersioning: + def test_version_writes_jsonl_and_manifest(self, tmp_path): + dv = version_dataset(_sft_dicts(3), name="sft", out_dir=tmp_path, sources=["a.jsonl"]) + assert dv.name == "sft" + assert dv.schema_version == "sft/v1" + assert dv.record_count == 3 + assert dv.version.startswith("v-") + assert dv.sources == ["a.jsonl"] + # files exist + assert (tmp_path / f"sft-{dv.version}.jsonl").exists() + assert (tmp_path / f"sft-{dv.version}.manifest.json").exists() + + def test_content_addressed_idempotent(self, tmp_path): + records = _sft_dicts(2) + dv1 = version_dataset(records, name="sft", out_dir=tmp_path) + dv2 = version_dataset(records, name="sft", out_dir=tmp_path) + assert dv1.content_digest == dv2.content_digest + assert dv1.version == dv2.version + assert dv1.path == dv2.path # same content → same versioned path + + def test_different_content_different_version(self, tmp_path): + dv1 = version_dataset(_sft_dicts(2), name="sft", out_dir=tmp_path) + dv2 = version_dataset(_sft_dicts(3), name="sft", out_dir=tmp_path) + assert dv1.content_digest != dv2.content_digest + assert dv1.version != dv2.version + + def test_digest_key_order_independent(self): + a = {"schema_version": "sft/v1", "messages": [], "a": 1, "b": 2} + b = {"b": 2, "a": 1, "messages": [], "schema_version": "sft/v1"} + assert compute_dataset_digest([a]) == compute_dataset_digest([b]) + + def test_manifest_round_trip(self, tmp_path): + dv = version_dataset(_sft_dicts(1), name="dpo", out_dir=tmp_path) + loaded = load_dataset_manifest(dv.manifest_path) + assert isinstance(loaded, DatasetVersion) + assert loaded == dv + + def test_accepts_pydantic_models(self, tmp_path): + dv = version_dataset([_sft_model()], name="sft", out_dir=tmp_path) + assert dv.record_count == 1 + assert dv.schema_version == "sft/v1" + + def test_rejects_unknown_mode(self, tmp_path): + with pytest.raises(ValueError, match="unknown training mode"): + version_dataset(_sft_dicts(1), name="grpo", out_dir=tmp_path) + + +# =========================================================================== +# 2. Checkpoint management +# =========================================================================== + +class TestCheckpointRegistry: + def _entry(self, ckpt_id, mode="sft", base_ref="Qwen/Qwen2.5-1.5B-Instruct", **kw): + return CheckpointEntry( + checkpoint_id=ckpt_id, + mode=mode, + dataset_version=kw.get("dataset_version", "digest-a"), + base_ref=base_ref, + status=kw.get("status", "ready"), + metrics=kw.get("metrics", {}), + ) + + def test_register_and_find(self, tmp_path): + reg = CheckpointRegistry(tmp_path / "checkpoints.jsonl") + e = reg.register(self._entry("sft-v1")) + assert reg.find("sft-v1") == e + assert reg.find("missing") is None + + def test_register_dedup(self, tmp_path): + reg = CheckpointRegistry(tmp_path / "checkpoints.jsonl") + reg.register(self._entry("sft-v1")) + reg.register(self._entry("sft-v1", status="failed")) + assert len(reg.all()) == 1 + # first write wins + assert reg.find("sft-v1").status == "ready" + + def test_latest_by_mode_and_time(self, tmp_path): + reg = CheckpointRegistry(tmp_path / "checkpoints.jsonl") + reg.register(CheckpointEntry("sft-1", "sft", "d1", "base", created_at_ms=100)) + reg.register(CheckpointEntry("dpo-1", "dpo", "d2", "sft-1", created_at_ms=200)) + reg.register(CheckpointEntry("sft-2", "sft", "d3", "base", created_at_ms=300)) + assert reg.latest().checkpoint_id == "sft-2" + assert reg.latest(mode="dpo").checkpoint_id == "dpo-1" + assert reg.latest(mode="ppo") is None + + def test_lineage_walks_base_ref(self, tmp_path): + reg = CheckpointRegistry(tmp_path / "checkpoints.jsonl") + reg.register(CheckpointEntry("sft-1", "sft", "d1", "Qwen-base", created_at_ms=100)) + reg.register(CheckpointEntry("dpo-1", "dpo", "d2", "sft-1", created_at_ms=200)) + chain = reg.lineage("dpo-1") + assert [c.checkpoint_id for c in chain] == ["dpo-1", "sft-1"] + # base model is not a registered checkpoint → chain stops + assert reg.lineage("sft-1")[0].base_ref == "Qwen-base" + + def test_lineage_unknown_is_empty(self, tmp_path): + reg = CheckpointRegistry(tmp_path / "checkpoints.jsonl") + assert reg.lineage("nope") == [] + + def test_persists_across_instances(self, tmp_path): + path = tmp_path / "checkpoints.jsonl" + CheckpointRegistry(path).register(self._entry("sft-v1")) + # new instance reading the same file sees the entry + assert CheckpointRegistry(path).find("sft-v1") is not None + + def test_rejects_unknown_mode(self, tmp_path): + reg = CheckpointRegistry(tmp_path / "checkpoints.jsonl") + with pytest.raises(ValueError, match="unknown mode"): + reg.register(CheckpointEntry("x", "grpo", "d", "base")) + + +# =========================================================================== +# 3. Loss telemetry → trust +# =========================================================================== + +class TestLossTelemetry: + def test_parse_json_log_lines(self): + lines = ['{"loss": 2.0, "step": 10}', "noise", '{"loss": 0.05, "step": 20}'] + t = parse_trainer_log(lines, mode="sft") + assert t.loss_history == [2.0, 0.05] + assert t.final_loss == 0.05 + assert t.initial_loss == 2.0 + assert t.converged is True + + def test_parse_keyvalue_log_line(self): + t = parse_trainer_log(["loss = 0.05", "loss=0.2"], mode="sft") + assert t.loss_history == [0.05, 0.2] + + def test_from_summary(self): + t = telemetry_from_summary( + {"final_loss": 0.05, "loss_history": [2.0, 0.05], "max_steps": 200}, mode="sft" + ) + assert t.final_loss == 0.05 + assert t.steps == 200 + assert t.converged is True + + def test_health_zero_without_evidence(self): + assert telemetry_to_training_health(LossTelemetry(mode="sft")) == 0.0 + assert telemetry_to_training_health(telemetry_from_loss_history([], mode="sft")) == 0.0 + + def test_health_converged_run_is_high(self): + t = telemetry_from_loss_history([2.0, 0.05], mode="sft") + health = telemetry_to_training_health(t) + assert 0.9 <= health <= 1.0 + assert t.converged is True + + def test_health_divergent_run_is_low(self): + t = telemetry_from_loss_history([0.5, 1.5], mode="sft") + health = telemetry_to_training_health(t) + assert 0.0 <= health < 0.3 + assert t.converged is False + + def test_trust_builder_folds_telemetry(self): + builder = AgentTrustScoreBuilder() + builder.add_task_success(True) + builder.add_training_telemetry(telemetry_from_loss_history([2.0, 0.05], mode="sft")) + score = builder.build() + assert "training_health" in score.breakdown + assert score.breakdown["training_health"] is not None + assert score.overall is not None + assert score.overall > 0.0 + + def test_trust_builder_accepts_float_and_ducktype(self): + b1 = AgentTrustScoreBuilder().add_training_telemetry(0.9) + assert b1.build().breakdown["training_health"] == pytest.approx(0.9) + + class _Duck: + mode = "sft" + steps = 2 + initial_loss = 2.0 + final_loss = 0.05 + loss_history = [2.0, 0.05] + converged = True + converged_below = 0.1 + + b2 = AgentTrustScoreBuilder().add_training_telemetry(_Duck()) + assert b2.build().breakdown["training_health"] > 0.9 + + +# =========================================================================== +# 4. Job manifest — wire exporters → orchestration +# =========================================================================== + +class TestTrainingJobManifest: + def test_build_writes_manifest_and_versioned_dataset(self, tmp_path): + manifest, dv = build_training_job_manifest( + records=_sft_dicts(2), + mode="sft", + out_dir=tmp_path, + base_ref="Qwen/Qwen2.5-1.5B-Instruct", + orchestrator="ray", + hyperparameters={"lr": 1e-4}, + sources=["rollouts.jsonl"], + ) + assert isinstance(manifest, TrainingJobManifest) + assert manifest.mode == "sft" + assert manifest.orchestrator == "ray" + assert manifest.base_ref == "Qwen/Qwen2.5-1.5B-Instruct" + assert manifest.dataset["content_digest"] == dv.content_digest + assert manifest.checkpoint_plan["save_steps"] == 50 + assert manifest.hyperparameters == {"lr": 1e-4} + assert manifest.telemetry_sink.endswith(".telemetry.json") + # files written + assert (tmp_path / "training-job.manifest.json").exists() + assert (tmp_path / f"sft-{dv.version}.jsonl").exists() + + def test_records_for_mode_dispatches_exporters(self): + branches = [ + _branch(branch_index=0, score=1, status="pass", answer="Good."), + _branch(branch_index=1, score=0, status="fail", answer="Bad."), + ] + sft = records_for_mode(branches, "sft") + dpo = records_for_mode(branches, "dpo") + ppo = records_for_mode(branches, "ppo") + # SFT: only the passing branch; DPO: one pair; PPO: both branches + assert len(sft) == 1 + assert len(dpo) == 1 + assert len(ppo) == 2 + + def test_records_for_mode_rejects_unknown(self): + with pytest.raises(ValueError, match="unknown mode"): + records_for_mode([], "grpo") + + def test_build_validates_mode_and_orchestrator(self, tmp_path): + with pytest.raises(ValueError, match="unknown mode"): + build_training_job_manifest(records=_sft_dicts(1), mode="grpo", out_dir=tmp_path) + with pytest.raises(ValueError, match="unknown orchestrator"): + build_training_job_manifest( + records=_sft_dicts(1), mode="sft", out_dir=tmp_path, orchestrator="slurm" + ) + + def test_manifest_round_trip_json(self, tmp_path): + manifest, _ = build_training_job_manifest(records=_sft_dicts(1), mode="ppo", out_dir=tmp_path) + data = json.loads((tmp_path / "training-job.manifest.json").read_text()) + assert data["mode"] == "ppo" + assert data["job_id"] == manifest.job_id + + def test_mode_and_orchestrator_constants(self): + assert TRAINING_MODES == ("sft", "dpo", "ppo") + assert "ray" in ORCHESTRATORS and "lightning" in ORCHESTRATORS and "local" in ORCHESTRATORS