--- a/training/fine_tune.py +++ b/training/fine_tune.py @@ -14,6 +14,8 @@ import time from datetime import datetime from pathlib import Path + +import trackio def parse_args() -> argparse.Namespace: @@ -43,6 +45,12 @@ random.seed(args.seed) args.out_dir.mkdir(parents=True, exist_ok=True) + # Initialize Trackio tracking + trackio.init( + project=args.project, + name=args.run_name, + ) + history: list[dict[str, float]] = [] print(f"Starting run: {args.run_name}") @@ -52,6 +60,10 @@ for step in range(1, args.steps + 1): metric = simulate_metrics(step, args.steps) history.append(metric) + + # Log metrics to Trackio + trackio.log(metric) + print( f"[step {metric['step']:03d}] " f"loss={metric['loss']:.4f} " @@ -73,9 +85,12 @@ summary_path = args.out_dir / f"{args.run_name}-stdout-summary.json" summary_path.write_text(json.dumps(summary, indent=2), encoding="utf-8") + # Finalize Trackio tracking + trackio.finish() + print("Training complete.") print(f"Local stdout summary saved to: {summary_path}") - print("Note: this script does NOT publish metrics to Trackio yet.") + print("Metrics tracked and synced to Trackio dashboard.") return 0