Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / trainer / callbacks / logging_callback.py: 92%

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"""Per-step structured terminal logging callback.""" 

16 

17from __future__ import annotations 

18 

19import logging 

20from typing import Any 

21 

22from .base import Callback, TrainerState 

23 

24 

25logger = logging.getLogger(__name__) 

26 

27 

28class LoggingCallback(Callback): 

29 """Log shared step metrics on global rank zero without breaking tqdm.""" 

30 

31 def __init__(self, trainer: Any) -> None: 

32 """Initialize the configured logging cadence. 

33 

34 Args: 

35 trainer: Trainer that owns the callback lifecycle. 

36 """ 

37 super().__init__(trainer) 

38 self.logging_steps = trainer.config.training.logging_steps 

39 self._last_logged_step: int | None = None 

40 

41 @staticmethod 

42 def _format_value(value: Any) -> str: 

43 """Format a scalar metric compactly for one terminal line.""" 

44 item = getattr(value, "item", None) 

45 if callable(item): 

46 value = item() 

47 try: 

48 return f"{float(value):.6g}" 

49 except (TypeError, ValueError): 

50 return str(value) 

51 

52 @classmethod 

53 def _format_message(cls, state: TrainerState, metrics: dict[str, Any]) -> str: 

54 """Build a deterministic, complete metric line for one optimizer step.""" 

55 fields = [f"step={state.global_step}", f"epoch={state.epoch}"] 

56 fields.extend( 

57 f"{name}={cls._format_value(value)}" 

58 for name, value in sorted(metrics.items()) 

59 ) 

60 return " ".join(fields) 

61 

62 def _write(self, message: str) -> None: 

63 """Write through the active tqdm bar or fall back to standard logging.""" 

64 tqdm_callback = getattr(self.trainer, "tqdm_callback", None) 

65 tqdm_write = getattr(tqdm_callback, "write", None) 

66 if callable(tqdm_write) and tqdm_write(message): 

67 return 

68 logger.info("%s", message) 

69 

70 def on_step_end( 

71 self, 

72 state: TrainerState, 

73 loss: float, 

74 loss_dict: dict[str, float], 

75 grad_norm: float, 

76 **kwargs: Any, 

77 ) -> None: 

78 """Log all shared environment metrics at the configured cadence.""" 

79 del loss, grad_norm, kwargs 

80 if self.logging_steps <= 0: 

81 return 

82 if getattr(self.trainer, "global_rank", 0) != 0: 

83 return 

84 if state.global_step % self.logging_steps != 0: 

85 return 

86 if self._last_logged_step == state.global_step: 

87 return 

88 

89 metrics = dict(getattr(self.trainer, "step_env_metrics", {})) 

90 if "training/router_aux_loss" not in metrics: 

91 # Without the meter, callers supply their own step aggregation. 

92 # The logger does not normalize local values across micro-batches or ranks. 

93 aux_loss = loss_dict.get("router_aux_loss") if loss_dict is not None else None 

94 if aux_loss is None: 

95 moe_monitor = getattr(self.trainer, "moe_monitor", None) 

96 if moe_monitor is not None: 

97 aux_loss = moe_monitor.last_mean_aux_loss 

98 if aux_loss is not None: 

99 metrics["training/router_aux_loss"] = float(aux_loss) 

100 if not metrics: 

101 return 

102 self._write(self._format_message(state, metrics)) 

103 self._last_logged_step = state.global_step