Skip to content

mlvbench.monitoring

Monitoring utilities for training loops.

ThroughputMonitor

Monitor wall-clock time per phase, throughput and GPU memory usage.

The monitor is driven by the training loop, which marks the phases to account for and registers each optimizer step:

monitor = ThroughputMonitor(device)

for batch in monitor.iterate(dataloader):
    with monitor.phase("features"):
        features = model.forward_features(batch["image"])
    with monitor.phase("probes"):
        ...
    monitor.record_samples(batch["image"].shape[0])
    stats = monitor.step(step)

Every record_every_n_steps steps, the accumulated timings are appended to records as a snapshot, returned by step(), and the accumulators are reset. Snapshots therefore cover disjoint intervals of the run and can be saved with save(). The monitor only collects the statistics; reporting them is up to the caller.

Phase timings are attributed correctly only if the device is synchronized, so the monitor synchronizes CUDA devices when a phase ends. This costs time, which is why monitoring can be disabled entirely by passing record_every_n_steps=None. All methods then become no-ops and no synchronization is performed.

enabled property

Whether the monitor accumulates timings.

records property

The snapshots recorded so far, one per logged step.

__init__(device, record_every_n_steps=50, clock=time.perf_counter)

Initialize the monitor.

Parameters:

Name Type Description Default
device device | str

The device the monitored work runs on. CUDA devices are synchronized at phase boundaries and reported with their memory usage.

 required
record_every_n_steps int | None

Record a snapshot every N steps. If None, monitoring is disabled entirely.

 50
clock Callable[[], float]

Function returning the current time in seconds. Intended for testing.

 perf_counter

iterate(iterable)

Iterate over iterable, accounting the time spent in it as "data".

Only the time needed to produce the next item is accounted for, not the time the loop body takes to process it.

Parameters:

Name Type Description Default
iterable Iterable[Any]

The iterable to monitor, typically a data loader.

 required

Yields:

Type Description
Any

The items of iterable, unchanged.

phase(name)

Account the time spent in the enclosed block to the phase name.

Nothing is recorded if the enclosed block raises.

Parameters:

Name Type Description Default
name str

Name of the phase. Phases are reported in the order first seen.

 required

Yields:

Type Description
None

None.

record_samples(count)

Register count processed samples for the throughput computation.

reset()

Reset the accumulated timings, sample count and peak memory usage.

step(step)

Register a completed step and record a snapshot every record_every_n_steps.

Parameters:

Name Type Description Default
step int

The current step, used for reporting and to determine whether a snapshot is due.

 required

Returns:

Type Description
dict[str, Any] | None

The recorded snapshot if one is due at this step, None otherwise. The snapshot holds the step, the accumulated per-phase durations (t_<name> in seconds), the throughput (samples_per_sec) and the current and peak GPU memory usage (gpu_mem_gb, gpu_mem_peak_gb).