"""Account-free TensorBoard tracking, with optional W&B and durable JSON history.""" import json import os import time import warnings from pathlib import Path from .common import atomic_json class Tracker: def __init__(self, cfg, run_dir, run_id, resume=False, disabled=False, resume_step=None): self.directory = Path(run_dir) self.directory.mkdir(parents=True, exist_ok=True) self.run = None self.wandb = None self.failure = None self.writer = None self.backend = cfg.get("tracking", {}).get("backend", "wandb") self.last_step = 0 if not disabled and self.backend == "tensorboard": from torch.utils.tensorboard import SummaryWriter self.writer = SummaryWriter( str(self.directory / "tensorboard"), flush_secs=10, purge_step=resume_step if resume else None, ) self.writer.add_text( "experiment/configuration", "```json\n" + json.dumps(cfg, indent=2) + "\n```", 0 ) atomic_json( self.directory / "dashboard.json", { "backend": "tensorboard", "url": "http://127.0.0.1:6006", "log_directory": "tensorboard", }, ) self.project = cfg["wandb"]["project"] self.entity = cfg["wandb"].get("entity") or os.environ.get("WANDB_ENTITY") if not disabled and self.backend == "wandb" and cfg["wandb"]["mode"] != "disabled": try: import wandb self.wandb = wandb self.run = wandb.init( project=self.project, entity=self.entity, id=run_id, name="themis-amos-v0-" + run_id, resume="must" if resume else "never", config=cfg, dir=str(self.directory), mode=cfg["wandb"]["mode"], tags=["amos-style", "qwen3.5-0.8b", "a100-40gb", "soft-labels"], settings=wandb.Settings(init_timeout=45), ) self.entity = self.run.entity self.run.define_metric("global_step") for section in ("train", "validation", "performance", "budget"): self.run.define_metric(section + "/*", step_metric="global_step") atomic_json( self.directory / "wandb-run.json", { "id": run_id, "url": self.run.url, "entity": self.entity, "project": self.project, }, ) except Exception as exc: self.failed("initialization", exc) def failed(self, operation, exc): # Exception text can contain connection details; persist only its class. self.failure = {"operation": operation, "error_type": type(exc).__name__} atomic_json(self.directory / "tracking-status.json", self.failure) warnings.warn( f"W&B {operation} unavailable; durable local tracking continues", stacklevel=2 ) def log(self, metrics, step): self.last_step = step value = {"global_step": step, "timestamp": time.time(), **metrics} with (self.directory / "metrics.jsonl").open("a") as stream: stream.write(json.dumps(value, allow_nan=False) + "\n") stream.flush() if self.writer is not None: for name, number in metrics.items(): if isinstance(number, (int, float)): self.writer.add_scalar(name, number, step) self.writer.flush() if self.run is not None: try: self.run.log(value) except Exception as exc: self.failed("logging", exc) def evaluation(self, metrics, step, section="validation"): flat = { section + "/" + k: v for k, v in metrics["overall"].items() if isinstance(v, (int, float)) } for group in ("category", "type"): for name, scores in metrics[group].items(): for metric in ("soft_ce", "argmax_agreement", "brier", "count"): flat[f"{section}_{group}/{name}_{metric}"] = scores[metric] self.log(flat, step) atomic_json(self.directory / f"{section}-latest.json", metrics) def final_plots(self, raw, calibrated, records, temperatures, limit=100): import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt from .evaluation import probabilities artifacts = self.directory / "figures" artifacts.mkdir(exist_ok=True) fig, axes = plt.subplots(1, 2, figsize=(11, 4), constrained_layout=True) for ax, field, title in zip( axes, ("hard_reliability", "soft_reliability"), ("One-hot reference calibration", "Soft-reference confidence"), strict=True, ): ax.plot([0, 1], [0, 1], "--", color="#9ca3af", linewidth=1) for metrics, label, color in [ (raw, "Before calibration", "#64748b"), (calibrated, "After calibration", "#4f46e5"), ]: bins = metrics["overall"][field] ax.plot( [b["confidence"] for b in bins], [b["reference_agreement"] for b in bins], "o-", label=label, color=color, ) ax.set( title=title, xlabel="Predicted confidence", ylabel="Reference agreement", xlim=(0, 1), ylim=(0, 1), ) ax.legend(fontsize=8) fig.savefig(artifacts / "calibration.png", dpi=170) plt.close(fig) names = sorted(calibrated["category"]) fig, ax = plt.subplots(figsize=(10, 5), constrained_layout=True) ax.barh( names, [calibrated["category"][k]["argmax_agreement"] for k in names], color="#4f46e5" ) ax.set(xlabel="Agreement with target argmax", title="Final test coverage", xlim=(0, 1)) fig.savefig(artifacts / "category-performance.png", dpi=170) plt.close(fig) table = [] for row in records: p = probabilities(row, temperatures) predicted = int(p.argmax()) reference = max(range(len(row["target"])), key=lambda i: row["target"][i]) if predicted != reference and len(table) < limit: table.append( [ row["id"], row["category"], row["type"], row["option_keys"][predicted], row["option_keys"][reference], float(p.max()), ] ) columns = [ "decision_id", "category", "type", "prediction", "reference_argmax", "confidence", ] atomic_json(self.directory / "error-examples.json", {"columns": columns, "rows": table}) if self.writer is not None: import numpy as np from PIL import Image for tag, filename in [ ("calibration/reliability", "calibration.png"), ("coverage/category_performance", "category-performance.png"), ]: with Image.open(artifacts / filename) as picture: self.writer.add_image( tag, np.asarray(picture.convert("RGB")), self.last_step, dataformats="HWC" ) lines = [" | ".join(columns), " | ".join(["---"] * len(columns))] lines += [" | ".join(str(value).replace("|", "\\|") for value in row) for row in table] self.writer.add_text("final/error_examples", "\n".join(lines), self.last_step) self.writer.flush() if self.run is not None: try: self.run.log( { "calibration/reliability": self.wandb.Image( str(artifacts / "calibration.png") ), "coverage/category_performance": self.wandb.Image( str(artifacts / "category-performance.png") ), "final/errors": self.wandb.Table(columns=columns, data=table), "coverage/categories": self.wandb.Table( columns=["category", "count", "agreement", "soft_ce", "brier"], data=[ [ k, calibrated["category"][k]["count"], calibrated["category"][k]["argmax_agreement"], calibrated["category"][k]["soft_ce"], calibrated["category"][k]["brier"], ] for k in names ], ), } ) except Exception as exc: self.failed("final plots", exc) def workspace(self): spec = { "Training": [ "train/soft_ce", "train/learning_rate", "train/gradient_norm", "train/epoch", ], "Validation": [ "validation/soft_ce", "validation/kl", "validation/brier", "validation/argmax_agreement", ], "Performance": [ "performance/decisions_second", "performance/tokens_second", "performance/step_seconds", "performance/peak_memory_gib", ], "Budget": [ "budget/estimated_usd", "budget/remaining_usd", "budget/worker_seconds", "performance/projected_remaining_seconds", ], "Calibration": [ "calibration/temperature_choice", "calibration/temperature_noul", "calibration/temperature_score", ], "Final evaluation": [ "test_raw/soft_ce", "test_calibrated/soft_ce", "test_calibrated/brier", "final/permutation_answer_flip_rate", ], } atomic_json(self.directory / "dashboard-spec.json", spec) if self.writer is not None: self.writer.add_custom_scalars( { section: {metric.split("/")[-1]: ["Multiline", [metric]] for metric in metrics} for section, metrics in spec.items() } ) self.writer.flush() if self.run is None: return try: from wandb_workspaces import workspaces as ws from wandb_workspaces.reports import v2 as panels sections = [ ws.Section( name=section, is_open=True, panels=[ panels.LinePlot( title=metric.split("/")[-1].replace("_", " ").title(), x="global_step", y=[metric], ) for metric in metrics ], ) for section, metrics in spec.items() ] categories = [ "engineering", "ai_agents_evals", "trust_safety", "customer_support", "sales_commerce", "finance", "legal", "product_management", "data_analytics", "documents_meetings", ] sections.append( ws.Section( name="Coverage", is_open=True, panels=[ panels.LinePlot( title="Validation agreement by category", x="global_step", y=[ f"validation_category/{name}_argmax_agreement" for name in categories ], ), panels.MediaBrowser( title="Final category comparison", media_keys=["coverage/category_performance"], ), ], ) ) sections.append( ws.Section( name="Reliability", is_open=True, panels=[ panels.MediaBrowser( title="Calibration before and after", media_keys=["calibration/reliability"], ), ], ) ) workspace = ws.Workspace( entity=self.entity, project=self.project, name="Themis AMOS training", sections=sections, ) workspace.save() atomic_json(self.directory / "dashboard.json", {"url": workspace.url}) except Exception as exc: self.failed("workspace layout", exc) def finish(self, export, summary): atomic_json(self.directory / "summary.json", summary) if self.writer is not None: self.writer.add_text( "final/summary", "```json\n" + json.dumps(summary, indent=2) + "\n```", self.last_step, ) self.writer.flush() self.writer.close() if self.run is not None: try: for key, value in summary.items(): if isinstance(value, (int, float, str, bool)): self.run.summary[key] = value if export: artifact = self.wandb.Artifact("themis-amos-v0", type="model", metadata=summary) artifact.add_dir(str(export)) for name in ("evaluation.json", "provenance.json", "summary.json"): if (self.directory / name).exists(): artifact.add_file(str(self.directory / name)) self.run.log_artifact(artifact, aliases=["selected"]) self.run.finish() except Exception as exc: self.failed("finalization", exc)