Add Python version of the basic benchmarks

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

rerun_py/pyproject.toml

@@ -130,3 +130,4 @@ python-packages = ["rerun_sdk/rerun"]
 filterwarnings = """
 error
 """
+norecursedirs = ".* venv* target* build"

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

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)

tests/python/requirements.txt

@@ -1,2 +1,3 @@
 -r test_api/requirements.txt
 -r nv12image/requirements.txt
+pytest-benchmark