Diff Coverage

# limitations under the License.
# ============================================================================
"""Minimal PuLP-compatible objects for SAPP static tests when PuLP is absent."""
# pylint: disable=invalid-name
from __future__ import annotations

from collections.abc import Iterable
from numbers import Number
from typing import Any, Optional

LpContinuous = "Continuous"
LpInteger = "Integer"
LpBinary = "Binary"
LpMinimize = 1
_PULP_IMPORT_ERROR = (
    "PuLP is required to solve SAPP-PPB ILP problems. Install the optional "
    "'pulp' package before calling solve()."
)

    "'pulp' package before calling solve()."
)


def _to_expr(value: Any) -> "LpAffineExpression":
    """Convert a scalar or variable into a linear expression."""
    if isinstance(value, LpAffineExpression):
        return value
    if isinstance(value, LpVariable):
        return LpAffineExpression({value: 1.0})
    if isinstance(value, Number):
        return LpAffineExpression(constant=float(value))
    raise TypeError(f"Unsupported linear expression term: {type(value)!r}")


class _LpConstraint:
    """Lightweight recorded linear constraint."""

    def __init__(self, lhs: Any, sense: str, rhs: Any) -> None:
        """Store the two sides and comparison sense for a constraint."""
        self.lhs = _to_expr(lhs)
        self.sense = sense
        self.rhs = _to_expr(rhs)

    def __str__(self) -> str:
        """Return a readable LP-like constraint string."""
        return f"{self.lhs} {self.sense} {self.rhs}"


class _LpAffineExpression:
    """Small linear-expression subset compatible with the SAPP static paths."""

    def __init__(self, terms: Optional[dict["_LpVariable", float]] = None, constant: float = 0.0) -> None:
        """Create a linear expression from variable coefficients and a constant."""
        self.terms = dict(terms or {})
        self.constant = float(constant)

    def value(self) -> float:
        """Evaluate the expression from each variable's ``varValue``."""
        total = self.constant
        for variable, coefficient in self.terms.items():
            total += coefficient * (0 if variable.varValue is None else variable.varValue)
        return total

    def __add__(self, other: Any) -> "_LpAffineExpression":
        rhs = _to_expr(other)
        terms = dict(self.terms)
        for variable, coefficient in rhs.terms.items():
            terms[variable] = terms.get(variable, 0.0) + coefficient
        return _LpAffineExpression(terms, self.constant + rhs.constant)

    def __radd__(self, other: Any) -> "_LpAffineExpression":
        return self + other

    def __iadd__(self, other: Any) -> "_LpAffineExpression":
        updated = self + other
        self.terms = updated.terms
        self.constant = updated.constant
        return self

    def __sub__(self, other: Any) -> "_LpAffineExpression":
        return self + (-_to_expr(other))

    def __rsub__(self, other: Any) -> "_LpAffineExpression":
        return _to_expr(other) - self

    def __neg__(self) -> "_LpAffineExpression":
        return _LpAffineExpression(
            {variable: -coefficient for variable, coefficient in self.terms.items()},
            -self.constant,
        )

    def __mul__(self, other: Any) -> "_LpAffineExpression":
        if not isinstance(other, Number):
            raise TypeError("Linear expressions only support scalar multiplication.")
        return _LpAffineExpression(
            {variable: coefficient * other for variable, coefficient in self.terms.items()},
            self.constant * other,
        )

    def __rmul__(self, other: Any) -> "_LpAffineExpression":
        return self * other

    def __truediv__(self, other: Any) -> "_LpAffineExpression":
        if not isinstance(other, Number):
            raise TypeError("Linear expressions only support scalar division.")
        return self * (1.0 / other)

    def __ge__(self, other: Any) -> "_LpConstraint":
        """Build a greater-or-equal constraint."""
        return _LpConstraint(self, ">=", other)

    def __le__(self, other: Any) -> "_LpConstraint":
        """Build a less-or-equal constraint."""
        return _LpConstraint(self, "<=", other)

    def __eq__(self, other: Any) -> "_LpConstraint":  # type: ignore[override]
        """Build an equality constraint."""
        return _LpConstraint(self, "==", other)

    def __str__(self) -> str:
        """Return a readable LP-like linear expression string."""
        parts = [
            f"{coefficient:g}*{variable.name}"
            for variable, coefficient in self.terms.items()
            if coefficient
        ]
        if self.constant or not parts:
            parts.append(f"{self.constant:g}")
        return " + ".join(parts)


class _LpVariable:
    """Minimal variable object exposing PuLP-like arithmetic and ``varValue``."""

    def __init__(
        self,
        name: str,
        lowBound: Optional[float] = None,  # pylint: disable=invalid-name
        upBound: Optional[float] = None,  # pylint: disable=invalid-name

        upBound: Optional[float] = None,  # pylint: disable=invalid-name
        cat: str = LpContinuous,
    ) -> None:
        """Create a PuLP-like variable placeholder."""
        self.name = name
        self.lowBound = lowBound
        self.upBound = upBound
        self.cat = cat
        self.varValue = None

    @staticmethod
    def dicts(
        name: str,
        indices: Any,
        lowBound: Optional[float] = None,  # pylint: disable=invalid-name
        upBound: Optional[float] = None,  # pylint: disable=invalid-name

        upBound: Optional[float] = None,  # pylint: disable=invalid-name
        cat: str = LpContinuous,
    ) -> Any:
        """Create nested dicts of variables over one or more index iterables."""
        if not isinstance(indices, tuple):
            indices = (indices,)
        index_lists = [list(index) for index in indices]

        def _make(level: int, prefix: list[str]):
            if level == len(index_lists):
                suffix = "_".join(prefix)
                variable_name = f"{name}_{suffix}" if suffix else name
                return _LpVariable(variable_name, lowBound, upBound, cat)
            return {
                index: _make(level + 1, [*prefix, str(index)])
                for index in index_lists[level]
            }

        return _make(0, [])

    def value(self) -> float:
        """Return ``varValue`` or zero before solving."""
        return 0 if self.varValue is None else self.varValue

    def _expr(self) -> _LpAffineExpression:
        return _LpAffineExpression({self: 1.0})

    def __hash__(self) -> int:
        """Return an identity hash so variables can key coefficient dicts."""
        return id(self)

    def __add__(self, other: Any) -> _LpAffineExpression:
        return self._expr() + other

    def __radd__(self, other: Any) -> _LpAffineExpression:
        return self._expr() + other

    def __sub__(self, other: Any) -> _LpAffineExpression:
        return self._expr() - other

    def __rsub__(self, other: Any) -> _LpAffineExpression:
        return _to_expr(other) - self._expr()

    def __mul__(self, other: Any) -> _LpAffineExpression:
        return self._expr() * other

    def __rmul__(self, other: Any) -> _LpAffineExpression:
        return self * other

    def __truediv__(self, other: Any) -> _LpAffineExpression:
        return self._expr() / other

    def __ge__(self, other: Any) -> _LpConstraint:
        """Build a greater-or-equal constraint."""
        return self._expr() >= other

    def __le__(self, other: Any) -> _LpConstraint:
        """Build a less-or-equal constraint."""
        return self._expr() <= other

    def __eq__(self, other: Any) -> _LpConstraint:  # type: ignore[override]
        """Build an equality constraint."""
        return self._expr() == other

    def __str__(self) -> str:
        """Return the variable name."""
        return self.name


class _LpProblem:
    """Recorded LP problem sufficient for static construction and LP dumps."""

    def __init__(self, name: str, sense: int = LpMinimize) -> None:
        """Create an LP problem record."""
        self.name = name
        self.sense = sense
        self.objective = None
        self.constraints = []

    def __iadd__(self, item: Any) -> "_LpProblem":
        if isinstance(item, _LpConstraint):
            self.constraints.append(item)
        else:
            self.objective = item
        return self

    def solve(self, solver: Any = None) -> None:
        """Reject real solving because the PuLP backend is unavailable."""
        del solver
        raise ImportError(_PULP_IMPORT_ERROR)

    def writeLP(self, filename: str) -> None:  # pylint: disable=invalid-name
        """Write a minimal LP-like dump for debugging/static tests."""
        with open(filename, "w", encoding="utf-8") as file:
            file.write(f"\\* {self.name} *\\\n")
            if self.objective is not None:
                file.write(f"Minimize\n obj: {self.objective}\n")
            file.write("Subject To\n")
            for idx, constraint in enumerate(self.constraints):
                file.write(f" c{idx}: {constraint}\n")
            file.write("End\n")


def _lp_sum(items: Any) -> _LpAffineExpression:
    """Return the linear sum of possibly nested items."""
    total = _LpAffineExpression()

    def _add(item: Any) -> None:
        nonlocal total
        if isinstance(item, dict):
            for value in item.values():
                _add(value)
        elif isinstance(item, Iterable) and not isinstance(item, (str, bytes)):
            for value in item:
                _add(value)
        else:
            total += item

    _add(items)
    return total


def _get_solver(*args: Any, **kwargs: Any) -> None:
    del args, kwargs
    raise ImportError(_PULP_IMPORT_ERROR)


LpAffineExpression = _LpAffineExpression
LpVariable = _LpVariable
LpProblem = _LpProblem
lpSum = _lp_sum
getSolver = _get_solver

from typing import Any, Dict, List, Optional

try:
    import pulp as lpSolver
except ImportError:
    from hyper_parallel.auto_parallel.sapp_ppb.sapp import _pulp_compat as lpSolver

import hyper_parallel.auto_parallel.sapp_ppb.utils.recompute as Recompute
from hyper_parallel.auto_parallel.sapp_ppb.utils.layer import Layer
from hyper_parallel.auto_parallel.sapp_ppb.utils.logger import logger

        return self._parallel_dim_size("cp")

    def _share_samples_across_dp(self) -> bool:
        """Return whether the visual validation path reuses samples across DP."""
        return get_vision_parallel_config(self.args.model).get(
            "share_samples_across_dp", False,
        )

    @staticmethod

    @staticmethod
    def _resolve_dataloader_cls():
        """Return the best available train dataloader class and its label."""
        try:
            stateful_dataloader = importlib.import_module(
                "torchdata.stateful_dataloader",
            )
            return (
                getattr(stateful_dataloader, "StatefulDataLoader"),
                "StatefulDataLoader",
            )
        except (AttributeError, ImportError):
            logger.warning_rank0(
                "torchdata is not installed; using torch.utils.data.DataLoader "
                "without checkpointable iterator state."
            )
            return DataLoader, "DataLoader"

    def _setup(self):
        """Step 1: Initialize distributed environment, device mesh, and seed.

        Each ``next()`` call yields a list of micro-batches (for gradient
        accumulation).
        """
        dataloader_cls, dataloader_name = self._resolve_dataloader_cls()

        micro_bs = self.args.train.micro_batch_size

        # Sampler uses DP rank/size — TP/CP/PP/EP peers share data.

        dp_size = self.parallel_dims.dp_size
        try:
            dp_rank = self.mesh["dp"].get_local_rank()
        except (KeyError, ValueError, RuntimeError):
            dp_rank = (
                platform.get_rank() // self.parallel_dims.non_dp_size
                if self.parallel_dims.non_dp_size > 1
                else platform.get_rank()
            )

        if num_workers > 0 and prefetch_factor is not None:
            loader_kwargs["prefetch_factor"] = prefetch_factor
        if self._deterministic:
            # Pin loader RNG to the trainer seed so shuffle order is stable.
            loader_kwargs["generator"] = torch.Generator().manual_seed(
                int(self.args.train.seed),
            )
        self.train_dataloader = dataloader_cls(
            self.train_dataset, **loader_kwargs,
        )

        # Use dp_size (not world_size) — TP/CP/PP ranks share data, not split it.

            ),
            1,
        )

        if self._share_samples_across_dp():
            logger.warning_rank0(
                "vision_parallel.share_samples_across_dp=true. Use this only for "
                "validation/self-consistency checks; normal training should keep "
                "distinct samples across DP ranks."
            )

                    local = getattr(hsdp_param.sharded_param, "_local_tensor", None)
                    if local is not None and local.is_meta:
                        new_local = torch.empty_like(local, device=device_type)
                        hsdp_param.sharded_param._local_tensor = new_local  # pylint: disable=W0212
                        hsdp_param._sharded_param_data = new_local.view(-1)  # pylint: disable=W0212

    def _init_local_shards(self) -> int:
        """Init local shard of every param (kaiming for >=2D, zero else); zero buffers."""
        param_count = 0

        """Load checkpoint safetensors via spec's ``state_dict_adapter``; drop shape mismatches."""
        # Cast loaded params down to the checkpoint's advertised dtype so the
        # fp32 master matches what forward consumes.
        load_dtype = self._resolve_hf_load_dtype(weights_path)
        strict_loading = bool(getattr(self.args.model, "strict_weight_loading", False))
        adapter = adapter_cls()
        hf_sd = adapter.load_hf_state_dict(
            weights_path, self.model.config, dtype=load_dtype,
        )

            logger.warning(
                "Dropped %d keys due to shape mismatch (first 5: %s)",
                len(dropped), dropped[:5],
            )
        logger.info_rank0(
            "HF state dict validation: missing_keys=%s unexpected_keys=%s",
            missing,
            unexpected,
        )
        if strict_loading and (dropped or missing or unexpected):
            raise ValueError(
                "Strict HF checkpoint loading failed: "
                f"dropped={len(dropped)}, missing={len(missing)}, unexpected={len(unexpected)}"
            )
        # Derive missing/unexpected ourselves — ``HSDPModule.load_state_dict``

        logger.info_rank0(
            "HF (%s) load: %d tensors into hyper model",
            model_name, len(valid_sd),
        )
        if dropped:
            logger.warning(
                "Dropped keys due to shape mismatch: %s", dropped[:5],
            )
        if missing:
            logger.warning(

            logical_to_real[_strip(name)] = name
            real_to_param[name] = param
        valid_sd: dict = {}
        dropped: list = []
        unexpected: list = []
        for hf_name, hf_tensor in hf_sd.items():
            real_name = logical_to_real.get(hf_name)
            if real_name is None:
                unexpected.append(hf_name)
                continue
            tgt = tuple(real_to_param[real_name].shape)
            src = tuple(hf_tensor.shape)
            if src == tgt:

            src = tuple(hf_tensor.shape)
            if src == tgt:
                valid_sd[real_name] = hf_tensor
            else:
                dropped.append((hf_name, real_name, src, tgt))
                unexpected.append(hf_name)
        param_names = set(real_to_param.keys())
        loaded_names = set(valid_sd.keys())
        missing = sorted(param_names - loaded_names)
        unexpected = sorted(set(unexpected))
        return valid_sd, dropped, missing, unexpected

    def _load_hyper_dcp(self, weights_path: str) -> None:
        """Load weights from hyper's own DCP checkpoint format."""