nvalchemi.hooks.TensorBoardReporter

class nvalchemi.hooks.TensorBoardReporter(log_dir, *, custom_scalars=None, include_losses=True, include_optimizer_lrs=True, rank_reduction=None, tag_prefix=None, flush=True, rank_zero_only=True, writer=None)

Write scalar reporting snapshots to TensorBoard.

Parameters:

• log_dir (str | Path) – TensorBoard log directory.

• custom_scalars (Mapping[str, ScalarCallback] | None, optional) – Additional scalar callbacks passed to collect_scalars().

• include_losses (bool, default True) – When True, include loss scalars from the hook context.

• include_optimizer_lrs (bool, default True) – When True, include optimizer learning rates from the hook context.

• rank_reduction (torch.distributed.ReduceOp | {"none", "mean", "sum", "min", "max"} | None, default None) – Optional distributed reduction applied to scalars before writing. String values are normalized to torch.distributed.ReduceOp. Reduction requires every rank to call this reporter; only rank zero writes the reduced snapshot.

• tag_prefix (str | None, optional) – Optional prefix prepended to every TensorBoard tag.

• flush (bool, default True) – Flush the writer after every report event.

• rank_zero_only (bool, default True) – Request rank-zero-only dispatch from ReportingOrchestrator. When False and rank_reduction="none", log_dir must contain "{rank}" or "{global_rank}" so every rank writes its own event directory.

• writer (TensorBoardWriter | None, optional) – Preconstructed writer. This is mainly useful for tests or integrations that own writer construction.

close()

Close the writer if it is open.

Return type:

None

report(ctx, stage, state)

Write one scalar snapshot to TensorBoard.

Parameters:

• ctx (HookContext) – Workflow hook context.

• stage (Enum) – Hook stage being reported.

• state (ReportingState) – Shared reporting state from the orchestrator.

Return type:

None