Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_multicore / diagnostics / common / schedule.py: 98%

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"""Shared step actions and capture-window scheduling for diagnostic tools.""" 

16 

17from __future__ import annotations 

18 

19from collections.abc import Callable 

20from dataclasses import dataclass 

21from enum import Enum 

22 

23 

24class DiagnosticAction(Enum): 

25 """Action selected for one user-visible training step.""" 

26 

27 NONE = 0 

28 WARMUP = 1 

29 RECORD = 2 

30 RECORD_AND_SAVE = 3 

31 

32 

33@dataclass(frozen=True) 

34class _DiagnosticSchedule: 

35 """Callable schedule with Torch-profiler-compatible window arguments.""" 

36 

37 skip_first: int 

38 wait: int 

39 warmup: int 

40 active: int 

41 repeat: int 

42 

43 def __call__(self, step: int) -> DiagnosticAction: 

44 """Return the capture action for one zero-based training step.""" 

45 if not isinstance(step, int) or isinstance(step, bool) or step < 0: 

46 raise ValueError(f"step must be a non-negative integer, got {step!r}") 

47 if step < self.skip_first: 

48 return DiagnosticAction.NONE 

49 

50 relative_step = step - self.skip_first 

51 cycle_length = self.wait + self.warmup + self.active 

52 cycle_index, cycle_step = divmod(relative_step, cycle_length) 

53 if self.repeat and cycle_index >= self.repeat: 

54 return DiagnosticAction.NONE 

55 if cycle_step < self.wait: 

56 return DiagnosticAction.NONE 

57 if cycle_step < self.wait + self.warmup: 

58 return DiagnosticAction.WARMUP 

59 if cycle_step == cycle_length - 1: 

60 return DiagnosticAction.RECORD_AND_SAVE 

61 return DiagnosticAction.RECORD 

62 

63 

64def schedule( 

65 *, 

66 wait: int, 

67 warmup: int, 

68 active: int, 

69 repeat: int = 0, 

70 skip_first: int = 0, 

71) -> Callable[[int], DiagnosticAction]: 

72 """Create a repeatable MegaKernel capture schedule. 

73 

74 Args: 

75 wait: Unrecorded steps between capture windows. 

76 warmup: Preparatory steps before each active window. Profiling records 

77 and discards cycles; Tensor dump does not copy values in these steps. 

78 active: Steps retained in every active window. 

79 repeat: Number of cycles. Zero repeats indefinitely. 

80 skip_first: Initial unrecorded steps before the first cycle. 

81 

82 Returns: 

83 Callable mapping a zero-based step number to DiagnosticAction. 

84 """ 

85 values = { 

86 "wait": wait, 

87 "warmup": warmup, 

88 "active": active, 

89 "repeat": repeat, 

90 "skip_first": skip_first, 

91 } 

92 for name, value in values.items(): 

93 if not isinstance(value, int) or isinstance(value, bool) or value < 0: 

94 raise ValueError(f"{name} must be a non-negative integer, got {value!r}") 

95 if active <= 0: 

96 raise ValueError(f"active must be positive, got {active}") 

97 return _DiagnosticSchedule( 

98 skip_first=skip_first, 

99 wait=wait, 

100 warmup=warmup, 

101 active=active, 

102 repeat=repeat, 

103 )