| """Structured [VISUAL-*] logging for checker / refiner / I_c extraction.""" |
| from __future__ import annotations |
|
|
| import json |
| import os |
| import time |
| from dataclasses import dataclass, field |
| from typing import Any, Optional |
|
|
| from opsd_utils import debug_log as opsd_debug |
|
|
|
|
| def _enabled(cfg: Optional[dict]) -> bool: |
| """Return the explicit YAML logging setting. |
| |
| Visual-supervision logging is part of the training recipe. It must not be |
| changed by a process environment variable, otherwise the resolved YAML |
| snapshot no longer describes the run that produced its artifacts. |
| """ |
| return bool((cfg or {}).get("enabled", True)) |
|
|
|
|
| def _sample_count(cfg: Optional[dict]) -> int: |
| if cfg and cfg.get("sample_count") is not None: |
| return int(cfg["sample_count"]) |
| return 3 |
|
|
|
|
| def _preview_chars(cfg: Optional[dict]) -> int: |
| if cfg and cfg.get("preview_chars") is not None: |
| return int(cfg["preview_chars"]) |
| return 400 |
|
|
|
|
| def _save_artifacts(cfg: Optional[dict]) -> bool: |
| if cfg is not None and "save_artifacts" in cfg: |
| return bool(cfg["save_artifacts"]) |
| return True |
|
|
|
|
| def _prefix(tag: str) -> str: |
| ts = time.strftime("%Y-%m-%d %H:%M:%S") |
| step = opsd_debug._DETAIL_STEP if opsd_debug._DETAIL_STEP is not None else "?" |
| return ( |
| f"[{tag}][{ts}][rank={opsd_debug._RANK}/{opsd_debug._WORLD_SIZE}]" |
| f"[global_step={step}]" |
| ) |
|
|
|
|
| def _fmt_preview(text: Any, max_len: int) -> str: |
| s = ("" if text is None else str(text)).replace("\n", "\\n") |
| if len(s) > max_len: |
| return s[: max_len - 3] + "..." |
| return s |
|
|
|
|
| def log_visual(tag: str, msg: str, *, cfg: Optional[dict] = None, **fields: Any) -> None: |
| if not _enabled(cfg): |
| return |
| if opsd_debug._RANK != 0: |
| return |
| preview = _preview_chars(cfg) |
| extra = "" |
| if fields: |
| extra = " | " + " | ".join( |
| f"{k}={_fmt_preview(v, preview)}" for k, v in fields.items() |
| ) |
| print(f"{_prefix(tag)} {msg}{extra}", flush=True) |
|
|
|
|
| @dataclass |
| class VisualBatchRecorder: |
| """Accumulates per-generate-batch visual supervision stats.""" |
|
|
| global_step: int |
| output_dir: str |
| log_cfg: dict = field(default_factory=dict) |
| ic_ok: int = 0 |
| ic_fail: int = 0 |
| ic_cache_hit: int = 0 |
| ic_chars: list[float] = field(default_factory=list) |
| checker_scores: list[float] = field(default_factory=list) |
| checker_high: int = 0 |
| checker_medium: int = 0 |
| checker_low: int = 0 |
| checker_skipped_no_thinking: int = 0 |
| checker_local_fallback: int = 0 |
| checker_image_missing: int = 0 |
| checker_aux_evidence_used: int = 0 |
| checker_parse_failure: int = 0 |
| refiner_changed: int = 0 |
| refiner_unchanged: int = 0 |
| refiner_in_lens: list[int] = field(default_factory=list) |
| refiner_out_lens: list[int] = field(default_factory=list) |
| refiner_fallback: int = 0 |
| pool_updates: int = 0 |
| artifacts: list[dict] = field(default_factory=list) |
| route_bindings: list[dict] = field(default_factory=list) |
| ic_latency_ms: float = 0.0 |
| checker_latency_ms: float = 0.0 |
| refiner_latency_ms: float = 0.0 |
| ic_calls: int = 0 |
| checker_calls: int = 0 |
| refiner_calls: int = 0 |
| teacher_batch_calls: int = 0 |
| ic_batch_calls: int = 0 |
| checker_batch_calls: int = 0 |
| refiner_batch_calls: int = 0 |
|
|
| def record_teacher_timing( |
| self, |
| kind: str, |
| *, |
| latency_ms: float, |
| n_calls: int = 1, |
| batch_size: int = 1, |
| ) -> None: |
| self.teacher_batch_calls += 1 |
| if kind == "ic": |
| self.ic_latency_ms += float(latency_ms) |
| self.ic_calls += int(n_calls) |
| self.ic_batch_calls += 1 |
| elif kind == "checker": |
| self.checker_latency_ms += float(latency_ms) |
| self.checker_calls += int(n_calls) |
| self.checker_batch_calls += 1 |
| elif kind == "refiner": |
| self.refiner_latency_ms += float(latency_ms) |
| self.refiner_calls += int(n_calls) |
| self.refiner_batch_calls += 1 |
| else: |
| self.checker_latency_ms += float(latency_ms) |
| self.checker_calls += int(n_calls) |
| self.checker_batch_calls += 1 |
|
|
| def record_ic(self, **fields: Any) -> None: |
| if fields.get("cache_hit"): |
| self.ic_cache_hit += 1 |
| if fields.get("parse_ok"): |
| self.ic_ok += 1 |
| self.ic_chars.append(float(fields.get("ic_chars", 0))) |
| else: |
| self.ic_fail += 1 |
| n = _sample_count(self.log_cfg) |
| if fields.get("sample_idx", 999) < n: |
| log_visual("VISUAL-IC", f"sample[{fields.get('sample_idx')}]", cfg=self.log_cfg, **fields) |
| if _save_artifacts(self.log_cfg): |
| self.artifacts.append({"kind": "ic", **fields}) |
|
|
| def record_checker(self, **fields: Any) -> None: |
| score = float(fields.get("score", 0.0)) |
| self.checker_scores.append(score) |
| label = fields.get("label", "low") |
| if label == "high": |
| self.checker_high += 1 |
| elif label == "medium": |
| self.checker_medium += 1 |
| else: |
| self.checker_low += 1 |
| if fields.get("skipped_no_thinking"): |
| self.checker_skipped_no_thinking += 1 |
| if fields.get("local_fallback"): |
| self.checker_local_fallback += 1 |
| if fields.get("image_missing"): |
| self.checker_image_missing += 1 |
| if fields.get("aux_evidence_used"): |
| self.checker_aux_evidence_used += 1 |
| if fields.get("parse_failure"): |
| self.checker_parse_failure += 1 |
| n = _sample_count(self.log_cfg) |
| idx = int(fields.get("sample_idx", 0)) |
| show = idx < n or fields.get("force_log") |
| if show: |
| log_visual("VISUAL-CHECKER", f"sample[{idx}]", cfg=self.log_cfg, **fields) |
| if _save_artifacts(self.log_cfg): |
| self.artifacts.append({"kind": "checker", **fields}) |
|
|
| def record_refiner(self, **fields: Any) -> None: |
| changed = bool(fields.get("changed")) |
| if changed: |
| self.refiner_changed += 1 |
| else: |
| self.refiner_unchanged += 1 |
| if fields.get("passthrough"): |
| self.refiner_fallback += 1 |
| in_len = int(fields.get("in_len", 0)) |
| out_len = int(fields.get("out_len", 0)) |
| self.refiner_in_lens.append(in_len) |
| self.refiner_out_lens.append(out_len) |
| n = _sample_count(self.log_cfg) |
| idx = int(fields.get("sample_idx", 0)) |
| if idx < n: |
| log_visual("VISUAL-REFINER", f"sample[{idx}]", cfg=self.log_cfg, **fields) |
| if _save_artifacts(self.log_cfg): |
| self.artifacts.append({"kind": "refiner", **fields}) |
|
|
| def record_pool(self, **fields: Any) -> None: |
| if fields.get("written"): |
| self.pool_updates += 1 |
| payload = dict(fields) |
| msg = payload.pop("msg", "template_event") |
| log_visual("VISUAL-POOL", msg, cfg=self.log_cfg, **payload) |
|
|
| def record_route(self, **fields: Any) -> None: |
| if not self.log_cfg.get("log_route_binding", True): |
| return |
| self.route_bindings.append(fields) |
| if opsd_debug.should_log_detail(self.global_step): |
| log_visual("VISUAL-ROUTE", f"sample[{fields.get('sample_idx')}]", cfg=self.log_cfg, **fields) |
| if _save_artifacts(self.log_cfg): |
| self.artifacts.append({"kind": "route", **fields}) |
|
|
| def finish(self) -> dict[str, Any]: |
| n_ic = self.ic_ok + self.ic_fail |
| ic_ok_rate = self.ic_ok / max(n_ic, 1) |
| n_chk = len(self.checker_scores) |
| checker_mean = sum(self.checker_scores) / max(n_chk, 1) |
| n_ref = self.refiner_changed + self.refiner_unchanged |
| refiner_changed_rate = self.refiner_changed / max(n_ref, 1) |
| mean_in = sum(self.refiner_in_lens) / max(len(self.refiner_in_lens), 1) |
| mean_out = sum(self.refiner_out_lens) / max(len(self.refiner_out_lens), 1) |
| mean_delta = mean_out - mean_in |
| mean_ic_chars = sum(self.ic_chars) / max(len(self.ic_chars), 1) |
|
|
| log_visual( |
| "VISUAL-IC", |
| "batch_summary", |
| cfg=self.log_cfg, |
| extract_ok=self.ic_ok, |
| extract_fail=self.ic_fail, |
| mean_ic_chars=round(mean_ic_chars, 1), |
| cache_hit=self.ic_cache_hit, |
| ) |
| log_visual( |
| "VISUAL-CHECKER", |
| "batch_summary", |
| cfg=self.log_cfg, |
| n=n_chk, |
| mean_score=round(checker_mean, 4), |
| high=self.checker_high, |
| medium=self.checker_medium, |
| low=self.checker_low, |
| skipped_no_thinking=self.checker_skipped_no_thinking, |
| local_fallback=self.checker_local_fallback, |
| image_missing=self.checker_image_missing, |
| aux_evidence_used=self.checker_aux_evidence_used, |
| parse_failure=self.checker_parse_failure, |
| ) |
| log_visual( |
| "VISUAL-REFINER", |
| "batch_summary", |
| cfg=self.log_cfg, |
| n=n_ref, |
| changed=self.refiner_changed, |
| unchanged=self.refiner_unchanged, |
| mean_in_len=round(mean_in, 1), |
| mean_out_len=round(mean_out, 1), |
| mean_delta=round(mean_delta, 1), |
| fallback_passthrough=self.refiner_fallback, |
| ) |
| summary = { |
| "visual/ic_ok_rate": ic_ok_rate, |
| "visual/ic_fail_count": float(self.ic_fail), |
| "visual/checker_mean": checker_mean, |
| "visual/checker_high": float(self.checker_high), |
| "visual/checker_medium": float(self.checker_medium), |
| "visual/checker_low": float(self.checker_low), |
| "visual/refiner_changed_rate": refiner_changed_rate, |
| "visual/refiner_mean_delta_len": mean_delta, |
| "visual/pool_updates": float(self.pool_updates), |
| "visual/fallback_checker": float(self.checker_local_fallback), |
| "visual/checker_image_missing": float(self.checker_image_missing), |
| "visual/checker_aux_evidence_used": float(self.checker_aux_evidence_used), |
| "visual/checker_parse_failure": float(self.checker_parse_failure), |
| "visual/fallback_refiner": float(self.refiner_fallback), |
| "visual/ic_latency_ms": round(self.ic_latency_ms, 1), |
| "visual/checker_latency_ms": round(self.checker_latency_ms, 1), |
| "visual/refiner_latency_ms": round(self.refiner_latency_ms, 1), |
| "visual/ic_calls": float(self.ic_calls), |
| "visual/checker_calls": float(self.checker_calls), |
| "visual/refiner_calls": float(self.refiner_calls), |
| "visual/teacher_batch_calls": float(self.teacher_batch_calls), |
| "visual/ic_batch_calls": float(self.ic_batch_calls), |
| "visual/checker_batch_calls": float(self.checker_batch_calls), |
| "visual/refiner_batch_calls": float(self.refiner_batch_calls), |
| } |
| log_visual("VISUAL-BATCH", "generate_summary", cfg=self.log_cfg, **summary) |
|
|
| if _save_artifacts(self.log_cfg) and self.artifacts: |
| art_dir = os.path.join( |
| self.output_dir, |
| "visual_supervision", |
| f"step_{self.global_step}", |
| ) |
| os.makedirs(art_dir, exist_ok=True) |
| path = os.path.join(art_dir, f"rank{opsd_debug._RANK}.jsonl") |
| with open(path, "a", encoding="utf-8") as f: |
| for row in self.artifacts: |
| f.write(json.dumps(row, ensure_ascii=False) + "\n") |
|
|
| return summary |
|
|
