Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / trainer / callbacks / tqdm_callback.py: 87%
118 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 07:18 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""Terminal progress callback for Trainer."""
17from __future__ import annotations
19import logging
20import sys
21import time
22from typing import Any
24from .base import Callback, TrainerState
27logger = logging.getLogger(__name__)
30class _LogProgressBar:
31 """Write one finalized tqdm-formatted snapshot per step to a non-TTY stream."""
33 def __init__(
34 self,
35 tqdm_type: Any,
36 total: int,
37 initial: int,
38 stream: Any,
39 ) -> None:
40 """Initialize non-interactive progress state without emitting a line."""
41 self._tqdm_type = tqdm_type
42 self._initial = initial
43 self._start_time = time.monotonic()
44 self._postfix: dict[str, str] = {}
45 self._pending_message: str | None = None
46 self.total = total
47 self.n = initial
48 self.fp = stream
50 def set_postfix(self, postfix: dict[str, str], refresh: bool = True) -> None:
51 """Store the final metrics for the next completed-step snapshot.
53 Args:
54 postfix: Formatted metric names and values.
55 refresh: Accepted for compatibility with the tqdm API.
56 """
57 del refresh
58 self._postfix = dict(postfix)
60 def update(self, amount: int) -> None:
61 """Advance progress and write exactly one newline-terminated snapshot.
63 Args:
64 amount: Number of newly completed steps.
65 """
66 self.n += amount
67 elapsed = max(time.monotonic() - self._start_time, 0.0)
68 completed = self.n - self._initial
69 rate = completed / elapsed if elapsed > 0 else None
70 postfix = ", ".join(f"{name}={value}" for name, value in self._postfix.items())
71 progress_line = self._tqdm_type.format_meter(
72 n=self.n,
73 total=self.total,
74 elapsed=elapsed,
75 prefix="Training",
76 unit="step",
77 rate=rate,
78 postfix=postfix or None,
79 )
80 if self._pending_message:
81 progress_line = f"{progress_line} | {self._pending_message}"
82 self.fp.write(f"{progress_line}\n")
83 self.fp.flush()
84 self._pending_message = None
86 def write(self, message: str, file: Any = None) -> None:
87 """Attach one structured message to the next progress snapshot.
89 Args:
90 message: Structured log message for the current step.
91 file: Accepted for compatibility with the tqdm API.
92 """
93 del file
94 self._pending_message = message
96 def close(self) -> None:
97 """Close without writing a duplicate final progress line."""
100class TqdmCallback(Callback):
101 """Display global training progress and shared metrics on global rank zero."""
103 _PRIMARY_METRICS = (
104 ("training/total_loss", "loss"),
105 ("training/grad_norm", "grad_norm"),
106 ("training/lr", "lr"),
107 )
108 _ENV_METRICS = (
109 ("performance/tokens_per_second", "tokens/s"),
110 ("performance/step_time", "step_time"),
111 ("performance/tflops_per_device", "TFLOP/s"),
112 )
114 def __init__(self, trainer: Any) -> None:
115 """Initialize progress state without producing terminal output.
117 Args:
118 trainer: Trainer that owns the callback lifecycle.
119 """
120 super().__init__(trainer)
121 self._progress_bar = None
122 self._last_updated_step = 0
123 self._missing_dependency_warned = False
125 @staticmethod
126 def _to_float(value: Any) -> float | None:
127 """Convert a scalar or scalar tensor-like value to a float."""
128 if value is None:
129 return None
130 item = getattr(value, "item", None)
131 if callable(item):
132 value = item()
133 return float(value)
135 @classmethod
136 def _format_metric(cls, name: str, value: Any) -> str:
137 """Format a shared metric using a unit-appropriate representation."""
138 scalar = cls._to_float(value)
139 if scalar is None:
140 return "n/a"
141 if name == "training/lr":
142 return f"{scalar:.3e}"
143 if name == "performance/tokens_per_second":
144 return f"{scalar:.3f}"
145 if name == "performance/step_time":
146 return f"{scalar:.3f}s"
147 if name == "performance/tflops_per_device":
148 return f"{scalar:.3f}"
149 return f"{scalar:.6g}"
151 def _create_progress_bar(self, total: int, initial: int) -> Any:
152 """Create tqdm lazily because it is an optional dependency."""
153 try:
154 from tqdm import tqdm # pylint: disable=C0415
155 except ImportError:
156 if not self._missing_dependency_warned:
157 logger.warning("TqdmCallback: 'tqdm' is not installed; progress bar disabled")
158 self._missing_dependency_warned = True
159 return None
160 stream = sys.stderr
161 isatty = getattr(stream, "isatty", None)
162 if not callable(isatty) or not isatty():
163 return _LogProgressBar(
164 tqdm_type=tqdm,
165 total=total,
166 initial=initial,
167 stream=stream,
168 )
169 return tqdm(
170 total=total,
171 initial=initial,
172 desc="Training",
173 unit="step",
174 dynamic_ncols=True,
175 file=stream,
176 )
178 def _postfix(self) -> dict[str, str]:
179 """Build the progress postfix exclusively from shared metric dictionaries."""
180 train_metrics = getattr(self.trainer, "step_train_metrics", {})
181 env_metrics = getattr(self.trainer, "step_env_metrics", {})
182 postfix = {
183 label: self._format_metric(name, train_metrics[name])
184 for name, label in self._PRIMARY_METRICS
185 if name in train_metrics
186 }
188 primary_names = {name for name, _ in self._PRIMARY_METRICS}
189 for name, value in sorted(train_metrics.items()):
190 if name in primary_names:
191 continue
192 label = name.split("/", 1)[-1]
193 postfix[label] = self._format_metric(name, value)
195 for name, label in self._ENV_METRICS:
196 if name in env_metrics:
197 postfix[label] = self._format_metric(name, env_metrics[name])
198 return postfix
200 def on_train_begin(self, state: TrainerState, **kwargs: Any) -> None:
201 """Create a single full-training progress bar on global rank zero.
203 Args:
204 state: Current training state.
205 **kwargs: Additional callback arguments.
206 """
207 del kwargs
208 self._last_updated_step = state.global_step
209 if getattr(self.trainer, "global_rank", 0) != 0:
210 return
211 self._progress_bar = self._create_progress_bar(
212 total=self.trainer.train_iters,
213 initial=state.global_step,
214 )
216 def on_step_end(
217 self,
218 state: TrainerState,
219 loss: float | None = None,
220 loss_dict: dict[str, float] | None = None,
221 grad_norm: float | None = None,
222 **kwargs: Any,
223 ) -> None:
224 """Publish shared metrics and advance once per completed global step.
226 Args:
227 state: Current training state.
228 loss: Aggregated loss for the completed step.
229 loss_dict: Named loss values for the completed step.
230 grad_norm: Gradient norm measured before the optimizer update.
231 **kwargs: Additional callback arguments.
232 """
233 del loss, loss_dict, grad_norm, kwargs
234 if self._progress_bar is None or state.global_step <= self._last_updated_step:
235 return
237 postfix = self._postfix()
238 if postfix:
239 self._progress_bar.set_postfix(postfix, refresh=False)
240 self._progress_bar.update(state.global_step - self._last_updated_step)
241 self._last_updated_step = state.global_step
243 def write(self, message: str) -> bool:
244 """Write above the active progress bar and redraw it.
246 Args:
247 message: Complete terminal line to display.
249 Returns:
250 Whether an active progress bar handled the message.
251 """
252 if self._progress_bar is None:
253 return False
254 self._progress_bar.write(message, file=getattr(self._progress_bar, "fp", None))
255 return True
257 def on_train_end(self, state: TrainerState, **kwargs: Any) -> None:
258 """Close the global progress bar idempotently.
260 Args:
261 state: Final training state.
262 **kwargs: Additional callback arguments.
263 """
264 del state, kwargs
265 if self._progress_bar is not None:
266 self._progress_bar.close()
267 self._progress_bar = None