Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / trainer / utils / logging.py: 99%

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"""Distributed-aware logging helpers for the trainer. 

16 

17Design (matches industry consensus across , Megatron-LM, DeepSpeed, 

18HF Transformers, ): 

19 

20- ``logger.info / warning / error`` always fire on **every rank** so that 

21 rank-local failures (OOM on rank 7, NCCL timeout on rank 3) are never 

22 silently dropped. 

23- Rank-0-only progress / status messages call **explicit helpers**: 

24 ``logger.info_rank0(...)`` / ``logger.warning_rank0(...)``. 

25- One-shot dedup helpers ``info_once`` / ``warning_once`` use ``lru_cache`` 

26 so the same message never spams every step. 

27 

28Design rejected: a global ``logging.Filter`` returning ``False`` on 

29non-rank-0 records — that's what hyper had before, and it silently dropped 

30rank-local errors. Don't reintroduce it. 

31 

32The functions here are also installed as bound methods on 

33``logging.Logger`` so any module that already does 

34``logger = logging.getLogger(__name__)`` gets the helpers for free. 

35""" 

36import functools 

37import logging 

38import os 

39import sys 

40from typing import Optional 

41 

42# --------------------------------------------------------------------------- 

43# Rank lookup — uses platform when available, falls back to env var. 

44# Going through platform avoids importing torch / mindspore here directly. 

45# --------------------------------------------------------------------------- 

46 

47 

48def _get_rank() -> int: 

49 """Best-effort rank lookup. 

50 

51 Order: 

52 1. ``hyper_parallel.platform.get_platform().get_rank()`` if the process 

53 group is initialised. 

54 2. ``RANK`` env var (set by ``torchrun`` / ``msrun``). 

55 3. ``LOCAL_RANK`` env var. 

56 4. ``0`` as the last resort. 

57 """ 

58 try: 

59 # pylint: disable=C0415 

60 from hyper_parallel import get_platform 

61 platform = get_platform() 

62 return int(platform.get_rank()) 

63 except (ImportError, RuntimeError, ValueError, AttributeError): 

64 pass 

65 return int(os.environ.get("RANK", os.environ.get("LOCAL_RANK", "0"))) 

66 

67# --------------------------------------------------------------------------- 

68# Configuration entry points. 

69# --------------------------------------------------------------------------- 

70 

71_DEFAULT_FORMAT = ( 

72 "[%(asctime)s][rank%(rank)s][%(levelname)s] %(name)s: %(message)s" 

73) 

74_DATE_FORMAT = "%H:%M:%S" 

75 

76 

77class _RankInjector(logging.Filter): 

78 """Inject the current rank into every record as ``record.rank``. 

79 

80 Unlike a "drop non-rank-0 record" filter, this **never returns False** — 

81 every rank's record is preserved. The format string can use 

82 ``%(rank)s`` to display it. 

83 """ 

84 

85 def filter(self, record: logging.LogRecord) -> bool: 

86 if not hasattr(record, "rank"): 

87 record.rank = _get_rank() 

88 return True 

89 

90 

91def init_logger( 

92 level: int = logging.INFO, 

93 fmt: str = _DEFAULT_FORMAT, 

94 datefmt: str = _DATE_FORMAT, 

95 stream=None, 

96) -> None: 

97 """Configure the root logger with rank-injecting formatter. 

98 

99 Idempotent: calling twice replaces handlers cleanly so re-importing or 

100 testing doesn't double-print. 

101 

102 Args: 

103 level: Root logger level (default ``INFO``). 

104 fmt: Format string. Must include ``%(rank)s`` if you want the rank 

105 displayed. 

106 datefmt: ``%(asctime)s`` format. 

107 stream: Output stream (default ``sys.stdout``). 

108 """ 

109 root = logging.getLogger() 

110 root.setLevel(level) 

111 # Replace handlers so re-init in tests / notebooks doesn't double-log. 

112 for handler in list(root.handlers): 

113 root.removeHandler(handler) 

114 handler = logging.StreamHandler(stream or sys.stdout) 

115 handler.setLevel(level) 

116 handler.setFormatter(logging.Formatter(fmt, datefmt=datefmt)) 

117 handler.addFilter(_RankInjector()) 

118 root.addHandler(handler) 

119 

120 

121def get_logger(name: Optional[str] = None) -> logging.Logger: 

122 """Return a logger with the rank-aware helpers attached. 

123 

124 Equivalent to ``logging.getLogger(name)`` plus binding 

125 ``info_rank0`` / ``warning_rank0`` / ``info_once`` / ``warning_once`` 

126 on the ``logging.Logger`` class (idempotent). 

127 """ 

128 _install_logger_methods() 

129 return logging.getLogger(name) 

130 

131# --------------------------------------------------------------------------- 

132# Standalone module-level functions. 

133# --------------------------------------------------------------------------- 

134 

135 

136def info_rank0(self, msg, *args, **kwargs) -> None: 

137 """``logger.info`` that fires only on rank 0.""" 

138 if _get_rank() == 0: 

139 kwargs.setdefault("stacklevel", 2) 

140 self.info(msg, *args, **kwargs) 

141 

142 

143def warning_rank0(self, msg, *args, **kwargs) -> None: 

144 """``logger.warning`` that fires only on rank 0.""" 

145 if _get_rank() == 0: 

146 kwargs.setdefault("stacklevel", 2) 

147 self.warning(msg, *args, **kwargs) 

148 

149 

150@functools.lru_cache(maxsize=None) 

151def _info_once_cached(name: str, msg: str) -> None: 

152 """LRU-cached one-shot info; key = (logger name, message).""" 

153 if _get_rank() == 0: 

154 logging.getLogger(name).info(msg) 

155 

156 

157def info_once(self, msg, *args, **kwargs) -> None: # pylint: disable=W0613 

158 """``logger.info`` that fires at most once across the whole run.""" 

159 if args: 

160 msg = msg % args 

161 _info_once_cached(self.name, str(msg)) 

162 

163 

164@functools.lru_cache(maxsize=None) 

165def _warning_once_cached(name: str, msg: str) -> None: 

166 if _get_rank() == 0: 

167 logging.getLogger(name).warning(msg) 

168 

169 

170def warning_once(self, msg, *args, **kwargs) -> None: # pylint: disable=W0613 

171 """``logger.warning`` that fires at most once across the whole run.""" 

172 if args: 

173 msg = msg % args 

174 _warning_once_cached(self.name, str(msg)) 

175 

176# --------------------------------------------------------------------------- 

177# Method installation — bind helpers onto Logger so existing 

178# ``logging.getLogger(__name__)`` consumers get them for free. 

179# --------------------------------------------------------------------------- 

180 

181_INSTALLED = False 

182 

183 

184def _install_logger_methods() -> None: 

185 """Idempotently install rank-aware helpers on ``logging.Logger``.""" 

186 global _INSTALLED 

187 if _INSTALLED: 

188 return 

189 logging.Logger.info_rank0 = info_rank0 

190 logging.Logger.warning_rank0 = warning_rank0 

191 logging.Logger.info_once = info_once 

192 logging.Logger.warning_once = warning_once 

193 _INSTALLED = True 

194 

195# Install at import time so any module that imports this package gets 

196# the methods automatically. 

197_install_logger_methods()