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