Skip to navigation
ReferenceFull Library ReferenceNemo AutomodelNemo AutomodelComponentsLoggersnemo_automodel.components.loggers.loggers

nemo_automodel.components.loggers.loggers

View as Markdown

Typed configs for remote loggers (WandB, Trackio, MLflow, Comet, SwanLab).

Each logger config is a plain dataclass exposing its YAML-configurable fields plus a build(...) method that initialises and returns the logger / run object. Loggers are a closed, section-named set (no _target_ dispatch), so there is no free builder function — config.build(...) is the entry point.

Module Contents

Classes

NameDescription
CometConfigUser-facing Comet ML configuration.
MLflowConfigUser-facing MLflow configuration (maps to the YAML mlflow: block).
SwanLabConfigAlso log the W&B run to SwanLab (maps to the YAML wandb.swanlab: block).
TrackioConfigUser-facing Trackio configuration (maps to the YAML trackio: block).
WandbConfigUser-facing WandB configuration (maps to the YAML wandb: block).

Functions

NameDescription
_builtin_scalarReturn a 0-d numpy / torch value as the equivalent Python scalar; anything else unchanged.
_patch_mirror_valuesMake swanlab.sync_wandb forward every value W&B logs.

Data

_SWANLAB_MIRROR_INSTALLED

logger

API

class nemo_automodel.components.loggers.loggers.CometConfig(
    project_name: str = 'automodel',
    workspace: str | None = None,
    api_key: str | None = None,
    experiment_name: str | None = None,
    tags: list[str] = list(),
    auto_metric_logging: bool = False
)
Dataclass

User-facing Comet ML configuration.

api_key
str | None = None

Comet API key. None reads from COMET_API_KEY env var.

auto_metric_logging
bool = False

Enable Comet’s automatic metric logging.

experiment_name
str | None = None

Display name for this experiment run.

project_name
str = 'automodel'

Comet project name.

tags
list[str] = field(default_factory=list)

List of string tags attached to the experiment.

workspace
str | None = None

Comet workspace. None uses the default.

nemo_automodel.components.loggers.loggers.CometConfig.build(
    model_name: str | None = None
) -> typing.Any

Initialise Comet ML and return the logger (active on rank 0).

When model_name is provided a model:<name> tag is appended and, if experiment_name is empty, one is derived from the model name.

Parameters:

model_name
str | NoneDefaults to None

Optional model name used to tag the run and derive an experiment name when none is set.

Returns: Any

A CometLogger instance.

class nemo_automodel.components.loggers.loggers.MLflowConfig(
    experiment_name: str = 'automodel-experiment',
    run_name: str = '',
    tracking_uri: str | None = None,
    artifact_location: str | None = None,
    tags: dict[str, str] = dict(),
    resume: bool = True,
    description: str | None = None,
    flatten_depth: int | None = 1
)
Dataclass

User-facing MLflow configuration (maps to the YAML mlflow: block).

artifact_location
str | None = None

Root artifact store URI for the experiment.

description
str | None = None

Free-text description shown in the MLflow UI (sets the mlflow.note.content tag).

experiment_name
str = 'automodel-experiment'

MLflow experiment name.

flatten_depth
int | None = 1

Nesting depth for mlflow.log_params. 1 (default) splits one level; None is fully recursive.

resume
bool = True

When True (default), look for a mlflow_run_id sidecar in the checkpoint dir and resume that run.

run_name
str = ''

Display name for the run.

tags
dict[str, str] = field(default_factory=dict)

Dictionary of string tags attached to the run.

tracking_uri
str | None = None

MLflow tracking server URI. None delegates to MLFLOW_TRACKING_URI and then MLflow’s local default (sqlite:///mlflow.db in current releases).

nemo_automodel.components.loggers.loggers.MLflowConfig.build(
    checkpoint_dir: str | None = None,
    run_config: collections.abc.Mapping[str, typing.Any] | None = None
) -> typing.Any

Initialise MLflow on rank 0 and start (or resume) a run.

Installs a sys.excepthook so crashed jobs report as FAILED rather than FINISHED. On non-rank-0 processes returns None.

Parameters:

checkpoint_dir
str | NoneDefaults to None

Checkpoint directory used to persist / read the mlflow_run_id sidecar for run resumption.

run_config
Mapping[str, Any] | NoneDefaults to None

Full training config dict logged as MLflow params and as a config.yaml artifact.

Returns: Any

Active mlflow.entities.Run on rank 0, or None.

class nemo_automodel.components.loggers.loggers.SwanLabConfig(
    mode: str = 'online',
    workspace: str | None = None,
    log_dir: str | None = None,
    upload_to_wandb: bool = True
)
Dataclass

Also log the W&B run to SwanLab (maps to the YAML wandb.swanlab: block).

SwanLab is an additional experiment tracker next to W&B; it is also useful where Weights & Biases is not reachable. Instead of a parallel set of logging calls in every recipe, it patches wandb through swanlab.sync_wandb so every wandb.init / wandb.log / wandb.config.update / finish a recipe makes is forwarded to SwanLab. SwanLab therefore records the items W&B records: the run name, project, tags, group, notes, the full run config, and every logged number, string, and image at the same step. See the SwanLab logging guide for the differences (e.g. W&B tables and histograms are not mirrored).

The SwanLab API key is read by the SDK from SWANLAB_API_KEY (or a prior swanlab login); it is never part of the YAML.

log_dir
str | None = None

Local directory for SwanLab run files. None uses the SDK default.

mode
str = 'online'

SwanLab mode: online, local, offline, or disabled.

upload_to_wandb
bool = True

Keep uploading to W&B as usual (default). Set it to False where W&B is not reachable: W&B then runs in offline mode and only SwanLab receives the run.

workspace
str | None = None

SwanLab workspace (user or organization). None uses the account’s default workspace.

nemo_automodel.components.loggers.loggers.SwanLabConfig.mirror_wandb() -> None

Patch wandb so the next wandb.init and all its logging also go to SwanLab.

Must run before wandb.init. The patch is process-wide and cannot be undone, so it is installed once: patching twice would forward every call to SwanLab twice, and a later call with different settings is ignored with a warning.

class nemo_automodel.components.loggers.loggers.TrackioConfig(
    project: str = 'automodel',
    name: str = '',
    group: str | None = None,
    space_id: str | None = None,
    extra: dict[str, typing.Any] = dict()
)
Dataclass

User-facing Trackio configuration (maps to the YAML trackio: block).

Trackio <https://github.com/gradio-app/trackio>_ is a local-first experiment tracker with a wandb-style API. Runs are stored locally (TRACKIO_DIR) unless space_id syncs them to a Hugging Face Space. Any other key under the YAML trackio: block (e.g. resume, private, server_url, bucket_id, auto_log_gpu) is a valid trackio.init() kwarg and is preserved verbatim in extra.

extra
dict[str, Any] = field(default_factory=dict)

Any additional trackio.init() kwargs carried through unfiltered.

group
str | None = None

Group name for related runs.

name
str = ''

Display name for the run. When empty, build derives one from the model name.

project
str = 'automodel'

Trackio project name.

space_id
str | None = None

Hugging Face Space ("user/space" or "space") to sync the project to. None keeps the run local (or uses TRACKIO_SPACE_ID).

nemo_automodel.components.loggers.loggers.TrackioConfig.build(
    run_config: collections.abc.Mapping[str, typing.Any] | None = None,
    model_name: str | None = None
) -> 'TrackioLogger | None'

Initialise a Trackio run on rank 0 and return its logger.

Parameters:

run_config
Mapping[str, Any] | NoneDefaults to None

JSON-compatible run config, already serialized at the recipe boundary (e.g. RecipeConfig.to_yaml_dict(use_orig_values=True)).

model_name
str | NoneDefaults to None

Optional model name used to derive the run name when name is empty.

Returns: 'TrackioLogger | None'

class:~nemo_automodel.components.loggers.trackio_utils.TrackioLogger on rank 0, None on other

nemo_automodel.components.loggers.loggers.TrackioConfig.from_kwargs(
    kwargs: typing.Any = {}
) -> 'TrackioConfig'
classmethod

Build from a flat kwargs dict, routing unknown keys to extra.

class nemo_automodel.components.loggers.loggers.WandbConfig(
    project: str = 'automodel',
    entity: str | None = None,
    name: str = '',
    group: str | None = None,
    tags: list[str] = list(),
    notes: str | None = None,
    swanlab: nemo_automodel.components.loggers.loggers.SwanLabConfig | None = None,
    extra: dict[str, typing.Any] = dict()
)
Dataclass

User-facing WandB configuration (maps to the YAML wandb: block).

The named fields are the common ones; any other key under the YAML wandb: block (e.g. mode, dir, resume) is a valid wandb.init() kwarg and is preserved verbatim in extra and forwarded to wandb.init().

entity
str | None = None

WandB team / entity. None uses the default from the wandb config or WANDB_ENTITY env var.

extra
dict[str, Any] = field(default_factory=dict)

Any additional wandb.init() kwargs (mode, dir, …) carried through unfiltered.

group
str | None = None

Group name for related runs.

name
str = ''

Display name for the run. When empty, build derives one from the model name.

notes
str | None = None

Free-text notes shown in the WandB UI.

project
str = 'automodel'

WandB project name.

swanlab
SwanLabConfig | None = None

Optional SwanLabConfig that mirrors this run to SwanLab (the wandb.swanlab: sub-block).

tags
list[str] = field(default_factory=list)

List of string tags attached to the run.

nemo_automodel.components.loggers.loggers.WandbConfig.build(
    run_config: collections.abc.Mapping[str, typing.Any] | None = None,
    model_name: str | None = None
) -> typing.Any

Initialise WandB and return the run.

Parameters:

run_config
Mapping[str, Any] | NoneDefaults to None

Full training config dict logged to the WandB run.

model_name
str | NoneDefaults to None

Optional model name used to derive the run name when name is empty.

Returns: Any

Initialised wandb.Run.

nemo_automodel.components.loggers.loggers.WandbConfig.from_kwargs(
    kwargs: typing.Any = {}
) -> 'WandbConfig'
classmethod

Build from a flat kwargs dict, routing unknown keys to extra.

Keys matching a named field are assigned directly; everything else (valid wandb.init params such as mode/dir that are not first-class fields here) is preserved in extra.

nemo_automodel.components.loggers.loggers._builtin_scalar(
    value: typing.Any
) -> typing.Any

Return a 0-d numpy / torch value as the equivalent Python scalar; anything else unchanged.

nemo_automodel.components.loggers.loggers._patch_mirror_values(
    swanlab: typing.Any
) -> None

Make swanlab.sync_wandb forward every value W&B logs.

The mirror (swanlab 0.10) forwards only Python int/float/bool/str values and SwanLab rejects strings as scalars, so a 0-d numpy or torch value (e.g. grad_norm) is dropped and a string (e.g. the per-step timestamp) fails to log. Unwrap 0-d values in Run.log before the mirror filters them (the same hook sync_wandb patches; W&B stores the same number either way) and log strings to SwanLab as text.

nemo_automodel.components.loggers.loggers._SWANLAB_MIRROR_INSTALLED: 'SwanLabConfig | None' = None
 
nemo_automodel.components.loggers.loggers.logger = logging.getLogger(__name__)