Benchmark callbacks

RxInfer provides a built-in callback structure called RxInferBenchmarkCallbacks for collecting timing information during the inference procedure. This structure aggregates timestamps across multiple inference runs, allowing you to track performance statistics (min/max/average/etc.) of your model's creation and inference procedure. For general information about the callbacks system, see Callbacks.

Basic usage

using RxInfer

@model function iid_normal(y)
    μ  ~ Normal(mean = 0.0, variance = 100.0)
    γ  ~ Gamma(shape = 1.0, rate = 1.0)
    y .~ Normal(mean = μ, precision = γ)
end

init = @initialization begin
    q(μ) = vague(NormalMeanVariance)
end

# Warm up to avoid measuring compilation time

# Create a benchmark callbacks instance to track performance
benchmark_callbacks = RxInferBenchmarkCallbacks()

# Run inference multiple times to gather statistics
for i in 1:5
    infer(
        model = iid_normal(),
        data = (y = randn(100),),
        constraints = MeanField(),
        iterations = 5,
        initialization = init,
        callbacks = benchmark_callbacks,
    )
end

The benchmark callbacks instance accumulates timestamps across multiple calls to infer, making it easy to collect performance statistics over many runs.

Displaying results

Install PrettyTables.jl to display the collected statistics in a nicely formatted table:

using PrettyTables

PrettyTables.pretty_table(benchmark_callbacks)

RxInfer inference benchmark statistics: 5 evaluations
╭────────────────┬────────────┬────────────┬────────────┬────────────┬──────────
│      Operation │        Min │        Max │       Mean │     Median │         ⋯
├────────────────┼────────────┼────────────┼────────────┼────────────┼──────────
│ Model creation │   1.836 ms │   2.113 ms │   1.950 ms │   1.922 ms │ 106.177 ⋯
│      Inference │ 422.330 μs │ 498.342 μs │ 452.202 μs │ 448.298 μs │  28.140 ⋯
│      Iteration │  70.411 μs │ 141.755 μs │  80.567 μs │  75.461 μs │  15.892 ⋯
╰────────────────┴────────────┴────────────┴────────────┴────────────┴──────────
                                                                1 column omitted

Accessing from model metadata

After model creation, the benchmark callbacks instance is automatically saved into the model's metadata under the :benchmark key. This makes it accessible from the inference result without needing to hold onto the callbacks object separately:

result = infer(
    model = iid_normal(),
    data = (y = randn(100),),
    constraints = MeanField(),
    iterations = 5,
    initialization = init,
    callbacks = RxInferBenchmarkCallbacks(),
)


benchmark = result.model.metadata[:benchmark]
println(benchmark)

RxInferBenchmarkCallbacks (1evaluations, use `pretty_table` from `PrettyTables.jl` to display the statistics in a tabular format)

Tracked events

The RxInferBenchmarkCallbacks structure collects timestamps at the following stages:

RxInferBenchmarkCallbacks declares these events to the engine with ReactiveMP.listens, so the engine builds no other event objects for it: the message rule calls, products and marginal computations cost nothing extra under benchmark = true. To inspect individual rule calls, with the algorithm each node ran under and the log scale of each message, use the Trace callbacks.

Buffer capacity

By default, the structure uses circular buffers with a capacity of RxInfer.DEFAULT_BENCHMARK_CALLBACKS_BUFFER_CAPACITY entries. This limits memory usage in long-running applications. You can change the capacity:

# Store up to 10000 benchmark entries
large_buffer_callbacks = RxInferBenchmarkCallbacks(capacity = 10_000)

Programmatic access to statistics

Use RxInfer.get_benchmark_stats to retrieve the raw statistics matrix:

# Use the previously populated benchmark_callbacks
stats = RxInfer.get_benchmark_stats(benchmark_callbacks)


for row in eachrow(stats)
    println(row[1], ": min=", round(row[2] / 1e6, digits=2), "ms, mean=", round(row[4] / 1e6, digits=2), "ms")
end

Model creation: min=1.84ms, mean=1.95ms
Inference: min=0.42ms, mean=0.45ms
Iteration: min=0.07ms, mean=0.08ms

The matrix contains the following columns:

1. Operation name (String)
2. Minimum time (Float64, nanoseconds)
3. Maximum time (Float64, nanoseconds)
4. Mean time (Float64, nanoseconds)
5. Median time (Float64, nanoseconds)
6. Standard deviation (Float64, nanoseconds)

API Reference

RxInferBenchmarkCallbacks(; capacity = RxInfer.DEFAULT_BENCHMARK_CALLBACKS_BUFFER_CAPACITY)

# Create a callbacks instance to track performance
callbacks = RxInferBenchmarkCallbacks()

# Run inference multiple times to gather statistics
for _ in 1:10
    result = infer(
        model = my_model(),
        data = my_data,
        callbacks = callbacks
    )
end

# Access the benchmark callbacks from the inference result
result.model.metadata[:benchmark] === callbacks # true

# Display the timing statistics (you need to install `PrettyTables.jl` to use `pretty_table` function)
using PrettyTables

PrettyTables.pretty_table(callbacks)

get_benchmark_stats(callbacks::RxInferBenchmarkCallbacks)

DEFAULT_BENCHMARK_CALLBACKS_BUFFER_CAPACITY