Diff Coverage

                desc["options"] = ["NONE", "FULL"]
                desc["forward_time"] = {"NONE": 1, "FULL": 1}
                desc["backward_time"] = {"NONE": 1, "FULL": 1}
            elif ctx.current_node == ctx.tail_node:
                desc["memory_activation"] = {"NONE": 0, "FULL": 0}
                d_out = self.mb(sum(self._inner_dynamic_mem(ppb=True)))
                desc["memory_parameter"] = self.mb(res_stat) + d_out
                desc["type"] = "TAIL"
                desc["options"] = ["NONE", "FULL"]
                desc["forward_time"] = {"NONE": 1, "FULL": 1}
                desc["backward_time"] = {"NONE": 1, "FULL": 1}
            else:
                desc["memory_activation"] = {"NONE": 0, "COMM": 0, "SLCT": 0, "BOTH": 0, "FULL": 0}
                original_current_node = ctx.current_node
                synthetic_rec_op = False

    from hyper_parallel.auto_parallel.sapp_ppb.sapp.sapp_pipeline import SappPipeline
    from hyper_parallel.auto_parallel.sapp_ppb.utils import recompute as Recompute
    from hyper_parallel.auto_parallel.sapp_ppb.utils.layer import generate_layers_list
    SAPP_PPB_AVAILABLE = True
except ImportError:
    SAPP_PPB_AVAILABLE = False
    SappPipeline = None
    Recompute = None
    generate_layers_list = None  # type: ignore[assignment]

logger = logging.getLogger(__name__)

    identity mismatches, we always retrieve ``Layer`` from
    ``SappPipeline``'s own module globals.
    """
    if SappPipeline is None:
        return None
    return SappPipeline.__init__.__globals__.get("Layer")


def _apply_recompute_considered(

            ValueError: If ``json_path`` is empty or layers cannot be
                parsed.
        """
        if not SAPP_PPB_AVAILABLE:
            raise ImportError(
                "sapp-ppb module is not available. "
                "Please ensure sapp-ppb is installed and accessible."
            )

                :func:`generate_layers_list`.
        """
        pipeline_layer = _get_pipeline_layer_class()
        if pipeline_layer is None:
            return

        for layer in layers:
            if layer.type_ in (pipeline_layer.type_enum.HEAD, pipeline_layer.type_enum.TAIL):
                _apply_recompute_considered(

                ``recompute_considered_`` masks.
        """
        pipeline_layer = _get_pipeline_layer_class()
        if pipeline_layer is None:
            return

        body_groups: Dict[str, Any] = {}
        for lay in layers:
            if lay.type_ == pipeline_layer.type_enum.BODY:

        """
        from hyper_parallel.auto_parallel.sapp_ppb.sapp.sapp_solver import SappSolver  # pylint: disable=C0415
        pipeline_layer = _get_pipeline_layer_class()
        if pipeline_layer is None:
            return

        reserved = {
            v for k, v in vars(SappSolver).items()
            if isinstance(v, str) and k.isupper()

            ValueError: If ``num_layer`` is provided and does not match the
                total body layers from JSON.
        """
        if self.yaml_config.num_layer is None:
            return

        pipeline_layer = _get_pipeline_layer_class()
        if pipeline_layer is None:
            return

        actual_body = sum(
            lay.nb_layer_ for lay in layers
            if lay.type_ == pipeline_layer.type_enum.BODY

            lay.nb_layer_ for lay in layers
            if lay.type_ == pipeline_layer.type_enum.BODY
        )
        if actual_body != self.yaml_config.num_layer:
            raise ValueError(
                f"num_layer in YAML ({self.yaml_config.num_layer}) does not match "
                f"the total body layers from JSON ({actual_body}). "
                f"Please ensure both sources agree or omit num_layer in YAML."
            )

        if value == 0:
            return False
        if value == 1:
            return True
        raise ValueError(
            f"{field_name}: cannot interpret {value} as boolean; "
            f"expected 0 or 1"
        )
    if isinstance(value, float) and value.is_integer():

        if int_val == 0:
            return False
        if int_val == 1:
            return True
        raise ValueError(
            f"{field_name}: cannot interpret {value} as boolean; "
            f"expected 0.0 or 1.0"
        )
    if isinstance(value, str):

        )
        for name in int_fields:
            val = getattr(self, name)
            if isinstance(val, bool) or not isinstance(val, int):
                raise ValueError(
                    f"{name} must be an integer, "
                    f"got {type(val).__name__} {val!r}"
                )
        if self.num_layer is not None:

                )
        if self.num_layer is not None:
            if (isinstance(self.num_layer, bool)
                    or not isinstance(self.num_layer, int)):
                raise ValueError(
                    f"num_layer must be an integer or None, "
                    f"got {type(self.num_layer).__name__} "
                    f"{self.num_layer!r}"
                )

                    f"{self.num_layer!r}"
                )
        for name in ("vpp_less_memory", "enable_simulation", "use_backward_time"):
            if not isinstance(getattr(self, name), bool):
                raise ValueError(
                    f"{name} must be a boolean, "
                    f"got {type(getattr(self, name)).__name__}"
                )
        if (isinstance(self.sim_comm_time, bool)

        )

    pipeline_cfg: dict = cfg.get("pipeline_config", {})
    if not isinstance(pipeline_cfg, dict):
        raise ValueError(
            f"YAML file {yaml_path} must contain a 'pipeline_config' section"
        )

    pp_degree, num_layer, micro_batch_num = _extract_required_fields(

    """Lazily import PPOptimizer to avoid circular imports."""
    if name not in _LAZY_EXPORTS:
        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

    module = _import_module(_LAZY_EXPORTS[name])
    value = getattr(module, name)
    globals()[name] = value
    return value


__all__ = [
    "PPStrategyResult",

        Raises:
            ImportError: If sapp-ppb module is not available.
        """
        if not SAPP_PPB_AVAILABLE:
            raise ImportError(
                "sapp-ppb module is not available. "
                "Please ensure sapp-ppb is installed and accessible."
            )

            PPBOutput with is_feasible=False and is_successful=False.
        """
        details: Dict[str, Any] = {"reason": reason}
        if error:
            details["error"] = error
        if solver_status is not None:
            details["solver_status"] = solver_status
        return PPBOutput(
            stage_partition=[],

        Returns:
            PPBOutput with is_feasible=True and is_successful=True.
        """
        if layer_offset is None:
            layer_offset = {}

        return PPBOutput(
            stage_partition=stage_partition,
            layer_offset=layer_offset,

            LpStatusNotSolved,
        )

        if not hasattr(self._pipeline, 'problem_'):
            return None
        pulp_problem = getattr(self._pipeline.problem_, 'problem_', None)
        if not pulp_problem or not hasattr(pulp_problem, 'status'):
            return None

        status = pulp_problem.status

        if status == LpStatusOptimal:

                "ILP solver timed out with no feasible solution found",
                solver_status=status,
            )

        return self._make_infeasible_output(
            f"ILP solver returned unbounded status {status}",
            solver_status=status,
        )

            >>> balancer = PPBalancer(builder)
            >>> result = balancer.balance_with_ilp(time_limit=60)
        """
        if not SAPP_PPB_AVAILABLE:
            raise ImportError("sapp-ppb module is not available")

        if not self._layer_builder.memory_limit:
            raise ValueError(
                "memory_limit is required for ILP load balancing. "
                "Please specify memory_limit in the JSON profile."
            )

        result = self._pipeline.get_result()

        try:
            stage_partition = self._extract_stage_partition(result)
        except RuntimeError as e:
            return self._make_infeasible_output(
                "Failed to extract partition from ILP solution",
                error=str(e),
            )

            )

        try:
            layer_offset = self._extract_layer_offset_from_ilp(stage_partition)
        except RuntimeError as e:
            return self._make_infeasible_output(
                "Failed to extract layer offset from ILP solution",
                error=str(e),
            )

            stage.  Length is ``vpp * pp_degree`` when VPP > 1,
            otherwise ``pp_degree``.
        """
        if self._pipeline is None or self._pipeline.problem_ is None:
            raise RuntimeError("Pipeline not constructed or solved yet")

        solver = self._pipeline.problem_
        pp = self.yaml_config.pp_degree
        vpp = self._pipeline.num_of_interleave_

                solver, body_group_names, pp, stage_partition,
                total_body,
            )
        else:
            self._assign_layers_with_vpp(
                solver, body_group_names, pp, vpp, stage_partition,
                total_body,
            )

        """
        current_layer_id = 1
        for group_name in body_group_names:
            if group_name not in solver.variables_:
                raise RuntimeError(
                    f"Group '{group_name}' not found in solver variables. "
                    f"Available groups: {list(solver.variables_.keys())}"
                )
            body_lay = self._get_body_layer_by_name(group_name)

                        stage_partition[stage_id].append((lid, rec_type))
                    current_layer_id += count

        if current_layer_id != total_body + 1:
            raise RuntimeError(
                f"ILP layer count mismatch: extracted {current_layer_id - 1} "
                f"body layers, expected {total_body}"
            )

            vpp: VPP interleaving factor.
            stage_partition: ``[vpp * pp]`` list to populate.
            total_body: Total number of BODY layers.
        """
        current_layer_id = 1
        for inter in range(vpp):
            for group_name in body_group_names:
                if group_name not in solver.variables_:
                    raise RuntimeError(
                        f"Group '{group_name}' not found in solver variables. "
                        f"Available groups: {list(solver.variables_.keys())}"
                    )
                body_lay = self._get_body_layer_by_name(group_name)
                chunk_stage_rec = self._extract_chunk_stage_recompute(
                    solver, group_name, inter, body_lay.recompute_considered_,
                )
                for stage_id, rec_counts in enumerate(chunk_stage_rec):
                    for rec_type, count in rec_counts:
                        vstage = inter * pp + stage_id
                        for lid in range(current_layer_id, current_layer_id + count):
                            stage_partition[vstage].append((lid, rec_type))
                        current_layer_id += count

        if current_layer_id != total_body + 1:
            raise RuntimeError(
                f"ILP layer count mismatch: extracted {current_layer_id - 1} "
                f"body layers, expected {total_body}"
            )

        stage_partition[0].insert(0, (0, RecomputeType.NONE))
        stage_partition[vpp * pp - 1].append((total_body + 1, RecomputeType.NONE))

    def _extract_chunk_stage_recompute(
        self,
        solver: Any,

        Returns:
            Per-stage list of ``(recompute_type, count)`` pairs for the
            specified interleave.
        """
        recompute_considered = layer_recompute_considered
        pp = self.yaml_config.pp_degree
        stage_rec: List[List[Tuple[Any, int]]] = [
            [] for _ in range(pp)
        ]

        for stage_id in range(pp):
            for rec in Recompute.TYPE:  # pylint: disable=E0606
                if recompute_considered and not recompute_considered.get(rec, False):
                    continue
                try:
                    var_value = solver.variables_[group_name][rec][interleave][stage_id].varValue
                    if var_value is not None:
                        count = round(var_value)
                        if count > 0:
                            stage_rec[stage_id].append((rec, count))
                except (KeyError, AttributeError):
                    continue

        return stage_rec

    def _get_body_group_names(self) -> List[str]:
        """Return ordered list of BODY group names from layers.

        pipeline_layer = _get_pipeline_layer_class()
        for lay in self._layer_builder.layers_sapp_ppb:
            if lay.type_ == pipeline_layer.type_enum.BODY and lay.name_ == name:
                return lay
        raise RuntimeError(f"BODY layer '{name}' not found in layers")

    def _extract_group_stage_recompute(
        self,
        solver: Any,

                    try:
                        var_value = solver.variables_[group_name][rec][inter][stage_id].varValue
                        if var_value is not None:
                            total_count += round(var_value)
                    except (KeyError, AttributeError):
                        continue
                if total_count > 0:
                    stage_rec[stage_id].append((rec, total_count))

        return stage_rec

            when the ILP deviates from the naive uniform baseline (see
            edge-case note above).
        """
        if self._pipeline is None or self._pipeline.problem_ is None:
            return {}

        solver = self._pipeline.problem_
        pp = self.yaml_config.pp_degree
        vpp = self._pipeline.num_of_interleave_

                continue

            group_name = lay.name_
            if group_name not in solver.variables_:
                raise RuntimeError(
                    f"Cannot extract layer offset from ILP: '{group_name}' not in solver variables"
                )

            raw_nass = (

        balancer = PPBalancer(builder)
        result = balancer.balance_with_ilp()

        if not result.is_feasible:
            details = result.infeasibility_details
            msg = f"PP optimization failed: {details.get('reason', 'unknown')}"
            if details.get("error"):
                msg += f" ({details['error']})"
            raise RuntimeError(msg)

        if yaml_config.enable_simulation:
            self._run_simulation(
                result,

                yaml_config.sim_comm_time,
            )

        if result.simulation_status == "failed":
            logging.getLogger(__name__).warning(
                "ILP succeeded but simulation failed: %s. "
                "Pipeline bubble and step time estimates are unavailable.",
                result.simulation_error,
            )

                "Pipeline bubble and step time estimates are unavailable.",
                result.simulation_error,
            )
        elif result.simulation_status == "not_run":
            logging.getLogger(__name__).info(
                "Simulation skipped (enable_simulation=False). "
                "Pipeline bubble and step time estimates are not available."
            )

        sim_result = pp_sim.simulate_from_ilp(sim_comm_time=sim_comm_time)

        if sim_result is None:
            result.simulation_status = "failed"
            result.simulation_error = (
                "ILP simulation returned None (e.g. micro_batch_num < pp_degree)"
            )
            return

        real_bubble_val = sim_result.simulator_bubbles.get("real", 0.0)
        if not math.isfinite(sim_result.simulator_end_time) or not math.isfinite(real_bubble_val):
            result.simulation_status = "failed"

        result.simulator_peak_memory = sim_result.simulator_peak_memory

        pipeline_bubble = sim_result.simulator_bubbles.get("real")
        if pipeline_bubble is None:
            logging.getLogger(__name__).warning(
                "Simulator bubbles dict missing 'real' key. "
                "Available keys: %s",
                list(sim_result.simulator_bubbles.keys()),
            )

            >>> result.simulator_end_time
            1234.5
        """
        if not SAPP_PPB_AVAILABLE or self._pipeline is None:
            return None

        if self._yaml_config.micro_batch_num < self._yaml_config.pp_degree:
            logger.warning(
                "micro_batch_num (%d) < pp_degree (%d); simulator skipped.",
                self._yaml_config.micro_batch_num, self._yaml_config.pp_degree,
            )
            return None

        try:
            end_time = self._pipeline.simulate(
                show=False, comm_time=sim_comm_time,

                show=False, comm_time=sim_comm_time,
            )

            if end_time is None or end_time <= 0:
                return None

            sim_instance = self._pipeline.simulator
            return PPBOutput(
                simulation_status="success",

                simulator_end_time=sim_instance.end_time,
                simulator_bubbles=dict(sim_instance.bubbles),
                simulator_peak_memory=list(sim_instance.peak_memory),
            )
        except ValueError as exc:
            logger.warning("Simulator failed with ValueError: %s", exc)
            return None
        except Exception as exc:
            from hyper_parallel.auto_parallel.sapp_ppb.simulator.causal_error import CausalCommError, CausalError  # pylint: disable=C0415
            if isinstance(exc, (CausalCommError, CausalError)):
                logger.warning("Simulator failed with %s: %s", type(exc).__name__, exc)
                return None
            raise

            self,
            each_layer_per_recompute: Dict[Layer, Dict[Recompute.TYPE, List[List[int]]]],
            interleave_num: int = 1) -> List[List[float]]:
        """Return the per-stage backward time for a user-supplied layer assignment."""
        time = []
        for i in range(interleave_num):
            time.append([])
            for s in range(self.num_of_stage_):
                time[i].append(0)
                for layer in self.layers_sorted_[Layer.type_enum.BODY]:
                    for r in Recompute.TYPE:
                        if (r not in Recompute.get_unused_list(each_layer_per_recompute[layer])
                            and each_layer_per_recompute[layer][r][i][s] > 0):
                            time[i][s] += each_layer_per_recompute[layer][r][i][s] * (
                                layer.backward_time_rec_[r])
        for head in self.layers_sorted_[Layer.type_enum.HEAD]:
            time[0][0] += head.backward_time_rec_[Recompute.TYPE.NONE]
        for tail in self.layers_sorted_[Layer.type_enum.TAIL]:
            time[interleave_num - 1][self.num_of_stage_ - 1] += (
                tail.backward_time_rec_[Recompute.TYPE.NONE]
            )
        return time

        Unlike :meth:`get_simulator_forward_time`, this returns the actual
        backward time derived from per-layer ``backward_time_rec_`` instead of
        relying on ``forward_time × backward_ratio``.
        """
        time = []
        for i in range(self.num_of_interleave_):
            time.append([])
            for s in range(self.num_of_stage_):
                time[i].append(0)
                for layer in self.layers_sorted_[Layer.type_enum.BODY]:
                    for rec in Recompute.TYPE:
                        if self.recompute_considered_[rec]:
                            time[i][s] += self.variables_[layer.name_][rec][i][
                                s].varValue * layer.backward_time_rec_[rec]
        for head in self.layers_sorted_[Layer.type_enum.HEAD]:
            time[0][0] += head.backward_time_rec_[Recompute.TYPE.NONE]
        for tail in self.layers_sorted_[Layer.type_enum.TAIL]:
            time[self.num_of_interleave_ - 1][self.num_of_stage_ -
                                              1] += tail.backward_time_rec_[Recompute.TYPE.NONE]
        return time

        if isinstance(backward_time, (int, float)) and backward_time == 0:
            self._provided_backward_time = None
        else:
            self._provided_backward_time = format_2d_inputs(backward_time, self.vp, self.pp)

    def _statistic_init(self) -> None:
        r"""init statistic info"""
        self.forward_time = self.block_time

    def _statistic_init(self) -> None:
        r"""init statistic info"""
        self.forward_time = self.block_time
        if self._provided_backward_time is not None:
            self.backward_time = self._provided_backward_time
        else:
            self.backward_time = self.block_time * self.backward_ratio + self.layer_recompute
        self.states = {'last_time': np.zeros(self.pp),
                       'warmup_time': np.zeros(self.pp),

    def _process_swap_gap3(self, block, lines, p, b, i_b):
        r"""process swap when gap == 3."""
        if p % 2 == 0 and lines[p][i_b + 1].type == 'r' and lines[p][i_b + 2].type == 's':
            lines[p][i_b + 1], lines[p][i_b + 2] = lines[p][i_b + 2], lines[p][i_b + 1]
        if p % 2 == 1 and lines[p][i_b + 1].type == 's' and lines[p][i_b + 2].type == 'r':
            if block.phase == 'warmup' and self.blocks[p][b + 1].phase == 'cooldown':
                return False
            lines[p][i_b + 1], lines[p][i_b + 2] = lines[p][i_b + 2], lines[p][i_b + 1]
        if lines[p][i_b + 1].dual.stage == lines[p][i_b + 2].dual.stage:
            pd = lines[p][i_b + 1].dual.stage
            j_b1 = lines[pd].index(lines[p][i_b + 1].dual)

        return True

    def _process_swap_gap4(self, lines, p, i_b):
        r"""process swap when gap == 4."""
        if lines[p][i_b + 1].dual.stage == lines[p][i_b + 2].dual.stage and \
            lines[p][i_b + 2].dual.stage == lines[p][i_b + 3].dual.stage:
            if lines[p][i_b + 1].type == 's' and lines[p][i_b + 2].type == 's' \
                and lines[p][i_b + 3].type == 'r':
                lines[p][i_b + 1], lines[p][i_b + 2] = lines[p][i_b + 2], lines[p][i_b + 1]

    def _process_swap(self, block, lines, p, b, i_b, i_bn) -> bool:
        r"""process swap in condition"""
        if i_bn - i_b == 3:

        r"""process swap in condition"""
        if i_bn - i_b == 3:
            return self._process_swap_gap3(block, lines, p, b, i_b)
        if i_bn - i_b == 4:
            self._process_swap_gap4(lines, p, i_b)
        return True

    def swap_send_rec(self, lines: list[list[BlockSim]]) -> list[list[BlockSim]]:
        """Adjust send blocks: swap adjacent send/receive pairs where ordering is ambiguous."""