train_example.py

"""Synthetic training loop that exercises every metric we are using so far"""

import argparse
import random
import time
import uuid

from monitoring.metrics import Monitor

        
def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--mode", choices=["scrape", "push"], default="scrape")
    ap.add_argument("--pushgateway", default="localhost:9091")
    ap.add_argument("--aim-repo", default=None)
    ap.add_argument("--prom-port", type=int, default=8000)
    ap.add_argument("--steps", type=int, default=100000)
    ap.add_argument("--step-time", type=float, default=1.0)
    ap.add_argument("--spike-at", type=int, default=-1)
    args = ap.parse_args()

    run_id = f"demo-{uuid.uuid4().hex[:8]}"
    hparams = {
        "learning_rate": round(random.choice([1e-4, 3e-4, 5e-5]), 6),
        "batch_size": random.choice([16, 32, 64]),
        "model": "demo-1b",
    }
    print(f"run_id={run_id}  hparams={hparams}")

    loss, acc = 2.5, 0.10
    with Monitor(
        run_id,
        mode=args.mode,
        prom_port=args.prom_port,
        pushgateway=args.pushgateway,
        aim_repo=args.aim_repo,
        hparams=hparams,
        experiment_name="Compose LLM throughput",
    ) as mon:
        try:
            for step in range(args.steps):
                loss = max(0.05, loss * 0.999 + random.uniform(-0.02, 0.02))
                acc = min(0.99, acc + random.uniform(-0.002, 0.006))
                kl = abs(random.gauss(0.05, 0.02))
                tokens = random.randint(800, 1200)
                checks = {
                    "compile": random.random() > 0.05,
                    "lint": random.random() > 0.15,
                    "style": random.random() > 0.10,
                }

                if step == args.spike_at:
                    loss = float("nan")
                    kl = 5.0
        
                mon.log(step, loss=loss, accuracy=acc, kl=kl,
                        tokens=tokens, checks=checks)

                if step % 10 == 0:
                    print(f"step {step:5d}  loss={loss:.3f}  "
                          f"acc={acc:.3f}  kl={kl:.3f}")
                time.sleep(args.step_time)
        except KeyboardInterrupt:
            print("\nstopped")


if __name__ == "__main__":
    main()