Repository: aai-institute/nnbench

Issue:
Make nnbench ready for concurrency/parallelism
From a conversation with @janfb. 

Unlike microbenchmarks, whose results become less accurate when parallelized due to resource contention, there is no reason not to run an ML benchmark suite in parallel. We should give an example of how to do that in the docs, or even make that the default (think something like `make -j $N`).

Maybe we can get started on a single MP backend, like `multiprocessing` or `joblib`.

Bonus: It could be wise to restructure (read: drop) the `BenchmarkRunner` before this, and instead give the collection and run loop APIs as stateless functions, which is then fairly easy to parallelize over.

Discussion:
@nicholasjng: This also has implications for our benchmark structuring recommendations. In the case of a multiprocessing approach, it might be best to structure benchmarks per-algorithm (model) and send one model's entire benchmark suite to one Python process to improve memory efficiency.

On the other hand, this clashes with our approach for parametrizing over a model input value in a benchmark, which we advertise as a best practice for avoiding code duplication.@nicholasjng: Parallel running is now supported on the CLI after #206, and for the SDK portion, we can cook something up for the docs.

joblib yes/no can be decided in another exploration, so I'm closing this as completed.

Solution patch:
diff --git a/docs/guides/customization.md b/docs/guides/customization.md
index abb1dd5..01a7ebf 100644
--- a/docs/guides/customization.md
+++ b/docs/guides/customization.md
@@ -79,8 +79,8 @@ To supply context to your benchmarks, you can give a sequence of context provide
 import nnbench
 
 # uses the `platinfo` context provider from above to log platform metadata.
-runner = nnbench.BenchmarkRunner()
-result = runner.run(__name__, params={}, context=[platinfo])
+benchmarks = nnbench.collect(__name__)
+result = nnbench.run(benchmarks, params={}, context=[platinfo])
 ```
 
 ## Being type safe by using `nnbench.Parameters`
@@ -104,8 +104,8 @@ def prod(a: int, b: int) -> int:
 
 
 params = MyParams(a=1, b=2)
-runner = nnbench.BenchmarkRunner()
-result = runner.run(__name__, params=params)
+benchmarks = nnbench.collect(__name__)
+result = nnbench.run(benchmarks, params=params)
 ```
 
 While this does not have a concrete advantage in terms of type safety over a raw dictionary, it guards against accidental modification of parameters breaking reproducibility.
diff --git a/docs/guides/organization.md b/docs/guides/organization.md
index 8558e1c..feebb57 100644
--- a/docs/guides/organization.md
+++ b/docs/guides/organization.md
@@ -75,8 +75,8 @@ Now, to only run data quality benchmarks marked "foo", pass the corresponding ta
 ```python
 import nnbench
 
-runner = nnbench.BenchmarkRunner()
-foo_data_metrics = runner.run("benchmarks/data_quality.py", params=..., tags=("foo",))
+benchmarks = nnbench.collect("benchmarks/data_quality.py", tags=("foo",))
+foo_data_metrics = nnbench.run(benchmarks, params=..., )
 ```
 
 !!!tip
diff --git a/docs/guides/runners.md b/docs/guides/runners.md
index 60f5b8f..4df48a6 100644
--- a/docs/guides/runners.md
+++ b/docs/guides/runners.md
@@ -1,17 +1,8 @@
 # Collecting and running benchmarks
 
-nnbench provides the `BenchmarkRunner` as a compact interface to collect and run benchmarks selectively.
+nnbench provides the `nnbench.collect` and `nnbench.run` APIs as a compact interface to collect and run benchmarks selectively.
 
-##  The abstract `BenchmarkRunner`  class
-Let's first instantiate and then walk through the base class.
-
-```python
-from nnbench import BenchmarkRunner
-
-runner = BenchmarkRunner()
-```
-
-Use the `BenchmarkRunner.collect()` method to collect benchmarks from files or directories.  
+Use the `nnbench.collect()` method to collect benchmarks from files or directories.  
 Assume we have the following benchmark setup:
 ```python
 # dir_a/bm1.py
@@ -46,26 +37,29 @@ def the_last_benchmark(d: int) -> int:
 Now we can collect benchmarks from files:
 
 ```python
-runner.collect('dir_a/bm1.py')
+import nnbench
+
+
+benchmarks = nnbench.collect('dir_a/bm1.py')
 ```
 Or directories:
 
 ```python
-runner.collect('dir_b')
+benchmarks = nnbench.collect('dir_b')
 ```
 
-This collection can happen iteratively. So, after executing the two collections our runner has all four benchmarks ready for execution.
-
-To remove the collected benchmarks again, use the `BenchmarkRunner.clear()` method.
 You can also supply tags to the runner to selectively collect only benchmarks with the appropriate tag.
 For example, after clearing the runner again, you can collect all benchmarks with the `"tag"` tag as such:
 
 ```python
-runner.collect('dir_b', tags=("tag",))
+import nnbench
+
+
+tagged_benchmarks = nnbench.collect('dir_b', tags=("tag",))
 ```
 
-To run the benchmarks, call the `BenchmarkRunner.run()` method and supply the necessary parameters required by the collected benchmarks.
+To run the benchmarks, call the `nnbench.run()` method and supply the necessary parameters required by the collected benchmarks.
 
 ```python
-runner.run("dir_b", params={"b": 1, "c": 2, "d": 3})
+result = nnbench.run(benchmarks, params={"b": 1, "c": 2, "d": 3})
 ```
diff --git a/docs/tutorials/duckdb.md b/docs/tutorials/duckdb.md
index 0ac7491..1878e75 100644
--- a/docs/tutorials/duckdb.md
+++ b/docs/tutorials/duckdb.md
@@ -23,8 +23,8 @@ import nnbench
 from nnbench.context import GitEnvironmentInfo
 from nnbench.reporter.file import FileReporter
 
-runner = nnbench.BenchmarkRunner()
-record = runner.run("benchmarks.py", params={"a": 1, "b": 1}, context=(GitEnvironmentInfo(),))
+benchmarks = nnbench.collect("benchmarks.py")
+record = nnbench.run(benchmarks, params={"a": 1, "b": 1}, context=(GitEnvironmentInfo(),))
 
 file_reporter = FileReporter()
 file_reporter.write(record, "record.json", driver="ndjson")
diff --git a/docs/tutorials/huggingface.md b/docs/tutorials/huggingface.md
index 40f36d3..096d628 100644
--- a/docs/tutorials/huggingface.md
+++ b/docs/tutorials/huggingface.md
@@ -91,7 +91,7 @@ In the following `IndexLabelMapMemo` class, we store a dictionary mapping the la
 
 !!! Info
     There is no need to type-hint `TokenClassificationModelMemo`s in the corresponding benchmarks -
-    the benchmark runner takes care of filling in the memoized values for the memos themselves.
+    the benchmark running method takes care of filling in the memoized values for the memos themselves.
 
 Because we implemented our memoized values as four different memo class types, this modularizes the benchmark input parameters -
 we only need to reference memos when they are actually used. Considering the recall benchmarks:
diff --git a/docs/tutorials/mnist.md b/docs/tutorials/mnist.md
index ce12e1a..b80fed8 100644
--- a/docs/tutorials/mnist.md
+++ b/docs/tutorials/mnist.md
@@ -16,14 +16,14 @@ To properly structure our project, we avoid mixing training pipeline code and be
 
 This definition is short and sweet, and contains a few important details:
 
-* Both functions are given the `@nnbench.benchmark` decorator - this enables our runner to find and collect them before starting the benchmark run.
+* Both functions are given the `@nnbench.benchmark` decorator - this allows us to find and collect them before starting the benchmark run.
 * The `modelsize` benchmark is given a custom name (`"Model size (MB)"`), indicating that the resulting number is the combined size of the model weights in megabytes.
 This is done for display purposes, to improve interpretability when reporting results.
 * The `params` argument is the same in both benchmarks, both in name and type. This is important, since it ensures that both benchmarks will be run with the same model weights.
 
 That's all - now we can shift over to our main pipeline code and see what is necessary to execute the benchmarks and visualize the results.
 
-## Setting up a benchmark runner and parameters
+## Setting up a benchmark run and parameters
 
 After finishing the benchmark setup, we only need a few more lines to augment our pipeline.
 
diff --git a/examples/bq/bq.py b/examples/bq/bq.py
index ff22fa6..bfeefbb 100644
--- a/examples/bq/bq.py
+++ b/examples/bq/bq.py
@@ -14,8 +14,8 @@ def main():
         autodetect=True, source_format=bigquery.SourceFormat.NEWLINE_DELIMITED_JSON
     )
 
-    runner = nnbench.BenchmarkRunner()
-    res = runner.run("benchmarks.py", params={"a": 1, "b": 1}, context=(GitEnvironmentInfo(),))
+    benchmarks = nnbench.collect("benchmarks.py")
+    res = nnbench.run(benchmarks, params={"a": 1, "b": 1}, context=(GitEnvironmentInfo(),))
 
     load_job = client.load_table_from_json(res.to_json(), table_id, job_config=job_config)
     load_job.result()
diff --git a/examples/huggingface/runner.py b/examples/huggingface/runner.py
index 18157f8..66651a8 100644
--- a/examples/huggingface/runner.py
+++ b/examples/huggingface/runner.py
@@ -2,9 +2,9 @@
 
 
 def main() -> None:
-    runner = nnbench.BenchmarkRunner()
+    benchmarks = nnbench.collect("benchmark.py", tags=("per-class",))
     reporter = nnbench.ConsoleReporter()
-    result = runner.run("benchmark.py", tags=("per-class",))
+    result = nnbench.run(benchmarks)
     reporter.display(result)
 
 
diff --git a/examples/mnist/mnist.py b/examples/mnist/mnist.py
index 1d19669..97faf9f 100644
--- a/examples/mnist/mnist.py
+++ b/examples/mnist/mnist.py
@@ -216,10 +216,10 @@ def mnist_jax():
     state, data = train(mnist)
 
     # the nnbench portion.
-    runner = nnbench.BenchmarkRunner()
+    benchmarks = nnbench.collect(HERE)
     reporter = nnbench.FileReporter()
     params = MNISTTestParameters(params=state.params, data=data)
-    result = runner.run(HERE, params=params)
+    result = nnbench.run(benchmarks, params=params)
     reporter.write(result, "result.json")
 
 
diff --git a/examples/prefect/src/runner.py b/examples/prefect/src/runner.py
index 351e8c8..569bbf7 100644
--- a/examples/prefect/src/runner.py
+++ b/examples/prefect/src/runner.py
@@ -31,10 +31,9 @@ async def write(
 def run_metric_benchmarks(
     model: base.BaseEstimator, X_test: np.ndarray, y_test: np.ndarray
 ) -> nnbench.types.BenchmarkRecord:
-    runner = nnbench.BenchmarkRunner()
-    results = runner.run(
-        os.path.join(dir_path, "benchmark.py"),
-        tags=("metric",),
+    benchmarks = nnbench.collect(os.path.join(dir_path, "benchmark.py"), tags=("metric",))
+    results = nnbench.run(
+        benchmarks,
         params={"model": model, "X_test": X_test, "y_test": y_test},
     )
     return results
@@ -44,10 +43,9 @@ def run_metric_benchmarks(
 def run_metadata_benchmarks(
     model: base.BaseEstimator, X: np.ndarray
 ) -> nnbench.types.BenchmarkRecord:
-    runner = nnbench.BenchmarkRunner()
-    result = runner.run(
-        os.path.join(dir_path, "benchmark.py"),
-        tags=("model-meta",),
+    benchmarks = nnbench.collect(os.path.join(dir_path, "benchmark.py"), tags=("model-meta",))
+    result = nnbench.run(
+        benchmarks,
         params={"model": model, "X": X},
     )
     return result
@@ -73,7 +71,7 @@ async def train_and_benchmark(
     metadata_results: types.BenchmarkRecord = run_metadata_benchmarks(model=model, X=X_test)
 
     metadata_results.context.update(data_params)
-    metadata_results.context.update(context.PythonInfo())
+    metadata_results.context.update(context.PythonInfo()())
 
     await reporter.write(
         record=metadata_results, key="model-attributes", description="Model Attributes"
@@ -84,7 +82,7 @@ async def train_and_benchmark(
     )
 
     metric_results.context.update(data_params)
-    metric_results.context.update(context.PythonInfo())
+    metric_results.context.update(context.PythonInfo()())
     await reporter.write(metric_results, key="model-performance", description="Model Performance")
     return metadata_results, metric_results
 
diff --git a/src/nnbench/__init__.py b/src/nnbench/__init__.py
index 0315052..fc35ed9 100644
--- a/src/nnbench/__init__.py
+++ b/src/nnbench/__init__.py
@@ -2,7 +2,7 @@
 
 from .core import benchmark, parametrize, product
 from .reporter import BenchmarkReporter, ConsoleReporter, FileReporter
-from .runner import BenchmarkRunner
+from .runner import collect, run
 from .types import Benchmark, BenchmarkRecord, Memo, Parameters
 
 __version__ = "0.4.0"
diff --git a/src/nnbench/cli.py b/src/nnbench/cli.py
index bee2d4a..a90cea8 100644
--- a/src/nnbench/cli.py
+++ b/src/nnbench/cli.py
@@ -6,7 +6,7 @@
 import sys
 from typing import Any
 
-from nnbench import BenchmarkRunner, ConsoleReporter, __version__
+from nnbench import ConsoleReporter, __version__, collect, run
 from nnbench.config import NNBenchConfig, parse_nnbench_config
 from nnbench.reporter import FileReporter
 
@@ -205,9 +205,9 @@ def main() -> int:
                 else:
                     context[k] = v
 
-            record = BenchmarkRunner().run(
-                args.benchmarks,
-                tags=tuple(args.tags),
+            benchmarks = collect(args.benchmarks, tags=tuple(args.tags))
+            record = run(
+                benchmarks,
                 context=[lambda: context],
             )
 
diff --git a/src/nnbench/context.py b/src/nnbench/context.py
index 27a3c33..be8cfb2 100644
--- a/src/nnbench/context.py
+++ b/src/nnbench/context.py
@@ -29,7 +29,7 @@ class PythonInfo:
 
     Parameters
     ----------
-    packages: str
+    packages: Sequence[str]
         Names of the requested packages under which they exist in the current environment.
         For packages installed through ``pip``, this equals the PyPI package name.
     """
@@ -37,7 +37,7 @@ class PythonInfo:
     key = "python"
 
     def __init__(self, packages: Sequence[str] = ()):
-        self.packages = packages
+        self.packages = tuple(packages)
 
     def __call__(self) -> dict[str, Any]:
         from importlib.metadata import PackageNotFoundError, version
diff --git a/src/nnbench/runner.py b/src/nnbench/runner.py
index 01266a7..d947813 100644
--- a/src/nnbench/runner.py
+++ b/src/nnbench/runner.py
@@ -9,7 +9,6 @@
 import sys
 import time
 import uuid
-import warnings
 from collections.abc import Callable, Generator, Sequence
 from dataclasses import asdict
 from datetime import datetime
@@ -94,213 +93,184 @@ def _jsonify(val):
     return json_params
 
 
-class BenchmarkRunner:
+def collect(path_or_module: str | os.PathLike[str], tags: tuple[str, ...] = ()) -> list[Benchmark]:
+    # TODO: functools.cache this guy
     """
-    An abstract benchmark runner class.
+    Discover benchmarks in a module and memoize them for later use.
 
-    Collects benchmarks from a module or file using the ``BenchmarkRunner.collect()``
-    method.
+    Parameters
+    ----------
+    path_or_module: str | os.PathLike[str]
+        Name or path of the module to discover benchmarks in. Can also be a directory,
+        in which case benchmarks are collected from the Python files therein.
+    tags: tuple[str, ...]
+        Tags to filter for when collecting benchmarks. Only benchmarks containing either of
+        these tags are collected.
+
+    Raises
+    ------
+    ValueError
+        If the given path is not a Python file, directory, or module name.
+    """
+    benchmarks: list[Benchmark] = []
+    ppath = Path(path_or_module)
+    if ppath.is_dir():
+        pythonpaths = (p for p in ppath.iterdir() if p.suffix == ".py")
+        for py in pythonpaths:
+            logger.debug(f"Collecting benchmarks from submodule {py.name!r}.")
+            benchmarks.extend(collect(py, tags))
+        return benchmarks
+    elif ppath.is_file():
+        logger.debug(f"Collecting benchmarks from file {ppath}.")
+        module = import_file_as_module(path_or_module)
+    elif ismodule(path_or_module):
+        module = sys.modules[str(path_or_module)]
+    else:
+        raise ValueError(
+            f"expected a module name, Python file, or directory, " f"got {str(path_or_module)!r}"
+        )
 
-    Runs a previously collected benchmark workload with parameters in the
-    ``BenchmarkRunner.run()`` method, outputting the results to a
-    ``BenchmarkRecord`` dataclass.
+    # iterate through the module dict members to register
+    for k, v in module.__dict__.items():
+        if k.startswith("__") and k.endswith("__"):
+            # dunder names are ignored.
+            continue
+        elif isinstance(v, Benchmark):
+            if not tags or set(tags) & set(v.tags):
+                benchmarks.append(v)
+        elif isinstance(v, list | tuple | set | frozenset):
+            for bm in v:
+                if isinstance(bm, Benchmark):
+                    if not tags or set(tags) & set(bm.tags):
+                        benchmarks.append(bm)
+
+    return benchmarks
+
+
+def run(
+    benchmarks: Benchmark | Sequence[Benchmark],  # TODO: Support benchmark iterators
+    name: str | None = None,
+    params: dict[str, Any] | Parameters | None = None,
+    context: Sequence[ContextProvider] = (),
+    jsonifier: Callable[[dict[str, Any]], dict[str, Any]] = jsonify_params,
+) -> BenchmarkRecord:
     """
+    Run a previously collected benchmark workload.
 
-    benchmark_type = Benchmark
-
-    def __init__(self):
-        self.benchmarks: list[Benchmark] = list()
-
-    def clear(self) -> None:
-        """Clear all registered benchmarks."""
-        self.benchmarks.clear()
-
-    def collect(self, path_or_module: str | os.PathLike[str], tags: tuple[str, ...] = ()) -> None:
-        # TODO: functools.cache this guy
-        """
-        Discover benchmarks in a module and memoize them for later use.
-
-        Parameters
-        ----------
-        path_or_module: str | os.PathLike[str]
-            Name or path of the module to discover benchmarks in. Can also be a directory,
-            in which case benchmarks are collected from the Python files therein.
-        tags: tuple[str, ...]
-            Tags to filter for when collecting benchmarks. Only benchmarks containing either of
-            these tags are collected.
-
-        Raises
-        ------
-        ValueError
-            If the given path is not a Python file, directory, or module name.
-        """
-        ppath = Path(path_or_module)
-        if ppath.is_dir():
-            pythonpaths = (p for p in ppath.iterdir() if p.suffix == ".py")
-            for py in pythonpaths:
-                logger.debug(f"Collecting benchmarks from submodule {py.name!r}.")
-                self.collect(py, tags)
-            return
-        elif ppath.is_file():
-            logger.debug(f"Collecting benchmarks from file {ppath}.")
-            module = import_file_as_module(path_or_module)
-        elif ismodule(path_or_module):
-            module = sys.modules[str(path_or_module)]
-        else:
-            raise ValueError(
-                f"expected a module name, Python file, or directory, "
-                f"got {str(path_or_module)!r}"
-            )
-
-        # iterate through the module dict members to register
-        for k, v in module.__dict__.items():
-            if k.startswith("__") and k.endswith("__"):
-                # dunder names are ignored.
-                continue
-            elif isinstance(v, self.benchmark_type):
-                if not tags or set(tags) & set(v.tags):
-                    self.benchmarks.append(v)
-            elif isinstance(v, list | tuple | set | frozenset):
-                for bm in v:
-                    if isinstance(bm, self.benchmark_type):
-                        if not tags or set(tags) & set(bm.tags):
-                            self.benchmarks.append(bm)
-
-    def run(
-        self,
-        path_or_module: str | os.PathLike[str],
-        name: str | None = None,
-        params: dict[str, Any] | Parameters | None = None,
-        tags: tuple[str, ...] = (),
-        context: Sequence[ContextProvider] = (),
-        jsonifier: Callable[[dict[str, Any]], dict[str, Any]] = jsonify_params,
-    ) -> BenchmarkRecord:
-        """
-        Run a previously collected benchmark workload.
-
-        Parameters
-        ----------
-        path_or_module: str | os.PathLike[str]
-            Name or path of the module to discover benchmarks in. Can also be a directory,
-            in which case benchmarks are collected from the Python files therein.
-        name: str | None
-            A name for the currently started run. If None, a name will be automatically generated.
-        params: dict[str, Any] | Parameters | None
-            Parameters to use for the benchmark run. Names have to match positional and keyword
-            argument names of the benchmark functions.
-        tags: tuple[str, ...]
-            Tags to filter for when collecting benchmarks. Only benchmarks containing either of
-            these tags are collected.
-        context: Sequence[ContextProvider]
-            Additional context to log with the benchmarks in the output JSON record. Useful for
-            obtaining environment information and configuration, like CPU/GPU hardware info,
-            ML model metadata, and more.
-        jsonifier: Callable[[dict[str, Any], dict[str, Any]]]
-            A function constructing a string representation from the input parameters.
-            Defaults to ``nnbench.runner.jsonify_params()``. Must produce a dictionary containing
-            only JSON-serializable values.
-
-        Returns
-        -------
-        BenchmarkRecord
-            A JSON output representing the benchmark results. Has three top-level keys,
-            "name" giving the benchmark run name, "context" holding the context information,
-            and "benchmarks", holding an array with the benchmark results.
-        """
-        run = name or "nnbench-" + platform.node() + "-" + uuid.uuid1().hex[:8]
-
-        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)
-
-        ctx: dict[str, Any] = {}
-        for provider in context:
-            val = provider()
-            duplicates = set(ctx.keys()) & set(val.keys())
-            if duplicates:
-                dupe, *_ = duplicates
-                raise ValueError(f"got multiple values for context key {dupe!r}")
-            ctx.update(val)
-
-        # if we didn't find any benchmarks, warn and return an empty record.
-        if not self.benchmarks:
-            warnings.warn(f"No benchmarks found in path/module {str(path_or_module)!r}.")
-            return BenchmarkRecord(run=run, context=ctx, benchmarks=[])
-
-        for bm in self.benchmarks:
-            family_sizes[bm.interface.funcname] += 1
-
-        if isinstance(params, Parameters):
-            dparams = asdict(params)
-        else:
-            dparams = params or {}
-
-        results: list[dict[str, Any]] = []
-
-        def _maybe_dememo(v, expected_type):
-            """Compute and memoize a value if a memo is given for a variable."""
-            if is_memo(v) and not is_memo_type(expected_type):
-                return v()
-            return v
-
-        for benchmark in self.benchmarks:
-            bm_family = benchmark.interface.funcname
-            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
-
-            # Assemble benchmark parameters. First grab all defaults from the interface,
-            bmparams = {
-                name: _maybe_dememo(val, typ)
-                for name, typ, val in benchmark.interface.variables
-                if val != inspect.Parameter.empty
-            }
-            # ... then hydrate with the appropriate subset of input parameters.
-            bmparams |= {k: v for k, v in dparams.items() if k in benchmark.interface.names}
-            # If any arguments are still unresolved, go look them up as fixtures.
-            if set(bmparams) < set(benchmark.interface.names):
-                # TODO: This breaks for a module name (like __main__).
-                # Since that only means that we cannot resolve fixtures when benchmarking
-                # a module name (which makes sense), and we can always pass extra
-                # parameters in the module case, fixing this is not as urgent.
-                p = Path(path_or_module)
-                if p.is_file():
-                    root = p.parent
-                elif p.is_dir():
-                    root = p
-                else:
-                    raise ValueError(f"expected a file or directory, got {path_or_module!r}")
-                fm = FixtureManager(root)
-                bmparams |= fm.resolve(benchmark)
-
-            res: dict[str, Any] = {
-                "name": benchmark.name,
-                "function": qualname(benchmark.fn),
-                "description": benchmark.fn.__doc__ or "",
-                "date": datetime.now().isoformat(timespec="seconds"),
-                "error_occurred": False,
-                "error_message": "",
-                "parameters": jsonifier(bmparams),
-            }
-            try:
-                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(state, bmparams)
-                results.append(res)
-
-        return BenchmarkRecord(
-            run=run,
-            context=ctx,
-            benchmarks=results,
+    Parameters
+    ----------
+    benchmarks: Sequence[Benchmark]
+        The list of discovered benchmarks to run.
+    name: str | None
+        A name for the currently started run. If None, a name will be automatically generated.
+    params: dict[str, Any] | Parameters | None
+        Parameters to use for the benchmark run. Names have to match positional and keyword
+        argument names of the benchmark functions.
+    context: Sequence[ContextProvider]
+        Additional context to log with the benchmarks in the output JSON record. Useful for
+        obtaining environment information and configuration, like CPU/GPU hardware info,
+        ML model metadata, and more.
+    jsonifier: Callable[[dict[str, Any], dict[str, Any]]]
+        A function constructing a string representation from the input parameters.
+        Defaults to ``nnbench.runner.jsonify_params()``. Must produce a dictionary containing
+        only JSON-serializable values.
+
+    Returns
+    -------
+    BenchmarkRecord
+        A JSON output representing the benchmark results. Has three top-level keys,
+        "name" giving the benchmark run name, "context" holding the context information,
+        and "benchmarks", holding an array with the benchmark results.
+    """
+    _run = name or "nnbench-" + platform.node() + "-" + uuid.uuid1().hex[:8]
+
+    family_sizes: dict[str, Any] = collections.defaultdict(int)
+    family_indices: dict[str, Any] = collections.defaultdict(int)
+
+    ctx: dict[str, Any] = {}
+    for provider in context:
+        val = provider()
+        duplicates = set(ctx.keys()) & set(val.keys())
+        if duplicates:
+            dupe, *_ = duplicates
+            raise ValueError(f"got multiple values for context key {dupe!r}")
+        ctx.update(val)
+
+    if isinstance(benchmarks, Benchmark):
+        benchmarks = [benchmarks]
+
+    # if we didn't find any benchmarks, return an empty record.
+    if not benchmarks:
+        return BenchmarkRecord(run=_run, context=ctx, benchmarks=[])
+
+    for bm in benchmarks:
+        family_sizes[bm.interface.funcname] += 1
+
+    if isinstance(params, Parameters):
+        dparams = asdict(params)
+    else:
+        dparams = params or {}
+
+    results: list[dict[str, Any]] = []
+
+    def _maybe_dememo(v, expected_type):
+        """Compute and memoize a value if a memo is given for a variable."""
+        if is_memo(v) and not is_memo_type(expected_type):
+            return v()
+        return v
+
+    for benchmark in benchmarks:
+        bm_family = benchmark.interface.funcname
+        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
+
+        # Assemble benchmark parameters. First grab all defaults from the interface,
+        bmparams = {
+            name: _maybe_dememo(val, typ)
+            for name, typ, val in benchmark.interface.variables
+            if val != inspect.Parameter.empty
+        }
+        # ... then hydrate with the appropriate subset of input parameters.
+        bmparams |= {k: v for k, v in dparams.items() if k in benchmark.interface.names}
+        # If any arguments are still unresolved, go look them up as fixtures.
+        if set(bmparams) < set(benchmark.interface.names):
+            # TODO: This breaks for a module name (like __main__).
+            # Since that only means that we cannot resolve fixtures when benchmarking
+            # a module name (which makes sense), and we can always pass extra
+            # parameters in the module case, fixing this is not as urgent.
+            mod = benchmark.__module__
+            file = sys.modules[mod].__file__
+            p = Path(file).parent
+            fm = FixtureManager(p)
+            bmparams |= fm.resolve(benchmark)
+
+        res: dict[str, Any] = {
+            "name": benchmark.name,
+            "function": qualname(benchmark.fn),
+            "description": benchmark.fn.__doc__ or "",
+            "date": datetime.now().isoformat(timespec="seconds"),
+            "error_occurred": False,
+            "error_message": "",
+            "parameters": jsonifier(bmparams),
+        }
+        try:
+            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(state, bmparams)
+            results.append(res)
+
+    return BenchmarkRecord(
+        run=_run,
+        context=ctx,
+        benchmarks=results,
+    )


Test patch:
diff --git a/tests/test_runner.py b/tests/test_runner.py
index c1aa2ca..3c138cc 100644
--- a/tests/test_runner.py
+++ b/tests/test_runner.py
@@ -7,45 +7,29 @@
 
 
 def test_runner_collection(testfolder: str) -> None:
-    r = nnbench.BenchmarkRunner()
+    benchmarks = nnbench.collect(os.path.join(testfolder, "standard.py"), tags=("runner-collect",))
+    assert len(benchmarks) == 1
 
-    r.collect(os.path.join(testfolder, "standard.py"), tags=("runner-collect",))
-    assert len(r.benchmarks) == 1
-    r.clear()
+    benchmarks = nnbench.collect(testfolder, tags=("non-existing-tag",))
+    assert len(benchmarks) == 0
 
-    r.collect(testfolder, tags=("non-existing-tag",))
-    assert len(r.benchmarks) == 0
-    r.clear()
-
-    r.collect(testfolder, tags=("runner-collect",))
-    assert len(r.benchmarks) == 1
+    benchmarks = nnbench.collect(testfolder, tags=("runner-collect",))
+    assert len(benchmarks) == 1
 
 
 def test_tag_selection(testfolder: str) -> None:
     PATH = os.path.join(testfolder, "tags.py")
 
-    r = nnbench.BenchmarkRunner()
-
-    r.collect(PATH)
-    assert len(r.benchmarks) == 3
-    r.clear()
-
-    r.collect(PATH, tags=("tag1",))
-    assert len(r.benchmarks) == 2
-    r.clear()
-
-    r.collect(PATH, tags=("tag2",))
-    assert len(r.benchmarks) == 1
-    r.clear()
+    assert len(nnbench.collect(PATH)) == 3
+    assert len(nnbench.collect(PATH, tags=("tag1",))) == 2
+    assert len(nnbench.collect(PATH, tags=("tag2",))) == 1
 
 
 def test_context_assembly(testfolder: str) -> None:
-    r = nnbench.BenchmarkRunner()
-
     context_providers = [system, cpuarch, python_version]
-    result = r.run(
-        testfolder,
-        tags=("standard",),
+    benchmarks = nnbench.collect(testfolder, tags=("standard",))
+    result = nnbench.run(
+        benchmarks,
         params={"x": 1, "y": 1},
         context=context_providers,
     )
@@ -57,17 +41,15 @@ def test_context_assembly(testfolder: str) -> None:
 
 
 def test_error_on_duplicate_context_keys_in_runner(testfolder: str) -> None:
-    r = nnbench.BenchmarkRunner()
-
     def duplicate_context_provider() -> dict[str, str]:
         return {"system": "DuplicateSystem"}
 
     context_providers = [system, duplicate_context_provider]
 
+    benchmarks = nnbench.collect(testfolder, tags=("standard",))
     with pytest.raises(ValueError, match="got multiple values for context key 'system'"):
-        r.run(
-            testfolder,
-            tags=("standard",),
+        nnbench.run(
+            benchmarks,
             params={"x": 1, "y": 1},
             context=context_providers,
         )
@@ -78,11 +60,9 @@ def test_filter_benchmarks_on_params(testfolder: str) -> None:
     def prod(a: int, b: int = 1) -> int:
         return a * b
 
-    r = nnbench.BenchmarkRunner()
-    r.benchmarks.append(prod)
-    # TODO (nicholasjng): This is hacky
-    rec1 = r.run("", params={"a": 1, "b": 2})
+    benchmarks = [prod]
+    rec1 = nnbench.run(benchmarks, params={"a": 1, "b": 2})
     assert rec1.benchmarks[0]["parameters"] == {"a": 1, "b": 2}
     # Assert that the defaults are also present if not overridden.
-    rec2 = r.run("", params={"a": 1})
+    rec2 = nnbench.run(benchmarks, params={"a": 1})
     assert rec2.benchmarks[0]["parameters"] == {"a": 1, "b": 1}