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 

17import logging 

18from typing import Any 

19 

20from .base import Callback, TrainerState 

21 

22 

23logger = logging.getLogger(__name__) 

24 

25 

26class LoggingCallback(Callback): 

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

28 

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

30 """Initialize the configured logging cadence. 

31 

32 Args: 

33 trainer: Trainer that owns the callback lifecycle. 

34 """ 

35 super().__init__(trainer) 

36 self.logging_steps = trainer.config.training.logging_steps 

37 self._last_logged_step: int | None = None 

38 

39 @staticmethod 

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

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

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

43 if callable(item): 

44 value = item() 

45 try: 

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

47 except (TypeError, ValueError): 

48 return str(value) 

49 

50 @classmethod 

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

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

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

54 fields.extend( 

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

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

57 ) 

58 return " ".join(fields) 

59 

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

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

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

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

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

65 return 

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

67 

68 def on_step_end( 

69 self, 

70 state: TrainerState, 

71 loss: float, 

72 loss_dict: dict[str, float], 

73 grad_norm: float, 

74 **kwargs: Any, 

75 ) -> None: 

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

77 del loss, loss_dict, grad_norm, kwargs 

78 if self.logging_steps <= 0: 

79 return 

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

81 return 

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

83 return 

84 if self._last_logged_step == state.global_step: 

85 return 

86 

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

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

89 aux_loss = moe_monitor.last_mean_aux_loss if moe_monitor is not None else None 

90 if aux_loss is not None: 

91 metrics["training/router_aux_loss"] = aux_loss 

92 if not metrics: 

93 return 

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

95 self._last_logged_step = state.global_step