chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -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)
|
||||
Reference in New Issue
Block a user