Repository: aai-institute/nnbench

Issue:
Add user guide: Examples of customizing benchmarks and user-defined `nnbench.Parameters` for type safety
Add user guide to the documentation:
- [ ] Examples of customizing benchmarks (setup/teardown, context info, ...) and user-defined `nnbench.Parameters` for type safety

Discussion:

Patch:
diff --git a/docs/guides/benchmarks.md b/docs/guides/benchmarks.md
index cfd1d0f..5190597 100644
--- a/docs/guides/benchmarks.md
+++ b/docs/guides/benchmarks.md
@@ -68,11 +68,11 @@ Lastly, we set up a benchmark runner in the `main.py`. There, we supply the para
 
 ```python
 # main.py
-from nnbench import runner
+import nnbench
 from nnbench.reporter import ConsoleReporter
 
 
-r = runner.BenchmarkRunner()
+r = nnbench.BenchmarkRunner()
 reporter = ConsoleReporter()
 
 result = r.run("./benchmarks.py", params={"n_estimators": 100, "max_depth": 5, "random_state": 42})
@@ -120,14 +120,13 @@ The unfilled arguments are given in `BenchmarkRunner.run()` via a dictionary pas
 
 ```python
 # main.py
-from nnbench import runner
+import nnbench
 from nnbench.reporter import ConsoleReporter
 
 
-r = runner.BenchmarkRunner()
+r = nnbench.BenchmarkRunner()
 reporter = ConsoleReporter()
 
-r = runner.BenchmarkRunner()
 result = r.run("./benchmarks.py", params={"random_state": 42})
 reporter.report(result)
 ```
@@ -170,7 +169,6 @@ It looks similar to this:
 ```bash
 python main.py  
 
-
 name                                                 value
 ------------------------------------------------  --------
 benchmark_accuracy_n_estimators=10_max_depth=2    0.933333
diff --git a/docs/guides/customization.md b/docs/guides/customization.md
new file mode 100644
index 0000000..9b6d261
--- /dev/null
+++ b/docs/guides/customization.md
@@ -0,0 +1,111 @@
+# Defining setup/teardown tasks, context, and `nnbench.Parameters`
+
+This page introduces some customization options for benchmark runs.
+These options can be helpful for tasks surrounding benchmark state management, such as automatic setup and cleanup, contextualizing results with context values, and defining typed parameters with the `nnbench.Parameters` class.
+
+## Defining setup and teardown tasks
+
+For some benchmarks, it is important to set certain configuration values and prepare the execution environment before running.
+To do this, you can pass a setup task to all of the nnbench decorators via the `setUp` keyword:
+
+```python
+import os
+
+import nnbench
+
+
+def set_envvar(**params):
+    os.environ["MY_ENV"] = "MY_VALUE"
+    
+
+@nnbench.benchmark(setUp=set_envvar)
+def prod(a: int, b: int) -> int:
+    return a * b
+```
+
+Similarly, to revert the environment state back to its previous form (or clean up any created resources), you can supply a finalization task with the `tearDown` keyword:
+
+```python
+import os
+
+import nnbench
+
+
+def set_envvar(**params):
+    os.environ["MY_ENV"] = "MY_VALUE"
+    
+
+def pop_envvar(**params):
+    os.environ.pop("MY_ENV")
+
+
+@nnbench.benchmark(setUp=set_envvar, tearDown=pop_envvar)
+def prod(a: int, b: int) -> int:
+    return a * b
+```
+
+Both the setup and teardown task must take the exact same set of parameters as the benchmark function. To simplify function declaration, it is easiest to use a variadic keyword-only interface, i.e. `setup(**kwargs)`, as shown.
+
+!!! tip
+    This facility works exactly the same for the `@nnbench.parametrize` and `@nnbench.product` decorators.
+    There, the specified setup and teardown tasks are run once before or after each of the resulting benchmarks respectively.
+
+## Enriching benchmark metadata with context values
+
+It is often useful to log specific environment metadata in addition to the benchmark's target metrics.
+Such metadata can give a clearer picture of how certain models perform on a given hardware, how model architectures compare in performance, and much more.
+In `nnbench`, you can give additional metadata to your benchmarks as **context values**.
+
+A _context value_ is defined here as a key-value pair where `key` is a string, and `value` is any valid JSON value holding the desired information.
+As an example, the context value `{"cpuarch": "arm64"}` gives information about the CPU architecture of the host machine running the benchmark.
+
+A _context provider_ is a function taking no arguments and returning a Python dictionary of context values. The following is a basic example of a context provider:
+
+```python
+import platform
+
+def platinfo() -> dict[str, str]:
+    """Returns CPU arch, system name (Windows/Linux/Darwin), and Python version."""
+    return {
+        "system": platform.system(),
+        "cpuarch": platform.machine(),
+        "python_version": platform.python_version(),
+    }
+```
+
+To supply context to your benchmarks, you can give a sequence of context providers to `BenchmarkRunner.run()`:
+
+```python
+import nnbench
+
+# uses the `platinfo` context provider from above to log platform metadata.
+runner = nnbench.BenchmarkRunner()
+result = runner.run(__name__, params={}, context=[platinfo])
+```
+
+## Being type safe by using `nnbench.Parameters`
+
+Instead of specifying your benchmark's parameters by using a raw Python dictionary, you can define a custom subclass of `nnbench.Parameters`:
+
+```python
+import nnbench
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class MyParams(nnbench.Parameters):
+    a: int
+    b: int
+
+
+@nnbench.benchmark
+def prod(a: int, b: int) -> int:
+    return a * b
+
+
+params = MyParams(a=1, b=2)
+runner = nnbench.BenchmarkRunner()
+result = runner.run(__name__, params=params)
+```
+
+While this does not have a concrete advantage in terms of type safety over a raw dictionary (all inputs will be checked against the types expected from the benchmark interfaces), it guards against accidental modification of parameters breaking reproducibility.
diff --git a/docs/guides/organization.md b/docs/guides/organization.md
index 24013e5..8558e1c 100644
--- a/docs/guides/organization.md
+++ b/docs/guides/organization.md
@@ -38,9 +38,9 @@ benchmarks/
 This is helpful when running multiple benchmark workloads separately, as you can just point your benchmark runner to each of these separate files:
 
 ```python
-from nnbench.runner import BenchmarkRunner
+import nnbench
 
-runner = BenchmarkRunner()
+runner = nnbench.BenchmarkRunner()
 data_metrics = runner.run("benchmarks/data_quality.py", params=...)
 # same for model metrics, where instead you pass benchmarks/model_perf.py.
 model_metrics = runner.run("benchmarks/model_perf.py", params=...)
@@ -73,9 +73,9 @@ def bar(data) -> int:
 Now, to only run data quality benchmarks marked "foo", pass the corresponding tag to `BenchmarkRunner.run()`:
 
 ```python
-from nnbench.runner import BenchmarkRunner
+import nnbench
 
-runner = BenchmarkRunner()
+runner = nnbench.BenchmarkRunner()
 foo_data_metrics = runner.run("benchmarks/data_quality.py", params=..., tags=("foo",))
 ```
 
diff --git a/docs/quickstart.md b/docs/quickstart.md
index 44674e4..93eef77 100644
--- a/docs/quickstart.md
+++ b/docs/quickstart.md
@@ -42,24 +42,23 @@ Now we can instantiate a benchmark runner to collect and run the accuracy benchm
 Then, using the `ConsoleReporter` we report the resulting accuracy metric by printing it to the terminal in a table.
 
 ```python
-from nnbench import runner
+import nnbench
 from nnbench.reporter import ConsoleReporter
 
 
-r = runner.BenchmarkRunner()
+r = nnbench.BenchmarkRunner()
 reporter = ConsoleReporter()
 
 # To collect in the current file, pass "__main__" as module name.
 result = r.run("__main__", params={"model": model, "X_test": X_test, "y_test": y_test})
-
 reporter.report(result)
 ```
+
 The resulting output might look like this:
 
 ```bash
 python benchmarks.py  
 
-
 name         value
 --------  --------
 accuracy  0.933333
diff --git a/examples/mnist/mnist.py b/examples/mnist/mnist.py
index b4fe8ef..d3b7bc5 100644
--- a/examples/mnist/mnist.py
+++ b/examples/mnist/mnist.py
@@ -22,7 +22,6 @@ from flax.training.train_state import TrainState
 
 import nnbench
 from nnbench.reporter import ConsoleReporter
-from nnbench.runner import BenchmarkRunner
 
 HERE = Path(__file__).parent
 
@@ -218,7 +217,7 @@ def mnist_jax():
     state, data = train(mnist)
 
     # the nnbench portion.
-    runner = BenchmarkRunner()
+    runner = nnbench.BenchmarkRunner()
     reporter = ConsoleReporter()
     params = MNISTTestParameters(params=state.params, data=data)
     result = runner.run(HERE, params=params)
diff --git a/mkdocs.yml b/mkdocs.yml
index 5750eab..c02f30d 100644
--- a/mkdocs.yml
+++ b/mkdocs.yml
@@ -19,6 +19,7 @@ nav:
   - User Guide:
     - guides/index.md
     - guides/benchmarks.md
+    - guides/customization.md
     - guides/organization.md
   - Examples:
     - tutorials/index.md
diff --git a/src/nnbench/__init__.py b/src/nnbench/__init__.py
index 41f45e3..5b8942c 100644
--- a/src/nnbench/__init__.py
+++ b/src/nnbench/__init__.py
@@ -11,6 +11,7 @@ except PackageNotFoundError:
 # TODO: This naming is unfortunate
 from .core import benchmark, parametrize, product
 from .reporter import BenchmarkReporter, register_reporter
+from .runner import BenchmarkRunner
 from .types import Parameters
 
 
diff --git a/src/nnbench/runner.py b/src/nnbench/runner.py
index d368aa3..656be11 100644
--- a/src/nnbench/runner.py
+++ b/src/nnbench/runner.py
@@ -5,6 +5,7 @@ import inspect
 import logging
 import os
 import sys
+import warnings
 from dataclasses import asdict
 from pathlib import Path
 from typing import Any, Sequence, get_origin
@@ -166,7 +167,7 @@ class BenchmarkRunner:
         params: dict[str, Any] | Parameters,
         tags: tuple[str, ...] = (),
         context: Sequence[ContextProvider] = (),
-    ) -> BenchmarkRecord | None:
+    ) -> BenchmarkRecord:
         """
         Run a previously collected benchmark workload.
 
@@ -188,7 +189,7 @@ class BenchmarkRunner:
 
         Returns
         -------
-        BenchmarkRecord | None
+        BenchmarkRecord
             A JSON output representing the benchmark results. Has two top-level keys, "context"
             holding the context information, and "benchmarks", holding an array with the
             benchmark results.
@@ -201,10 +202,10 @@ class BenchmarkRunner:
         if not self.benchmarks:
             self.collect(path_or_module, tags)
 
-        # if we still have no benchmarks after collection, warn and return early.
+        # if we still have no benchmarks after collection, warn and return an empty record.
         if not self.benchmarks:
-            logger.warning(f"No benchmarks found in path/module {str(path_or_module)!r}.")
-            return None  # TODO: Return empty result to preserve strong typing
+            warnings.warn(f"No benchmarks found in path/module {str(path_or_module)!r}.")
+            return BenchmarkRecord(context={}, benchmarks=[])
 
         if isinstance(params, Parameters):
             dparams = asdict(params)


Test patch:
diff --git a/tests/test_argcheck.py b/tests/test_argcheck.py
index d16e26e..a9dfc67 100644
--- a/tests/test_argcheck.py
+++ b/tests/test_argcheck.py
@@ -3,12 +3,12 @@ import os
 
 import pytest
 
-from nnbench import runner
+import nnbench
 
 
 def test_argcheck(testfolder: str) -> None:
     benchmarks = os.path.join(testfolder, "standard_benchmarks.py")
-    r = runner.BenchmarkRunner()
+    r = nnbench.BenchmarkRunner()
     with pytest.raises(TypeError, match="expected type <class 'int'>.*"):
         r.run(benchmarks, params={"x": 1, "y": "1"}, tags=("standard",))
     with pytest.raises(ValueError, match="missing value for required parameter.*"):
@@ -21,13 +21,13 @@ def test_error_on_duplicate_params(testfolder: str) -> None:
     benchmarks = os.path.join(testfolder, "argument_check_benchmarks.py")
 
     with pytest.raises(TypeError, match="got incompatible types.*"):
-        r = runner.BenchmarkRunner()
+        r = nnbench.BenchmarkRunner()
         r.run(benchmarks, params={"x": 1, "y": 1}, tags=("duplicate",))
 
 
 def test_log_warn_on_overwrite_default(testfolder: str, caplog: pytest.LogCaptureFixture) -> None:
     benchmark = os.path.join(testfolder, "argument_check_benchmarks.py")
-    r = runner.BenchmarkRunner()
+    r = nnbench.BenchmarkRunner()
     with caplog.at_level(logging.DEBUG):
         r.run(benchmark, params={"a": 1}, tags=("with_default",))
     assert "using given value 1 over default value" in caplog.text
@@ -36,5 +36,5 @@ def test_log_warn_on_overwrite_default(testfolder: str, caplog: pytest.LogCaptur
 def test_untyped_interface(testfolder: str) -> None:
     benchmarks = os.path.join(testfolder, "argument_check_benchmarks.py")
 
-    r = runner.BenchmarkRunner()
+    r = nnbench.BenchmarkRunner()
     r.run(benchmarks, params={"value": 2}, tags=("untyped",))
diff --git a/tests/test_runner.py b/tests/test_runner.py
index ada7f98..73648c8 100644
--- a/tests/test_runner.py
+++ b/tests/test_runner.py
@@ -1,10 +1,10 @@
 import os
 
-from nnbench.runner import BenchmarkRunner
+import nnbench
 
 
 def test_runner_discovery(testfolder: str) -> None:
-    r = BenchmarkRunner()
+    r = nnbench.BenchmarkRunner()
 
     r.collect(os.path.join(testfolder, "standard_benchmarks.py"), tags=("runner-collect",))
     assert len(r.benchmarks) == 1
@@ -18,7 +18,7 @@ def test_runner_discovery(testfolder: str) -> None:
 def test_tag_selection(testfolder: str) -> None:
     PATH = os.path.join(testfolder, "tag_selection_benchmark.py")
 
-    r = BenchmarkRunner()
+    r = nnbench.BenchmarkRunner()
 
     r.collect(PATH, tags=())
     assert len(r.benchmarks) == 3