chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,79 @@
|
||||
# Python SDK logging benchmarks
|
||||
|
||||
Manual performance benchmarks for the Rerun Python SDK logging pipeline.
|
||||
These are **not** run in CI — they are intended for local profiling and regression checks.
|
||||
|
||||
## Running benchmarks
|
||||
|
||||
From the `rerun/` directory:
|
||||
|
||||
```bash
|
||||
# Run all benchmarks:
|
||||
pixi run py-bench
|
||||
|
||||
# Run only throughput benchmarks:
|
||||
pixi run py-bench -k "not micro"
|
||||
|
||||
# Run only micro-benchmarks:
|
||||
pixi run py-bench -k micro
|
||||
|
||||
# Run a specific benchmark:
|
||||
pixi run py-bench -k "micro_log-Points3D"
|
||||
```
|
||||
|
||||
## Running standalone (for profiling)
|
||||
|
||||
Enter the pixi shell first:
|
||||
|
||||
```bash
|
||||
pixi shell
|
||||
```
|
||||
|
||||
Then run a benchmark directly:
|
||||
|
||||
```bash
|
||||
# Run the throughput benchmark standalone:
|
||||
uvpy -m tests.python.log_benchmark.test_log_benchmark transform3d
|
||||
|
||||
# With options:
|
||||
uvpy -m tests.python.log_benchmark.test_log_benchmark transform3d --num-entities 10 --num-time-steps 10000 --static
|
||||
|
||||
# Connect to a running Rerun viewer (start `rerun` first):
|
||||
uvpy -m tests.python.log_benchmark.test_log_benchmark transform3d --connect
|
||||
```
|
||||
|
||||
### Profiling with py-spy
|
||||
|
||||
```bash
|
||||
# Generate a flamegraph (on Linux, add --native for native stack traces):
|
||||
sudo PYTHONPATH=rerun_py/rerun_sdk:rerun_py py-spy record -o flamegraph.svg -- \
|
||||
.venv/bin/python -m tests.python.log_benchmark.test_log_benchmark transform3d
|
||||
|
||||
# Then open flamegraph.svg in a browser
|
||||
```
|
||||
|
||||
## Comparing benchmark runs
|
||||
|
||||
Use `--benchmark-save` to save benchmark results:
|
||||
|
||||
```bash
|
||||
# Save a baseline on the current branch:
|
||||
pixi run py-bench -k micro --benchmark-save=before
|
||||
|
||||
# Make changes, rebuild, then save again:
|
||||
pixi run py-bench -k micro --benchmark-save=after
|
||||
```
|
||||
|
||||
Saved results are stored in `.benchmarks/` under the project root and are automatically numbered, e.g. `0001_before` and `0002_after`.
|
||||
|
||||
You can then compare using the `pytest-benchmark` CLI:
|
||||
```
|
||||
uv run pytest-benchmark compare 0001 0002
|
||||
```
|
||||
|
||||
|
||||
## Test files
|
||||
|
||||
- `__init__.py` — Shared data classes (`Point3DInput`, `Transform3DInput`)
|
||||
- `test_log_benchmark.py` — Throughput benchmarks
|
||||
- `test_micro_benchmark.py` — Per-call overhead micro-benchmarks
|
||||
@@ -0,0 +1,54 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import dataclasses
|
||||
|
||||
import numpy as np
|
||||
import numpy.typing as npt
|
||||
|
||||
MAX_INT64 = 2**63 - 1
|
||||
MAX_INT32 = 2**31 - 1
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class Point3DInput:
|
||||
positions: npt.NDArray[np.float32]
|
||||
colors: npt.NDArray[np.uint32]
|
||||
radii: npt.NDArray[np.float32]
|
||||
label: str = "some label"
|
||||
|
||||
@classmethod
|
||||
def prepare(cls, seed: int, num_points: int) -> Point3DInput:
|
||||
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),
|
||||
)
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class Transform3DInput:
|
||||
"""Input data for Transform3D benchmark with translation and mat3x3."""
|
||||
|
||||
translations: npt.NDArray[np.float32] # Shape: (num_time_steps, num_entities, 3)
|
||||
mat3x3s: npt.NDArray[np.float32] # Shape: (num_time_steps, num_entities, 3, 3)
|
||||
num_entities: int
|
||||
num_time_steps: int
|
||||
|
||||
@classmethod
|
||||
def prepare(cls, seed: int, num_entities: int, num_time_steps: int) -> Transform3DInput:
|
||||
rng = np.random.default_rng(seed=seed)
|
||||
|
||||
# Generate translations in range [0, 10)
|
||||
translations = rng.random((num_time_steps, num_entities, 3), dtype=np.float32) * 10.0
|
||||
|
||||
# Generate mat3x3 values in range [-1, 1)
|
||||
mat3x3s = rng.random((num_time_steps, num_entities, 3, 3), dtype=np.float32) * 2.0 - 1.0
|
||||
|
||||
return cls(
|
||||
translations=translations,
|
||||
mat3x3s=mat3x3s,
|
||||
num_entities=num_entities,
|
||||
num_time_steps=num_time_steps,
|
||||
)
|
||||
@@ -0,0 +1,253 @@
|
||||
"""Python SDK logging throughput benchmarks. See README.md for usage."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import numpy.typing as npt
|
||||
import pytest
|
||||
|
||||
import rerun as rr
|
||||
|
||||
from . import Point3DInput, Transform3DInput
|
||||
|
||||
|
||||
def log_points3d_large_batch(data: Point3DInput) -> None:
|
||||
# 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: Any, num_points: int) -> None:
|
||||
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) -> None:
|
||||
# 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: Any, num_points: int) -> None:
|
||||
rr.init("rerun_example_benchmark_points3d_many_individual")
|
||||
data = Point3DInput.prepare(1337, num_points)
|
||||
benchmark(log_points3d_many_individual, data)
|
||||
|
||||
|
||||
def log_image(image: npt.NDArray[np.uint8], num_log_calls: int) -> None:
|
||||
# create a new, empty memory sink for the current recording
|
||||
rr.memory_recording()
|
||||
|
||||
for _ in range(num_log_calls):
|
||||
rr.log("test_image", rr.Tensor(image))
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
["image_dimension", "image_channels", "num_log_calls"],
|
||||
[pytest.param(1024, 4, 20_000, id="1024^2px-4channels-20000calls")],
|
||||
)
|
||||
def test_bench_image(benchmark: Any, image_dimension: int, image_channels: int, num_log_calls: int) -> None:
|
||||
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)
|
||||
|
||||
|
||||
def test_bench_transforms_over_time_individual(
|
||||
rand_trans: npt.NDArray[np.float32],
|
||||
rand_quats: npt.NDArray[np.float32],
|
||||
rand_scales: npt.NDArray[np.float32],
|
||||
) -> None:
|
||||
# create a new, empty memory sink for the current recording
|
||||
rr.memory_recording()
|
||||
|
||||
num_transforms = rand_trans.shape[0]
|
||||
for i in range(num_transforms):
|
||||
rr.set_time("frame", sequence=i)
|
||||
rr.log(
|
||||
"test_transform",
|
||||
rr.Transform3D(translation=rand_trans[i], rotation=rr.Quaternion(xyzw=rand_quats[i]), scale=rand_scales[i]),
|
||||
)
|
||||
|
||||
|
||||
def test_bench_transforms_over_time_batched(
|
||||
rand_trans: npt.NDArray[np.float32],
|
||||
rand_quats: npt.NDArray[np.float32],
|
||||
rand_scales: npt.NDArray[np.float32],
|
||||
num_transforms_per_batch: int,
|
||||
) -> None:
|
||||
# create a new, empty memory sink for the current recording
|
||||
rr.memory_recording()
|
||||
|
||||
num_transforms = rand_trans.shape[0]
|
||||
num_log_calls = num_transforms // num_transforms_per_batch
|
||||
for i in range(num_log_calls):
|
||||
start = i * num_transforms_per_batch
|
||||
end = (i + 1) * num_transforms_per_batch
|
||||
times = np.arange(start, end)
|
||||
|
||||
rr.send_columns(
|
||||
"test_transform",
|
||||
indexes=[rr.TimeColumn("frame", sequence=times)],
|
||||
columns=rr.Transform3D.columns(
|
||||
translation=rand_trans[start:end],
|
||||
quaternion=rand_quats[start:end],
|
||||
scale=rand_scales[start:end],
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
["num_transforms", "num_transforms_per_batch"],
|
||||
[
|
||||
pytest.param(10_000, 1),
|
||||
pytest.param(10_000, 100),
|
||||
pytest.param(10_000, 1_000),
|
||||
],
|
||||
)
|
||||
def test_bench_transforms_over_time(benchmark: Any, num_transforms: int, num_transforms_per_batch: int) -> None:
|
||||
rr.init("rerun_example_benchmark_transforms_individual")
|
||||
|
||||
rand_trans = np.array(np.random.rand(num_transforms, 3), dtype=np.float32)
|
||||
rand_quats = np.array(np.random.rand(num_transforms, 4), dtype=np.float32)
|
||||
rand_scales = np.array(np.random.rand(num_transforms, 3), dtype=np.float32)
|
||||
|
||||
print(rand_trans.shape)
|
||||
|
||||
if num_transforms_per_batch > 1:
|
||||
benchmark(
|
||||
test_bench_transforms_over_time_batched,
|
||||
rand_trans,
|
||||
rand_quats,
|
||||
rand_scales,
|
||||
num_transforms_per_batch,
|
||||
)
|
||||
else:
|
||||
benchmark(test_bench_transforms_over_time_individual, rand_trans, rand_quats, rand_scales)
|
||||
|
||||
|
||||
def log_transform3d_translation_mat3x3(data: Transform3DInput, static: bool) -> None:
|
||||
"""Log Transform3D with translation and mat3x3 for each entity at each time step."""
|
||||
# create a new, empty memory sink for the current recording
|
||||
rr.memory_recording()
|
||||
|
||||
start = time.perf_counter()
|
||||
|
||||
for time_index in range(data.num_time_steps):
|
||||
for entity_index in range(data.num_entities):
|
||||
entity_path = f"transform_{entity_index}"
|
||||
transform = rr.Transform3D(
|
||||
translation=data.translations[time_index, entity_index].tolist(),
|
||||
mat3x3=np.array(data.mat3x3s[time_index, entity_index], dtype=np.float32),
|
||||
)
|
||||
|
||||
if static:
|
||||
rr.log(entity_path, transform, static=True)
|
||||
else:
|
||||
rr.set_time("frame", sequence=time_index)
|
||||
rr.set_time("sim_time", duration=time_index * 0.01)
|
||||
rr.log(entity_path, transform)
|
||||
|
||||
elapsed = time.perf_counter() - start
|
||||
total_log_calls = data.num_entities * data.num_time_steps
|
||||
transforms_per_second = total_log_calls / elapsed
|
||||
print(f"Logged {total_log_calls} transforms in {elapsed:.2f}s ({transforms_per_second:.0f} transforms/second)")
|
||||
|
||||
|
||||
def test_bench_create_transform3d_translation_mat3x3(benchmark: Any) -> None:
|
||||
data = Transform3DInput.prepare(42, 1000, 100)
|
||||
benchmark(create_transform3d_translation_mat3x3, data, False)
|
||||
|
||||
|
||||
def create_transform3d_translation_mat3x3(data: Transform3DInput, static: bool) -> None:
|
||||
"""Just create Transform3D with translation and mat3x3 for each entity at each time step."""
|
||||
start = time.perf_counter()
|
||||
|
||||
transforms = []
|
||||
for time_index in range(data.num_time_steps):
|
||||
for entity_index in range(data.num_entities):
|
||||
transforms.append(
|
||||
rr.Transform3D(
|
||||
translation=data.translations[time_index, entity_index],
|
||||
mat3x3=np.array(data.mat3x3s[time_index, entity_index], dtype=np.float32),
|
||||
)
|
||||
)
|
||||
|
||||
elapsed = time.perf_counter() - start
|
||||
total_log_calls = data.num_entities * data.num_time_steps
|
||||
transforms_per_second = total_log_calls / elapsed
|
||||
print(f"Logged {total_log_calls} transforms in {elapsed:.2f}s ({transforms_per_second:.0f} transforms/second)")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
["num_entities", "num_time_steps", "static"],
|
||||
[
|
||||
pytest.param(10, 10_000, False, id="10entities-10000steps-temporal"),
|
||||
pytest.param(10, 10_000, True, id="10entities-10000steps-static"),
|
||||
],
|
||||
)
|
||||
def test_bench_transform3d_translation_mat3x3(
|
||||
benchmark: Any, num_entities: int, num_time_steps: int, static: bool
|
||||
) -> None:
|
||||
rr.init("rerun_example_benchmark_transform3d_translation_mat3x3")
|
||||
data = Transform3DInput.prepare(42, num_entities, num_time_steps)
|
||||
benchmark(log_transform3d_translation_mat3x3, data, static)
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Standalone execution (for profiling with py-spy, etc.)
|
||||
# -----------------------------------------------------------------------------
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="Run logging benchmarks standalone (useful for profiling)")
|
||||
parser.add_argument(
|
||||
"benchmark",
|
||||
choices=["transform3d"],
|
||||
help="Which benchmark to run",
|
||||
)
|
||||
parser.add_argument("--num-entities", type=int, default=10, help="Number of entities")
|
||||
parser.add_argument("--num-time-steps", type=int, default=1_000, help="Number of time steps")
|
||||
parser.add_argument("--static", action="store_true", help="Log as static data")
|
||||
parser.add_argument(
|
||||
"--connect",
|
||||
action="store_true",
|
||||
help="Connect to a running Rerun viewer via gRPC instead of using memory recording",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--create-only", action="store_true", help="Only create Transform3D instances without logging them"
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.benchmark == "transform3d":
|
||||
total_log_calls = args.num_entities * args.num_time_steps
|
||||
print(
|
||||
f"Preparing {total_log_calls} transforms ({args.num_entities} entities x {args.num_time_steps} time steps)…"
|
||||
)
|
||||
rr.init("rerun_example_benchmark_transform3d_translation_mat3x3", spawn=False)
|
||||
if args.connect:
|
||||
print("Connecting to Rerun viewer…")
|
||||
rr.connect_grpc()
|
||||
else:
|
||||
rr.memory_recording()
|
||||
data = Transform3DInput.prepare(42, args.num_entities, args.num_time_steps)
|
||||
print("Logging…")
|
||||
if args.create_only:
|
||||
create_transform3d_translation_mat3x3(data, args.static)
|
||||
else:
|
||||
log_transform3d_translation_mat3x3(data, args.static)
|
||||
@@ -0,0 +1,77 @@
|
||||
"""Per-call overhead micro-benchmarks for the rr.log() pipeline.
|
||||
|
||||
See README.md for usage.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
import rerun as rr
|
||||
import rerun_bindings as bindings
|
||||
from rerun._log import _log_components
|
||||
|
||||
ARCHETYPE_CASES = [
|
||||
pytest.param(lambda: rr.Scalars(42.0), id="Scalars"),
|
||||
pytest.param(
|
||||
lambda: rr.Points3D([[1, 2, 3]], colors=[0xFF0000FF], radii=[0.1]),
|
||||
id="Points3D",
|
||||
),
|
||||
pytest.param(
|
||||
lambda: rr.Transform3D(translation=[1, 2, 3], mat3x3=np.eye(3, dtype=np.float32)),
|
||||
id="Transform3D",
|
||||
),
|
||||
pytest.param(
|
||||
lambda: rr.Boxes3D(half_sizes=[[1, 2, 3]], colors=[0xFF0000FF]),
|
||||
id="Boxes3D",
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def _init() -> None:
|
||||
"""Common setup: init rerun + memory recording."""
|
||||
rr.init("rerun_example_micro_benchmark", spawn=False)
|
||||
rr.memory_recording()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("make_archetype", ARCHETYPE_CASES)
|
||||
def test_bench_micro_construct(benchmark: Any, make_archetype: Any) -> None:
|
||||
_init()
|
||||
benchmark(make_archetype)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("make_archetype", ARCHETYPE_CASES)
|
||||
def test_bench_micro_as_component_batches(benchmark: Any, make_archetype: Any) -> None:
|
||||
_init()
|
||||
archetype = make_archetype()
|
||||
benchmark(archetype.as_component_batches)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("make_archetype", ARCHETYPE_CASES)
|
||||
def test_bench_micro_log_components(benchmark: Any, make_archetype: Any) -> None:
|
||||
_init()
|
||||
batches = make_archetype().as_component_batches()
|
||||
benchmark(_log_components, "test_entity", batches)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("make_archetype", ARCHETYPE_CASES)
|
||||
def test_bench_micro_log_arrow_msg(benchmark: Any, make_archetype: Any) -> None:
|
||||
_init()
|
||||
batches = make_archetype().as_component_batches()
|
||||
instanced = {b.component_descriptor(): b.as_arrow_array() for b in batches if b.as_arrow_array() is not None}
|
||||
benchmark(bindings.log_arrow_msg, "test_entity", components=instanced, static_=False, recording=None)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("make_archetype", ARCHETYPE_CASES)
|
||||
def test_bench_micro_log(benchmark: Any, make_archetype: Any) -> None:
|
||||
_init()
|
||||
archetype = make_archetype()
|
||||
benchmark(rr.log, "test_entity", archetype)
|
||||
|
||||
|
||||
def test_bench_micro_set_time(benchmark: Any) -> None:
|
||||
_init()
|
||||
benchmark(rr.set_time, "frame", sequence=42)
|
||||
Reference in New Issue
Block a user