« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +0800
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"""Layer loader — load sapp-ppb Layer objects from native JSON and post-process."""
17from __future__ import annotations
19import logging
20import os
21from typing import Any, Dict, List, Optional
23from hyper_parallel.auto_parallel.sapp_ppb.pp_config_builder.yaml_parser import (
24 YamlOptimizationConfig,
25)
27try:
28 from hyper_parallel.auto_parallel.sapp_ppb.sapp.sapp_pipeline import SappPipeline
29 from hyper_parallel.auto_parallel.sapp_ppb.utils import recompute as Recompute
30 from hyper_parallel.auto_parallel.sapp_ppb.utils.layer import generate_layers_list
31except ImportError:
32 SAPP_PPB_AVAILABLE = False
33 SappPipeline = None
34 Recompute = None
35 generate_layers_list = None # type: ignore[assignment]
36else:
37 SAPP_PPB_AVAILABLE = True
39logger = logging.getLogger(__name__)
42def _get_pipeline_layer_class() -> Any:
43 """Return the ``Layer`` class that ``SappPipeline`` actually uses.
45 ``SappPipeline`` imports ``Layer`` from ``sapp_ppb.utils.layer``.
46 Depending on import order, the top-level ``sapp_ppb`` package may
47 resolve to the ``hyper_parallel.auto_parallel.sapp_ppb`` namespace,
48 creating a *different* ``Layer`` class (with a different
49 ``type_enum``) than what ``SappPipeline`` uses internally. To
50 avoid ``filter_layer_type`` returning empty lists due to enum
51 identity mismatches, we always retrieve ``Layer`` from
52 ``SappPipeline``'s own module globals.
53 """
54 if SappPipeline is None:
55 return None
56 return SappPipeline.__init__.__globals__.get("Layer")
59def _apply_recompute_considered(
60 layer: Any,
61 recompute_considered: Dict[Any, bool],
62) -> None:
63 """Override ``recompute_considered_`` on a sapp-ppb Layer object.
65 The ``Layer`` constructor ignores the ``None`` semantics in
66 ``backward_time_rec`` and ``memory_activation_rec`` — it treats
67 ``None`` as "auto-fill with DEFAULT_COEF" rather than "disable". This
68 method forces the Layer's ``recompute_considered_`` to match the
69 desired inference so that the ILP solver only creates decision
70 variables for genuinely configured recompute types.
72 Args:
73 layer: A sapp-ppb Layer object with a ``recompute_considered_``
74 dict attribute.
75 recompute_considered: Desired mapping from
76 :class:`Recompute.TYPE` to bool.
77 """
78 for rec_type in Recompute.TYPE:
79 layer.recompute_considered_[rec_type] = recompute_considered.get(
80 rec_type, False,
81 )
82 layer.compute_internal_time()
85def _head_tail_recompute_considered() -> Dict[Any, bool]:
86 """Return ``recompute_considered`` for HEAD/TAIL layers.
88 HEAD and TAIL layers never participate in recompute decisions; only
89 ``NONE`` is considered.
91 Returns:
92 Dict mapping each :class:`Recompute.TYPE` to ``False`` except
93 ``NONE`` which is ``True``.
94 """
95 return {r: (r == Recompute.TYPE.NONE) for r in Recompute.TYPE}
98class LayerBuilder:
99 """Build sapp-ppb Layer objects from YAML config + native JSON.
101 Uses :func:`generate_layers_list` from the native sapp-ppb parser
102 to read the ``layers_description`` from the JSON file, then applies
103 recompute-considered overrides and consistency validation to produce
104 a list of sapp-ppb ``Layer`` objects that can be consumed by
105 :class:`PPBalancer`.
107 Args:
108 yaml_config: YAML configuration with pipeline topology and
109 ILP constraints.
110 json_path: Path to the native sapp-ppb JSON file (containing
111 ``layers_description``).
113 Example:
114 >>> builder = LayerBuilder(yaml_config, json_path)
115 >>> layers = builder.layers_sapp_ppb
116 """
118 def __init__(
119 self,
120 yaml_config: YamlOptimizationConfig,
121 json_path: str,
122 ) -> None:
123 """Initialize LayerBuilder.
125 Args:
126 yaml_config: YAML configuration with pipeline topology.
127 json_path: Path to the native sapp-ppb JSON file.
129 Raises:
130 ImportError: If sapp-ppb module is not available.
131 ValueError: If ``json_path`` is empty or layers cannot be
132 parsed.
133 """
134 if not SAPP_PPB_AVAILABLE:
135 raise ImportError(
136 "sapp-ppb module is not available. "
137 "Please ensure sapp-ppb is installed and accessible."
138 )
140 if not json_path:
141 raise ValueError(
142 "LayerBuilder requires a json_path. "
143 "Please provide a valid JSON profile path."
144 )
146 yaml_config.validate()
148 self.yaml_config = yaml_config
149 self._memory_limit = yaml_config.memory_limit
150 self._constant_memory = yaml_config.constant_memory
151 self._enable_simulation = yaml_config.enable_simulation
152 self._use_backward_time = yaml_config.use_backward_time
154 layer_folder = os.path.dirname(json_path)
155 model_name = os.path.splitext(os.path.basename(json_path))[0]
157 layers = generate_layers_list(layer_folder, model_name)
158 if not layers:
159 raise ValueError(
160 f"No layers parsed from '{json_path}'. "
161 f"Ensure the JSON file contains a 'layers_description' section."
162 )
164 self._post_process_layers(layers)
165 self._validate_recompute_consistency(layers)
166 self._validate_group_names(layers)
167 self._validate_num_layer_consistency(layers)
169 self.layers_sapp_ppb = layers
171 @property
172 def memory_limit(self) -> Optional[int]:
173 """Maximum memory per stage in MB (from YAML config)."""
174 return self._memory_limit
176 @property
177 def constant_memory(self) -> int:
178 """Constant memory overhead per stage in MB (from YAML config)."""
179 return self._constant_memory
181 @property
182 def use_backward_time(self) -> bool:
183 """Whether to use backward time in ILP optimization."""
184 return self._use_backward_time
186 @staticmethod
187 def _post_process_layers(layers: List[Any]) -> None:
188 """Override recompute_considered on HEAD/TAIL layers.
190 For HEAD/TAIL layers, only NONE recompute is considered.
191 BODY layers retain the recompute_considered mask inferred by
192 the native ``Layer.find_recompute_considered()`` result.
194 Args:
195 layers: List of sapp-ppb Layer objects from
196 :func:`generate_layers_list`.
197 """
198 pipeline_layer = _get_pipeline_layer_class()
199 if pipeline_layer is None:
200 return
202 for layer in layers:
203 if layer.type_ in (pipeline_layer.type_enum.HEAD, pipeline_layer.type_enum.TAIL):
204 _apply_recompute_considered(
205 layer, _head_tail_recompute_considered(),
206 )
208 @staticmethod
209 def _validate_recompute_consistency(layers: List[Any]) -> None:
210 """Validate that all BODY groups share the same recompute_considered mask.
212 In a multi-body-group configuration, all BODY groups must enable
213 the same set of recompute types so that the ILP solver can
214 construct a coherent global mask. If different groups enable
215 different recompute types, the solver would produce incorrect
216 decision variable assignments.
218 Args:
219 layers: List of sapp-ppb Layer objects with
220 ``.recompute_considered_`` and ``.type_`` attributes.
222 Raises:
223 ValueError: If two or more BODY groups have different
224 ``recompute_considered_`` masks.
225 """
226 pipeline_layer = _get_pipeline_layer_class()
227 if pipeline_layer is None:
228 return
230 body_groups: Dict[str, Any] = {}
231 for lay in layers:
232 if lay.type_ == pipeline_layer.type_enum.BODY:
233 body_groups[lay.name_] = lay
235 if len(body_groups) <= 1:
236 return
238 ref_name = next(iter(body_groups))
239 ref_mask = body_groups[ref_name].recompute_considered_
241 inconsistent = []
242 for name, lay in body_groups.items():
243 if lay.recompute_considered_ != ref_mask:
244 inconsistent.append(name)
246 if inconsistent:
247 raise ValueError(
248 f"All BODY groups must share the same recompute_considered mask, "
249 f"but groups {inconsistent} differ from '{ref_name}'. "
250 f"Please ensure all BODY groups enable the same set of recompute types."
251 )
253 @staticmethod
254 def _validate_group_names(layers: List[Any]) -> None:
255 """Validate that body group names do not conflict with ILP solver internal variables.
257 ``SappSolver._create_variables_to_solve_`` stores internal variables
258 (e.g. ``max_stage_time``) and body-layer variables in the same
259 ``variables_`` dict keyed by name. A body group whose name matches
260 a solver internal variable would overwrite the internal entry,
261 causing ``'list' object has no attribute 'varValue'`` errors during
262 result extraction.
264 Args:
265 layers: List of sapp-ppb Layer objects with ``.name_`` attribute.
267 Raises:
268 ValueError: If any group name conflicts with a solver reserved name.
269 """
270 from hyper_parallel.auto_parallel.sapp_ppb.sapp.sapp_solver import SappSolver # pylint: disable=C0415
271 pipeline_layer = _get_pipeline_layer_class()
272 if pipeline_layer is None:
273 return
275 reserved = {
276 v for k, v in vars(SappSolver).items()
277 if isinstance(v, str) and k.isupper()
278 }
279 body_names = {
280 lay.name_ for lay in layers
281 if lay.type_ == pipeline_layer.type_enum.BODY
282 }
283 conflicts = body_names & reserved
284 if conflicts:
285 raise ValueError(
286 f"body_groups names {conflicts} conflict with ILP solver "
287 f"internal variables; please use different names"
288 )
290 def _validate_num_layer_consistency(self, layers: List[Any]) -> None:
291 """Validate that YAML ``num_layer`` matches the actual body layer count from JSON.
293 When ``num_layer`` is specified in the YAML config, it must agree with
294 the total ``nb_layer_`` of all BODY layers in the JSON profile. A
295 mismatch would silently produce incorrect ILP results because the
296 solver uses the JSON-derived count while the user expects the YAML
297 value to be authoritative.
299 Args:
300 layers: List of sapp-ppb Layer objects with ``.type_`` and
301 ``.nb_layer_`` attributes.
303 Raises:
304 ValueError: If ``num_layer`` is provided and does not match the
305 total body layers from JSON.
306 """
307 if self.yaml_config.num_layer is None:
308 return
310 pipeline_layer = _get_pipeline_layer_class()
311 if pipeline_layer is None:
312 return
314 actual_body = sum(
315 lay.nb_layer_ for lay in layers
316 if lay.type_ == pipeline_layer.type_enum.BODY
317 )
318 if actual_body != self.yaml_config.num_layer:
319 raise ValueError(
320 f"num_layer in YAML ({self.yaml_config.num_layer}) does not match "
321 f"the total body layers from JSON ({actual_body}). "
322 f"Please ensure both sources agree or omit num_layer in YAML."
323 )