rerun-io--rerun/tests/python/plot_dashboard_stress/main.py

#!/usr/bin/env python3
"""
Plot dashboard stress test.

Usage:
-----
```
pixi run py-plot-dashboard --help
```

Example:
-------
```
pixi run py-plot-dashboard --num-plots 10 --num-series-per-plot 5 --num-points-per-series 5000 --freq 1000
```

"""

from __future__ import annotations

import argparse
import math
import time
from typing import Any, cast

import numpy as np

import rerun as rr  # pip install rerun-sdk
import rerun.blueprint as rrb

parser = argparse.ArgumentParser(description="Plot dashboard stress test")
rr.script_add_args(parser)

parser.add_argument("--num-plots", type=int, default=1, help="How many different plots?")
parser.add_argument(
    "--num-series-per-plot",
    type=int,
    default=1,
    help="How many series in each single plot?",
)
parser.add_argument(
    "--num-points-per-series",
    type=int,
    default=100000,
    help="How many points in each single series?",
)
parser.add_argument(
    "--freq",
    type=float,
    default=1000,
    help="Frequency of logging (applies to all series)",
)
parser.add_argument(
    "--temporal-batch-size",
    type=int,
    default=None,
    help="Number of rows to include in each log call",
)
parser.add_argument(
    "--blueprint",
    action="store_true",
    help="Setup a blueprint for a 5s window",
)

order = [
    "forwards",
    "backwards",
    "random",
]
parser.add_argument(
    "--order",
    type=str,
    default=order[0],
    help="What order to log the data in (applies to all series)",
    choices=order,
)

series_type = [
    "gaussian-random-walk",
    "sin-uniform",
]
parser.add_argument(
    "--series-type",
    type=str,
    default=series_type[0],
    choices=series_type,
    help="The method used to generate time series",
)


# TODO(cmc): could have flags to add attributes (color, radius...) to put some more stress
# on the line fragmenter.

args = parser.parse_args()


def main() -> None:
    rr.script_setup(args, "rerun_example_plot_dashboard_stress")

    plot_paths = [f"plot_{i}" for i in range(args.num_plots)]
    series_paths = [f"series_{i}" for i in range(args.num_series_per_plot)]

    if args.blueprint:
        print("logging blueprint!")
        rr.send_blueprint(
            rrb.Blueprint(
                rrb.Grid(*[
                    rrb.TimeSeriesView(
                        name=p,
                        origin=f"/{p}",
                        time_ranges=rrb.VisibleTimeRanges(
                            timeline="sim_time",
                            start=rrb.TimeRangeBoundary.cursor_relative(offset=rr.TimeInt(seconds=-2.5)),
                            end=rrb.TimeRangeBoundary.cursor_relative(offset=rr.TimeInt(seconds=2.5)),
                        ),
                    )
                    for p in plot_paths
                ]),
                rrb.BlueprintPanel(state="collapsed"),
                rrb.SelectionPanel(state="collapsed"),
            ),
        )

    time_per_sim_step = 1.0 / args.freq
    stop_time = args.num_points_per_series * time_per_sim_step

    if args.order == "forwards":
        sim_times = np.arange(0, stop_time, time_per_sim_step)
    elif args.order == "backwards":
        sim_times = np.arange(0, stop_time, time_per_sim_step)[::-1]
    else:
        sim_times = np.random.randint(0, args.num_points_per_series)

    num_series = len(plot_paths) * len(series_paths)
    time_per_tick = time_per_sim_step
    scalars_per_tick = num_series
    if args.temporal_batch_size is not None:
        time_per_tick *= args.temporal_batch_size
        scalars_per_tick *= args.temporal_batch_size

    expected_total_freq = args.freq * num_series

    values_shape = (
        len(sim_times),
        len(plot_paths),
        len(series_paths),
    )
    if args.series_type == "gaussian-random-walk":
        values = np.cumsum(np.random.normal(size=values_shape), axis=0)
    elif args.series_type == "sin-uniform":
        values = np.sin(np.random.uniform(0, math.pi, size=values_shape))
    else:
        # Just generate random numbers rather than crash
        values = np.random.normal(size=values_shape)

    if args.temporal_batch_size is None:
        ticks: Any = enumerate(sim_times)
    else:
        offsets = range(0, len(sim_times), args.temporal_batch_size)
        ticks = zip(
            offsets,
            (sim_times[offset : offset + args.temporal_batch_size] for offset in offsets),
            strict=False,
        )

    time_column = None

    total_start_time = time.time()
    total_num_scalars = 0

    tick_start_time = time.time()
    max_load = 0.0

    for index, sim_time in ticks:
        if args.temporal_batch_size is None:
            rr.set_time("sim_time", duration=sim_time)
        else:
            time_column = rr.TimeColumn("sim_time", duration=sim_time)

        # Log
        for plot_idx, plot_path in enumerate(plot_paths):
            for series_idx, series_path in enumerate(series_paths):
                if args.temporal_batch_size is None:
                    value = values[index, plot_idx, series_idx]
                    rr.log(f"{plot_path}/{series_path}", rr.Scalars(value))
                else:
                    value_index = slice(index, index + args.temporal_batch_size)
                    rr.send_columns(
                        f"{plot_path}/{series_path}",
                        indexes=[cast("rr.TimeColumn", time_column)],
                        columns=rr.Scalars.columns(scalars=values[value_index, plot_idx, series_idx]),
                    )

        # Measure how long this took and how high the load was.

        elapsed = time.time() - tick_start_time
        max_load = max(max_load, elapsed / time_per_tick)

        # Throttle

        sleep_duration = time_per_tick - elapsed
        if sleep_duration > 0.0:
            sleep_start_time = time.time()
            time.sleep(sleep_duration)
            sleep_elapsed = time.time() - sleep_start_time

            # We will very likely be put to sleep for more than we asked for, and therefore need
            # to pay off that debt in order to meet our frequency goal.
            sleep_debt = sleep_elapsed - sleep_duration
            tick_start_time = time.time() - sleep_debt
        else:
            tick_start_time = time.time()

        # Progress report
        #
        # Must come after throttle since we report every wall-clock second:
        # If ticks are large & fast, then after each send we run into throttle.
        # So if this was before throttle, we'd not report the first tick no matter how large it was.

        total_num_scalars += scalars_per_tick
        total_elapsed = time.time() - total_start_time

        if total_elapsed >= 1.0:
            print(
                f"logged {total_num_scalars} scalars over {round(total_elapsed, 3)}s \
(freq={round(total_num_scalars / total_elapsed, 3)}Hz, expected={round(expected_total_freq, 3)}Hz, \
load={round(max_load * 100.0, 3)}%)",
            )

            elapsed_debt = total_elapsed % 1  # just keep the fractional part
            total_start_time = time.time() - elapsed_debt
            total_num_scalars = 0
            max_load = 0.0

    if total_num_scalars > 0:
        total_elapsed = time.time() - total_start_time
        print(
            f"logged {total_num_scalars} scalars over {round(total_elapsed, 3)}s \
(freq={round(total_num_scalars / total_elapsed, 3)}Hz, expected={round(expected_total_freq, 3)}Hz, \
load={round(max_load * 100.0, 3)}%)",
        )

    rr.script_teardown(args)


if __name__ == "__main__":
    main()