Fix benchmark defects and flatten the benchmark package

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")
 
 

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__":

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__":

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,

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",

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)"
    ]
   }
  ],

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
 

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
 

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),

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