Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / trainer / callbacks / tqdm_callback.py: 86%
116 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +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 )
113 def __init__(self, trainer: Any) -> None:
114 """Initialize progress state without producing terminal output.
116 Args:
117 trainer: Trainer that owns the callback lifecycle.
118 """
119 super().__init__(trainer)
120 self._progress_bar = None
121 self._last_updated_step = 0
122 self._missing_dependency_warned = False
124 @staticmethod
125 def _to_float(value: Any) -> float | None:
126 """Convert a scalar or scalar tensor-like value to a float."""
127 if value is None:
128 return None
129 item = getattr(value, "item", None)
130 if callable(item):
131 value = item()
132 return float(value)
134 @classmethod
135 def _format_metric(cls, name: str, value: Any) -> str:
136 """Format a shared metric using a unit-appropriate representation."""
137 scalar = cls._to_float(value)
138 if scalar is None:
139 return "n/a"
140 if name == "training/lr":
141 return f"{scalar:.3e}"
142 if name == "performance/tokens_per_second":
143 return f"{scalar:.3f}"
144 if name == "performance/step_time":
145 return f"{scalar:.3f}s"
146 return f"{scalar:.6g}"
148 def _create_progress_bar(self, total: int, initial: int) -> Any:
149 """Create tqdm lazily because it is an optional dependency."""
150 try:
151 from tqdm import tqdm # pylint: disable=C0415
152 except ImportError:
153 if not self._missing_dependency_warned:
154 logger.warning("TqdmCallback: 'tqdm' is not installed; progress bar disabled")
155 self._missing_dependency_warned = True
156 return None
157 stream = sys.stderr
158 isatty = getattr(stream, "isatty", None)
159 if not callable(isatty) or not isatty():
160 return _LogProgressBar(
161 tqdm_type=tqdm,
162 total=total,
163 initial=initial,
164 stream=stream,
165 )
166 return tqdm(
167 total=total,
168 initial=initial,
169 desc="Training",
170 unit="step",
171 dynamic_ncols=True,
172 file=stream,
173 )
175 def _postfix(self) -> dict[str, str]:
176 """Build the progress postfix exclusively from shared metric dictionaries."""
177 train_metrics = getattr(self.trainer, "step_train_metrics", {})
178 env_metrics = getattr(self.trainer, "step_env_metrics", {})
179 postfix = {
180 label: self._format_metric(name, train_metrics[name])
181 for name, label in self._PRIMARY_METRICS
182 if name in train_metrics
183 }
185 primary_names = {name for name, _ in self._PRIMARY_METRICS}
186 for name, value in sorted(train_metrics.items()):
187 if name in primary_names:
188 continue
189 label = name.split("/", 1)[-1]
190 postfix[label] = self._format_metric(name, value)
192 for name, label in self._ENV_METRICS:
193 if name in env_metrics:
194 postfix[label] = self._format_metric(name, env_metrics[name])
195 return postfix
197 def on_train_begin(self, state: TrainerState, **kwargs: Any) -> None:
198 """Create a single full-training progress bar on global rank zero.
200 Args:
201 state: Current training state.
202 **kwargs: Additional callback arguments.
203 """
204 del kwargs
205 self._last_updated_step = state.global_step
206 if getattr(self.trainer, "global_rank", 0) != 0:
207 return
208 self._progress_bar = self._create_progress_bar(
209 total=self.trainer.train_iters,
210 initial=state.global_step,
211 )
213 def on_step_end(
214 self,
215 state: TrainerState,
216 loss: float | None = None,
217 loss_dict: dict[str, float] | None = None,
218 grad_norm: float | None = None,
219 **kwargs: Any,
220 ) -> None:
221 """Publish shared metrics and advance once per completed global step.
223 Args:
224 state: Current training state.
225 loss: Aggregated loss for the completed step.
226 loss_dict: Named loss values for the completed step.
227 grad_norm: Gradient norm measured before the optimizer update.
228 **kwargs: Additional callback arguments.
229 """
230 del loss, loss_dict, grad_norm, kwargs
231 if self._progress_bar is None or state.global_step <= self._last_updated_step:
232 return
234 postfix = self._postfix()
235 if postfix:
236 self._progress_bar.set_postfix(postfix, refresh=False)
237 self._progress_bar.update(state.global_step - self._last_updated_step)
238 self._last_updated_step = state.global_step
240 def write(self, message: str) -> bool:
241 """Write above the active progress bar and redraw it.
243 Args:
244 message: Complete terminal line to display.
246 Returns:
247 Whether an active progress bar handled the message.
248 """
249 if self._progress_bar is None:
250 return False
251 self._progress_bar.write(message, file=getattr(self._progress_bar, "fp", None))
252 return True
254 def on_train_end(self, state: TrainerState, **kwargs: Any) -> None:
255 """Close the global progress bar idempotently.
257 Args:
258 state: Final training state.
259 **kwargs: Additional callback arguments.
260 """
261 del state, kwargs
262 if self._progress_bar is not None:
263 self._progress_bar.close()
264 self._progress_bar = None