diff --git a/bluemath_tk/core/resources.py b/bluemath_tk/core/resources.py new file mode 100644 index 0000000..0991fc1 --- /dev/null +++ b/bluemath_tk/core/resources.py @@ -0,0 +1,231 @@ +""" +Project: BlueMath_tk +Sub-Module: core.resources +Author: GeoOcean Research Group, Universidad de Cantabria +Repository: https://github.com/GeoOcean/BlueMath_tk.git +Status: Under development (Working) + +What this machine has, and how to share it out: cores, free memory and GPU +(:func:`machine`), the part a job may take (:data:`USAGE`), thread limits for +this process and for child processes, and a pool of single-purpose processes +(:func:`run_processes`). +""" + +from __future__ import annotations + +import contextlib +import logging +import multiprocessing +import os +from collections.abc import Callable, Iterator, Sequence +from concurrent.futures import ProcessPoolExecutor, as_completed +from dataclasses import dataclass +from typing import Any + +import psutil + +logger = logging.getLogger(__name__) + +#: Fraction of the machine a job may take (cores, free memory, GPU memory). +USAGE = 0.75 + +GIB = 1024**3 + +_THREAD_ENV_VARS = ( + "OMP_NUM_THREADS", + "OPENBLAS_NUM_THREADS", + "MKL_NUM_THREADS", + "NUMEXPR_NUM_THREADS", + "VECLIB_MAXIMUM_THREADS", +) + + +@dataclass(frozen=True) +class Machine: + """ + Resources of this machine, as seen now. + + Attributes + ---------- + cores : int + CPU cores this process may run on. + memory : int + Free memory (bytes). + memory_total : int + Total memory (bytes). + gpu : str or None + CUDA device name, ``None`` without one. + gpu_memory : int + Memory of the CUDA device (bytes), 0 without one. + """ + + cores: int + memory: int + memory_total: int + gpu: str | None = None + gpu_memory: int = 0 + + def usable(self, usage: float = USAGE) -> tuple[int, int]: + """ + Cores and memory a job may take. + + Parameters + ---------- + usage : float, optional + Fraction of the cores and of the *free* memory. Default + :data:`USAGE`. + + Returns + ------- + tuple[int, int] + Cores (at least 1) and memory (bytes). + """ + return max(1, int(self.cores * usage)), int(self.memory * usage) + + def __str__(self) -> str: + gpu = f"GPU {self.gpu} ({self.gpu_memory / GIB:.0f} GiB)" if self.gpu else "no GPU" + return ( + f"{self.cores} cores, {self.memory / GIB:.0f} GiB free of " + f"{self.memory_total / GIB:.0f} GiB, {gpu}" + ) + + +def machine() -> Machine: + """ + Cores, free memory and CUDA device of this machine. + + Returns + ------- + Machine + The resources seen now (cores this process may use, e.g. inside a + SLURM allocation). + """ + + cores = len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else os.cpu_count() + vm = psutil.virtual_memory() + gpu, gpu_memory = None, 0 + with contextlib.suppress(ImportError): + import torch + + if torch.cuda.is_available(): + props = torch.cuda.get_device_properties(0) + gpu, gpu_memory = props.name, int(props.total_memory) + + return Machine(int(cores or 1), int(vm.available), int(vm.total), gpu, gpu_memory) + + +@contextlib.contextmanager +def child_threads(n: int) -> Iterator[None]: + """Limit the BLAS/OpenMP threads of the processes started inside the block to *n*.""" + + previous = {name: os.environ.get(name) for name in _THREAD_ENV_VARS} + try: + for name in _THREAD_ENV_VARS: + os.environ[name] = str(n) + yield + finally: + for name, value in previous.items(): + if value is None: + os.environ.pop(name, None) + else: + os.environ[name] = value + + +@contextlib.contextmanager +def limit_threads(n: int) -> Iterator[None]: + """Limit the BLAS/OpenMP (and PyTorch, if loaded) threads of this process to *n* inside the block.""" + + from threadpoolctl import threadpool_limits + + torch = None + with contextlib.suppress(ImportError): + import torch + previous = torch.get_num_threads() if torch else None + try: + with threadpool_limits(limits=n): + if torch: + torch.set_num_threads(n) + yield + finally: + if torch: + torch.set_num_threads(previous) + + +def _init_child(threads: int, log_file: str | None) -> None: + """Thread limit (PyTorch included) and log file of a pool process.""" + + with contextlib.suppress(ImportError): + import torch + + torch.set_num_threads(threads) + if log_file: + handler = logging.FileHandler(log_file) + handler.setFormatter( + logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s") + ) + logging.getLogger().addHandler(handler) + logging.getLogger().setLevel(logging.INFO) + + +def run_processes( + fn: Callable[..., Any], + jobs: Sequence[tuple[Any, ...]], + *, + workers: int, + threads: int = 1, + log_file: str | None = None, +) -> list[Any]: + """ + ``fn(*job)`` for every job, in *workers* processes of *threads* threads each. + + Processes are started with ``spawn`` (safe after CUDA or threads were + used in this process); *fn* must be importable (a module-level function). + With one worker the jobs run here, in order, limited to *threads*. + + ``spawn`` re-imports the main module in every process: a *script* that + calls this (directly or through a library) must run its work under + ``if __name__ == "__main__":``, or each process runs the script again. + Notebooks need nothing. + + Parameters + ---------- + fn : callable + Module-level function. + jobs : sequence of tuple + Positional arguments of each call. + workers : int + Processes. + threads : int, optional + BLAS/OpenMP/PyTorch threads per process. Default is 1. + log_file : str, optional + File the processes log to (``INFO``). + + Returns + ------- + list + Results, in the order of *jobs*. + """ + + if workers <= 1 or len(jobs) <= 1: + with limit_threads(threads): + return [fn(*job) for job in jobs] + + results: list[Any] = [None] * len(jobs) + with ( + child_threads(threads), + ProcessPoolExecutor( + max_workers=min(workers, len(jobs)), + mp_context=multiprocessing.get_context("spawn"), + initializer=_init_child, + initargs=(threads, log_file), + ) as pool, + ): + futures = {pool.submit(fn, *job): i for i, job in enumerate(jobs)} + for done, future in enumerate(as_completed(futures), start=1): + results[futures[future]] = future.result() + logger.info("%s: %d/%d jobs done", getattr(fn, "__name__", "job"), done, len(jobs)) + + return results + + +__all__ = ["GIB", "USAGE", "Machine", "child_threads", "limit_threads", "machine", "run_processes"] diff --git a/bluemath_tk/waves/hywaves/__init__.py b/bluemath_tk/waves/hywaves/__init__.py index 189c1d3..86a5741 100644 --- a/bluemath_tk/waves/hywaves/__init__.py +++ b/bluemath_tk/waves/hywaves/__init__.py @@ -12,7 +12,17 @@ contributions. """ -from . import metamodel, offshore, output_points, plots, reconstruction, sectors, wind_correction +from . import ( + metamodel, + offshore, + output_points, + plots, + reconstruction, + resources, + sectors, + store, + wind_correction, +) __all__ = [ "metamodel", @@ -20,6 +30,8 @@ "output_points", "plots", "reconstruction", + "resources", "sectors", + "store", "wind_correction", ] diff --git a/bluemath_tk/waves/hywaves/metamodel.py b/bluemath_tk/waves/hywaves/metamodel.py index e84bc8d..00b0d33 100644 --- a/bluemath_tk/waves/hywaves/metamodel.py +++ b/bluemath_tk/waves/hywaves/metamodel.py @@ -483,6 +483,7 @@ def fit_gp( pcas: Mapping[str, PCA], directional_variables: list[str], epochs: int = 1000, + device: str | None = None, ) -> ExactGPInterpolation: """ Exact GP from forcing (MDA centroids) to all stacked PCs. @@ -497,6 +498,8 @@ def fit_gp( Directional columns of *centroids*. epochs : int, optional GP optimisation iterations. Default is 1000. + device : str, optional + PyTorch device; default CUDA when available. Returns ------- @@ -504,7 +507,7 @@ def fit_gp( The fitted GP. """ - gp = ExactGPInterpolation(epochs=epochs) + gp = ExactGPInterpolation(epochs=epochs, device=device) gp.fit( subset_data=centroids, subset_directional_variables=directional_variables, diff --git a/bluemath_tk/waves/hywaves/reconstruction.py b/bluemath_tk/waves/hywaves/reconstruction.py index ecd44f4..43529bf 100644 --- a/bluemath_tk/waves/hywaves/reconstruction.py +++ b/bluemath_tk/waves/hywaves/reconstruction.py @@ -7,7 +7,7 @@ contribution becomes a JONSWAP x Cartwright spectrum, the spectra are summed over goals and partitions, and bulk parameters are computed from the sum (:func:`sum_partition_spectra`). With a wind correction, one more -contribution, the local wind sea, is added per site batch before the sum +contribution, the local wind sea, is added per block before the sum (see :mod:`.wind_correction`). The sum reproduces a model run forced by every goal at once only if each @@ -21,16 +21,16 @@ :func:`bluemath_tk.wrappers.swan.swan_utils.boundary_par_table`). The spectral construction allocates the full ``(time, sites, goal x partition, -freq, dir)`` array, so long periods or many sites are processed in site -batches sized to a memory budget (:func:`site_batch_size`), optionally across -worker processes (:func:`sum_partition_spectra_batched`). +freq, dir)`` array (its peak was measured at 6.55x the float32 array, see +:func:`site_sum_bytes`). Every site and hour is independent, so the sum runs +in site x time blocks sized to a memory budget (:func:`batch_shape`), +optionally across worker processes (:func:`sum_partition_spectra_batched`); +the budget and the workers come from :mod:`.resources`. """ from __future__ import annotations -import contextlib import logging -import os import time from collections.abc import Mapping from typing import Any @@ -38,6 +38,8 @@ import numpy as np import xarray as xr +from ...core.resources import child_threads + logger = logging.getLogger(__name__) #: Frequency (Hz) and direction (deg) grids for the summed spectra. @@ -47,14 +49,13 @@ #: Spectral partitions expected in an offshore dataset (``phs0``...). SPECTRAL_PARTITIONS = (0, 1, 2, 3) -#: Default memory budget for one :func:`sum_partition_spectra` call. -DEFAULT_MEMORY_BUDGET_BYTES = 12 * 1024**3 - # wavespectra's construct_partition (jonswap x cartwright, then the sum) keeps -# several temporaries of the size of the full (time, sites, parts, freq, dir) -# array alive at once; peaks of at least ~5x that array were measured (a run -# sized for 12 GiB reached 55+ GiB), so batches are sized with this margin. -PEAK_MEMORY_MULTIPLIER = 8 +# several temporaries of the size of the float32 (time, sites, parts, freq, +# dir) array alive at once, and the stats add a few of the summed spectrum: +# measured peaks (2026-10-15) were 6.35x per part + 2.3x per summed spectrum, +# sized here with a margin. +PEAK_MEMORY_PER_PART = 8 +PEAK_MEMORY_PER_SUM = 4 def partition_ids(offshore: xr.Dataset, spectral: bool) -> list[int]: @@ -93,39 +94,63 @@ def partition_ids(offshore: xr.Dataset, spectral: bool) -> list[int]: return list(SPECTRAL_PARTITIONS) -def site_batch_size( - n_times: int, - n_parts: int, - budget_bytes: int = DEFAULT_MEMORY_BUDGET_BYTES, - n_freq: int = len(SPECTRAL_FREQ), - n_dir: int = len(SPECTRAL_DIR), -) -> int: +def site_sum_bytes(n_times: int, n_parts: int) -> int: """ - Sites per :func:`sum_partition_spectra` call that fit a memory budget. + Peak memory of :func:`sum_partition_spectra` per site (bytes). + + :data:`PEAK_MEMORY_PER_PART` float32 spectra per contribution plus + :data:`PEAK_MEMORY_PER_SUM` for the summed spectrum and its stats. Parameters ---------- n_times : int - Time steps in the call. + Time steps summed. n_parts : int - Goals x partitions summed. - budget_bytes : int, optional - Memory budget. Default is :data:`DEFAULT_MEMORY_BUDGET_BYTES`. - n_freq, n_dir : int, optional - Spectral grid size. + Contributions summed per site. Returns ------- int - At least 1. + Bytes. """ - # float32 spectra (see sum_partition_spectra), times the peak multiplier - bytes_per_site = ( - max(1, n_times) * max(1, n_parts) * n_freq * n_dir * 4 * PEAK_MEMORY_MULTIPLIER - ) + spectrum = max(1, n_times) * len(SPECTRAL_FREQ) * len(SPECTRAL_DIR) * 4 + return spectrum * (PEAK_MEMORY_PER_PART * max(1, n_parts) + PEAK_MEMORY_PER_SUM) - return max(1, budget_bytes // bytes_per_site) + +def spectra_bytes(n_times: int, n_sites: int) -> int: + """Size of the summed float32 spectra ``efth`` of *n_sites* over *n_times* steps.""" + + return n_times * n_sites * len(SPECTRAL_FREQ) * len(SPECTRAL_DIR) * 4 + + +def batch_shape(n_times: int, n_parts: int, budget_bytes: int) -> tuple[int, int]: + """ + Sites and time steps per :func:`sum_partition_spectra` call within a memory budget. + + Whole time series of as many sites as fit; when one site does not fit, + one site over as many steps as fit (every step is independent). + + Parameters + ---------- + n_times : int + Time steps to sum. + n_parts : int + Contributions summed per site. + budget_bytes : int + Memory budget. + + Returns + ------- + tuple[int, int] + Sites and steps per call, each at least 1. + """ + + per_site = site_sum_bytes(n_times, n_parts) + if per_site <= budget_bytes: + return budget_bytes // per_site, max(1, n_times) + + return 1, max(1, budget_bytes // site_sum_bytes(1, n_parts)) def sum_partition_spectra( @@ -162,7 +187,7 @@ def sum_partition_spectra( Bulk parameters on ``(time, sites)`` and the optional spectra. """ - from wavespectra.construct import STATS, construct_partition + from wavespectra.construct import construct_partition parts = cube.sel(sites=site_ids).stack(part=("goal", "partition")) # contributions never predicted (e.g. a goal x partition combination @@ -199,7 +224,7 @@ def sum_partition_spectra( if load: spectra = spectra.load() - stats = combined.spec.stats(list(STATS)) + stats = combined.spec.stats(sorted(set(variables.values()))) result = xr.Dataset({name: stats[stat] for name, stat in variables.items()}) if load: result = result.load() @@ -222,11 +247,19 @@ def _add_wind_sea( return xr.merge([merged, contrib.drop_vars("k")]), k -def _sum_site_batches( +def _time_axis(chunks: list[xr.Dataset]) -> np.ndarray: + return np.unique(np.concatenate([c["time"].values for c in chunks])) + + +def _n_parts(chunks: list[xr.Dataset], wind_correction: Any) -> int: + return sum(c.sizes["goal"] * c.sizes["partition"] for c in chunks) + int( + wind_correction is not None + ) + + +def _sum_blocks( chunks: list[xr.Dataset], site_ids: list[str], - n_times: int, - n_parts: int, variables: Mapping[str, str], save_spectra: bool, budget_bytes: int, @@ -234,39 +267,49 @@ def _sum_site_batches( label: str = "", wind_correction: Mapping[str, Any] | None = None, ) -> tuple[xr.Dataset, xr.Dataset | None]: - """Merge contribution chunks and sum them, site batch by site batch.""" + """Merge contribution chunks and sum them, in site x time blocks within the budget.""" - if wind_correction is not None: - n_parts += 1 - batch = site_batch_size(n_times, n_parts, budget_bytes) - n_batches = -(-len(site_ids) // batch) + times = _time_axis(chunks) + n_sites, n_steps = batch_shape(len(times), _n_parts(chunks, wind_correction), budget_bytes) + n_blocks = -(-len(site_ids) // n_sites) * -(-len(times) // n_steps) results, spectras = [], [] t0 = time.monotonic() - for k, i in enumerate(range(0, len(site_ids), batch), start=1): - t_batch = time.monotonic() - sites = site_ids[i : i + batch] - merged = xr.merge([c.sel(sites=sites) for c in chunks], join="outer", compat="no_conflicts") - wind_k = None - if wind_correction is not None: - merged, wind_k = _add_wind_sea(merged, wind_correction) - result, spectra = sum_partition_spectra( - merged, sites, variables, save_spectra=save_spectra, load=load + k = 0 + for i in range(0, len(site_ids), n_sites): + sites = site_ids[i : i + n_sites] + merged = xr.merge( + [c.sel(sites=sites) for c in chunks], join="outer", compat="no_conflicts" ) - if wind_k is not None: - result["wind_k"] = wind_k - results.append(result) - if save_spectra and spectra is not None: - spectras.append(spectra) + site_results, site_spectra = [], [] + for j in range(0, merged.sizes["time"], n_steps): + t_block = time.monotonic() + block = merged.isel(time=slice(j, j + n_steps)) + wind_k = None + if wind_correction is not None: + block, wind_k = _add_wind_sea(block, wind_correction) + result, spectra = sum_partition_spectra( + block, sites, variables, save_spectra=save_spectra, load=load + ) + if wind_k is not None: + result["wind_k"] = wind_k + site_results.append(result) + if save_spectra and spectra is not None: + site_spectra.append(spectra) + k += 1 + logger.info( + "%ssum block %d/%d (%d sites x %d steps) in %.1fs (%.1fs elapsed)", + f"[{label}] " if label else "", + k, + n_blocks, + len(sites), + block.sizes["time"], + time.monotonic() - t_block, + time.monotonic() - t0, + ) + results.append(xr.concat(site_results, dim="time")) + if site_spectra: + spectras.append(xr.concat(site_spectra, dim="time")) del merged - logger.info( - "%ssum batch %d/%d (%d sites) in %.1fs (%.1fs elapsed)", - f"[{label}] " if label else "", - k, - n_batches, - len(sites), - time.monotonic() - t_batch, - time.monotonic() - t0, - ) return ( xr.concat(results, dim="sites"), @@ -275,40 +318,15 @@ def _sum_site_batches( # --------------------------------------------------------------------------- -# Parallel site batches +# Parallel blocks # --------------------------------------------------------------------------- # -# Summation is embarrassingly parallel across sites and pure CPU. Worker +# Summation is independent per site and per time step, and pure CPU. Worker # processes use the "spawn" start method: by this point the caller has usually # run GP inference on CUDA, and fork()ed children inheriting a CUDA context # hang. Each worker is limited to one BLAS/OpenMP thread (the parallelism is # the processes), and gets an equal share of the memory budget. -_BLAS_THREAD_ENV_VARS = ( - "OMP_NUM_THREADS", - "OPENBLAS_NUM_THREADS", - "MKL_NUM_THREADS", - "NUMEXPR_NUM_THREADS", - "VECLIB_MAXIMUM_THREADS", -) - - -@contextlib.contextmanager -def _single_threaded_blas_for_children(): - """Set BLAS/OpenMP threads to 1 in the environment workers inherit.""" - - previous = {name: os.environ.get(name) for name in _BLAS_THREAD_ENV_VARS} - try: - for name in _BLAS_THREAD_ENV_VARS: - os.environ[name] = "1" - yield - finally: - for name, value in previous.items(): - if value is None: - os.environ.pop(name, None) - else: - os.environ[name] = value - def _init_worker(log_file: str | None) -> None: """Send this module's log records to *log_file* in a worker process.""" @@ -322,39 +340,51 @@ def _init_worker(log_file: str | None) -> None: logger.setLevel(logging.INFO) +def _work_items( + site_ids: list[str], times: np.ndarray, workers: int +) -> list[tuple[list[str], np.ndarray]]: + """Site groups x time spans for *workers*: sites first, time too when sites are fewer.""" + + groups = [g.tolist() for g in np.array_split(np.array(site_ids, dtype=object), min(workers, len(site_ids)))] + spans = np.array_split(times, max(1, min(len(times), workers // len(groups)))) + + return [(g, span) for g in groups for span in spans if len(span)] + + def sum_partition_spectra_batched( chunks: list[xr.Dataset], site_ids: list[str], - n_times: int, - n_parts: int, variables: Mapping[str, str], - save_spectra: bool = False, - budget_bytes: int = DEFAULT_MEMORY_BUDGET_BYTES, + *, + budget_bytes: int, workers: int = 1, + save_spectra: bool = False, load: bool = True, log_file: str | None = None, wind_correction: Mapping[str, Any] | None = None, ) -> tuple[xr.Dataset, xr.Dataset | None]: """ - :func:`sum_partition_spectra` over many sites, in memory-sized batches. + :func:`sum_partition_spectra` over many sites and steps, in memory-sized blocks. + + The work is split into site groups (and time spans when there are fewer + sites than workers), one per worker; each worker sums its share in site x + time blocks of its part of the budget (:func:`batch_shape`). Parameters ---------- chunks : list of xr.Dataset One contribution per goal x partition, each on - ``(goal, partition, time, sites)``; merged per site batch. + ``(goal, partition, time, sites)``; merged per block. site_ids : list of str All sites. - n_times, n_parts : int - See :func:`site_batch_size`. variables : mapping Output name -> wavespectra stat. - save_spectra : bool, optional - Also return summed spectra. - budget_bytes : int, optional - Total memory budget, shared by the workers. + budget_bytes : int + Memory budget of the sum, shared by the workers. workers : int, optional Worker processes; 1 (default) runs in this process. + save_spectra : bool, optional + Also return summed spectra. load : bool, optional Load results into memory (always True in workers). log_file : str, optional @@ -363,8 +393,8 @@ def sum_partition_spectra_batched( Keyword arguments of :func:`~bluemath_tk.waves.hywaves.wind_correction.wind_sea_contribution` (``correction``, ``fetch``, optional ``wind``, ``spr``, ``tp``): adds - the local wind sea per site batch, and ``wind_k`` (``K``, 1 where - there is no correction) to the bulk output. Chunks then also carry + the local wind sea per block, and ``wind_k`` (``K``, 1 where there is + no correction) to the bulk output. Chunks then also carry ``ws_weight`` (and ``u10``, ``u10dir`` without ``wind``). Returns @@ -374,67 +404,58 @@ def sum_partition_spectra_batched( """ if workers <= 1: - return _sum_site_batches( - chunks, - site_ids, - n_times, - n_parts, - variables, - save_spectra, - budget_bytes, - load, + return _sum_blocks( + chunks, site_ids, variables, save_spectra, budget_bytes, load, wind_correction=wind_correction, ) import multiprocessing from concurrent.futures import ProcessPoolExecutor, as_completed - per_worker = max(1, budget_bytes // workers) - groups = [ - g.tolist() - for g in np.array_split(np.array(site_ids, dtype=object), workers) - if len(g) - ] + items = _work_items(site_ids, _time_axis(chunks), workers) + per_worker = max(1, budget_bytes // len(items)) logger.info( "sum: %d sites across %d workers (%.2f GiB each)", len(site_ids), - len(groups), + len(items), per_worker / 1024**3, ) - results: list[xr.Dataset] = [] - spectras: list[xr.Dataset] = [] + out: dict[int, tuple[xr.Dataset, xr.Dataset | None]] = {} with ( - _single_threaded_blas_for_children(), + child_threads(1), ProcessPoolExecutor( - max_workers=len(groups), + max_workers=len(items), mp_context=multiprocessing.get_context("spawn"), initializer=_init_worker, initargs=(log_file,), ) as executor, ): - futures: dict[Any, str] = {} - for k, sites in enumerate(groups, start=1): - label = f"group {k}/{len(groups)}" + futures = {} + for k, (sites, span) in enumerate(items): + span_slice = slice(span[0], span[-1]) future = executor.submit( - _sum_site_batches, - [c.sel(sites=sites) for c in chunks], + _sum_blocks, + [c.sel(sites=sites, time=span_slice) for c in chunks], sites, - n_times, - n_parts, variables, save_spectra, per_worker, True, - label, + f"item {k + 1}/{len(items)}", wind_correction, ) - futures[future] = label + futures[future] = k for future in as_completed(futures): - result, spectra = future.result() - results.append(result) - if save_spectra and spectra is not None: - spectras.append(spectra) - logger.info("%s done", futures[future]) + out[futures[future]] = future.result() + + by_group: dict[str, list[int]] = {} + for k, (sites, _) in enumerate(items): + by_group.setdefault(sites[0], []).append(k) + results, spectras = [], [] + for ks in by_group.values(): + results.append(xr.concat([out[k][0] for k in ks], dim="time")) + if save_spectra: + spectras.append(xr.concat([out[k][1] for k in ks], dim="time")) return ( xr.concat(results, dim="sites"), diff --git a/bluemath_tk/waves/hywaves/resources.py b/bluemath_tk/waves/hywaves/resources.py new file mode 100644 index 0000000..134b195 --- /dev/null +++ b/bluemath_tk/waves/hywaves/resources.py @@ -0,0 +1,388 @@ +""" +Project: BlueMath_tk +Sub-Module: waves.hywaves.resources +Author: GeoOcean Research Group, Universidad de Cantabria +Repository: https://github.com/GeoOcean/BlueMath_tk.git +Status: Under development (Working) + +How much of the machine a HyWaves fit or reconstruction uses. + +The plans are **proposals**: :func:`plan_fit` and :func:`plan_reconstruction` +size the work to at most :data:`~bluemath_tk.core.resources.USAGE` (75%) of +the cores and free memory of this machine +(:func:`~bluemath_tk.core.resources.machine`) and return a small dataclass. +The caller shows it, may edit any field, and passes it to the run:: + + res = plan_reconstruction(sites=4971, months=12, parts=40) + res.workers = 20 # the user decides + # ... goal contributions with res.device / res.threads, + # sum_partition_spectra_batched(..., workers=res.workers, + # budget_bytes=sum_budget(res.memory, held, res.workers)) + +How the work is split, and why (measured on a 48-core Xeon Gold 5418Y node, +2026-10-15): + +- **Fit**: every (goal, metamodel set, MDA size) PCA + GP fit is independent + and small (one exact GP per principal component on a few hundred points). + PyTorch uses threads poorly on such small problems (36 threads were 1.9x + faster than one), so on CPU the fits run as parallel single-thread + processes (:class:`FitResources`); a fit needs <= ~1.2 GiB. +- **Reconstruction**, one call = *months* x *sites*: + + 1. goal contributions: GP + PCA inverse per goal and partition, threaded, + in the calling process (``device``, ``threads``); + 2. spectral sum: JONSWAP x Cartwright spectra summed per site and hour in + ``workers`` single-thread processes. Every site and hour is + independent, so the sum is split into site x time blocks + (:func:`~bluemath_tk.waves.hywaves.reconstruction.batch_shape`) and + each worker needs 0.5-2 GiB whatever the call. The sum is bound by + memory bandwidth: small blocks are fastest, and 8 / 18 / 36 workers + are 7.1 / 10.9 / 14.1x faster than one. + + :func:`call_memory` is the memory of one call, :func:`hours` a rough wall + time from the measured rates. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +from ...core.resources import GIB, USAGE, Machine, machine +from .reconstruction import spectra_bytes + +#: Longest calendar month of an hourly hindcast (time steps). +STEPS_PER_MONTH = 744 + +#: Most months per reconstruction call: a year on CPU (one GP call and one +#: worker pool, ~6 s to start, per call); on a GPU, the value checked in +#: spectral mode with all goals. +MONTHS_PER_CALL = {"cpu": 12, "cuda": 4} + +# Goal contributions (hs, tp, dir, spr; float64) are held by this process and, +# split by site, by the sum workers: twice their size. +_CONTRIBUTION_BYTES = 4 * 8 * 2 + +# Saved spectra: the workers' results, their concatenation and the copy made +# when stored. +_SPECTRA_COPIES = 3 + +#: Spectral-sum block memory per worker when there is room: 1-5 sites per +#: block summed fastest (one month, 10 contributions; 17 sites was 40% slower). +SUM_BYTES_PER_WORKER = 2 * GIB + +#: The least a worker gets (one site over ~190 steps of 40 contributions). +MIN_SUM_BYTES_PER_WORKER = GIB // 2 + +#: Memory of a sum worker process itself (imports, its share of the inputs). +WORKER_BYTES = 3 * GIB // 4 + +# Xeon Gold 5418Y (2026-10-15). One core: building and summing the spectra per +# step x site x contribution, plus the bulk stats per step x site. Workers +# slow each other down (memory bandwidth): speedup = w / (1 + _CONTENTION w), +# measured 7.1 / 10.9 / 14.1 with 8 / 18 / 36 workers. GP + PCA (36 threads): +# a fixed part per month x contribution (models, kernel) plus one per site. +_SUM_SECONDS_PER_PART = 2.2e-5 +_SUM_SECONDS_PER_SITE = 1.0e-4 +_CONTENTION = 0.04 +_GP_SECONDS_PER_MONTH_PART = 0.125 +_GP_SECONDS_PER_SITE_MONTH_PART = 1.9e-4 + + +@dataclass +class ReconstructionResources: + """ + Resources of a reconstruction, from :func:`plan_reconstruction`. + + Edit any field before using it, e.g. fewer ``workers`` to leave cores to + other users of a shared machine. + + Attributes + ---------- + device : str + PyTorch device of the GP prediction (``"cpu"`` or ``"cuda"``). + threads : int + Threads of the GP prediction and PCA inverse (calling process). + workers : int + Processes of the spectral sum, one thread each. + months : int + Calendar months per call (more months: fewer GP calls and worker + pools, more memory for the contributions). + sites : int or None + Sites per call; ``None`` for every site (fewer only when memory is + short, e.g. many sites with ``save_spectra``). + memory : int + Most memory a call may take (bytes). The contributions and saved + spectra are held first; the sum workers get up to + :data:`SUM_BYTES_PER_WORKER` each of what is left. + """ + + device: str = "cpu" + threads: int = 1 + workers: int = 1 + months: int = 1 + sites: int | None = None + memory: int = 8 * GIB + + def __str__(self) -> str: + sites = "every site" if self.sites is None else f"{self.sites} sites" + return ( + f"device {self.device}, {self.threads} threads (GP + PCA), {self.workers} workers " + f"(spectral sum), {self.months} month(s) x {sites} per call, " + f"{self.memory / GIB:.0f} GiB" + ) + + +@dataclass +class FitResources: + """ + Resources of the metamodel fits, from :func:`plan_fit`. + + Edit any field before using it. + + Attributes + ---------- + device : str + PyTorch device of the GP fits (``"cpu"`` or ``"cuda"``). + workers : int + Fits run at the same time, one process each. + threads : int + Threads per process (1 on CPU: the parallelism is the processes). + """ + + device: str = "cpu" + workers: int = 1 + threads: int = 1 + + def __str__(self) -> str: + return f"device {self.device}, {self.workers} process(es) x {self.threads} thread(s)" + + +def call_memory( + months: int, + sites: int, + parts: int, + *, + workers: int, + save_spectra: bool = False, + sum_per_worker: int = MIN_SUM_BYTES_PER_WORKER, +) -> dict[str, int]: + """ + Memory of one reconstruction call (bytes), by part. + + - ``contributions``: the predicted ``hs``, ``tp``, ``dir``, ``spr`` of + every contribution (float64), held by the caller and, split, by the + workers; + - ``spectra``: the summed spectra when kept (float32; workers' results, + their concatenation and the stored copy); + - ``sum``: each worker process and its spectral blocks. + + Parameters + ---------- + months, sites : int + Calendar months and sites of the call. + parts : int + Contributions summed per site (goals x partitions, + 1 with the wind + correction). + workers : int + Sum processes. + save_spectra : bool, optional + The summed spectra are kept. + sum_per_worker : int, optional + Block memory per worker; default the least it can get. + + Returns + ------- + dict + ``contributions``, ``spectra``, ``sum`` (workers and their blocks) + and their ``total``. + """ + + steps = months * STEPS_PER_MONTH + out = { + "contributions": steps * sites * parts * _CONTRIBUTION_BYTES, + "spectra": _SPECTRA_COPIES * spectra_bytes(steps, sites) if save_spectra else 0, + "sum": workers * (WORKER_BYTES + sum_per_worker), + } + out["total"] = sum(out.values()) + + return out + + +def sum_budget(memory: int, held: int, workers: int) -> int: + """ + Memory budget of one spectral sum (bytes). + + Up to :data:`SUM_BYTES_PER_WORKER` per worker (blocks this size summed + fastest), within what *held* and the worker processes leave of *memory*. + + Parameters + ---------- + memory : int + Memory of the whole call (``ReconstructionResources.memory``). + held : int + Memory already taken by the contributions (and saved spectra). + workers : int + Sum processes. + + Returns + ------- + int + Bytes for ``sum_partition_spectra_batched(budget_bytes=...)``; 0 + when nothing is left (the blocks then shrink to one site-hour). + """ + + return max(0, min(workers * SUM_BYTES_PER_WORKER, memory - held - workers * WORKER_BYTES)) + + +def hours(site_months: int, months: int, parts: int, workers: int) -> float: + """ + Rough wall time (hours) of a reconstruction. + + From rates measured on one 48-core Xeon Gold 5418Y node (GP + PCA with + 36 threads, the sum slowed by memory bandwidth as workers grow); other + machines differ, so take it as an order of magnitude. + + Parameters + ---------- + site_months : int + Site-months to compute. + months : int + Calendar months they span. + parts : int + Contributions per site (goals x partitions, + 1 with wind). + workers : int + Sum processes. + + Returns + ------- + float + Hours. + """ + + workers = max(1, workers) + speedup = workers / (1 + _CONTENTION * workers) + per_site_month = STEPS_PER_MONTH * (_SUM_SECONDS_PER_PART * parts + _SUM_SECONDS_PER_SITE) + seconds = site_months * (per_site_month / speedup + _GP_SECONDS_PER_SITE_MONTH_PART * parts) + + return (seconds + months * parts * _GP_SECONDS_PER_MONTH_PART) / 3600 + + +def _halvings(n: int) -> list[int]: + out = [n] + while out[-1] > 1: + out.append(-(-out[-1] // 2)) + return out + + +def plan_reconstruction( + sites: int, + months: int, + parts: int, + *, + save_spectra: bool = False, + host: Machine | None = None, + usage: float = USAGE, +) -> ReconstructionResources: + """ + Proposed resources to reconstruct *months* x *sites* on this machine. + + One sum worker per usable core, and the largest call that fits *usage* + of the free memory (:func:`call_memory`, with the least block memory per + worker): every site and up to :data:`MONTHS_PER_CALL` months, then fewer + months, then fewer sites, and only then fewer workers. + + Parameters + ---------- + sites, months : int + Sites and calendar months to compute. + parts : int + Contributions per site (see :func:`call_memory`). + save_spectra : bool, optional + The summed spectra are kept. + host : Machine, optional + The machine; default :func:`bluemath_tk.core.resources.machine`. + usage : float, optional + Fraction of the machine to use. Default :data:`USAGE`. + + Returns + ------- + ReconstructionResources + + Raises + ------ + MemoryError + If one site and one month do not fit. + """ + + host = host or machine() + cores, memory = host.usable(usage) + device = "cuda" if host.gpu else "cpu" + sites, months = max(1, sites), max(1, months) + + def fits(n_sites: int, n_months: int, workers: int) -> bool: + need = call_memory(n_months, n_sites, parts, workers=workers, save_spectra=save_spectra) + return need["total"] <= memory + + for n_sites in _halvings(sites): + for n_months in range(min(MONTHS_PER_CALL[device], months), 0, -1): + if fits(n_sites, n_months, cores): + return ReconstructionResources( + device, cores, cores, n_months, None if n_sites == sites else n_sites, memory + ) + for workers in range(cores - 1, 0, -1): + if fits(1, 1, workers): + return ReconstructionResources(device, cores, workers, 1, 1, memory) + need = call_memory(1, 1, parts, workers=1, save_spectra=save_spectra)["total"] + raise MemoryError( + f"One site and one month need {need / GIB:.1f} GiB; {memory / GIB:.1f} GiB usable" + ) + + +def plan_fit( + jobs: int, job_memory: int, *, host: Machine | None = None, usage: float = USAGE +) -> FitResources: + """ + Proposed resources for *jobs* independent fits of *job_memory* bytes each. + + On CPU: one single-thread process per fit, as many at a time as *usage* + of the cores and of the free memory allow. With a GPU: one process on it, + with *usage* of the cores as threads. + + Parameters + ---------- + jobs : int + Fits to run. + job_memory : int + Memory of one fit (bytes). + host : Machine, optional + The machine; default :func:`bluemath_tk.core.resources.machine`. + usage : float, optional + Fraction of the machine to use. Default :data:`USAGE`. + + Returns + ------- + FitResources + """ + + host = host or machine() + cores, memory = host.usable(usage) + if host.gpu: + return FitResources("cuda", 1, cores) + + return FitResources("cpu", max(1, min(cores, jobs, memory // max(1, job_memory))), 1) + + +__all__ = [ + "MONTHS_PER_CALL", + "STEPS_PER_MONTH", + "FitResources", + "MIN_SUM_BYTES_PER_WORKER", + "ReconstructionResources", + "SUM_BYTES_PER_WORKER", + "WORKER_BYTES", + "call_memory", + "hours", + "plan_fit", + "plan_reconstruction", + "sum_budget", +] diff --git a/bluemath_tk/waves/hywaves/store.py b/bluemath_tk/waves/hywaves/store.py new file mode 100644 index 0000000..ac98aa9 --- /dev/null +++ b/bluemath_tk/waves/hywaves/store.py @@ -0,0 +1,246 @@ +""" +Project: BlueMath_tk +Sub-Module: waves.hywaves.store +Author: GeoOcean Research Group, Universidad de Cantabria +Repository: https://github.com/GeoOcean/BlueMath_tk.git +Status: Under development (Working) + +Cache of reconstructed waves: one zarr store per reconstruction setup, +filled one site x calendar month at a time. + +Every (site, hour) of a HyWaves reconstruction is independent of the others, +so results can be computed in any pieces and kept together. The store holds +the whole time axis of the offshore hindcast and every site of the setup: + +- arrays on ``(month, step, sites)`` (plus e.g. ``freq``, ``dir`` for + spectra), ``step`` = position inside the calendar month (padded with NaT to + the longest month), so **one chunk is one site and one calendar month**; + bulk arrays put all the sites of a month in one file (zarr sharding); +- ``done`` on ``(month, sites)``: which site-months are stored; +- the setup (``key``) and a ``fingerprint`` of its input files in the + attributes: a store whose setup, inputs, sites or time axis changed is + rebuilt, never reused. + +A request reads the site-months it needs; :meth:`ReconstructionStore.missing` +lists the ones to compute (always whole months), and +:meth:`ReconstructionStore.write` stores them. +""" + +from __future__ import annotations + +import json +import logging +import shutil +from collections.abc import Callable, Iterable, Mapping +from pathlib import Path +from typing import Any + +import numpy as np +import pandas as pd +import xarray as xr +import zarr +from filelock import FileLock + +logger = logging.getLogger(__name__) + +_NAT = np.iinfo(np.int64).min + + +class ReconstructionStore: + """ + Site x month cache of a reconstruction on the time axis of its hindcast. + + Parameters + ---------- + path : str or Path + The zarr store (created if missing, rebuilt if it does not match). + times : array-like of datetime64 + The full time axis of the offshore hindcast. + sites : sequence of str + Every site the store can hold. + variables : mapping + Variable name -> names of its extra dimensions (``()`` for bulk + parameters, ``("freq", "dir")`` for spectra). + coords : mapping, optional + Values of the extra dimensions, e.g. ``{"freq": ..., "dir": ...}``. + key : mapping, optional + The setup (JSON-serialisable); part of the validity check. + fingerprint : str, optional + Fingerprint of the input files; part of the validity check. + attrs : mapping or callable, optional + More attributes to record when the store is created (not checked), + e.g. the PCs per goal; a callable is only called then. + """ + + def __init__( + self, + path: str | Path, + times: Any, + sites: Iterable[str], + variables: Mapping[str, tuple[str, ...]], + *, + coords: Mapping[str, Any] | None = None, + key: Mapping[str, Any] | None = None, + fingerprint: str = "", + attrs: Mapping[str, Any] | Callable[[], Mapping[str, Any]] | None = None, + ) -> None: + self.path = Path(path) + self.sites = [str(s) for s in sites] + self.variables = {name: tuple(dims) for name, dims in variables.items()} + self.coords = {name: np.asarray(v) for name, v in (coords or {}).items()} + index = pd.DatetimeIndex(times) + self.month_labels = index.to_period("M") + self.months = [str(m) for m in self.month_labels.unique()] + steps = pd.Series(np.arange(len(index))).groupby(np.asarray(self.month_labels)).cumcount() + self._month_of = pd.Index(self.months).get_indexer(self.month_labels.astype(str)) + self._step_of = steps.to_numpy() + n_steps = int(self._step_of.max()) + 1 + grid = np.full((len(self.months), n_steps), _NAT, dtype=np.int64) + grid[self._month_of, self._step_of] = index.asi8 + self._time_grid = grid + self._site_index = pd.Index(self.sites) + meta = { + "key": json.dumps(key or {}, sort_keys=True, default=str), + "fingerprint": fingerprint, + "sites": self.sites, + "variables": {k: list(v) for k, v in self.variables.items()}, + } + self.group = self._open(meta, attrs) + + # ------------------------------------------------------------------ + # Creation and validity + # ------------------------------------------------------------------ + + def _matches(self, group: zarr.Group, meta: Mapping[str, Any]) -> bool: + stored = dict(group.attrs) + return all(stored.get(k) == v for k, v in meta.items()) and np.array_equal( + group["time"][:], self._time_grid + ) + + def _open(self, meta: dict[str, Any], attrs: Any) -> zarr.Group: + if self.path.exists(): + group = zarr.open_group(self.path, mode="r+") + if self._matches(group, meta): + return group + logger.warning("%s does not match this setup or its inputs: rebuilt", self.path.name) + shutil.rmtree(self.path) + self.path.parent.mkdir(parents=True, exist_ok=True) + group = zarr.open_group(self.path, mode="w") + n_months, n_steps = self._time_grid.shape + n_sites = len(self.sites) + group.create_array("time", data=self._time_grid, chunks=self._time_grid.shape) + group.create_array( + "done", shape=(n_months, n_sites), chunks=(n_months, n_sites), dtype="bool", fill_value=False + ) + for name, values in self.coords.items(): + group.create_array(name, data=values) + for name, extra in self.variables.items(): + extra_shape = tuple(len(self.coords[d]) for d in extra) + shape = (n_months, n_steps, n_sites, *extra_shape) + chunks = (1, n_steps, 1, *extra_shape) + group.create_array( + name, + shape=shape, + chunks=chunks, + shards=None if extra else (1, n_steps, n_sites), + dtype="float32", + fill_value=np.nan, + ) + extra = attrs() if callable(attrs) else (attrs or {}) + group.attrs.update({**meta, **extra}) + return group + + def _lock(self) -> FileLock: + return FileLock(self.path.with_suffix(".lock")) + + # ------------------------------------------------------------------ + # Requests + # ------------------------------------------------------------------ + + def months_of(self, time: Any) -> list[str]: + """Calendar months (``"2013-12"``) touched by a time selection.""" + index = pd.DatetimeIndex(self._time_grid[self._time_grid != _NAT].view("datetime64[ns]")) + selected = xr.DataArray(index, dims="time", coords={"time": index}).sel(time=time) + return [str(m) for m in pd.DatetimeIndex(np.atleast_1d(selected.values)).to_period("M").unique()] + + def missing(self, time: Any, sites: Iterable[str]) -> dict[str, list[str]]: + """Sites not stored yet for each month of the selection (months with none left omitted).""" + site_idx = self._site_index.get_indexer(list(sites)) + if (site_idx < 0).any(): + raise KeyError("Sites not in this store") + done = self.group["done"][:] + out = {} + for month in self.months_of(time): + m = self.months.index(month) + todo = [self.sites[i] for i in site_idx if not done[m, i]] + if todo: + out[month] = todo + return out + + def write(self, ds: xr.Dataset, months: Iterable[str]) -> None: + """ + Store whole months of *ds* (on ``time``, ``sites``) and mark them done. + + *ds* must cover every step of each month in *months* for its sites; + steps it lacks (e.g. hours with no wave energy) are stored as NaN. + """ + site_idx = self._site_index.get_indexer([str(s) for s in ds.sites.values]) + with self._lock(): + done = self.group["done"][:] + for month in months: + m = self.months.index(month) + steps = self._time_grid[m][self._time_grid[m] != _NAT].view("datetime64[ns]") + part = ds.reindex(time=steps) + for name, extra in self.variables.items(): + values = part[name].transpose("time", "sites", *extra).to_numpy().astype("float32", copy=False) + if extra: # one chunk per site: write each site's month + for j, i in enumerate(site_idx): + self.group[name][m, : len(steps), i] = values[:, j] + else: # one shard per month: read, update, write back + block = self.group[name][m] + block[: len(steps)][:, site_idx] = values + self.group[name][m] = block + done[m, site_idx] = True + self.group["done"][:] = done + + def invalidate(self, time: Any, sites: Iterable[str]) -> None: + """Mark the months of *time* as not stored for *sites* (they are recomputed).""" + site_idx = self._site_index.get_indexer(list(sites)) + with self._lock(): + done = self.group["done"][:] + for month in self.months_of(time): + done[self.months.index(month), site_idx] = False + self.group["done"][:] = done + + def read(self, time: Any, sites: Iterable[str]) -> xr.Dataset: + """ + Stored values of *sites* for a time selection, on ``(time, sites)``. + + Raises + ------ + KeyError + If a requested site-month is not stored. + """ + sites = [str(s) for s in sites] + if self.missing(time, sites): + raise KeyError("Requested site-months are not stored; compute them first") + site_idx = self._site_index.get_indexer(sites) + month_idx = [self.months.index(m) for m in self.months_of(time)] + times = self._time_grid[month_idx] + valid = times != _NAT + data = {} + for name, extra in self.variables.items(): + values = np.stack([self.group[name].oindex[m, :, site_idx] for m in month_idx]) + data[name] = (("time", "sites", *extra), values[valid]) + coords = {"time": times[valid].view("datetime64[ns]"), "sites": sites} + coords.update({d: self.coords[d] for extra in self.variables.values() for d in extra}) + return xr.Dataset(data, coords=coords).sel(time=time) + + def coverage(self) -> pd.DataFrame: + """Stored sites per month (months with at least one).""" + done = self.group["done"][:] + counts = pd.Series(done.sum(axis=1), index=self.months, name="sites") + return counts[counts > 0].to_frame() + + +__all__ = ["ReconstructionStore"] diff --git a/pyproject.toml b/pyproject.toml index 7c00b91..ce582c2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -48,6 +48,7 @@ dependencies = [ "prettytable", "numba", "csaps", + "filelock", ] [project.optional-dependencies] diff --git a/tests/core/test_resources.py b/tests/core/test_resources.py new file mode 100644 index 0000000..29a63c4 --- /dev/null +++ b/tests/core/test_resources.py @@ -0,0 +1,38 @@ +import os +import unittest + +from bluemath_tk.core.resources import GIB, Machine, child_threads, limit_threads, machine, run_processes + + +def _threads_seen(x: int) -> tuple[int, str | None]: + return x * 2, os.environ.get("OMP_NUM_THREADS") + + +class TestResources(unittest.TestCase): + def test_machine_and_usable(self): + host = machine() + self.assertGreaterEqual(host.cores, 1) + self.assertGreater(host.memory_total, 0) + cores, memory = Machine(48, 400 * GIB, 512 * GIB).usable(0.75) + self.assertEqual((cores, memory), (36, 300 * GIB)) + self.assertEqual(Machine(1, GIB, GIB).usable()[0], 1) + + def test_child_threads_restores(self): + before = os.environ.get("OMP_NUM_THREADS") + with child_threads(3): + self.assertEqual(os.environ["OMP_NUM_THREADS"], "3") + self.assertEqual(os.environ.get("OMP_NUM_THREADS"), before) + with limit_threads(1): + pass + + def test_run_processes_keeps_order(self): + jobs = [(i,) for i in range(5)] + here = run_processes(_threads_seen, jobs, workers=1) + self.assertEqual([r[0] for r in here], [0, 2, 4, 6, 8]) + pooled = run_processes(_threads_seen, jobs, workers=2, threads=1) + self.assertEqual([r[0] for r in pooled], [0, 2, 4, 6, 8]) + self.assertTrue(all(r[1] == "1" for r in pooled)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/waves/test_hywaves.py b/tests/waves/test_hywaves.py index 7c1edf4..a0e6d9d 100644 --- a/tests/waves/test_hywaves.py +++ b/tests/waves/test_hywaves.py @@ -377,10 +377,15 @@ def test_batched_equals_single_call(self): cube = self._cube([1.0, 2.0]) chunks = [cube.isel(goal=[i]) for i in range(2)] batched, _ = reconstruction.sum_partition_spectra_batched( - chunks, ["a", "b"], 2, 2, {"hs": "hs"}, budget_bytes=1 + chunks, ["a", "b"], {"hs": "hs"}, budget_bytes=1 ) single, _ = reconstruction.sum_partition_spectra(cube, ["a", "b"], {"hs": "hs"}) xr.testing.assert_allclose(batched, single) + pooled, spectra = reconstruction.sum_partition_spectra_batched( + chunks, ["a", "b"], {"hs": "hs"}, budget_bytes=1, workers=4, save_spectra=True + ) + xr.testing.assert_allclose(pooled.transpose(*single.dims), single) + self.assertEqual(dict(spectra.efth.sizes)["time"], single.sizes["time"]) def test_partitions_and_batch_size(self): """Partition ids per mode, and a batch size that shrinks with the problem.""" @@ -390,11 +395,13 @@ def test_partitions_and_batch_size(self): ) with self.assertRaises(ValueError): reconstruction.partition_ids(xr.Dataset(), spectral=True) - big = reconstruction.site_batch_size(24, 4) - self.assertGreater(big, reconstruction.site_batch_size(24 * 30, 40)) - self.assertEqual( - reconstruction.site_batch_size(10**6, 10**3, budget_bytes=1), 1 - ) + budget = 2**34 + big, steps = reconstruction.batch_shape(24, 4, budget) + self.assertEqual(steps, 24) + self.assertGreater(big, reconstruction.batch_shape(24 * 30, 40, budget)[0]) + one_site = reconstruction.site_sum_bytes(1, 10) + self.assertEqual(reconstruction.batch_shape(100, 10, 30 * one_site), (1, 30)) + self.assertEqual(reconstruction.batch_shape(10**6, 10**3, 1), (1, 1)) if __name__ == "__main__": diff --git a/tests/waves/test_hywaves_resources.py b/tests/waves/test_hywaves_resources.py new file mode 100644 index 0000000..4e3ad67 --- /dev/null +++ b/tests/waves/test_hywaves_resources.py @@ -0,0 +1,48 @@ +import unittest + +from bluemath_tk.core.resources import GIB, Machine +from bluemath_tk.waves.hywaves import resources + +BIG = Machine(48, 480 * GIB, 512 * GIB) +LAPTOP = Machine(6, 12 * GIB, 16 * GIB) + + +class TestPlans(unittest.TestCase): + def test_big_machine_takes_a_year_with_every_site(self): + res = resources.plan_reconstruction(38, 24, 10, host=BIG) + self.assertEqual((res.device, res.threads, res.workers, res.months, res.sites), + ("cpu", 36, 36, 12, None)) + self.assertEqual(res.memory, 360 * GIB) + + def test_fewer_months_then_fewer_sites(self): + res = resources.plan_reconstruction(5000, 12, 40, save_spectra=True, host=LAPTOP) + self.assertEqual(res.months, 1) + self.assertIsNotNone(res.sites) + self.assertEqual(resources.plan_reconstruction(38, 24, 40, host=LAPTOP).months, 12) + need = resources.call_memory(res.months, res.sites, 40, workers=res.workers, save_spectra=True) + self.assertLessEqual(need["total"], res.memory) + + def test_workers_capped_by_memory(self): + self.assertEqual(resources.plan_reconstruction(3, 1, 10, host=BIG).workers, 36) + tiny = Machine(48, int(4 * GIB / 0.75), GIB * 64) + res = resources.plan_reconstruction(1, 1, 10, host=tiny) + self.assertEqual(res.workers, 3) # 3 x (0.75 + 0.5) GiB in 4 GiB + with self.assertRaises(MemoryError): + resources.plan_reconstruction(1, 1, 10, host=Machine(4, GIB // 4, GIB)) + + def test_sum_budget_and_hours(self): + self.assertEqual(resources.sum_budget(100 * GIB, 0, 4), 4 * resources.SUM_BYTES_PER_WORKER) + self.assertEqual(resources.sum_budget(GIB, 2 * GIB, 4), 0) + self.assertLess(resources.hours(100, 1, 10, 36), resources.hours(100, 1, 10, 1)) + + def test_gpu_months_and_fits(self): + gpu = Machine(16, 100 * GIB, 128 * GIB, "A100", 40 * GIB) + self.assertEqual(resources.plan_reconstruction(38, 24, 10, host=gpu).months, 4) + self.assertEqual(resources.plan_fit(60, GIB, host=gpu), resources.FitResources("cuda", 1, 12)) + self.assertEqual(resources.plan_fit(60, 2 * GIB, host=BIG).workers, 36) + self.assertEqual(resources.plan_fit(60, 100 * GIB, host=BIG).workers, 3) + self.assertEqual(resources.plan_fit(5, GIB, host=BIG).workers, 5) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/waves/test_hywaves_store.py b/tests/waves/test_hywaves_store.py new file mode 100644 index 0000000..aff6b0d --- /dev/null +++ b/tests/waves/test_hywaves_store.py @@ -0,0 +1,81 @@ +import tempfile +import unittest +from pathlib import Path + +import numpy as np +import pandas as pd +import xarray as xr + +from bluemath_tk.waves.hywaves.store import ReconstructionStore + +# hourly, 2 Jan -> 3 Feb 2020, with the nanosecond jitter real hindcasts carry +TIMES = pd.date_range("2020-01-02", "2020-02-03 23:00", freq="h") + pd.to_timedelta( + np.tile([0, 1], 792), "ns" +)[: len(pd.date_range("2020-01-02", "2020-02-03 23:00", freq="h"))] +SITES = ["a", "b", "c"] + + +def _bulk(times, sites, value): + return xr.Dataset( + {"hs": (("time", "sites"), np.full((len(times), len(sites)), value))}, + coords={"time": times, "sites": sites}, + ) + + +class TestReconstructionStore(unittest.TestCase): + def setUp(self): + self._tmp = tempfile.TemporaryDirectory() + self.addCleanup(self._tmp.cleanup) + self.path = Path(self._tmp.name) / "s.zarr" + + def store(self, **kw): + return ReconstructionStore(self.path, TIMES, SITES, {"hs": ()}, key={"n": 1}, fingerprint="f1", **kw) + + def test_months_missing_write_read(self): + store = self.store() + self.assertEqual(store.months, ["2020-01", "2020-02"]) + self.assertEqual(store.months_of("2020-01-31"), ["2020-01"]) + self.assertEqual(store.missing(slice("2020-01-31", "2020-02-01"), ["b"]), + {"2020-01": ["b"], "2020-02": ["b"]}) + jan = TIMES[TIMES.month == 1] + store.write(_bulk(jan[:-5], ["a", "c"], 2.0), ["2020-01"]) # last 5 h lacking -> NaN + self.assertEqual(store.missing("2020-01", SITES), {"2020-01": ["b"]}) + day = store.read("2020-01-10", ["c"]) + self.assertEqual(day.sizes["time"], 24) + self.assertTrue((day.hs == 2.0).all()) + tail = store.read("2020-01-31", ["a"]) + self.assertEqual(int(tail.hs.isnull().sum()), 5) + with self.assertRaises(KeyError): + store.read("2020-01", ["b"]) + self.assertEqual(store.coverage().loc["2020-01", "sites"], 2) + + def test_reopen_keeps_and_mismatch_rebuilds(self): + self.store().write(_bulk(TIMES[TIMES.month == 2], SITES, 1.0), ["2020-02"]) + self.assertEqual(self.store().missing("2020-02", SITES), {}) + changed = ReconstructionStore(self.path, TIMES, SITES, {"hs": ()}, key={"n": 1}, fingerprint="f2") + self.assertEqual(changed.missing("2020-02", ["a"]), {"2020-02": ["a"]}) + + def test_invalidate(self): + store = self.store() + store.write(_bulk(TIMES[TIMES.month == 2], SITES, 1.0), ["2020-02"]) + store.invalidate("2020-02-01", ["b"]) + self.assertEqual(store.missing("2020-02", SITES), {"2020-02": ["b"]}) + + def test_spectra_variable(self): + freq, dirs = np.array([0.1, 0.2]), np.array([0.0, 90.0, 180.0]) + store = ReconstructionStore( + self.path, TIMES, SITES, {"efth": ("freq", "dir")}, coords={"freq": freq, "dir": dirs} + ) + feb = TIMES[TIMES.month == 2] + ds = xr.Dataset( + {"efth": (("time", "sites", "freq", "dir"), np.ones((len(feb), 2, 2, 3)))}, + coords={"time": feb, "sites": ["a", "b"], "freq": freq, "dir": dirs}, + ) + store.write(ds, ["2020-02"]) + out = store.read("2020-02-02", ["b"]) + self.assertEqual(out.efth.dims, ("time", "sites", "freq", "dir")) + self.assertTrue((out.efth == 1).all()) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/waves/test_wind_correction.py b/tests/waves/test_wind_correction.py index 7060f44..b5ad83a 100644 --- a/tests/waves/test_wind_correction.py +++ b/tests/waves/test_wind_correction.py @@ -126,10 +126,10 @@ def test_reconstruction_adds_the_energy(self): cube = _cube() chunks = [cube.sel(goal=[g]) for g in (1, 2)] variables = {"hs": "hs"} - base, _ = sum_partition_spectra_batched(chunks, SITES, 3, 2, variables) + base, _ = sum_partition_spectra_batched(chunks, SITES, variables, budget_bytes=2**30) corr = _fitted(lambda h: np.full_like(h, 1.5)) wind, _ = sum_partition_spectra_batched( - chunks, SITES, 3, 2, variables, + chunks, SITES, variables, budget_bytes=2**30, wind_correction={"correction": corr, "fetch": _fetch()}, ) k = float(wind["wind_k"].mean())