wehub-resource-sync
254 lines
9.3 KiB
Python
254 lines
9.3 KiB
Python
"""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)
|