State injection into setup and teardown

src/nnbench/core.py

@@ -10,6 +10,7 @@
 from typing import Any, Callable, Iterable, Union, get_args, get_origin, overload
 
 from nnbench.types import Benchmark
+from nnbench.types.types import NoOp
 from nnbench.types.util import is_memo, is_memo_type
 
 
@@ -52,10 +53,6 @@ def _default_namegen(fn: Callable, **kwargs: Any) -> str:
     return fn.__name__ + "_" + "_".join(f"{k}={v}" for k, v in kwargs.items())
 
 
-def NoOp(**kwargs: Any) -> None:
-    pass
-
-
 # Overloads for the ``benchmark`` decorator.
 # Case #1: Bare application without parentheses
 # @nnbench.benchmark
@@ -178,7 +175,14 @@ def decorator(fn: Callable) -> list[Benchmark]:
                 )
             names.add(name)
 
-            bm = Benchmark(fn, name=name, params=params, setUp=setUp, tearDown=tearDown, tags=tags)
+            bm = Benchmark(
+                fn,
+                name=name,
+                params=params,
+                setUp=setUp,
+                tearDown=tearDown,
+                tags=tags,
+            )
             benchmarks.append(bm)
         return benchmarks
 
@@ -236,7 +240,14 @@ def decorator(fn: Callable) -> list[Benchmark]:
                 )
             names.add(name)
 
-            bm = Benchmark(fn, name=name, params=params, setUp=setUp, tearDown=tearDown, tags=tags)
+            bm = Benchmark(
+                fn,
+                name=name,
+                params=params,
+                setUp=setUp,
+                tearDown=tearDown,
+                tags=tags,
+            )
             benchmarks.append(bm)
         return benchmarks
 

src/nnbench/runner.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import collections
 import contextlib
 import inspect
 import logging
@@ -15,7 +16,7 @@
 from typing import Any, Callable, Generator, Sequence, get_origin
 
 from nnbench.context import Context, ContextProvider
-from nnbench.types import Benchmark, BenchmarkRecord, Parameters
+from nnbench.types import Benchmark, BenchmarkRecord, Parameters, State
 from nnbench.types.util import is_memo, is_memo_type
 from nnbench.util import import_file_as_module, ismodule
 
@@ -247,6 +248,9 @@ def run(
         if not self.benchmarks:
             self.collect(path_or_module, tags)
 
+        family_sizes: dict[str, Any] = collections.defaultdict(int)
+        family_indices: dict[str, Any] = collections.defaultdict(int)
+
         if isinstance(context, Context):
             ctx = context
         else:
@@ -259,6 +263,9 @@ def run(
             warnings.warn(f"No benchmarks found in path/module {str(path_or_module)!r}.")
             return BenchmarkRecord(context=ctx, benchmarks=[])
 
+        for bm in self.benchmarks:
+            family_sizes[bm.fn.__name__] += 1
+
         if isinstance(params, Parameters):
             dparams = asdict(params)
         else:
@@ -274,6 +281,14 @@ def _maybe_dememo(v, expected_type):
             return v
 
         for benchmark in self.benchmarks:
+            bm_family = benchmark.fn.__name__
+            state = State(
+                name=benchmark.name,
+                family=bm_family,
+                family_size=family_sizes[bm_family],
+                family_index=family_indices[bm_family],
+            )
+            family_indices[bm_family] += 1
             bmtypes = dict(zip(benchmark.interface.names, benchmark.interface.types))
             bmparams = dict(zip(benchmark.interface.names, benchmark.interface.defaults))
             # TODO: Does this need a copy.deepcopy()?
@@ -291,14 +306,14 @@ def _maybe_dememo(v, expected_type):
                 "parameters": bmparams,
             }
             try:
-                benchmark.setUp(**bmparams)
+                benchmark.setUp(state, bmparams)
                 with timer(res):
                     res["value"] = benchmark.fn(**bmparams)
             except Exception as e:
                 res["error_occurred"] = True
                 res["error_message"] = str(e)
             finally:
-                benchmark.tearDown(**bmparams)
+                benchmark.tearDown(state, bmparams)
                 results.append(res)
 
         return BenchmarkRecord(

src/nnbench/types/__init__.py

@@ -1 +1 @@
-from .types import Benchmark, BenchmarkRecord, Memo, Parameters
+from .types import Benchmark, BenchmarkRecord, Memo, Parameters, State

src/nnbench/types/types.py

@@ -6,24 +6,23 @@
 import functools
 import inspect
 from dataclasses import dataclass, field
-from typing import (
-    Any,
-    Callable,
-    Generic,
-    Literal,
-    TypeVar,
-)
+from types import MappingProxyType
+from typing import Any, Callable, Generic, Literal, Mapping, Protocol, TypeVar
 
 from nnbench.context import Context
 
 T = TypeVar("T")
 Variable = tuple[str, type, Any]
 
 
-def NoOp(**kwargs: Any) -> None:
+def NoOp(state: State, params: Mapping[str, Any] = MappingProxyType({})) -> None:
     pass
 
 
+class CallbackProtocol(Protocol):
+    def __call__(self, state: State, params: Mapping[str, Any]) -> None: ...
+
+
 @dataclass(frozen=True)
 class BenchmarkRecord:
     context: Context
@@ -101,6 +100,14 @@ def expand(cls, bms: list[dict[str, Any]]) -> BenchmarkRecord:
     #  context data.
 
 
+@dataclass(frozen=True)
+class State:
+    name: str
+    family: str
+    family_size: int
+    family_index: int
+
+
 class Memo(Generic[T]):
     @functools.cache
     # TODO: Swap this out for a local type-wide memo cache.
@@ -158,10 +165,10 @@ class Benchmark:
     """
 
     fn: Callable[..., Any]
-    name: str | None = field(default=None)
+    name: str = field(default="")
     params: dict[str, Any] = field(default_factory=dict)
-    setUp: Callable[..., None] = field(repr=False, default=NoOp)
-    tearDown: Callable[..., None] = field(repr=False, default=NoOp)
+    setUp: Callable[[State, Mapping[str, Any]], None] = field(repr=False, default=NoOp)
+    tearDown: Callable[[State, Mapping[str, Any]], None] = field(repr=False, default=NoOp)
     tags: tuple[str, ...] = field(repr=False, default=())
     interface: Interface = field(init=False, repr=False)