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}
            synthetic_rec_op = False
            if not hasattr(ccfg, 'rec_op'):

                provided, other per-layer parameters are ignored.
        """
        if layers is not None:
            if not layers:
                raise ValueError(
                    "For LayerRecomputeModel, layers should not be empty."
                )
            self._layers = layers
            self._num_layer = len(layers)

            }
            self._homogeneous = False
        else:
            if num_layer <= 0:
                raise ValueError(
                    f"For LayerRecomputeModel, num_layer should be positive, got {num_layer}."
                )
            self._layers = None
            self._num_layer = num_layer

    @property
    def num_layer(self) -> int:
        """Return the number of body layers."""
        return self._num_layer

    def _get_layer_can_recompute(self, layer_id: int) -> bool:
        """Get whether a layer supports recompute."""
        if not self._homogeneous:

        feasible_results = [r for r in results if r.is_feasible]

        if not feasible_results:
            raise RuntimeError(
                f"No feasible PP strategy found among {len(results)} candidate(s). "
                "All strategies exceed the memory limit."
            )

            )

        if optimize_for == "throughput":
            feasible_results.sort(key=lambda r: r.estimated_step_time)
        elif optimize_for == "memory":
            feasible_results.sort(
                key=lambda r: max(
                    s.total_memory if s.total_memory is not None else 0.0
                    for s in r.stages
                )

                "Please provide a valid dry-run YAML path."
            )

        if num_layer <= 0:
            raise ValueError(
                "PSStrategySearcher requires positive num_layer."
            )

        self.num_layer = num_layer

            >>> result = searcher.evaluate_strategy(strategy)
        """
        validation_result = self._validate_strategy(strategy)
        if not validation_result[0]:
            return PPStrategyResult(
                strategy=strategy,
                is_feasible=False,
                infeasibility_reason=validation_result[1],
            )

                infeasibility_reason=validation_result[1],
            )

        if not strategy.stage_partition:
            partition = StagePartition(self.num_layers, strategy.pp_degree)
            strategy.stage_partition = partition.uniform_partition()

        stages = self._build_stage_info_from_ppb(
            strategy, ppb_output,
        )

            Tuple of (is_feasible, infeasibility_reason).
        """
        if ppb_output is not None:
            if not ppb_output.is_feasible:
                return False, str(ppb_output.infeasibility_details)
            return True, ""

        if memory_limit is not None:
            for stage in stages:
                if (
                    stage.total_memory is not None
                    and stage.total_memory > memory_limit
                ):
                    return False, (
                        f"Stage {stage.stage_id} memory "
                        f"({stage.total_memory:.1f} MB) "
                        f"exceeds limit ({memory_limit:.1f} MB)."
                    )

        if ppb_output is not None and ppb_output.stage_compute_cost:
            compute_costs = ppb_output.stage_compute_cost
            if not ppb_output.stage_memory_cost:
                raise ValueError(
                    "ppb_output.stage_memory_cost is required but empty. "
                    "Ensure balance_with_ilp() produces memory cost output."
                )
            memory_costs = ppb_output.stage_memory_cost

        else:
            for stage_id, layer_ids in enumerate(strategy.stage_partition):
                recompute_layers = set()
                if strategy.layer_recompute_config:
                    recompute_layers = (
                        strategy.layer_recompute_config.recompute_layers & set(layer_ids)
                    )

                stages.append(StageInfo(

            >>> searcher = PSStrategySearcher(layers, dry_run_path="dry_run.yaml")
            >>> is_valid, msg = searcher._validate_strategy(strategy)
        """
        if strategy.pp_degree <= 0:
            return False, f"pp_degree must be positive, got {strategy.pp_degree}."

        if strategy.micro_batch_num <= 0:
            return False, f"micro_batch_num must be positive, got {strategy.micro_batch_num}."

        if strategy.pp_degree > self.num_layers:
            return (
                False,
                f"pp_degree ({strategy.pp_degree}) exceeds number of layers ({self.num_layers}).",
            )

        if strategy.stage_partition:
            partition = StagePartition(self.num_layers, strategy.pp_degree)
            is_valid, error_msg = partition.validate_partition(strategy.stage_partition)
            if not is_valid:
                return False, error_msg

        if strategy.layer_recompute_config:
            is_valid, error_msg = self._recompute_model.validate_config(
                strategy.layer_recompute_config

            is_valid, error_msg = self._recompute_model.validate_config(
                strategy.layer_recompute_config
            )
            if not is_valid:
                return False, error_msg

        return True, ""

    def _generate_visualization_data(

        results = []

        for pp_degree in pp_degrees:
            if pp_degree > self.num_layers:
                continue

            for vpp in vpps:
                for micro_batch_num in micro_batch_nums:
                    ppb_input = PPBInput(

            ValueError: If num_layer is not positive, or if pp_degree /
                micro_batch_num are invalid.
        """
        if ppb_input.num_layer <= 0:
            raise ValueError(
                "For PPBLoadBalancer, ppb_input.num_layer must be positive."
            )

        if ppb_input.pp_degree <= 0:

        if SAPP_PPB_AVAILABLE:
            try:
                self._sapp_ppb_adapter = SappPPBAdapter(ppb_input)
            except ImportError:
                self._sapp_ppb_adapter = None
            except ValueError:  # pylint: disable=W0706
                raise

            This method uses ILP (Integer Linear Programming) solver exclusively.
            No fallback to greedy/dp algorithms.
        """
        if not SAPP_PPB_AVAILABLE or self._sapp_ppb_adapter is None:
            raise ImportError(
                "sapp-ppb module is required for ILP-based load balancing. "
                "Please ensure sapp-ppb is installed and accessible."
            )

    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.simulator.pp_simulator import PipelineSimulator
    SAPP_PPB_AVAILABLE = True
except ImportError:
    SAPP_PPB_AVAILABLE = False
    SappPipeline = None
    Recompute = None
    PipelineSimulator = None


@dataclass
class YamlConstraints:

    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 _build_backward_time_rec(

            ImportError: If sapp-ppb module is not available.
            ValueError: If ppb_input.dry_run_path is empty or required fields are missing.
        """
        if not SAPP_PPB_AVAILABLE:
            raise ImportError(
                "sapp-ppb module is not available. "
                "Please ensure sapp-ppb is installed and accessible."
            )

                "Please provide a valid dry-run YAML path."
            )

        if ppb_input.num_layer <= 0:
            raise ValueError(
                "ppb_input.num_layer must be positive."
            )

        self.ppb_input = ppb_input

            pipeline_cfg = cfg.get("pipeline_config", {})
            num_layer = pipeline_cfg.get("num_layer", 0)
            if num_layer <= 0:
                raise ValueError(
                    "pipeline_config.num_layer is required when using "
                    "memory_parameter_config"
                )
        else:

            Integer PP degree used during dry-run, or 0 if not specified.
        """
        if "pipeline_num" in pipeline_cfg:
            return int(pipeline_cfg["pipeline_num"])
        if "pipeline_num_range" in pipeline_cfg:
            val = pipeline_cfg["pipeline_num_range"]
            if isinstance(val, list) and len(val) > 0:
                return int(val[0])
            if isinstance(val, (int, float)):
                return int(val)
        return 0

    @staticmethod
    def _make_default_constraints(
        num_of_interleave: int, target_pp: int, layer_per_stage_uniform: int,

            if val is None:
                continue
            if isinstance(val, bool):
                if val:
                    return True
            elif isinstance(val, (int, float)) and val != 0:
                return True
            elif isinstance(val, list):
                flat_vals: List[int] = []
                for item in val:
                    if isinstance(item, list):

                    yaml_format[yaml_key] = val
                else:
                    yaml_format[yaml_key] = [list(val) for _ in range(num_of_interleave)]
            else:
                yaml_format[yaml_key] = val
        return yaml_format

    def _parse_yaml_constraints(self, num_of_interleave: int) -> YamlConstraints:
        """Parse the dry-run YAML for user-specified offset/recompute values.

            ImportError: If sapp-ppb module is not available.
            ValueError: If YAML parsing or memory computation fails.
        """
        if not SAPP_PPB_AVAILABLE:
            raise ImportError("sapp-ppb module is not available")

        pipeline_layer = _get_pipeline_layer_class()
        if pipeline_layer is None:
            raise ImportError("Cannot resolve the Layer class used by SappPipeline")

        comp_mem, mem_par, mem_head, mem_tail, mem_act, _ = self._parse_dry_run_memory()
        head_time, body_time, tail_time = self._parse_dry_run_timing()

            PPBOutput with is_feasible=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 all computed data.
        """
        if layer_offset is None:
            layer_offset = []
        if layer_recompute is None:
            layer_recompute = self.ppb_input.layer_recompute_config

        limitations = []
        if not comm_cost_in_objective:
            limitations.append(

            limitations.append(
                "Memory parameters from memory_parameter_config (ComputeMemory decomposition bypassed)"
            )
        else:
            limitations.append(
                "Layer grouping driven by dry-run ComputeMemory (HEAD/BODY/TAIL)"
            )

        return PPBOutput(

            Cumulative per-interleave per-stage counts.
        """
        zeros_2d = [[0] * pp for _ in range(num_of_interleave)]
        if rec not in (Recompute.TYPE.SLCT, Recompute.TYPE.COMM):
            return [list(inter_exclusive) for inter_exclusive in exclusive]
        full_counts = constraints.recompute_per_type_per_stage.get(
            Recompute.TYPE.FULL, zeros_2d,
        )
        both_counts = constraints.recompute_per_type_per_stage.get(

            lay for lay in self.layers_sapp_ppb
            if lay.type_ == _get_pipeline_layer_class().type_enum.BODY
        ]
        if not body_layers:
            raise ValueError(
                "_compute_pre_opt_baseline requires BODY layers in the pipeline. "
                "Ensure the dry-run YAML contains at least one BODY layer group."
            )

        """
        from pulp import LpStatusInfeasible, LpStatusUndefined  # pylint: disable=C0415

        if not hasattr(self._pipeline, 'problem_'):
            return None
        pulp_problem = getattr(self._pipeline.problem_, 'problem_', None)
        if pulp_problem and hasattr(pulp_problem, 'status'):
            if pulp_problem.status in [LpStatusInfeasible, LpStatusUndefined]:
                return self._make_infeasible_output(

            >>> adapter = SappPPBAdapter(ppb_input)
            >>> output = adapter.balance_with_ilp(time_limit=60)
        """
        if not SAPP_PPB_AVAILABLE:
            raise ImportError("sapp-ppb module is not available")

        if not self.ppb_input.memory_limit:
            raise ValueError(
                "memory_limit is required for ILP load balancing. "
                "Please specify memory_limit in PPBInput."
            )

        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 valid partition from ILP solution",
                error=str(e),
            )

        """
        stage_partition = [[] for _ in range(self.ppb_input.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_

        body_layer_ids = list(range(1, self.ppb_input.num_layer + 1))

                stage_partition,
                allow_empty_stages=self.ppb_input.allow_empty_stages,
            )
            if not is_valid:
                raise RuntimeError(f"Invalid partition: {error_msg}")
        except (RuntimeError, ValueError) as e:
            raise RuntimeError(
                f"Invalid partition extracted from ILP solver: {str(e)}"
            ) from e

        return stage_partition

                        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
                stage_counts[stage_id] += total_count

        return stage_counts

        Returns:
            List of offsets per stage boundary (length = pp_degree - 1).
        """
        if not stage_partition or len(stage_partition) <= 1:
            return []

        uniform = StagePartition(
            self.ppb_input.num_layer + 2, self.ppb_input.pp_degree,
        ).uniform_partition()

                for inter in range(num_interleave):
                    try:
                        var_value = solver.variables_[group_name][rec][inter][stage_id].varValue
                        total += round(var_value) if var_value is not None else 0
                    except (KeyError, AttributeError):
                        pass
                if total > 0:
                    rec_counts[rec.name] = total
            if rec_counts:
                stage_detail[stage_id] = rec_counts

            RecomputeMode,
        )

        if self._pipeline is None or self._pipeline.problem_ is None:
            raise RuntimeError("Cannot extract recompute from ILP: pipeline or problem is None")

        solver = self._pipeline.problem_
        if not hasattr(solver, 'variables_') or not solver.variables_:
            raise RuntimeError("Cannot extract recompute from ILP: solver has no variables")

        recompute_type_map = {
            Recompute.TYPE.NONE: RecomputeMode.NONE,
            Recompute.TYPE.SLCT: RecomputeMode.SLCT,

        }

        group_name = "layer_group_1"
        if group_name not in solver.variables_:
            raise RuntimeError(f"Cannot extract recompute from ILP: '{group_name}' not in solver variables")

        stage_detail = self._extract_per_stage_recompute(
            solver, group_name,
            self._pipeline.num_of_interleave_,

            self.ppb_input.pp_degree,
        )

        if not stage_detail:
            raise RuntimeError("Cannot extract recompute from ILP: stage_detail is empty after extraction")

        stage_recompute_modes: Dict[int, RecomputeMode] = {}
        global_max_count = 0
        global_dominant = RecomputeMode.NONE

            >>> adapter = SappPPBAdapter(ppb_input)
            >>> score = adapter._calculate_imbalance_score([100, 110, 105])
        """
        if not costs or all(not c for c in costs):
            return 0.0

        avg_cost = sum(costs) / len(costs)
        if not avg_cost:
            return 0.0

        max_cost = max(costs)
        return max_cost / avg_cost - 1.0

        Returns:
            :class:`SimulatorResult` or ``None`` if simulation fails.
        """
        if not SAPP_PPB_AVAILABLE or self._pipeline is None:
            return None

        if self.ppb_input.micro_batch_num < self.ppb_input.pp_degree:
            return None

        try:
            if not fw_time_2d or not fw_time_2d[0]:
                return None

            num_interleave = self._pipeline.num_of_interleave_

            if num_interleave == 1:

                layer_recompute_input = list(rec_time_2d[0])
                block_mem_input = list(mem_act_2d[0])
                block_mem_par_input = list(mem_par_2d[0])
            else:
                block_time_input = [list(fw_time_2d[inter]) for inter in range(num_interleave)]
                layer_recompute_input = [list(rec_time_2d[inter]) for inter in range(num_interleave)]
                block_mem_input = [list(mem_act_2d[inter]) for inter in range(num_interleave)]
                block_mem_par_input = [list(mem_par_2d[inter]) for inter in range(num_interleave)]

            use_comm = comm_time > 0.0

            simulator = PipelineSimulator(

                end_time=simulator.end_time,
                bubbles=dict(simulator.bubbles),
                peak_memory=list(simulator.peak_memory),
            )
        except ValueError:
            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)):
                return None
            raise

    def simulate_from_ilp(
        self,
        comm_time: float = 0.0,

            >>> sim.end_time
            1234.5
        """
        if not SAPP_PPB_AVAILABLE or self._pipeline is None:
            return None

        try:
            fw_time = self._pipeline.get_fw_time()
            rec_time = self._pipeline.get_recompute_time()

            mem_act = self._pipeline.get_memory_activation()
            mem_par = self._pipeline.get_memory_parameter()

            if not fw_time or not fw_time[0]:
                return None

            return self._run_simulator(
                fw_time_2d=fw_time,
                rec_time_2d=rec_time,

                mem_act_2d=mem_act,
                mem_par_2d=mem_par,
                comm_time=comm_time,
            )
        except ValueError:
            return None

    def simulate_pre_opt_baseline(
        self,
        comm_time: float = 0.0,

            or self._pre_opt_rec_time_2d is None
            or self._pre_opt_mem_act_2d is None
            or self._pre_opt_mem_par_2d is None
        ):
            return None

        return self._run_simulator(
            fw_time_2d=self._pre_opt_fw_time_2d,
            rec_time_2d=self._pre_opt_rec_time_2d,

            num_layers: Total number of layers (HEAD + body layers + TAIL).
            pp_degree: Number of pipeline stages.
        """
        if pp_degree <= 0:
            raise ValueError(
                f"For StagePartition, pp_degree should be positive, but got {pp_degree}."
            )
        if num_layers <= 0:
            raise ValueError(
                f"For StagePartition, num_layers should be positive, but got {num_layers}."
            )
        if pp_degree > num_layers:
            raise ValueError(

            >>> partition = StagePartition(num_layers=4, pp_degree=2)
            >>> is_valid, msg = partition.validate_partition([[0, 1], [2, 3]])
        """
        if len(stage_partition) != self.pp_degree:
            return (
                False,
                f"Expected {self.pp_degree} stages, got {len(stage_partition)}.",
            )

        all_layer_ids = set()
        for stage_id, stage_layers in enumerate(stage_partition):
            if not stage_layers and not allow_empty_stages:
                return False, f"Stage {stage_id} is empty."

            for layer_id in stage_layers:
                if layer_id < 0 or layer_id >= self.num_layers:
                    return (
                        False,
                        f"Invalid layer ID {layer_id} in stage {stage_id}.",
                    )
                if layer_id in all_layer_ids:
                    return (
                        False,
                        f"Layer {layer_id} appears in multiple stages.",
                    )
                all_layer_ids.add(layer_id)

            )

        for stage_id in range(self.pp_degree - 1):
            if not stage_partition[stage_id] or not stage_partition[stage_id + 1]:
                continue
            last_layer_in_stage = max(stage_partition[stage_id])
            first_layer_in_next = min(stage_partition[stage_id + 1])
            if last_layer_in_stage >= first_layer_in_next:
                return (
                    False,
                    f"Stage {stage_id} and {stage_id + 1} are not ordered correctly.",
                )

    if range_key in cfg:
        val = cfg[range_key]
        if isinstance(val, list):
            if len(val) != 2:
                raise ValueError(
                    f"{range_key} must be a 2-element list [min, max], got {val}"
                )
            lo, hi = int(val[0]), int(val[1])
        elif isinstance(val, (int, float)):
            v = int(val)
            lo, hi = v, v
        else:
            raise ValueError(
                f"{range_key} must be a list or int, got {type(val).__name__}"
            )
    elif single_key in cfg:
        v = int(cfg[single_key])

        lo, hi = v, v
    elif default is not None:
        lo, hi = default, default
    else:
        raise ValueError(
            f"YAML pipeline_config must specify either '{range_key}' "
            f"or '{single_key}'.  Neither was found."
        )

            f"or '{single_key}'.  Neither was found."
        )

    if lo <= 0 or hi <= 0:
        raise ValueError(
            f"{field_name} values must be positive, got range [{lo}, {hi}]"
        )
    if lo > hi:
        raise ValueError(
            f"{field_name} min ({lo}) must not exceed max ({hi})"
        )

    return (lo, hi)

        ValueError: If required fields are missing or values are invalid.
    """
    num_layer = pipeline_cfg.get("num_layer")
    if num_layer is None:
        raise ValueError("YAML pipeline_config.num_layer is required")
    num_layer = int(num_layer)
    if num_layer <= 0:
        raise ValueError(f"pipeline_config.num_layer must be positive, got {num_layer}")

    num_of_interleave_range = _resolve_range(
        pipeline_cfg, "num_of_interleave_range", "num_of_interleave",
        "num_of_interleave", default=1,

    )

    memory_limit = pipeline_cfg.get("memory_limit")
    if memory_limit is None:
        raise ValueError("YAML pipeline_config.memory_limit is required")
    memory_limit = float(memory_limit)
    if memory_limit <= 0:
        raise ValueError(f"pipeline_config.memory_limit must be positive, got {memory_limit}")

    comm_time = float(pipeline_cfg.get("comm_time", 0.0))
    if comm_time < 0.0:
        raise ValueError(f"pipeline_config.comm_time must be non-negative, got {comm_time}")

    constant_memory = int(pipeline_cfg.get("constant_memory", 0))
    if constant_memory < 0:
        raise ValueError(f"pipeline_config.constant_memory must be non-negative, got {constant_memory}")

    with open(yaml_path, encoding="utf-8") as fp:
        cfg = yaml.safe_load(fp)

    if not isinstance(cfg, dict):
        raise ValueError(
            f"YAML file {yaml_path} must contain a top-level mapping, "
            f"got {type(cfg).__name__}"
        )

        )

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

    pipeline_num_single = pipeline_cfg.get("pipeline_num")

        Raises:
            ValueError: If any parameter is invalid.
        """
        if pipeline_num_range[0] <= 0 or pipeline_num_range[1] <= 0:
            raise ValueError(
                f"pipeline_num_range values must be positive, got {pipeline_num_range}"
            )
        if pipeline_num_range[0] > pipeline_num_range[1]:
            raise ValueError(

            raise ValueError(
                f"pipeline_num_range min must not exceed max, got {pipeline_num_range}"
            )
        if micro_batch_num_range[0] <= 0 or micro_batch_num_range[1] <= 0:
            raise ValueError(
                f"micro_batch_num_range values must be positive, got {micro_batch_num_range}"
            )
        if micro_batch_num_range[0] > micro_batch_num_range[1]:
            raise ValueError(
                f"micro_batch_num_range min must not exceed max, got {micro_batch_num_range}"
            )
        if memory_limit <= 0:
            raise ValueError(

            raise ValueError(
                f"memory_limit must be positive, got {memory_limit}"
            )
        if num_of_interleave <= 0:
            raise ValueError(
                f"num_of_interleave must be positive, got {num_of_interleave}"
            )

    def _resolve_num_layer(self, num_layer: Optional[int]) -> int:

        Returns:
            Resolved number of body layers.
        """
        if num_layer is not None:
            return int(num_layer)
        if self._profiled_layers:
            body_layers = [
                l for l in self._profiled_layers
                if l.layer_type not in ("head", "tail", "embedding", "output")