diff --git a/justfile b/justfile index 713b4180b865..5ea4a5bde46e 100644 --- a/justfile +++ b/justfile @@ -130,6 +130,11 @@ py-test: py-test-allpy: nox -s tests +# Run all Python benchmarks +py-bench *ARGS: + python -m pytest -c rerun_py/pyproject.toml --benchmark-only {{ARGS}} + + # Serve the python docs locally py-docs-serve: mkdocs serve -f rerun_py/mkdocs.yml -w rerun_py diff --git a/rerun_py/pyproject.toml b/rerun_py/pyproject.toml index 0b57e5d89531..81aa060951fc 100644 --- a/rerun_py/pyproject.toml +++ b/rerun_py/pyproject.toml @@ -130,3 +130,4 @@ python-packages = ["rerun_sdk/rerun"] filterwarnings = """ error """ +norecursedirs = ".* venv* target* build" diff --git a/tests/python/log_benchmark/__init__.py b/tests/python/log_benchmark/__init__.py new file mode 100644 index 000000000000..35bcee02e22c --- /dev/null +++ b/tests/python/log_benchmark/__init__.py @@ -0,0 +1,26 @@ +from __future__ import annotations + +import dataclasses + +import numpy as np + +MAX_INT64 = 2**63 - 1 +MAX_INT32 = 2**31 - 1 + + +@dataclasses.dataclass +class Point3DInput: + positions: np.ndarray + colors: np.ndarray + radii: np.ndarray + label: str = "some label" + + @classmethod + def prepare(cls, seed: int, num_points: int): + rng = np.random.default_rng(seed=seed) + + return cls( + positions=rng.integers(0, MAX_INT64, (num_points, 3)).astype(dtype=np.float32), + colors=rng.integers(0, MAX_INT32, num_points, dtype=np.uint32), + radii=rng.integers(0, MAX_INT64, num_points).astype(dtype=np.float32), + ) diff --git a/tests/python/log_benchmark/test_log_benchmark.py b/tests/python/log_benchmark/test_log_benchmark.py new file mode 100644 index 000000000000..5a26d0426169 --- /dev/null +++ b/tests/python/log_benchmark/test_log_benchmark.py @@ -0,0 +1,63 @@ +"""Python logging benchmarks. Use `just py-bench` to run.""" + +from __future__ import annotations + +import numpy as np +import pytest +import rerun as rr + +from . import Point3DInput + + +def log_points3d_large_batch(data: Point3DInput): + # create a new, empty memory sink for the current recording + rr.memory_recording() + + rr.log( + "large_batch", + rr.Points3D(positions=data.positions, colors=data.colors, radii=data.radii, labels=data.label), + ) + + +@pytest.mark.parametrize("num_points", [50_000_000]) +def test_bench_points3d_large_batch(benchmark, num_points): + rr.init("rerun_example_benchmark_points3d_large_batch") + data = Point3DInput.prepare(42, num_points) + benchmark(log_points3d_large_batch, data) + + +def log_points3d_many_individual(data: Point3DInput): + # create a new, empty memory sink for the current recording + rr.memory_recording() + + for i in range(data.positions.shape[0]): + rr.log( + "single_point", + rr.Points3D(positions=data.positions[i], colors=data.colors[i], radii=data.radii[i]), + ) + + +@pytest.mark.parametrize("num_points", [100_000]) +def test_bench_points3d_many_individual(benchmark, num_points): + rr.init("rerun_example_benchmark_points3d_many_individual") + data = Point3DInput.prepare(1337, num_points) + benchmark(log_points3d_many_individual, data) + + +def log_image(image: np.ndarray, num_log_calls): + # create a new, empty memory sink for the current recording + rr.memory_recording() + + for i in range(num_log_calls): + rr.log("test_image", rr.Tensor(image)) + + +@pytest.mark.parametrize( + ["image_dimension", "image_channels", "num_log_calls"], + [pytest.param(16_384, 4, 4, id="16384^2px-4channels-4calls")], +) +def test_bench_image(benchmark, image_dimension, image_channels, num_log_calls): + rr.init("rerun_example_benchmark_image") + + image = np.zeros((image_dimension, image_dimension, image_channels), dtype=np.uint8) + benchmark(log_image, image, num_log_calls) diff --git a/tests/python/requirements.txt b/tests/python/requirements.txt index a9fce2d3701c..2a2f51e525aa 100644 --- a/tests/python/requirements.txt +++ b/tests/python/requirements.txt @@ -1,2 +1,3 @@ -r test_api/requirements.txt -r nv12image/requirements.txt +pytest-benchmark