Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / sapp_nd / nd / _pool.py: 91%

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"""Price a search's candidates in worker processes, for run_nd -j. 

16 

17A worker builds its own runner once, from the command-line values run_nd built 

18the main process's from (the recipe), with the censuses the main process ran 

19handed over, so that no worker runs one. The main process walks the loop nest, 

20refuses and sets each candidate as a single process does, and sends the rest 

21as plain degrees in chunks; a worker sets each on its own config, prices it and 

22returns the numbers and the debug parts. The main process writes them in the 

23search's order, so every output reads as one process writes it. 

24""" 

25import multiprocessing 

26import os 

27from collections import Counter 

28from functools import partial 

29from typing import Any, Callable, Dict, List, Optional, Tuple 

30 

31from hyper_parallel.auto_parallel import _hf_model_spec 

32import hyper_parallel.auto_parallel.sapp_nd.nd.debug as Debug 

33from hyper_parallel.auto_parallel.sapp_nd.nd.common import cost_model_preprocess 

34from hyper_parallel.auto_parallel.sapp_nd.perf_estimation.estimate import estimate_performance 

35 

36# This worker's runner, built on its first task. 

37_RUNNER = None 

38 

39 

40def default_workers(verbosity: int) -> int: 

41 """The workers a search uses when run_nd -j states none: ND_JOBS, else half the logical cores, at most 4. 

42 

43 On Windows a worker is a new process that imports torch and parses the run 

44 again, about 7 s of CPU, so on an 8-core laptop four workers took the 

45 Qwen3.5 example to 25.5 s, eight to 27.0 and sixteen to 35.7. Above 

46 verbosity 2 a search prints a line per estimate, which only the main 

47 process can, so it prices in one process. 

48 """ 

49 stated = os.environ.get("ND_JOBS") 

50 if stated: 

51 return max(1, int(stated)) 

52 if verbosity > 2: 

53 return 1 

54 return max(1, min(4, (os.cpu_count() or 2) // 2)) 

55 

56 

57class Workers: 

58 """The processes that price a search's candidates, and the recipe each rebuilds the search's runner from. 

59 

60 run_nd starts them before it parses the run, so that they import while it 

61 does. The search hands them its candidates as plain degrees and gets back 

62 each candidate's number and debug parts, in the order it sent them. 

63 """ 

64 

65 def __init__(self, count: int, recipe: Dict[str, Any], pool: Any = None) -> None: 

66 """Start *count* spawned workers for *recipe*, or use *pool*, which a test passes in their place.""" 

67 self.count = count 

68 self.recipe = recipe 

69 self._pool = pool if pool is not None else multiprocessing.get_context("spawn").Pool( 

70 count, initializer=_warm_up) 

71 

72 def memory(self, items: List[Tuple[int, tuple]]) -> List[Tuple[Any, Any]]: 

73 """Each item's peak and debug parts, in order; an item is its position and the candidate's loop degrees.""" 

74 return self._run(memory_task, items) 

75 

76 def performance(self, items: List[Tuple[int, tuple, Any]], cache_file: Optional[str]) -> List[Tuple[Any, Any]]: 

77 """Each item's score and debug parts, in order; an item adds the peak its generation found.""" 

78 return self._run(partial(perf_task, cache_file=cache_file), items) 

79 

80 def _run(self, func: Callable, items: list) -> List[Tuple[Any, Any]]: 

81 """*func* over *items* in chunks, each with the recipe and the censuses this process ran.""" 

82 censuses = dict(_hf_model_spec._CENSUSES) # pylint: disable=protected-access 

83 tasks = [(self.recipe, censuses, chunk) for chunk in chunks(items, self.count)] 

84 priced = [None] * len(items) 

85 for out, reads in self._pool.imap(func, tasks): 

86 merge_unset_reads(reads) 

87 for index, value, info in out: 

88 priced[index] = (value, info) 

89 return priced 

90 

91 def close(self, abort: bool = False) -> None: 

92 """Stop the workers: once their tasks are done, or at once where the search failed.""" 

93 if abort: 

94 self._pool.terminate() 

95 else: 

96 self._pool.close() 

97 self._pool.join() 

98 

99 

100def _warm_up() -> None: 

101 """A worker's start: its output silenced, the search's modules imported; the runner comes with a task.""" 

102 devnull = os.open(os.devnull, os.O_WRONLY) 

103 os.dup2(devnull, 1) 

104 os.dup2(devnull, 2) 

105 # pylint: disable=import-outside-toplevel,unused-import 

106 import hyper_parallel.auto_parallel.sapp_nd.nd.parallelize # noqa: F401 

107 

108 

109def build_runner(recipe: Dict[str, Any]) -> Any: 

110 """The search runner run_nd builds from *recipe*, the command-line values a pooled search takes.""" 

111 # pylint: disable=import-outside-toplevel 

112 from hyper_parallel.auto_parallel.sapp_nd.memory_estimation.size import Memory 

113 from hyper_parallel.auto_parallel.sapp_nd.nd.logger import set_verbose_level 

114 import hyper_parallel.auto_parallel.sapp_nd.nd.common.hardware as Hard 

115 import hyper_parallel.auto_parallel.sapp_nd.nd.dimensions as Dim 

116 import hyper_parallel.auto_parallel.sapp_nd.nd.parallelize as Par 

117 

118 cost_model_preprocess.set_strict(recipe["strict"]) 

119 set_verbose_level(recipe["verbosity"]) 

120 Debug.set_output_dir(recipe["output_dir"]) 

121 machine = Hard.Machine(recipe["devices"], recipe["device_type"]) 

122 max_mem = Memory.from_string(recipe["max_mem"]) if recipe["max_mem"] is not None else None 

123 return Par.Parallelize( 

124 recipe["framework"], recipe["yaml_config"], machine, 

125 global_batch_size=recipe["global_batch_size"], dimensions=Dim.get_dims(recipe["dimensions"]), 

126 swap_os=recipe["swap_os"], mppb=recipe["mppb"], auto_recompute=False, auto_offload=False, 

127 host_link=None, model=recipe["model"], max_mem=max_mem, 

128 mem_for_ppb=Memory.from_string(recipe["mem_for_ppb"]), 

129 ).instance 

130 

131 

132def _runner(recipe: Dict[str, Any], censuses: Dict[Any, Any]) -> Any: 

133 """This worker's runner, built on its first task with the main process's censuses so that it runs none.""" 

134 global _RUNNER # pylint: disable=global-statement 

135 if _RUNNER is None: 

136 _hf_model_spec._CENSUSES.update(censuses) # pylint: disable=protected-access 

137 _RUNNER = build_runner(recipe) 

138 # The main process reports the fields its own parse read; a worker reports what its estimates read. 

139 cost_model_preprocess.UNSET_READS.clear() 

140 return _RUNNER 

141 

142 

143def _unset_reads() -> Dict[str, Dict[str, int]]: 

144 """The unset fields this worker's estimates read since its last task, which the main process reports.""" 

145 reads = {attr: dict(readers) for attr, readers in cost_model_preprocess.UNSET_READS.items()} 

146 cost_model_preprocess.UNSET_READS.clear() 

147 return reads 

148 

149 

150def memory_task(task: Tuple[dict, dict, List[Tuple[int, tuple]]]) -> Tuple[list, dict]: 

151 """Each candidate's peak, and its debug parts where the search writes them, as inside_loop_nest prices it.""" 

152 recipe, censuses, chunk = task 

153 runner = _runner(recipe, censuses) 

154 out = [] 

155 for index, dims in chunk: 

156 config = runner.config.make_parallel_config(*dims) 

157 runner.config.set_parallel_config(config) 

158 if runner.enable_debug: 

159 debugger = Debug.Debug(config, info_type=Debug.MemParts, enable=True, output_file="debug_mem.csv") 

160 out.append((index, runner.memory_estim(debugger), dict(debugger.info))) 

161 else: 

162 out.append((index, runner.memory_estim(), None)) 

163 return out, _unset_reads() 

164 

165 

166def perf_task(task: Tuple[dict, dict, List[Tuple[int, tuple, Any]]], cache_file: Optional[str]) -> Tuple[list, dict]: 

167 """Each candidate's score, and its debug parts where the search writes them, as order_search_space prices it.""" 

168 recipe, censuses, chunk = task 

169 runner = _runner(recipe, censuses) 

170 out = [] 

171 for index, dims, memory in chunk: 

172 config = runner.config.make_parallel_config(*dims) 

173 runner.config.set_parallel_config(config) 

174 if runner.enable_debug: 

175 debugger = Debug.Debug(config, info_type=Debug.PerfParts, enable=True) 

176 score = estimate_performance(runner.priced(), debugger=debugger, device_type=runner.machine.device, 

177 memory=memory, cache_file=cache_file, stage_savings=None) 

178 out.append((index, score, dict(debugger.info))) 

179 else: 

180 score = estimate_performance(runner.priced(), device_type=runner.machine.device, memory=memory, 

181 cache_file=cache_file, stage_savings=None) 

182 out.append((index, score, None)) 

183 return out, _unset_reads() 

184 

185 

186def merge_unset_reads(reads: Dict[str, Dict[str, int]]) -> None: 

187 """Add a worker's unset reads to this process's, which the search reports under its ranking.""" 

188 for attr, readers in reads.items(): 

189 cost_model_preprocess.UNSET_READS[attr].update(Counter(readers)) 

190 

191 

192def chunks(items: list, workers: int) -> List[list]: 

193 """*items* in contiguous chunks, about eight a worker, so that a slow chunk does not hold the others.""" 

194 size = max(1, len(items) // (workers * 8)) 

195 return [items[i:i + size] for i in range(0, len(items), size)]