Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed_data / balance_logging.py: 100%

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"""Format balancing statistics without inspecting tensors or communicating.""" 

16 

17from __future__ import annotations 

18 

19import logging 

20from collections.abc import Sequence 

21from typing import Any 

22 

23logger = logging.getLogger(__name__) 

24 

25 

26def _layout_lines(label: str, ranks: Sequence[int], layout: Sequence[Any], bin_label: str = "mb") -> list[str]: 

27 """Format per-rank pack lengths, sample counts, costs and optional padding.""" 

28 lines = [f" {label}:"] 

29 for rank, bins in zip(ranks, layout): 

30 microbatches = [] 

31 for index, stats in enumerate(bins): 

32 microbatches.append( 

33 f"{bin_label}{index}: seq_len={stats['seq_len']}, samples={stats['samples']}, cost={stats['cost']:.12g}" 

34 ) 

35 if "padding" in stats: 

36 microbatches[-1] += f", padding={stats['padding']}" 

37 lines.append(f" dp{rank}: " + " | ".join(microbatches)) 

38 return lines 

39 

40 

41def _placement_lines(stats: dict[str, Any], ranks: Sequence[int], bin_label: str = "mb") -> list[str]: 

42 lines = [] 

43 before, after = stats.get("bins_before"), stats.get("bins_after") 

44 if before is not None: 

45 lines.extend(_layout_lines("before balance", ranks, before, bin_label)) 

46 for rank, sent, received in zip(ranks, stats.get("send_samples", ()), stats.get("recv_samples", ())): 

47 lines.append(f"dp{rank}: send {sent} samples, recv {received} samples") 

48 if after is not None: 

49 lines.extend(_layout_lines("after balance", ranks, after, bin_label)) 

50 return lines 

51 

52 

53def format_balance_stats(stats: dict[str, Any], step: int) -> str: 

54 """Format original/accepted microbatch layouts with raw predicted LLM costs. 

55 

56 Args: 

57 stats: Sample counts, sequence lengths, predicted costs and transfers 

58 produced by a balancing loader. Native windows include per-round layouts. 

59 step: One-based iteration number, within the epoch for native loaders. 

60 

61 Returns: 

62 One multiline message. Costs retain the planner's units; they are not milliseconds. 

63 """ 

64 lines = [f"[dp-balance] iteration {step} :"] 

65 if "rounds" in stats: 

66 rounds = stats["rounds"] 

67 lines.append( 

68 f" epoch={stats['epoch']}, micro_batches={rounds[0]['micro_batch']}-{rounds[-1]['micro_batch']}, " 

69 f"planner_rank={stats['planner_rank']}, constructor_ranks={stats['constructor_ranks']}" 

70 ) 

71 if "relative_gain" in stats: 

72 lines.append( 

73 f" cost_component=llm, objective={stats['objective']}, original={stats['original_objective']:.12g}, " 

74 f"candidate={stats['candidate_objective']:.12g}, relative_gain={stats['relative_gain']:.12g}, " 

75 f"min_balance_gain={stats['min_balance_gain']:.12g}, accepted={stats['accepted']}" 

76 ) 

77 if "accepted_objective" in stats: 

78 lines.append(f" accepted_objective={stats['accepted_objective']:.12g}") 

79 ranks = stats["group_ranks"] 

80 if "rounds" in stats: 

81 for round_stats in stats["rounds"]: 

82 lines.append(f" micro_batch {round_stats['micro_batch']}:") 

83 lines.extend(_placement_lines(round_stats, ranks, bin_label="pack")) 

84 else: 

85 lines.extend(_placement_lines(stats, ranks)) 

86 return "\n".join(lines) 

87 

88 

89def log_balance_stats(stats: dict[str, Any], step: int) -> None: 

90 """Default log callback; the loader invokes it only on global rank zero. 

91 

92 Args: 

93 stats: Before/after costs, bin counters and transfer counts. 

94 step: One-based delivered step number. 

95 """ 

96 logger.info("%s", format_balance_stats(stats, step)) 

97 

98 

99__all__ = ["format_balance_stats", "log_balance_stats"]