Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / sapp_ppb / pp_modeling / pp_balancer.py: 75%
208 statements
« 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"""PP balancer — ILP-based pipeline parallelism load balancing."""
17from __future__ import annotations
19import logging
20from dataclasses import asdict
21from typing import Any, Dict, List, Optional, Tuple
23from hyper_parallel.auto_parallel.sapp_ppb.pp_config_builder.layer_loader import (
24 SAPP_PPB_AVAILABLE,
25 _get_pipeline_layer_class,
26)
27from hyper_parallel.auto_parallel.sapp_ppb.pp_config_builder.yaml_parser import (
28 YamlOptimizationConfig,
29)
30from hyper_parallel.auto_parallel.sapp_ppb.pp_modeling.pp_structs import (
31 PPBOutput,
32 PPStrategyResult,
33)
34from hyper_parallel.auto_parallel.sapp_ppb.utils.recompute import TYPE as RecomputeType
36if SAPP_PPB_AVAILABLE:
37 from hyper_parallel.auto_parallel.sapp_ppb.sapp.sapp_pipeline import SappPipeline # pylint: disable=C0415
38 from hyper_parallel.auto_parallel.sapp_ppb.utils import recompute as Recompute # pylint: disable=C0415
40logger = logging.getLogger(__name__)
43class PPBalancer:
44 """ILP-based pipeline parallelism load balancer.
46 Receives a :class:`LayerBuilder` (which holds the converted
47 sapp-ppb Layer objects and YAML/JSON configs), constructs and
48 solves the ILP, extracts the balanced partition, and builds
49 the strategy result.
51 Args:
52 layer_builder: A :class:`LayerBuilder` instance with validated
53 configuration and constructed Layer objects.
55 Example:
56 >>> builder = LayerBuilder(yaml_config, json_path)
57 >>> balancer = PPBalancer(builder)
58 >>> output = balancer.balance_with_ilp(time_limit=60)
59 """
61 def __init__(
62 self,
63 layer_builder: Any,
64 ) -> None:
65 """Initialize PPBalancer.
67 Args:
68 layer_builder: A :class:`LayerBuilder` instance.
70 Raises:
71 ImportError: If sapp-ppb module is not available.
72 """
73 if not SAPP_PPB_AVAILABLE:
74 raise ImportError(
75 "sapp-ppb module is not available. "
76 "Please ensure sapp-ppb is installed and accessible."
77 )
79 self._layer_builder = layer_builder
80 self._pipeline: Optional[Any] = None
81 self._is_successful: bool = False
83 @property
84 def yaml_config(self) -> YamlOptimizationConfig:
85 """YAML configuration with pipeline topology."""
86 return self._layer_builder.yaml_config
88 @property
89 def pipeline(self) -> Optional[Any]:
90 """Solved SappPipeline instance (available after :meth:`balance_with_ilp`)."""
91 return self._pipeline
93 @staticmethod
94 def _make_infeasible_output(
95 reason: str,
96 error: str = "",
97 solver_status: Any = None,
98 ) -> PPBOutput:
99 """Create a PPBOutput indicating infeasibility.
101 Args:
102 reason: Human-readable infeasibility reason.
103 error: Optional error message string.
104 solver_status: Optional solver status code.
106 Returns:
107 PPBOutput with is_feasible=False and is_successful=False.
108 """
109 details: Dict[str, Any] = {"reason": reason}
110 if error:
111 details["error"] = error
112 if solver_status is not None:
113 details["solver_status"] = solver_status
114 return PPBOutput(
115 stage_partition=[],
116 layer_offset={},
117 is_feasible=False,
118 infeasibility_details=details,
119 )
121 @staticmethod
122 def _build_feasible_output(
123 stage_partition: List[List[Tuple[int, RecomputeType]]],
124 layer_offset: Optional[Dict[str, List[List[int]]]] = None,
125 num_of_interleave: int = 1,
126 ) -> PPBOutput:
127 """Build a PPBOutput from the ILP solution.
129 Args:
130 stage_partition: Per-(virtual-)stage list of ``(layer_id, RecomputeType)`` tuples.
131 Length is ``vpp * pp`` when VPP > 1, otherwise ``pp``.
132 layer_offset: Per-group layer offset (key=group_name, value shape ``[vpp][pp]``).
133 num_of_interleave: VPP interleaving factor (1 = no VPP).
135 Returns:
136 PPBOutput with is_feasible=True and is_successful=True.
137 """
138 if layer_offset is None:
139 layer_offset = {}
141 return PPBOutput(
142 stage_partition=stage_partition,
143 layer_offset=layer_offset,
144 is_feasible=True,
145 is_successful=True,
146 infeasibility_details={},
147 num_of_interleave=num_of_interleave,
148 )
150 @staticmethod
151 def _pulp_has_feasible_solution(pulp_problem: Any) -> bool:
152 """Check whether a PuLP problem has a feasible solution despite non-optimal status.
154 When CBC times out (status=0 / Not Solved), it may have found a
155 feasible incumbent without proving optimality. PuLP exposes this
156 via ``sol_status``: ``LpSolutionOptimal`` (1) or
157 ``LpSolutionIntegerFeasible`` (2) indicate a genuine feasible
158 solution was found. ``LpSolutionNoSolutionFound`` (0) means no
159 incumbent exists even though decision variables may have default
160 values assigned.
162 Args:
163 pulp_problem: The PuLP LpProblem instance.
165 Returns:
166 True if ``sol_status`` indicates a feasible or optimal
167 solution was found.
168 """
169 from pulp import ( # pylint: disable=C0415
170 LpSolutionOptimal,
171 LpSolutionIntegerFeasible,
172 )
173 try:
174 return pulp_problem.sol_status in (LpSolutionOptimal, LpSolutionIntegerFeasible)
175 except (AttributeError, TypeError):
176 return False
178 def _check_ilp_solve_status(self) -> Optional[PPBOutput]:
179 """Check ILP solver status and determine feasibility.
181 Distinguishes three categories:
183 1. **Optimal** (status=1) → return ``None`` (proceed to extract).
184 2. **Not Solved** (status=0) with feasible incumbent
185 (``sol_status`` is ``LpSolutionOptimal`` or
186 ``LpSolutionIntegerFeasible``) → return ``None`` (proceed to
187 extract, ``_is_successful`` set to ``True``).
188 3. **Infeasible / Undefined / Unbounded / Not Solved without
189 incumbent** (``sol_status`` is ``LpSolutionNoSolutionFound``)
190 → return infeasible PPBOutput.
192 Returns:
193 PPBOutput with ``is_feasible=False`` if no usable solution
194 exists, otherwise ``None``.
195 """
196 from pulp import ( # pylint: disable=C0415
197 LpStatusOptimal,
198 LpStatusInfeasible,
199 LpStatusUndefined,
200 LpStatusNotSolved,
201 )
203 if not hasattr(self._pipeline, 'problem_'):
204 return None
205 pulp_problem = getattr(self._pipeline.problem_, 'problem_', None)
206 if not pulp_problem or not hasattr(pulp_problem, 'status'):
207 return None
209 status = pulp_problem.status
211 if status == LpStatusOptimal:
212 self._is_successful = True
213 return None
215 if status == LpStatusInfeasible:
216 return self._make_infeasible_output(
217 "ILP solver returned infeasible status",
218 solver_status=status,
219 )
221 if status == LpStatusUndefined:
222 return self._make_infeasible_output(
223 "ILP solver returned undefined status",
224 solver_status=status,
225 )
227 if status == LpStatusNotSolved:
228 if self._pulp_has_feasible_solution(pulp_problem):
229 self._is_successful = True
230 return None
231 return self._make_infeasible_output(
232 "ILP solver timed out with no feasible solution found",
233 solver_status=status,
234 )
236 return self._make_infeasible_output(
237 f"ILP solver returned unbounded status {status}",
238 solver_status=status,
239 )
241 def balance_with_ilp(
242 self,
243 time_limit: int = 90,
244 solver: str = "pulp",
245 ) -> PPStrategyResult:
246 """Run ILP-based load balancing using sapp-ppb.
248 Solves the ILP and builds the :class:`PPStrategyResult`.
249 Simulation is *not* run here — that responsibility belongs to
250 :class:`PPOptimizer`.
252 Pipeline topology parameters (``num_of_interleave``,
253 ``vpp_less_memory``, ``optimization_level``) are read from
254 the ``yaml_config`` stored in the :class:`LayerBuilder` passed
255 at construction.
257 Args:
258 time_limit: Solver time limit in seconds.
259 solver: Solver backend ("pulp" or "gurobi").
261 Returns:
262 PP strategy result with balanced partition (without
263 simulation metrics; simulation is the caller's
264 responsibility).
266 Example:
267 >>> builder = LayerBuilder(yaml_config, json_path)
268 >>> balancer = PPBalancer(builder)
269 >>> result = balancer.balance_with_ilp(time_limit=60)
270 """
271 if not SAPP_PPB_AVAILABLE:
272 raise ImportError("sapp-ppb module is not available")
274 if not self._layer_builder.memory_limit:
275 raise ValueError(
276 "memory_limit is required for ILP load balancing. "
277 "Please specify memory_limit in the JSON profile."
278 )
280 num_of_interleave = self.yaml_config.num_of_interleave
281 vpp_less_memory = self.yaml_config.vpp_less_memory
282 optimization_level = self.yaml_config.optimization_level
284 self._pipeline = SappPipeline( # pylint: disable=E0606
285 model_name="sapp_nd_model",
286 num_of_stage=self.yaml_config.pp_degree,
287 num_of_micro_batch=self.yaml_config.micro_batch_num,
288 max_memory=self._layer_builder.memory_limit,
289 layers=self._layer_builder.layers_sapp_ppb,
290 num_of_interleave=num_of_interleave,
291 vpp_less_memory=vpp_less_memory,
292 constant_memory=self._layer_builder.constant_memory,
293 optimization_level=optimization_level,
294 use_backward_time=self._layer_builder.use_backward_time,
295 )
297 self._pipeline.construct_problem(solver=solver)
299 self._pipeline.solve_problem(time_limit=time_limit)
301 ppb_output = self._build_ilp_result()
303 return self._build_strategy_result(ppb_output)
305 def _build_strategy_result(
306 self,
307 ppb_output: PPBOutput,
308 ) -> PPStrategyResult:
309 """Build a :class:`PPStrategyResult` from the ILP output.
311 Translates the raw ILP result into the strategy result that
312 downstream consumers use. Simulation metrics are left at
313 their defaults (``simulation_status="not_run"``, etc.);
314 :class:`PPOptimizer` is responsible for running the simulator
315 and merging those fields.
317 Args:
318 ppb_output: Output from :meth:`_build_ilp_result`.
320 Returns:
321 Assembled strategy result (without simulation metrics).
322 """
323 return PPStrategyResult(
324 pp_degree=self.yaml_config.pp_degree,
325 micro_batch_num=self.yaml_config.micro_batch_num,
326 vpp_less_memory=self.yaml_config.vpp_less_memory,
327 pipeline_bubble=None,
328 **asdict(ppb_output),
329 )
331 def _build_ilp_result(
332 self,
333 ) -> PPBOutput:
334 """Extract results from solved ILP and build PPBOutput.
336 Returns:
337 PPB output with balanced partition.
338 """
339 infeasible_output = self._check_ilp_solve_status()
340 if infeasible_output is not None:
341 return infeasible_output
343 result = self._pipeline.get_result()
345 try:
346 stage_partition = self._extract_stage_partition(result)
347 except RuntimeError as e:
348 return self._make_infeasible_output(
349 "Failed to extract partition from ILP solution",
350 error=str(e),
351 )
353 try:
354 layer_offset = self._extract_layer_offset_from_ilp(stage_partition)
355 except RuntimeError as e:
356 return self._make_infeasible_output(
357 "Failed to extract layer offset from ILP solution",
358 error=str(e),
359 )
361 vpp = self._pipeline.num_of_interleave_
363 return self._build_feasible_output(
364 stage_partition,
365 layer_offset=layer_offset,
366 num_of_interleave=vpp,
367 )
369 def _extract_stage_partition(self, result: Dict[str, List[List[str]]]) -> List[List[Tuple[int, RecomputeType]]]: # pylint: disable=unused-argument
370 """Extract stage partition from sapp-ppb result.
372 Each BODY group has its own entry in the ILP variables (keyed by
373 ``layer.name_``). This method extracts per-group
374 per-interleave-per-stage per-recompute layer counts, maps them
375 to contiguous layer-ID ranges (HEAD=0, then BODY groups in
376 order, then TAIL=total_body+1), and attaches the per-layer
377 ``RecomputeType`` from the ILP decision variables.
379 When ``num_of_interleave == 1`` (no VPP), the output has
380 ``pp_degree`` entries indexed by physical stage.
382 When ``num_of_interleave > 1`` (VPP enabled), the output has
383 ``vpp * pp_degree`` entries treated as **virtual stages**:
384 virtual stage ``v * pp_degree + s`` corresponds to VPP chunk
385 ``v`` on physical stage ``s``. Layer IDs are assigned
386 per-VPP-chunk: each chunk gets a contiguous range of body
387 layer IDs (chunk 0 gets the first range, chunk 1 the next,
388 etc.), preserving the true layer-to-chunk mapping.
390 HEAD (layer_id 0) is placed in virtual stage 0 (chunk 0,
391 stage 0). TAIL (layer_id ``total_body + 1``) is placed in
392 the last virtual stage (chunk ``vpp-1``, stage
393 ``pp_degree-1``).
395 HEAD and TAIL layers are always annotated as
396 ``RecomputeType.NONE``.
398 Args:
399 result: sapp-ppb result dictionary (unused, kept for API compatibility).
401 Returns:
402 List of ``(layer_id, RecomputeType)`` tuples per virtual
403 stage. Length is ``vpp * pp_degree`` when VPP > 1,
404 otherwise ``pp_degree``.
405 """
406 if self._pipeline is None or self._pipeline.problem_ is None:
407 raise RuntimeError("Pipeline not constructed or solved yet")
409 solver = self._pipeline.problem_
410 pp = self.yaml_config.pp_degree
411 vpp = self._pipeline.num_of_interleave_
412 num_virtual_stages = vpp * pp
414 body_group_names = self._get_body_group_names()
415 total_body = self._total_body_layers()
417 stage_partition: List[List[Tuple[int, RecomputeType]]] = [
418 [] for _ in range(num_virtual_stages)
419 ]
421 if vpp <= 1:
422 self._assign_layers_no_vpp(
423 solver, body_group_names, pp, stage_partition,
424 total_body,
425 )
426 else:
427 self._assign_layers_with_vpp(
428 solver, body_group_names, pp, vpp, stage_partition,
429 total_body,
430 )
432 for vstage in stage_partition:
433 vstage.sort(key=lambda entry: entry[0])
435 return stage_partition
437 def _assign_layers_no_vpp(
438 self,
439 solver: Any,
440 body_group_names: List[str],
441 pp: int, # pylint: disable=unused-argument
442 stage_partition: List[List[Tuple[int, RecomputeType]]],
443 total_body: int,
444 ) -> None:
445 """Assign layer IDs when VPP is disabled (vpp <= 1).
447 Iterates over BODY groups in order and assigns contiguous
448 layer IDs to each physical stage based on the summed
449 (across recompute types) ILP variable values.
451 Args:
452 solver: ILP solver object with ``variables_`` attribute.
453 body_group_names: Ordered list of BODY group names.
454 _pp: Pipeline parallel degree (unused, kept for API
455 consistency with :meth:`_assign_layers_with_vpp`).
456 stage_partition: ``[pp]`` list to populate.
457 total_body: Total number of BODY layers.
458 """
459 current_layer_id = 1
460 for group_name in body_group_names:
461 if group_name not in solver.variables_:
462 raise RuntimeError(
463 f"Group '{group_name}' not found in solver variables. "
464 f"Available groups: {list(solver.variables_.keys())}"
465 )
466 body_lay = self._get_body_layer_by_name(group_name)
467 stage_rec_counts = self._extract_group_stage_recompute(
468 solver, group_name, body_lay.recompute_considered_,
469 )
470 for stage_id, rec_counts in enumerate(stage_rec_counts):
471 for rec_type, count in rec_counts:
472 for lid in range(current_layer_id, current_layer_id + count):
473 stage_partition[stage_id].append((lid, rec_type))
474 current_layer_id += count
476 if current_layer_id != total_body + 1:
477 raise RuntimeError(
478 f"ILP layer count mismatch: extracted {current_layer_id - 1} "
479 f"body layers, expected {total_body}"
480 )
482 stage_partition[0].insert(0, (0, RecomputeType.NONE))
483 stage_partition[-1].append((total_body + 1, RecomputeType.NONE))
485 def _assign_layers_with_vpp(
486 self,
487 solver: Any,
488 body_group_names: List[str],
489 pp: int,
490 vpp: int,
491 stage_partition: List[List[Tuple[int, RecomputeType]]],
492 total_body: int,
493 ) -> None:
494 """Assign layer IDs when VPP is enabled (vpp > 1).
496 Each VPP chunk receives a contiguous block of body layer IDs.
497 Chunk 0 gets the first block, chunk 1 the next, and so on.
498 Within each chunk, BODY groups are iterated in order; for each
499 group the per-stage per-recompute counts are read from the ILP
500 variables for that specific interleave index.
502 HEAD (layer_id 0) is placed in virtual stage 0 (chunk 0,
503 stage 0). TAIL (layer_id ``total_body + 1``) is placed in
504 the last virtual stage (chunk ``vpp-1``, stage ``pp-1``).
506 Args:
507 solver: ILP solver object with ``variables_`` attribute.
508 body_group_names: Ordered list of BODY group names.
509 pp: Pipeline parallel degree.
510 vpp: VPP interleaving factor.
511 stage_partition: ``[vpp * pp]`` list to populate.
512 total_body: Total number of BODY layers.
513 """
514 current_layer_id = 1
515 for inter in range(vpp):
516 for group_name in body_group_names:
517 if group_name not in solver.variables_:
518 raise RuntimeError(
519 f"Group '{group_name}' not found in solver variables. "
520 f"Available groups: {list(solver.variables_.keys())}"
521 )
522 chunk_stage_rec = self._extract_chunk_stage_recompute(
523 solver, group_name, inter, self._get_body_layer_by_name(group_name).recompute_considered_,
524 )
525 for stage_id, rec_counts in enumerate(chunk_stage_rec):
526 for rec_type, count in rec_counts:
527 stage_partition[inter * pp + stage_id].extend(
528 (layer_id, rec_type) for layer_id in range(current_layer_id, current_layer_id + count)
529 )
530 current_layer_id += count
532 if current_layer_id != total_body + 1:
533 raise RuntimeError(
534 f"ILP layer count mismatch: extracted {current_layer_id - 1} "
535 f"body layers, expected {total_body}"
536 )
538 stage_partition[0].insert(0, (0, RecomputeType.NONE))
539 stage_partition[vpp * pp - 1].append((total_body + 1, RecomputeType.NONE))
541 def _extract_chunk_stage_recompute(
542 self,
543 solver: Any,
544 group_name: str,
545 interleave: int,
546 layer_recompute_considered: Optional[Dict[Any, bool]] = None,
547 ) -> List[List[Tuple[Any, int]]]:
548 """Extract per-stage per-recompute layer counts for a single VPP chunk.
550 Unlike :meth:`_extract_group_stage_recompute` which sums across
551 all interleaves, this method reads only the specified interleave
552 index from the ILP variables.
554 Args:
555 solver: ILP solver object with ``variables_`` attribute.
556 group_name: Name of the layer group.
557 interleave: The VPP chunk (interleave) index to extract.
558 layer_recompute_considered: Optional per-layer recompute
559 considered dict.
561 Returns:
562 Per-stage list of ``(recompute_type, count)`` pairs for the
563 specified interleave.
564 """
565 recompute_considered = layer_recompute_considered
566 pp = self.yaml_config.pp_degree
567 stage_rec: List[List[Tuple[Any, int]]] = [
568 [] for _ in range(pp)
569 ]
571 for stage_id in range(pp):
572 for rec in Recompute.TYPE: # pylint: disable=E0606
573 if recompute_considered and not recompute_considered.get(rec, False):
574 continue
575 try:
576 var_value = solver.variables_[group_name][rec][interleave][stage_id].varValue
577 if var_value is not None:
578 count = round(var_value)
579 if count > 0:
580 stage_rec[stage_id].append((rec, count))
581 except (KeyError, AttributeError):
582 continue
584 return stage_rec
586 def _get_body_group_names(self) -> List[str]:
587 """Return ordered list of BODY group names from layers.
589 Returns:
590 List of group name strings, in the order the BODY layers were
591 appended to the layer list.
592 """
593 pipeline_layer = _get_pipeline_layer_class()
594 return [
595 lay.name_ for lay in self._layer_builder.layers_sapp_ppb
596 if lay.type_ == pipeline_layer.type_enum.BODY
597 ]
599 def _total_body_layers(self) -> int:
600 """Return total number of BODY layers across all groups.
602 Returns:
603 Sum of ``nb_layer_`` for all BODY layer objects.
604 """
605 pipeline_layer = _get_pipeline_layer_class()
606 return sum(
607 lay.nb_layer_ for lay in self._layer_builder.layers_sapp_ppb
608 if lay.type_ == pipeline_layer.type_enum.BODY
609 )
611 def _get_body_layer_by_name(self, name: str) -> Any:
612 """Find a BODY layer by name from layers.
614 Args:
615 name: The ``name_`` of the BODY layer to find.
617 Returns:
618 The matching :class:`Layer` object.
620 Raises:
621 RuntimeError: If no BODY layer with the given name exists.
622 """
623 pipeline_layer = _get_pipeline_layer_class()
624 for lay in self._layer_builder.layers_sapp_ppb:
625 if lay.type_ == pipeline_layer.type_enum.BODY and lay.name_ == name:
626 return lay
627 raise RuntimeError(f"BODY layer '{name}' not found in layers")
629 def _extract_group_stage_recompute(
630 self,
631 solver: Any,
632 group_name: str,
633 layer_recompute_considered: Optional[Dict[Any, bool]] = None,
634 ) -> List[List[Tuple[Any, int]]]:
635 """Extract per-stage per-recompute layer counts for a group.
637 For each stage, returns a list of ``(Recompute.TYPE, count)``
638 pairs for recompute types that have non-zero layer counts in
639 that stage (summed across all interleaves).
641 Args:
642 solver: ILP solver object with ``variables_`` attribute.
643 group_name: Name of the layer group.
644 layer_recompute_considered: Optional per-layer recompute
645 considered dict. When provided, used instead of the
646 solver's global ``recompute_considered_`` so that each
647 BODY group is filtered by its own supported types.
649 Returns:
650 Per-stage list of ``(recompute_type, count)`` pairs.
651 """
652 recompute_considered = layer_recompute_considered
653 stage_rec: List[List[Tuple[Any, int]]] = [
654 [] for _ in range(self.yaml_config.pp_degree)
655 ]
657 for stage_id in range(self.yaml_config.pp_degree):
658 for rec in Recompute.TYPE:
659 if recompute_considered and not recompute_considered.get(rec, False):
660 continue
661 total_count = 0
662 for inter in range(self._pipeline.num_of_interleave_):
663 try:
664 var_value = solver.variables_[group_name][rec][inter][stage_id].varValue
665 if var_value is not None:
666 total_count += round(var_value)
667 except (KeyError, AttributeError):
668 continue
669 if total_count > 0:
670 stage_rec[stage_id].append((rec, total_count))
672 return stage_rec
674 def _extract_layer_offset_from_ilp(
675 self,
676 stage_partition: List[List[Tuple[int, RecomputeType]]], # pylint: disable=unused-argument
677 ) -> Dict[str, List[List[int]]]:
678 """Extract per-group layer offset from ILP solution using sapp-ppb native semantics.
680 For each BODY group, delegates to
681 :func:`Recompute.yaml_from_internal` which computes the offset
682 following the sapp-ppb convention::
684 offset[group_name][i][s] = actual_group[i][s] - nass[i][s]
686 where ``nass`` is the naive layer assignment (pure integer division
687 ``nb_layer_ // (pp * vpp)``) per ``(interleave, stage)`` cell for
688 that group, and ``actual_group[i][s]`` is the total body layers the
689 ILP assigned to that cell for that group (summed across recompute
690 types).
692 Using naive (uncorrected) nass ensures round-trip consistency with
693 :func:`Recompute.internal_from_yaml` and
694 :meth:`SappPipeline.print_yaml_results`, which both use the same
695 naive nass baseline.
697 **Edge case —** ``nb_layer_ < pp * vpp``: the naive nass is 0 for
698 every cell, so the offset equals the total ILP assignment for that
699 cell. This is consistent with the round-trip.
701 Args:
702 stage_partition: The extracted stage partition from ILP
703 (unused; offset is computed directly from solver
704 variables to preserve the VPP dimension).
706 Returns:
707 Per-group offset dict. Key is the BODY group name; value
708 is a list of shape ``[vpp][pp]``. Values may be negative
709 when the ILP deviates from the naive uniform baseline (see
710 edge-case note above).
711 """
712 if self._pipeline is None or self._pipeline.problem_ is None:
713 return {}
715 solver = self._pipeline.problem_
716 pp = self.yaml_config.pp_degree
717 vpp = self._pipeline.num_of_interleave_
719 pipeline_layer = _get_pipeline_layer_class()
721 result: Dict[str, List[List[int]]] = {}
723 for lay in self._layer_builder.layers_sapp_ppb:
724 if lay.type_ != pipeline_layer.type_enum.BODY:
725 continue
727 group_name = lay.name_
728 if group_name not in solver.variables_:
729 raise RuntimeError(
730 f"Cannot extract layer offset from ILP: '{group_name}' not in solver variables"
731 )
733 raw_nass = (
734 [[lay.nb_layer_ // (pp * vpp)] * pp for _ in range(vpp)]
735 if (pp * vpp) > 0
736 else [[0] * pp for _ in range(vpp)]
737 )
739 yaml_out = Recompute.yaml_from_internal( # pylint: disable=E0606
740 vpp, pp, solver.variables_[group_name], raw_nass,
741 )
742 result[group_name] = yaml_out[Recompute.OFFSET] # pylint: disable=E0606
744 return result