Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / data / dataset_logging.py: 36%

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"""Dataset logging configuration with optional per-record rank selection.""" 

16 

17from __future__ import annotations 

18 

19import logging 

20from collections.abc import Callable, Iterable 

21from datetime import datetime 

22from typing import Any, Literal 

23 

24import torch.distributed as dist 

25 

26from hyper_parallel.data.constants import DATASET_DATE_FORMAT, DATASET_LOG_FORMAT, DATASET_LOGGER_NAME 

27 

28RankCondition = bool | Callable[[], bool] 

29DatasetLogLevel = Literal["debug", "info", "warn"] 

30 

31 

32def _get_rank() -> int: 

33 try: 

34 rank = int(dist.get_rank()) 

35 except (RuntimeError, ValueError): 

36 rank = 0 

37 return rank 

38 

39 

40class _DatasetDebugRankFilter(logging.Filter): 

41 """Filter Dataset DEBUG records by distributed rank.""" 

42 

43 def __init__(self) -> None: 

44 """Default Dataset DEBUG output to rank zero.""" 

45 super().__init__() 

46 self.ranks: frozenset[int] | None = frozenset({0}) 

47 

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

49 """Keep regular logs and Dataset DEBUG logs from selected ranks.""" 

50 rank = _get_rank() 

51 record.rank_id = rank 

52 if record.levelno != logging.DEBUG or not record.name.startswith(DATASET_LOGGER_NAME): 

53 return True 

54 rank_enabled = getattr(record, "dataset_rank_enabled", None) 

55 if rank_enabled is not None: 

56 return bool(rank_enabled) 

57 return self.ranks is None or rank in self.ranks 

58 

59 

60class _DatasetLogFormatter(logging.Formatter): 

61 """Format Dataset log timestamps with microsecond precision.""" 

62 

63 def formatTime(self, record: logging.LogRecord, datefmt: str | None = None) -> str: # pylint: disable=C0103 

64 """Format a log record timestamp in local time.""" 

65 record_time = datetime.fromtimestamp(record.created).astimezone() 

66 return record_time.strftime(datefmt) if datefmt else record_time.isoformat(timespec="microseconds") 

67 

68 

69_DEBUG_RANK_FILTER = _DatasetDebugRankFilter() 

70 

71 

72class DatasetLogger(logging.LoggerAdapter): 

73 """Add an optional rank condition to Dataset DEBUG records.""" 

74 

75 def debug(self, msg: object, *args: object, **kwargs: Any) -> None: 

76 """Log on default ranks, or on ranks selected by ``enabled`` when provided.""" 

77 if not self.isEnabledFor(logging.DEBUG): 

78 return 

79 enabled: RankCondition | None = kwargs.pop("enabled", None) 

80 if enabled is not None: 

81 rank_enabled = enabled() if callable(enabled) else enabled 

82 extra = dict(kwargs.get("extra", {})) 

83 extra["dataset_rank_enabled"] = rank_enabled 

84 kwargs["extra"] = extra 

85 kwargs.setdefault("stacklevel", 2) 

86 self.logger.debug(msg, *args, **kwargs) 

87 

88 

89def get_dataset_logger(name: str) -> DatasetLogger: 

90 """Return a Dataset logger supporting the ``enabled`` DEBUG argument.""" 

91 return DatasetLogger(logging.getLogger(name), {}) 

92 

93 

94def enable_dataset_logging(level: DatasetLogLevel, ranks: Iterable[int] | None = (0,)) -> None: 

95 """Enable uniformly formatted Dataset logs at the configured level. 

96 

97 Args: 

98 level: Minimum Dataset log level: ``debug``, ``info``, or ``warn``. 

99 ranks: Global ranks allowed to emit Dataset DEBUG records. Pass 

100 ``None`` to allow every rank. Dataset records above DEBUG remain 

101 visible on every rank. 

102 

103 Raises: 

104 ValueError: If ``level`` is unsupported, or ``ranks`` is invalid. 

105 """ 

106 log_levels = {"debug": logging.DEBUG, "info": logging.INFO, "warn": logging.WARNING} 

107 if level not in log_levels: 

108 raise ValueError(f"Unsupported Dataset log level: {level!r}") 

109 

110 if ranks is None: 

111 selected_ranks = None 

112 else: 

113 selected_ranks = frozenset(ranks) 

114 if not selected_ranks or any(isinstance(rank, bool) or not isinstance(rank, int) or rank < 0 

115 for rank in selected_ranks): 

116 raise ValueError("ranks must contain non-negative integers, or be None for all ranks") 

117 

118 _DEBUG_RANK_FILTER.ranks = selected_ranks 

119 dataset_logger = logging.getLogger(DATASET_LOGGER_NAME) 

120 dataset_logger.setLevel(log_levels[level]) 

121 dataset_logger.propagate = False 

122 dataset_handlers = [handler for handler in dataset_logger.handlers if _DEBUG_RANK_FILTER in handler.filters] 

123 if not dataset_handlers: 

124 dataset_handler = logging.StreamHandler() 

125 dataset_handler.setFormatter(_DatasetLogFormatter(DATASET_LOG_FORMAT, datefmt=DATASET_DATE_FORMAT)) 

126 dataset_handler.addFilter(_DEBUG_RANK_FILTER) 

127 dataset_logger.addHandler(dataset_handler) 

128 dataset_handlers.append(dataset_handler) 

129 for dataset_handler in dataset_handlers: 

130 dataset_handler.setLevel(log_levels[level])