From d957b49f5e75d488241567a8ac3053359c475a05 Mon Sep 17 00:00:00 2001 From: Fabian Peddinghaus Date: Tue, 29 Sep 2026 14:37:12 +0000 Subject: [PATCH 1/4] Fix benchmark source and converter defects and share a fixed-source base --- src/python/omnimalloc/benchmark/benchmark.py | 7 +- .../benchmark/converters/__init__.py | 2 - .../omnimalloc/benchmark/converters/model.py | 86 +-- .../omnimalloc/benchmark/converters/onnx.py | 132 ++-- .../benchmark/sources/adversarial.py | 14 +- .../omnimalloc/benchmark/sources/base.py | 99 ++- .../benchmark/sources/concurrent_tiling.py | 15 +- .../omnimalloc/benchmark/sources/generator.py | 41 +- .../benchmark/sources/huggingface.py | 218 ++---- .../benchmark/sources/minimalloc.py | 119 +-- .../omnimalloc/benchmark/sources/sizes.py | 7 +- .../benchmark/sources/sync_patterns.py | 21 +- .../omnimalloc/benchmark/sources/tiling.py | 6 +- .../benchmark/sources/tiling_base.py | 22 +- tests/unit/benchmark/converters/test_model.py | 708 ++---------------- tests/unit/benchmark/converters/test_onnx.py | 126 +--- tests/unit/benchmark/sources/test_base.py | 2 +- .../sources/test_concurrent_tiling.py | 14 +- .../unit/benchmark/sources/test_generator.py | 63 +- .../benchmark/sources/test_huggingface.py | 173 ++--- .../unit/benchmark/sources/test_minimalloc.py | 53 +- .../benchmark/sources/test_sync_patterns.py | 8 + tests/unit/benchmark/sources/test_tiling.py | 12 +- 23 files changed, 485 insertions(+), 1463 deletions(-) diff --git a/src/python/omnimalloc/benchmark/benchmark.py b/src/python/omnimalloc/benchmark/benchmark.py index 1ea79a1..7039b89 100644 --- a/src/python/omnimalloc/benchmark/benchmark.py +++ b/src/python/omnimalloc/benchmark/benchmark.py @@ -57,10 +57,7 @@ def _resolve_parameterizable_variants( def _resolve_fixed_variants( source: BaseSource, variants: int | tuple[IdType, ...] | None ) -> tuple[str, ...]: - variant_count = ( - variants if isinstance(variants, int) else len(variants) if variants else None - ) - available = source.get_available_variants(variant_count) + available = source.get_available_variants() if available is None: return () if variants is None: @@ -183,7 +180,7 @@ def _benchmark_report( # The ground truth is a property of the instance, not the allocator, and # the tiling sources rebuild their whole construction to read it if variant_id not in known_optima: - known_optima[variant_id] = source.get_known_optimum(variant_id) + known_optima[variant_id] = source.get_known_optimum() return BenchmarkReport( id=report_id, diff --git a/src/python/omnimalloc/benchmark/converters/__init__.py b/src/python/omnimalloc/benchmark/converters/__init__.py index 0cd75d8..dbbc96a 100644 --- a/src/python/omnimalloc/benchmark/converters/__init__.py +++ b/src/python/omnimalloc/benchmark/converters/__init__.py @@ -6,6 +6,4 @@ from .model import Model as Model from .model import Op as Op from .model import model_to_allocations as model_to_allocations -from .model import model_to_pools as model_to_pools -from .model import model_to_system as model_to_system from .onnx import from_onnx as from_onnx diff --git a/src/python/omnimalloc/benchmark/converters/model.py b/src/python/omnimalloc/benchmark/converters/model.py index 984da8f..7bbbb7e 100644 --- a/src/python/omnimalloc/benchmark/converters/model.py +++ b/src/python/omnimalloc/benchmark/converters/model.py @@ -6,14 +6,7 @@ from dataclasses import dataclass, field from typing import Final -from omnimalloc.primitives import ( - Allocation, - AllocationKind, - IdType, - Memory, - Pool, - System, -) +from omnimalloc.primitives import Allocation, AllocationKind, IdType # Bits, not bytes: the sub-byte types are packed several to a byte, which is # what the tensors actually occupy and what numpy's itemsize gets wrong. @@ -60,10 +53,6 @@ def __post_init__(self) -> None: if self.dtype not in ITEMBITS: raise ValueError(f"unknown dtype {self.dtype!r}") - @property - def ndim(self) -> int: - return len(self.shape) - @property def size(self) -> int: return math.ceil(ITEMBITS[self.dtype] * math.prod(self.shape) / 8) @@ -74,7 +63,6 @@ class Op: id: IdType inputs: set[Buffer] = field(default_factory=set) outputs: set[Buffer] = field(default_factory=set) - op_type: str | None = None def __post_init__(self) -> None: if not isinstance(self.id, (int, str)): @@ -101,10 +89,8 @@ def __post_init__(self) -> None: def _compute_buffer_lifetimes( model: Model, - const_inf_lifetime: bool, - io_inf_lifetime: bool, ) -> tuple[dict[Buffer, int], dict[Buffer, int]]: - """Compute first and last usage indices for each buffer.""" + """First and last op index using each buffer; constants and IO span the model.""" buffer_to_first_index: dict[Buffer, int] = {} buffer_to_last_index: dict[Buffer, int] = {} @@ -114,26 +100,23 @@ def _compute_buffer_lifetimes( buffer_to_first_index[buffer] = idx buffer_to_last_index[buffer] = idx - # Apply infinite lifetime constraints (at least one step for op-less models) + # At least one step for op-less models max_index = max(len(model.ops) - 1, 0) for buffer in model.buffers.values(): - if (buffer.kind == AllocationKind.CONSTANT and const_inf_lifetime) or ( - buffer.kind.is_io and io_inf_lifetime - ): + if buffer.kind == AllocationKind.CONSTANT or buffer.kind.is_io: buffer_to_first_index[buffer] = 0 buffer_to_last_index[buffer] = max_index return buffer_to_first_index, buffer_to_last_index -def _create_allocations( +def model_to_allocations( model: Model, - include_const: bool, - include_io: bool, - buffer_to_first_index: dict[Buffer, int], - buffer_to_last_index: dict[Buffer, int], + include_const: bool = False, + include_io: bool = False, ) -> list[Allocation]: - """Create allocations from buffers and their lifetimes.""" + """Extract Allocations from Model buffers.""" + buffer_to_first_index, buffer_to_last_index = _compute_buffer_lifetimes(model) return [ Allocation( id=buffer.id, @@ -150,54 +133,3 @@ def _create_allocations( and buffer in buffer_to_first_index ) ] - - -def model_to_allocations( - model: Model, - include_const: bool = False, - include_io: bool = False, - const_inf_lifetime: bool = True, - io_inf_lifetime: bool = True, -) -> list[Allocation]: - """Extract Allocations from Model buffers.""" - buffer_to_first_index, buffer_to_last_index = _compute_buffer_lifetimes( - model, const_inf_lifetime, io_inf_lifetime - ) - allocations = _create_allocations( - model, include_const, include_io, buffer_to_first_index, buffer_to_last_index - ) - return allocations - - -def model_to_pools( - model: Model, - include_const: bool = True, - include_io: bool = True, - const_inf_lifetime: bool = True, - io_inf_lifetime: bool = True, -) -> tuple[Pool, ...]: - """Extract Pools grouped by buffer kind.""" - buffer_to_first_index, buffer_to_last_index = _compute_buffer_lifetimes( - model, const_inf_lifetime, io_inf_lifetime - ) - allocations = _create_allocations( - model, include_const, include_io, buffer_to_first_index, buffer_to_last_index - ) - - # Group allocations by kind - allocations_by_kind: dict[AllocationKind, list[Allocation]] = {} - for alloc in allocations: - kind = alloc.kind if alloc.kind is not None else AllocationKind.WORKSPACE - allocations_by_kind.setdefault(kind, []).append(alloc) - - return tuple( - Pool(id=str(kind), allocations=tuple(allocs)) - for kind, allocs in allocations_by_kind.items() - ) - - -def model_to_system(model: Model) -> System: - """Convert an model to a system with a single memory and pools.""" - pools = model_to_pools(model) - memory = Memory(id=0, pools=pools) - return System(id=model.id, memories=(memory,)) diff --git a/src/python/omnimalloc/benchmark/converters/onnx.py b/src/python/omnimalloc/benchmark/converters/onnx.py index 84cbdc7..44b8f3f 100644 --- a/src/python/omnimalloc/benchmark/converters/onnx.py +++ b/src/python/omnimalloc/benchmark/converters/onnx.py @@ -3,7 +3,9 @@ # import logging +from collections.abc import Iterable from pathlib import Path +from typing import Any, cast from omnimalloc.common.optional import require_optional from omnimalloc.primitives import AllocationKind @@ -15,24 +17,13 @@ HAS_ONNX = True except ImportError: - from types import SimpleNamespace - HAS_ONNX = False - onnx = SimpleNamespace( # ty: ignore[invalid-assignment] - checker=SimpleNamespace(check_model=None), - shape_inference=SimpleNamespace(infer_shapes=None), - load_model=None, - helper=SimpleNamespace(tensor_dtype_to_np_dtype=None), - ModelProto=None, - TensorProto=None, - ValueInfoProto=None, - NodeProto=None, - ) + onnx = cast("Any", None) logger = logging.getLogger(__name__) -def _from_onnx_model(onnx_model: onnx.ModelProto) -> Model: +def _from_onnx_model(onnx_model: "onnx.ModelProto") -> Model: onnx.checker.check_model(onnx_model, full_check=True) onnx_model = onnx.shape_inference.infer_shapes( onnx_model, @@ -42,28 +33,26 @@ def _from_onnx_model(onnx_model: onnx.ModelProto) -> Model: ) graph = onnx_model.graph + initializers = {tensor.name for tensor in graph.initializer} + candidates = [ + *(_tensor_proto_to_buffer(tensor) for tensor in graph.initializer), + # Legacy IR re-lists initializers under graph inputs; those are constants. + *( + _value_info_to_buffer(value, AllocationKind.INPUT) + for value in graph.input + if value.name not in initializers + ), + *(_value_info_to_buffer(v, AllocationKind.OUTPUT) for v in graph.output), + *(_value_info_to_buffer(v, AllocationKind.WORKSPACE) for v in graph.value_info), + ] buffers: dict[str | int, Buffer] = {} - - def _add_buffer(buffer: Buffer) -> None: + for buffer in candidates: + if buffer is None: + continue if buffer.id in buffers: raise ValueError(f"Buffer {buffer.id} already exists") buffers[buffer.id] = buffer - for init in graph.initializer: - _add_buffer(_tensor_proto_to_buffer(init)) - - for inp in graph.input: - # Legacy IR re-lists initializers under graph inputs; those are constants. - if inp.name in buffers: - continue - _add_buffer(_value_info_to_buffer(inp, AllocationKind.INPUT)) - - for out in graph.output: - _add_buffer(_value_info_to_buffer(out, AllocationKind.OUTPUT)) - - for val in graph.value_info: - _add_buffer(_value_info_to_buffer(val, AllocationKind.WORKSPACE)) - ops = {} for idx, node in enumerate(graph.node): # Node names are optional in ONNX; synthesize unique ids for unnamed nodes. @@ -72,71 +61,48 @@ def _add_buffer(buffer: Buffer) -> None: raise ValueError(f"Node {op.id} already exists in ops") ops[op.id] = op - name = onnx_model.doc_string or graph.name or "unnamed_model" - return Model(id=name, ops=ops, buffers=buffers) - - -def _tensor_proto_to_buffer(tensor: onnx.TensorProto) -> Buffer: - original_shape = tuple(tensor.dims) - shape = tuple(dim for dim in original_shape if dim > 0) - if len(shape) != len(original_shape): - logger.debug( - f"Dropped dimensions with size <=0 in tensor '{tensor.name}': " - f"{original_shape} -> {shape}" - ) - return Buffer( - id=tensor.name, - shape=shape, - dtype=onnx.helper.tensor_dtype_to_np_dtype(tensor.data_type).name, - kind=AllocationKind.CONSTANT, - ) + return Model(id=graph.name or "unnamed_model", ops=ops, buffers=buffers) + + +def _buffer( + name: str, dims: Iterable[int], elem_type: int, kind: AllocationKind +) -> Buffer | None: + """The tensor's buffer, or None for a zero-size tensor, which needs no memory.""" + shape = tuple(dims) + if 0 in shape: + logger.debug(f"Skipping zero-size tensor '{name}' of shape {shape}") + return None + dtype = onnx.helper.tensor_dtype_to_np_dtype(elem_type).name + return Buffer(id=name, shape=shape, dtype=dtype, kind=kind) + + +def _tensor_proto_to_buffer(tensor: "onnx.TensorProto") -> Buffer | None: + return _buffer(tensor.name, tensor.dims, tensor.data_type, AllocationKind.CONSTANT) def _value_info_to_buffer( - value_info: onnx.ValueInfoProto, kind: AllocationKind -) -> Buffer: + value_info: "onnx.ValueInfoProto", kind: AllocationKind +) -> Buffer | None: tt = value_info.type.tensor_type - original_shape = tuple(int(dim.dim_value) for dim in tt.shape.dim) - shape = tuple(dim for dim in original_shape if dim > 0) - if len(shape) != len(original_shape): - logger.debug( - f"Dropped dimensions with size <=0 in value '{value_info.name}': " - f"{original_shape} -> {shape}" - ) - return Buffer( - id=value_info.name, - shape=shape, - dtype=onnx.helper.tensor_dtype_to_np_dtype(tt.elem_type).name, - kind=kind, - ) + # A symbolic dim (e.g. the batch) has no value; take it as 1 + dims = (d.dim_value if d.HasField("dim_value") else 1 for d in tt.shape.dim) + return _buffer(value_info.name, dims, tt.elem_type, kind) def _node_to_op( - node: onnx.NodeProto, buffers: dict[str | int, Buffer], op_id: str + node: "onnx.NodeProto", buffers: dict[str | int, Buffer], op_id: str ) -> Op: - input_buffers = [] - for name in node.input: - if name not in buffers: - logger.debug(f"Input buffer '{name}' not found for node '{op_id}'") - continue - input_buffers.append(buffers[name]) - - output_buffers = [] - for name in node.output: - if name not in buffers: - logger.debug(f"Output buffer '{name}' not found for node '{op_id}'") - continue - output_buffers.append(buffers[name]) - + missing = [name for name in (*node.input, *node.output) if name not in buffers] + if missing: + logger.debug(f"Buffers {missing} not found for node '{op_id}'") return Op( id=op_id, - inputs=set(input_buffers), - outputs=set(output_buffers), - op_type=node.op_type, + inputs={buffers[name] for name in node.input if name in buffers}, + outputs={buffers[name] for name in node.output if name in buffers}, ) -def from_onnx(onnx_input: onnx.ModelProto | str | Path) -> Model: +def from_onnx(onnx_input: "onnx.ModelProto | str | Path") -> Model: """Convert ONNX model or file path to Model.""" if not HAS_ONNX: require_optional("onnx", "ONNX model conversion") @@ -146,5 +112,5 @@ def from_onnx(onnx_input: onnx.ModelProto | str | Path) -> Model: if isinstance(onnx_input, onnx.ModelProto): return _from_onnx_model(onnx_input) raise TypeError( - f"onnx_input must be an onnx.ModelProto or str, got {type(onnx_input)}" + f"onnx_input must be an onnx.ModelProto, str or Path, got {type(onnx_input)}" ) diff --git a/src/python/omnimalloc/benchmark/sources/adversarial.py b/src/python/omnimalloc/benchmark/sources/adversarial.py index 8cd366c..b12b844 100644 --- a/src/python/omnimalloc/benchmark/sources/adversarial.py +++ b/src/python/omnimalloc/benchmark/sources/adversarial.py @@ -6,6 +6,7 @@ from typing import ClassVar from omnimalloc.common.constants import DEFAULT_SEED, KB, MB +from omnimalloc.common.validation import ensure_positive from omnimalloc.primitives import Allocation from .base import BaseSource @@ -40,12 +41,10 @@ def __init__( duration_max: int = 64, seed: int | None = DEFAULT_SEED, ) -> None: - if size_min <= 0: - raise ValueError("size_min must be positive") + ensure_positive(size_min, "size_min") if size_max < size_min: raise ValueError("size_max must be >= size_min") - if duration_min <= 0: - raise ValueError("duration_min must be positive") + ensure_positive(duration_min, "duration_min") if duration_max < duration_min: raise ValueError("duration_max must be >= duration_min") if time_max <= duration_max: @@ -96,8 +95,7 @@ def __init__( ) -> None: if not 0.0 <= noise < 1.0: raise ValueError("noise must be in [0, 1)") - if size_min <= 0: - raise ValueError("size_min must be positive") + ensure_positive(size_min, "size_min") if size_max < size_min: raise ValueError("size_max must be >= size_min") super().__init__(num_allocations=num_allocations) @@ -134,14 +132,14 @@ def get_allocations( while len(allocations) < num: lane = rng.randrange(2) - step = rng.randrange(self._GROUP * obstructions or 1) + step = rng.randrange(self._GROUP * obstructions) start = (step, 0) if lane == 0 else (0, step) end = (step + 1, 0) if lane == 0 else (0, step + 1) allocations.append( self._allocation(rng, skip + len(allocations), start, end) ) - return tuple(allocations[:num]) + return tuple(allocations) def _allocation( self, diff --git a/src/python/omnimalloc/benchmark/sources/base.py b/src/python/omnimalloc/benchmark/sources/base.py index e5baec3..c7f8c07 100644 --- a/src/python/omnimalloc/benchmark/sources/base.py +++ b/src/python/omnimalloc/benchmark/sources/base.py @@ -4,6 +4,7 @@ import inspect from abc import abstractmethod +from functools import cached_property from typing import ClassVar from omnimalloc.common.registry import Registered @@ -68,22 +69,16 @@ def is_parameterizable(self) -> bool: """Whether this source can generate arbitrary allocation counts.""" return True - def get_known_optimum(self, variant_id: IdType | None = None) -> int | None: - """Provably achievable peak size for a variant, or None if unknown. + def get_known_optimum(self) -> int | None: + """Provably achievable peak size of every variant, or None if unknown. Sources that reverse-construct their instances from a packing know the optimum and override this; everyone else leaves it unknown. """ + return None - def get_available_variants( - self, - count: int | None = None, # noqa: ARG002 - ) -> tuple[str, ...] | None: - """Variant names offered by a fixed source; None from parameterizable ones. - - `count` is how many the caller needs; a source materializing variants - lazily may provision at least that many. - """ + def get_available_variants(self) -> tuple[str, ...] | None: + """Variant names offered by a fixed source; None from parameterizable ones.""" return None def get_variant(self, variant_id: IdType) -> Pool: @@ -135,6 +130,7 @@ def get_memories( num_pools=self.num_pools, skip=(skip + i) * self.num_pools, ) + # A fixed source runs out of pools; a generated one never does if not pools: raise ValueError(f"source {self.name()} returned no pools") memories.append( @@ -151,16 +147,16 @@ def get_systems( ) -> tuple[System, ...]: num_systems = self.num_systems if num_systems is None else num_systems ensure_positive(num_systems, "num_systems") - systems = [] - for i in range(num_systems): - memories = self.get_memories( - num_memories=self.num_memories, - skip=(skip + i) * self.num_memories, + return tuple( + System( + id=f"{self.name()}_system_{i}", + memories=self.get_memories( + num_memories=self.num_memories, + skip=(skip + i) * self.num_memories, + ), ) - if not memories: - raise ValueError(f"source {self.name()} returned no memories") - systems.append(System(id=f"{self.name()}_system_{i}", memories=memories)) - return tuple(systems) + for i in range(num_systems) + ) def get_allocation(self) -> Allocation: allocations = self.get_allocations(num_allocations=1) @@ -175,9 +171,64 @@ def get_pool(self) -> Pool: return pools[0] def get_memory(self) -> Memory: - memories = self.get_memories(num_memories=1) - return memories[0] + return self.get_memories(num_memories=1)[0] def get_system(self) -> System: - systems = self.get_systems(num_systems=1) - return systems[0] + return self.get_systems(num_systems=1)[0] + + +class FixedSource(BaseSource): + """A fixed collection of pools, one variant per pool, loaded on first use.""" + + @abstractmethod + def _load_pools(self) -> tuple[Pool, ...]: ... + + @cached_property + def pools(self) -> tuple[Pool, ...]: + return self._load_pools() + + def is_parameterizable(self) -> bool: + return False + + def get_available_variants(self) -> tuple[str, ...]: + return tuple(str(pool.id) for pool in self.pools) + + def get_variant(self, variant_id: IdType) -> Pool: + """The pool at an index, or the pool with that id.""" + if isinstance(variant_id, int) and 0 <= variant_id < len(self.pools): + return self.pools[variant_id] + for pool in self.pools: + if pool.id == variant_id: + return pool + raise ValueError(f"Variant {variant_id!r} not found in {self.label()}") + + def get_allocations( + self, num_allocations: int | None = None, skip: int = 0 + ) -> tuple[Allocation, ...]: + allocations = tuple(a for pool in self.pools for a in pool.allocations) + end = None if num_allocations is None else skip + num_allocations + return allocations[skip:end] + + def get_pools( + self, num_pools: int | None = None, skip: int = 0 + ) -> tuple[Pool, ...]: + ensure_positive(num_pools, "num_pools", allow_none=True) + end = None if num_pools is None else skip + num_pools + return self.pools[skip:end] + + +def prefix_ids(pool: Pool) -> Pool: + """Qualify allocation ids with the pool id, keeping them unique across pools.""" + return pool.with_allocations( + tuple( + Allocation( + id=f"{pool.id}_{alloc.id}", + size=alloc.size, + start=alloc.start, + end=alloc.end, + offset=alloc.offset, + kind=alloc.kind, + ) + for alloc in pool.allocations + ) + ) diff --git a/src/python/omnimalloc/benchmark/sources/concurrent_tiling.py b/src/python/omnimalloc/benchmark/sources/concurrent_tiling.py index 9d89ed5..322aa77 100644 --- a/src/python/omnimalloc/benchmark/sources/concurrent_tiling.py +++ b/src/python/omnimalloc/benchmark/sources/concurrent_tiling.py @@ -7,6 +7,7 @@ from typing import ClassVar from omnimalloc.common.constants import DEFAULT_SEED, KB, MB +from omnimalloc.common.validation import ensure_positive from omnimalloc.primitives import TimePoint, VectorClock from .tiling import TilingSource @@ -24,12 +25,7 @@ class ConcurrentTilingSource(TilingSource): """ _label_fields: ClassVar[tuple[str, ...]] = ( - "capacity", - "makespan", - "size_min", - "duration_min", - "seed", - "mem_cut_prob", + *TilingSource._label_fields, # noqa: SLF001 "num_threads", "num_syncs", ) @@ -46,10 +42,7 @@ def __init__( mem_cut_prob: float = 0.5, seed: int | None = DEFAULT_SEED, ) -> None: - if num_threads <= 0: - raise ValueError("num_threads must be positive") - if num_allocations < num_threads: - raise ValueError("num_allocations must be >= num_threads") + ensure_positive(num_threads, "num_threads") if capacity % num_threads: raise ValueError("capacity must be divisible by num_threads") if capacity // num_threads < size_min: @@ -96,7 +89,7 @@ def _simulate_syncs(self, rng: random.Random) -> list[_SyncHistory]: """Deliver random sync messages, max-merging the receiver's clock.""" knowledge = [[0] * self.num_threads for _ in range(self.num_threads)] histories: list[_SyncHistory] = [([], []) for _ in range(self.num_threads)] - if self.num_threads < 2 or not self.num_syncs: + if self.num_threads < 2: return histories # All workers share the local step scale, so delivering messages in # instant order is a causally consistent execution. diff --git a/src/python/omnimalloc/benchmark/sources/generator.py b/src/python/omnimalloc/benchmark/sources/generator.py index 0598ba4..c64a8d6 100644 --- a/src/python/omnimalloc/benchmark/sources/generator.py +++ b/src/python/omnimalloc/benchmark/sources/generator.py @@ -7,6 +7,7 @@ from typing import ClassVar from omnimalloc.common.constants import DEFAULT_SEED, KB, MB +from omnimalloc.common.validation import ensure_positive from omnimalloc.primitives import Allocation, AllocationKind from .base import BaseSource @@ -61,16 +62,14 @@ def __init__( seed: int | None = DEFAULT_SEED, ) -> None: super().__init__(num_allocations=num_allocations) - if size_min <= 0: - raise ValueError("size_min must be positive") + ensure_positive(size_min, "size_min") if size_max < size_min: raise ValueError("size_max must be >= size_min") if time_min < 0: raise ValueError("time_min must be non-negative") if time_max <= time_min: raise ValueError("time_max must be > time_min") - if duration_min <= 0: - raise ValueError("duration_min must be positive") + ensure_positive(duration_min, "duration_min") if duration_max < duration_min: raise ValueError("duration_max must be >= duration_min") if duration_max > (time_max - time_min): @@ -91,8 +90,7 @@ def __init__( def _generate_one(self, rng: random.Random, alloc_id: int) -> Allocation: size = rng.randint(self.size_min, self.size_max) duration = rng.randint(self.duration_min, self.duration_max) - max_start = self.time_max - duration - start = rng.randint(self.time_min, max(self.time_min, max_start)) + start = rng.randint(self.time_min, self.time_max - duration) kind = None if self.kinds: @@ -126,12 +124,9 @@ def __init__( seed: int | None = DEFAULT_SEED, ) -> None: super().__init__(num_allocations=num_allocations) - if size <= 0: - raise ValueError("size must be positive") - if duration <= 0: - raise ValueError("duration must be positive") - if time_max <= 0: - raise ValueError("time_max must be positive") + ensure_positive(size, "size") + ensure_positive(duration, "duration") + ensure_positive(time_max, "time_max") if duration > time_max: raise ValueError("duration must be <= time_max") @@ -141,8 +136,7 @@ def __init__( self.seed = seed def _generate_one(self, rng: random.Random, alloc_id: int) -> Allocation: - max_start = max(0, self.time_max - self.duration) - start = rng.randint(0, max_start) + start = rng.randint(0, self.time_max - self.duration) return Allocation( id=alloc_id, @@ -179,12 +173,12 @@ def __init__( raise ValueError("size_exponent_min must be non-negative") if size_exponent_max < size_exponent_min: raise ValueError("size_exponent_max must be >= size_exponent_min") - if time_max <= 0: - raise ValueError("time_max must be positive") - if duration_min <= 0: - raise ValueError("duration_min must be positive") + ensure_positive(time_max, "time_max") + ensure_positive(duration_min, "duration_min") if duration_max < duration_min: raise ValueError("duration_max must be >= duration_min") + if duration_max > time_max: + raise ValueError("duration_max must be <= time_max") self.size_exponent_min = size_exponent_min self.size_exponent_max = size_exponent_max @@ -196,7 +190,7 @@ def __init__( def _generate_one(self, rng: random.Random, alloc_id: int) -> Allocation: exponent = rng.randint(self.size_exponent_min, self.size_exponent_max) duration = rng.randint(self.duration_min, self.duration_max) - start = rng.randint(0, max(0, self.time_max - duration)) + start = rng.randint(0, self.time_max - duration) return Allocation( id=alloc_id, @@ -225,8 +219,7 @@ def __init__( seed: int | None = DEFAULT_SEED, ) -> None: super().__init__(num_allocations=num_allocations) - if size_min <= 0: - raise ValueError("size_min must be positive") + ensure_positive(size_min, "size_min") if size_max < size_min: raise ValueError("size_max must be >= size_min") if time_window < 2: @@ -271,12 +264,10 @@ def __init__( seed: int | None = DEFAULT_SEED, ) -> None: super().__init__(num_allocations=num_allocations) - if size_min <= 0: - raise ValueError("size_min must be positive") + ensure_positive(size_min, "size_min") if size_max < size_min: raise ValueError("size_max must be >= size_min") - if duration_min <= 0: - raise ValueError("duration_min must be positive") + ensure_positive(duration_min, "duration_min") if duration_max < duration_min: raise ValueError("duration_max must be >= duration_min") diff --git a/src/python/omnimalloc/benchmark/sources/huggingface.py b/src/python/omnimalloc/benchmark/sources/huggingface.py index 4e8a411..5924a3c 100644 --- a/src/python/omnimalloc/benchmark/sources/huggingface.py +++ b/src/python/omnimalloc/benchmark/sources/huggingface.py @@ -3,131 +3,65 @@ # import re -from collections import defaultdict -from importlib.util import find_spec +from collections.abc import Iterable from pathlib import Path -from typing import ClassVar +from typing import Any, ClassVar, Final, cast from omnimalloc.benchmark.converters.model import model_to_allocations -from omnimalloc.benchmark.converters.onnx import from_onnx +from omnimalloc.benchmark.converters.onnx import HAS_ONNX, from_onnx +from omnimalloc.common.constants import MB from omnimalloc.common.optional import require_optional -from omnimalloc.primitives import Allocation, IdType, Pool +from omnimalloc.primitives import Pool -from ..utils import tqdm # noqa: TID252 -from .base import BaseSource +from .base import FixedSource, prefix_ids try: - from huggingface_hub import HfApi, ModelInfo + from huggingface_hub import HfApi, RepoFile HAS_HUGGINGFACE_HUB = True except ImportError: HAS_HUGGINGFACE_HUB = False - HfApi = None # ty: ignore[invalid-assignment] - ModelInfo = None # ty: ignore[invalid-assignment] + HfApi = RepoFile = cast("Any", None) -HAS_ONNX = find_spec("onnx") is not None +_MIN_OPSET: Final[int] = 16 +_MAX_FILE_SIZE: Final[int] = 200 * MB -def _get_hf_api() -> HfApi: - """Get HfApi instance, checking that dependency is available.""" - if not HAS_HUGGINGFACE_HUB: - require_optional("huggingface-hub", "HuggingfaceSource") - return HfApi() - - -def _list_onnx_models(limit: int = 10) -> list[ModelInfo]: - """Return ONNX models from Hugging Face Hub, excluding the legacy repository.""" - hf_api = _get_hf_api() - models = hf_api.list_models(author="onnxmodelzoo", limit=limit + 1) - return [m for m in models if m.id != "onnxmodelzoo/legacy_models"][:limit] - - -def _filter_onnx_opsets( - model_infos: list[ModelInfo], min_opset: int = 16 -) -> list[ModelInfo]: - """Filter ONNX models to only include the highest opset per base model name.""" - model_groups = defaultdict(list) - - for model_info in model_infos: - match = re.search(r"Opset(\d+)", model_info.id) - if not match: +def _latest_opsets(repo_ids: Iterable[str]) -> list[str]: + """Keep the highest opset of each model; the zoo has one repo per opset.""" + latest: dict[str, tuple[int, str]] = {} + for repo_id in repo_ids: + match = re.search(r"Opset(\d+)", repo_id) + if match is None or int(match.group(1)) < _MIN_OPSET: continue opset = int(match.group(1)) - if opset < min_opset: - continue - base_name = re.sub(r"Opset\d+", "", model_info.id) - model_groups[base_name].append((opset, model_info)) - - return [max(models, key=lambda x: x[0])[1] for models in model_groups.values()] - - -def _gather_download_info( - model_infos: list[ModelInfo], - filename_filter: str, - max_file_size_mb: float | None = 200, -) -> dict[str, str]: - """Gather information about which models to download, filtering by size.""" - hf_api = _get_hf_api() - id_file_map = {} - - for model_info in model_infos: - repo_files = hf_api.list_repo_tree(model_info.id, recursive=True) - onnx_files = [ + base_name = repo_id.replace(match.group(0), "") + if base_name not in latest or opset > latest[base_name][0]: + latest[base_name] = (opset, repo_id) + return [repo_id for _, repo_id in latest.values()] + + +def _download_onnx_models(num_models: int, output_dir: str | Path | None) -> list[Path]: + """Download the first `num_models` single-file ONNX zoo models of bounded size.""" + api = HfApi() + listed = api.list_models(author="onnxmodelzoo", limit=5 * num_models) + paths: list[Path] = [] + for repo_id in _latest_opsets(model.id for model in listed): + if len(paths) == num_models: + break + files = [ f - for f in repo_files - if f.path.endswith(filename_filter) and hasattr(f, "size") + for f in api.list_repo_tree(repo_id, recursive=True) + if isinstance(f, RepoFile) and f.path.endswith(".onnx") ] - - if len(onnx_files) != 1: + if len(files) != 1 or not 0 < files[0].size <= _MAX_FILE_SIZE: continue + path = api.hf_hub_download(repo_id, files[0].path, local_dir=output_dir) + paths.append(Path(path)) + return paths - file_info = onnx_files[0] - if file_info.size is None: - continue - - # ty does not narrow file_info.size from the None check above - size_mb = file_info.size / (1024 * 1024) # ty: ignore[unsupported-operator] - if max_file_size_mb is not None and size_mb > max_file_size_mb: - continue - - id_file_map[model_info.id] = file_info.path - - return id_file_map - - -def _download_files( - id_file_map: dict[str, str], - output_dir: str | Path | None = None, - filename_filter: str = ".onnx", -) -> list[Path]: - """Download files from Hugging Face Hub and return their local paths.""" - hf_api = _get_hf_api() - desc = f"Downloading {len(id_file_map)} '{filename_filter}' models from HuggingFace" - local_paths = [] - for repo_id, filename in tqdm(id_file_map.items(), desc=desc, leave=False): - local_path = hf_api.hf_hub_download( - repo_id=repo_id, - filename=filename, - local_dir=output_dir, - ) - local_paths.append(Path(local_path)) - - return local_paths - - -def _download_onnx_models( - num_models: int = 10, output_dir: str | Path | None = None -) -> list[Path]: - """Download ONNX models and return local file paths.""" - models = _list_onnx_models(limit=num_models * 5) - filtered = _filter_onnx_opsets(models) - id_file_map = _gather_download_info(filtered, ".onnx", max_file_size_mb=200) - id_file_map_limited = dict(list(id_file_map.items())[:num_models]) - return _download_files(id_file_map_limited, output_dir, ".onnx") - - -class HuggingfaceSource(BaseSource): +class HuggingfaceSource(FixedSource): """Fixed source of Huggingface ONNX model allocations, one variant per model.""" _label_fields: ClassVar[tuple[str, ...]] = ("num_models", "output_dir") @@ -141,74 +75,18 @@ def __init__( require_optional("onnx", "HuggingfaceSource") if not HAS_HUGGINGFACE_HUB: require_optional("huggingface-hub", "HuggingfaceSource") - super().__init__() self.num_models = num_models self.output_dir = output_dir - self._model_paths: list[Path] | None = None - self._model_pools: dict[str, Pool] | None = None - self._downloaded_num_models: int | None = None - - def _ensure_downloaded(self) -> None: - # Re-download when num_models changed (e.g. grown via - # get_available_variants) since the cached download. - if ( - self._model_pools is not None - and self._downloaded_num_models == self.num_models - ): - return - - self._downloaded_num_models = self.num_models - self._model_paths = _download_onnx_models(self.num_models, self.output_dir) - self._model_pools = {} - for model_path in self._model_paths: - model = from_onnx(model_path) - allocations = model_to_allocations(model) - model_name = model_path.stem - pool = Pool(id=f"hf_{model_name}", allocations=tuple(allocations)) - self._model_pools[model_name] = pool - - def is_parameterizable(self) -> bool: - return False - - def get_available_variants(self, count: int | None = None) -> tuple[str, ...]: - if count is not None: - self.num_models = max(self.num_models, count) - self._ensure_downloaded() - assert self._model_pools is not None - return tuple(self._model_pools.keys()) - - def get_variant(self, variant_id: IdType) -> Pool: - self._ensure_downloaded() - assert self._model_pools is not None - - if isinstance(variant_id, int): - model_names = list(self._model_pools.keys()) - if not (0 <= variant_id < len(model_names)): - msg = f"Model index {variant_id} out of range [0, {len(model_names)})" - raise ValueError(msg) - return self._model_pools[model_names[variant_id]] - - if variant_id not in self._model_pools: - msg = f"Model '{variant_id}' not found in Huggingface source" - raise ValueError(msg) - - return self._model_pools[variant_id] - - def get_allocations( - self, num_allocations: int | None = None, skip: int = 0 - ) -> tuple[Allocation, ...]: - self._ensure_downloaded() - assert self._model_pools is not None - - all_allocations: list[Allocation] = [] - for pool in self._model_pools.values(): - all_allocations.extend(pool.allocations) - - if skip >= len(all_allocations): - return () - end = ( - len(all_allocations) if num_allocations is None else skip + num_allocations + def _load_pools(self) -> tuple[Pool, ...]: + # Models share tensor names, so their allocation ids need qualifying + return tuple( + prefix_ids( + Pool( + id=path.stem, + allocations=tuple(model_to_allocations(from_onnx(path))), + ) + ) + for path in _download_onnx_models(self.num_models, self.output_dir) ) - return tuple(all_allocations[skip:end]) diff --git a/src/python/omnimalloc/benchmark/sources/minimalloc.py b/src/python/omnimalloc/benchmark/sources/minimalloc.py index 2f964b7..0dd35f7 100644 --- a/src/python/omnimalloc/benchmark/sources/minimalloc.py +++ b/src/python/omnimalloc/benchmark/sources/minimalloc.py @@ -7,11 +7,10 @@ from pathlib import Path from typing import ClassVar -from omnimalloc.common.validation import ensure_positive from omnimalloc.io import load_allocation -from omnimalloc.primitives import Allocation, IdType, Pool +from omnimalloc.primitives import Pool -from .base import BaseSource +from .base import FixedSource, prefix_ids logger = logging.getLogger(__name__) @@ -39,24 +38,7 @@ def _checkout_csv_dir(subset: MinimallocSubset) -> Path | None: return None -def _prefix_ids(pool: Pool) -> Pool: - """CSV ids restart at 0 per file, so qualify them with the pool id.""" - return pool.with_allocations( - tuple( - Allocation( - id=f"{pool.id}_{alloc.id}", - size=alloc.size, - start=alloc.start, - end=alloc.end, - offset=alloc.offset, - kind=alloc.kind, - ) - for alloc in pool.allocations - ) - ) - - -class MinimallocSource(BaseSource): +class MinimallocSource(FixedSource): """Fixed source loading pools from a directory of Minimalloc CSV files. `csv_dir` defaults to the `subset`'s directory in a source checkout, which @@ -70,85 +52,26 @@ def __init__( subset: MinimallocSubset | str = MinimallocSubset.CHALLENGING, csv_dir: str | Path | None = None, ) -> None: + super().__init__() self.subset = MinimallocSubset(subset) # The label must carry an explicit csv_dir but not the checkout default self.csv_dir = Path(csv_dir) if csv_dir is not None else None - self._cached_pools: list[Pool] | None = None - - # An empty dataset keeps the base invariant; the accessors raise instead - num_allocs = sum(len(p.allocations) for p in self._pools) - super().__init__(num_allocations=max(num_allocs, 1)) - - @property - def _pools(self) -> list[Pool]: - if self._cached_pools is None: - csv_dir = ( - self.csv_dir - if self.csv_dir is not None - else _checkout_csv_dir(self.subset) + + def _load_pools(self) -> tuple[Pool, ...]: + csv_dir = self.csv_dir or _checkout_csv_dir(self.subset) + if csv_dir is None: + logger.warning( + f"Not running from a source checkout, so the " + f"{self.subset.value!r} subset has no datasets; pass " + "csv_dir to read them from an install." ) - # Sort for a filesystem-independent, reproducible variant order - files = sorted(csv_dir.glob("*.csv")) if csv_dir is not None else [] - if csv_dir is None: - logger.warning( - f"Not running from a source checkout, so the " - f"{self.subset.value!r} subset has no datasets; pass " - "csv_dir to read them from an install." - ) - elif not files: - logger.warning( - f"No Minimalloc CSVs found in {csv_dir}; the " - f"{self.subset.value!r} subset yields no variants." - ) - self._cached_pools = [_prefix_ids(load_allocation(f)) for f in files] - return self._cached_pools - - def _all_allocations(self) -> tuple[Allocation, ...]: - return tuple(alloc for pool in self._pools for alloc in pool.allocations) - - def is_parameterizable(self) -> bool: - """Minimalloc has fixed pools, not parameterizable.""" - return False - - def get_available_variants( - self, - count: int | None = None, # noqa: ARG002 - ) -> tuple[str, ...]: - """Return pool IDs from Minimalloc benchmarks.""" - return tuple(str(pool.id) for pool in self._pools) - - def get_variant(self, variant_id: IdType) -> Pool: - """Get a specific Minimalloc pool by name.""" - if isinstance(variant_id, int): - # Support integer indexing - if 0 <= variant_id < len(self._pools): - return self._pools[variant_id] - msg = f"Pool index {variant_id} out of range [0, {len(self._pools)})" - raise ValueError(msg) - - # String lookup by pool ID - for pool in self._pools: - if pool.id == variant_id: - return pool - - raise ValueError(f"Pool with ID '{variant_id}' not found in Minimalloc source") - - def get_allocations( - self, num_allocations: int | None = None, skip: int = 0 - ) -> tuple[Allocation, ...]: - all_allocations = self._all_allocations() - if skip >= len(all_allocations): - return () - if num_allocations is None: - return all_allocations[skip:] - return all_allocations[skip : skip + num_allocations] - - def get_pools( - self, num_pools: int | None = None, skip: int = 0 - ) -> tuple[Pool, ...]: - ensure_positive(num_pools, "num_pools", allow_none=True) - if skip >= len(self._pools): return () - if num_pools is None: - return tuple(self._pools[skip:]) - return tuple(self._pools[skip : skip + num_pools]) + # Sort for a filesystem-independent, reproducible variant order + files = sorted(csv_dir.glob("*.csv")) + if not files: + logger.warning( + f"No Minimalloc CSVs found in {csv_dir}; the " + f"{self.subset.value!r} subset yields no variants." + ) + # CSV ids restart at 0 per file + return tuple(prefix_ids(load_allocation(f)) for f in files) diff --git a/src/python/omnimalloc/benchmark/sources/sizes.py b/src/python/omnimalloc/benchmark/sources/sizes.py index 9cefb2c..195d532 100644 --- a/src/python/omnimalloc/benchmark/sources/sizes.py +++ b/src/python/omnimalloc/benchmark/sources/sizes.py @@ -7,6 +7,8 @@ from enum import Enum from typing import Final +from omnimalloc.common.validation import ensure_positive + class SizeDistribution(str, Enum): """Size distribution families, ordered from flattest to most skewed. @@ -46,8 +48,7 @@ def sample_sizes( `bimodal` the 90/10 accelerator mix, `dominant` one buffer at 90%. """ distribution = SizeDistribution(distribution) - if size_min <= 0: - raise ValueError("size_min must be positive") + ensure_positive(size_min, "size_min") if size_max < size_min: raise ValueError("size_max must be >= size_min") if count <= 0: @@ -81,7 +82,7 @@ def _zipf(rng: random.Random, size_min: int, size_max: int) -> int: def _bimodal(rng: random.Random, count: int, size_min: int, size_max: int) -> list[int]: - boundary = max(size_min, math.isqrt(size_min * size_max)) + boundary = math.isqrt(size_min * size_max) sizes = [] for _ in range(count): if rng.random() < _SMALL_FRACTION: diff --git a/src/python/omnimalloc/benchmark/sources/sync_patterns.py b/src/python/omnimalloc/benchmark/sources/sync_patterns.py index e3eebdd..24998f3 100644 --- a/src/python/omnimalloc/benchmark/sources/sync_patterns.py +++ b/src/python/omnimalloc/benchmark/sources/sync_patterns.py @@ -9,6 +9,7 @@ from typing import ClassVar, Final from omnimalloc.common.constants import DEFAULT_SEED, KB, MB +from omnimalloc.common.validation import ensure_positive from omnimalloc.primitives import Allocation, VectorClock from .base import BaseSource @@ -72,22 +73,16 @@ def __init__( max_lifetime: int | None = None, seed: int | None = DEFAULT_SEED, ) -> None: - if num_threads <= 0: - raise ValueError("num_threads must be positive") - if speed_skew <= 0: - raise ValueError("speed_skew must be positive") + ensure_positive(num_threads, "num_threads") + ensure_positive(speed_skew, "speed_skew") if steps is not None and steps < 2: raise ValueError("steps must be >= 2") - if sync_period <= 0: - raise ValueError("sync_period must be positive") - if group_size is not None and group_size <= 0: - raise ValueError("group_size must be positive") - if size_min <= 0: - raise ValueError("size_min must be positive") + ensure_positive(sync_period, "sync_period") + ensure_positive(group_size, "group_size", allow_none=True) + ensure_positive(size_min, "size_min") if size_max < size_min: raise ValueError("size_max must be >= size_min") - if max_lifetime is not None and max_lifetime <= 0: - raise ValueError("max_lifetime must be positive") + ensure_positive(max_lifetime, "max_lifetime", allow_none=True) super().__init__(num_allocations=num_allocations) self.num_threads = num_threads self.pattern = SyncPattern(pattern) @@ -192,7 +187,7 @@ def _sync_subset( self, clocks: list[list[int]], step: int, _rng: random.Random ) -> None: if step % self.sync_period == 0: - _merge(clocks, range(self.num_threads // 2 or 1)) + _merge(clocks, range(max(2, self.num_threads // 2))) def _sync_sparse( self, clocks: list[list[int]], step: int, rng: random.Random diff --git a/src/python/omnimalloc/benchmark/sources/tiling.py b/src/python/omnimalloc/benchmark/sources/tiling.py index 31a8271..b6a3bbe 100644 --- a/src/python/omnimalloc/benchmark/sources/tiling.py +++ b/src/python/omnimalloc/benchmark/sources/tiling.py @@ -18,11 +18,7 @@ class TilingSource(TilingBase): """ _label_fields: ClassVar[tuple[str, ...]] = ( - "capacity", - "makespan", - "size_min", - "duration_min", - "seed", + *TilingBase._label_fields, # noqa: SLF001 "mem_cut_prob", ) diff --git a/src/python/omnimalloc/benchmark/sources/tiling_base.py b/src/python/omnimalloc/benchmark/sources/tiling_base.py index 6696c33..3608818 100644 --- a/src/python/omnimalloc/benchmark/sources/tiling_base.py +++ b/src/python/omnimalloc/benchmark/sources/tiling_base.py @@ -8,7 +8,8 @@ from dataclasses import dataclass from typing import ClassVar, Generic, TypeVar -from omnimalloc.primitives import Allocation, IdType, Pool, TimePoint +from omnimalloc.common.validation import ensure_positive +from omnimalloc.primitives import Allocation, Pool, TimePoint from .base import BaseSource @@ -54,10 +55,8 @@ def __init__( seed: int | None, ) -> None: super().__init__(num_allocations=num_allocations) - if size_min <= 0: - raise ValueError("size_min must be positive") - if duration_min <= 0: - raise ValueError("duration_min must be positive") + ensure_positive(size_min, "size_min") + ensure_positive(duration_min, "duration_min") self.capacity = capacity self.makespan = makespan self.size_min = size_min @@ -140,13 +139,6 @@ def get_ground_truth_pool( allocations = self._tile_allocations(num_allocations, skip, with_offsets=True) return Pool(id=f"{self.name()}_ground_truth", allocations=allocations) - def get_known_optimum(self, variant_id: IdType | None = None) -> int | None: - """Peak of the construction packing, which is optimal by tiling. - - Unknown without a fixed seed: the ground truth is then a different - random packing than the one that was benchmarked. - """ - if self.seed is None: - return None - num = variant_id if isinstance(variant_id, int) else None - return self.get_ground_truth_pool(num).size + def get_known_optimum(self) -> int: + """The capacity: the leaves tile it exactly, whatever the seed.""" + return self.capacity diff --git a/tests/unit/benchmark/converters/test_model.py b/tests/unit/benchmark/converters/test_model.py index e5baa36..d204da7 100644 --- a/tests/unit/benchmark/converters/test_model.py +++ b/tests/unit/benchmark/converters/test_model.py @@ -7,687 +7,129 @@ Buffer, Model, Op, - _compute_buffer_lifetimes, - _create_allocations, model_to_allocations, - model_to_pools, - model_to_system, ) from omnimalloc.primitives import AllocationKind +WORKSPACE = AllocationKind.WORKSPACE +CONSTANT = AllocationKind.CONSTANT +INPUT = AllocationKind.INPUT +OUTPUT = AllocationKind.OUTPUT + def _buf( buf_id: int | str, - kind: AllocationKind = AllocationKind.WORKSPACE, + kind: AllocationKind = WORKSPACE, shape: tuple[int, ...] = (10,), + dtype: str = "float32", ) -> Buffer: - return Buffer(id=buf_id, shape=shape, dtype="float32", kind=kind) - - -def test_buffer_basic_creation_int_id() -> None: - buffer = Buffer( - id=0, shape=(10, 20), dtype="float32", kind=AllocationKind.WORKSPACE - ) - assert buffer.id == 0 - assert buffer.shape == (10, 20) - assert buffer.dtype == "float32" - assert buffer.kind == AllocationKind.WORKSPACE - - -def test_buffer_basic_creation_str_id() -> None: - buffer = Buffer(id="buf_0", shape=(10,), dtype="int8", kind=AllocationKind.CONSTANT) - assert buffer.id == "buf_0" - assert buffer.shape == (10,) - assert buffer.dtype == "int8" - assert buffer.kind == AllocationKind.CONSTANT - - -def test_buffer_ndim_1d() -> None: - assert _buf(0, shape=(10,)).ndim == 1 - - -def test_buffer_ndim_2d() -> None: - assert _buf(0, shape=(10, 20)).ndim == 2 - - -def test_buffer_ndim_4d() -> None: - assert _buf(0, shape=(1, 3, 224, 224)).ndim == 4 - - -def test_buffer_size_float32() -> None: - buffer = Buffer( - id=0, shape=(10, 20), dtype="float32", kind=AllocationKind.WORKSPACE - ) - assert buffer.size == 10 * 20 * 4 - - -def test_buffer_size_int8() -> None: - buffer = Buffer(id=0, shape=(100,), dtype="int8", kind=AllocationKind.WORKSPACE) - assert buffer.size == 100 - - -def test_buffer_size_float64() -> None: - buffer = Buffer(id=0, shape=(5, 5), dtype="float64", kind=AllocationKind.WORKSPACE) - assert buffer.size == 5 * 5 * 8 - - -def test_buffer_size_complex_shape() -> None: - buffer = Buffer( - id=0, - shape=(2, 3, 4, 5), - dtype="float32", - kind=AllocationKind.WORKSPACE, - ) - assert buffer.size == 2 * 3 * 4 * 5 * 4 - - -def test_buffer_size_packs_sub_byte_dtypes() -> None: - buffer = Buffer(id=0, shape=(10, 20), dtype="int4", kind=AllocationKind.WORKSPACE) - assert buffer.size == 10 * 20 // 2 - - -def test_buffer_size_rounds_an_odd_sub_byte_count_up() -> None: - buffer = Buffer(id=0, shape=(7,), dtype="uint4", kind=AllocationKind.WORKSPACE) - assert buffer.size == 4 + return Buffer(id=buf_id, shape=shape, dtype=dtype, kind=kind) + + +@pytest.mark.parametrize( + ("shape", "dtype", "size"), + [ + ((10, 20), "float32", 800), + ((100,), "int8", 100), + ((5, 5), "float64", 200), + ((2, 3, 4, 5), "float32", 480), + ((10, 20), "int4", 100), + ((7,), "uint4", 4), + ], +) +def test_buffer_size(shape: tuple[int, ...], dtype: str, size: int) -> None: + assert _buf(0, shape=shape, dtype=dtype).size == size -def test_buffer_unknown_dtype() -> None: +def test_buffer_rejects_unknown_dtype() -> None: with pytest.raises(ValueError, match="unknown dtype 'float24'"): - Buffer(id=0, shape=(10,), dtype="float24", kind=AllocationKind.WORKSPACE) - - -def test_buffer_invalid_shape_zero() -> None: - with pytest.raises(ValueError, match="shape dimensions must be positive integers"): - _buf(0, shape=(10, 0)) - - -def test_buffer_invalid_shape_negative() -> None: - with pytest.raises(ValueError, match="shape dimensions must be positive integers"): - _buf(0, shape=(10, -5)) + _buf(0, dtype="float24") -def test_buffer_invalid_shape_float() -> None: +@pytest.mark.parametrize("shape", [(10, 0), (10, -5), (10.5, 20)]) +def test_buffer_rejects_non_positive_integer_shapes(shape: tuple[int, ...]) -> None: with pytest.raises(ValueError, match="shape dimensions must be positive integers"): - _buf(0, shape=(10.5, 20)) # type: ignore[arg-type] - - -def test_buffer_various_kinds() -> None: - for kind in AllocationKind: - assert _buf(0, kind).kind == kind - - -def test_op_basic_creation_int_id() -> None: - op = Op(id=0) - assert op.id == 0 - assert op.inputs == set() - assert op.outputs == set() - assert op.op_type is None - - -def test_op_basic_creation_str_id() -> None: - op = Op(id="conv1") - assert op.id == "conv1" - assert op.inputs == set() - assert op.outputs == set() - - -def test_op_with_inputs_and_outputs() -> None: - buf_in = _buf(0, AllocationKind.INPUT) - buf_out = _buf(1, AllocationKind.OUTPUT) - op = Op(id=0, inputs={buf_in}, outputs={buf_out}) - assert buf_in in op.inputs - assert buf_out in op.outputs - - -def test_op_with_op_type() -> None: - op = Op(id=0, op_type="Conv2D") - assert op.op_type == "Conv2D" + _buf(0, shape=shape) -def test_op_invalid_id_type() -> None: +def test_op_rejects_invalid_id_type() -> None: with pytest.raises(TypeError, match="id must be int or str"): Op(id=3.14) # type: ignore[arg-type] -def test_op_duplicate_buffer_ids_input_output() -> None: - buf = _buf(0) +def test_op_rejects_a_buffer_id_used_twice() -> None: with pytest.raises(ValueError, match="buffer ids must be unique"): - Op(id=0, inputs={buf}, outputs={buf}) - - -def test_op_duplicate_buffer_ids_multiple_inputs() -> None: - buf1 = _buf(0) - buf2 = _buf(0, shape=(20,)) + Op(id=0, inputs={_buf(0)}, outputs={_buf(0)}) with pytest.raises(ValueError, match="buffer ids must be unique"): - Op(id=0, inputs={buf1, buf2}) - - -def test_op_unique_buffer_ids_different_buffers() -> None: - op = Op(id=0, inputs={_buf(0)}, outputs={_buf(1, shape=(20,))}) - assert len(op.inputs) == 1 - assert len(op.outputs) == 1 - - -def test_op_multiple_inputs_outputs() -> None: - op = Op( - id=0, - inputs={_buf(0, AllocationKind.INPUT), _buf(1, AllocationKind.INPUT)}, - outputs={_buf(2, AllocationKind.OUTPUT), _buf(3, AllocationKind.OUTPUT)}, - ) - assert len(op.inputs) == 2 - assert len(op.outputs) == 2 - - -def test_model_basic_creation_int_id() -> None: - model = Model(id=0) - assert model.id == 0 - assert model.ops == {} - assert model.buffers == {} - + Op(id=0, inputs={_buf(0), _buf(0, shape=(20,))}) -def test_model_basic_creation_str_id() -> None: - model = Model(id="resnet50") - assert model.id == "resnet50" - -def test_model_with_ops_and_buffers() -> None: - buf = _buf(0) - op = Op(id=0, outputs={buf}) - model = Model(id=0, ops={0: op}, buffers={0: buf}) - assert 0 in model.ops - assert 0 in model.buffers - - -def test_model_invalid_id_type() -> None: +def test_model_rejects_invalid_id_type() -> None: with pytest.raises(TypeError, match="id must be int or str"): Model(id=3.14) # type: ignore[arg-type] -def test_model_duplicate_op_ids() -> None: +def test_model_rejects_duplicate_ids() -> None: with pytest.raises(ValueError, match="op ids must be unique"): Model(id=0, ops={0: Op(id=0), 1: Op(id=0)}) - - -def test_model_duplicate_buffer_ids() -> None: with pytest.raises(ValueError, match="buffer ids must be unique"): Model(id=0, buffers={0: _buf(0), 1: _buf(0, shape=(20,))}) -def test_model_unique_op_ids() -> None: - model = Model(id=0, ops={0: Op(id=0), 1: Op(id=1)}) - assert len(model.ops) == 2 - - -def test_model_unique_buffer_ids() -> None: - model = Model(id=0, buffers={0: _buf(0), 1: _buf(1, shape=(20,))}) - assert len(model.buffers) == 2 - - -def test_compute_buffer_lifetimes_basic() -> None: - buf1 = _buf(0) - buf2 = _buf(1) - op1 = Op(id=0, outputs={buf1}) - op2 = Op(id=1, inputs={buf1}, outputs={buf2}) - model = Model(id=0, ops={0: op1, 1: op2}, buffers={0: buf1, 1: buf2}) - - first_index, last_index = _compute_buffer_lifetimes( - model, const_inf_lifetime=False, io_inf_lifetime=False - ) - assert first_index[buf1] == 0 - assert last_index[buf1] == 1 - assert first_index[buf2] == 1 - assert last_index[buf2] == 1 - - -def test_compute_buffer_lifetimes_const_inf_lifetime() -> None: - buf_const = _buf(0, AllocationKind.CONSTANT) - buf_work = _buf(1) - op1 = Op(id=0, outputs={buf_const}) - op2 = Op(id=1, inputs={buf_const}, outputs={buf_work}) - model = Model(id=0, ops={0: op1, 1: op2}, buffers={0: buf_const, 1: buf_work}) - - first_index, last_index = _compute_buffer_lifetimes( - model, const_inf_lifetime=True, io_inf_lifetime=False - ) - assert first_index[buf_const] == 0 - assert last_index[buf_const] == 1 - assert first_index[buf_work] == 1 - assert last_index[buf_work] == 1 - - -def test_compute_buffer_lifetimes_io_inf_lifetime() -> None: - buf_input = _buf(0, AllocationKind.INPUT) - buf_output = _buf(1, AllocationKind.OUTPUT) - op1 = Op(id=0, inputs={buf_input}, outputs={buf_output}) - model = Model(id=0, ops={0: op1}, buffers={0: buf_input, 1: buf_output}) - - first_index, last_index = _compute_buffer_lifetimes( - model, const_inf_lifetime=False, io_inf_lifetime=True - ) - assert first_index[buf_input] == 0 - assert last_index[buf_input] == 0 - assert first_index[buf_output] == 0 - assert last_index[buf_output] == 0 - - -def test_compute_buffer_lifetimes_multiple_uses() -> None: - buf = _buf(0) - op1 = Op(id=0, outputs={buf}) - op2 = Op(id=1, inputs={buf}) - op3 = Op(id=2, inputs={buf}) - model = Model(id=0, ops={0: op1, 1: op2, 2: op3}, buffers={0: buf}) - - first_index, last_index = _compute_buffer_lifetimes( - model, const_inf_lifetime=False, io_inf_lifetime=False - ) - assert first_index[buf] == 0 - assert last_index[buf] == 2 +def _chain() -> Model: + """Input -> a -> b -> output through three ops, reading one constant.""" + bufs = { + "input": _buf("input", INPUT), + "const": _buf("const", CONSTANT), + "a": _buf("a"), + "b": _buf("b"), + "output": _buf("output", OUTPUT), + } + ops = { + 0: Op(id=0, inputs={bufs["input"]}, outputs={bufs["a"]}), + 1: Op(id=1, inputs={bufs["a"]}, outputs={bufs["b"]}), + 2: Op(id=2, inputs={bufs["b"], bufs["const"]}, outputs={bufs["output"]}), + } + return Model(id="chain", ops=ops, buffers=bufs) -def test_compute_buffer_lifetimes_all_kinds() -> None: - buf_work = _buf(0) - buf_const = _buf(1, AllocationKind.CONSTANT) - buf_input = _buf(2, AllocationKind.INPUT) - buf_output = _buf(3, AllocationKind.OUTPUT) - op = Op(id=0, inputs={buf_input, buf_const}, outputs={buf_work, buf_output}) - model = Model( - id=0, - ops={0: op}, - buffers={0: buf_work, 1: buf_const, 2: buf_input, 3: buf_output}, - ) +def _lifetimes(**kwargs: bool) -> dict[object, tuple[object, object]]: + return {a.id: (a.start, a.end) for a in model_to_allocations(_chain(), **kwargs)} - first_index, last_index = _compute_buffer_lifetimes( - model, const_inf_lifetime=True, io_inf_lifetime=True - ) - for buf in (buf_work, buf_const, buf_input, buf_output): - assert first_index[buf] == 0 - assert last_index[buf] == 0 - - -def test_create_allocations_basic() -> None: - buf = _buf(0) - op = Op(id=0, outputs={buf}) - model = Model(id=0, ops={0: op}, buffers={0: buf}) - - allocations = _create_allocations( - model, - include_const=True, - include_io=True, - buffer_to_first_index={buf: 0}, - buffer_to_last_index={buf: 0}, - ) - assert len(allocations) == 1 - assert allocations[0].id == 0 - assert allocations[0].size == 40 - assert allocations[0].start == 0 - assert allocations[0].end == 1 - assert allocations[0].kind == AllocationKind.WORKSPACE - - -def test_create_allocations_exclude_const() -> None: - buf_work = _buf(0) - buf_const = _buf(1, AllocationKind.CONSTANT) - model = Model(id=0, buffers={0: buf_work, 1: buf_const}) - - allocations = _create_allocations( - model, - include_const=False, - include_io=True, - buffer_to_first_index={buf_work: 0, buf_const: 0}, - buffer_to_last_index={buf_work: 0, buf_const: 0}, - ) - assert len(allocations) == 1 - assert allocations[0].kind == AllocationKind.WORKSPACE - - -def test_create_allocations_exclude_io() -> None: - buf_work = _buf(0) - buf_input = _buf(1, AllocationKind.INPUT) - buf_output = _buf(2, AllocationKind.OUTPUT) - model = Model(id=0, buffers={0: buf_work, 1: buf_input, 2: buf_output}) - indices = {buf_work: 0, buf_input: 0, buf_output: 0} - - allocations = _create_allocations( - model, - include_const=True, - include_io=False, - buffer_to_first_index=indices, - buffer_to_last_index=indices, - ) - assert len(allocations) == 1 - assert allocations[0].kind == AllocationKind.WORKSPACE +def test_model_to_allocations_keeps_workspace_by_default() -> None: + assert _lifetimes() == {"a": (0, 2), "b": (1, 3)} -def test_create_allocations_include_all() -> None: - buf_work = _buf(0) - buf_const = _buf(1, AllocationKind.CONSTANT) - buf_input = _buf(2, AllocationKind.INPUT) - buf_output = _buf(3, AllocationKind.OUTPUT) - model = Model( - id=0, buffers={0: buf_work, 1: buf_const, 2: buf_input, 3: buf_output} - ) - indices = {buf_work: 0, buf_const: 0, buf_input: 0, buf_output: 0} - - allocations = _create_allocations( - model, - include_const=True, - include_io=True, - buffer_to_first_index=indices, - buffer_to_last_index=indices, - ) - assert len(allocations) == 4 +def test_model_to_allocations_spans_constants_and_io_over_the_model() -> None: + lifetimes = _lifetimes(include_const=True, include_io=True) + assert lifetimes["const"] == lifetimes["input"] == lifetimes["output"] == (0, 3) -def test_create_allocations_exclude_all() -> None: - buf_const = _buf(0, AllocationKind.CONSTANT) - buf_input = _buf(1, AllocationKind.INPUT) - model = Model(id=0, buffers={0: buf_const, 1: buf_input}) - indices = {buf_const: 0, buf_input: 0} - allocations = _create_allocations( - model, - include_const=False, - include_io=False, - buffer_to_first_index=indices, - buffer_to_last_index=indices, - ) - assert len(allocations) == 0 - - -def test_create_allocations_end_index_increment() -> None: - buf = _buf(0) - model = Model(id=0, buffers={0: buf}) - - allocations = _create_allocations( - model, - include_const=True, - include_io=True, - buffer_to_first_index={buf: 5}, - buffer_to_last_index={buf: 10}, - ) - assert allocations[0].start == 5 - assert allocations[0].end == 11 - - -def test_model_to_allocations_basic() -> None: - buf = _buf(0) - op = Op(id=0, outputs={buf}) - model = Model(id=0, ops={0: op}, buffers={0: buf}) - - allocations = model_to_allocations(model) - assert len(allocations) == 1 - assert allocations[0].size == 40 - - -def test_model_to_allocations_exclude_const() -> None: - buf_work = _buf(0) - buf_const = _buf(1, AllocationKind.CONSTANT) - op = Op(id=0, inputs={buf_const}, outputs={buf_work}) - model = Model(id=0, ops={0: op}, buffers={0: buf_work, 1: buf_const}) - - allocations = model_to_allocations(model, include_const=False) - assert len(allocations) == 1 - assert all(a.kind != AllocationKind.CONSTANT for a in allocations) - - -def test_model_to_allocations_include_const() -> None: - buf_work = _buf(0) - buf_const = _buf(1, AllocationKind.CONSTANT) - op = Op(id=0, inputs={buf_const}, outputs={buf_work}) - model = Model(id=0, ops={0: op}, buffers={0: buf_work, 1: buf_const}) - - allocations = model_to_allocations(model, include_const=True) - assert len(allocations) == 2 - - -def test_model_to_allocations_exclude_io() -> None: - buf_work = _buf(0) - buf_input = _buf(1, AllocationKind.INPUT) - op = Op(id=0, inputs={buf_input}, outputs={buf_work}) - model = Model(id=0, ops={0: op}, buffers={0: buf_work, 1: buf_input}) - - allocations = model_to_allocations(model, include_io=False) - assert len(allocations) == 1 - assert all(not a.kind.is_io for a in allocations if a.kind) - - -def test_model_to_allocations_include_io() -> None: - buf_work = _buf(0) - buf_input = _buf(1, AllocationKind.INPUT) - op = Op(id=0, inputs={buf_input}, outputs={buf_work}) - model = Model(id=0, ops={0: op}, buffers={0: buf_work, 1: buf_input}) - - allocations = model_to_allocations(model, include_io=True) - assert len(allocations) == 2 - - -def test_model_to_allocations_const_inf_lifetime() -> None: - buf_const = _buf(0, AllocationKind.CONSTANT) - op1 = Op(id=0, outputs={buf_const}) - op2 = Op(id=1) - op3 = Op(id=2, inputs={buf_const}) - model = Model(id=0, ops={0: op1, 1: op2, 2: op3}, buffers={0: buf_const}) - - allocations = model_to_allocations( - model, include_const=True, const_inf_lifetime=True - ) - assert len(allocations) == 1 - assert allocations[0].start == 0 - assert allocations[0].end == 3 - - -def test_model_to_allocations_io_inf_lifetime() -> None: - buf_input = _buf(0, AllocationKind.INPUT) - op1 = Op(id=0, inputs={buf_input}) - op2 = Op(id=1) - model = Model(id=0, ops={0: op1, 1: op2}, buffers={0: buf_input}) - - allocations = model_to_allocations(model, include_io=True, io_inf_lifetime=True) - assert len(allocations) == 1 - assert allocations[0].start == 0 - assert allocations[0].end == 2 - - -def test_model_to_allocations_complex_model() -> None: - buf1 = _buf(0, AllocationKind.INPUT) - buf2 = _buf(1) - buf3 = _buf(2) - buf4 = _buf(3, AllocationKind.OUTPUT) - op1 = Op(id=0, inputs={buf1}, outputs={buf2}) - op2 = Op(id=1, inputs={buf2}, outputs={buf3}) - op3 = Op(id=2, inputs={buf3}, outputs={buf4}) - model = Model( - id=0, - ops={0: op1, 1: op2, 2: op3}, - buffers={0: buf1, 1: buf2, 2: buf3, 3: buf4}, - ) - - allocations = model_to_allocations(model, include_io=True) - assert len(allocations) == 4 - - -def test_model_to_pools_basic() -> None: - buf = _buf(0) - op = Op(id=0, outputs={buf}) - model = Model(id=0, ops={0: op}, buffers={0: buf}) - - pools = model_to_pools(model) - assert len(pools) == 1 - assert pools[0].id == "workspace" - - -def test_model_to_pools_grouped_by_kind() -> None: - buf_work1 = _buf(0) - buf_work2 = _buf(1) - buf_const = _buf(2, AllocationKind.CONSTANT) - op = Op(id=0, inputs={buf_const}, outputs={buf_work1, buf_work2}) - model = Model(id=0, ops={0: op}, buffers={0: buf_work1, 1: buf_work2, 2: buf_const}) - - pools = model_to_pools(model) - assert len(pools) == 2 - pool_ids = {pool.id for pool in pools} - assert "workspace" in pool_ids - assert "constant" in pool_ids - - -def test_model_to_pools_all_kinds() -> None: - buf_work = _buf(0) - buf_const = _buf(1, AllocationKind.CONSTANT) - buf_input = _buf(2, AllocationKind.INPUT) - buf_output = _buf(3, AllocationKind.OUTPUT) - op = Op(id=0, inputs={buf_input, buf_const}, outputs={buf_work, buf_output}) - model = Model( - id=0, - ops={0: op}, - buffers={0: buf_work, 1: buf_const, 2: buf_input, 3: buf_output}, - ) - - pools = model_to_pools(model) - assert len(pools) == 4 - assert {pool.id for pool in pools} == {"workspace", "constant", "input", "output"} - - -def test_model_to_pools_include_const_false() -> None: - buf_work = _buf(0) - buf_const = _buf(1, AllocationKind.CONSTANT) - op = Op(id=0, inputs={buf_const}, outputs={buf_work}) - model = Model(id=0, ops={0: op}, buffers={0: buf_work, 1: buf_const}) - - pools = model_to_pools(model, include_const=False) - assert "constant" not in {pool.id for pool in pools} - - -def test_model_to_pools_include_io_true() -> None: - buf_work = _buf(0) - buf_input = _buf(1, AllocationKind.INPUT) - op = Op(id=0, inputs={buf_input}, outputs={buf_work}) - model = Model(id=0, ops={0: op}, buffers={0: buf_work, 1: buf_input}) - - pools = model_to_pools(model, include_io=True) - assert "input" in {pool.id for pool in pools} - - -def test_model_to_pools_include_io_false() -> None: - buf_work = _buf(0) - buf_input = _buf(1, AllocationKind.INPUT) - op = Op(id=0, inputs={buf_input}, outputs={buf_work}) - model = Model(id=0, ops={0: op}, buffers={0: buf_work, 1: buf_input}) - - pools = model_to_pools(model, include_io=False) - assert "input" not in {pool.id for pool in pools} - - -def test_model_to_pools_allocations_count_per_pool() -> None: - buf_work1 = _buf(0) - buf_work2 = _buf(1) - buf_work3 = _buf(2) - op = Op(id=0, outputs={buf_work1, buf_work2, buf_work3}) - model = Model(id=0, ops={0: op}, buffers={0: buf_work1, 1: buf_work2, 2: buf_work3}) - - pools = model_to_pools(model) - assert len(pools) == 1 - assert len(pools[0].allocations) == 3 - - -def test_model_to_system_basic() -> None: - buf = _buf(0) - op = Op(id=0, outputs={buf}) - model = Model(id=0, ops={0: op}, buffers={0: buf}) - - system = model_to_system(model) - assert system.id == 0 - assert len(system.memories) == 1 - - -def test_model_to_system_model_id_preserved() -> None: - buf = _buf(0) - op = Op(id=0, outputs={buf}) - model = Model(id="resnet50", ops={0: op}, buffers={0: buf}) - - system = model_to_system(model) - assert system.id == "resnet50" - - -def test_model_to_system_memory_structure() -> None: - buf = _buf(0) - op = Op(id=0, outputs={buf}) - model = Model(id=0, ops={0: op}, buffers={0: buf}) - - system = model_to_system(model) - assert len(system.memories) == 1 - assert system.memories[0].id == 0 - assert len(system.memories[0].pools) >= 1 - - -def test_model_to_system_pools_included() -> None: - buf_work = _buf(0) - buf_const = _buf(1, AllocationKind.CONSTANT) - buf_input = _buf(2, AllocationKind.INPUT) - buf_output = _buf(3, AllocationKind.OUTPUT) - op = Op(id=0, inputs={buf_input, buf_const}, outputs={buf_work, buf_output}) - model = Model( - id=0, - ops={0: op}, - buffers={0: buf_work, 1: buf_const, 2: buf_input, 3: buf_output}, - ) - - system = model_to_system(model) - pools = system.memories[0].pools - assert len(pools) == 4 - assert {pool.id for pool in pools} == {"workspace", "constant", "input", "output"} - - -def test_model_to_system_empty_model() -> None: - system = model_to_system(Model(id=0)) - assert len(system.memories) == 1 - assert len(system.memories[0].pools) == 0 - - -def test_model_to_pools_skips_unreferenced_workspace_buffers() -> None: - used = _buf("used") - unused = _buf("unused") - op = Op(id=0, outputs={used}) - model = Model(id=0, ops={0: op}, buffers={"used": used, "unused": unused}) - - pools = model_to_pools(model) - - allocation_ids = {a.id for pool in pools for a in pool.allocations} - assert allocation_ids == {"used"} +@pytest.mark.parametrize( + ("kwargs", "extra"), + [ + ({"include_const": True}, {"const"}), + ({"include_io": True}, {"input", "output"}), + ], +) +def test_model_to_allocations_includes_by_kind( + kwargs: dict[str, bool], extra: set[str] +) -> None: + assert set(_lifetimes(**kwargs)) == {"a", "b"} | extra -def test_model_to_pools_keeps_unreferenced_io_and_constant_buffers() -> None: - used = _buf("used") - spare_input = _buf("spare_input", AllocationKind.INPUT) - spare_const = _buf("spare_const", AllocationKind.CONSTANT) - op = Op(id=0, outputs={used}) +def test_model_to_allocations_skips_unreferenced_buffers() -> None: + used, unused = _buf("used"), _buf("unused") model = Model( id=0, - ops={0: op}, - buffers={"used": used, "spare_input": spare_input, "spare_const": spare_const}, + ops={0: Op(id=0, outputs={used})}, + buffers={"used": used, "unused": unused}, ) - - pools = model_to_pools(model) - - allocation_ids = {a.id for pool in pools for a in pool.allocations} - assert allocation_ids == {"used", "spare_input", "spare_const"} - - -def test_model_to_allocations_skips_unreferenced_buffers() -> None: - used = _buf("used") - unused = _buf("unused") - op = Op(id=0, outputs={used}) - model = Model(id=0, ops={0: op}, buffers={"used": used, "unused": unused}) - - allocations = model_to_allocations(model, const_inf_lifetime=False) - - assert [a.id for a in allocations] == ["used"] + assert [a.id for a in model_to_allocations(model)] == ["used"] def test_model_to_allocations_handles_model_without_ops() -> None: - buf = _buf("io", AllocationKind.INPUT) - model = Model(id=0, buffers={"io": buf}) - - allocations = model_to_allocations(model, include_io=True, io_inf_lifetime=True) - - assert len(allocations) == 1 - assert (allocations[0].start, allocations[0].end) == (0, 1) + model = Model(id=0, buffers={"io": _buf("io", INPUT)}) + allocations = model_to_allocations(model, include_io=True) + assert [(a.start, a.end) for a in allocations] == [(0, 1)] diff --git a/tests/unit/benchmark/converters/test_onnx.py b/tests/unit/benchmark/converters/test_onnx.py index cc9565e..850fa12 100644 --- a/tests/unit/benchmark/converters/test_onnx.py +++ b/tests/unit/benchmark/converters/test_onnx.py @@ -11,7 +11,7 @@ if HAS_ONNX: import numpy as np import onnx - from omnimalloc.benchmark.converters.model import ITEMBITS + from omnimalloc.benchmark.converters.model import ITEMBITS, Buffer, Op from omnimalloc.benchmark.converters.onnx import ( _node_to_op, _tensor_proto_to_buffer, @@ -96,98 +96,29 @@ def test_sub_byte_buffers_are_sized_packed(dtype_name: str, count: int) -> None: assert _tensor_proto_to_buffer(packed).size == len(packed.raw_data) -def test_tensor_proto_to_buffer() -> None: - tensor = helper.make_tensor( - "test_tensor", - TensorProto.FLOAT, - [2, 3, 4], - np.zeros([2, 3, 4], dtype=np.float32).tobytes(), - raw=True, - ) - - buffer = _tensor_proto_to_buffer(tensor) - - assert buffer.id == "test_tensor" - assert buffer.shape == (2, 3, 4) - assert buffer.dtype == "float32" - assert buffer.kind == AllocationKind.CONSTANT - - -def test_tensor_proto_to_buffer_different_dtype() -> None: +@pytest.mark.parametrize( + ("data_type", "dtype"), + [(TensorProto.FLOAT, "float32"), (TensorProto.INT64, "int64")], +) +def test_tensor_proto_to_buffer(data_type: int, dtype: str) -> None: tensor = helper.make_tensor( - "int_tensor", - TensorProto.INT64, - [3, 5], - np.ones([3, 5], dtype=np.int64).tobytes(), - raw=True, + "tensor", data_type, [2, 3], np.zeros(6, dtype=dtype).tobytes(), raw=True ) buffer = _tensor_proto_to_buffer(tensor) - assert buffer.id == "int_tensor" - assert buffer.shape == (3, 5) - assert buffer.dtype == "int64" - assert buffer.kind == AllocationKind.CONSTANT - - -def test_value_info_to_buffer() -> None: - value_info = helper.make_tensor_value_info("test_value", TensorProto.INT32, [5, 10]) - - buffer = _value_info_to_buffer(value_info, AllocationKind.WORKSPACE) - - assert buffer.id == "test_value" - assert buffer.shape == (5, 10) - assert buffer.dtype == "int32" - assert buffer.kind == AllocationKind.WORKSPACE - - -def test_value_info_to_buffer_input_kind() -> None: - value_info = helper.make_tensor_value_info("input", TensorProto.FLOAT, [1, 3, 224]) - - buffer = _value_info_to_buffer(value_info, AllocationKind.INPUT) - - assert buffer.kind == AllocationKind.INPUT - - -def test_value_info_to_buffer_output_kind() -> None: - value_info = helper.make_tensor_value_info("output", TensorProto.FLOAT, [1, 1000]) - - buffer = _value_info_to_buffer(value_info, AllocationKind.OUTPUT) + assert buffer == Buffer("tensor", (2, 3), dtype, AllocationKind.CONSTANT) - assert buffer.kind == AllocationKind.OUTPUT +@pytest.mark.parametrize( + "kind", [AllocationKind.WORKSPACE, AllocationKind.INPUT, AllocationKind.OUTPUT] +) +def test_value_info_to_buffer(kind: AllocationKind) -> None: + value_info = helper.make_tensor_value_info("value", TensorProto.INT32, [5, 10]) -def test_value_info_to_buffer_filters_zero_dims() -> None: - value_info = helper.make_tensor_value_info( - "test_value", TensorProto.FLOAT, [3, 0, 5] - ) - - buffer = _value_info_to_buffer(value_info, AllocationKind.WORKSPACE) - - assert buffer.shape == (3, 5) - - -def test_node_to_op(simple_onnx_model: "onnx.ModelProto") -> None: - graph = simple_onnx_model.graph - node = graph.node[0] - - buffers = {} - for init in graph.initializer: - buf = _tensor_proto_to_buffer(init) - buffers[buf.id] = buf - for inp in graph.input: - buf = _value_info_to_buffer(inp, AllocationKind.INPUT) - buffers[buf.id] = buf - for val in graph.value_info: - buf = _value_info_to_buffer(val, AllocationKind.WORKSPACE) - buffers[buf.id] = buf - - op = _node_to_op(node, buffers, node.name) + buffer = _value_info_to_buffer(value_info, kind) - assert op.id == "matmul_node" - assert op.op_type == "MatMul" - assert len(op.inputs) == 2 - assert len(op.outputs) == 1 + assert buffer == Buffer("value", (5, 10), "int32", kind) def test_node_to_op_handles_missing_buffers( @@ -197,9 +128,7 @@ def test_node_to_op_handles_missing_buffers( op = _node_to_op(node, {}, node.name) - assert op.id == "matmul_node" - assert len(op.inputs) == 0 - assert len(op.outputs) == 0 + assert op == Op(id="matmul_node") def test_from_onnx_model_proto(simple_onnx_model: "onnx.ModelProto") -> None: @@ -276,3 +205,26 @@ def test_from_onnx_skips_initializers_relisted_as_inputs( model = from_onnx(simple_onnx_model) assert model.buffers["weights"].kind == AllocationKind.CONSTANT + + +def test_from_onnx_names_the_model_after_its_graph( + simple_onnx_model: "onnx.ModelProto", +) -> None: + simple_onnx_model.doc_string = "A long free-text description." + assert from_onnx(simple_onnx_model).id == "test_model" + + +def test_from_onnx_skips_zero_size_tensors( + simple_onnx_model: "onnx.ModelProto", +) -> None: + empty = helper.make_tensor_value_info("empty", TensorProto.FLOAT, [3, 0]) + simple_onnx_model.graph.value_info.append(empty) + assert "empty" not in from_onnx(simple_onnx_model).buffers + + +def test_from_onnx_takes_symbolic_dims_as_one( + simple_onnx_model: "onnx.ModelProto", +) -> None: + batch = helper.make_tensor_value_info("batch", TensorProto.FLOAT, ["N", 4]) + simple_onnx_model.graph.value_info.append(batch) + assert from_onnx(simple_onnx_model).buffers["batch"].shape == (1, 4) diff --git a/tests/unit/benchmark/sources/test_base.py b/tests/unit/benchmark/sources/test_base.py index b4f5e7d..7226ac4 100644 --- a/tests/unit/benchmark/sources/test_base.py +++ b/tests/unit/benchmark/sources/test_base.py @@ -202,7 +202,7 @@ def test_base_source_label_separates_instances() -> None: def test_base_source_known_optimum_is_unknown_by_default() -> None: - assert RandomSource(num_allocations=10, seed=42).get_known_optimum(10) is None + assert RandomSource().get_known_optimum() is None def test_get_memories_with_skip_returns_requested_count() -> None: diff --git a/tests/unit/benchmark/sources/test_concurrent_tiling.py b/tests/unit/benchmark/sources/test_concurrent_tiling.py index 396376e..5d5c718 100644 --- a/tests/unit/benchmark/sources/test_concurrent_tiling.py +++ b/tests/unit/benchmark/sources/test_concurrent_tiling.py @@ -116,13 +116,7 @@ def test_concurrent_tiling_rejects_nonpositive_threads() -> None: def test_concurrent_tiling_rejects_fewer_allocations_than_threads() -> None: - with pytest.raises(ValueError, match="num_threads"): - ConcurrentTilingSource( - num_allocations=2, num_threads=4, capacity=4096, size_min=1 - ) - source = ConcurrentTilingSource( - num_allocations=8, num_threads=4, capacity=4096, size_min=1 - ) + source = ConcurrentTilingSource(num_threads=4, capacity=4096, size_min=1) with pytest.raises(ValueError, match="num_threads"): source.get_allocations(num_allocations=2) @@ -159,7 +153,5 @@ def test_concurrent_tiling_label_carries_thread_count() -> None: def test_concurrent_tiling_known_optimum_is_the_capacity() -> None: capacity = 1024 * 1024 - source = ConcurrentTilingSource( - num_allocations=32, num_threads=4, capacity=capacity - ) - assert source.get_known_optimum(32) == capacity + source = ConcurrentTilingSource(num_threads=4, capacity=capacity) + assert source.get_known_optimum() == capacity diff --git a/tests/unit/benchmark/sources/test_generator.py b/tests/unit/benchmark/sources/test_generator.py index 0553057..bea7e60 100644 --- a/tests/unit/benchmark/sources/test_generator.py +++ b/tests/unit/benchmark/sources/test_generator.py @@ -97,46 +97,29 @@ def test_random_source_allocation_kinds_with_weights() -> None: assert workspace_count > 60 -def test_random_source_validation_count() -> None: - with pytest.raises(ValueError, match="num_allocations must be positive"): - RandomSource(num_allocations=0) - - -def test_random_source_validation_size_min() -> None: - with pytest.raises(ValueError, match="size_min must be positive"): - RandomSource(size_min=0) - - -def test_random_source_validation_size_max() -> None: - with pytest.raises(ValueError, match="size_max must be >= size_min"): - RandomSource(size_min=100, size_max=50) - - -def test_random_source_validation_time_min() -> None: - with pytest.raises(ValueError, match="time_min must be non-negative"): - RandomSource(time_min=-1) - - -def test_random_source_validation_time_max() -> None: - with pytest.raises(ValueError, match="time_max must be > time_min"): - RandomSource(time_min=100, time_max=100) - - -def test_random_source_validation_duration_min() -> None: - with pytest.raises(ValueError, match="duration_min must be positive"): - RandomSource(duration_min=0) - - -def test_random_source_validation_duration_max() -> None: - with pytest.raises(ValueError, match="duration_max must be >= duration_min"): - RandomSource(duration_min=10, duration_max=5) - - -def test_random_source_validation_kind_weights() -> None: - with pytest.raises( - ValueError, match="kinds and kind_weights must have same length" - ): - RandomSource(kinds=(AllocationKind.WORKSPACE,), kind_weights=(0.5, 0.5)) +@pytest.mark.parametrize( + ("source_cls", "kwargs", "message"), + [ + (RandomSource, {"num_allocations": 0}, "num_allocations must be positive"), + (RandomSource, {"size_min": 0}, "size_min must be positive"), + (RandomSource, {"size_min": 100, "size_max": 50}, "size_max must be >="), + (RandomSource, {"time_min": -1}, "time_min must be non-negative"), + (RandomSource, {"time_min": 100, "time_max": 100}, "time_max must be >"), + (RandomSource, {"duration_min": 0}, "duration_min must be positive"), + (RandomSource, {"duration_min": 10, "duration_max": 5}, "duration_max must"), + ( + RandomSource, + {"kinds": (AllocationKind.WORKSPACE,), "kind_weights": (0.5, 0.5)}, + "same length", + ), + (PowerOf2Source, {"time_max": 10}, "duration_max must be <= time_max"), + ], +) +def test_source_rejects_invalid_parameters( + source_cls: type[BaseSource], kwargs: dict[str, object], message: str +) -> None: + with pytest.raises(ValueError, match=message): + source_cls(**kwargs) def test_uniform_source_fixed_size_and_duration() -> None: diff --git a/tests/unit/benchmark/sources/test_huggingface.py b/tests/unit/benchmark/sources/test_huggingface.py index c6c6e38..a9e8f34 100644 --- a/tests/unit/benchmark/sources/test_huggingface.py +++ b/tests/unit/benchmark/sources/test_huggingface.py @@ -2,11 +2,10 @@ # SPDX-License-Identifier: Apache-2.0 # -import functools -from collections.abc import Callable from pathlib import Path import pytest +from omnimalloc.benchmark.sources import huggingface from omnimalloc.benchmark.sources.huggingface import ( HAS_HUGGINGFACE_HUB, HAS_ONNX, @@ -19,101 +18,75 @@ ) -def skip_when_rate_limited(test: Callable[..., None]) -> Callable[..., None]: - @functools.wraps(test) - def wrapper(*args: object, **kwargs: object) -> None: - # Local import so collection works without huggingface_hub installed - from huggingface_hub.errors import HfHubHTTPError - - try: - test(*args, **kwargs) - except HfHubHTTPError as error: - if error.response.status_code == 429: - pytest.skip("Hugging Face Hub rate limited the request") - raise - - return wrapper - - -def test_huggingface_source_creation() -> None: - """Test basic HuggingfaceSource instantiation.""" - source = HuggingfaceSource(num_models=1) - assert source.num_models == 1 - assert source._model_paths is None # noqa: SLF001 - assert source._model_pools is None # noqa: SLF001 - # HuggingfaceSource is a fixed source, so num_allocations is from base class - assert source.num_allocations == 100 # Default from BaseSource - - -def test_huggingface_source_creation_with_params(artifacts_dir: Path) -> None: - """Test HuggingfaceSource with custom parameters.""" - source = HuggingfaceSource(num_models=2, output_dir=str(artifacts_dir)) - assert source.num_models == 2 - assert source.output_dir == str(artifacts_dir) - - -@skip_when_rate_limited -def test_huggingface_source_get_allocations_single_model(artifacts_dir: Path) -> None: - """Test downloading and extracting allocations from a single model. - - This test downloads from Huggingface and may be slow. - """ - source = HuggingfaceSource(num_models=1, output_dir=str(artifacts_dir)) - allocations = source.get_allocations() - - # Should have allocations from the model - assert len(allocations) > 0 - assert all(hasattr(alloc, "id") for alloc in allocations) - assert all(hasattr(alloc, "size") for alloc in allocations) - assert all(hasattr(alloc, "start") for alloc in allocations) - assert all(hasattr(alloc, "end") for alloc in allocations) - - -@skip_when_rate_limited -def test_huggingface_source_get_allocations_with_count(artifacts_dir: Path) -> None: - """Test getting a limited number of allocations.""" - source = HuggingfaceSource(num_models=1, output_dir=str(artifacts_dir)) - allocations = source.get_allocations(num_allocations=5) - - # Should have at most 5 allocations - assert len(allocations) <= 5 - - -@skip_when_rate_limited -def test_huggingface_source_get_allocations_with_skip(artifacts_dir: Path) -> None: - """Test skipping allocations.""" - source = HuggingfaceSource(num_models=1, output_dir=str(artifacts_dir)) - all_allocations = source.get_allocations() - skipped_allocations = source.get_allocations(skip=2) - - # Should skip first 2 - if len(all_allocations) > 2: - assert len(skipped_allocations) == len(all_allocations) - 2 - assert skipped_allocations[0] == all_allocations[2] - - -@skip_when_rate_limited -def test_huggingface_source_caching(artifacts_dir: Path) -> None: - """Test that models are cached after first download.""" - source = HuggingfaceSource(num_models=1, output_dir=str(artifacts_dir)) - - # First call downloads - allocations1 = source.get_allocations() - assert source._model_paths is not None # noqa: SLF001 - assert source._model_pools is not None # noqa: SLF001 - - # Second call uses cache - allocations2 = source.get_allocations() - assert allocations1 == allocations2 - - -@skip_when_rate_limited -def test_huggingface_source_allocation_kinds(artifacts_dir: Path) -> None: - """Test that allocations have appropriate allocation kinds.""" - source = HuggingfaceSource(num_models=1, output_dir=str(artifacts_dir)) - allocations = source.get_allocations() - - # Should have various allocation kinds - kinds = {alloc.kind for alloc in allocations if alloc.kind is not None} - # At least workspace or input/output/constant should be present - assert len(kinds) > 0 +def _save_tiny_model(path: Path) -> Path: + """Input -> Relu -> hidden -> Relu -> output, one workspace tensor.""" + import onnx + from onnx import TensorProto, helper + + def value(name: str) -> "onnx.ValueInfoProto": + return helper.make_tensor_value_info(name, TensorProto.FLOAT, [1, 8]) + + graph = helper.make_graph( + [ + helper.make_node("Relu", ["input"], ["hidden"], name="relu_0"), + helper.make_node("Relu", ["hidden"], ["output"], name="relu_1"), + ], + path.stem, + [value("input")], + [value("output")], + value_info=[value("hidden")], + ) + onnx.save(helper.make_model(graph), path) + return path + + +@pytest.fixture +def local_models(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + """Serve two tiny local models in place of the Hub download.""" + paths = [_save_tiny_model(tmp_path / f"model_{i}.onnx") for i in range(2)] + monkeypatch.setattr( + huggingface, "_download_onnx_models", lambda num_models, _: paths[:num_models] + ) + + +@pytest.mark.usefixtures("local_models") +def test_huggingface_source_downloads_lazily() -> None: + source = HuggingfaceSource(num_models=2) + assert "pools" not in vars(source) + assert source.get_available_variants() == ("model_0", "model_1") + + +@pytest.mark.usefixtures("local_models") +def test_huggingface_source_qualifies_ids_across_models() -> None: + allocations = HuggingfaceSource(num_models=2).get_allocations() + assert [a.id for a in allocations] == ["model_0_hidden", "model_1_hidden"] + + +@pytest.mark.usefixtures("local_models") +def test_huggingface_source_get_variant_by_name_and_index() -> None: + source = HuggingfaceSource(num_models=2) + assert source.get_variant("model_1") is source.get_variant(1) + with pytest.raises(ValueError, match="not found"): + source.get_variant("model_2") + + +@pytest.mark.usefixtures("local_models") +def test_huggingface_source_get_allocations_slices() -> None: + source = HuggingfaceSource(num_models=2) + assert source.get_allocations(num_allocations=1, skip=1) == ( + source.get_allocations()[1], + ) + + +def test_huggingface_source_downloads_a_real_model(tmp_path: Path) -> None: + from huggingface_hub.errors import HfHubHTTPError + + source = HuggingfaceSource(num_models=1, output_dir=tmp_path) + try: + pools = source.pools + except HfHubHTTPError as error: + if error.response.status_code == 429: + pytest.skip("Hugging Face Hub rate limited the request") + raise + assert len(pools) == 1 + assert pools[0].allocations diff --git a/tests/unit/benchmark/sources/test_minimalloc.py b/tests/unit/benchmark/sources/test_minimalloc.py index 9d86cfe..9fc8779 100644 --- a/tests/unit/benchmark/sources/test_minimalloc.py +++ b/tests/unit/benchmark/sources/test_minimalloc.py @@ -31,7 +31,7 @@ def test_minimalloc_source_default_subset_is_challenging() -> None: reason="dataset discovery only resolves from a source checkout", ) def test_minimalloc_source_discovers_the_checkout_datasets() -> None: - assert MinimallocSource().num_allocations > 0 + assert MinimallocSource().get_available_variants() def test_checkout_csv_dir_stops_at_the_project_root( @@ -49,43 +49,16 @@ def test_checkout_csv_dir_stops_at_the_project_root( assert minimalloc._checkout_csv_dir(MinimallocSubset.SMALL) == checkout # noqa: SLF001 -def test_minimalloc_source_accepts_enum_member() -> None: - source = MinimallocSource(MinimallocSubset.SMALL) - assert source.subset is MinimallocSubset.SMALL +@pytest.mark.parametrize("subset", [MinimallocSubset.SMALL, "small"]) +def test_minimalloc_source_accepts_enum_member_or_string(subset: str) -> None: + assert MinimallocSource(subset).subset is MinimallocSubset.SMALL -def test_minimalloc_source_accepts_string_alias() -> None: - source = MinimallocSource("small") - assert source.subset is MinimallocSubset.SMALL - assert source.subset == "small" - - -def test_minimalloc_source_examples_subset() -> None: - source = _source("examples") - assert source.subset == "examples" - variants = source.get_available_variants() - assert len(variants) == 1 # Only one example pool - - -def test_minimalloc_source_small_subset() -> None: - source = _source("small") - assert source.subset == "small" - variants = source.get_available_variants() - assert len(variants) > 0 - assert all(v[0].islower() for v in variants) - - -def test_minimalloc_source_challenging_subset() -> None: - source = _source("challenging") - variants = source.get_available_variants() - assert len(variants) > 0 - - -def test_minimalloc_source_subsets_are_disjoint() -> None: +def test_minimalloc_source_subsets_are_non_empty_and_disjoint() -> None: examples = set(_source("examples").get_available_variants()) small = set(_source("small").get_available_variants()) challenging = set(_source("challenging").get_available_variants()) - assert examples + assert len(examples) == 1 assert small assert challenging assert examples.isdisjoint(small) @@ -122,17 +95,11 @@ def test_minimalloc_source_get_allocation_keeps_kind_none() -> None: assert allocation.kind is None -def test_minimalloc_source_get_variant_by_id() -> None: - source = _source("small") - variants = source.get_available_variants() - pool = source.get_variant(variants[0]) - assert pool.id == variants[0] - - -def test_minimalloc_source_get_variant_by_index() -> None: +def test_minimalloc_source_get_variant_by_id_and_index() -> None: source = _source("small") - pool = source.get_variant(0) - assert pool.id in source.get_available_variants() + first = source.get_available_variants()[0] + assert source.get_variant(first).id == first + assert source.get_variant(0).id == first def test_minimalloc_source_get_variant_unknown_id() -> None: diff --git a/tests/unit/benchmark/sources/test_sync_patterns.py b/tests/unit/benchmark/sources/test_sync_patterns.py index b635dcf..c0bc39a 100644 --- a/tests/unit/benchmark/sources/test_sync_patterns.py +++ b/tests/unit/benchmark/sources/test_sync_patterns.py @@ -193,3 +193,11 @@ def test_every_step_barrier_linearizes_while_a_sparser_one_does_not() -> None: ).get_allocations() assert try_linearize(tight, work_budget=None) is not None assert try_linearize(loose, work_budget=None) is None + + +def test_subset_pattern_couples_threads_at_three_threads() -> None: + def signatures(pattern: str) -> list[tuple[object, object, int]]: + source = SyncPatternSource(num_allocations=32, num_threads=3, pattern=pattern) + return _signatures(source.get_allocations()) + + assert signatures("subset") != signatures("independent") diff --git a/tests/unit/benchmark/sources/test_tiling.py b/tests/unit/benchmark/sources/test_tiling.py index f1c0ca9..33703c7 100644 --- a/tests/unit/benchmark/sources/test_tiling.py +++ b/tests/unit/benchmark/sources/test_tiling.py @@ -119,15 +119,9 @@ def test_tiling_ground_truth_requires_seed() -> None: TilingSource(seed=None).get_ground_truth_pool() -def test_tiling_known_optimum_is_the_capacity() -> None: - capacity = 1024 * 1024 - source = TilingSource(num_allocations=64, capacity=capacity) - assert source.get_known_optimum() == capacity - assert source.get_known_optimum(128) == capacity - - -def test_tiling_known_optimum_unknown_without_seed() -> None: - assert TilingSource(num_allocations=64, seed=None).get_known_optimum() is None +@pytest.mark.parametrize("seed", [0, None]) +def test_tiling_known_optimum_is_the_capacity(seed: int | None) -> None: + assert TilingSource(capacity=4096, seed=seed).get_known_optimum() == 4096 def test_tiling_variant_sweep_builds_ladder() -> None: From 95fed6b5392498996fb3661a592d2af79131065e Mon Sep 17 00:00:00 2001 From: Fabian Peddinghaus Date: Tue, 29 Sep 2026 14:42:04 +0000 Subject: [PATCH 2/4] Flatten run_benchmark and the results package --- scripts/benchmark_allocation.py | 10 +- scripts/benchmark_pressure.py | 10 +- scripts/stress_omni.py | 9 +- src/python/omnimalloc/benchmark/benchmark.py | 336 +++++++---------- .../omnimalloc/benchmark/results/campaign.py | 70 ---- .../omnimalloc/benchmark/results/export.py | 132 ++----- .../omnimalloc/benchmark/results/report.py | 22 +- .../omnimalloc/benchmark/results/visualize.py | 95 ++--- src/python/omnimalloc/benchmark/timer.py | 61 +--- src/python/omnimalloc/benchmark/utils.py | 17 +- tests/unit/benchmark/results/test_campaign.py | 158 ++------ tests/unit/benchmark/results/test_export.py | 319 +++++------------ tests/unit/benchmark/results/test_report.py | 191 +++------- .../unit/benchmark/results/test_visualize.py | 20 +- tests/unit/benchmark/test_benchmark.py | 338 +++++------------- tests/unit/benchmark/test_timer.py | 145 +------- 16 files changed, 468 insertions(+), 1465 deletions(-) diff --git a/scripts/benchmark_allocation.py b/scripts/benchmark_allocation.py index f92ccdf..2c78721 100644 --- a/scripts/benchmark_allocation.py +++ b/scripts/benchmark_allocation.py @@ -54,15 +54,13 @@ def _timed( allocator: BaseAllocator, allocations: "tuple[Allocation, ...]" ) -> "tuple[float, tuple[Allocation, ...]]": - timer = Timer().start() - placed = allocator.allocate(allocations) - timer.stop() + with Timer() as timer: + placed = allocator.allocate(allocations) seconds = timer.elapsed_s if seconds < 1e-3: for _ in range(4): - timer = Timer().start() - allocator.allocate(allocations) - timer.stop() + with Timer() as timer: + allocator.allocate(allocations) seconds = min(seconds, timer.elapsed_s) return seconds, placed diff --git a/scripts/benchmark_pressure.py b/scripts/benchmark_pressure.py index ec1baac..9ccb81b 100644 --- a/scripts/benchmark_pressure.py +++ b/scripts/benchmark_pressure.py @@ -122,15 +122,13 @@ def _sample_runners( def _timed(runner: "Runner") -> "tuple[float, Value]": - timer = Timer().start() - value = runner() - timer.stop() + with Timer() as timer: + value = runner() seconds = timer.elapsed_s if seconds < 1e-3: for _ in range(4): - timer = Timer().start() - runner() - timer.stop() + with Timer() as timer: + runner() seconds = min(seconds, timer.elapsed_s) return seconds, value diff --git a/scripts/stress_omni.py b/scripts/stress_omni.py index 2b1c61a..f72f1cc 100644 --- a/scripts/stress_omni.py +++ b/scripts/stress_omni.py @@ -52,9 +52,8 @@ def _timed_allocate( seconds = [] placed: tuple[Allocation, ...] = () for _ in range(repeats): - timer = Timer().start() - placed = OmniAllocator().allocate(allocations) - timer.stop() + with Timer() as timer: + placed = OmniAllocator().allocate(allocations) seconds.append(timer.elapsed_s) return min(seconds), placed @@ -111,10 +110,8 @@ def run_once(_: int) -> list[int | None]: throughput: dict[int, float] = {} for callers in args.callers: - with ThreadPoolExecutor(max_workers=callers) as executor: - timer = Timer().start() + with ThreadPoolExecutor(max_workers=callers) as executor, Timer() as timer: offsets = list(executor.map(run_once, range(args.calls))) - timer.stop() if any(result != expected for result in offsets): raise AssertionError(f"non-deterministic placement with {callers} callers") throughput[callers] = args.calls / timer.elapsed_s diff --git a/src/python/omnimalloc/benchmark/benchmark.py b/src/python/omnimalloc/benchmark/benchmark.py index 7039b89..d517200 100644 --- a/src/python/omnimalloc/benchmark/benchmark.py +++ b/src/python/omnimalloc/benchmark/benchmark.py @@ -3,7 +3,6 @@ # import logging -from dataclasses import asdict, dataclass from omnimalloc import allocate, validate_allocation from omnimalloc.allocators import BaseAllocator, available_allocators @@ -11,184 +10,135 @@ from omnimalloc.primitives import IdType, Pool from .results import BenchmarkCampaign, BenchmarkReport, BenchmarkResult -from .results.utils import get_date_time_snake_case +from .results.utils import get_date_time_snake_case, get_environment_metadata from .sources import DEFAULT_SOURCE, BaseSource from .timer import Timer from .utils import tqdm logger = logging.getLogger(__name__) +VariantSpec = int | tuple[IdType, ...] | None -@dataclass(frozen=True) -class SkippedAllocator: - """An allocator left out of a campaign, with the reason it was skipped.""" - - source: str - allocator: str - reason: str - - -@dataclass(frozen=True) -class SkippedVariant: - """A workload a source could not express, with the reason it was skipped.""" - source: str - variant: str - reason: str +def _ensure_known_variant_keys( + sources: tuple[BaseSource, ...], + variants: VariantSpec | dict[str, VariantSpec], +) -> None: + if not isinstance(variants, dict): + return + known = {s.label() for s in sources} | {s.name() for s in sources} + unknown = sorted(set(variants) - known) + if unknown: + raise ValueError(f"Variants keys {unknown} match no source in this campaign") -def _resolve_parameterizable_variants( - source: BaseSource, variants: int | tuple[IdType, ...] | None -) -> tuple[int, ...]: +def _parameterizable_variants( + source: BaseSource, variants: VariantSpec +) -> tuple[IdType, ...]: if variants is None: return (source.num_allocations,) if isinstance(variants, int): return (variants,) - resolved_variants = [] for v in variants: if not isinstance(v, int): raise TypeError( f"Non-integer variant {v!r} for parameterizable source {source.name()}" ) - resolved_variants.append(v) - return tuple(resolved_variants) + return variants -def _resolve_fixed_variants( - source: BaseSource, variants: int | tuple[IdType, ...] | None -) -> tuple[str, ...]: - available = source.get_available_variants() - if available is None: - return () +def _fixed_variants(source: BaseSource, variants: VariantSpec) -> tuple[str, ...]: + available = source.get_available_variants() or () if variants is None: return available if isinstance(variants, int): return available[:variants] - resolved_variants = [] + resolved = [] for v in variants: if isinstance(v, str) and v in available: - resolved_variants.append(v) + resolved.append(v) # Int variants index into the available variants elif isinstance(v, int) and 0 <= v < len(available): - resolved_variants.append(available[v]) + resolved.append(available[v]) else: raise ValueError(f"Unknown variant {v!r} for source {source.name()}") - return tuple(resolved_variants) - - -VariantSpec = int | tuple[IdType, ...] | None - - -def _ensure_known_variant_keys( - sources: tuple[BaseSource, ...], - variants: VariantSpec | dict[str, VariantSpec], -) -> None: - if not isinstance(variants, dict): - return - known = {s.label() for s in sources} | {s.name() for s in sources} - unknown = sorted(set(variants) - known) - if unknown: - raise ValueError(f"Variants keys {unknown} match no source in this campaign") + return tuple(resolved) def _get_variant_ids( - source_inst: BaseSource, + source: BaseSource, variants: VariantSpec | dict[str, VariantSpec], ) -> tuple[IdType, ...]: if isinstance(variants, dict): # Labelled instances can be addressed individually; the class name # keeps working and covers every instance of that source - label = source_inst.label() - variants = ( - variants[label] if label in variants else variants.get(source_inst.name()) - ) - if source_inst.is_parameterizable(): - return _resolve_parameterizable_variants(source_inst, variants) - return _resolve_fixed_variants(source_inst, variants) - - -def _benchmark_result( - allocator: BaseAllocator, - source: BaseSource, - pool: Pool, - result_id: IdType, - validate: bool, -) -> BenchmarkResult: - with Timer() as timer: - allocated_pool = allocate(pool, allocator, validate=False) - - if validate: - validate_allocation(allocated_pool) - - return BenchmarkResult( - id=result_id, - allocator=allocator, - source=source, - entity=allocated_pool, - duration=timer.elapsed_s, - ) + label = source.label() + variants = variants[label] if label in variants else variants.get(source.name()) + if source.is_parameterizable(): + return _parameterizable_variants(source, variants) + return _fixed_variants(source, variants) + + +def _resolve_allocators( + allocators: tuple[BaseAllocator | type[BaseAllocator] | str, ...], + skipped: list[dict[str, str]], +) -> list[BaseAllocator]: + # An allocator wrapping an uninstalled library is a skip, not an abort: + # `available_allocators()` lists every registered name, so the default + # campaign would otherwise die on the first optional one + resolved = [] + for allocator in allocators: + try: + resolved.append(BaseAllocator.resolve(allocator)) + except ImportError as error: + name = allocator if isinstance(allocator, str) else allocator.name() + _skip(skipped, str(error).splitlines()[0], allocator=name) + return resolved + + +def _skip(skipped: list[dict[str, str]], reason: str, **where: str) -> None: + """Record a combination left out, so a shrunken comparison is visible.""" + logger.warning(f"Skipping {'/'.join(where.values())}: {reason}") + skipped.append(where | {"reason": reason}) def _benchmark_report( + report_id: int, allocator: BaseAllocator, source: BaseSource, - iterations: int, variant_id: IdType, - report_id: int, - result_id: int, + pool: Pool, + iterations: int, validate: bool, - known_optima: dict[IdType, int | None], -) -> BenchmarkReport | SkippedAllocator | SkippedVariant: - """Time one allocator/source/variant, or report why it was skipped.""" - variant_desc = variant_id if isinstance(variant_id, str) else f"{variant_id} allocs" - - # Validate and error out early; a variant the source cannot express - # (e.g. fewer allocations than threads) skips instead of aborting the - # whole campaign - try: - pool = source.get_variant(variant_id) - except ValueError as error: - logger.warning(f"Skipping {source.label()}[{variant_desc}]: {error}") - return SkippedVariant( - source=source.label(), variant=str(variant_id), reason=str(error) - ) - if pool is None: - raise ValueError(f"source {source.name()} returned no pool") - try: - allocator.ensure_supported(pool.allocations) - except ValueError as error: - logger.warning( - f"Skipping {allocator.name()} on {source.label()}[{variant_desc}]: {error}" - ) - return SkippedAllocator( - source=source.label(), - allocator=allocator.name(), - reason=str(error), - ) - +) -> BenchmarkReport: results = [] - for _ in tqdm( + for i in tqdm( range(iterations), - desc=f"Iterations [{variant_desc}]", + desc=f"Iterations [{allocator.name()}]", position=3, leave=False, ): - result = _benchmark_result(allocator, source, pool, result_id, validate) - results.append(result) - result_id += 1 - - # The ground truth is a property of the instance, not the allocator, and - # the tiling sources rebuild their whole construction to read it - if variant_id not in known_optima: - known_optima[variant_id] = source.get_known_optimum() - + # Validation runs outside the timer: it is quadratic and would skew timings + with Timer() as timer: + allocated_pool = allocate(pool, allocator, validate=False) + if validate: + validate_allocation(allocated_pool) + results.append( + BenchmarkResult( + id=i, + allocator=allocator, + source=source, + entity=allocated_pool, + duration=timer.elapsed_s, + ) + ) return BenchmarkReport( id=report_id, results=tuple(results), allocator=allocator, source=source, variant_id=variant_id, - known_optimum=known_optima[variant_id], + known_optimum=source.get_known_optimum(), ) @@ -207,85 +157,64 @@ def run_benchmark( Unlike `allocate`, `validate` defaults to True here. """ ensure_positive(iterations, "iterations") - allocators = allocators or available_allocators() - sources = sources or (DEFAULT_SOURCE,) - source_insts = tuple(BaseSource.resolve(source) for source in sources) + source_insts = tuple(BaseSource.resolve(s) for s in sources or (DEFAULT_SOURCE,)) _ensure_known_variant_keys(source_insts, variants) - campaign_id = campaign_id or "campaign_" + get_date_time_snake_case() - - reports = [] - skipped: list[SkippedAllocator] = [] - skipped_variants: list[SkippedVariant] = [] - report_id = 0 - result_id = 0 - - timer = Timer() - timer.start() - - for source_inst in tqdm( - source_insts, - desc="Sources", - position=0, - leave=False, - ): - if getattr(source_inst, "seed", 0) is None: - logger.warning( - f"Source {source_inst.name()} has seed=None; each allocator " - f"gets a different random problem, so results are not comparable" - ) + if campaign_id is None: + campaign_id = "campaign_" + get_date_time_snake_case() - variant_ids = _get_variant_ids(source_inst, variants) - known_optima: dict[IdType, int | None] = {} + skipped: list[dict[str, str]] = [] + allocator_insts = _resolve_allocators(allocators or available_allocators(), skipped) + reports: list[BenchmarkReport] = [] - for allocator in tqdm( - allocators, - desc=f"Allocators [{source_inst.label()}]", - position=1, - leave=False, - ): - # An allocator wrapping an uninstalled library is a skip, not an - # abort: `available_allocators()` lists every registered name, so - # the default campaign would otherwise die on the first optional one - try: - allocator_inst = BaseAllocator.resolve(allocator) - except ImportError as error: - reason = str(error).splitlines()[0] - name = allocator if isinstance(allocator, str) else allocator.name() - logger.warning(f"Skipping {name} on {source_inst.label()}: {reason}") - skipped.append( - SkippedAllocator( - source=source_inst.label(), allocator=name, reason=reason - ) + with Timer() as timer: + for source in tqdm(source_insts, desc="Sources", position=0, leave=False): + label = source.label() + if getattr(source, "seed", 0) is None: + logger.warning( + f"Source {label} has seed=None; each allocator gets a " + f"different random problem, so results are not comparable" ) - continue - for variant_id in tqdm( - variant_ids, - desc=f"Variants [{allocator}]", - position=2, + _get_variant_ids(source, variants), + desc=f"Variants [{label}]", + position=1, leave=False, ): - report = _benchmark_report( - allocator_inst, - source_inst, - iterations, - variant_id, - report_id, - result_id, - validate, - known_optima, - ) - if isinstance(report, SkippedAllocator): - skipped.append(report) + # A variant the source cannot express (e.g. fewer allocations + # than threads) skips instead of aborting the whole campaign + try: + pool = source.get_variant(variant_id) + except ValueError as error: + _skip(skipped, str(error), source=label, variant=str(variant_id)) continue - if isinstance(report, SkippedVariant): - skipped_variants.append(report) - continue - reports.append(report) - report_id += 1 - result_id += iterations - - timer.stop() + for allocator in tqdm( + allocator_insts, + desc=f"Allocators [{variant_id}]", + position=2, + leave=False, + ): + try: + allocator.ensure_supported(pool.allocations) + except ValueError as error: + _skip( + skipped, + str(error), + source=label, + variant=str(variant_id), + allocator=allocator.name(), + ) + continue + reports.append( + _benchmark_report( + len(reports), + allocator, + source, + variant_id, + pool, + iterations, + validate, + ) + ) if not reports: raise ValueError( @@ -293,17 +222,10 @@ def run_benchmark( "combination was skipped or empty. Double-check your setup." ) - campaign = BenchmarkCampaign( - id=campaign_id, - reports=tuple(reports), - metadata={ - "total_duration": timer.elapsed, - # Same allocator and source repeat once per variant, and a dropped - # variant repeats once per allocator; report the distinct omissions - # so a shrunken comparison is visible - "skipped_allocators": [asdict(s) for s in dict.fromkeys(skipped)], - "skipped_variants": [asdict(s) for s in dict.fromkeys(skipped_variants)], - }, - ) - campaign = campaign.finalize_metadata() - return campaign + metadata = get_environment_metadata() | { + "total_duration": f"{timer.elapsed_s:.2f} s", + "num_reports": len(reports), + "num_results": len(reports) * iterations, + "skipped": skipped, + } + return BenchmarkCampaign(id=campaign_id, reports=tuple(reports), metadata=metadata) diff --git a/src/python/omnimalloc/benchmark/results/campaign.py b/src/python/omnimalloc/benchmark/results/campaign.py index 48dc7e0..235698b 100644 --- a/src/python/omnimalloc/benchmark/results/campaign.py +++ b/src/python/omnimalloc/benchmark/results/campaign.py @@ -2,8 +2,6 @@ # SPDX-License-Identifier: Apache-2.0 # -import copy -from collections import defaultdict from dataclasses import dataclass, field from typing import Any @@ -11,7 +9,6 @@ from omnimalloc.primitives.utils import ensure_unique_ids from .report import BenchmarkReport -from .utils import get_environment_metadata @dataclass(frozen=True) @@ -26,12 +23,6 @@ def __post_init__(self) -> None: if not self.reports: raise ValueError("BenchmarkCampaign must contain at least one report") ensure_unique_ids(self.reports, "report") - try: - object.__setattr__(self, "metadata", copy.deepcopy(self.metadata)) - except TypeError as error: - raise TypeError( - f"Campaign metadata must be deep-copyable: {error}" - ) from error @property def num_reports(self) -> int: @@ -41,22 +32,6 @@ def num_reports(self) -> int: def num_results(self) -> int: return sum(r.num_results for r in self.reports) - @property - def num_results_per_report(self) -> float: - return self.num_results / self.num_reports - - @property - def num_allocations(self) -> int: - return sum(r.total_num_allocations for r in self.reports) - - @property - def num_allocations_per_report(self) -> float: - return self.num_allocations / self.num_reports - - @property - def num_allocations_per_result(self) -> float: - return self.num_allocations / self.num_results - @property def num_allocators(self) -> int: return len(self.allocator_names) @@ -72,48 +47,3 @@ def allocator_names(self) -> tuple[str, ...]: @property def source_names(self) -> tuple[str, ...]: return tuple(sorted({r.source_name for r in self.reports})) - - @property - def reports_by_source_allocator_variant( - self, - ) -> dict[str, dict[str, dict[str, tuple[BenchmarkReport, ...]]]]: - """Group reports by source -> allocator -> variant_label.""" - - result: dict[str, dict[str, dict[str, list[BenchmarkReport]]]] = defaultdict( - lambda: defaultdict(lambda: defaultdict(list)) - ) - - for report in self.reports: - source = report.source_name - allocator = report.allocator_name - variant = report.variant_label - result[source][allocator][variant].append(report) - - return { - source: { - alloc: { - variant: tuple(reports) - for variant, reports in sorted(variants.items()) - } - for alloc, variants in sorted(allocs.items()) - } - for source, allocs in sorted(result.items()) - } - - @property - def default_metadata(self) -> dict[str, Any]: - """Environment plus campaign shape; `metadata` wins on a key clash.""" - return get_environment_metadata() | { - "num_reports": self.num_reports, - "num_results_per_report": round(self.num_results_per_report, 2), - "num_results_total": self.num_results, - "num_allocations_per_report": round(self.num_allocations_per_report, 2), - "num_allocations_per_result": round(self.num_allocations_per_result, 2), - "num_allocations_total": self.num_allocations, - "num_allocators": self.num_allocators, - "num_sources": self.num_sources, - } - - def finalize_metadata(self) -> "BenchmarkCampaign": - m = self.default_metadata | self.metadata - return BenchmarkCampaign(id=self.id, reports=self.reports, metadata=m) diff --git a/src/python/omnimalloc/benchmark/results/export.py b/src/python/omnimalloc/benchmark/results/export.py index afd0801..69acf05 100644 --- a/src/python/omnimalloc/benchmark/results/export.py +++ b/src/python/omnimalloc/benchmark/results/export.py @@ -5,10 +5,11 @@ import csv import json import logging +import re import shutil import tempfile from pathlib import Path -from typing import Any, Final, Literal, Protocol +from typing import Any, Final, Literal from ..utils import tqdm # noqa: TID252 from .campaign import BenchmarkCampaign @@ -37,15 +38,9 @@ ) -class ProgressBar(Protocol): - """Update-only view of a tqdm(-like) progress bar.""" - - def update(self, n: int = 1) -> None: ... - - -def _prepare_base_dir(output_path: Path, output_format: str, overwrite: bool) -> Path: +def _prepare_base_dir(output_path: Path, output_format: str) -> Path: if output_format == "dir": - output_path.mkdir(parents=True, exist_ok=overwrite) + output_path.mkdir(parents=True) return output_path base_dir = Path(tempfile.mkdtemp(prefix="omnimalloc_dump_")) / output_path.stem base_dir.mkdir() @@ -87,8 +82,6 @@ def _write_results_csv(base_dir: Path, campaign: BenchmarkCampaign) -> None: def _create_zip_archive(base_dir: Path, final_path: Path) -> Path: - if final_path.exists(): - final_path.unlink() return Path( shutil.make_archive( str(final_path.with_suffix("")), @@ -99,99 +92,25 @@ def _create_zip_archive(base_dir: Path, final_path: Path) -> Path: ) -def _write_iterations( - report_dir: Path, - report: BenchmarkReport, - pbar: ProgressBar, -) -> None: - iterations_dir = report_dir / "iterations" - iterations_dir.mkdir(exist_ok=True) - - for i, result in enumerate(report.results): - iteration_file = iterations_dir / f"iteration_{i}.pdf" - result.visualize(iteration_file) - pbar.update(1) - - -def _write_allocator_reports( - source_dir: Path, - allocator_name: str, - variant_dict: dict[str, tuple[BenchmarkReport, ...]], - visualize_iterations: bool, - pbar: ProgressBar, -) -> None: - allocator_dir = source_dir / allocator_name - allocator_dir.mkdir(parents=True, exist_ok=True) - - for variant_label in sorted(variant_dict.keys()): - reports = variant_dict[variant_label] - variant_dir = allocator_dir / variant_label - variant_dir.mkdir(parents=True, exist_ok=True) - - for report_idx, report in enumerate(reports): - report_dir = ( - variant_dir / f"report_{report_idx}" - if len(reports) > 1 - else variant_dir - ) - report_dir.mkdir(parents=True, exist_ok=True) - - if visualize_iterations: - _write_iterations(report_dir, report, pbar) - else: - pbar.update(1) - - -def _write_source_reports( - base_dir: Path, - source_name: str, - allocator_dict: dict[str, dict[str, tuple[BenchmarkReport, ...]]], - visualize_iterations: bool, - pbar: ProgressBar, -) -> None: - source_dir = base_dir / "sources" / source_name / "allocators" - source_dir.mkdir(parents=True, exist_ok=True) - - for allocator_name in sorted(allocator_dict.keys()): - _write_allocator_reports( - source_dir, - allocator_name, - allocator_dict[allocator_name], - visualize_iterations, - pbar, - ) - - -def _write_nested_reports( - base_dir: Path, campaign: BenchmarkCampaign, visualize_iterations: bool -) -> None: - reports_by_source = campaign.reports_by_source_allocator_variant - - total_iterations = ( - sum(report.num_results for report in campaign.reports) - if visualize_iterations - else len(campaign.reports) - ) - - unit = "iteration" if visualize_iterations else "report" - - with tqdm( - total=total_iterations, - desc="Saving campaign", - unit=unit, - leave=False, - ) as pbar: - for source_name in sorted(reports_by_source.keys()): - _write_source_reports( - base_dir, - source_name, - reports_by_source[source_name], - visualize_iterations, - pbar, - ) +def _dir_name(label: str) -> str: + """A label as one path component; labels may carry an explicit path.""" + return re.sub(r"[/\\]", "_", label) -# TODO(fpedd): Optionally timestamp the campaign name so saves cannot collide +def _write_iterations(base_dir: Path, campaign: BenchmarkCampaign) -> None: + """Plot every iteration under sources//allocators//.""" + for report in tqdm(campaign.reports, desc="Saving iterations", leave=False): + iterations_dir = ( + base_dir + / "sources" + / _dir_name(report.source_name) + / "allocators" + / _dir_name(report.allocator_name) + / _dir_name(report.variant_label) + / "iterations" + ) + for i, result in enumerate(report.results): + result.visualize(iterations_dir / f"iteration_{i}.pdf") def save_benchmark( @@ -201,13 +120,13 @@ def save_benchmark( visualize_iterations: bool = True, overwrite: bool = True, ) -> Path: - """Save a campaign, defaulting to `artifacts/campaign_` under the cwd.""" + """Save a campaign, defaulting to `artifacts/` under the cwd.""" if not isinstance(campaign, BenchmarkCampaign): raise TypeError(f"Expected a BenchmarkCampaign, got {type(campaign)!r}") if output_path is None: - output_path = Path.cwd() / "artifacts" / f"campaign_{campaign.id}" + output_path = Path.cwd() / "artifacts" / str(campaign.id) output_path = Path(output_path) @@ -226,13 +145,14 @@ def save_benchmark( else: final_path.unlink() - base_dir = _prepare_base_dir(output_path, output_format, overwrite) + base_dir = _prepare_base_dir(output_path, output_format) try: _write_metadata(base_dir, campaign) _write_results_csv(base_dir, campaign) plot_benchmark(campaign, base_dir / "campaign_overview.pdf") - _write_nested_reports(base_dir, campaign, visualize_iterations) + if visualize_iterations: + _write_iterations(base_dir, campaign) if output_format == "zip": final_path = _create_zip_archive(base_dir, final_path) diff --git a/src/python/omnimalloc/benchmark/results/report.py b/src/python/omnimalloc/benchmark/results/report.py index b94f290..1134902 100644 --- a/src/python/omnimalloc/benchmark/results/report.py +++ b/src/python/omnimalloc/benchmark/results/report.py @@ -63,12 +63,8 @@ def source_name(self) -> str: @property def variant_label(self) -> str: - """Human-readable label for this variant.""" - if self.variant_id is None: - return f"{self.num_allocations}" - if isinstance(self.variant_id, str): - return self.variant_id - return f"{self.variant_id}" + """The variant id, or the allocation count for a report without one.""" + return str(self.num_allocations if self.variant_id is None else self.variant_id) @property def is_categorical(self) -> bool: @@ -79,10 +75,6 @@ def is_categorical(self) -> bool: def num_allocations(self) -> int: return self.results[0].num_allocations - @property - def total_num_allocations(self) -> int: - return sum(r.num_allocations for r in self.results) - @property def num_results(self) -> int: return len(self.results) @@ -135,13 +127,3 @@ def optimum_ratio(self) -> float | None: if not self.known_optimum: return None return self.mean_peak_size / self.known_optimum - - def with_results(self, results: tuple[BenchmarkResult, ...]) -> "BenchmarkReport": - return BenchmarkReport( - id=self.id, - allocator=self.allocator, - source=self.source, - variant_id=self.variant_id, - known_optimum=self.known_optimum, - results=self.results + results, - ) diff --git a/src/python/omnimalloc/benchmark/results/visualize.py b/src/python/omnimalloc/benchmark/results/visualize.py index e0c4bed..2846718 100644 --- a/src/python/omnimalloc/benchmark/results/visualize.py +++ b/src/python/omnimalloc/benchmark/results/visualize.py @@ -4,7 +4,7 @@ import logging from pathlib import Path -from typing import Any, Final +from typing import TYPE_CHECKING, Any, Final, cast from omnimalloc.common.optional import require_optional @@ -14,24 +14,16 @@ try: import matplotlib.pyplot as plt - from matplotlib.axes import Axes - from matplotlib.figure import Figure from matplotlib.lines import Line2D HAS_MATPLOTLIB = True except ImportError: - from types import SimpleNamespace - HAS_MATPLOTLIB = False - plt = SimpleNamespace( # ty: ignore[invalid-assignment] - subplots=None, - savefig=None, - show=None, - close=None, - ) - Line2D = None # ty: ignore[invalid-assignment] - Axes = None # ty: ignore[invalid-assignment] - Figure = None # ty: ignore[invalid-assignment] + plt = Line2D = cast("Any", None) + +if TYPE_CHECKING: + from matplotlib.axes import Axes + from matplotlib.figure import Figure logger = logging.getLogger(__name__) @@ -50,30 +42,17 @@ def _format_metadata(metadata: dict[str, Any] | None) -> str: ) -def _is_categorical(data: dict[str, dict[str, tuple[BenchmarkReport, ...]]]) -> bool: - return any( - data[name][alloc_name][0].is_categorical - for name in data - for alloc_name in data[name] - ) - - -def _get_sorted_reports( - allocator_data: dict[str, tuple[BenchmarkReport, ...]], -) -> list[BenchmarkReport]: - reports = [r for rs in allocator_data.values() for r in rs] +def _sorted_reports(reports: list[BenchmarkReport]) -> list[BenchmarkReport]: # One key type for the whole group: mixing a categorical variant_id with a # numeric fallback compares str against int and raises mid-sort if any(r.is_categorical for r in reports): - reports.sort(key=lambda r: str(r.variant_id)) - else: - reports.sort(key=lambda r: r.num_allocations) - return reports + return sorted(reports, key=lambda r: str(r.variant_id)) + return sorted(reports, key=lambda r: r.num_allocations) def _draw_graphs( - ax: Axes, - ax2: Axes, + ax: "Axes", + ax2: "Axes", name: str, color: str, is_categorical: bool, @@ -160,35 +139,31 @@ def _draw_graphs( def _draw_subplot( - ax: Axes, + ax: "Axes", source_name: str, - source_data: dict[str, dict[str, tuple[BenchmarkReport, ...]]], + reports: list[BenchmarkReport], allocator_names: tuple[str, ...], ) -> None: ax2 = ax.twinx() - - is_categorical = _is_categorical(source_data) - - for allocator_name, allocator_data in source_data.items(): - # Color by campaign-wide allocator index so colors match the legend - # even when a source lacks some allocators. - color = _get_allocator_color(allocator_names.index(allocator_name)) - reports = _get_sorted_reports(allocator_data) - - _draw_graphs(ax, ax2, allocator_name, color, is_categorical, reports) + is_categorical = any(r.is_categorical for r in reports) + + # Color by campaign-wide allocator index so colors match the legend + # even when a source lacks some allocators. + for index, allocator_name in enumerate(allocator_names): + series = [r for r in reports if r.allocator_name == allocator_name] + if series: + color = _get_allocator_color(index) + _draw_graphs( + ax, ax2, allocator_name, color, is_categorical, _sorted_reports(series) + ) if is_categorical: ax.set_xlabel("Model / Variant", fontsize=10) plt.setp(ax.xaxis.get_majorticklabels(), rotation=45, ha="right") else: ax.set_xlabel("Number of Allocations", fontsize=10) - num_allocations = [ - r.num_allocations - for allocator_data in source_data.values() - for reports in allocator_data.values() - for r in reports - ] - if num_allocations and max(num_allocations) / min(num_allocations) > 10: + num_allocations = [r.num_allocations for r in reports] + if max(num_allocations) / min(num_allocations) > 10: ax.set_xscale("log") ax.set_ylabel("Time (s)", fontsize=10, color="black") @@ -202,7 +177,7 @@ def _draw_subplot( ax.set_title(f"Source: {source_name}", fontsize=12, fontweight="bold", pad=10) -def _add_footer(campaign: BenchmarkCampaign, fig: Figure) -> None: +def _add_footer(campaign: BenchmarkCampaign, fig: "Figure") -> None: metadata_text = _format_metadata(campaign.metadata) txt = fig.text( 0.5, @@ -217,7 +192,7 @@ def _add_footer(campaign: BenchmarkCampaign, fig: Figure) -> None: txt._get_wrap_line_width = lambda: fig.bbox.width * 0.90 # ty: ignore[unresolved-attribute] # noqa: SLF001 -def _add_legend(fig: Figure, allocator_names: tuple[str, ...]) -> None: +def _add_legend(fig: "Figure", allocator_names: tuple[str, ...]) -> None: handles = [ Line2D( [], @@ -241,7 +216,7 @@ def _add_legend(fig: Figure, allocator_names: tuple[str, ...]) -> None: ) -def _create_figure(num_sources: int) -> tuple[Figure, list[Axes]]: +def _create_figure(num_sources: int) -> "tuple[Figure, list[Axes]]": fig, axs = plt.subplots( nrows=num_sources, ncols=1, @@ -256,19 +231,11 @@ def _visualize_campaign( ) -> None: source_names = campaign.source_names allocator_names = campaign.allocator_names - reports_by_source = campaign.reports_by_source_allocator_variant - - if not source_names: - raise ValueError("Campaign has no sources to visualize") - if not allocator_names: - raise ValueError("Campaign has no allocators to visualize") - if not reports_by_source: - raise ValueError("Campaign has no reports to visualize") - fig, axs = _create_figure(len(source_names)) for ax, source_name in zip(axs, source_names, strict=True): - _draw_subplot(ax, source_name, reports_by_source[source_name], allocator_names) + reports = [r for r in campaign.reports if r.source_name == source_name] + _draw_subplot(ax, source_name, reports, allocator_names) fig.tight_layout(rect=(0.01, 0.05, 0.99, 0.92)) # l, b, r, t diff --git a/src/python/omnimalloc/benchmark/timer.py b/src/python/omnimalloc/benchmark/timer.py index 17f3a78..1db37b4 100644 --- a/src/python/omnimalloc/benchmark/timer.py +++ b/src/python/omnimalloc/benchmark/timer.py @@ -7,19 +7,12 @@ class Timer: - """Lightweight timer for performance measurement. + """Context manager measuring the wall time of its block.""" - TODO(fpedd): This class is not yet thread-safe. Concurrent access from - multiple threads may result in race conditions and inconsistent state. - """ - - def __init__(self) -> None: - self._start_ns: int | None = None - self._stop_ns: int | None = None + elapsed_ns: int = 0 def __enter__(self) -> "Timer": - if not self.is_running: - self.start() + self._start_ns = time.perf_counter_ns() return self def __exit__( @@ -28,52 +21,8 @@ def __exit__( exc_val: BaseException | None, exc_tb: TracebackType | None, ) -> None: - if self.is_running: - self.stop() - - def start(self) -> "Timer": - if self.is_running: - raise RuntimeError("Timer is already running") - self._start_ns = time.perf_counter_ns() - self._stop_ns = None - return self - - def stop(self) -> "Timer": - if not self.is_running: - raise RuntimeError("Timer is not running") - self._stop_ns = time.perf_counter_ns() - return self - - @property - def is_running(self) -> bool: - return self._start_ns is not None and self._stop_ns is None - - @property - def elapsed_ns(self) -> int: - if self._start_ns is None: - raise RuntimeError("Timer has no start time") - if self._stop_ns is None: - return time.perf_counter_ns() - self._start_ns - return self._stop_ns - self._start_ns + self.elapsed_ns = time.perf_counter_ns() - self._start_ns @property def elapsed_s(self) -> float: - return self.elapsed_ns / 1_000_000_000 - - @property - def elapsed(self) -> str: - return _format_time(self.elapsed_ns) - - -def _format_time(ns: int) -> str: - if ns < 1_000: - return f"{ns} ns" - if ns < 1_000_000: - return f"{ns / 1_000:.2f} us" - if ns < 1_000_000_000: - return f"{ns / 1_000_000:.2f} ms" - if ns < 60 * 1_000_000_000: - return f"{ns / 1_000_000_000:.2f} s" - if ns < 3_600 * 1_000_000_000: - return f"{ns / (60 * 1_000_000_000):.2f} min" - return f"{ns / (3_600 * 1_000_000_000):.2f} h" + return self.elapsed_ns / 1e9 diff --git a/src/python/omnimalloc/benchmark/utils.py b/src/python/omnimalloc/benchmark/utils.py index ee8513e..99fe2d7 100644 --- a/src/python/omnimalloc/benchmark/utils.py +++ b/src/python/omnimalloc/benchmark/utils.py @@ -8,21 +8,6 @@ from tqdm.auto import tqdm except ImportError: - class _DummyProgressBar: - """No-op stand-in for tqdm's total=... progress bar.""" - - def __enter__(self) -> "_DummyProgressBar": - return self - - def __exit__(self, *args: object) -> None: - pass - - def update(self, n: int = 1) -> None: - pass - - def tqdm(iterable: Any = None, **kwargs: Any) -> Any: # noqa: ARG001, ANN401 + def tqdm(iterable: Any, **_: Any) -> Any: # noqa: ANN401 """No-op tqdm fallback when tqdm is not installed.""" - if iterable is None: - # When called with total= instead of an iterable - return _DummyProgressBar() return iterable diff --git a/tests/unit/benchmark/results/test_campaign.py b/tests/unit/benchmark/results/test_campaign.py index c5e8868..b292bff 100644 --- a/tests/unit/benchmark/results/test_campaign.py +++ b/tests/unit/benchmark/results/test_campaign.py @@ -2,7 +2,6 @@ # SPDX-License-Identifier: Apache-2.0 # - import pytest from omnimalloc import allocate from omnimalloc.allocators import GreedyAllocator @@ -11,165 +10,50 @@ BenchmarkReport, BenchmarkResult, ) -from omnimalloc.benchmark.sources.generator import RandomSource -from omnimalloc.benchmark.sources.sync_patterns import SyncPatternSource +from omnimalloc.benchmark.sources import BaseSource, RandomSource, SyncPatternSource -def test_benchmark_campaign_creation() -> None: - source = RandomSource(num_allocations=10, seed=42) +def _report(report_id: int, source: BaseSource | None = None) -> BenchmarkReport: + source = source or RandomSource(num_allocations=10, seed=42) allocator = GreedyAllocator() - pool = allocate(source.get_pool(), allocator) - result = BenchmarkResult( - id=0, allocator=allocator, source=source, entity=pool, duration=0.5 + id=0, + allocator=allocator, + source=source, + entity=allocate(source.get_pool(), allocator), + duration=0.5, ) - report = BenchmarkReport(id=0, results=(result,)) - campaign = BenchmarkCampaign(id="campaign_0", reports=(report,)) + return BenchmarkReport(id=report_id, results=(result,), source=source) + - assert campaign.num_reports == 1 - assert campaign.num_results == 1 +def test_benchmark_campaign_counts() -> None: + campaign = BenchmarkCampaign(id="c", reports=(_report(0), _report(1))) + + assert campaign.num_reports == 2 + assert campaign.num_results == 2 + assert campaign.num_allocators == 1 + assert campaign.allocator_names == ("greedy",) def test_benchmark_campaign_empty_reports_raises_error() -> None: with pytest.raises(ValueError, match="must contain at least one report"): - BenchmarkCampaign(id="campaign_0", reports=()) + BenchmarkCampaign(id="c", reports=()) def test_benchmark_campaign_duplicate_report_ids_raises_error() -> None: - source = RandomSource(num_allocations=10, seed=42) - allocator = GreedyAllocator() - pool = allocate(source.get_pool(), allocator) - - result1 = BenchmarkResult( - id=0, allocator=allocator, source=source, entity=pool, duration=0.5 - ) - result2 = BenchmarkResult( - id=1, allocator=allocator, source=source, entity=pool, duration=0.6 - ) - - report1 = BenchmarkReport(id=0, results=(result1,)) - report2 = BenchmarkReport(id=0, results=(result2,)) - with pytest.raises(ValueError, match="report ids must be unique"): - BenchmarkCampaign(id="campaign_0", reports=(report1, report2)) - - -def test_benchmark_campaign_properties() -> None: - source = RandomSource(num_allocations=10, seed=42) - allocator = GreedyAllocator() - pool = allocate(source.get_pool(), allocator) - - results = tuple( - BenchmarkResult( - id=i, allocator=allocator, source=source, entity=pool, duration=0.5 - ) - for i in range(3) - ) - report = BenchmarkReport(id=0, results=results) - campaign = BenchmarkCampaign(id="campaign_0", reports=(report,)) - - assert campaign.num_results == 3 - assert campaign.num_allocations == 30 - assert campaign.num_allocations_per_result == 10 - assert campaign.num_allocators == 1 - assert campaign.num_sources == 1 - - -def test_benchmark_campaign_metadata_deep_copies_nested_values() -> None: - source = RandomSource(num_allocations=10, seed=42) - allocator = GreedyAllocator() - pool = allocate(source.get_pool(), allocator) - result = BenchmarkResult( - id=0, allocator=allocator, source=source, entity=pool, duration=0.5 - ) - report = BenchmarkReport(id=0, results=(result,)) - metadata = {"tags": ["a"]} - campaign = BenchmarkCampaign(id="campaign_0", reports=(report,), metadata=metadata) - - metadata["tags"].append("b") - - assert campaign.metadata == {"tags": ["a"]} - - -def test_benchmark_campaign_finalize_metadata() -> None: - source = RandomSource(num_allocations=10, seed=42) - allocator = GreedyAllocator() - pool = allocate(source.get_pool(), allocator) - - result = BenchmarkResult( - id=0, allocator=allocator, source=source, entity=pool, duration=0.5 - ) - report = BenchmarkReport(id=0, results=(result,)) - campaign = BenchmarkCampaign( - id="campaign_0", reports=(report,), metadata={"custom": "value"} - ) - finalized = campaign.finalize_metadata() - - assert "custom" in finalized.metadata - assert "num_reports" in finalized.metadata - - -def _report( - source: SyncPatternSource, report_id: int, result_id: int -) -> BenchmarkReport: - pool = allocate(source.get_pool(), GreedyAllocator()) - result = BenchmarkResult( - id=result_id, - allocator=GreedyAllocator(), - source=source, - entity=pool, - duration=0.5, - ) - return BenchmarkReport(id=report_id, results=(result,), source=source) + BenchmarkCampaign(id="c", reports=(_report(0), _report(0))) def test_benchmark_campaign_keeps_thread_counts_apart() -> None: few = SyncPatternSource(num_allocations=16, num_threads=2) many = SyncPatternSource(num_allocations=16, num_threads=8) - campaign = BenchmarkCampaign( - id="campaign_0", reports=(_report(few, 0, 0), _report(many, 1, 1)) - ) + campaign = BenchmarkCampaign(id="c", reports=(_report(0, few), _report(1, many))) assert campaign.num_sources == 2 assert campaign.source_names == tuple(sorted((few.label(), many.label()))) - assert set(campaign.reports_by_source_allocator_variant) == set( - campaign.source_names - ) def test_benchmark_campaign_groups_unlabelled_sources_together() -> None: - source = RandomSource(num_allocations=10, seed=42) - pool = allocate(source.get_pool(), GreedyAllocator()) - reports = tuple( - BenchmarkReport( - id=i, - results=( - BenchmarkResult( - id=i, - allocator=GreedyAllocator(), - source=source, - entity=pool, - duration=0.5, - ), - ), - source=source, - ) - for i in range(2) - ) - campaign = BenchmarkCampaign(id="campaign_0", reports=reports) - + campaign = BenchmarkCampaign(id="c", reports=(_report(0), _report(1))) assert campaign.source_names == ("random",) - - -def test_benchmark_campaign_rejects_non_copyable_metadata_with_clear_error() -> None: - source = RandomSource(num_allocations=10, seed=42) - allocator = GreedyAllocator() - pool = allocate(source.get_pool(), allocator) - result = BenchmarkResult( - id=0, allocator=allocator, source=source, entity=pool, duration=0.5 - ) - report = BenchmarkReport(id=0, results=(result,)) - with pytest.raises(TypeError, match="deep-copyable"): - BenchmarkCampaign( - id="c", reports=(report,), metadata={"gen": (i for i in range(1))} - ) diff --git a/tests/unit/benchmark/results/test_export.py b/tests/unit/benchmark/results/test_export.py index fba5d14..3fcb20b 100644 --- a/tests/unit/benchmark/results/test_export.py +++ b/tests/unit/benchmark/results/test_export.py @@ -2,7 +2,6 @@ # SPDX-License-Identifier: Apache-2.0 # - import csv import json from pathlib import Path @@ -11,22 +10,22 @@ import pytest from omnimalloc import allocate from omnimalloc.allocators import GreedyAllocator +from omnimalloc.benchmark import run_benchmark from omnimalloc.benchmark.results import ( BenchmarkCampaign, BenchmarkReport, BenchmarkResult, ) -from omnimalloc.benchmark.results.export import ( - RESULTS_CSV_COLUMNS, - _prepare_base_dir, - _write_metadata, - save_benchmark, -) -from omnimalloc.benchmark.sources.generator import RandomSource +from omnimalloc.benchmark.results.export import RESULTS_CSV_COLUMNS, save_benchmark +from omnimalloc.benchmark.sources import MinimallocSource, RandomSource from omnimalloc.benchmark.sources.sync_patterns import SyncPatternSource +from omnimalloc.io import save_allocation +from omnimalloc.primitives import Allocation, Pool from tests.markers import needs_matplotlib +pytestmark = needs_matplotlib + @pytest.fixture def simple_campaign() -> BenchmarkCampaign: @@ -43,285 +42,137 @@ def simple_campaign() -> BenchmarkCampaign: ) -@needs_matplotlib +def _rows(output_path: Path) -> list[dict[str, str]]: + with (output_path / "results.csv").open(newline="") as f: + return list(csv.DictReader(f)) + + def test_save_benchmark_creates_directory( - simple_campaign: BenchmarkCampaign, artifacts_dir: Path + simple_campaign: BenchmarkCampaign, tmp_path: Path ) -> None: - output_path = artifacts_dir / "campaign_output" - result_path = save_benchmark( - simple_campaign, - output_path=output_path, - output_format="dir", - visualize_iterations=False, + simple_campaign, tmp_path / "out", visualize_iterations=False ) - assert result_path.exists() - assert result_path.is_dir() - assert (result_path / "metadata.json").exists() - assert (result_path / "campaign_overview.pdf").exists() + assert result_path == tmp_path / "out" + assert sorted(p.name for p in result_path.iterdir()) == [ + "campaign_overview.pdf", + "metadata.json", + "results.csv", + ] + assert json.loads((result_path / "metadata.json").read_text()) == {"test": "value"} -@needs_matplotlib +@pytest.mark.parametrize("name", ["out", "out.zip"]) def test_save_benchmark_creates_zip( - simple_campaign: BenchmarkCampaign, artifacts_dir: Path + simple_campaign: BenchmarkCampaign, tmp_path: Path, name: str ) -> None: - output_path = artifacts_dir / "campaign_output" - result_path = save_benchmark( simple_campaign, - output_path=output_path, + tmp_path / name, output_format="zip", visualize_iterations=False, ) - assert result_path.exists() - assert result_path.suffix == ".zip" - assert result_path.is_file() - + assert result_path == tmp_path / "out.zip" with ZipFile(result_path, "r") as zip_file: names = zip_file.namelist() - assert any("metadata.json" in name for name in names) - assert any("campaign_overview.pdf" in name for name in names) + assert "out/metadata.json" in names + assert all(name.startswith("out/") for name in names) -@needs_matplotlib -def test_save_benchmark_with_none_path( +def test_save_benchmark_defaults_to_artifacts_under_cwd( simple_campaign: BenchmarkCampaign, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.chdir(tmp_path) - - result_path = save_benchmark( - simple_campaign, - output_path=None, - output_format="dir", - visualize_iterations=False, - ) - - assert result_path.exists() - assert result_path.is_relative_to(tmp_path) - assert "campaign_test_campaign" in str(result_path) + result_path = save_benchmark(simple_campaign, visualize_iterations=False) + assert result_path == tmp_path / "artifacts" / "test_campaign" -def test_save_benchmark_raises_typeerror_for_non_campaign(artifacts_dir: Path) -> None: - source = RandomSource(num_allocations=10, seed=42) - allocator = GreedyAllocator() - pool = allocate(source.get_pool(), allocator) - result = BenchmarkResult( - id=0, allocator=allocator, source=source, entity=pool, duration=0.5 - ) - - with pytest.raises(TypeError, match="Expected a BenchmarkCampaign"): - save_benchmark(result, output_path=artifacts_dir / "output") # type: ignore[arg-type] - - -def test_save_benchmark_raises_valueerror_for_invalid_format( - simple_campaign: BenchmarkCampaign, artifacts_dir: Path +def test_save_benchmark_rejects_invalid_arguments( + simple_campaign: BenchmarkCampaign, tmp_path: Path ) -> None: + with pytest.raises(TypeError, match="Expected a BenchmarkCampaign"): + save_benchmark(simple_campaign.reports[0], tmp_path) # type: ignore[arg-type] with pytest.raises(ValueError, match="output_format must be 'dir' or 'zip'"): - save_benchmark( - simple_campaign, - output_path=artifacts_dir / "output", - output_format="invalid", # type: ignore[arg-type] - ) + save_benchmark(simple_campaign, tmp_path, output_format="tar") # type: ignore[arg-type] -@needs_matplotlib -def test_save_benchmark_raises_fileexistserror_when_not_overwriting( - simple_campaign: BenchmarkCampaign, artifacts_dir: Path +def test_save_benchmark_overwrites_only_when_allowed( + simple_campaign: BenchmarkCampaign, tmp_path: Path ) -> None: - output_path = artifacts_dir / "campaign_output" - - save_benchmark( - simple_campaign, output_path=output_path, output_format="dir", overwrite=True - ) + output_path = tmp_path / "out" + save_benchmark(simple_campaign, output_path) + (output_path / "stale").touch() + save_benchmark(simple_campaign, output_path) + assert not (output_path / "stale").exists() with pytest.raises(FileExistsError, match="already exists"): - save_benchmark( - simple_campaign, - output_path=output_path, - output_format="dir", - overwrite=False, - ) - - -@needs_matplotlib -def test_save_benchmark_overwrites_existing_directory( - simple_campaign: BenchmarkCampaign, artifacts_dir: Path -) -> None: - output_path = artifacts_dir / "campaign_output" - - result1 = save_benchmark( - simple_campaign, output_path=output_path, output_format="dir", overwrite=True - ) - result2 = save_benchmark( - simple_campaign, output_path=output_path, output_format="dir", overwrite=True - ) - - assert result1 == result2 - assert result2.exists() - - -def test_write_metadata_creates_json_file( - simple_campaign: BenchmarkCampaign, artifacts_dir: Path -) -> None: - _write_metadata(artifacts_dir, simple_campaign) - - metadata_file = artifacts_dir / "metadata.json" - assert metadata_file.exists() - - with metadata_file.open("r") as f: - metadata = json.load(f) - assert "test" in metadata - assert metadata["test"] == "value" - + save_benchmark(simple_campaign, output_path, overwrite=False) -def test_prepare_base_dir_creates_directory(artifacts_dir: Path) -> None: - output_path = artifacts_dir / "test_dir" - base_dir = _prepare_base_dir(output_path, output_format="dir", overwrite=True) - - assert base_dir == output_path - assert base_dir.exists() - assert base_dir.is_dir() - - -def test_prepare_base_dir_creates_temp_for_zip(artifacts_dir: Path) -> None: - output_path = artifacts_dir / "test_zip" - - base_dir = _prepare_base_dir(output_path, output_format="zip", overwrite=True) - - assert base_dir != output_path - assert base_dir.exists() - assert "omnimalloc_dump_" in str(base_dir) - - -@needs_matplotlib def test_save_benchmark_writes_results_csv( - simple_campaign: BenchmarkCampaign, artifacts_dir: Path + simple_campaign: BenchmarkCampaign, tmp_path: Path ) -> None: - output_path = save_benchmark( - simple_campaign, - output_path=artifacts_dir / "campaign_output", - output_format="dir", - visualize_iterations=False, + rows = _rows( + save_benchmark(simple_campaign, tmp_path / "out", visualize_iterations=False) ) - with (output_path / "results.csv").open(newline="") as f: - rows = list(csv.DictReader(f)) - - assert len(rows) == simple_campaign.num_reports assert tuple(rows[0]) == RESULTS_CSV_COLUMNS - assert rows[0]["source"] == "random" - assert rows[0]["allocator"] == "greedy" - assert rows[0]["num_allocations"] == "10" - assert rows[0]["iterations"] == "1" - assert float(rows[0]["mean_seconds"]) == 0.5 - assert rows[0]["stdev_seconds"] == "" - assert rows[0]["known_optimum"] == "" - assert rows[0]["optimum_ratio"] == "" - - -@needs_matplotlib -def test_results_csv_has_one_row_per_report(artifacts_dir: Path) -> None: - source = RandomSource(num_allocations=10, seed=42) - allocator = GreedyAllocator() - pool = allocate(source.get_pool(), allocator) - reports = tuple( - BenchmarkReport( - id=i, - results=tuple( - BenchmarkResult( - id=10 * i + j, - allocator=allocator, - source=source, - entity=pool, - duration=float(j + 1), - ) - for j in range(2) - ), - variant_id=10 * (i + 1), - known_optimum=pool.pressure, - ) - for i in range(3) - ) - campaign = BenchmarkCampaign(id="multi", reports=reports) - - output_path = save_benchmark( - campaign, - output_path=artifacts_dir / "multi_output", - output_format="dir", - visualize_iterations=False, - ) - - with (output_path / "results.csv").open(newline="") as f: - rows = list(csv.DictReader(f)) + assert rows == [ + rows[0] + | { + "source": "random", + "allocator": "greedy", + "variant": "10", + "num_allocations": "10", + "iterations": "1", + "mean_seconds": "0.5", + "stdev_seconds": "", + "known_optimum": "", + "optimum_ratio": "", + } + ] + + +def test_results_csv_has_one_row_per_report(tmp_path: Path) -> None: + campaign = run_benchmark( + allocators=("greedy",), + sources=(RandomSource(num_allocations=10, seed=42),), + variants=(10, 20, 30), + iterations=2, + ) + rows = _rows(save_benchmark(campaign, tmp_path / "out", visualize_iterations=False)) assert [row["variant"] for row in rows] == ["10", "20", "30"] - assert all(float(row["stdev_seconds"]) > 0 for row in rows) - assert all(float(row["optimum_ratio"]) >= 1.0 for row in rows) - assert all(int(row["lower_bound"]) == pool.pressure for row in rows) + assert all(float(row["stdev_seconds"]) >= 0 for row in rows) -@needs_matplotlib -def test_results_csv_leaves_unmeasurable_efficiency_empty(artifacts_dir: Path) -> None: +def test_results_csv_leaves_unmeasurable_efficiency_empty(tmp_path: Path) -> None: source = SyncPatternSource( num_allocations=2000, num_threads=64, pattern="independent" ) - allocator = GreedyAllocator() - pool = allocate(source.get_pool(), allocator) - result = BenchmarkResult( - id=0, allocator=allocator, source=source, entity=pool, duration=0.5 - ) - campaign = BenchmarkCampaign( - id="wide", reports=(BenchmarkReport(id=0, results=(result,), source=source),) - ) + campaign = run_benchmark(allocators=("greedy",), sources=(source,)) + rows = _rows(save_benchmark(campaign, tmp_path / "out", visualize_iterations=False)) - output_path = save_benchmark( - campaign, - output_path=artifacts_dir / "wide_output", - output_format="dir", - visualize_iterations=False, - ) - - with (output_path / "results.csv").open(newline="") as f: - rows = list(csv.DictReader(f)) - - assert len(rows) == 1 assert rows[0]["mean_efficiency"] == "" assert rows[0]["lower_bound"] == "" - assert float(rows[0]["mean_seconds"]) == 0.5 assert int(rows[0]["mean_peak_size"]) > 0 - assert (output_path / "campaign_overview.pdf").exists() -@needs_matplotlib -def test_save_benchmark_zip_path_lands_exactly_there( - simple_campaign: BenchmarkCampaign, artifacts_dir: Path -) -> None: - output_path = artifacts_dir / "campaign_output.zip" - result_path = save_benchmark( - simple_campaign, - output_path=output_path, - output_format="zip", - visualize_iterations=False, +def test_save_benchmark_plots_iterations_one_dir_per_label(tmp_path: Path) -> None: + csv_dir = tmp_path / "csvs" + csv_dir.mkdir() + pool = Pool(id="p", allocations=(Allocation(id=0, size=8, start=0, end=4),)) + save_allocation(pool, csv_dir / "p.csv") + campaign = run_benchmark( + allocators=("greedy",), sources=(MinimallocSource(csv_dir=csv_dir),) ) - assert result_path == output_path - assert result_path.is_file() - assert not output_path.with_suffix(".zip.zip").exists() + output_path = save_benchmark(campaign, tmp_path / "out") -def test_zip_internal_folder_uses_campaign_name( - simple_campaign: BenchmarkCampaign, artifacts_dir: Path -) -> None: - output_path = artifacts_dir / "named_campaign" - result_path = save_benchmark( - simple_campaign, - output_path=output_path, - output_format="zip", - visualize_iterations=False, - ) - with ZipFile(result_path, "r") as zip_file: - assert all(name.startswith("named_campaign/") for name in zip_file.namelist()) + (source_dir,) = (output_path / "sources").iterdir() + iteration = source_dir / "allocators" / "greedy" / "p" / "iterations" + assert (iteration / "iteration_0.pdf").is_file() diff --git a/tests/unit/benchmark/results/test_report.py b/tests/unit/benchmark/results/test_report.py index fd1c054..8bc5516 100644 --- a/tests/unit/benchmark/results/test_report.py +++ b/tests/unit/benchmark/results/test_report.py @@ -2,179 +2,84 @@ # SPDX-License-Identifier: Apache-2.0 # - import pytest from omnimalloc import allocate -from omnimalloc.allocators import GreedyAllocator, NaiveAllocator +from omnimalloc.allocators import BaseAllocator, GreedyAllocator, NaiveAllocator from omnimalloc.benchmark.results import BenchmarkReport, BenchmarkResult -from omnimalloc.benchmark.sources.generator import RandomSource -from omnimalloc.benchmark.sources.sync_patterns import SyncPatternSource - - -def test_benchmark_report_creation() -> None: - source = RandomSource(num_allocations=10, seed=42) - allocator = GreedyAllocator() - pool = allocate(source.get_pool(), allocator) - result = BenchmarkResult( - id=0, allocator=allocator, source=source, entity=pool, duration=0.5 - ) - - report = BenchmarkReport(id=0, results=(result,)) - assert report.num_results == 1 - assert report.num_allocations == 10 - - -def test_benchmark_report_empty_results_raises_error() -> None: - with pytest.raises(ValueError, match="must contain at least one result"): - BenchmarkReport(id=0, results=()) - - -def test_benchmark_report_duplicate_ids_raises_error() -> None: - source = RandomSource(num_allocations=10, seed=42) - allocator = GreedyAllocator() - pool = allocate(source.get_pool(), allocator) - - result1 = BenchmarkResult( - id=0, allocator=allocator, source=source, entity=pool, duration=0.5 - ) - result2 = BenchmarkResult( - id=0, allocator=allocator, source=source, entity=pool, duration=0.6 - ) - - with pytest.raises(ValueError, match="result ids must be unique"): - BenchmarkReport(id=0, results=(result1, result2)) - - -def test_benchmark_report_statistics() -> None: - source = RandomSource(num_allocations=10, seed=42) - allocator = GreedyAllocator() - pool = allocate(source.get_pool(), allocator) - - results = tuple( - BenchmarkResult( - id=i, allocator=allocator, source=source, entity=pool, duration=float(i) - ) - for i in range(3) +from omnimalloc.benchmark.sources import BaseSource, RandomSource, SyncPatternSource + + +def _result( + result_id: int = 0, + duration: float = 1.0, + allocator: BaseAllocator | None = None, + source: BaseSource | None = None, +) -> BenchmarkResult: + allocator = allocator or GreedyAllocator() + source = source or RandomSource(num_allocations=10, seed=42) + return BenchmarkResult( + id=result_id, + allocator=allocator, + source=source, + entity=allocate(source.get_pool(), allocator), + duration=duration, ) - report = BenchmarkReport(id=0, results=results) - assert report.mean_seconds > 0 - assert report.median_seconds > 0 - assert 0.0 <= report.mean_allocation_efficiency <= 1.0 - -def test_benchmark_report_allocator_mismatch_raises_error() -> None: - source = RandomSource(num_allocations=10, seed=42) - allocator1 = GreedyAllocator() - allocator2 = NaiveAllocator() - - pool1 = allocate(source.get_pool(), allocator1) - pool2 = allocate(source.get_pool(), allocator2) - - result1 = BenchmarkResult( - id=0, allocator=allocator1, source=source, entity=pool1, duration=0.5 - ) - result2 = BenchmarkResult( - id=1, allocator=allocator2, source=source, entity=pool2, duration=0.6 - ) - - with pytest.raises(ValueError, match="Allocator mismatch"): - BenchmarkReport(id=0, results=(result1, result2), allocator=allocator1) - - -def test_benchmark_report_with_results() -> None: - source = RandomSource(num_allocations=10, seed=42) - allocator = GreedyAllocator() - pool = allocate(source.get_pool(), allocator) - - result1 = BenchmarkResult( - id=0, allocator=allocator, source=source, entity=pool, duration=0.5 - ) - result2 = BenchmarkResult( - id=1, allocator=allocator, source=source, entity=pool, duration=0.6 - ) - - report1 = BenchmarkReport(id=0, results=(result1,)) - report2 = report1.with_results((result2,)) - - assert len(report1.results) == 1 - assert len(report2.results) == 2 - - -def _report(durations: tuple[float, ...], **overrides: object) -> BenchmarkReport: - source = RandomSource(num_allocations=10, seed=42) - allocator = GreedyAllocator() - pool = allocate(source.get_pool(), allocator) - results = tuple( - BenchmarkResult( - id=i, allocator=allocator, source=source, entity=pool, duration=d - ) - for i, d in enumerate(durations) - ) +def _report(*durations: float, **overrides: object) -> BenchmarkReport: + results = tuple(_result(i, d) for i, d in enumerate(durations)) return BenchmarkReport(id=0, results=results, **overrides) # type: ignore[arg-type] -def test_benchmark_report_dispersion_statistics() -> None: - report = _report((1.0, 2.0, 3.0)) +def test_benchmark_report_statistics() -> None: + report = _report(1.0, 2.0, 3.0) + assert report.num_results == 3 + assert report.num_allocations == 10 + assert report.mean_seconds == report.median_seconds == 2.0 assert report.min_seconds == 1.0 assert report.max_seconds == 3.0 assert report.stdev_seconds == pytest.approx(1.0) + assert 0.0 <= report.mean_allocation_efficiency <= 1.0 + assert report.mean_peak_size >= report.lower_bound > 0 def test_benchmark_report_stdev_is_none_for_single_iteration() -> None: - report = _report((1.0,)) + assert _report(1.0).stdev_seconds is None - assert report.stdev_seconds is None - assert report.min_seconds == report.max_seconds == 1.0 +def test_benchmark_report_empty_results_raises_error() -> None: + with pytest.raises(ValueError, match="must contain at least one result"): + BenchmarkReport(id=0, results=()) -def test_benchmark_report_peak_size_and_lower_bound() -> None: - report = _report((1.0,)) - assert report.lower_bound > 0 - assert report.mean_peak_size >= report.lower_bound +def test_benchmark_report_duplicate_ids_raises_error() -> None: + with pytest.raises(ValueError, match="result ids must be unique"): + BenchmarkReport(id=0, results=(_result(0), _result(0))) -def test_benchmark_report_optimum_ratio_absent_without_known_optimum() -> None: - assert _report((1.0,)).known_optimum is None - assert _report((1.0,)).optimum_ratio is None +def test_benchmark_report_allocator_mismatch_raises_error() -> None: + results = (_result(0), _result(1, allocator=NaiveAllocator())) + with pytest.raises(ValueError, match="Allocator mismatch"): + BenchmarkReport(id=0, results=results, allocator=GreedyAllocator()) -def test_benchmark_report_optimum_ratio_compares_peak_to_optimum() -> None: - report = _report((1.0,)) - with_optimum = BenchmarkReport( - id=0, results=report.results, known_optimum=int(report.mean_peak_size) // 2 - ) +def test_benchmark_report_optimum_ratio() -> None: + report = _report(1.0) + half = int(report.mean_peak_size) // 2 + + assert report.optimum_ratio is None + assert _report(1.0, known_optimum=half).optimum_ratio == pytest.approx(2.0, 0.01) - assert with_optimum.optimum_ratio == pytest.approx(2.0, rel=0.01) - - -def test_benchmark_report_with_results_keeps_known_optimum() -> None: - report = _report((1.0,), known_optimum=1024) - extended = report.with_results( - ( - BenchmarkResult( - id=99, - allocator=report.results[0].allocator, - source=report.results[0].source, - entity=report.results[0].entity, - duration=2.0, - ), - ) - ) - assert extended.known_optimum == 1024 - assert extended.stdev_seconds is not None +@pytest.mark.parametrize(("variant_id", "label"), [(None, "10"), (7, "7"), ("m", "m")]) +def test_benchmark_report_variant_label(variant_id: object, label: str) -> None: + assert _report(1.0, variant_id=variant_id).variant_label == label def test_benchmark_report_source_name_uses_instance_label() -> None: source = SyncPatternSource(num_allocations=16, num_threads=8) - pool = allocate(source.get_pool(), GreedyAllocator()) - result = BenchmarkResult( - id=0, allocator=GreedyAllocator(), source=source, entity=pool, duration=0.5 - ) - report = BenchmarkReport(id=0, results=(result,), source=source) + report = BenchmarkReport(id=0, results=(_result(source=source),), source=source) assert report.source_name == source.label() assert "num_threads=8" in report.source_name diff --git a/tests/unit/benchmark/results/test_visualize.py b/tests/unit/benchmark/results/test_visualize.py index bbeba4d..75300c5 100644 --- a/tests/unit/benchmark/results/test_visualize.py +++ b/tests/unit/benchmark/results/test_visualize.py @@ -12,7 +12,7 @@ _canonicalize_artifact, _format_metadata, _get_allocator_color, - _get_sorted_reports, + _sorted_reports, plot_benchmark, ) from omnimalloc.benchmark.sources.generator import RandomSource @@ -99,21 +99,15 @@ def test_plot_benchmark_without_path_shows_figure( def test_sorted_reports_handles_mixed_variant_id_types() -> None: - mixed = {"a": (_report("r0", "small"), _report("r1", 100), _report("r2", None))} - assert len(_get_sorted_reports(mixed)) == 3 + mixed = [_report("r0", "small"), _report("r1", 100), _report("r2", None)] + assert len(_sorted_reports(mixed)) == 3 def test_sorted_reports_orders_numeric_variants_by_size() -> None: - numeric = { - "a": ( - _report("r0", 30, size=30), - _report("r1", 10, size=10), - _report("r2", 20, size=20), - ) - } - assert [r.variant_id for r in _get_sorted_reports(numeric)] == [10, 20, 30] + numeric = [_report(f"r{n}", n, size=n) for n in (30, 10, 20)] + assert [r.variant_id for r in _sorted_reports(numeric)] == [10, 20, 30] def test_sorted_reports_orders_categorical_variants_by_name() -> None: - categorical = {"a": (_report("r0", "c"), _report("r1", "a"), _report("r2", "b"))} - assert [r.variant_id for r in _get_sorted_reports(categorical)] == ["a", "b", "c"] + categorical = [_report(f"r{name}", name) for name in "cab"] + assert [r.variant_id for r in _sorted_reports(categorical)] == ["a", "b", "c"] diff --git a/tests/unit/benchmark/test_benchmark.py b/tests/unit/benchmark/test_benchmark.py index b68a4a5..84f7495 100644 --- a/tests/unit/benchmark/test_benchmark.py +++ b/tests/unit/benchmark/test_benchmark.py @@ -2,175 +2,88 @@ # SPDX-License-Identifier: Apache-2.0 # - import inspect +from pathlib import Path import pytest from omnimalloc.allocators import GreedyAllocator, NaiveAllocator from omnimalloc.allocators.supermalloc import SupermallocAllocator from omnimalloc.benchmark.benchmark import run_benchmark -from omnimalloc.benchmark.sources.concurrent_tiling import ConcurrentTilingSource -from omnimalloc.benchmark.sources.generator import RandomSource -from omnimalloc.benchmark.sources.pinwheel import PinwheelSource -from omnimalloc.benchmark.sources.sync_patterns import SyncPatternSource -from omnimalloc.benchmark.sources.tiling import TilingSource - - -def test_run_benchmark_basic() -> None: - source = RandomSource(num_allocations=10, seed=42) - allocator = GreedyAllocator() +from omnimalloc.benchmark.results import BenchmarkCampaign +from omnimalloc.benchmark.sources import ( + ConcurrentTilingSource, + MinimallocSource, + PinwheelSource, + RandomSource, + SyncPatternSource, + TilingSource, +) +from omnimalloc.io import save_allocation +from omnimalloc.primitives import Allocation, Pool - campaign = run_benchmark( - allocators=(allocator,), - sources=(source,), - iterations=1, - variants=10, - ) - assert campaign.num_reports >= 1 - assert campaign.num_results >= 1 +def _run(**kwargs: object) -> BenchmarkCampaign: + """One greedy run on ten random allocations unless overridden.""" + defaults = { + "allocators": (GreedyAllocator(),), + "sources": (RandomSource(num_allocations=10, seed=42),), + } + return run_benchmark(**(defaults | kwargs)) # type: ignore[arg-type] -def test_run_benchmark_multiple_allocators() -> None: - source = RandomSource(num_allocations=10, seed=42) - allocator1 = GreedyAllocator() - allocator2 = NaiveAllocator() +def _variants(campaign: BenchmarkCampaign, source: object) -> set[object]: + return {r.variant_id for r in campaign.reports if r.source == source} - campaign = run_benchmark( - allocators=(allocator1, allocator2), - sources=(source,), - iterations=1, - variants=10, - ) - assert campaign.num_allocators == 2 +def _vector_source() -> ConcurrentTilingSource: + return ConcurrentTilingSource(num_allocations=16, num_threads=2, num_syncs=8) -def test_run_benchmark_multiple_iterations() -> None: - source = RandomSource(num_allocations=10, seed=42) - allocator = GreedyAllocator() +def test_run_benchmark_defaults() -> None: + campaign = _run() - campaign = run_benchmark( - allocators=(allocator,), - sources=(source,), - iterations=3, - variants=10, + assert inspect.signature(run_benchmark).parameters["validate"].default is True + assert campaign.num_reports == campaign.num_results == 1 + assert campaign.reports[0].num_allocations == 10 + assert campaign.reports[0].known_optimum is None + assert campaign.metadata["skipped"] == [] + assert {"total_duration", "num_reports", "omnimalloc_version"} <= set( + campaign.metadata ) - assert all(report.num_results == 3 for report in campaign.reports) +def test_run_benchmark_repeats_iterations_per_allocator() -> None: + campaign = _run(allocators=(GreedyAllocator(), NaiveAllocator()), iterations=3) -def test_run_benchmark_metadata() -> None: - source = RandomSource(num_allocations=10, seed=42) - allocator = GreedyAllocator() + assert campaign.num_allocators == 2 + assert all(report.num_results == 3 for report in campaign.reports) - campaign = run_benchmark( - allocators=(allocator,), - sources=(source,), - iterations=1, - variants=10, - ) - assert "total_duration" in campaign.metadata - assert "num_reports" in campaign.metadata +def test_run_benchmark_keeps_a_falsy_campaign_id() -> None: + assert _run(campaign_id=0).id == 0 def test_run_benchmark_per_source_variants() -> None: - source = RandomSource(num_allocations=10, seed=42) - allocator = GreedyAllocator() - - campaign = run_benchmark( - allocators=(allocator,), - sources=(source,), - iterations=1, - variants={"random": (5, 10)}, - ) - - assert campaign.num_reports == 2 + campaign = _run(variants={"random": (5, 10)}) assert {r.variant_id for r in campaign.reports} == {5, 10} def test_run_benchmark_on_vector_clock_source() -> None: - source = ConcurrentTilingSource(num_allocations=16, num_threads=2, num_syncs=8) - - campaign = run_benchmark( - allocators=(GreedyAllocator(),), - sources=(source,), - iterations=1, - variants=16, - validate=True, - ) - - assert campaign.num_reports == 1 + campaign = _run(sources=(_vector_source(),)) assert campaign.reports[0].mean_allocation_efficiency > 0 -def test_run_benchmark_skips_scalar_only_allocators_on_vector_source() -> None: - source = ConcurrentTilingSource(num_allocations=16, num_threads=2, num_syncs=8) - - campaign = run_benchmark( - allocators=(SupermallocAllocator(), GreedyAllocator()), - sources=(source,), - iterations=1, - variants=16, - ) - - assert campaign.num_reports == 1 - assert campaign.reports[0].allocator_name == "greedy" - - -def test_run_benchmark_skips_unsupported_variants() -> None: - source = ConcurrentTilingSource(num_allocations=16, num_threads=4, num_syncs=8) - - campaign = run_benchmark( - allocators=(GreedyAllocator(),), - sources=(source,), - iterations=1, - variants=(2, 16), - ) - - assert campaign.num_reports == 1 - assert campaign.reports[0].variant_id == 16 - - -def test_run_benchmark_raises_when_all_pairs_skipped() -> None: - source = ConcurrentTilingSource(num_allocations=16, num_threads=2, num_syncs=8) - - with pytest.raises(ValueError, match="No benchmark reports"): - run_benchmark( - allocators=(SupermallocAllocator(),), - sources=(source,), - iterations=1, - variants=16, - ) - - -def test_run_benchmark_validates_by_default() -> None: - campaign = run_benchmark( - allocators=(GreedyAllocator(),), - sources=(RandomSource(num_allocations=10, seed=42),), - iterations=1, - variants=10, - ) - - assert inspect.signature(run_benchmark).parameters["validate"].default is True - assert campaign.num_reports == 1 - - -def test_run_benchmark_records_skipped_allocators_in_metadata() -> None: - source = ConcurrentTilingSource(num_allocations=16, num_threads=2, num_syncs=8) - - campaign = run_benchmark( - allocators=(SupermallocAllocator(), GreedyAllocator()), - sources=(source,), - iterations=1, - variants=(16, 32), +def test_run_benchmark_skips_unsupported_allocators() -> None: + source = _vector_source() + campaign = _run( + allocators=(SupermallocAllocator(), GreedyAllocator()), sources=(source,) ) - skipped = campaign.metadata["skipped_allocators"] - assert skipped == [ + assert [r.allocator_name for r in campaign.reports] == ["greedy"] + assert campaign.metadata["skipped"] == [ { "source": source.label(), + "variant": "16", "allocator": "supermalloc", "reason": ( "supermalloc requires scalar (interval) lifetimes, " @@ -180,30 +93,21 @@ def test_run_benchmark_records_skipped_allocators_in_metadata() -> None: ] -def test_run_benchmark_metadata_lists_no_skips_when_all_supported() -> None: - campaign = run_benchmark( - allocators=(GreedyAllocator(),), - sources=(RandomSource(num_allocations=10, seed=42),), - iterations=1, - variants=10, - ) - - assert campaign.metadata["skipped_allocators"] == [] - assert campaign.metadata["skipped_variants"] == [] +def test_run_benchmark_raises_when_all_pairs_skipped() -> None: + with pytest.raises(ValueError, match="No benchmark reports"): + _run(allocators=(SupermallocAllocator(),), sources=(_vector_source(),)) -def test_run_benchmark_records_unreachable_variants_in_metadata() -> None: +def test_run_benchmark_skips_unreachable_variants_once() -> None: source = PinwheelSource(num_allocations=65) - - campaign = run_benchmark( - allocators=(GreedyAllocator(),), + campaign = _run( + allocators=(GreedyAllocator(), NaiveAllocator()), sources=(source,), - iterations=1, variants=(64, 65), ) - assert campaign.num_reports == 1 - assert campaign.metadata["skipped_variants"] == [ + assert _variants(campaign, source) == {65} + assert campaign.metadata["skipped"] == [ { "source": source.label(), "variant": "64", @@ -212,122 +116,62 @@ def test_run_benchmark_records_unreachable_variants_in_metadata() -> None: ] -def test_run_benchmark_reports_an_unreachable_variant_once_per_source() -> None: - campaign = run_benchmark( - allocators=(GreedyAllocator(), NaiveAllocator()), - sources=(PinwheelSource(num_allocations=65),), - iterations=1, - variants=(64, 65), - ) - - assert len(campaign.metadata["skipped_variants"]) == 1 - - -def test_run_benchmark_records_known_optimum_when_available() -> None: - capacity = 1024 * 1024 - campaign = run_benchmark( - allocators=(GreedyAllocator(),), - sources=(TilingSource(num_allocations=32, capacity=capacity),), - iterations=1, - variants=32, - ) +def test_run_benchmark_records_known_optimum() -> None: + campaign = _run(sources=(TilingSource(num_allocations=32, capacity=4096),)) - report = campaign.reports[0] - assert report.known_optimum == capacity - assert report.optimum_ratio is not None - assert report.optimum_ratio >= 1.0 + assert campaign.reports[0].known_optimum == 4096 + assert campaign.reports[0].optimum_ratio >= 1.0 -def test_run_benchmark_leaves_known_optimum_empty_without_ground_truth() -> None: - campaign = run_benchmark( - allocators=(GreedyAllocator(),), - sources=(RandomSource(num_allocations=10, seed=42),), - iterations=1, - variants=10, - ) - - assert campaign.reports[0].known_optimum is None - assert campaign.reports[0].optimum_ratio is None - - -def test_run_benchmark_keeps_thread_counts_as_separate_series() -> None: +def test_run_benchmark_variants_can_be_keyed_by_label() -> None: few = SyncPatternSource(num_allocations=16, num_threads=2) many = SyncPatternSource(num_allocations=16, num_threads=8) - campaign = run_benchmark( - allocators=(GreedyAllocator(),), - sources=(few, many), - iterations=1, - variants=16, + campaign = _run( + sources=(few, many), variants={few.label(): 16, "sync_pattern": (24, 32)} ) - assert campaign.num_sources == 2 assert campaign.source_names == tuple(sorted((few.label(), many.label()))) + assert _variants(campaign, few) == {16} + assert _variants(campaign, many) == {24, 32} -def test_run_benchmark_variants_can_be_keyed_by_label() -> None: - few = SyncPatternSource(num_allocations=16, num_threads=2) - many = SyncPatternSource(num_allocations=16, num_threads=8) +def test_run_benchmark_resolves_fixed_variants_by_name_and_index( + tmp_path: Path, +) -> None: + for name in ("a", "b", "c"): + pool = Pool(id=name, allocations=(Allocation(id=0, size=8, start=0, end=4),)) + save_allocation(pool, tmp_path / f"{name}.csv") + source = MinimallocSource(csv_dir=tmp_path) - campaign = run_benchmark( - allocators=(GreedyAllocator(),), - sources=(few, many), - iterations=1, - variants={few.label(): 16, "sync_pattern": (24, 32)}, - ) - - by_source = campaign.reports_by_source_allocator_variant - assert set(by_source[few.label()]["greedy"]) == {"16"} - assert set(by_source[many.label()]["greedy"]) == {"24", "32"} + assert _variants(_run(sources=(source,), variants=("c", 0)), source) == {"a", "c"} + assert _variants(_run(sources=(source,), variants=2), source) == {"a", "b"} + with pytest.raises(ValueError, match="Unknown variant 'd'"): + _run(sources=(source,), variants=("d",)) def test_run_benchmark_tolerates_unmeasurable_pressure() -> None: - campaign = run_benchmark( - allocators=(GreedyAllocator(),), - sources=( - SyncPatternSource( - num_allocations=2000, num_threads=64, pattern="independent" - ), - ), - iterations=1, - variants=2000, + source = SyncPatternSource( + num_allocations=2000, num_threads=64, pattern="independent" ) + report = _run(sources=(source,)).reports[0] - report = campaign.reports[0] assert report.num_allocations == 2000 assert report.mean_seconds > 0 assert report.mean_allocation_efficiency is None assert report.lower_bound is None -def test_run_benchmark_defaults_to_source_configured_size() -> None: - source = RandomSource(num_allocations=12, seed=42) - campaign = run_benchmark( - allocators=(GreedyAllocator(),), sources=(source,), iterations=1 - ) - assert campaign.reports[0].num_allocations == 12 - - -def test_run_benchmark_rejects_unknown_variants_key() -> None: - source = RandomSource(num_allocations=10, seed=42) - with pytest.raises(ValueError, match="match no source"): - run_benchmark( - allocators=(GreedyAllocator(),), sources=(source,), variants={"randm": 10} - ) - - -def test_run_benchmark_rejects_non_integer_variant_for_parameterizable_source() -> None: - source = RandomSource(num_allocations=10, seed=42) - with pytest.raises(TypeError, match="Non-integer variant"): - run_benchmark( - allocators=(GreedyAllocator(),), sources=(source,), variants=("small",) - ) - - -def test_run_benchmark_rejects_non_positive_iterations() -> None: - with pytest.raises(ValueError, match="iterations must be positive"): - run_benchmark( - allocators=(GreedyAllocator(),), - sources=(RandomSource(num_allocations=10, seed=42),), - iterations=0, - ) +@pytest.mark.parametrize( + ("kwargs", "error", "match"), + [ + ({"variants": {"randm": 10}}, ValueError, "match no source"), + ({"variants": ("small",)}, TypeError, "Non-integer variant"), + ({"iterations": 0}, ValueError, "iterations must be positive"), + ], +) +def test_run_benchmark_rejects_invalid_arguments( + kwargs: dict[str, object], error: type[Exception], match: str +) -> None: + with pytest.raises(error, match=match): + _run(**kwargs) diff --git a/tests/unit/benchmark/test_timer.py b/tests/unit/benchmark/test_timer.py index 5387836..794468f 100644 --- a/tests/unit/benchmark/test_timer.py +++ b/tests/unit/benchmark/test_timer.py @@ -2,150 +2,27 @@ # SPDX-License-Identifier: Apache-2.0 # - import time -import pytest -from omnimalloc.benchmark.timer import Timer, _format_time - - -def test_init_default() -> None: - timer = Timer() - assert not timer.is_running - assert timer._start_ns is None # noqa: SLF001 - assert timer._stop_ns is None # noqa: SLF001 - - -def test_start() -> None: - timer = Timer() - result = timer.start() - assert timer.is_running - assert timer._start_ns is not None # noqa: SLF001 - assert result is timer - - -def test_start_already_running() -> None: - timer = Timer() - timer.start() - with pytest.raises(RuntimeError, match="Timer is already running"): - timer.start() - - -def test_stop() -> None: - timer = Timer() - timer.start() - time.sleep(0.001) - result = timer.stop() - assert not timer.is_running - assert timer._stop_ns is not None # noqa: SLF001 - assert timer.elapsed_ns > 0 - assert result is timer - +from omnimalloc.benchmark.timer import Timer -def test_stop_not_running() -> None: - timer = Timer() - with pytest.raises(RuntimeError, match="Timer is not running"): - timer.stop() - - -def test_is_running_property() -> None: - timer = Timer() - assert not timer.is_running - timer.start() - assert timer.is_running - timer.stop() - assert not timer.is_running - - -def test_elapsed_ns_stopped(monkeypatch: pytest.MonkeyPatch) -> None: - ticks = iter([1_000, 2_000_000]) - monkeypatch.setattr(time, "perf_counter_ns", lambda: next(ticks)) - timer = Timer() - timer.start() - timer.stop() - elapsed = timer.elapsed_ns - assert elapsed == 1_999_000 - assert timer.elapsed_ns == elapsed - - -def test_elapsed_ns_running() -> None: - timer = Timer() - timer.start() - time.sleep(0.001) - elapsed1 = timer.elapsed_ns - time.sleep(0.001) - elapsed2 = timer.elapsed_ns - assert elapsed1 > 0 - assert elapsed2 > elapsed1 - -def test_elapsed_ns_never_started() -> None: - timer = Timer() - with pytest.raises(RuntimeError, match="Timer has no start time"): - _ = timer.elapsed_ns - - -def test_elapsed_s() -> None: - timer = Timer() - timer.start() - time.sleep(0.001) - timer.stop() - assert timer.elapsed_s == pytest.approx(timer.elapsed_ns / 1_000_000_000) - - -def test_elapsed_formatted() -> None: - timer = Timer() - timer.start() - time.sleep(0.001) - timer.stop() - elapsed = timer.elapsed - assert isinstance(elapsed, str) - assert "ms" in elapsed or "us" in elapsed - - -def test_context_manager_basic() -> None: +def test_timer_measures_its_block() -> None: with Timer() as timer: - assert timer.is_running - time.sleep(0.001) - assert not timer.is_running - assert timer.elapsed_ns > 0 + time.sleep(0.01) + assert timer.elapsed_s >= 0.01 + assert timer.elapsed_s == timer.elapsed_ns / 1e9 -def test_context_manager_access_after() -> None: +def test_timer_freezes_on_exit() -> None: with Timer() as timer: - time.sleep(0.001) + pass elapsed = timer.elapsed_ns - assert elapsed > 0 - assert timer.elapsed_ns == elapsed - - -def test_restart_without_reset() -> None: - timer = Timer() - timer.start() - time.sleep(0.001) - timer.stop() - elapsed1 = timer.elapsed_ns - timer.start() time.sleep(0.001) - timer.stop() - assert elapsed1 > 0 - assert timer.elapsed_ns > 0 + assert timer.elapsed_ns == elapsed -def test_nested_context_managers() -> None: - with Timer() as outer: - time.sleep(0.001) - with Timer() as inner: - time.sleep(0.001) - inner_elapsed = inner.elapsed_ns +def test_timer_nests() -> None: + with Timer() as outer, Timer() as inner: time.sleep(0.001) - outer_elapsed = outer.elapsed_ns - - assert inner_elapsed > 0 - assert outer_elapsed > inner_elapsed - - -def test_format_time_minutes_and_hours() -> None: - assert _format_time(90 * 10**9) == "1.50 min" - assert _format_time(30 * 60 * 10**9) == "30.00 min" - assert _format_time(2 * 3600 * 10**9) == "2.00 h" + assert outer.elapsed_ns >= inner.elapsed_ns > 0 From da5e6935c94b5060a57606ba04abafc16271d958 Mon Sep 17 00:00:00 2001 From: Fabian Peddinghaus Date: Tue, 29 Sep 2026 14:45:06 +0000 Subject: [PATCH 3/4] Trim the examples, notebooks and README asset script --- examples/03_allocators.py | 25 +++++++---------- examples/04_sources.py | 28 +++++-------------- examples/05_benchmark.py | 42 +++++++---------------------- examples/README.md | 6 ++--- notebooks/custom_source.ipynb | 7 ++--- notebooks/extensive_benchmark.ipynb | 12 ++------- scripts/generate_readme_assets.py | 6 ++--- 7 files changed, 36 insertions(+), 90 deletions(-) diff --git a/examples/03_allocators.py b/examples/03_allocators.py index e294bc2..4ed757b 100644 --- a/examples/03_allocators.py +++ b/examples/03_allocators.py @@ -6,7 +6,6 @@ import omnimalloc as om from omnimalloc.allocators import DEFAULT_ALLOCATOR, available_allocators -from omnimalloc.allocators.minimalloc import HAS_MINIMALLOC def main() -> None: @@ -18,24 +17,20 @@ def main() -> None: alloc_2 = om.Allocation(id="alloc_2", size=4, start=5, end=15) alloc_3 = om.Allocation(id="alloc_3", size=5, start=15, end=23) - # Create pool and allocate + # Create pool pool = om.Pool(id="pool_0", allocations=(alloc_0, alloc_1, alloc_2, alloc_3)) - # Get and run the default allocator - print(f"Running allocation with default allocator: {DEFAULT_ALLOCATOR}") - placed = om.allocate(pool, allocator=DEFAULT_ALLOCATOR, validate=True) - print(f"Pool {placed.id!r} size: {placed.size}") - om.plot_allocation(placed, example_dir / f"{DEFAULT_ALLOCATOR}_default.pdf") - - # Run allocation with all available allocators + # Run allocation with every registered allocator; without one, + # om.allocate uses the default + print(f"Default allocator: {DEFAULT_ALLOCATOR}") for allocator_name in available_allocators(): - # minimalloc is an optional dependency that only builds on some platforms - if "minimalloc" in allocator_name and not HAS_MINIMALLOC: - print(f"Skipping unavailable allocator: {allocator_name}") + try: + placed = om.allocate(pool, allocator_name, validate=True) + except ImportError as error: + # Optional allocators wrap libraries that may not be installed + print(f"Skipping {allocator_name}: {str(error).splitlines()[0]}") continue - print(f"Running allocation with allocator: {allocator_name}") - placed = om.allocate(pool, allocator_name, validate=True) - print(f"Pool {placed.id!r} size: {placed.size}") + print(f"Pool {placed.id!r} size with {allocator_name}: {placed.size}") om.plot_allocation(placed, example_dir / f"{allocator_name}.pdf") diff --git a/examples/04_sources.py b/examples/04_sources.py index 9b9203a..b573866 100644 --- a/examples/04_sources.py +++ b/examples/04_sources.py @@ -5,34 +5,18 @@ from pathlib import Path import omnimalloc as om -from omnimalloc.benchmark.sources import ( - DEFAULT_SOURCE, - BaseSource, - available_sources, -) - - -def allocate_and_plot(source: BaseSource, output: Path) -> None: - pool = source.get_pool() - pool = om.allocate(pool, validate=True) - print(f"Pool {pool.id!r} size: {pool.size}") - om.plot_allocation(pool, output) +from omnimalloc.benchmark.sources import BaseSource def main() -> None: example_dir = Path("04_example_output") - # Get and use the default source - default_source = BaseSource.get(DEFAULT_SOURCE)() - print(f"Using default source: {DEFAULT_SOURCE}") - allocate_and_plot( - default_source, example_dir / f"source_{DEFAULT_SOURCE}_default.pdf" - ) - - for source_name in available_sources(): + # A few synthetic sources, each generating a seeded, reproducible pool + for source_name in ("random", "skewed", "tiling", "pinwheel"): source = BaseSource.get(source_name)() - print(f"Using source: {source_name}") - allocate_and_plot(source, example_dir / f"source_{source_name}.pdf") + pool = om.allocate(source.get_pool(), validate=True) + print(f"Source {source_name!r} pool size: {pool.size}") + om.plot_allocation(pool, example_dir / f"source_{source_name}.pdf") if __name__ == "__main__": diff --git a/examples/05_benchmark.py b/examples/05_benchmark.py index 7275aa4..4027ad3 100644 --- a/examples/05_benchmark.py +++ b/examples/05_benchmark.py @@ -4,54 +4,30 @@ from pathlib import Path -from omnimalloc.allocators.minimalloc import HAS_MINIMALLOC -from omnimalloc.benchmark import ( - VariantSpec, - plot_benchmark, - run_benchmark, - save_benchmark, -) +from omnimalloc.benchmark import run_benchmark, save_benchmark def main() -> None: - example_dir = Path("05_example_output") - - # Define allocators, sources, and variants to benchmark + # Allocators whose optional library is missing are skipped, not fatal allocators = ( "greedy_by_size", "greedy_by_all", "omni", "best_fit", "telamalloc", - ) - # minimalloc is an optional dependency that only builds on some platforms - if HAS_MINIMALLOC: - allocators += ("minimalloc",) - sources = ( - "random", "minimalloc", - "huggingface", ) - # Counts for the parameterizable source, "first 5" for the fixed ones - variants: dict[str, VariantSpec] = { - "random": (10, 50, 100, 250, 500), - "minimalloc": 5, - "huggingface": 5, - } - - # Run benchmark campaign + campaign = run_benchmark( allocators=allocators, - sources=sources, - variants=variants, - validate=True, + sources=("random", "minimalloc", "huggingface"), + # Counts for the parameterizable source, "first 5" for the Minimalloc + # one; the Hugging Face source downloads a single model by default + variants={"random": (10, 50, 100, 250, 500), "minimalloc": 5}, ) - # Visualize - plot_benchmark(campaign, example_dir / "benchmark_results.pdf") - - # Save results (contains overview and individual allocation plots) - save_benchmark(campaign, example_dir / "benchmark_results") + # Writes the overview plot, a results CSV, and one plot per iteration + save_benchmark(campaign, Path("05_example_output") / "benchmark_results") if __name__ == "__main__": diff --git a/examples/README.md b/examples/README.md index 3577fb8..26b7e0a 100644 --- a/examples/README.md +++ b/examples/README.md @@ -1,9 +1,9 @@ # OmniMalloc Examples -This directory contains example scripts demonstrating the usage of the OmniMalloc. +This directory contains example scripts demonstrating the usage of OmniMalloc. -1. `01_basic.py`: A basic example showcasing how to define allocations, create a pool, run the allocation algorithm. -2. `02_plotting.py`: Same as `basic.py` but with additional visualization of the allocation using `plot_allocation()`. +1. `01_basic.py`: A basic example showcasing how to define allocations, create a pool, and run the allocation algorithm. +2. `02_plotting.py`: Same as `01_basic.py` but with additional visualization of the allocation using `plot_allocation()`. 3. `03_allocators.py`: An example using different allocators from the `omnimalloc.allocators` module. 4. `04_sources.py`: An example demonstrating how to use different allocation sources for generating workloads. 5. `05_benchmark.py`: A full benchmarking script that compares different allocation strategies on various workloads, diff --git a/notebooks/custom_source.ipynb b/notebooks/custom_source.ipynb index a818126..ba2a2d0 100644 --- a/notebooks/custom_source.ipynb +++ b/notebooks/custom_source.ipynb @@ -42,9 +42,10 @@ "id": "3", "metadata": {}, "source": [ - "`get_allocations(num_allocations, skip)` returns a deterministic tuple;\n", - "`skip` discards that many leading allocations so the base class can carve\n", - "consecutive, non-repeating batches for multi-pool workloads.\n", + "`get_allocations(num_allocations, skip)` returns a deterministic tuple.\n", + "`get_pools` asks pool `i` for `skip = i * num_allocations`, so `skip` must\n", + "select a different batch: this source discards `skip` leading draws and\n", + "continues one stream, so consecutive pools never repeat.\n", "\n", "This source models the activations of a sequential layer chain: layer `i`\n", "writes its output at time `i` and the consumer at layer `c` reads it while\n", diff --git a/notebooks/extensive_benchmark.ipynb b/notebooks/extensive_benchmark.ipynb index 4164126..d3935f9 100644 --- a/notebooks/extensive_benchmark.ipynb +++ b/notebooks/extensive_benchmark.ipynb @@ -18,11 +18,7 @@ "metadata": {}, "outputs": [], "source": [ - "from omnimalloc.benchmark import (\n", - " run_benchmark,\n", - " plot_benchmark,\n", - " save_benchmark,\n", - ")" + "from omnimalloc.benchmark import plot_benchmark, run_benchmark" ] }, { @@ -36,13 +32,11 @@ " \"greedy_by_size\",\n", " \"greedy_by_area\",\n", " \"greedy_by_conflict\",\n", - " # \"minimalloc\",\n", ")\n", "\n", "sources = (\n", " \"random\",\n", " \"minimalloc\",\n", - " # \"huggingface\",\n", ")" ] }, @@ -60,10 +54,8 @@ " \"random\": (10, 100, 500, 1_000, 2_000, 5_000, 10_000),\n", " \"minimalloc\": None, # all bundled pools\n", " },\n", - " validate=True,\n", ")\n", - "plot_benchmark(campaign)\n", - "# save_benchmark(campaign, \"benchmark_results\")" + "plot_benchmark(campaign)" ] } ], diff --git a/scripts/generate_readme_assets.py b/scripts/generate_readme_assets.py index 13d0618..4c79925 100644 --- a/scripts/generate_readme_assets.py +++ b/scripts/generate_readme_assets.py @@ -9,10 +9,10 @@ import argparse import json -import random import shutil import subprocess import sys +from bisect import bisect_left from dataclasses import dataclass from pathlib import Path from statistics import mean @@ -34,7 +34,6 @@ from matplotlib.figure import Figure from omnimalloc.primitives import Pool -SEED = 0 MINIMALLOC_URL = "git+https://github.com/google/minimalloc.git" SCALING_SIZES = (10, 32, 100, 316, 1000, 3162, 10000) SCALING_SIZES_SLOW = SCALING_SIZES[:-1] # minimalloc cannot solve 10k in budget @@ -231,7 +230,6 @@ def _hard_suite() -> "dict[str, Pool]": def collect_data() -> dict[str, Any]: _ensure_minimalloc() - random.seed(SEED) suite = _hard_suite() hard = [k for k in suite if k != "random"] @@ -602,7 +600,7 @@ def render_allocation(data: dict[str, Any], theme: Theme, preview: Path | None) ) ordered_sizes = sorted(r[3] for r in rects) for start, duration, offset, height in rects: - quantile = ordered_sizes.index(height) / max(len(ordered_sizes) - 1, 1) + quantile = bisect_left(ordered_sizes, height) / max(len(ordered_sizes) - 1, 1) ax.add_patch( Rectangle( (start, offset), From cf93fae3b6778ec1371c14f408ee940541325ba7 Mon Sep 17 00:00:00 2001 From: Fabian Peddinghaus Date: Tue, 29 Sep 2026 17:19:52 +0000 Subject: [PATCH 4/4] Replace every character Windows forbids when naming benchmark export directories --- src/python/omnimalloc/benchmark/results/export.py | 4 ++-- tests/unit/benchmark/results/test_export.py | 10 +++++++++- 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/src/python/omnimalloc/benchmark/results/export.py b/src/python/omnimalloc/benchmark/results/export.py index 69acf05..83e29ed 100644 --- a/src/python/omnimalloc/benchmark/results/export.py +++ b/src/python/omnimalloc/benchmark/results/export.py @@ -93,8 +93,8 @@ def _create_zip_archive(base_dir: Path, final_path: Path) -> Path: def _dir_name(label: str) -> str: - """A label as one path component; labels may carry an explicit path.""" - return re.sub(r"[/\\]", "_", label) + """A label as one path component, safe on every OS; labels may carry paths.""" + return re.sub(r'[<>:"/\\|?*]', "_", label) def _write_iterations(base_dir: Path, campaign: BenchmarkCampaign) -> None: diff --git a/tests/unit/benchmark/results/test_export.py b/tests/unit/benchmark/results/test_export.py index 3fcb20b..d3f53d3 100644 --- a/tests/unit/benchmark/results/test_export.py +++ b/tests/unit/benchmark/results/test_export.py @@ -16,7 +16,11 @@ BenchmarkReport, BenchmarkResult, ) -from omnimalloc.benchmark.results.export import RESULTS_CSV_COLUMNS, save_benchmark +from omnimalloc.benchmark.results.export import ( + RESULTS_CSV_COLUMNS, + _dir_name, + save_benchmark, +) from omnimalloc.benchmark.sources import MinimallocSource, RandomSource from omnimalloc.benchmark.sources.sync_patterns import SyncPatternSource from omnimalloc.io import save_allocation @@ -176,3 +180,7 @@ def test_save_benchmark_plots_iterations_one_dir_per_label(tmp_path: Path) -> No (source_dir,) = (output_path / "sources").iterdir() iteration = source_dir / "allocators" / "greedy" / "p" / "iterations" assert (iteration / "iteration_0.pdf").is_file() + + +def test_dir_name_is_one_component_on_every_os() -> None: + assert _dir_name(r'src[dir=C:\a/b] <"x|y?*>') == "src[dir=C__a_b] __x_y___"