Diff Coverage

def __getattr__(name):  # pylint: disable=C0103
    if name not in _EXPORTS:
        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
    module = _import_module(_EXPORTS[name], __name__)
    value = getattr(module, name)
    globals()[name] = value
    return value

            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'):

                    desc["name"] += "_" + target["name"]
                    desc["memory_parameter"] += target["memory_parameter"]
                    desc["memory_activation"]["NONE"] += target["memory_activation"]["NONE"]
                    desc["memory_activation"]["COMM"] += target["memory_activation"]["COMM"]
                    desc["memory_activation"]["SLCT"] += target["memory_activation"]["SLCT"]
                    desc["memory_activation"]["BOTH"] += target["memory_activation"]["BOTH"]
                    desc["memory_activation"]["FULL"] += target["memory_activation"]["FULL"]
                    target_idx = ppb_lay_desc.index(target)
                    idx = target_idx if idx < 0 else min(idx, target_idx)
                    del ppb_lay_desc[target_idx]
            idx = max(idx, 0)

        if user_global_batch_size is not None:
            calculated_gbs = self.global_batch_size()
            if user_global_batch_size != calculated_gbs:
                dp = self.dims_val.get(DP, 1)
                mbs = self.dims_val.get(MBS, 1)
                mbn = self.dims_val.get(MBN, 1)
                logger.warning(
                    "Batch size constraint violated: user_global_batch_size (%d) != "
                    "DP (%d) * MBS (%d) * MBN (%d) = %d",
                    user_global_batch_size, dp, mbs, mbn, calculated_gbs
                )
                return False

        return True

    def val(self, dim: Dimension) -> int:

                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_act(self, layer_id: int, mode: RecomputeMode) -> float:
        """Get activation memory for a layer under a recompute mode."""
        if not self._homogeneous:

            layer = self._layers[self._layer_id_to_idx[layer_id]]
            if mode == RecomputeMode.SLCT:
                return layer.activation_memory_slct
            if mode == RecomputeMode.COMM:
                return layer.activation_memory_comm
            if mode == RecomputeMode.BOTH:
                return layer.activation_memory_both
            if mode == RecomputeMode.FULL:
                return layer.activation_memory_full
            return layer.activation_memory

        if mode == RecomputeMode.SLCT:
            return self._activation_memory_slct
        if mode == RecomputeMode.COMM:
            return self._activation_memory_comm
        if mode == RecomputeMode.BOTH:
            return self._activation_memory_both
        if mode == RecomputeMode.FULL:
            return self._activation_memory_full
        return self._activation_memory

    def _get_layer_act_none(self, layer_id: int) -> float:
        """Get activation memory without recompute for a layer."""
        if not self._homogeneous:

    def _get_layer_act_none(self, layer_id: int) -> float:
        """Get activation memory without recompute for a layer."""
        if not self._homogeneous:
            return self._layers[self._layer_id_to_idx[layer_id]].activation_memory
        return self._activation_memory

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

    def _get_layer_forward_time(self, layer_id: int) -> float:
        """Get forward time for a layer."""
        if not self._homogeneous:
            return self._layers[self._layer_id_to_idx[layer_id]].forward_time
        return self._forward_time

    def _get_layer_recompute_forward_overhead(self, layer_id: int) -> float:
        """Get recompute forward overhead for a layer."""
        if not self._homogeneous:

    def _get_layer_recompute_forward_overhead(self, layer_id: int) -> float:
        """Get recompute forward overhead for a layer."""
        if not self._homogeneous:
            return self._layers[self._layer_id_to_idx[layer_id]].recompute_forward_overhead
        return self._recompute_forward_overhead

    def validate_config(
        self,
        config: LayerRecomputeConfig,

            ... )
            >>> saving = model.estimate_memory_saving(config)
        """
        if not config.enabled or config.recompute_mode == RecomputeMode.NONE:
            return 0.0

        total_saving = 0.0

        for layer_id in config.recompute_layers:

        total_saving = 0.0

        for layer_id in config.recompute_layers:
            if layer_id not in self._layer_id_to_idx:
                continue

            act_none = self._get_layer_act_none(layer_id)
            act_mode = self._get_layer_act(layer_id, config.recompute_mode)
            saving = act_none - act_mode

            >>> model = LayerRecomputeModel(layers)
            >>> config = LayerRecomputeConfig(recompute_mode=RecomputeMode.FULL)
            >>> mem = model.get_layer_activation_memory(0, config)
        """
        if layer_id not in self._layer_id_to_idx:
            raise ValueError(
                f"For get_layer_activation_memory, layer_id {layer_id} is invalid."
            )

        if not config.enabled or config.recompute_mode == RecomputeMode.NONE:
            return self._get_layer_act_none(layer_id)

        if layer_id not in config.recompute_layers:
            return self._get_layer_act_none(layer_id)

        return self._get_layer_act(layer_id, config.recompute_mode)

    def generate_recompute_config_for_memory_target(
        self,
        memory_target: float,

            ...     memory_target=1000.0,
            ...     current_activation_memory=2000.0,
            ... )
        """
        if current_activation_memory <= memory_target:
            return LayerRecomputeConfig(enabled=False)

        memory_to_save = current_activation_memory - memory_target
        recompute_layers = set()

        layer_savings = []
        for layer_id in range(self._num_layer):
            act_none = self._get_layer_act_none(layer_id)
            act_mode = self._get_layer_act(layer_id, recompute_mode)
            saving = act_none - act_mode
            if saving > 0:
                layer_savings.append((layer_id, saving))

        if strategy == "greedy":
            layer_savings.sort(key=lambda x: x[1], reverse=True)
        else:
            layer_savings.sort(key=lambda x: x[1])

        accumulated_saving = 0.0
        for layer_id, saving in layer_savings:
            if accumulated_saving >= memory_to_save:
                break
            recompute_layers.add(layer_id)
            accumulated_saving += saving

        return LayerRecomputeConfig(
            recompute_layers=recompute_layers,
            recompute_mode=recompute_mode,
            enabled=True,
        )

                rec_fwd = self._get_layer_recompute_forward_overhead(layer_id)
                fwd = self._get_layer_forward_time(layer_id)
                overhead = rec_fwd if rec_fwd > 0.0 else fwd * 0.5
            else:
                overhead = 0.0

            total_overhead += max(0.0, overhead)

        return total_overhead

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

        if not feasible_results:
            if results:
                return results[0]
            raise RuntimeError("No feasible PP strategy found.")

        if optimize_for == "throughput":
            feasible_results.sort(key=lambda r: r.estimated_step_time)
        elif optimize_for == "memory":

        output_files = {}

        if output_format in ["png", "all"]:
            try:
                viz_data = self._visualizer.visualize_with_sapp_ppb(
                    result=result,
                    file_name=str(output_path / "pipeline_timeline.png"),
                    show=False,
                )
                if viz_data is not None:
                    output_files["pipeline_timeline"] = str(output_path / "pipeline_timeline.png")
            except (ImportError, ValueError, RuntimeError):
                pass

        if output_format in ["json", "all", "both"]:
            json_data = self._visualizer.export_to_json(result)
            json_file = output_path / "pp_strategy_visualization.json"

            >>> optimizer = PPOptimizer()
            >>> result = optimizer.optimize(dry_run_path="dry_run.yaml")
            >>> viz_data = optimizer.generate_pipeline_visualization(result, show=True)
        """
        return self._visualizer.generate_pipeline_timeline_plot(
            result=result,
            file_name=file_name,
            show=show,
        )

        num_layers = yaml_config.num_layer

        recommended_strategy = "greedy"
        if num_layers > 100:
            recommended_strategy = "dp"

        recommended_mode = "fast"
        if pp_range[1] - pp_range[0] <= 4 and num_layers <= 50:
            recommended_mode = "precise"

        recommendations = {
            "pp_degree_range": pp_range,
            "micro_batch_num_range": mb_range,

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

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

        self.num_layer = num_layer

                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,
        )

        infeasibility_reason = ""
        if ppb_output is not None:
            is_feasible = ppb_output.is_feasible
            if not is_feasible:
                infeasibility_reason = str(ppb_output.infeasibility_details)
        else:
            if memory_limit is not None:
                for stage in stages:
                    if (
                        stage.memory_breakdown
                        and stage.memory_breakdown.total_memory > memory_limit
                    ):
                        is_feasible = False
                        infeasibility_reason = (
                            f"Stage {stage.stage_id} memory "
                            f"({stage.memory_breakdown.total_memory:.1f} MB) "
                            f"exceeds limit ({memory_limit:.1f} MB)."
                        )
                        break

        pipeline_bubble = 0.0
        estimated_step_time = 0.0
        imbalance_score = 0.0

            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)
                    )

                total_mem = memory_costs[stage_id] if stage_id < len(memory_costs) else 0.0

                    recompute_layers=recompute_layers,
                    is_valid=True,
                ))
        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(
                    stage_id=stage_id,
                    layer_ids=layer_ids,
                    memory_breakdown=None,
                    performance_breakdown=None,

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

            ...     pp_degrees=[2, 4],
            ...     micro_batch_nums=[4, 8],
            ... )
        """
        if schedule_types is None:
            schedule_types = [PipelineScheduleType.ONE_F_ONE_B]

        results = []

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

            partition = StagePartition(self.num_layers, pp_degree)
            base_partition = partition.uniform_partition()

            for micro_batch_num in micro_batch_nums:
                for schedule_type in schedule_types:
                    strategy = PPStrategy(
                        pp_degree=pp_degree,
                        micro_batch_num=micro_batch_num,
                        stage_partition=base_partition,
                        schedule_type=schedule_type,

                        stage_partition=base_partition,
                        schedule_type=schedule_type,
                    )

                    result = self.evaluate_strategy(strategy, memory_limit)
                    results.append(result)

        results.sort(key=lambda r: (not r.is_feasible, r.estimated_step_time))
        return results

    def search_with_load_balancing(
        self,
        pp_degrees: List[int],

        results = []

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

            for micro_batch_num in micro_batch_nums:
                ppb_input = PPBInput(
                    num_layer=self.num_layer,

            ...     pp_degrees=[2, 4],
            ...     micro_batch_nums=[4, 8],
            ... )
        """
        results = self.search_with_load_balancing(
            pp_degrees=pp_degrees,
            micro_batch_nums=micro_batch_nums,
            memory_limit=memory_limit,
        )

            micro_batch_nums=micro_batch_nums,
            memory_limit=memory_limit,
        )

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

        if not feasible_results:
            return None

        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.memory_breakdown.total_memory if s.memory_breakdown else 0.0
                    for s in r.stages
                )

                    for s in r.stages
                )
            )
        else:
            feasible_results.sort(key=lambda r: r.estimated_step_time)

        return feasible_results[0]

    def search_with_adaptive_optimization(
        self,
        pp_degrees: List[int],

        results = []

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

            for micro_batch_num in micro_batch_nums:
                ppb_input = PPBInput(
                    num_layer=self.num_layer,

any active code path within ``pp_modeling`` since the dual-path estimator
architecture was replaced by the single ILP-path design.
"""

from __future__ import annotations

from typing import List, Dict, Any, Optional, Tuple

from hyper_parallel.auto_parallel.sapp_nd.pp_modeling.pp_types import (
    LayerInfo,
    LayerRecomputeConfig,
    RecomputeMode,
)

    RecomputeMode,
)


def recompute_mode_to_layer_type(
    mode: RecomputeMode,
) -> "LayerType":
    """Map a :class:`RecomputeMode` to the corresponding :class:`LayerType`.

    Returns:
        LayerType enum value understood by EvaluatorV2 and perf_estimation.
    """
    from hyper_parallel.auto_parallel.sapp_nd.nd.common.layer_type import LayerType  # pylint: disable=C0415

    if mode in (RecomputeMode.SLCT, RecomputeMode.COMM, RecomputeMode.BOTH):
        return LayerType.SEL_REC_LAYER
    if mode == RecomputeMode.FULL:
        return LayerType.FULL_REC_LAYER
    return LayerType.NOT_REC_LAYER


REC_OP_KEYS = ['attBMM', 'headCast', 'dropout', 'softmax', 'normOp', 'gather', 'ffAct']
REC_OP_OPERATOR_KEYS = ['attBMM', 'headCast', 'dropout', 'softmax', 'normOp', 'ffAct']


def apply_recompute_to_rec_op(ccfg: Any, mode: RecomputeMode) -> Dict[str, int]:
    """Set ``ccfg.rec_op`` flags according to *mode* and return the original values.

    For :class:`EvaluatorV2`, ``rec_op`` flags control which sub-operators
    are recomputed under ``SEL_REC_LAYER``:

    Returns:
        Dictionary of original ``rec_op`` values for later restoration.
    """
    original: Dict[str, int] = {}
    for key in REC_OP_KEYS:
        original[key] = getattr(ccfg.rec_op, key, 1)

    if mode == RecomputeMode.NONE:
        for key in REC_OP_KEYS:
            setattr(ccfg.rec_op, key, 1)
    elif mode == RecomputeMode.SLCT:
        for key in REC_OP_OPERATOR_KEYS:
            setattr(ccfg.rec_op, key, 0)
        setattr(ccfg.rec_op, 'gather', 1)
    elif mode == RecomputeMode.COMM:
        for key in REC_OP_OPERATOR_KEYS:
            setattr(ccfg.rec_op, key, 1)
        setattr(ccfg.rec_op, 'gather', 0)
    elif mode == RecomputeMode.BOTH:
        for key in REC_OP_KEYS:
            setattr(ccfg.rec_op, key, 0)

    return original


def restore_rec_op(ccfg: Any, original: Dict[str, int]) -> None:
    """Restore ``ccfg.rec_op`` flags from *original*.

    Args:
        ccfg: EvaluatorV2's ``_ccfg`` object.

    Args:
        ccfg: EvaluatorV2's ``_ccfg`` object.
        original: Original values returned by :func:`apply_recompute_to_rec_op`.
    """
    for key, val in original.items():
        setattr(ccfg.rec_op, key, val)


def _extract_body_ids(
    layer_ids: List[int],
    is_first_stage: bool,
    is_last_stage: bool,
    head_layer_id: Any,

    has_head_in_partition: bool,
    has_tail_in_partition: bool,
) -> List[int]:
    """Strip HEAD/TAIL layer IDs from *layer_ids* when they are already in the partition."""
    body_ids = list(layer_ids)
    if is_first_stage and has_head_in_partition and head_layer_id in body_ids:
        body_ids.remove(head_layer_id)
    if is_last_stage and has_tail_in_partition and tail_layer_id in body_ids:
        body_ids.remove(tail_layer_id)
    return body_ids


def _build_non_vpp_chunk(
    n_body: int,
    is_first_stage: bool,
    is_last_stage: bool,
    body_layer_type: "LayerType",

    is_last_stage: bool,
    body_layer_type: "LayerType",
) -> List[List["LayerType"]]:
    """Build a single-chunk stage (non-VPP)."""
    from hyper_parallel.auto_parallel.sapp_nd.nd.common.layer_type import LayerType  # pylint: disable=C0415
    chunk: List[LayerType] = []
    if is_first_stage:
        chunk.append(LayerType.EMBEDDING_LAYER)
    chunk.extend([body_layer_type] * n_body)
    if is_last_stage:
        chunk.append(LayerType.OUTPUT_LAYER)
    return [chunk]


def _build_vpp_chunks(
    n_body: int,
    is_first_stage: bool,
    is_last_stage: bool,
    num_model_chunks: int,

    num_model_chunks: int,
    body_layer_type: "LayerType",
) -> List[List["LayerType"]]:
    """Build multi-chunk stage (VPP) by splitting layers across virtual chunks."""
    from hyper_parallel.auto_parallel.sapp_nd.nd.common.layer_type import LayerType  # pylint: disable=C0415
    extra = int(is_first_stage) + int(is_last_stage)
    total_layers_in_stage = n_body + extra
    base = total_layers_in_stage // num_model_chunks
    remainder = total_layers_in_stage % num_model_chunks
    chunk_sizes = [base + (1 if c < remainder else 0) for c in range(num_model_chunks)]

    all_layers: List[LayerType] = []
    if is_first_stage:
        all_layers.append(LayerType.EMBEDDING_LAYER)
    all_layers.extend([body_layer_type] * n_body)
    if is_last_stage:
        all_layers.append(LayerType.OUTPUT_LAYER)

    stage_chunks: List[List[LayerType]] = []
    offset = 0
    for size in chunk_sizes:
        stage_chunks.append(all_layers[offset:offset + size])
        offset += size
    return stage_chunks


def _check_head_tail_in_partition(
    layers: List[LayerInfo],
    stage_partition: List[List[int]],
) -> Tuple[Optional[int], Optional[int], bool, bool]:
    """Return (head_layer_id, tail_layer_id, has_head, has_tail)."""
    head_layer_id = layers[0].layer_id if layers else None
    tail_layer_id = layers[-1].layer_id if len(layers) > 1 else None
    pp_degree = len(stage_partition)
    has_head = (
        head_layer_id is not None
        and pp_degree > 0
        and head_layer_id in stage_partition[0]
    )
    has_tail = (
        tail_layer_id is not None
        and pp_degree > 0
        and tail_layer_id in stage_partition[-1]
    )
    return head_layer_id, tail_layer_id, has_head, has_tail


def _validate_body_layer_count(
    stage_partition: List[List[int]],
    has_head: bool,
    has_tail: bool,
    evaluator_body_layers: int,

    has_tail: bool,
    evaluator_body_layers: int,
) -> None:
    """Raise ValueError if body-layer count does not match evaluator config."""
    total = sum(len(s) for s in stage_partition)
    if has_head:
        total -= 1
    if has_tail:
        total -= 1
    if total != evaluator_body_layers:
        raise ValueError(
            f"Stage partition has {total} body layers but the "
            f"model config defines {evaluator_body_layers}."
        )

            f"model config defines {evaluator_body_layers}."
        )


def _build_vpp_stage_chunks(
    body_layer_type: "LayerType",
    n_body: int,
    is_first_stage: bool,
    is_last_stage: bool,

    is_last_stage: bool,
    num_model_chunks: int,
) -> List[List["LayerType"]]:
    """Build VPP (interleaved) stage chunks for a single stage."""
    from hyper_parallel.auto_parallel.sapp_nd.nd.common.layer_type import LayerType  # pylint: disable=C0415

    extra = int(is_first_stage) + int(is_last_stage)
    total_layers_in_stage = n_body + extra
    base = total_layers_in_stage // num_model_chunks
    remainder = total_layers_in_stage % num_model_chunks
    chunk_sizes = [
        base + (1 if c < remainder else 0)
        for c in range(num_model_chunks)
    ]

    all_layers: List[LayerType] = []
    if is_first_stage:
        all_layers.append(LayerType.EMBEDDING_LAYER)
    all_layers.extend([body_layer_type] * n_body)
    if is_last_stage:
        all_layers.append(LayerType.OUTPUT_LAYER)

    stage_chunks: List[List[LayerType]] = []
    offset = 0
    for size in chunk_sizes:
        stage_chunks.append(all_layers[offset:offset + size])
        offset += size
    return stage_chunks


def build_stages_from_partition(
    layers: List[LayerInfo],
    evaluator_body_layers: int,
    stage_partition: List[List[int]],
    recompute_config: LayerRecomputeConfig,

    Raises:
        ValueError: If the total number of body layers in *stage_partition*
            does not match *evaluator_body_layers*.
    """
    from hyper_parallel.auto_parallel.sapp_nd.nd.common.layer_type import LayerType  # pylint: disable=C0415

    body_layer_type = recompute_mode_to_layer_type(recompute_config.recompute_mode)
    pp_degree = len(stage_partition)

    head_layer_id, tail_layer_id, has_head, has_tail = _check_head_tail_in_partition(
        layers, stage_partition
    )
    _validate_body_layer_count(stage_partition, has_head, has_tail, evaluator_body_layers)

    stages: List[List[List[LayerType]]] = []
    for stage_id, layer_ids in enumerate(stage_partition):
        is_first_stage = not stage_id
        is_last_stage = stage_id == pp_degree - 1

        body_ids = list(layer_ids)
        if is_first_stage and has_head and head_layer_id in body_ids:
            body_ids.remove(head_layer_id)
        if is_last_stage and has_tail and tail_layer_id in body_ids:
            body_ids.remove(tail_layer_id)

        n_body = len(body_ids)

        if num_model_chunks <= 1:
            chunk: List[LayerType] = []
            if is_first_stage:
                chunk.append(LayerType.EMBEDDING_LAYER)
            chunk.extend([body_layer_type] * n_body)
            if is_last_stage:
                chunk.append(LayerType.OUTPUT_LAYER)
            stages.append([chunk])
        else:
            stages.append(_build_vpp_stage_chunks(
                body_layer_type, n_body, is_first_stage, is_last_stage, num_model_chunks
            ))

    return stages

    from hyper_parallel.auto_parallel.sapp_nd.pp_modeling.sapp_ppb_adapter import (
        SappPPBAdapter,
        SAPP_PPB_AVAILABLE,
    )
except ImportError:
    SAPP_PPB_AVAILABLE = False
    SappPPBAdapter = None


class PPVisualizer:
    """Visualization utilities for Pipeline Parallelism.

            >>> visualizer = PPVisualizer()
            >>> chart = visualizer.generate_stage_load_chart(result)
        """
        if not result.stages:
            return {
                "chart_type": "bar",
                "title": "Stage Load Distribution",
                "x_axis": {"label": "Pipeline Stage", "data": []},
                "y_axes": [],

        for stage in result.stages:
            if stage.memory_breakdown:
                memory_data.append(stage.memory_breakdown.total_memory)
            else:
                memory_data.append(0.0)

            if stage.performance_breakdown:
                compute_data.append(stage.performance_breakdown.total_time)
            else:
                compute_data.append(0.0)

        return {
            "chart_type": "bar",
            "title": "Stage Load Distribution",

    def _build_gpipe_schedule(
        stages: List[int], micro_batch_num: int,
    ) -> List[Dict[str, Any]]:
        """Build GPipe schedule entries."""
        schedule = []
        for mb_id in range(micro_batch_num):
            for stage_id in stages:
                schedule.append({
                    "micro_batch": mb_id,
                    "stage": stage_id,
                    "type": "forward",
                    "time_step": mb_id * len(stages) + stage_id,

                    "stage": stage_id,
                    "type": "forward",
                    "time_step": mb_id * len(stages) + stage_id,
                })
        for mb_id in range(micro_batch_num):
            for stage_id in reversed(stages):
                schedule.append({
                    "micro_batch": mb_id,
                    "stage": stage_id,
                    "type": "backward",
                    "time_step": (micro_batch_num + mb_id) * len(stages) + stage_id,

                    "stage": stage_id,
                    "type": "backward",
                    "time_step": (micro_batch_num + mb_id) * len(stages) + stage_id,
                })
        return schedule

    @staticmethod
    def _build_1f1b_schedule(
        stages: List[int], micro_batch_num: int,

    def _build_interleaved_schedule(
        stages: List[int], micro_batch_num: int, num_model_chunks: int = 2,
    ) -> List[Dict[str, Any]]:
        """Build interleaved (VPP) schedule entries."""
        schedule = []
        for mb_id in range(micro_batch_num):
            for chunk_id in range(num_model_chunks):
                for stage_id in stages:
                    schedule.append({
                        "micro_batch": mb_id,
                        "stage": stage_id,
                        "type": "forward",
                        "chunk": chunk_id,

                        "type": "forward",
                        "chunk": chunk_id,
                        "time_step": (mb_id * num_model_chunks + chunk_id) * len(stages) + stage_id,
                    })
        for mb_id in range(micro_batch_num):
            for chunk_id in range(num_model_chunks):
                for stage_id in reversed(stages):
                    schedule.append({
                        "micro_batch": mb_id,
                        "stage": stage_id,
                        "type": "backward",
                        "chunk": chunk_id,

                            + mb_id * num_model_chunks
                            + chunk_id
                        ) * len(stages) + stage_id,
                    })
        return schedule

    def _build_schedule(
        self,
        stages: List[int],

            >>> visualizer = PPVisualizer()
            >>> schedule = visualizer._build_schedule([0, 1], 4, PipelineScheduleType.ONE_F_ONE_B)
        """
        if schedule_type == PipelineScheduleType.GPipe:
            return self._build_gpipe_schedule(stages, micro_batch_num)
        if schedule_type == PipelineScheduleType.ONE_F_ONE_B:
            return self._build_1f1b_schedule(stages, micro_batch_num)
        return self._build_interleaved_schedule(stages, micro_batch_num)

    def generate_memory_breakdown_chart(
        self,
        result: PPStrategyResult,

            >>> visualizer = PPVisualizer()
            >>> chart = visualizer.generate_memory_breakdown_chart(result)
        """
        if not result.stages:
            return {
                "chart_type": "stacked_bar",
                "title": "Memory Breakdown by Stage",
                "x_axis": {"label": "Pipeline Stage", "data": []},
                "y_axis": {"label": "Memory (MB)"},

                grad_data.append(stage.memory_breakdown.grad_memory)
                optimizer_data.append(stage.memory_breakdown.optimizer_state_memory)
                activation_data.append(stage.memory_breakdown.activation_memory)
            else:
                param_data.append(0.0)
                grad_data.append(0.0)
                optimizer_data.append(0.0)
                activation_data.append(0.0)

        return {
            "chart_type": "stacked_bar",
            "title": "Memory Breakdown by Stage",

            >>> visualizer = PPVisualizer()
            >>> chart = visualizer.generate_performance_breakdown_chart(result)
        """
        if not result.stages:
            return {
                "chart_type": "stacked_bar",
                "title": "Performance Breakdown by Stage",
                "x_axis": {"label": "Pipeline Stage", "data": []},
                "y_axis": {"label": "Time (ms)"},

                backward_data.append(stage.performance_breakdown.backward_time)
                comm_data.append(stage.performance_breakdown.communication_time)
                recompute_data.append(stage.performance_breakdown.recompute_overhead)
            else:
                forward_data.append(0.0)
                backward_data.append(0.0)
                comm_data.append(0.0)
                recompute_data.append(0.0)

        return {
            "chart_type": "stacked_bar",
            "title": "Performance Breakdown by Stage",

        bubble_stages = pp_degree - 1

        if not result.stages:
            ideal_step_time = 0.0
        else:
            ideal_step_time = micro_batch_num * max(
                s.performance_breakdown.total_time if s.performance_breakdown else 0.0
                for s in result.stages

        Example:
            >>> visualizer = PPVisualizer()
            >>> heatmap = visualizer.generate_imbalance_heatmap(result)
        """
        if not result.stages:
            return {
                "chart_type": "heatmap",
                "title": "Stage Load Imbalance Heatmap",
                "x_axis": {"label": "Stage", "data": []},
                "y_axis": {"label": "Stage", "data": []},

                "data": [],
                "colorscale": "Viridis",
            }

        num_stages = len(result.stages)
        heatmap_data = []

        for i, stage_i in enumerate(result.stages):
            row = []
            for j, stage_j in enumerate(result.stages):
                if i == j:
                    row.append(0.0)
                else:
                    mem_i = (
                        stage_i.memory_breakdown.total_memory
                        if stage_i.memory_breakdown
                        else 0.0
                    )
                    mem_j = (
                        stage_j.memory_breakdown.total_memory
                        if stage_j.memory_breakdown
                        else 0.0
                    )

                        if stage_j.memory_breakdown
                        else 0.0
                    )

                    if mem_i + mem_j > 0:
                        imbalance = abs(mem_i - mem_j) / (mem_i + mem_j)
                    else:
                        imbalance = 0.0
                    row.append(imbalance)
            heatmap_data.append(row)

        return {
            "chart_type": "heatmap",
            "title": "Stage Load Imbalance Heatmap",
            "x_axis": {"label": "Stage", "data": list(range(num_stages))},
            "y_axis": {"label": "Stage", "data": list(range(num_stages))},

            >>> visualizer = PPVisualizer()
            >>> viz_data = visualizer.visualize_with_sapp_ppb(
            ...     result, dry_run_path="dry_run.yaml", show=True)
        """
        if not self.use_sapp_ppb or not SAPP_PPB_AVAILABLE:
            return None

        if not dry_run_path:
            return None

        try:
            from hyper_parallel.auto_parallel.sapp_nd.pp_modeling.pp_types import PPBInput  # pylint: disable=C0415

            if not result.stages or not result.strategy.stage_partition:
                return None

            if not self._sapp_ppb_adapter:
                ppb_input = PPBInput(
                    num_layer=len(result.stages) * 2,
                    pp_degree=result.strategy.pp_degree,
                    micro_batch_num=result.strategy.micro_batch_num,
                    dry_run_path=dry_run_path,

                    pp_degree=result.strategy.pp_degree,
                    micro_batch_num=result.strategy.micro_batch_num,
                    dry_run_path=dry_run_path,
                )
                self._sapp_ppb_adapter = SappPPBAdapter(ppb_input)

            return self._sapp_ppb_adapter.visualize(
                result=result,
                file_name=file_name,
                show=show,
            )
        except (ImportError, ValueError, RuntimeError):
            return None

    def generate_pipeline_timeline_plot(
        self,
        result: PPStrategyResult,

            >>> visualizer = PPVisualizer()
            >>> data = visualizer.generate_pipeline_timeline_plot(
            ...     result, dry_run_path="dry_run.yaml")
        """
        if self.use_sapp_ppb:
            sapp_ppb_result = self.visualize_with_sapp_ppb(
                result=result,
                file_name=file_name,
                show=show,
                dry_run_path=dry_run_path,

                file_name=file_name,
                show=show,
                dry_run_path=dry_run_path,
            )
            if sapp_ppb_result is not None:
                return sapp_ppb_result

        basic_data = self.generate_pipeline_schedule_diagram(result)

        return basic_data

            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

    def validate_dimensions(

        Example:
            >>> balancer = PPBLoadBalancer(ppb_input)
            >>> is_valid, msg = balancer.validate_pp_degree(machine_device_num=8, tp=2)
        """
        pp = self.pp_degree

        devices_per_node = machine_device_num // tp
        if devices_per_node <= 0:
            return (
                False,
                f"Invalid device configuration: machine_device_num ({machine_device_num}) // TP ({tp}) <= 0.",
            )

        if pp > devices_per_node:
            return (
                False,
                f"PP ({pp}) exceeds maximum pipeline bound ({devices_per_node}).",
            )

        if num_layers is not None:
            max_pp_by_layers = num_layers // vpp
            if pp > max_pp_by_layers:
                return (
                    False,
                    f"PP ({pp}) exceeds layer count limit ({max_pp_by_layers} = {num_layers} // {vpp}).",
                )

        if devices_per_node % pp:
            return (
                False,
                f"PP ({pp}) must divide devices_per_node ({devices_per_node}).",
            )

        return True, ""

    def _ilp_balance(self) -> PPBOutput:
        """Create ILP-optimized load-balanced partition using sapp-ppb.