helios.core.loggers.wandb

Classes

WandbArgs	Arguments for constructing a WandbWriter.
WandbWriter	Wrapper for the Weights & Biases wandb.init run.

Module Contents

class helios.core.loggers.wandb.WandbArgs

Bases: TypedDict

Arguments for constructing a WandbWriter.

project is the only required key; all others are optional.

Parameters:

• project – W&B project name.

• name – display name for the run shown in the W&B UI. If not provided, defaults to run_name provided by setup().

• config – hyper-parameter dictionary to associate with the run.

• extra_args – additional keyword arguments forwarded verbatim to wandb.init().

project: Required[str]

name: str

config: dict[str, Any]

extra_args: dict[str, Any]

class helios.core.loggers.wandb.WandbWriter(project: str, name: str | None = None, config: dict[str, Any] | None = None, extra_args: dict[str, Any] | None = None)

Bases: helios.core.loggers.base.Logger

Wrapper for the Weights & Biases wandb.init run.

Data for the logger will be placed under log_root/wandb. When resuming, the original run ID is restored and new data is appended to it.

Requires the wandb package. Install it with:

pip install wandb

Parameters:

• project – W&B project name.

• name – (optional) display name for the run.

• config – (optional) hyper-parameter configuration dictionary to associate with the run.

• extra_args – (optional) extra keyword arguments forwarded verbatim to wandb.init(). Defaults to {}.

setup(run_name: str, log_root: pathlib.Path | None, is_resume: bool) → None

Finish configuring the WandbWriter.

In particular, this function will call wandb.init(). If a run ID was previously saved, then it will be forwarded to W&B so the run continues in place. In distributed training, the writer will only be created on rank 0.

Parameters:

• run_name – the name of the current run; used as the W&B run name when no explicit name was given in __init__.

• log_root – root directory for logs. W&B data will be written under log_root/wandb/. None lets W&B choose its own default directory.

• is_resume – True when continuing a previous run.

flush() → None

No-op, W&B syncs its own data automatically.

close() → None

Finish and close the W&B run.

state_dict() → dict[str, Any]

Return a dictionary containing the writer state.

The state will be saved under a key called "run_id" holding the current run ID. If W&B was disabled, then None is stored.

Returns:

A dictionary with the logger state.

load_state_dict(state_dict: dict[str, Any]) → None

Restore the writer state from a previously saved dictionary.

Parameters:

state_dict – the state dictionary returned by state_dict().