Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / sapp_ppb / utils / interactive.py: 97%

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"""Interactive CLI flow used when ``run_pipeline_balance.py`` is launched without arguments.""" 

16from collections import namedtuple 

17from typing import Any, List 

18 

19import hyper_parallel.auto_parallel.sapp_ppb.utils.recompute as Recompute 

20from hyper_parallel.auto_parallel.sapp_ppb.sapp.sapp_pipeline import SappPipeline 

21from hyper_parallel.auto_parallel.sapp_ppb.utils.config import generate_solvable_config, print_dryrun_config 

22from hyper_parallel.auto_parallel.sapp_ppb.utils.error import check_in_bounds 

23from hyper_parallel.auto_parallel.sapp_ppb.utils.layer import Layer 

24from hyper_parallel.auto_parallel.sapp_ppb.utils.logger import logger 

25 

26YES_OR_NO = "[y/n]? " 

27 

28OPTIONAL = " if you wish: " 

29 

30GLOBALARGUMENTS = namedtuple('GLOBALARGUMENTS', ['stage_num', 'micro_batch', 'interleave', 'max_memory']) 

31 

32 

33def default_v(d: Any) -> str: 

34 """Format an inline ``(<default> if none)`` hint for ``input()`` prompts.""" 

35 return " (" + str(d) + " if none): " 

36 

37 

38def is_yes(s: str) -> bool: 

39 """Return True when ``s`` is a yes-like response (``y``, ``yes``, ``1``).""" 

40 return s.lower().startswith('y') or s == "1" 

41 

42 

43def is_empty(s: str) -> bool: 

44 """Return True when ``s`` is blank or equal to ``*``.""" 

45 return len(s.strip()) == 0 or (s.strip() == '*') 

46 

47 

48def global_arguments() -> GLOBALARGUMENTS: 

49 """Prompt the user for the global pipeline arguments and return them as a named tuple.""" 

50 stage_num = 4 

51 micro_batch = 8 

52 interleave = 1 

53 max_memory = 56000 

54 s = input("Please enter the pipeline stage number" + default_v(stage_num)) 

55 if not is_empty(s): 

56 stage_num = int(s) 

57 check_in_bounds(stage_num, "Pipeline stage number", 1, 10000) 

58 

59 s = input("Please enter the micro batch number" + default_v(micro_batch)) 

60 if not is_empty(s): 

61 micro_batch = int(s) 

62 check_in_bounds(micro_batch, "Micro batch number", 1, 10000) 

63 

64 s = input("Please enter the pipeline interleave number" + default_v(interleave)) 

65 if not is_empty(s): 

66 interleave = int(s) 

67 check_in_bounds(interleave, "Interleave", 1, 10) 

68 

69 s = input("Please enter maximum memory" + default_v(max_memory)) 

70 if not is_empty(s): 

71 max_memory = int(s) 

72 check_in_bounds(max_memory, "Maximum memory", 1, 1000000) 

73 

74 return GLOBALARGUMENTS(stage_num, micro_batch, interleave, max_memory) 

75 

76 

77def make_layer(t: Layer.type_enum, model_name: str) -> Layer: 

78 """Prompt the user for one layer's metadata and return the resulting :class:`Layer`.""" 

79 nb_layer = 1 

80 layer_time = 0 

81 memory_parameter = 0 

82 memory_activation_rec = {r: None for r in Recompute.TYPE} 

83 layer_name = "misc_" + t.name 

84 s = input("\tEnter the layer name" + OPTIONAL) 

85 if not is_empty(s): 

86 layer_name = s 

87 s = input("\tEnter the layer execution time: ") 

88 if not is_empty(s): 

89 layer_time = int(s) 

90 if t is Layer.type_enum.BODY: 

91 s = input("\tEnter the number of such layer: ") 

92 if not is_empty(s): 

93 nb_layer = int(s) 

94 s = input("\tEnter the layer parameter memory (MB): ") 

95 if not is_empty(s): 

96 memory_parameter = int(s) 

97 for r in Recompute.TYPE: 

98 s = input("\tEnter the layer " + Recompute.JSON_MEMORY_NAME[r] + OPTIONAL) 

99 if not is_empty(s): 

100 memory_activation_rec[r] = int(s) 

101 else: 

102 s = input("\tEnter the layer memory (MB): ") 

103 if not is_empty(s): 

104 memory_parameter = int(s) 

105 

106 return Layer(name=layer_name, ltype=t, nb_layer=nb_layer, time=layer_time, 

107 model_name=model_name, memory_activation_rec=memory_activation_rec, 

108 memory_parameter=memory_parameter,) 

109 

110 

111def dryrun_guide() -> None: 

112 """Prompt the user through a dry-run configuration and print a candidate layout.""" 

113 considered_rec: List[Recompute.TYPE] = [] 

114 stage_num = 0 

115 num_layers = 0 

116 s = input("Please enter the pipeline stage number" + default_v(stage_num)) 

117 if not is_empty(s): 

118 stage_num = int(s) 

119 check_in_bounds(stage_num, "Pipeline stage number", 1, 10000) 

120 else: 

121 return 

122 

123 s = input("Please enter the number of layers" + default_v(num_layers)) 

124 if not is_empty(s): 

125 num_layers = int(s) 

126 check_in_bounds(num_layers, "Micro batch number", 1, 10000) 

127 else: 

128 return 

129 

130 s = input("Do you consider full recomputation?" + YES_OR_NO) 

131 if is_yes(s): 

132 considered_rec.append(Recompute.TYPE.FULL) 

133 

134 s = input("Do you consider select recomputation?" + YES_OR_NO) 

135 if is_yes(s): 

136 considered_rec.append(Recompute.TYPE.SLCT) 

137 

138 s = input("Does your communication recomputation co-work with select recomputation?" + YES_OR_NO) 

139 if is_yes(s): 

140 considered_rec.append(Recompute.TYPE.BOTH) 

141 

142 s = input("Do you consider extra communication recomputation?" + YES_OR_NO) 

143 if is_yes(s): 

144 considered_rec.append(Recompute.TYPE.COMM) 

145 

146 offset_config_list, rec_config_list = generate_solvable_config(stage_num, num_layers, considered_rec) 

147 print_dryrun_config(offset_config_list, rec_config_list) 

148 

149 

150def main() -> None: 

151 """Entry point for the interactive session launched without CLI arguments.""" 

152 s = input( 

153 "No arguments were given. Would you like to proceed to the interactive mode " + YES_OR_NO) 

154 if not is_yes(s): 

155 return 

156 

157 global_args = global_arguments() 

158 number_of_stage = global_args.stage_num 

159 number_of_micro_batch = global_args.micro_batch 

160 interleave_degree = global_args.interleave 

161 max_memory = global_args.max_memory 

162 

163 model_name = "misc" 

164 s = input("\tEnter the model name" + OPTIONAL) 

165 if not is_empty(s): 

166 model_name = s 

167 

168 layers = [] 

169 for ltype in Layer.type_enum: 

170 if ltype is not Layer.type_enum.UNKNOWN: 

171 logger.info("Please enter information of your network %s", ltype.name) 

172 layers.append(make_layer(ltype, model_name)) 

173 

174 pipe = SappPipeline(model_name=model_name, num_of_stage=number_of_stage, 

175 num_of_micro_batch=number_of_micro_batch, max_memory=max_memory, 

176 layers=layers, num_of_interleave=interleave_degree,) 

177 

178 for layer in layers: 

179 logger.info("%s", layer) 

180 

181 pipe.construct_problem(solver="pulp") 

182 pipe.solve_problem(time_limit=40, dump_folder="output") 

183 pipe.print_yaml_results() 

184 pipe.simulate(show=True)