train_matryoshka.py

#!/usr/bin/env python3
"""Train a resource-guarded Turkish Matryoshka E5 retriever."""

from __future__ import annotations

import argparse
import json
import os
import platform
import random
import re
import subprocess
import time
from pathlib import Path

import datasets
import numpy as np
import sentence_transformers
import torch
import transformers
import yaml
from datasets import load_dataset
from sentence_transformers import (
    SentenceTransformer,
    SentenceTransformerTrainer,
    SentenceTransformerTrainingArguments,
    losses,
)
from sentence_transformers.evaluation import TripletEvaluator
from sentence_transformers.training_args import BatchSamplers
from transformers import TrainerCallback


def seed_everything(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)


def e5_query(text: str) -> str:
    return f"query: {text}"


def e5_passage(text: str) -> str:
    return f"passage: {text}"


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", type=Path, default=Path("configs/train_matryoshka.yaml"))
    parser.add_argument("--data-dir", type=Path, default=Path("data/hard-negatives-v1"))
    parser.add_argument(
        "--output-dir", type=Path, default=Path("outputs/goktugtr-e5-matryoshka-278m-v1")
    )
    parser.add_argument("--max-train-rows", type=int)
    parser.add_argument("--max-validation-rows", type=int)
    parser.add_argument("--max-steps", type=int, default=-1)
    parser.add_argument("--protected-pid", type=int)
    parser.add_argument("--protected-command-marker", default="kamera/.venv/bin/python")
    return parser.parse_args()


def latest_complete_checkpoint(output_dir: Path) -> Path | None:
    candidates: list[tuple[int, Path]] = []
    for path in output_dir.glob("checkpoint-*"):
        match = re.fullmatch(r"checkpoint-(\d+)", path.name)
        if not match:
            continue
        required = (
            "model.safetensors",
            "optimizer.pt",
            "scheduler.pt",
            "trainer_state.json",
            "rng_state.pth",
        )
        if all((path / filename).is_file() for filename in required):
            candidates.append((int(match.group(1)), path))
    return max(candidates, default=(0, None), key=lambda item: item[0])[1]


def process_command(pid: int) -> str | None:
    try:
        return Path(f"/proc/{pid}/cmdline").read_bytes().replace(b"\0", b" ").decode().strip()
    except (FileNotFoundError, PermissionError, ProcessLookupError):
        return None


def require_protected_process(pid: int | None, marker: str) -> str | None:
    if pid is None:
        return None
    command = process_command(pid)
    if command is None:
        raise RuntimeError(f"Protected camera process {pid} is not running; refusing to train")
    if marker not in command:
        raise RuntimeError(
            f"PID {pid} no longer matches protected command marker {marker!r}; refusing to train"
        )
    return command


def gpu_snapshot() -> dict[str, int | str] | None:
    command = [
        "nvidia-smi",
        "--query-gpu=name,memory.total,memory.used,memory.free,utilization.gpu,temperature.gpu",
        "--format=csv,noheader,nounits",
    ]
    try:
        values = subprocess.check_output(command, text=True, timeout=10).strip().split(", ")
    except (FileNotFoundError, subprocess.SubprocessError):
        return None
    if len(values) != 6:
        return None
    return {
        "name": values[0],
        "memory_total_mib": int(values[1]),
        "memory_used_mib": int(values[2]),
        "memory_free_mib": int(values[3]),
        "utilization_percent": int(values[4]),
        "temperature_c": int(values[5]),
    }


class CameraSafetyCallback(TrainerCallback):
    """Yield between optimizer steps and stop if the protected camera exits."""

    def __init__(self, pid: int | None, marker: str, pause_seconds: float) -> None:
        self.pid = pid
        self.marker = marker
        self.pause_seconds = pause_seconds

    def on_step_end(self, args, state, control, **kwargs):
        require_protected_process(self.pid, self.marker)
        if self.pause_seconds > 0:
            time.sleep(self.pause_seconds)
        return control


def main() -> None:
    args = parse_args()
    config = yaml.safe_load(args.config.read_text(encoding="utf-8"))
    seed_everything(int(config["seed"]))
    args.output_dir.mkdir(parents=True, exist_ok=True)

    protected_command = require_protected_process(
        args.protected_pid, args.protected_command_marker
    )
    before_gpu = gpu_snapshot()
    if before_gpu and before_gpu["memory_free_mib"] < 8192:
        raise RuntimeError(
            f"Only {before_gpu['memory_free_mib']} MiB GPU memory is free; refusing to train"
        )
    safety = {
        "protected_pid": args.protected_pid,
        "protected_command": protected_command,
        "cuda_memory_fraction": float(config["cuda_memory_fraction"]),
        "step_pause_seconds": float(config["step_pause_seconds"]),
        "gpu_before": before_gpu,
        "launcher_pid": os.getpid(),
    }
    (args.output_dir / "safety.json").write_text(
        json.dumps(safety, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
    )

    dataset = load_dataset(
        "parquet",
        data_files={
            "train": str(args.data_dir / "train.parquet"),
            "validation": str(args.data_dir / "validation.parquet"),
        },
    )
    train = dataset["train"]
    validation = dataset["validation"]
    if args.max_train_rows:
        train = train.select(range(min(args.max_train_rows, len(train))))
    if args.max_validation_rows:
        validation = validation.select(range(min(args.max_validation_rows, len(validation))))

    def format_triplet(row: dict[str, str]) -> dict[str, str]:
        return {
            "anchor": e5_query(row["query"]),
            "positive": e5_passage(row["positive"]),
            "negative": e5_passage(row["negative"]),
        }

    remove_columns = dataset["train"].column_names
    train = train.map(format_triplet, remove_columns=remove_columns, desc="Formatting E5 train")
    validation = validation.map(
        format_triplet, remove_columns=remove_columns, desc="Formatting E5 validation"
    )
    columns = ["anchor", "positive", "negative"]
    train = train.select_columns(columns)
    validation = validation.select_columns(columns)

    if not torch.cuda.is_available():
        raise RuntimeError("CUDA is required for this training run")
    torch.cuda.set_per_process_memory_fraction(float(config["cuda_memory_fraction"]), 0)
    model = SentenceTransformer(
        config["base_model"],
        revision=config["base_model_revision"],
        model_kwargs={"dtype": torch.bfloat16},
    )
    model.max_seq_length = int(config["max_seq_length"])

    evaluator = TripletEvaluator(
        anchors=validation["anchor"],
        positives=validation["positive"],
        negatives=validation["negative"],
        name="goktugtr-matryoshka-validation",
        batch_size=8,
        show_progress_bar=True,
        truncate_dim=min(int(dim) for dim in config["matryoshka_dims"]),
    )
    baseline = evaluator(model, output_path=str(args.output_dir), epoch=0, steps=0)
    (args.output_dir / "baseline_triplet.json").write_text(
        json.dumps(baseline, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
    )

    training_args = SentenceTransformerTrainingArguments(
        output_dir=str(args.output_dir),
        num_train_epochs=float(config["epochs"]),
        max_steps=args.max_steps,
        per_device_train_batch_size=int(config["per_device_batch_size"]),
        per_device_eval_batch_size=8,
        gradient_accumulation_steps=int(config["gradient_accumulation_steps"]),
        learning_rate=float(config["learning_rate"]),
        warmup_ratio=float(config["warmup_ratio"]),
        bf16=bool(config["bf16"]),
        tf32=True,
        gradient_checkpointing=bool(config["gradient_checkpointing"]),
        gradient_checkpointing_kwargs={"use_reentrant": False},
        optim="adamw_torch_fused",
        batch_sampler=BatchSamplers.NO_DUPLICATES,
        eval_strategy="steps",
        eval_steps=int(config["eval_steps"]),
        save_strategy="steps",
        save_steps=int(config["save_steps"]),
        save_total_limit=2,
        logging_steps=int(config["logging_steps"]),
        dataloader_num_workers=0,
        dataloader_pin_memory=True,
        report_to="none",
        run_name=config["project_name"],
        seed=int(config["seed"]),
    )
    base_loss = losses.CachedMultipleNegativesRankingLoss(
        model,
        mini_batch_size=int(config["loss_mini_batch_size"]),
        scale=20.0,
    )
    loss = losses.MatryoshkaLoss(
        model,
        base_loss,
        matryoshka_dims=[int(dim) for dim in config["matryoshka_dims"]],
        n_dims_per_step=int(config["matryoshka_dims_per_step"]),
    )
    trainer = SentenceTransformerTrainer(
        model=model,
        args=training_args,
        train_dataset=train,
        eval_dataset=validation,
        loss=loss,
        evaluator=evaluator,
        callbacks=[
            CameraSafetyCallback(
                args.protected_pid,
                args.protected_command_marker,
                float(config["step_pause_seconds"]),
            )
        ],
    )
    resume_checkpoint = latest_complete_checkpoint(args.output_dir)
    torch.cuda.reset_peak_memory_stats()
    started = time.perf_counter()
    train_output = trainer.train(
        resume_from_checkpoint=str(resume_checkpoint) if resume_checkpoint else None
    )
    training_seconds = time.perf_counter() - started
    trainer.state.save_to_json(str(args.output_dir / "trainer_state.json"))
    final_dir = args.output_dir / "final"
    model.save_pretrained(str(final_dir), safe_serialization=True)
    final_metrics = evaluator(model, output_path=str(args.output_dir), epoch=1, steps=-1)
    (args.output_dir / "final_triplet.json").write_text(
        json.dumps(final_metrics, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
    )
    environment = {
        "python": platform.python_version(),
        "torch": torch.__version__,
        "transformers": transformers.__version__,
        "sentence_transformers": sentence_transformers.__version__,
        "datasets": datasets.__version__,
        "cuda": torch.version.cuda,
        "gpu": torch.cuda.get_device_name(0),
        "config": config,
        "train_rows": len(train),
        "validation_rows": len(validation),
        "training_seconds": training_seconds,
        "resumed_from_checkpoint": str(resume_checkpoint) if resume_checkpoint else None,
        "training_metrics": train_output.metrics,
        "max_gpu_memory_gb": round(torch.cuda.max_memory_allocated() / 2**30, 3),
        "gpu_after": gpu_snapshot(),
        "protected_process_alive_after": process_command(args.protected_pid)
        if args.protected_pid
        else None,
    }
    (args.output_dir / "environment.json").write_text(
        json.dumps(environment, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
    )
    print(json.dumps({"baseline": baseline, "final": final_metrics}, indent=2))


if __name__ == "__main__":
    main()