Accumulates labelled GPU memory snapshots and prints a summary.
| 127 | |
| 128 | |
| 129 | class GpuMonitor: |
| 130 | """Accumulates labelled GPU memory snapshots and prints a summary.""" |
| 131 | |
| 132 | def __init__(self, enabled: bool = True): |
| 133 | self.enabled = enabled |
| 134 | self._snaps: list[_Snap] = [] |
| 135 | self._gpu_name: str = "" |
| 136 | self._gpu_total_mb: float = 0.0 |
| 137 | if enabled: |
| 138 | self._gpu_name, self._gpu_total_mb = gpu_name_and_total() |
| 139 | |
| 140 | def snapshot(self, label: str) -> float: |
| 141 | """Take a snapshot. Returns current device-level GPU used MB.""" |
| 142 | if not self.enabled: |
| 143 | return 0.0 |
| 144 | _sync() |
| 145 | used = gpu_device_used_mb() |
| 146 | talloc = _torch_allocated_mb() |
| 147 | self._snaps.append(_Snap(label=label, gpu_used_mb=used, |
| 148 | torch_alloc_mb=talloc, wall_time=time.time())) |
| 149 | return used |
| 150 | |
| 151 | @property |
| 152 | def peak_gpu_mb(self) -> float: |
| 153 | if not self._snaps: |
| 154 | return 0.0 |
| 155 | return max(s.gpu_used_mb for s in self._snaps) |
| 156 | |
| 157 | @property |
| 158 | def snapshots(self) -> list[tuple[str, float]]: |
| 159 | return [(s.label, s.gpu_used_mb) for s in self._snaps] |
| 160 | |
| 161 | def format_summary(self) -> str: |
| 162 | if not self._snaps: |
| 163 | return "(no GPU snapshots)" |
| 164 | |
| 165 | base = self._snaps[0].gpu_used_mb |
| 166 | lines: list[str] = [] |
| 167 | lines.append(f" GPU: {self._gpu_name} ({self._gpu_total_mb / 1024:.1f} GB total)") |
| 168 | lines.append(f" Note: values are device-level used memory (like nvidia-smi).") |
| 169 | lines.append(f" Δ columns show change from previous snapshot.") |
| 170 | lines.append(f" {'Stage':<35s} {'Used MB':>8s} {'Δ MB':>8s} {'This proc':>10s}") |
| 171 | lines.append(f" {'─' * 35} {'─' * 8} {'─' * 8} {'─' * 10}") |
| 172 | prev = base |
| 173 | for s in self._snaps: |
| 174 | delta = s.gpu_used_mb - prev |
| 175 | from_base = s.gpu_used_mb - base |
| 176 | ds = f"+{delta:.0f}" if delta >= 0 else f"{delta:.0f}" |
| 177 | lines.append(f" {s.label:<35s} {s.gpu_used_mb:>8.0f} {ds:>8s} {from_base:>+10.0f}") |
| 178 | prev = s.gpu_used_mb |
| 179 | lines.append(f" {'─' * 35} {'─' * 8} {'─' * 10}") |
| 180 | total_alloc = max(s.gpu_used_mb for s in self._snaps) - base |
| 181 | lines.append(f" {'TOTAL (peak − baseline)':<35s} {'':>8s} {total_alloc:>+10.0f}") |
| 182 | return "\n".join(lines) |
| 183 | |
| 184 | def as_dict(self) -> dict: |
| 185 | base = self._snaps[0].gpu_used_mb if self._snaps else 0 |
| 186 | return { |