diff --git a/pyproject.toml b/pyproject.toml index a7aef3c..532bc29 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -50,6 +50,7 @@ all = [ dev = [ "ipykernel>=6.30.0", "ipywidgets>=8.1.0", + "nanobind>=2.0", "nbconvert>=7.16.0", "pre-commit>=4.0.0", "pytest>=8.0", diff --git a/src/python/omnimalloc/_allocate.py b/src/python/omnimalloc/_allocate.py index 11846a4..1b2fecb 100644 --- a/src/python/omnimalloc/_allocate.py +++ b/src/python/omnimalloc/_allocate.py @@ -3,37 +3,59 @@ # -from typing import TypeVar, cast +from collections.abc import Sequence +from typing import TypeAlias, TypeVar, cast, overload from .allocators import DEFAULT_ALLOCATOR, BaseAllocator -from .primitives import Memory, Pool, System +from .primitives import Allocation, Memory, Pool, System from .validate import validate_allocation +AllocatorLike: TypeAlias = BaseAllocator | type[BaseAllocator] | str | None + T = TypeVar("T", System, Memory, Pool) +@overload def allocate( entity: T, - allocator: BaseAllocator | type[BaseAllocator] | str | None = None, - *, + allocator: AllocatorLike = None, + validate: bool = False, +) -> T: ... + + +@overload +def allocate( + entity: Sequence[Allocation], + allocator: AllocatorLike = None, validate: bool = False, -) -> T: - """Return the entity (System, Memory, or Pool) with offsets assigned. +) -> tuple[Allocation, ...]: ... + - `allocator` accepts an instance, a class, a registry name, or `None` - (the default allocator); `validate=True` additionally runs - `validate_allocation` on the result. +def allocate( + entity: System | Memory | Pool | Sequence[Allocation], + allocator: AllocatorLike = None, + validate: bool = False, +) -> System | Memory | Pool | tuple[Allocation, ...]: + """Return the entity with offsets assigned. + + Accepts a System, Memory, or Pool (returned as the same type) or a raw + sequence of Allocations (returned as a tuple in input order). `allocator` + accepts an instance, a class, a registry name, or `None` (the default + allocator); `validate=True` additionally runs `validate_allocation` on + the result. """ if allocator is None: allocator = DEFAULT_ALLOCATOR - resolved = BaseAllocator.resolve(allocator) + resolved = cast("BaseAllocator", BaseAllocator.resolve(allocator)) - # ty doesn't understand that TypeVar T (System|Memory|Pool) all have allocate method - allocated = entity.allocate(resolved) # type: ignore[invalid-argument-type] + if isinstance(entity, System | Memory | Pool): + allocated = entity.allocate(resolved) + else: + allocated = Pool.from_allocations(entity).allocate(resolved).allocations if validate: validate_allocation(allocated) - return cast("T", allocated) + return allocated diff --git a/src/python/omnimalloc/allocators/genetic.py b/src/python/omnimalloc/allocators/genetic.py index 70d93d2..f4d51fa 100644 --- a/src/python/omnimalloc/allocators/genetic.py +++ b/src/python/omnimalloc/allocators/genetic.py @@ -44,7 +44,6 @@ class GeneticAllocator(GreedyAllocator): def __init__( self, - *, seed: int = DEFAULT_SEED, population_size: int = 100, max_generations: int = 50, diff --git a/src/python/omnimalloc/allocators/greedy.py b/src/python/omnimalloc/allocators/greedy.py index 6950343..084840c 100644 --- a/src/python/omnimalloc/allocators/greedy.py +++ b/src/python/omnimalloc/allocators/greedy.py @@ -77,7 +77,7 @@ class GreedyByAllAllocator(GreedyAllocator): `num_threads=None` uses all cores. """ - def __init__(self, *, num_threads: int | None = None) -> None: + def __init__(self, num_threads: int | None = None) -> None: ensure_valid_num_threads(num_threads) self._num_threads = num_threads diff --git a/src/python/omnimalloc/allocators/greedy_base.py b/src/python/omnimalloc/allocators/greedy_base.py index e00f2b8..2dfddde 100644 --- a/src/python/omnimalloc/allocators/greedy_base.py +++ b/src/python/omnimalloc/allocators/greedy_base.py @@ -72,7 +72,6 @@ def _allocate( def allocate_parallel( allocations: tuple[Allocation, ...], variants: tuple[BaseAllocator, ...], - *, num_threads: int | None = None, ) -> tuple[Allocation, ...]: """Run each variant and return the lowest peak memory results. diff --git a/src/python/omnimalloc/allocators/hillclimb.py b/src/python/omnimalloc/allocators/hillclimb.py index 2eb9f1c..9f1b4a2 100644 --- a/src/python/omnimalloc/allocators/hillclimb.py +++ b/src/python/omnimalloc/allocators/hillclimb.py @@ -33,7 +33,6 @@ class HillClimbAllocator(GreedyAllocator): def __init__( self, - *, max_iterations: int = 100, seed: int = DEFAULT_SEED, acceptance_temperature: float = 2.0, diff --git a/src/python/omnimalloc/allocators/minimalloc.py b/src/python/omnimalloc/allocators/minimalloc.py index b2256af..e430c0f 100644 --- a/src/python/omnimalloc/allocators/minimalloc.py +++ b/src/python/omnimalloc/allocators/minimalloc.py @@ -51,7 +51,7 @@ class MinimallocAllocator(BaseAllocator): supports_vector_time = False def __init__( - self, *, timeout: float | None = DEFAULT_TIMEOUT, capacity: int = 1 * TB + self, timeout: float | None = DEFAULT_TIMEOUT, capacity: int = 1 * TB ) -> None: _require_minimalloc() ensure_valid_timeout(timeout) diff --git a/src/python/omnimalloc/allocators/omni.py b/src/python/omnimalloc/allocators/omni.py index aeb6817..cfe703b 100644 --- a/src/python/omnimalloc/allocators/omni.py +++ b/src/python/omnimalloc/allocators/omni.py @@ -23,7 +23,7 @@ class OmniAllocator(BaseAllocator): supports_vector_time = True - def __init__(self, *, linearize_budget: int | None = DEFAULT_WORK_BUDGET) -> None: + def __init__(self, linearize_budget: int | None = DEFAULT_WORK_BUDGET) -> None: ensure_valid_budget(linearize_budget, name="linearize_budget") self._linearize_budget = linearize_budget diff --git a/src/python/omnimalloc/allocators/random.py b/src/python/omnimalloc/allocators/random.py index 3b8f43d..ab111a0 100644 --- a/src/python/omnimalloc/allocators/random.py +++ b/src/python/omnimalloc/allocators/random.py @@ -15,7 +15,7 @@ class RandomAllocator(GreedyAllocator): """Randomized allocator that tries multiple random orders and picks the best.""" - def __init__(self, *, num_trials: int = 100, seed: int = DEFAULT_SEED) -> None: + def __init__(self, num_trials: int = 100, seed: int = DEFAULT_SEED) -> None: ensure_non_negative(num_trials, "num_trials") self._seed = seed self._num_trials = num_trials diff --git a/src/python/omnimalloc/allocators/simulated_annealing.py b/src/python/omnimalloc/allocators/simulated_annealing.py index 63f6234..b66ad64 100644 --- a/src/python/omnimalloc/allocators/simulated_annealing.py +++ b/src/python/omnimalloc/allocators/simulated_annealing.py @@ -31,7 +31,6 @@ class SimulatedAnnealingAllocator(BaseAllocator): def __init__( self, - *, seed: int = DEFAULT_SEED, max_iterations: int = 3000, initial_temperature: float = 3.0, diff --git a/src/python/omnimalloc/allocators/supermalloc.py b/src/python/omnimalloc/allocators/supermalloc.py index 99f226b..31a6775 100644 --- a/src/python/omnimalloc/allocators/supermalloc.py +++ b/src/python/omnimalloc/allocators/supermalloc.py @@ -133,7 +133,6 @@ class SupermallocAllocator(BaseAllocator): def __init__( self, - *, timeout: float | None = DEFAULT_TIMEOUT, heuristics: tuple[Heuristic, ...] = DEFAULT_HEURISTICS, num_threads: int | None = None, diff --git a/src/python/omnimalloc/allocators/tabu_search.py b/src/python/omnimalloc/allocators/tabu_search.py index 208084a..d9ff1c3 100644 --- a/src/python/omnimalloc/allocators/tabu_search.py +++ b/src/python/omnimalloc/allocators/tabu_search.py @@ -30,7 +30,6 @@ class TabuSearchAllocator(BaseAllocator): def __init__( self, - *, seed: int = DEFAULT_SEED, max_iterations: int = 500, neighborhood_size: int = 20, diff --git a/src/python/omnimalloc/allocators/telamalloc.py b/src/python/omnimalloc/allocators/telamalloc.py index fab890b..fcbcdaa 100644 --- a/src/python/omnimalloc/allocators/telamalloc.py +++ b/src/python/omnimalloc/allocators/telamalloc.py @@ -36,7 +36,6 @@ class TelamallocAllocator(BaseAllocator): def __init__( self, - *, seed: int = DEFAULT_SEED, max_backtracks: int = 10000, timeout: float | None = DEFAULT_TIMEOUT, diff --git a/src/python/omnimalloc/analysis/__init__.py b/src/python/omnimalloc/analysis/__init__.py index b2f6cf0..f2f9975 100644 --- a/src/python/omnimalloc/analysis/__init__.py +++ b/src/python/omnimalloc/analysis/__init__.py @@ -4,6 +4,10 @@ from ._conflicts import conflict_degrees as conflict_degrees from ._conflicts import conflicts as conflicts +from ._pressure import antichain_pressure as antichain_pressure +from ._pressure import ( + antichain_pressure_per_allocation as antichain_pressure_per_allocation, +) from ._pressure import closure_pressure as closure_pressure from ._pressure import ( closure_pressure_per_allocation as closure_pressure_per_allocation, diff --git a/src/python/omnimalloc/analysis/_pressure.py b/src/python/omnimalloc/analysis/_pressure.py index 247dde1..41a3fa8 100644 --- a/src/python/omnimalloc/analysis/_pressure.py +++ b/src/python/omnimalloc/analysis/_pressure.py @@ -4,7 +4,7 @@ from collections.abc import Sequence -from omnimalloc._cpp import antichain_pressure +from omnimalloc._cpp import antichain_pressure as _antichain_pressure from omnimalloc._cpp import ( antichain_pressure_per_allocation as _antichain_pressure_per_allocation, ) @@ -21,7 +21,7 @@ from omnimalloc.primitives.utils import ensure_unique_ids -def pressure( +def antichain_pressure( allocations: Sequence[Allocation], work_budget: int | None = DEFAULT_WORK_BUDGET ) -> int: """Peak memory pressure: exact max-weight antichain of the happens-before order. @@ -36,7 +36,7 @@ def pressure( the flow work exceeds `work_budget`; pass `None` to always compute. """ ensure_valid_budget(work_budget) - return antichain_pressure(allocations, work_budget) + return _antichain_pressure(allocations, work_budget) def closure_pressure( @@ -46,7 +46,7 @@ def closure_pressure( C++ enumeration of the join-closure of the birth clocks. Pairwise- concurrent allocations need not share a cut, so this can sit strictly - below `pressure`; both soundly lower-bound any placement's peak. Raises + below `antichain_pressure`; both soundly lower-bound any placement's peak. Raises `RuntimeError` once the closure exceeds `closure_cap`; pass `None` to always enumerate — memory then grows with the closure itself. """ @@ -58,7 +58,7 @@ def placement_pressure(allocations: Sequence[Allocation]) -> int: """Peak of a placement: the highest occupied address, max(offset + size). Simply the pressure the placement realizes after allocation — an upper - bound on `pressure` (and so on `closure_pressure`), equal to the max + bound on `antichain_pressure` (and so on `closure_pressure`), equal to the max entry of `placement_pressure_per_allocation`. Raises `ValueError` on unplaced input. """ @@ -71,14 +71,14 @@ def placement_pressure(allocations: Sequence[Allocation]) -> int: return max(heights, default=0) -def pressure_per_allocation( +def antichain_pressure_per_allocation( allocations: Sequence[Allocation], work_budget: int | None = DEFAULT_WORK_BUDGET ) -> dict[IdType, int]: """Peak pressure over each allocation's own lifetime, keyed by id. The max-weight antichain through each allocation: the tightest order-derived lower bound on the pressure any placement can exhibit - while that allocation is live; the max entry equals `pressure`. + while that allocation is live; the max entry equals `antichain_pressure`. Linearizable instances take one O(N log N) window sweep, genuinely partial orders solve one pinned antichain (min flow over the conflict neighborhood) per distinct lifetime — exact, but built to certify @@ -98,7 +98,7 @@ def closure_pressure_per_allocation( """Exact realizable peak while each allocation is live, keyed by id. The max total size at any join-closure cut where the allocation is - live. Can sit elementwise strictly below `pressure_per_allocation`, + live. Can sit elementwise strictly below `antichain_pressure_per_allocation`, since pairwise-concurrent allocations need not share a cut; the max entry equals `closure_pressure`. Raises `RuntimeError` once the closure exceeds `closure_cap`; pass `None` to always enumerate — @@ -128,6 +128,11 @@ def placement_pressure_per_allocation( return _keyed_by_id(allocations, peaks) +# The antichain bound is the canonical/default pressure metric +pressure = antichain_pressure +pressure_per_allocation = antichain_pressure_per_allocation + + def _keyed_by_id( allocations: Sequence[Allocation], peaks: list[int] ) -> dict[IdType, int]: diff --git a/src/python/omnimalloc/benchmark/timer.py b/src/python/omnimalloc/benchmark/timer.py index 90fc2c9..930457d 100644 --- a/src/python/omnimalloc/benchmark/timer.py +++ b/src/python/omnimalloc/benchmark/timer.py @@ -20,7 +20,7 @@ class Timer: multiple threads may result in race conditions and inconsistent state. """ - def __init__(self, *, auto_start: bool = False) -> None: + def __init__(self, auto_start: bool = False) -> None: self._start_ns: int | None = None self._stop_ns: int | None = None self._is_running: bool = False diff --git a/src/python/omnimalloc/io.py b/src/python/omnimalloc/io.py index ef0fc02..4d35b41 100644 --- a/src/python/omnimalloc/io.py +++ b/src/python/omnimalloc/io.py @@ -3,6 +3,7 @@ # import csv +from collections.abc import Sequence from pathlib import Path from .analysis.clock import time_components @@ -39,7 +40,7 @@ def _write_pool(pool: Pool, path: Path) -> Path: # Any placed allocation brings in the offset column (minimalloc's # solution format; unplaced rows leave the cell blank), so save/load # round-trips placements instead of silently stripping them. - with_offsets = any(alloc.is_allocated for alloc in pool.allocations) + with_offsets = pool.any_allocated fields = ("id", "lower", "upper", "size") if with_offsets: fields = (*fields, "offset") @@ -56,11 +57,12 @@ def _write_pool(pool: Pool, path: Path) -> Path: def save_allocation( - entity: System | Memory | Pool, path: str | Path + entity: System | Memory | Pool | Sequence[Allocation], path: str | Path ) -> tuple[Path, ...]: """Save the entity's pools to disk as minimalloc-format CSV files. - Saving a `Pool` writes exactly `path`; a `Memory` or `System` fans out + Saving a `Pool` (or a raw sequence of Allocations, saved as one pool) + writes exactly `path`; a `Memory` or `System` fans out to one `_.csv` per pool (the CSV is pool-level minimalloc interchange). Returns the tuple of paths actually written. Pools with any placed allocation include an `offset` column (minimalloc's solution @@ -72,10 +74,10 @@ def save_allocation( """ path_ = Path(path) path_.parent.mkdir(parents=True, exist_ok=True) + if not isinstance(entity, System | Memory | Pool): + entity = Pool.from_allocations(entity) if isinstance(entity, Pool): return (_write_pool(entity, path_),) - if not isinstance(entity, (Memory, System)): - raise TypeError(f"Unsupported entity type: {type(entity)!r}") return tuple( _write_pool(pool, path_.with_name(f"{path_.stem}_{name}.csv")) for name, pool in _collect_pools(entity).items() diff --git a/src/python/omnimalloc/primitives/memory.py b/src/python/omnimalloc/primitives/memory.py index 6c4abb6..b4fa25d 100644 --- a/src/python/omnimalloc/primitives/memory.py +++ b/src/python/omnimalloc/primitives/memory.py @@ -19,19 +19,19 @@ class Memory: """A physical memory unit containing one or more pools. - `capacity` is the declared memory limit (an input); the computed extent - of a placement is `used_size`. `capacity=None` means unbounded. + `size` is the declared memory size (an input); the computed extent of a + placement is `used_size`. `size=None` means unbounded. """ id: IdType pools: tuple[Pool, ...] - capacity: int | None = None + size: int | None = None def __post_init__(self) -> None: if len({pool.id for pool in self.pools}) != len(self.pools): raise ValueError("pool ids must be unique") - if self.capacity is not None: - ensure_non_negative(self.capacity, "capacity") + if self.size is not None: + ensure_non_negative(self.size, "size") @cached_property def used_size(self) -> int: @@ -43,9 +43,14 @@ def is_allocated(self) -> bool: """True if all pools have been allocated.""" return all(pool.is_allocated for pool in self.pools) + @cached_property + def any_allocated(self) -> bool: + """True if any pool has a placed allocation.""" + return any(pool.any_allocated for pool in self.pools) + def with_pools(self, pools: tuple[Pool, ...]) -> "Memory": """Return new Memory with specified pools.""" - return Memory(id=self.id, capacity=self.capacity, pools=pools) + return Memory(id=self.id, size=self.size, pools=pools) def allocate(self, allocator: "BaseAllocator") -> "Memory": """Apply allocator to all pools.""" diff --git a/src/python/omnimalloc/primitives/pool.py b/src/python/omnimalloc/primitives/pool.py index 5ae52c1..a2f898e 100644 --- a/src/python/omnimalloc/primitives/pool.py +++ b/src/python/omnimalloc/primitives/pool.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 # +from collections.abc import Sequence from dataclasses import dataclass from functools import cached_property from typing import TYPE_CHECKING @@ -13,6 +14,7 @@ from omnimalloc.common.validation import ensure_non_negative from .allocation import Allocation, IdType +from .utils import ensure_allocations @dataclass(frozen=True) @@ -29,13 +31,14 @@ def __post_init__(self) -> None: if self.offset is not None: ensure_non_negative(self.offset, "offset") + @classmethod + def from_allocations(cls, allocations: Sequence[Allocation]) -> "Pool": + """Wrap a raw sequence of allocations in an anonymous pool.""" + return cls(id=0, allocations=ensure_allocations(allocations)) + @cached_property def size(self) -> int: - """Memory extent from the pool base (offset 0) to the highest allocated end. - - Counts any gap below the lowest allocation, so it is consistent with - how `overlaps`, pool stacking, and visualization interpret pool extent. - """ + """Memory extent from the pool base (offset 0) to the highest allocated end.""" if not self.is_allocated: raise ValueError("cannot compute size of unallocated pool") ends = [ @@ -47,11 +50,7 @@ def size(self) -> int: @cached_property def pressure(self) -> int: - """Peak memory pressure (max cut through all buffer lifetimes). - - Exact: scalar sweep for linearizable lifetimes, max-weight antichain - for non-linearizable vector-clock instances (see `analysis.pressure`). - """ + """Peak memory pressure (max cut through all buffer lifetimes).""" return _pressure(self.allocations) @cached_property @@ -68,6 +67,11 @@ def is_allocated(self) -> bool: """True if all allocations have been assigned memory offsets.""" return all(alloc.offset is not None for alloc in self.allocations) + @cached_property + def any_allocated(self) -> bool: + """True if any allocation has been assigned a memory offset.""" + return any(alloc.offset is not None for alloc in self.allocations) + def overlaps(self, other: "Pool") -> bool: """True if pools overlap in memory space.""" if self.offset is None or other.offset is None: @@ -82,10 +86,17 @@ def with_allocations(self, allocations: tuple[Allocation, ...]) -> "Pool": return Pool(id=self.id, offset=self.offset, allocations=allocations) def allocate(self, allocator: "BaseAllocator") -> "Pool": - """Apply allocator to assign memory offsets to all allocations.""" + """Apply allocator to assign memory offsets, preserving allocation order.""" allocated = allocator.allocate(self.allocations) - if {a.id for a in allocated} != {a.id for a in self.allocations}: + if len(allocated) != len(self.allocations) or {a.id for a in allocated} != { + a.id for a in self.allocations + }: raise ValueError( f"allocator {allocator!s} returned a different allocation set" ) - return self.with_allocations(allocated) + # Allocators may reorder internally; restore the pool's input order so + # positions in the returned tuple keep corresponding to the request + placed_by_id = {a.id: a for a in allocated} + return self.with_allocations( + tuple(placed_by_id[a.id] for a in self.allocations) + ) diff --git a/src/python/omnimalloc/primitives/system.py b/src/python/omnimalloc/primitives/system.py index 7cd704a..2439d72 100644 --- a/src/python/omnimalloc/primitives/system.py +++ b/src/python/omnimalloc/primitives/system.py @@ -29,6 +29,11 @@ def is_allocated(self) -> bool: """True if all memories have been allocated.""" return all(memory.is_allocated for memory in self.memories) + @property + def any_allocated(self) -> bool: + """True if any memory has a placed allocation.""" + return any(memory.any_allocated for memory in self.memories) + def with_memories(self, memories: tuple[Memory, ...]) -> "System": """Return new System with specified memories.""" return System(id=self.id, memories=memories) diff --git a/src/python/omnimalloc/primitives/utils.py b/src/python/omnimalloc/primitives/utils.py index cca746e..d116ae2 100644 --- a/src/python/omnimalloc/primitives/utils.py +++ b/src/python/omnimalloc/primitives/utils.py @@ -11,3 +11,15 @@ def ensure_unique_ids(allocations: Sequence[Allocation]) -> None: """Raise if any allocation id repeats; id-keyed placement assumes uniqueness.""" if len({alloc.id for alloc in allocations}) != len(allocations): raise ValueError("allocation ids must be unique") + + +def ensure_allocations(allocations: object) -> tuple[Allocation, ...]: + """Coerce a raw sequence to a tuple, requiring every element be an Allocation.""" + if isinstance(allocations, str | bytes) or not isinstance(allocations, Sequence): + raise TypeError(f"Unsupported entity type: {type(allocations)!r}") + checked: list[Allocation] = [] + for alloc in allocations: + if not isinstance(alloc, Allocation): + raise TypeError(f"Expected Allocation, got {type(alloc)!r}") + checked.append(alloc) + return tuple(checked) diff --git a/src/python/omnimalloc/validate.py b/src/python/omnimalloc/validate.py index 1e66b6f..92588ee 100644 --- a/src/python/omnimalloc/validate.py +++ b/src/python/omnimalloc/validate.py @@ -2,8 +2,11 @@ # SPDX-License-Identifier: Apache-2.0 # +from collections.abc import Sequence + from .analysis.clock import uniform_dim from .primitives import Allocation, IdType, Memory, Pool, System +from .primitives.utils import ensure_allocations def _check_unique_ids(entities: tuple[Memory | Pool | Allocation, ...]) -> None: @@ -51,12 +54,12 @@ def _validate_pools(pools: tuple[Pool, ...]) -> None: raise ValueError(f"in pool {pool.id!r}, {e}") from e -def _check_capacity(memory: Memory) -> None: - if memory.capacity is None or not memory.is_allocated: +def _check_size(memory: Memory) -> None: + if memory.size is None or not memory.is_allocated: return - if memory.used_size > memory.capacity: + if memory.used_size > memory.size: raise ValueError( - f"used size {memory.used_size} exceeds capacity {memory.capacity}" + f"used size {memory.used_size} exceeds memory size {memory.size}" ) @@ -65,29 +68,35 @@ def _validate_memories(memories: tuple[Memory, ...]) -> None: for memory in memories: try: _validate_pools(memory.pools) - _check_capacity(memory) + _check_size(memory) except ValueError as e: raise ValueError(f"in memory {memory.id!r}, {e}") from e -def validate_allocation(entity: System | Memory | Pool) -> None: +def validate_allocation(entity: System | Memory | Pool | Sequence[Allocation]) -> None: """Raise ValueError unless the entity is fully placed with no collisions. + Accepts a System, Memory, or Pool, or a raw sequence of Allocations. Checks unique ids, that every allocation and pool carries an offset, that no placed rectangles collide, and that each memory's used size - stays within its declared capacity. + stays within its declared size. """ + if isinstance(entity, System | Memory | Pool): + described = f"{type(entity).__name__} {entity.id!r}" + elif isinstance(entity, Sequence) and not isinstance(entity, str | bytes): + described = f"{len(entity)} allocations" + else: + raise TypeError(f"Unsupported entity type: {type(entity)!r}") + try: if isinstance(entity, System): _validate_memories(entity.memories) elif isinstance(entity, Memory): _validate_pools(entity.pools) - _check_capacity(entity) + _check_size(entity) elif isinstance(entity, Pool): _validate_allocations(entity.allocations) else: - raise TypeError(f"Unsupported entity type: {type(entity)!r}") + _validate_allocations(ensure_allocations(entity)) except ValueError as e: - raise ValueError( - f"Validation of {type(entity).__name__} {entity.id!r} failed, {e}." - ) from e + raise ValueError(f"Validation of {described} failed, {e}.") from e diff --git a/src/python/omnimalloc/visualize.py b/src/python/omnimalloc/visualize.py index 69b8b5f..9d88540 100644 --- a/src/python/omnimalloc/visualize.py +++ b/src/python/omnimalloc/visualize.py @@ -4,6 +4,7 @@ import logging from collections import defaultdict +from collections.abc import Sequence from pathlib import Path from typing import Final, Literal, NamedTuple @@ -219,23 +220,23 @@ def _select_lanes( def _get_y_limits(system: System) -> dict[Memory, tuple[int, int]]: limits: dict[Memory, tuple[int, int]] = {} for memory in system.memories: - capacity = memory.capacity + size = memory.size used = memory.used_size - if capacity is None: - # No capacity declared, scale to 1.2x used + if size is None: + # No size declared, scale to 1.2x used y_limit = used * 1.2 - elif used > capacity: - # Usage exceeds capacity, scale to 1.2x usage + elif used > size: + # Usage exceeds the declared size, scale to 1.2x usage y_limit = used * 1.2 - elif used >= capacity * 0.5: - # Usage is 50-100% of capacity, use capacity as limit - y_limit = capacity + elif used >= size * 0.5: + # Usage is 50-100% of the declared size, use the size as limit + y_limit = size else: - # Usage below 50% of capacity, scale to 2x usage + # Usage below 50% of the declared size, scale to 2x usage y_limit = used * 2 # Clamp to at least 1 so downstream tick spacing stays positive @@ -307,7 +308,7 @@ def _draw_pool_background( def _draw_limit_lines(ax: Axes, limits: dict[str, int]) -> None: - """Draw annotated horizontal lines for used size, capacity, and extras.""" + """Draw annotated horizontal lines for used size, declared size, and extras.""" _, x_max = ax.get_xlim() for name, value in limits.items(): ax.axhline(value, color="black", linestyle="--", linewidth=1, alpha=0.8) @@ -343,11 +344,9 @@ def _set_axes_ticks(ax: Axes, y_limit: int, num_ticks: int = 8) -> None: def _memory_title(memory: Memory, threads: str) -> str: - capacity = memory.capacity - capacity_str = ( - _format_bytes(capacity) if capacity is not None else "Unbounded Capacity" - ) - return f"{memory.id} ({capacity_str}, {len(memory.pools)} pools{threads})" + size = memory.size + size_str = _format_bytes(size) if size is not None else "Unbounded Size" + return f"{memory.id} ({size_str}, {len(memory.pools)} pools{threads})" def _lane_panels( @@ -412,16 +411,31 @@ def _projection_panels(system: System) -> tuple[list[_Panel], str | None]: return panels, caveat +def _add_legend(fig: Figure) -> None: + """Add figure legend for allocation kinds.""" + handles = [ + Patch(color=color, label=kind.name, alpha=0.8) + for kind, color in KIND_COLOR_MAP.items() + ] + fig.legend( + handles=handles, + loc="outside lower center", + ncol=len(handles), + fontsize=8, + title="Allocation Kinds", + ) + + def _set_axes_limits( ax: Axes, x_limits: tuple[int, int], y_limits: tuple[int, int], - capacity: int | None, + size: int | None, ) -> None: """Set axis limits and add scaling notice if needed.""" ax.set_xlim(x_limits) ax.set_ylim(y_limits) - if capacity is not None and y_limits[1] < capacity: + if size is not None and y_limits[1] < size: ax.text( 0.02, 0.98, @@ -433,21 +447,6 @@ def _set_axes_limits( ) -def _add_legend(fig: Figure) -> None: - """Add figure legend for allocation kinds.""" - handles = [ - Patch(color=color, label=kind.name, alpha=0.8) - for kind, color in KIND_COLOR_MAP.items() - ] - fig.legend( - handles=handles, - loc="outside lower center", - ncol=len(handles), - fontsize=8, - title="Allocation Kinds", - ) - - # TODO(fpedd): Add a pools size descriptor on the right side of each pool @@ -462,7 +461,7 @@ def _draw_panel( if panel.title is not None: ax.set_title(panel.title) ax.set_xlabel(panel.xlabel) - _set_axes_limits(ax, panel.x_limits, y_limits, memory.capacity) + _set_axes_limits(ax, panel.x_limits, y_limits, memory.size) _set_axes_ticks(ax, y_limits[1]) if panel.note is not None: ax.text( @@ -488,10 +487,10 @@ def _draw_panel( _draw_allocation(ax, alloc, y_offset, color, extent) _draw_pool_background(ax, y_offset, pool.size, colors) - # Draw used-size, capacity, and extra capacity lines + # Draw used-size, declared-size, and extra capacity lines limits: dict[str, int] = {"used": memory.used_size} - if memory.capacity is not None: - limits["capacity"] = memory.capacity + if memory.size is not None: + limits["size"] = memory.size for label, per_memory_capacity in capacities.items(): if memory.id in per_memory_capacity: limits[label] = per_memory_capacity[memory.id] @@ -551,9 +550,8 @@ def _visualize_system( def plot_allocation( - entity: System | Memory | Pool, + entity: System | Memory | Pool | Sequence[Allocation], path: Path | str | None = None, - *, capacities: dict[str, dict[IdType, int]] | None = None, view: Literal["panel", "lanes"] = "panel", max_lanes: int | None = None, @@ -561,7 +559,8 @@ def plot_allocation( """Plot an allocated entity: `path=None` displays the figure, `path=...` saves it. Args: - entity: The entity to plot (System, Memory, or Pool). + entity: The entity to plot (System, Memory, Pool, or a raw sequence + of Allocations, plotted as a single pool). path: Where to save the figure; `None` displays it instead. capacities: Extra capacity lines to draw, keyed by label then by memory id (byte limits, drawn as horizontal lines). @@ -588,6 +587,9 @@ def plot_allocation( if not HAS_MATPLOTLIB: require_optional("matplotlib", "visualization") + if not isinstance(entity, System | Memory | Pool): + entity = Pool.from_allocations(entity) + if isinstance(entity, Pool): entity = Memory(id=f"memory_{entity.id}", pools=(entity,)) diff --git a/tests/integration/test_omni_torture.py b/tests/integration/test_omni_torture.py index 70762ef..e4d2bba 100644 --- a/tests/integration/test_omni_torture.py +++ b/tests/integration/test_omni_torture.py @@ -61,7 +61,9 @@ def _certify( peak = _peak(placed) assert pressure(allocations, work_budget=None) <= peak assert peak <= sum(a.size for a in allocations) - assert max(placement_pressure_per_allocation(placed).values()) == peak + per_allocation = placement_pressure_per_allocation(placed) + assert per_allocation == placement_pressure_per_allocation(placed, None) + assert max(per_allocation.values()) == peak _assert_conflicting_pairs_disjoint(placed) return peak diff --git a/tests/unit/analysis/test_pressure.py b/tests/unit/analysis/test_pressure.py index e2b1afd..43f6250 100644 --- a/tests/unit/analysis/test_pressure.py +++ b/tests/unit/analysis/test_pressure.py @@ -8,6 +8,8 @@ import pytest from omnimalloc.allocators.omni import OmniAllocator from omnimalloc.analysis import ( + antichain_pressure, + antichain_pressure_per_allocation, closure_pressure, closure_pressure_per_allocation, placement_pressure, @@ -18,6 +20,11 @@ from omnimalloc.primitives import Allocation +def test_default_names_alias_antichain() -> None: + assert pressure is antichain_pressure + assert pressure_per_allocation is antichain_pressure_per_allocation + + def test_pressure_empty_is_zero() -> None: assert pressure(()) == 0 diff --git a/tests/unit/primitives/test_memory.py b/tests/unit/primitives/test_memory.py index 35d392f..0a888a4 100644 --- a/tests/unit/primitives/test_memory.py +++ b/tests/unit/primitives/test_memory.py @@ -16,7 +16,7 @@ def test_basic_creation_with_int_id_simple() -> None: memory = Memory(id=301, pools=(pool,)) assert memory.id == 301 assert len(memory.pools) == 1 - assert memory.capacity is None + assert memory.size is None def test_basic_creation_with_int_id() -> None: @@ -34,15 +34,15 @@ def test_basic_creation_with_str_id() -> None: memory = Memory(id="mem_ddr", pools=(pool,)) assert memory.id == "mem_ddr" assert len(memory.pools) == 1 - assert memory.capacity is None + assert memory.size is None -def test_creation_with_capacity() -> None: - """Test creating a memory with declared capacity.""" +def test_creation_with_size() -> None: + """Test creating a memory with declared size.""" alloc = Allocation(id=101, size=100, start=0, end=10, offset=0) pool = Pool(id=211, allocations=(alloc,)) - memory = Memory(id=301, pools=(pool,), capacity=1000) - assert memory.capacity == 1000 + memory = Memory(id=301, pools=(pool,), size=1000) + assert memory.size == 1000 def test_creation_with_multiple_pools() -> None: @@ -63,20 +63,20 @@ def test_empty_memory() -> None: assert memory.is_allocated is True -def test_negative_capacity() -> None: - """Test that negative capacity raises ValueError.""" +def test_negative_size() -> None: + """Test that negative size raises ValueError.""" alloc = Allocation(id=101, size=100, start=0, end=10, offset=0) pool = Pool(id=211, allocations=(alloc,)) - with pytest.raises(ValueError, match="capacity must be non-negative"): - Memory(id=301, pools=(pool,), capacity=-1) + with pytest.raises(ValueError, match="size must be non-negative"): + Memory(id=301, pools=(pool,), size=-1) -def test_zero_capacity() -> None: - """Test that zero capacity is valid.""" +def test_zero_size() -> None: + """Test that zero size is valid.""" alloc = Allocation(id=101, size=100, start=0, end=10, offset=0) pool = Pool(id=211, allocations=(alloc,)) - memory = Memory(id=301, pools=(pool,), capacity=0) - assert memory.capacity == 0 + memory = Memory(id=301, pools=(pool,), size=0) + assert memory.size == 0 def test_duplicate_pool_ids() -> None: @@ -113,11 +113,11 @@ def test_used_size_empty_memory() -> None: assert memory.used_size == 0 -def test_used_size_can_exceed_capacity() -> None: - """Capacity is a declared limit, not a constructor constraint.""" +def test_used_size_can_exceed_size() -> None: + """Size is a declared limit, not a constructor constraint.""" alloc = Allocation(id=101, size=100, start=0, end=10, offset=0) pool = Pool(id=211, allocations=(alloc,)) - memory = Memory(id=301, pools=(pool,), capacity=10) + memory = Memory(id=301, pools=(pool,), size=10) assert memory.used_size == 100 @@ -163,12 +163,12 @@ def test_with_pools_replace() -> None: alloc2 = Allocation(id=102, size=200, start=0, end=10, offset=0) pool1 = Pool(id=211, allocations=(alloc1,)) pool2 = Pool(id=212, allocations=(alloc2,)) - memory = Memory(id=301, pools=(pool1,), capacity=1000) + memory = Memory(id=301, pools=(pool1,), size=1000) new_memory = memory.with_pools((pool2,)) assert len(new_memory.pools) == 1 assert new_memory.pools[0].id == 212 assert new_memory.id == memory.id - assert new_memory.capacity == memory.capacity + assert new_memory.size == memory.size def test_with_pools_immutability() -> None: @@ -211,11 +211,11 @@ def test_cannot_modify_pools() -> None: memory.pools = () # type: ignore[misc] -def test_cannot_modify_capacity() -> None: - """Test that capacity cannot be modified.""" - memory = Memory(id=301, pools=(), capacity=1000) +def test_cannot_modify_size() -> None: + """Test that size cannot be modified.""" + memory = Memory(id=301, pools=(), size=1000) with pytest.raises(AttributeError): - memory.capacity = 2000 # type: ignore[misc] + memory.size = 2000 # type: ignore[misc] def test_large_values() -> None: @@ -224,9 +224,25 @@ def test_large_values() -> None: alloc2 = Allocation(id=102, size=10**11, start=0, end=100, offset=0) pool1 = Pool(id=211, allocations=(alloc1,)) pool2 = Pool(id=212, allocations=(alloc2,)) - memory = Memory(id=999, pools=(pool1, pool2), capacity=10**15) + memory = Memory(id=999, pools=(pool1, pool2), size=10**15) assert memory.used_size == 10**12 + 10**11 - assert memory.capacity == 10**15 + assert memory.size == 10**15 + + +def test_any_allocated_across_pools() -> None: + placed = Pool( + id=1, allocations=(Allocation(id=1, size=10, start=0, end=5, offset=0),) + ) + unplaced = Pool(id=2, allocations=(Allocation(id=2, size=10, start=0, end=5),)) + memory = Memory(id=1, pools=(placed, unplaced)) + assert memory.any_allocated is True + assert memory.is_allocated is False + + +def test_any_allocated_empty_memory() -> None: + memory = Memory(id=1, pools=()) + assert memory.any_allocated is False + assert memory.is_allocated is True def test_complex_memory_structure() -> None: @@ -236,7 +252,7 @@ def test_complex_memory_structure() -> None: alloc3 = Allocation(id=103, size=75, start=10, end=20, offset=0) pool1 = Pool(id=211, allocations=(alloc1, alloc2)) pool2 = Pool(id=212, allocations=(alloc3,)) - memory = Memory(id=400, pools=(pool1, pool2), capacity=500) + memory = Memory(id=400, pools=(pool1, pool2), size=500) assert memory.used_size == pool1.size + pool2.size assert memory.is_allocated is True @@ -247,7 +263,7 @@ def test_allocate_with_allocator() -> None: alloc2 = Allocation(id=102, size=50, start=5, end=15) pool1 = Pool(id=211, allocations=(alloc1,)) pool2 = Pool(id=212, allocations=(alloc2,)) - memory = Memory(id=301, pools=(pool1, pool2), capacity=1000) + memory = Memory(id=301, pools=(pool1, pool2), size=1000) assert memory.is_allocated is False allocator = NaiveAllocator() @@ -255,7 +271,7 @@ def test_allocate_with_allocator() -> None: assert allocated_memory.is_allocated is True assert allocated_memory.id == memory.id - assert allocated_memory.capacity == memory.capacity + assert allocated_memory.size == memory.size assert len(allocated_memory.pools) == 2 # Original memory should be unchanged assert memory.is_allocated is False diff --git a/tests/unit/primitives/test_pool.py b/tests/unit/primitives/test_pool.py index e8c7f39..1a74da6 100644 --- a/tests/unit/primitives/test_pool.py +++ b/tests/unit/primitives/test_pool.py @@ -3,6 +3,7 @@ # import pytest +from omnimalloc.allocators.greedy import GreedyBySizeAllocator from omnimalloc.allocators.naive import NaiveAllocator from omnimalloc.primitives import Allocation from omnimalloc.primitives.pool import Pool @@ -385,3 +386,51 @@ def test_size_counts_gap_below_lowest_allocation() -> None: alloc = Allocation(id=1, size=100, start=0, end=10, offset=1000) pool = Pool(id="p", allocations=(alloc,)) assert pool.size == 1100 + + +def test_any_allocated_empty_pool() -> None: + pool = Pool(id=1, allocations=()) + assert pool.any_allocated is False + assert pool.is_allocated is True + + +def test_any_allocated_none_placed() -> None: + pool = Pool(id=1, allocations=(Allocation(id=1, size=10, start=0, end=5),)) + assert pool.any_allocated is False + + +def test_any_allocated_partially_placed() -> None: + alloc1 = Allocation(id=1, size=10, start=0, end=5, offset=0) + alloc2 = Allocation(id=2, size=10, start=0, end=5) + pool = Pool(id=1, allocations=(alloc1, alloc2)) + assert pool.any_allocated is True + assert pool.is_allocated is False + + +def test_from_allocations_wraps_sequence() -> None: + allocations = [Allocation(id=1, size=10, start=0, end=5)] + pool = Pool.from_allocations(allocations) + assert pool.allocations == tuple(allocations) + + +def test_from_allocations_rejects_non_allocation_elements() -> None: + with pytest.raises(TypeError, match="Expected Allocation"): + Pool.from_allocations((1, 2)) + + +def test_from_allocations_rejects_non_sequence() -> None: + with pytest.raises(TypeError, match="Unsupported entity type"): + Pool.from_allocations("abc") + + +def test_allocate_preserves_allocation_order() -> None: + allocations = tuple( + Allocation(id=i, size=10 * (i + 1), start=0, end=5) for i in range(8) + ) + pool = Pool(id=1, allocations=allocations) + allocated = pool.allocate(GreedyBySizeAllocator()) + assert allocated.is_allocated is True + assert tuple(a.id for a in allocated.allocations) == tuple(range(8)) + assert tuple(a.size for a in allocated.allocations) == tuple( + a.size for a in allocations + ) diff --git a/tests/unit/primitives/test_system.py b/tests/unit/primitives/test_system.py index b929482..f468fd0 100644 --- a/tests/unit/primitives/test_system.py +++ b/tests/unit/primitives/test_system.py @@ -168,8 +168,8 @@ def test_complex_system_structure() -> None: pool1 = Pool(id=211, allocations=(alloc1, alloc2)) pool2 = Pool(id=212, allocations=(alloc3,)) pool3 = Pool(id=213, allocations=(alloc4,)) - memory1 = Memory(id=311, pools=(pool1, pool2), capacity=500) - memory2 = Memory(id=312, pools=(pool3,), capacity=300) + memory1 = Memory(id=311, pools=(pool1, pool2), size=500) + memory2 = Memory(id=312, pools=(pool3,), size=300) system = System(id=401, memories=(memory1, memory2)) assert len(system.memories) == 2 assert system.is_allocated is True @@ -193,7 +193,7 @@ def test_single_memory_system() -> None: """Test system with single memory.""" alloc = Allocation(id=101, size=100, start=0, end=10, offset=0) pool = Pool(id=211, allocations=(alloc,)) - memory = Memory(id=311, pools=(pool,), capacity=1000) + memory = Memory(id=311, pools=(pool,), size=1000) system = System(id=401, memories=(memory,)) assert len(system.memories) == 1 assert system.is_allocated is True @@ -215,14 +215,32 @@ def test_multiple_pools_per_memory() -> None: assert system.is_allocated is True +def test_any_allocated_across_memories() -> None: + alloc = Allocation(id=1, size=10, start=0, end=5, offset=0) + placed = Memory(id=1, pools=(Pool(id=1, allocations=(alloc,)),)) + unplaced = Memory( + id=2, + pools=(Pool(id=2, allocations=(Allocation(id=2, size=10, start=0, end=5),)),), + ) + system = System(id=1, memories=(placed, unplaced)) + assert system.any_allocated is True + assert system.is_allocated is False + + +def test_any_allocated_empty_system() -> None: + system = System(id=1, memories=()) + assert system.any_allocated is False + assert system.is_allocated is True + + def test_large_system() -> None: """Test system with large values.""" alloc1 = Allocation(id=101, size=10**12, start=0, end=100, offset=0) alloc2 = Allocation(id=102, size=10**11, start=0, end=100, offset=0) pool1 = Pool(id=211, allocations=(alloc1,)) pool2 = Pool(id=212, allocations=(alloc2,)) - memory1 = Memory(id=311, pools=(pool1,), capacity=10**15) - memory2 = Memory(id=312, pools=(pool2,), capacity=10**14) + memory1 = Memory(id=311, pools=(pool1,), size=10**15) + memory2 = Memory(id=312, pools=(pool2,), size=10**14) system = System(id=999, memories=(memory1, memory2)) assert memory1.used_size == 10**12 assert memory2.used_size == 10**11 @@ -237,8 +255,8 @@ def test_allocate_with_allocator() -> None: alloc3 = Allocation(id=103, size=75, start=10, end=20) pool1 = Pool(id=211, allocations=(alloc1,)) pool2 = Pool(id=212, allocations=(alloc2, alloc3)) - memory1 = Memory(id=311, pools=(pool1,), capacity=1000) - memory2 = Memory(id=312, pools=(pool2,), capacity=500) + memory1 = Memory(id=311, pools=(pool1,), size=1000) + memory2 = Memory(id=312, pools=(pool2,), size=500) system = System(id=401, memories=(memory1, memory2)) assert system.is_allocated is False diff --git a/tests/unit/test_allocate.py b/tests/unit/test_allocate.py index da2849d..dbdea6e 100644 --- a/tests/unit/test_allocate.py +++ b/tests/unit/test_allocate.py @@ -2,8 +2,9 @@ # SPDX-License-Identifier: Apache-2.0 # +import pytest from omnimalloc import allocate -from omnimalloc.allocators.greedy import GreedyAllocator +from omnimalloc.allocators.greedy import GreedyAllocator, GreedyBySizeAllocator from omnimalloc.allocators.naive import NaiveAllocator from omnimalloc.primitives import Allocation, AllocationKind, Memory, Pool, System @@ -140,14 +141,14 @@ def test_allocate_complex_hierarchy() -> None: pool1 = Pool(id=1, allocations=tuple(allocations_pool1)) pool2 = Pool(id=2, allocations=tuple(allocations_pool2)) - memory1 = Memory(id=1, pools=(pool1, pool2), capacity=2048) + memory1 = Memory(id=1, pools=(pool1, pool2), size=2048) # Second memory allocations_pool3 = [ Allocation(id=i + 10, size=75, start=i * 5, end=(i + 1) * 5) for i in range(4) ] pool3 = Pool(id=3, allocations=tuple(allocations_pool3)) - memory2 = Memory(id=2, pools=(pool3,), capacity=1024) + memory2 = Memory(id=2, pools=(pool3,), size=1024) system = System(id=1, memories=(memory1, memory2)) @@ -246,10 +247,70 @@ def test_allocate_memory_calculates_used_size() -> None: alloc2 = Allocation(id=2, size=150, start=0, end=10) pool1 = Pool(id=1, allocations=(alloc1,)) pool2 = Pool(id=2, allocations=(alloc2,)) - memory = Memory(id=1, pools=(pool1, pool2), capacity=1000) + memory = Memory(id=1, pools=(pool1, pool2), size=1000) allocator = NaiveAllocator() allocated_memory = allocate(memory, allocator) # Each pool gets its allocations placed assert allocated_memory.used_size == 250 # 100 + 150 + + +def test_allocate_raw_allocations_returns_tuple() -> None: + allocations = ( + Allocation(id=1, size=100, start=0, end=5), + Allocation(id=2, size=150, start=5, end=10), + ) + allocated = allocate(allocations, NaiveAllocator(), validate=True) + assert isinstance(allocated, tuple) + assert all(a.offset is not None for a in allocated) + assert {a.id for a in allocated} == {1, 2} + + +def test_allocate_raw_allocations_accepts_list_and_registry_name() -> None: + allocations = [ + Allocation(id="a", size=100, start=0, end=5), + Allocation(id="b", size=150, start=0, end=5), + ] + allocated = allocate(allocations, "naive") + assert isinstance(allocated, tuple) + assert all(a.is_allocated for a in allocated) + + +def test_allocate_raw_allocations_does_not_mutate_input() -> None: + allocations = (Allocation(id=1, size=100, start=0, end=5),) + allocate(allocations, NaiveAllocator()) + assert allocations[0].offset is None + + +def test_allocate_raw_allocations_duplicate_ids_raise() -> None: + allocations = ( + Allocation(id=1, size=100, start=0, end=5), + Allocation(id=1, size=150, start=5, end=10), + ) + with pytest.raises(ValueError, match="allocation ids must be unique"): + allocate(allocations, NaiveAllocator()) + + +def test_allocate_raw_allocations_empty() -> None: + assert allocate((), NaiveAllocator()) == () + + +def test_allocate_raw_allocations_rejects_non_allocation_elements() -> None: + with pytest.raises(TypeError, match="Expected Allocation"): + allocate((1, 2, 3), NaiveAllocator()) + + +def test_allocate_rejects_unsupported_entity() -> None: + with pytest.raises(TypeError, match="Unsupported entity type"): + allocate("naive") + + +def test_allocate_raw_allocations_preserves_input_order() -> None: + allocations = tuple( + Allocation(id=i, size=10 * (i + 1), start=0, end=5) for i in range(10) + ) + allocated = allocate(allocations, GreedyBySizeAllocator(), validate=True) + assert tuple(a.id for a in allocated) == tuple(a.id for a in allocations) + assert all(a.offset is not None for a in allocated) + assert tuple(a.size for a in allocated) == tuple(a.size for a in allocations) diff --git a/tests/unit/test_io.py b/tests/unit/test_io.py index 8f77ef5..1335800 100644 --- a/tests/unit/test_io.py +++ b/tests/unit/test_io.py @@ -32,6 +32,20 @@ def test_save_pool_writes_exactly_path(tmp_path: Path) -> None: assert file_path.read_text() == "id,lower,upper,size\na,0,3,4\nb,2,9,8\n" +def test_save_raw_allocations_writes_one_pool(tmp_path: Path) -> None: + file_path = tmp_path / "problem.csv" + written = save_allocation(list(make_pool().allocations), file_path) + assert written == (file_path,) + assert file_path.read_text() == "id,lower,upper,size\na,0,3,4\nb,2,9,8\n" + + +def test_save_raw_allocations_rejects_non_allocation_elements( + tmp_path: Path, +) -> None: + with pytest.raises(TypeError, match="Expected Allocation"): + save_allocation((1, 2), tmp_path / "problem.csv") + + def test_save_creates_missing_parent_directories(tmp_path: Path) -> None: file_path = tmp_path / "nested" / "dir" / "problem.csv" written = save_allocation(make_pool(), file_path) diff --git a/tests/unit/test_validate.py b/tests/unit/test_validate.py index da7d05f..fd7e921 100644 --- a/tests/unit/test_validate.py +++ b/tests/unit/test_validate.py @@ -202,20 +202,20 @@ def test_validate_complex_hierarchy_with_error_raises() -> None: validate_allocation(system) -def test_validate_memory_over_capacity_fails() -> None: +def test_validate_memory_over_size_fails() -> None: pool = Pool( id="p", allocations=(Allocation(id=1, size=100, start=0, end=5, offset=0),) ) - memory = Memory(id="m", capacity=50, pools=(pool,)) - with pytest.raises(ValueError, match="exceeds capacity"): + memory = Memory(id="m", size=50, pools=(pool,)) + with pytest.raises(ValueError, match="exceeds memory size"): validate_allocation(memory) -def test_validate_memory_within_capacity_passes() -> None: +def test_validate_memory_within_size_passes() -> None: pool = Pool( id="p", allocations=(Allocation(id=1, size=100, start=0, end=5, offset=0),) ) - memory = Memory(id="m", capacity=100, pools=(pool,)) + memory = Memory(id="m", size=100, pools=(pool,)) validate_allocation(memory) @@ -240,3 +240,58 @@ def test_validate_pool_mixed_dimensions_rejected() -> None: pool = Pool(id=1, allocations=(alloc1, alloc2)) with pytest.raises(ValueError, match="share one clock dimension"): validate_allocation(pool) + + +def test_validate_raw_allocations_valid() -> None: + allocations = ( + Allocation(id=1, size=100, start=0, end=5, offset=0), + Allocation(id=2, size=100, start=5, end=10, offset=0), + ) + validate_allocation(allocations) + + +def test_validate_raw_allocations_list() -> None: + validate_allocation([Allocation(id=1, size=10, start=0, end=5, offset=0)]) + + +def test_validate_raw_allocations_empty() -> None: + validate_allocation(()) + + +def test_validate_raw_allocations_overlapping_raises() -> None: + allocations = ( + Allocation(id=1, size=100, start=0, end=10, offset=0), + Allocation(id=2, size=100, start=5, end=15, offset=50), + ) + with pytest.raises(ValueError, match=r"Validation of 2 allocations failed"): + validate_allocation(allocations) + + +def test_validate_raw_allocations_unallocated_raises() -> None: + with pytest.raises(ValueError, match="is not allocated"): + validate_allocation((Allocation(id=1, size=10, start=0, end=5),)) + + +def test_validate_raw_allocations_duplicate_ids_raise() -> None: + allocations = ( + Allocation(id=1, size=10, start=0, end=5, offset=0), + Allocation(id=1, size=10, start=5, end=10, offset=0), + ) + with pytest.raises(ValueError, match="duplicate id"): + validate_allocation(allocations) + + +def test_validate_raw_allocations_rejects_non_allocation_elements() -> None: + with pytest.raises(TypeError, match="Expected Allocation"): + validate_allocation([1, 2, 3]) + + +def test_validate_raw_allocations_complex_mix() -> None: + valid = tuple( + Allocation(id=i, size=50, start=5 * i, end=5 * (i + 1), offset=0) + for i in range(10) + ) + validate_allocation(valid) + overlapping = (*valid, Allocation(id=99, size=50, start=12, end=18, offset=25)) + with pytest.raises(ValueError, match=r"Validation of 11 allocations failed"): + validate_allocation(overlapping) diff --git a/tests/unit/test_visualize.py b/tests/unit/test_visualize.py index 6819ee8..216f2d4 100644 --- a/tests/unit/test_visualize.py +++ b/tests/unit/test_visualize.py @@ -77,7 +77,7 @@ def test_visualize_memory_with_multiple_pools(artifacts_dir: Path) -> None: pool2 = Pool(id=2, allocations=(alloc2,), offset=200) pool3 = Pool(id=3, allocations=(alloc3,), offset=500) - memory = Memory(id="mem_1", pools=(pool1, pool2, pool3), capacity=1000) + memory = Memory(id="mem_1", pools=(pool1, pool2, pool3), size=1000) output_path = artifacts_dir / "test_memory_pools.pdf" plot_allocation(memory, output_path) @@ -89,14 +89,14 @@ def test_visualize_system_with_multiple_memories(artifacts_dir: Path) -> None: # Memory 1: Simple pool alloc1 = Allocation(id=1, size=100, start=0, end=5, offset=0) pool1 = Pool(id=1, allocations=(alloc1,), offset=0) - memory1 = Memory(id="ddr4_1", pools=(pool1,), capacity=500) + memory1 = Memory(id="ddr4_1", pools=(pool1,), size=500) # Memory 2: Multiple pools alloc2 = Allocation(id=2, size=150, start=0, end=10, offset=0) alloc3 = Allocation(id=3, size=75, start=5, end=15, offset=0) pool2 = Pool(id=2, allocations=(alloc2,), offset=0) pool3 = Pool(id=3, allocations=(alloc3,), offset=200) - memory2 = Memory(id="ddr4_2", pools=(pool2, pool3), capacity=1000) + memory2 = Memory(id="ddr4_2", pools=(pool2, pool3), size=1000) system = System(id="test_system", memories=(memory1, memory2)) @@ -127,7 +127,7 @@ def test_visualize_complex_hierarchy(artifacts_dir: Path) -> None: pool1 = Pool(id="pool_1", allocations=tuple(allocations_mem1_pool1), offset=0) pool2 = Pool(id="pool_2", allocations=tuple(allocations_mem1_pool2), offset=300) - memory1 = Memory(id="main_memory", pools=(pool1, pool2), capacity=2048) + memory1 = Memory(id="main_memory", pools=(pool1, pool2), size=2048) # Second memory with different allocation patterns allocations_mem2 = [ @@ -142,7 +142,7 @@ def test_visualize_complex_hierarchy(artifacts_dir: Path) -> None: for i in range(0, 20, 5) ] pool3 = Pool(id="pool_3", allocations=tuple(allocations_mem2), offset=0) - memory2 = Memory(id="cache_memory", pools=(pool3,), capacity=1024) + memory2 = Memory(id="cache_memory", pools=(pool3,), size=1024) system = System(id="complex_system", memories=(memory1, memory2)) @@ -159,7 +159,7 @@ def test_visualize_with_string_ids(artifacts_dir: Path) -> None: alloc1 = Allocation(id="workspace_buf", size=100, start=0, end=5, offset=0) alloc2 = Allocation(id="temp_buf", size=150, start=5, end=10, offset=0) pool = Pool(id="tensor_pool", allocations=(alloc1, alloc2), offset=0) - memory = Memory(id="ddr_ram", pools=(pool,), capacity=512) + memory = Memory(id="ddr_ram", pools=(pool,), size=512) output_path = artifacts_dir / "test_string_ids.pdf" plot_allocation(memory, output_path) @@ -193,19 +193,30 @@ def test_visualize_memory_converts_to_system(artifacts_dir: Path) -> None: """Test that visualizing a Memory creates appropriate System wrapper.""" alloc = Allocation(id=1, size=100, start=0, end=10, offset=0) pool = Pool(id=1, allocations=(alloc,), offset=0) - memory = Memory(id=1, pools=(pool,), capacity=500) + memory = Memory(id=1, pools=(pool,), size=500) output_path = artifacts_dir / "test_memory_wrapper.pdf" plot_allocation(memory, output_path) assert output_path.exists() +def test_visualize_raw_allocations_converts_to_pool(artifacts_dir: Path) -> None: + allocations = ( + Allocation(id=1, size=100, start=0, end=5, offset=0), + Allocation(id=2, size=150, start=5, end=10, offset=0), + ) + + output_path = artifacts_dir / "test_raw_allocations.pdf" + plot_allocation(allocations, output_path) + assert output_path.exists() + + def test_visualize_with_capacities(artifacts_dir: Path) -> None: """Test visualization with extra capacity lines.""" alloc1 = Allocation(id=1, size=100, start=0, end=5, offset=0) alloc2 = Allocation(id=2, size=150, start=5, end=10, offset=100) pool = Pool(id=1, allocations=(alloc1, alloc2), offset=0) - memory = Memory(id="ddr_mem", pools=(pool,), capacity=1000) + memory = Memory(id="ddr_mem", pools=(pool,), size=1000) system = System(id="test_sys", memories=(memory,)) custom_limits = { @@ -283,7 +294,7 @@ def test_visualize_mixed_dimension_pools(artifacts_dir: Path) -> None: allocations=(Allocation(id=2, size=50, start=(0, 1), end=(2, 3), offset=0),), offset=200, ) - memory = Memory(id="mem", pools=(scalar_pool, vector_pool), capacity=1000) + memory = Memory(id="mem", pools=(scalar_pool, vector_pool), size=1000) output_path = artifacts_dir / "test_mixed_dim_pools.pdf" plot_allocation(memory, output_path) @@ -298,7 +309,7 @@ def test_visualize_empty_pool(artifacts_dir: Path) -> None: allocations=(Allocation(id=1, size=100, start=0, end=4, offset=0),), offset=100, ) - memory = Memory(id="mem", pools=(empty, filled), capacity=500) + memory = Memory(id="mem", pools=(empty, filled), size=500) output_path = artifacts_dir / "test_empty_pool.pdf" plot_allocation(memory, output_path) diff --git a/uv.lock b/uv.lock index e264d07..664bcd4 100644 --- a/uv.lock +++ b/uv.lock @@ -1334,6 +1334,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ad/3f/3d42e9a78fe5edf792a83c074b13b9b770092a4fbf3462872f4303135f09/ml_dtypes-0.5.4-cp314-cp314t-win_arm64.whl", hash = "sha256:11942cbf2cf92157db91e5022633c0d9474d4dfd813a909383bd23ce828a4b7d", size = 168825, upload-time = "2025-11-17T22:32:23.766Z" }, ] +[[package]] +name = "nanobind" +version = "2.13.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/07/e5/a6fc2f024eaef17c68b94b5f8f56098d31859a7840a11a3dd45da2d3908a/nanobind-2.13.0.tar.gz", hash = "sha256:c7b04d6a6a4cd57985571e605539399b51331ae455d7fce576a5e2fcb89b1dcf", size = 1052101, upload-time = "2026-06-18T06:03:49.491Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/39/8d/cec69a0a804711fd89f47fad92cd2a438a1ceaafd9dabaf3b8b403b5cbba/nanobind-2.13.0-py3-none-any.whl", hash = "sha256:9268d7201eaf5519ae61b7a83175f2af77088cb4b99dd3ce5b00550f697da7b7", size = 269387, upload-time = "2026-06-18T06:03:47.508Z" }, +] + [[package]] name = "nbclient" version = "0.10.2" @@ -1580,6 +1589,7 @@ all = [ dev = [ { name = "ipykernel" }, { name = "ipywidgets" }, + { name = "nanobind" }, { name = "nbconvert" }, { name = "pre-commit" }, { name = "pytest" }, @@ -1605,6 +1615,7 @@ provides-extras = ["all"] dev = [ { name = "ipykernel", specifier = ">=6.30.0" }, { name = "ipywidgets", specifier = ">=8.1.0" }, + { name = "nanobind", specifier = ">=2.0" }, { name = "nbconvert", specifier = ">=7.16.0" }, { name = "pre-commit", specifier = ">=4.0.0" }, { name = "pytest", specifier = ">=8.0" },