Coverage for cafl4ds/run_log.py: 100%

1"""Run logging — one run log, two separate series (SSL loss and health). 

2 

3The Phase-0 exit criterion is a single run log in which the **SSL loss** and the **health 

4metrics** (rankme, drift, probe) can be read *side by side over steps*. This module writes 

5one JSONL file per run, each line a record tagged with its ``series`` (``"loss"`` or 

6``"health"``), and mirrors a human-readable line to loguru. Every health record also carries 

7the most recent loss, so :func:`tabulate` can render the required side-by-side table from the 

8health series alone. 

9""" 

10 

11from __future__ import annotations 

12 

13import json 

14from pathlib import Path 

15from typing import Any 

16 

17from loguru import logger 

18 

19# Columns rendered, in order, by :func:`tabulate` (only those present are shown). 

20_TABLE_COLUMNS = ("step", "era", "loss", "rankme", "cka_drift", "cosine_drift", "knn_acc", "linear_acc") 

21 

22 

23class RunLogger: 

24 """Append-only JSONL logger with distinct loss and health series.""" 

25 

26 def __init__(self, path: str | Path, run_name: str = "run") -> None: 

27 """Open (truncate) the run log. 

28 

29 Args: 

30 path: Destination JSONL file (parent directories are created). 

31 run_name: A short name recorded on every line (for grouping across runs). 

32 """ 

33 self.path = Path(path) 

34 self.run_name = run_name 

35 self.path.parent.mkdir(parents=True, exist_ok=True) 

36 self._file = self.path.open("w", encoding="utf-8") 

37 self._last_loss: float | None = None 

38 self._health: list[dict[str, Any]] = [] 

39 

40 def _write(self, record: dict[str, Any]) -> None: 

41 """Write one record as a JSON line and flush. 

42 

43 Args: 

44 record: The record to serialize. 

45 """ 

46 record = {"run": self.run_name, **record} 

47 self._file.write(json.dumps(record) + "\n") 

48 self._file.flush() 

49 

50 def log_loss(self, step: int, era: int, loss: float) -> None: 

51 """Record one point of the SSL-loss series. 

52 

53 Args: 

54 step: Global step index. 

55 era: Current era (class block). 

56 loss: The SSL loss value at this step. 

57 """ 

58 self._last_loss = loss 

59 self._write({"series": "loss", "step": step, "era": era, "loss": loss}) 

60 

61 def log_health(self, step: int, era: int, metrics: dict[str, float]) -> None: 

62 """Record one point of the health series (with the most recent loss attached). 

63 

64 Args: 

65 step: Global step index. 

66 era: Current era (class block). 

67 metrics: The metric dictionary from :meth:`cafl4ds.monitor.HealthMonitor.measure`. 

68 """ 

69 record = {"series": "health", "step": step, "era": era, "loss": self._last_loss, **metrics} 

70 self._health.append(record) 

71 self._write(record) 

72 shown = " ".join( 

73 f"{k}={record[k]:.4f}" for k in _TABLE_COLUMNS if k in record and isinstance(record[k], int | float) 

74 ) 

75 logger.info(f"[{self.run_name}] health {shown}") 

76 

77 def close(self) -> None: 

78 """Close the underlying file.""" 

79 self._file.close() 

80 

81 def tabulate(self) -> str: 

82 """Render the health series (loss + metrics) as a side-by-side text table. 

83 

84 Returns: 

85 A fixed-width table with one row per health checkpoint, or a placeholder string 

86 if no health was logged. 

87 """ 

88 return tabulate(self._health) 

89 

90 

91def tabulate(health_records: list[dict[str, Any]]) -> str: 

92 """Render health records as a fixed-width side-by-side table. 

93 

94 Args: 

95 health_records: Health-series records (as written by :meth:`RunLogger.log_health`). 

96 

97 Returns: 

98 The formatted table, or a placeholder if there are no records. 

99 """ 

100 if not health_records: 

101 return "(no health records)" 

102 cols = [c for c in _TABLE_COLUMNS if any(c in r and r[c] is not None for r in health_records)] 

103 header = " ".join(f"{c:>12}" for c in cols) 

104 lines = [header, " ".join("-" * 12 for _ in cols)] 

105 for r in health_records: 

106 cells = [] 

107 for c in cols: 

108 v = r.get(c) 

109 cells.append(f"{v:>12.4f}" if isinstance(v, int | float) else f"{'-':>12}") 

110 lines.append(" ".join(cells)) 

111 return "\n".join(lines) 

112 

113 

114def read_run(path: str | Path) -> list[dict[str, Any]]: 

115 """Read a run log back into a list of records. 

116 

117 Args: 

118 path: Path to a JSONL run log. 

119 

120 Returns: 

121 The records, in file order. 

122 """ 

123 with Path(path).open(encoding="utf-8") as f: 

124 return [json.loads(line) for line in f if line.strip()]