"""Measure throughput and memory use to recommend a worker count.
spaCR's worker defaults are arithmetic on the core count --
``cpu_count() - 4``, ``cpu_count() // 2``, ``-1``. On a measurement run the
core count is rarely the limiting factor: a field is hundreds of megabytes
decompressed, and eight workers on a 16-core laptop with 16 GB will swap long
before they saturate the CPU. The limiting quantity is usually the memory one
worker needs for one representative field, which depends on the plate and the
machine, not on the core count.
The benchmark therefore runs the real work over a few fields, records peak
RSS and peak VRAM, and divides the available memory by the memory one worker
used.
**This measures; it does not tune.** It returns a recommendation and the
evidence behind it, and never writes a setting. A benchmark that silently
changed `n_jobs` would make a run's speed depend on when the benchmark last
ran, which is the opposite of reproducible.
Qt-free and torch-optional: VRAM is reported when torch is present with a
CUDA device and reported as ``None`` otherwise, which is not the same as 0.
"""
from __future__ import annotations
import os
import time
from dataclasses import dataclass, field
from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence
@dataclass(frozen=True)
[docs]
class Measurement:
"""Item count, wall time and peak memory observed in one benchmark run.
:param items: how many representative units were processed.
:param seconds: wall clock for all of them.
:param peak_rss_bytes: the highest resident set size seen, for the whole
process. Not per worker -- the benchmark runs serially on purpose, so
this is one worker's requirement.
:param peak_vram_bytes: peak CUDA allocation, or ``None`` when there is
no CUDA device. ``None`` and ``0`` are different answers: one means
"not measured", the other "measured, and it used none".
:param baseline_rss_bytes: RSS before the work started, so the caller can
tell the interpreter's own footprint from the work's.
:param notes: caveats about warm-up or unavailable measurements that must
accompany the numeric result.
"""
items: int
seconds: float
peak_rss_bytes: int
peak_vram_bytes: Optional[int] = None
baseline_rss_bytes: int = 0
notes: List[str] = field(default_factory=list)
@property
[docs]
def per_item_seconds(self) -> float:
"""Return mean wall time per measured item, or NaN when none ran."""
return self.seconds / self.items if self.items else float("nan")
@property
[docs]
def items_per_second(self) -> float:
"""Return measured throughput, or NaN for a nonpositive duration."""
return self.items / self.seconds if self.seconds > 0 else float("nan")
@property
@dataclass(frozen=True)
[docs]
class Recommendation:
"""A worker count and the retained evidence supporting it.
:param workers: final recommended parallel-worker count; always at least
one and bounded by the normalized core count, measured memory capacity
when usable, and configured maximum.
:param reason: human-readable explanation of the branch that set
``workers``: missing measurement, one-worker fallback, memory bound,
core bound, or configured maximum.
:param measurement: exact :class:`Measurement` supplied to
:func:`recommend_workers`; ``None`` means no benchmark was available,
while a zero-work-footprint measurement is retained but triggers the
core-count fallback.
:param cores: effective logical-core ceiling after defaulting from
:func:`os.cpu_count` and clamping to at least one.
:param available_bytes: available-memory snapshot before the configured
reserve is subtracted; supplied by the caller or measured by spaCR,
and zero when unavailable.
"""
workers: int
reason: str
measurement: Optional[Measurement] = None
cores: int = 0
available_bytes: int = 0
[docs]
def __str__(self) -> str:
"""Return the worker count followed by its sizing explanation."""
return f"{self.workers} worker(s): {self.reason}"
def _rss_bytes() -> int:
"""Resident set size now, in bytes, or 0 where it cannot be read."""
try:
import resource
except ImportError:
return 0
usage = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
import sys
return int(usage) if sys.platform == "darwin" else int(usage) * 1024
def _vram_bytes() -> Optional[int]:
"""Peak CUDA allocation since the counter was reset, or None."""
try:
import torch
except Exception:
return None
try:
if not torch.cuda.is_available():
return None
return int(torch.cuda.max_memory_allocated())
except Exception:
return None
def _reset_vram() -> None:
"""Reset PyTorch's peak CUDA counter when that facility is available."""
try:
import torch
if torch.cuda.is_available():
torch.cuda.reset_peak_memory_stats()
except Exception:
pass
[docs]
def available_memory_bytes() -> int:
"""Return ``MemAvailable`` from ``/proc/meminfo`` in bytes, or 0 if unreadable.
Available rather than total memory is used because total includes memory
already in use; sizing workers against it can exhaust memory mid-run.
"""
try:
with open("/proc/meminfo", encoding="utf-8") as handle:
for line in handle:
if line.startswith("MemAvailable:"):
return int(line.split()[1]) * 1024
except (OSError, ValueError, IndexError):
pass
return 0
[docs]
def benchmark(work: Callable[[Any], Any], items: Sequence[Any], *,
warmup: int = 1) -> Measurement:
"""Run ``work`` over ``items`` serially and record time and peak memory.
Items run serially because the recommendation divides by what one worker
needs; a parallel run would measure the sum and hide that figure.
:param work: callable invoked once for every warm-up and measured item;
its return value is ignored because only resource use is measured.
:param items: ordered workload to process. At least one item is
required, and the final item is always kept in the measured set.
:param warmup: items processed before the clock starts. The first field
pays for imports, CUDA context creation and page faults that no later
field pays again, and counting it makes a short run look far slower
than the plate it is predicting.
:raises ValueError: no items to measure -- a benchmark over nothing would
return a per-item cost of NaN and a recommendation built on it.
"""
items = list(items)
if not items:
raise ValueError("benchmark needs at least one item to measure")
notes: List[str] = []
warmup = max(0, min(int(warmup), len(items) - 1))
for item in items[:warmup]:
work(item)
if warmup:
notes.append(f"{warmup} warm-up item(s) excluded from the timing")
measured = items[warmup:]
baseline = _rss_bytes()
_reset_vram()
started = time.perf_counter()
for item in measured:
work(item)
elapsed = time.perf_counter() - started
vram = _vram_bytes()
if vram is None:
notes.append("no CUDA device, so VRAM was not measured")
return Measurement(
items=len(measured),
seconds=elapsed,
peak_rss_bytes=_rss_bytes(),
peak_vram_bytes=vram,
baseline_rss_bytes=baseline,
notes=notes,
)
[docs]
def recommend_workers(measurement: Optional[Measurement] = None, *,
cores: Optional[int] = None,
available_bytes: Optional[int] = None,
reserve_bytes: int = 2 * 1024 ** 3,
maximum: int = 32) -> Recommendation:
"""Recommend a worker count from cores, free memory and measured usage.
The rule, in order:
1. Never more than ``cores``. More workers than cores is contention.
2. Never more than ``available memory - reserve`` divided by what one
worker measurably needed. This is the term the core-count defaults
omit, and it is usually the binding one.
3. Never fewer than 1, and never more than ``maximum``.
:param reserve_bytes: memory left for everything that is not a worker --
the GUI, the page cache the readers depend on, and the operating
system. Defaults to 2 GiB.
:returns: a :class:`Recommendation` carrying the reason, which states
the limit that set the worker count.
"""
cores = int(cores if cores is not None else (os.cpu_count() or 1))
cores = max(1, cores)
available = int(available_bytes if available_bytes is not None
else available_memory_bytes())
if measurement is None or measurement.work_rss_bytes <= 0:
workers = max(1, min(cores, maximum))
return Recommendation(
workers=workers,
reason=(f"{cores} core(s); no usable memory measurement, so the "
f"core count is the only bound"),
measurement=measurement, cores=cores, available_bytes=available)
per_worker = measurement.work_rss_bytes
budget = max(0, available - int(reserve_bytes))
by_memory = int(budget // per_worker)
if by_memory < 1:
return Recommendation(
workers=1,
reason=(f"one worker needs {per_worker / 1024 ** 3:.2f} GiB and "
f"only {budget / 1024 ** 3:.2f} GiB is free after the "
f"reserve; one worker at a time is what fits"),
measurement=measurement, cores=cores, available_bytes=available)
workers = max(1, min(cores, by_memory, maximum))
if workers == by_memory < cores:
reason = (f"memory-bound: {budget / 1024 ** 3:.2f} GiB free after the "
f"reserve, {per_worker / 1024 ** 3:.2f} GiB per worker, "
f"{cores} cores available")
elif workers == cores:
reason = (f"core-bound: {cores} core(s), and memory would allow "
f"{by_memory}")
else:
reason = f"capped at the {maximum}-worker maximum"
return Recommendation(workers=workers, reason=reason,
measurement=measurement, cores=cores,
available_bytes=available)