diff --git a/doc/api.rst b/doc/api.rst index 973915893..87ef772a0 100644 --- a/doc/api.rst +++ b/doc/api.rst @@ -612,6 +612,25 @@ Remote solving remote.RemoteHandler +PIPS block decomposition +======================== + +Helpers for the block-structured PIPS-IPM++ backend. See the +:doc:`pips-installation` guide for usage. ``assign_blocks`` is also bound as +:meth:`Model.assign_blocks `. + +.. autosummary:: + :toctree: generated/ + + backends.pips.assign_blocks + backends.pips.diagnose + backends.pips.BlockReport + backends.pips.PipsConfig + backends.pips.build_pips_command + backends.pips.write_job + backends.pips.doctor + + Solver status and result types ============================== diff --git a/doc/pips-installation.rst b/doc/pips-installation.rst index 4a6aa2dab..73e209bcc 100644 --- a/doc/pips-installation.rst +++ b/doc/pips-installation.rst @@ -58,6 +58,62 @@ structure via ``m.blocks`` (a ``xarray.DataArray`` of block ids over one dimension); everything independent of that dimension becomes global (block 0). A model where everything ends up linking defeats the purpose. +Diagnosing block quality +======================== + +Before committing to an HPC run, check *whether* a K-way split will actually +scale on PIPS. ``linopy.backends.pips.assign_blocks`` builds the ``m.blocks`` assignment +for you by splitting one dimension into contiguous blocks, and +``linopy.backends.pips.diagnose`` reports the resulting arrowhead structure — a cheap, +pure-Python analysis that needs no PIPS build: + +.. code-block:: python + + import linopy.backends.pips + + linopy.backends.pips.assign_blocks( + m, "time", 50 + ) # 50 contiguous blocks over "time" + report = linopy.backends.pips.diagnose(m) + print(report) + +.. code-block:: text + + BlockReport: 50 blocks | 412 340 vars | 388 900 cons | 2 140 552 nnz + columns global=1 240 per-block min/med/max = 8 180 / 8 220 / 8 260 + block nnz min/med/max = 41 002 / 42 780 / 44 190 (max/med ratio 1.03) + rows local=386 400 global=12 linking=2 488 (adjacent=2 450 border=38) + border nnz=214 300 / 2 140 552 = 10.0% + parallel max_ranks=50 target_cores=200 -> ranks=50 threads=4 + warnings (none) + +``assign_blocks(m, dim, n_blocks)`` is also available as a bound method, +``m.assign_blocks("time", 50)``. It returns the model and only supports the +default ``boundary="contiguous"`` split for now. + +How to read the report: + +- **border_fraction** is the share of matrix nonzeros that sit in linking rows + or global columns — the part PIPS handles through the root Schur complement. + It must stay small: above ~15% the root work dominates and parallel speedup + collapses, so reduce ``K`` or reformulate. +- **balance** (the block-nnz ``max/median`` ratio) measures how evenly work is + spread across blocks. Synchronous interior-point iterations move at the pace + of the slowest block, so a ratio far above 1 (roughly > 3) means stragglers + stall every iteration. +- **adjacent vs. border linking rows**: *adjacent* rows touch exactly two + neighbouring local blocks (e.g. storage state-of-charge continuity at block + boundaries) and are cheap; *border* rows span many blocks (global budget / + CO₂ / energy caps) and feed the root complement. Many border rows are the + expensive case. +- **max_ranks = n_blocks**: MPI width is capped by the number of blocks, exactly + as PIPS enforces at solve time. ``diagnose(m, target_cores=...)`` folds any + cores beyond that cap into threads-per-rank and warns when the target exceeds + the block count. + +``diagnose`` also emits warnings for a non-decomposed model (``n_blocks == 1``), +high border fraction, block imbalance, and empty blocks. + .. _pips-install: System dependencies @@ -118,6 +174,66 @@ The callback driver (``pips_driver.cpp``) and a scripted end-to-end build (``build.sh``, which performs all of the above) live under ``dev-scripts/pips/`` in the linopy repository. +On HPC (module-based build) +--------------------------- + +A cluster gives you no ``sudo``: instead of ``apt-get`` you ``module load`` a +compiler + MPI toolchain (plus MKL if you want PARDISO), then run the same +``dev-scripts/pips/build.sh``. It honours the compilers the modules export +(``CC``/``CXX``/``FC``, never overriding them) and takes three optional env +knobs: ``PIPS_PREFIX`` (install prefix, default ``dev-scripts/pips/pips-install``), +``NPROC`` (parallel build jobs) and ``CMAKE_EXTRA_ARGS`` (appended verbatim to +the PIPS-IPM++ ``cmake`` configure line). + +OpenMPI + MUMPS — the fully-open path, if your site ships those modules: + +.. code-block:: bash + + module load gcc openmpi + git clone --depth 1 https://gitlab.com/pips-ipmpp/pips-ipmpp.git \ + dev-scripts/pips/pips-ipmpp + NPROC=16 bash dev-scripts/pips/build.sh + +Vendor toolchain (Intel/Cray MPI + MKL) — sketch only; the exact module names +and cmake flags are site- and version-specific: + +.. code-block:: bash + + module load intel impi mkl # or the Cray PrgEnv equivalent + export CC=mpiicc CXX=mpiicpc FC=mpiifort + export PIPS_PREFIX=$HOME/opt/pips + export CMAKE_EXTRA_ARGS="-D" + bash dev-scripts/pips/build.sh + +MUMPS needs nothing beyond the toolchain. PARDISO (Panua ≥7.2) and HSL/MA57 are +**user-supplied** and their cmake flags are not invented here — pass them +through ``CMAKE_EXTRA_ARGS`` following the upstream PIPS-IPM++ documentation. + +.. note:: + + CMake reads ``CC``/``CXX``/``FC`` only on the **first** configure of a build + directory, so ``module load`` your compiler *before* the first ``build.sh`` + run. If ``dev-scripts/pips/{pips-ipmpp/build,build-driver}`` already exist + from an earlier toolchain, delete them so the new compilers take effect. + +Preflight: ``linopy.backends.pips.doctor()`` +----------------------------------- + +Once the driver is built, run :func:`linopy.backends.pips.doctor` on the login node +before you queue anything: + +.. code-block:: python + + import linopy.backends.pips + + print(linopy.backends.pips.doctor()) # PIPS-IPM++ OK: objective=3.0000 ... + +It builds a tiny 2-block LP with a known optimum and solves it end-to-end +through the resolved launcher, driver binary and callback backend, then checks +the objective. It raises if anything in the chain is broken (missing binary or +launcher, a link/runtime failure, a wrong optimum) — so a broken toolchain +surfaces in seconds on the login node instead of mid-allocation. + Exporting and solving a model ============================= @@ -151,6 +267,64 @@ Export in-process, then hand the directory to the driver: The driver writes the primal, objective and status back into ``export-dir``. Read them with ``linopy.io.read_pips_solution``. +Running on a cluster (detached / SLURM) +======================================= + +On an HPC system you do not hold a Python process for a multi-hour, multi-node +solve. Split the work into three steps — export, submit, ingest — controlled by +a :class:`linopy.backends.pips.PipsConfig`: + +.. code-block:: python + + import linopy.backends.pips as pips + + # 1. build on a login/build node + pips.assign_blocks(m, "time", 50) + m.to_pips_files("/lustre/run42") # put the export on the parallel FS + cfg = pips.PipsConfig(threads_per_rank=8, linear_solver="pardiso") + pips.write_job( + "/lustre/run42", + cfg, + binary="/opt/pips/pips_driver", + nodes=13, + time="04:00:00", + partition="fat", + account="psa", + modules=["gcc", "openmpi", "mkl"], + env_setup=["source /opt/pips/env.sh"], + output="/lustre/run42/pips.%j.log", + ) + +.. code-block:: bash + + # 2. submit the generated job + sbatch /lustre/run42/pips.slurm + +.. code-block:: python + + # 3. later, in a fresh session, load the result onto the model + from linopy.io import read_pips_solution + + read_pips_solution("/lustre/run42", model=m) + +``write_job`` writes a SLURM script that sets ``--ntasks`` to the block count +(or ``config.n_ranks``, capped at the number of blocks), ``--cpus-per-task`` to +``threads_per_rank``, ``--output`` to the log path (``output=``, default +``/pips.%j.log``), exports ``OMP_NUM_THREADS``/``MKL_NUM_THREADS``, +and runs the driver under ``srun`` with the chosen ``linear_solver`` and any +extra options. The script body starts with ``set -euo pipefail``; when +``modules=`` is given it emits a ``module purge`` + ``module load`` preamble, +and any ``env_setup=`` lines (e.g. ``source`` a site profile) are run verbatim +before the launch. For an interactive allocation, the inline path +(``m.solve(solver_name="pips", solver_options={...})``) uses the same +``PipsConfig`` keys and honours ``launcher="srun"``. + +Both ``write_job`` and the inline solver drop a ``pips.run.json`` provenance +manifest next to the export — the linopy version, a timestamp, the resolved +``PipsConfig``, the exact launch command and its environment (and, for +``write_job``, the job settings) — so a run stays reproducible and auditable +long after the allocation ends. + Validating the export without PIPS ================================== diff --git a/doc/release_notes.rst b/doc/release_notes.rst index 99e875c36..5f6fddb19 100644 --- a/doc/release_notes.rst +++ b/doc/release_notes.rst @@ -21,6 +21,7 @@ Upcoming Version * New ``Model.to_pips_files`` exports a model with block structure (assigned via ``Model.blocks``) into the doubly-bordered block-diagonal ("arrowhead") format consumed by the distributed solver `PIPS-IPM++ `_, deriving the submatrices from linopy's internal CSR. ``linopy.io.read_pips_files`` reconstructs an equivalent flat model (usable to validate the export against any solver with no PIPS build), and ``read_pips_solution`` reads the solver output back. See :doc:`pips-installation`. * An env-gated ``pips`` solver (``solver_name="pips"``, enabled through the ``PIPS_BINARY`` environment variable) runs the export through PIPS under MPI and maps the primal, duals and objective back onto the model. It stays absent from ``available_solvers`` unless the driver binary is present, so ordinary installs are unaffected. +* New ``linopy.backends.pips`` HPC helpers: ``assign_blocks(m, dim, n_blocks)`` (also ``Model.assign_blocks``) splits a dimension into contiguous blocks, and ``diagnose(m)`` reports the arrowhead structure — border fraction, block balance, adjacent-vs-border linking rows and a rank/thread recommendation — to judge whether a decomposition will scale before running. ``PipsConfig`` controls the parallel launch (``launcher``, ``n_ranks``, ``threads_per_rank``, ``linear_solver``) via ``solver_options`` or environment variables, and ``write_job`` emits a SLURM batch script for detached cluster runs. *Other* diff --git a/linopy/backends/__init__.py b/linopy/backends/__init__.py new file mode 100644 index 000000000..eb282a5ff --- /dev/null +++ b/linopy/backends/__init__.py @@ -0,0 +1 @@ +"""Backend-specific integrations that need more than a solver class in `solvers.py`.""" diff --git a/linopy/backends/pips.py b/linopy/backends/pips.py new file mode 100644 index 000000000..7a2505477 --- /dev/null +++ b/linopy/backends/pips.py @@ -0,0 +1,463 @@ +from __future__ import annotations + +import json +import os +import shlex +from dataclasses import asdict, dataclass, field, replace +from datetime import datetime +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import numpy as np +import xarray as xr + +if TYPE_CHECKING: + from linopy.model import Model + +LAUNCHER_RANK_FLAG = {"mpirun": "-np", "srun": "-n"} + + +def assign_blocks( + m: Model, dim: str, n_blocks: int, boundary: str = "contiguous" +) -> Model: + if boundary != "contiguous": + raise NotImplementedError(f"boundary {boundary!r} not supported") + if dim not in m.variables.indexes: + raise ValueError(f"dimension {dim!r} not found in model variables") + index = m.variables.indexes[dim] + if n_blocks < 1: + raise ValueError("n_blocks must be >= 1") + if n_blocks > len(index): + raise ValueError( + f"n_blocks ({n_blocks}) exceeds length of {dim!r} ({len(index)})" + ) + ids = np.empty(len(index), dtype=np.int64) + for i, part in enumerate(np.array_split(np.arange(len(index)), n_blocks)): + ids[part] = i + 1 + m.blocks = xr.DataArray(ids, coords=[index], dims=[dim]) + return m + + +def _fmt(value: int) -> str: + return f"{value:,}".replace(",", " ") + + +@dataclass +class BlockReport: + n_blocks: int + n_vars: int + n_cons: int + nnz: int + + n_global_cols: int + block_cols: dict[int, int] + + block_nnz: dict[int, int] + balance_min: int + balance_median: float + balance_max: int + balance_ratio: float + + n_local_rows: int + n_global_rows: int + n_linking_rows: int + n_adjacent_rows: int + n_border_rows: int + + border_nnz: int + border_fraction: float + + max_ranks: int + target_cores: int | None + rec_ranks: int + rec_threads: int + + warnings: list[str] = field(default_factory=list) + + def __str__(self) -> str: + cols_med = ( + float(np.median(list(self.block_cols.values()))) if self.block_cols else 0.0 + ) + col_min = min(self.block_cols.values()) if self.block_cols else 0 + col_max = max(self.block_cols.values()) if self.block_cols else 0 + header = ( + f"BlockReport: {self.n_blocks} blocks | {_fmt(self.n_vars)} vars | " + f"{_fmt(self.n_cons)} cons | {_fmt(self.nnz)} nnz" + ) + columns = ( + f" columns global={_fmt(self.n_global_cols)} per-block min/med/max " + f"= {_fmt(col_min)} / {_fmt(int(cols_med))} / {_fmt(col_max)}" + ) + block = ( + f" block nnz min/med/max = {_fmt(self.balance_min)} / " + f"{_fmt(int(self.balance_median))} / {_fmt(self.balance_max)} " + f"(max/med ratio {self.balance_ratio:.2f})" + ) + rows = ( + f" rows local={_fmt(self.n_local_rows)} " + f"global={_fmt(self.n_global_rows)} linking={_fmt(self.n_linking_rows)} " + f"(adjacent={_fmt(self.n_adjacent_rows)} border={_fmt(self.n_border_rows)})" + ) + border = ( + f" border nnz={_fmt(self.border_nnz)} / {_fmt(self.nnz)} = " + f"{self.border_fraction:.1%}" + ) + if self.target_cores is None: + parallel = ( + f" parallel max_ranks={self.max_ranks} " + f"ranks={self.rec_ranks} threads={self.rec_threads}" + ) + else: + parallel = ( + f" parallel max_ranks={self.max_ranks} " + f"target_cores={self.target_cores} -> ranks={self.rec_ranks} " + f"threads={self.rec_threads}" + ) + if self.warnings: + warn = "\n".join(f" {w}" for w in self.warnings) + warnings = f" warnings\n{warn}" + else: + warnings = " warnings (none)" + return "\n".join([header, columns, block, rows, border, parallel, warnings]) + + __repr__ = __str__ + + +def diagnose(m: Model, target_cores: int | None = None) -> BlockReport: + if m.blocks is None: + raise ValueError( + "no blocks assigned; call assign_blocks(m, dim, n_blocks) first" + ) + + m.calculate_block_maps() + + if m.matrices.A is None: + raise ValueError("model has no regular constraints to diagnose") + + N = int(m.blocks.max()) + block_map = m.variables.get_blockmap(m.blocks.dtype.type) + vlabels = m.matrices.vlabels + col_blocks = block_map[vlabels] + A = m.matrices.A + + row_blocks = np.concatenate( + [ + c.data["blocks"].values.ravel()[c.active_row_mask()] + for _, c in m.constraints.items() + if not c.is_indicator + ] + ) + assert len(row_blocks) == A.shape[0] + + coo = A.tocoo() + rb = row_blocks[coo.row] + cb = col_blocks[coo.col] + + is_border = (rb == N + 1) | (cb == 0) + border_nnz = int(is_border.sum()) + border_fraction = border_nnz / A.nnz if A.nnz else 0.0 + + counts = np.bincount(np.clip(rb, 0, N + 1), minlength=N + 2) + block_nnz = {n: int(counts[n]) for n in range(1, N + 1)} + nonempty = np.array([v for v in block_nnz.values() if v > 0]) + if nonempty.size: + balance_min = int(nonempty.min()) + balance_max = int(nonempty.max()) + balance_median = float(np.median(nonempty)) + balance_ratio = balance_max / balance_median if balance_median else 0.0 + else: + balance_min = balance_max = 0 + balance_median = balance_ratio = 0.0 + + n_global_cols = int((col_blocks == 0).sum()) + col_counts = np.bincount(np.clip(col_blocks, 0, N), minlength=N + 1) + block_cols = {n: int(col_counts[n]) for n in range(1, N + 1)} + + n_local_rows = int(((row_blocks >= 1) & (row_blocks <= N)).sum()) + n_global_rows = int((row_blocks == 0).sum()) + linking_mask = row_blocks == N + 1 + n_linking_rows = int(linking_mask.sum()) + + indptr, indices = A.indptr, A.indices + n_adjacent_rows = 0 + n_border_rows = 0 + for i in np.nonzero(linking_mask)[0]: + cols = indices[indptr[i] : indptr[i + 1]] + local = np.unique(col_blocks[cols]) + n_local = int((local >= 1).sum()) + if n_local == 2: + n_adjacent_rows += 1 + else: + n_border_rows += 1 + + max_ranks = N + if target_cores is None: + rec_ranks, rec_threads = max_ranks, 1 + else: + rec_ranks = min(max_ranks, target_cores) + rec_threads = max(1, target_cores // rec_ranks) + + warnings: list[str] = [] + if N == 1: + warnings.append("model is not decomposed (n_blocks == 1)") + if border_fraction > 0.15: + warnings.append( + f"high border fraction {border_fraction:.1%} (> 15%): root Schur " + "complement will dominate; reduce K or reformulate" + ) + if balance_ratio > 3: + warnings.append( + f"block imbalance max/median = {balance_ratio:.1f} (> 3): stragglers " + "will stall synchronous IPM iterations" + ) + empty = [n for n, v in block_nnz.items() if v == 0] + if empty: + warnings.append(f"empty local blocks (no rows): {empty}") + if target_cores is not None and target_cores > max_ranks: + warnings.append( + f"target_cores {target_cores} > n_blocks {max_ranks}: MPI width is " + "capped by blocks; raise K or add threads_per_rank" + ) + + return BlockReport( + n_blocks=N, + n_vars=len(vlabels), + n_cons=A.shape[0], + nnz=A.nnz, + n_global_cols=n_global_cols, + block_cols=block_cols, + block_nnz=block_nnz, + balance_min=balance_min, + balance_median=balance_median, + balance_max=balance_max, + balance_ratio=balance_ratio, + n_local_rows=n_local_rows, + n_global_rows=n_global_rows, + n_linking_rows=n_linking_rows, + n_adjacent_rows=n_adjacent_rows, + n_border_rows=n_border_rows, + border_nnz=border_nnz, + border_fraction=border_fraction, + max_ranks=max_ranks, + target_cores=target_cores, + rec_ranks=rec_ranks, + rec_threads=rec_threads, + warnings=warnings, + ) + + +@dataclass +class PipsConfig: + launcher: str = "mpirun" + n_ranks: int | None = None + threads_per_rank: int = 1 + launcher_args: list[str] = field(default_factory=list) + linear_solver: str | None = None + options: dict[str, Any] = field(default_factory=dict) + + +def _resolve_ranks(config: PipsConfig, n_blocks: int | None) -> int: + ranks = config.n_ranks if config.n_ranks is not None else (n_blocks or 1) + if ranks < 1: + raise ValueError("n_ranks must be >= 1") + if n_blocks is not None and ranks > n_blocks: + ranks = n_blocks + return ranks + + +def build_pips_command( + binary: str, + export_dir: str, + config: PipsConfig, + n_blocks: int | None = None, +) -> tuple[list[str], dict[str, str]]: + if config.launcher not in LAUNCHER_RANK_FLAG: + raise ValueError( + f"launcher {config.launcher!r} not supported; use one of " + f"{sorted(LAUNCHER_RANK_FLAG)}" + ) + ranks = _resolve_ranks(config, n_blocks) + command = [ + config.launcher, + LAUNCHER_RANK_FLAG[config.launcher], + str(ranks), + *config.launcher_args, + binary, + export_dir, + ] + driver_options = dict(config.options) + if config.linear_solver is not None: + driver_options.setdefault("linear-solver", config.linear_solver) + for key, value in driver_options.items(): + command += [f"--{key}", str(value)] + threads = str(config.threads_per_rank) + env = {"OMP_NUM_THREADS": threads, "MKL_NUM_THREADS": threads} + return command, env + + +def _linopy_version() -> str: + from importlib.metadata import PackageNotFoundError, version + + try: + return version("linopy") + except PackageNotFoundError: + return "unknown" + + +def write_run_manifest( + export_dir: str | Path, + *, + config: PipsConfig, + command: list[str], + env: dict[str, str], + n_blocks: int | None, + job: dict[str, Any] | None = None, +) -> Path: + manifest = { + "linopy_version": _linopy_version(), + "created": datetime.now().isoformat(timespec="seconds"), + "n_blocks": n_blocks, + "config": asdict(config), + "command": command, + "env": env, + } + if job is not None: + manifest["job"] = job + path = Path(export_dir) / "pips.run.json" + path.write_text(json.dumps(manifest, indent=2, default=str)) + return path + + +def write_job( + export_dir: str | Path, + config: PipsConfig | None = None, + *, + binary: str | None = None, + scheduler: str = "slurm", + nodes: int | None = None, + time: str | None = None, + partition: str | None = None, + account: str | None = None, + job_name: str = "pips-ipmpp", + output: str | None = None, + modules: list[str] | None = None, + env_setup: list[str] | None = None, + sbatch_args: list[str] | None = None, + path: str | Path | None = None, +) -> Path: + if scheduler != "slurm": + raise NotImplementedError(f"scheduler {scheduler!r} not supported; use 'slurm'") + config = config or PipsConfig() + export_dir = Path(export_dir).resolve() + n_blocks = json.loads((export_dir / "pips.json").read_text()).get("n_blocks") + binary = binary or os.environ.get("PIPS_BINARY") + if not binary: + raise ValueError("no PIPS driver binary given and $PIPS_BINARY is not set") + resolved = Path(binary) + binary = str(resolved.resolve()) if resolved.exists() else binary + + launch_config = replace(config, launcher="srun") + command, env = build_pips_command(binary, str(export_dir), launch_config, n_blocks) + ranks = _resolve_ranks(config, n_blocks) + output = output or str(export_dir / "pips.%j.log") + + directives = [ + f"#SBATCH --job-name={job_name}", + f"#SBATCH --ntasks={ranks}", + f"#SBATCH --cpus-per-task={config.threads_per_rank}", + f"#SBATCH --output={output}", + ] + if nodes is not None: + directives.append(f"#SBATCH --nodes={nodes}") + if time is not None: + directives.append(f"#SBATCH --time={time}") + if partition is not None: + directives.append(f"#SBATCH --partition={partition}") + if account is not None: + directives.append(f"#SBATCH --account={account}") + directives += [f"#SBATCH {arg}" for arg in sbatch_args or []] + + setup = ["set -euo pipefail"] + env_block = [] + if modules: + env_block.append("module purge") + env_block.append("module load " + " ".join(modules)) + env_block += list(env_setup or []) + if env_block: + setup += ["set +u", *env_block, "set -u"] + + exports = [f"export {key}={value}" for key, value in env.items()] + echoes = [f'echo "PIPS-IPM++ driver: {binary}"', f'echo "export dir: {export_dir}"'] + lines = [ + "#!/bin/bash", + *directives, + "", + *setup, + "", + *exports, + "", + *echoes, + "", + shlex.join(command), + "", + ] + + path = Path(path) if path is not None else export_dir / "pips.slurm" + path.write_text("\n".join(lines)) + + job = { + "scheduler": scheduler, + "job_name": job_name, + "nodes": nodes, + "time": time, + "partition": partition, + "account": account, + "output": output, + "modules": list(modules or []), + "env_setup": list(env_setup or []), + "sbatch_args": list(sbatch_args or []), + "script": str(path), + } + write_run_manifest( + export_dir, + config=launch_config, + command=command, + env=env, + n_blocks=n_blocks, + job=job, + ) + return path + + +def doctor(solver_options: dict[str, Any] | None = None) -> str: + """ + Solve a tiny 2-block LP through the resolved PIPS launcher+binary+backend. + + A login-node preflight: builds a synthetic separable model with a known + optimum (3.0), solves it via ``solver_name="pips"`` and checks the objective. + Raises on any failure (missing binary/launcher, non-optimal, wrong optimum); + returns a one-line OK report otherwise. + """ + import pandas as pd + + from linopy import Model + + m = Model() + time = pd.RangeIndex(4, name="time") + x = m.add_variables(lower=0, coords=[time], name="x") + g = m.add_variables(lower=0, name="g") + m.add_constraints(x - g <= 0, name="cap") + m.add_constraints(x.sum() <= 4, name="budget") + m.add_objective(x.sum() - g, sense="max") + m.assign_blocks("time", 2) + + m.solve(solver_name="pips", solver_options=solver_options or {}) + + objective = m.objective.value + expected = 3.0 + if objective is None or abs(float(objective) - expected) > 1e-4: + raise RuntimeError( + f"PIPS doctor failed: objective {objective} != expected {expected}; " + "the driver ran but produced a wrong optimum" + ) + return f"PIPS-IPM++ OK: objective={float(objective):.4f} (expected {expected:.4f})" diff --git a/linopy/model.py b/linopy/model.py index 5c55fdb72..0a833e3cf 100644 --- a/linopy/model.py +++ b/linopy/model.py @@ -30,6 +30,7 @@ from linopy import solvers from linopy.alignment import as_dataarray, broadcast_to_coords +from linopy.backends.pips import assign_blocks from linopy.common import ( assign_multiindex_safe, best_int, @@ -2395,4 +2396,6 @@ def reset_solution(self) -> None: to_pips_files = to_pips_files + assign_blocks = assign_blocks + dualize = dualize diff --git a/linopy/solvers.py b/linopy/solvers.py index caede2b0f..6ae235255 100644 --- a/linopy/solvers.py +++ b/linopy/solvers.py @@ -9,6 +9,7 @@ import enum import functools import io +import json import logging import os import re @@ -34,6 +35,7 @@ from scipy.sparse import tril, triu import linopy.io +from linopy.backends.pips import PipsConfig, build_pips_command, write_run_manifest from linopy.common import count_initial_letters, values_to_lookup_array from linopy.constants import ( EQUAL, @@ -4112,6 +4114,9 @@ def get_solver_solution() -> Solution: PIPS_BINARY_ENV = "PIPS_BINARY" PIPS_MPI_RANKS_ENV = "PIPS_MPI_RANKS" +PIPS_LAUNCHER_ENV = "PIPS_LAUNCHER" +PIPS_THREADS_ENV = "PIPS_THREADS" +PIPS_LINEAR_SOLVER_ENV = "PIPS_LINEAR_SOLVER" def _pips_binary() -> str | None: @@ -4120,22 +4125,38 @@ def _pips_binary() -> str | None: return shutil.which(binary) if binary else None +def _pips_config(solver_options: dict[str, Any]) -> PipsConfig: + """Build a :class:`PipsConfig` from env-var defaults overridden by solver options.""" + opts = dict(solver_options) + n_ranks = opts.pop("n_ranks", os.environ.get(PIPS_MPI_RANKS_ENV)) + threads = opts.pop("threads_per_rank", os.environ.get(PIPS_THREADS_ENV, "1")) + return PipsConfig( + launcher=opts.pop("launcher", os.environ.get(PIPS_LAUNCHER_ENV, "mpirun")), + n_ranks=int(n_ranks) if n_ranks is not None else None, + threads_per_rank=int(threads), + launcher_args=list(opts.pop("launcher_args", [])), + linear_solver=opts.pop("linear_solver", os.environ.get(PIPS_LINEAR_SOLVER_ENV)), + options=opts, + ) + + class PIPS(Solver[None]): """ Solver subclass for the distributed PIPS-IPM++ solver. The model is exported to the arrowhead block format via :func:`linopy.io.to_pips_files`, the driver is launched through - ``mpirun -np `` and the solution is read back - with :func:`linopy.io.read_pips_solution`. The driver binary is taken from - ``$PIPS_BINARY`` and the MPI rank count from ``$PIPS_MPI_RANKS`` (default 1). - PIPS only appears in ``available_solvers`` when both the binary and - ``mpirun`` are found. - - Attributes - ---------- - **solver_options - forwarded to the driver as ``-- `` command-line arguments + `` <-np|-n> `` and the solution is + read back with :func:`linopy.io.read_pips_solution`. The driver binary is + taken from ``$PIPS_BINARY``; PIPS only appears in ``available_solvers`` when + the binary and a launcher (``mpirun`` or ``srun``) are found. + + Launch is controlled through ``solver_options`` (overriding env defaults): + ``launcher`` (``mpirun``/``srun``; ``$PIPS_LAUNCHER``), ``n_ranks`` + (``$PIPS_MPI_RANKS``, default = number of blocks), ``threads_per_rank`` + (``$PIPS_THREADS``, sets ``OMP_NUM_THREADS``/``MKL_NUM_THREADS``), + ``linear_solver`` (``$PIPS_LINEAR_SOLVER``) and ``launcher_args``. Any other + option is forwarded to the driver as ``-- ``. """ display_name: ClassVar[str] = "PIPS-IPM++" @@ -4146,7 +4167,10 @@ class PIPS(Solver[None]): @classmethod @functools.cache def is_available(cls) -> bool: - return _pips_binary() is not None and shutil.which("mpirun") is not None + launcher = ( + shutil.which("mpirun") is not None or shutil.which("srun") is not None + ) + return _pips_binary() is not None and launcher def _build_file(self, **build_kwargs: Any) -> None: model = self.model @@ -4172,19 +4196,37 @@ def _run_file( f"PIPS driver binary not found. Set the {PIPS_BINARY_ENV} " "environment variable to the PIPS-IPM++ executable." ) - if shutil.which("mpirun") is None: - raise RuntimeError("`mpirun` not found on PATH; cannot launch PIPS-IPM++.") + config = _pips_config(self.solver_options) + if shutil.which(config.launcher) is None: + raise RuntimeError( + f"launcher {config.launcher!r} not found on PATH; cannot launch " + "PIPS-IPM++." + ) - ranks = int(os.environ.get(PIPS_MPI_RANKS_ENV, "1")) - command = ["mpirun", "-np", str(ranks), binary, path_to_string(export_dir)] - for k, v in self.solver_options.items(): - command += [f"--{k}", str(v)] + manifest = json.loads( + (Path(path_to_string(export_dir)) / "pips.json").read_text() + ) + command, launch_env = build_pips_command( + binary, path_to_string(export_dir), config, manifest.get("n_blocks") + ) + write_run_manifest( + path_to_string(export_dir), + config=config, + command=command, + env=launch_env, + n_blocks=manifest.get("n_blocks"), + ) + run_env = {**os.environ, **launch_env} if log_fn is not None: with open(log_fn, "w") as log_f: - proc = sub.run(command, stdout=log_f, stderr=sub.STDOUT, text=True) + proc = sub.run( + command, stdout=log_f, stderr=sub.STDOUT, text=True, env=run_env + ) else: - proc = sub.run(command, stdout=sub.PIPE, stderr=sub.STDOUT, text=True) + proc = sub.run( + command, stdout=sub.PIPE, stderr=sub.STDOUT, text=True, env=run_env + ) logger.info(proc.stdout) if proc.returncode != 0: raise RuntimeError( diff --git a/test/test_pips.py b/test/test_pips.py new file mode 100644 index 000000000..b0db4d0a9 --- /dev/null +++ b/test/test_pips.py @@ -0,0 +1,441 @@ +from __future__ import annotations + +import json +from dataclasses import asdict, replace + +import numpy as np +import pandas as pd +import pytest +from test_io import _pips_time_plant_model + +import linopy.backends.pips as pips +from linopy import Model, available_solvers + + +def _time_model(n_time: int) -> Model: + m = Model() + time = pd.RangeIndex(n_time, name="time") + x = m.add_variables(coords=[time], name="x") + m.add_constraints(x <= 1, name="c") + m.add_objective(x.sum()) + return m + + +def _synthetic_storage_model() -> Model: + m = Model() + time = pd.RangeIndex(6, name="time") + plant = pd.Index([0, 1], name="plant") + x = m.add_variables(lower=0, coords=[time, plant], name="x") + soc = m.add_variables(coords=[time], name="soc") + cap = m.add_variables(lower=0, coords=[plant], name="cap") + m.add_constraints(x - cap <= 0, name="capacity") + continuity = (soc - soc.shift(time=1) - x.sum("plant")).isel(time=slice(1, None)) + m.add_constraints(continuity == 0, name="storage") + m.add_constraints(x.sum() <= 100, name="budget") + m.add_objective(x.sum() + soc.sum() + cap.sum()) + return m + + +@pytest.mark.parametrize( + "n_time, n_blocks, sizes", + [(6, 2, [3, 3]), (6, 4, [2, 2, 1, 1]), (10, 3, [4, 3, 3])], +) +def test_assign_blocks_contiguous(n_time: int, n_blocks: int, sizes: list[int]) -> None: + m = _time_model(n_time) + out = pips.assign_blocks(m, "time", n_blocks) + expected = np.repeat(np.arange(1, n_blocks + 1), sizes) + assert out is m + assert m.blocks.dims == ("time",) + assert list(m.blocks.coords["time"].values) == list(range(n_time)) + np.testing.assert_array_equal(m.blocks.values, expected) + + +def test_assign_blocks_dtype_is_best_int() -> None: + m = _time_model(6) + m.assign_blocks("time", 2) + assert m.blocks.dtype == np.dtype("int8") + + +def test_assign_blocks_method_matches_function() -> None: + a = _time_model(6).assign_blocks("time", 3) + b = pips.assign_blocks(_time_model(6), "time", 3) + np.testing.assert_array_equal(a.blocks.values, b.blocks.values) + + +@pytest.mark.parametrize( + "kwargs, exc", + [ + ({"dim": "nope", "n_blocks": 2}, ValueError), + ({"dim": "time", "n_blocks": 0}, ValueError), + ({"dim": "time", "n_blocks": 7}, ValueError), + ({"dim": "time", "n_blocks": 2, "boundary": "custom"}, NotImplementedError), + ], +) +def test_assign_blocks_fail_fast(kwargs: dict, exc: type[Exception]) -> None: + m = _time_model(6) + with pytest.raises(exc): + pips.assign_blocks(m, **kwargs) + + +def test_diagnose_synthetic_exact_counts() -> None: + m = _synthetic_storage_model() + m.assign_blocks("time", 3) + r = pips.diagnose(m) + + assert r.n_blocks == 3 + assert r.n_vars == 20 + assert r.n_cons == 18 + assert r.nnz == 56 + + assert r.n_global_cols == 2 + assert r.block_cols == {1: 6, 2: 6, 3: 6} + assert r.n_global_cols + sum(r.block_cols.values()) == r.n_vars + + assert r.block_nnz == {1: 12, 2: 12, 3: 12} + assert r.balance_ratio == 1.0 + + assert (r.n_local_rows, r.n_global_rows, r.n_linking_rows) == (15, 0, 3) + assert r.n_local_rows + r.n_global_rows + r.n_linking_rows == r.n_cons + assert (r.n_adjacent_rows, r.n_border_rows) == (2, 1) + assert r.n_adjacent_rows + r.n_border_rows == r.n_linking_rows + + assert r.border_nnz == 32 + assert r.border_fraction == pytest.approx(32 / 56) + + +def test_diagnose_border_nnz_bruteforce() -> None: + m = _synthetic_storage_model() + m.assign_blocks("time", 3) + r = pips.diagnose(m) + + N = int(m.blocks.max()) + block_map = m.variables.get_blockmap(m.blocks.dtype.type) + col_blocks = block_map[m.matrices.vlabels] + row_blocks = np.concatenate( + [ + c.data["blocks"].values.ravel()[c.active_row_mask()] + for _, c in m.constraints.items() + if not c.is_indicator + ] + ) + coo = m.matrices.A.tocoo() + is_border = (row_blocks[coo.row] == N + 1) | (col_blocks[coo.col] == 0) + assert r.border_nnz == int(is_border.sum()) + assert 0.0 <= r.border_fraction <= 1.0 + + +def test_diagnose_preconditions() -> None: + m = _synthetic_storage_model() + with pytest.raises(ValueError, match="no blocks assigned"): + pips.diagnose(m) + + +def test_diagnose_no_regular_constraints() -> None: + m = Model() + time = pd.RangeIndex(4, name="time") + x = m.add_variables(coords=[time], name="x") + m.add_objective(x.sum()) + m.assign_blocks("time", 2) + with pytest.raises(ValueError, match="no regular constraints"): + pips.diagnose(m) + + +@pytest.mark.parametrize( + "target_cores, rec_ranks, rec_threads, capped", + [(None, 50, 1, False), (50, 50, 1, False), (200, 50, 4, True), (30, 30, 1, False)], +) +def test_diagnose_recommendation( + target_cores: int | None, rec_ranks: int, rec_threads: int, capped: bool +) -> None: + m = _time_model(50) + m.assign_blocks("time", 50) + r = pips.diagnose(m, target_cores=target_cores) + assert r.max_ranks == 50 + assert r.rec_ranks == rec_ranks + assert r.rec_threads == rec_threads + assert any("capped" in w for w in r.warnings) == capped + assert r.rec_ranks <= r.max_ranks + assert r.rec_threads >= 1 + + +def _imbalanced_model() -> Model: + m = Model() + time = pd.RangeIndex(9, name="time") + x = m.add_variables(coords=[time], name="x") + m.add_constraints(x <= 1, name="c") + for i in range(3): + m.add_constraints(x.isel(time=slice(0, 3)) >= -5, name=f"extra{i}") + m.add_objective(x.sum()) + m.assign_blocks("time", 3) + return m + + +def _empty_block_model() -> Model: + m = Model() + time = pd.RangeIndex(9, name="time") + x = m.add_variables(coords=[time], name="x") + m.add_constraints(x.isel(time=slice(0, 6)) <= 1, name="c") + m.add_objective(x.sum()) + m.assign_blocks("time", 3) + return m + + +@pytest.mark.parametrize( + "builder, substring", + [ + (_synthetic_storage_model, "high border fraction"), + (_imbalanced_model, "block imbalance"), + (_empty_block_model, "empty local blocks"), + (lambda: _time_model(6).assign_blocks("time", 1), "not decomposed"), + ], +) +def test_diagnose_warnings(builder, substring: str) -> None: + m = builder() + if m.blocks is None: + m.assign_blocks("time", 3) + r = pips.diagnose(m) + assert any(substring in w for w in r.warnings) + + +def test_diagnose_str_renders_all_groups() -> None: + m = _synthetic_storage_model() + m.assign_blocks("time", 3) + text = str(pips.diagnose(m, target_cores=8)) + for token in [ + "BlockReport", + "columns", + "block nnz", + "rows", + "border", + "parallel", + "warnings", + ]: + assert token in text + + +@pytest.mark.parametrize("masked", [False, True]) +def test_diagnose_realistic_consistency(masked: bool) -> None: + m = _pips_time_plant_model(masked=masked) + m.assign_blocks("time", 2) + r = pips.diagnose(m, target_cores=8) + assert r.n_global_cols + sum(r.block_cols.values()) == r.n_vars + assert r.n_local_rows + r.n_global_rows + r.n_linking_rows == r.n_cons + assert r.n_adjacent_rows + r.n_border_rows == r.n_linking_rows + assert r.border_nnz == pytest.approx(r.border_fraction * r.nnz) + assert r.rec_ranks <= r.max_ranks == r.n_blocks + + +def test_build_pips_command_defaults_ranks_to_blocks() -> None: + cmd, env = pips.build_pips_command("drv", "dir", pips.PipsConfig(), n_blocks=8) + assert cmd == ["mpirun", "-np", "8", "drv", "dir"] + assert env == {"OMP_NUM_THREADS": "1", "MKL_NUM_THREADS": "1"} + + +def test_build_pips_command_caps_ranks_at_blocks() -> None: + cfg = pips.PipsConfig(n_ranks=200) + cmd, _ = pips.build_pips_command("drv", "dir", cfg, n_blocks=50) + assert cmd[:3] == ["mpirun", "-np", "50"] + + +def test_build_pips_command_srun_threads_solver_and_args() -> None: + cfg = pips.PipsConfig( + launcher="srun", + n_ranks=4, + threads_per_rank=8, + launcher_args=["--exclusive"], + linear_solver="pardiso", + options={"tol": 1e-8}, + ) + cmd, env = pips.build_pips_command("drv", "dir", cfg) + assert cmd[:4] == ["srun", "-n", "4", "--exclusive"] + assert cmd[4:6] == ["drv", "dir"] + assert cmd[6:] == ["--tol", "1e-08", "--linear-solver", "pardiso"] + assert env == {"OMP_NUM_THREADS": "8", "MKL_NUM_THREADS": "8"} + + +@pytest.mark.parametrize( + "cfg, exc", + [ + (pips.PipsConfig(launcher="poe"), ValueError), + (pips.PipsConfig(n_ranks=0), ValueError), + ], +) +def test_build_pips_command_fail_fast( + cfg: pips.PipsConfig, exc: type[Exception] +) -> None: + with pytest.raises(exc): + pips.build_pips_command("drv", "dir", cfg, n_blocks=4) + + +def test_pips_config_env_defaults_and_option_overrides(monkeypatch) -> None: + from linopy.solvers import _pips_config + + monkeypatch.setenv("PIPS_LAUNCHER", "srun") + monkeypatch.setenv("PIPS_THREADS", "4") + monkeypatch.setenv("PIPS_LINEAR_SOLVER", "mumps") + env_cfg = _pips_config({"presolve": "on"}) + assert env_cfg.launcher == "srun" + assert env_cfg.threads_per_rank == 4 + assert env_cfg.linear_solver == "mumps" + assert env_cfg.options == {"presolve": "on"} + + override = _pips_config( + {"launcher": "mpirun", "n_ranks": 12, "linear_solver": "ma57"} + ) + assert override.launcher == "mpirun" + assert override.n_ranks == 12 + assert override.linear_solver == "ma57" + + +def _exported(tmp_path, n_blocks: int = 3) -> Model: + m = _time_model(6) + m.assign_blocks("time", n_blocks) + m.to_pips_files(tmp_path) + return m + + +def test_write_job_slurm_script(tmp_path) -> None: + _exported(tmp_path, 3) + cfg = pips.PipsConfig(threads_per_rank=4, linear_solver="mumps") + script = pips.write_job( + tmp_path, + cfg, + binary="/opt/pips/drv", + nodes=2, + time="01:00:00", + partition="p", + account="a", + ) + assert script == tmp_path / "pips.slurm" + text = script.read_text() + assert text.startswith("#!/bin/bash") + assert "#SBATCH --ntasks=3" in text + assert "#SBATCH --cpus-per-task=4" in text + assert "#SBATCH --nodes=2" in text + assert "#SBATCH --time=01:00:00" in text + assert "#SBATCH --partition=p" in text + assert "#SBATCH --account=a" in text + assert "export OMP_NUM_THREADS=4" in text + assert "export MKL_NUM_THREADS=4" in text + assert "srun -n 3" in text + assert "/opt/pips/drv" in text + assert str(tmp_path.resolve()) in text + assert "--linear-solver mumps" in text + + +def test_write_job_binary_from_env(tmp_path, monkeypatch) -> None: + _exported(tmp_path, 2) + monkeypatch.setenv("PIPS_BINARY", "/env/drv") + text = pips.write_job(tmp_path).read_text() + assert "/env/drv" in text + assert "#SBATCH --ntasks=2" in text + + +def test_write_job_custom_path_and_sbatch_args(tmp_path) -> None: + _exported(tmp_path, 2) + out = tmp_path / "custom.slurm" + script = pips.write_job( + tmp_path, binary="/d", path=out, sbatch_args=["--qos=long", "--mem=0"] + ) + assert script == out + text = out.read_text() + assert "#SBATCH --qos=long" in text + assert "#SBATCH --mem=0" in text + + +@pytest.mark.parametrize( + "kwargs, exc", + [ + ({"scheduler": "pbs", "binary": "/d"}, NotImplementedError), + ({}, ValueError), + ], +) +def test_write_job_fail_fast(tmp_path, monkeypatch, kwargs: dict, exc) -> None: + _exported(tmp_path, 2) + monkeypatch.delenv("PIPS_BINARY", raising=False) + with pytest.raises(exc): + pips.write_job(tmp_path, **kwargs) + + +def test_write_job_emits_output_and_strict_mode_and_manifest(tmp_path) -> None: + _exported(tmp_path, 2) + text = pips.write_job(tmp_path, binary="/d").read_text() + assert "#SBATCH --output=" in text + assert "set -euo pipefail" in text + assert "module load" not in text + assert (tmp_path / "pips.run.json").exists() + + +def test_write_job_modules_env_setup_and_custom_output(tmp_path) -> None: + _exported(tmp_path, 2) + text = pips.write_job( + tmp_path, + binary="/d", + modules=["gcc", "openmpi"], + env_setup=["source /opt/x/setvars.sh"], + output="/x/%j.out", + ).read_text() + assert "module purge" in text + assert "module load gcc openmpi" in text + assert "source /opt/x/setvars.sh" in text + assert "#SBATCH --output=/x/%j.out" in text + assert "set +u" in text and "set -u" in text + + +def test_write_job_run_manifest_content(tmp_path) -> None: + _exported(tmp_path, 2) + cfg = pips.PipsConfig(threads_per_rank=4, linear_solver="mumps") + pips.write_job( + tmp_path, + cfg, + binary="/d", + modules=["gcc"], + env_setup=["source /opt/x/setvars.sh"], + output="/x/%j.out", + ) + manifest = json.loads((tmp_path / "pips.run.json").read_text()) + assert set(manifest) == { + "linopy_version", + "created", + "n_blocks", + "config", + "command", + "env", + "job", + } + assert manifest["config"] == asdict(replace(cfg, launcher="srun")) + assert manifest["command"][0] == "srun" + assert manifest["env"]["OMP_NUM_THREADS"] == "4" + assert manifest["n_blocks"] == 2 + job = manifest["job"] + assert job["modules"] == ["gcc"] + assert job["env_setup"] == ["source /opt/x/setvars.sh"] + assert job["output"] == "/x/%j.out" + + +def test_write_run_manifest_inline_shape(tmp_path) -> None: + cfg = pips.PipsConfig(threads_per_rank=2) + command, env = pips.build_pips_command("/d", str(tmp_path), cfg, n_blocks=2) + pips.write_run_manifest(tmp_path, config=cfg, command=command, env=env, n_blocks=2) + manifest = json.loads((tmp_path / "pips.run.json").read_text()) + assert set(manifest) == { + "linopy_version", + "created", + "n_blocks", + "config", + "command", + "env", + } + assert manifest["config"] == asdict(cfg) + assert manifest["command"] == command + assert manifest["env"] == env + + +@pytest.mark.skipif( + "pips" not in available_solvers, + reason="requires the PIPS driver (set PIPS_BINARY)", +) +def test_doctor_ok() -> None: + report = pips.doctor() + assert "OK" in report + assert "objective" in report