`nemo_rl.utils.logger`#

Module Contents#

Classes#

`WandbConfig`
`SwanlabConfig`
`TensorboardConfig`
`MLflowConfig`
`GPUMonitoringConfig`
`LoggerConfig`
`LoggerInterface`	Abstract base class for logger backends.
`TensorboardLogger`	Tensorboard logger backend.
`WandbLogger`	Weights & Biases logger backend.
`SwanlabLogger`	SwanLab logger backend.
`GpuMetricSnapshot`
`RayGpuMonitorLogger`	Monitor GPU utilization across a Ray cluster and log metrics to a parent logger.
`MLflowLogger`	MLflow logger backend.
`Logger`	Main logger class that delegates to multiple backend loggers.

Functions#

`_merge_generation_logger_workers`	Merge per-worker async-vLLM generation-logger lists into one list per metric.
`_summarize_list`	Reduce a list to bounded summary stats for scalar-only backends (MLflow).
`flatten_dict`	Flatten a nested dictionary.
`configure_rich_logging`	Configure rich logging for more visually appealing log output.
`print_message_log_samples`	Visualization for message logs and rewards using a more visual approach with emoji indicators and horizontal layout.
`get_next_experiment_dir`	Create a new experiment directory with an incremented ID.
`log_container_init_timing`	Log pre-Python container timing from environment variables.

Data#

_rich_logging_configured

API#

nemo_rl.utils.logger._rich_logging_configured#: False

class nemo_rl.utils.logger.WandbConfig#

Bases: typing.TypedDict

project: NotRequired[str]#: None

name: NotRequired[str]#: None

entity: NotRequired[str]#: None

class nemo_rl.utils.logger.SwanlabConfig#

Bases: typing.TypedDict

project: NotRequired[str]#: None

name: NotRequired[str]#: None

class nemo_rl.utils.logger.TensorboardConfig#

Bases: typing.TypedDict

log_dir: NotRequired[str]#: None

class nemo_rl.utils.logger.MLflowConfig#

Bases: typing.TypedDict

experiment_name: NotRequired[str | None]#: None

run_id: NotRequired[str | None]#: None

run_name: NotRequired[str | None]#: None

tracking_uri: NotRequired[str | None]#: None

artifact_location: NotRequired[str | None]#: None

class nemo_rl.utils.logger.GPUMonitoringConfig#

Bases: typing.TypedDict

collection_interval: int | float#: None

flush_interval: int | float#: None

class nemo_rl.utils.logger.LoggerConfig#

Bases: typing.TypedDict

log_dir: str#: None

wandb_enabled: bool#: None

swanlab_enabled: bool#: None

tensorboard_enabled: bool#: None

mlflow_enabled: bool#: None

wandb: nemo_rl.utils.logger.WandbConfig#: None

tensorboard: NotRequired[nemo_rl.utils.logger.TensorboardConfig]#: None

swanlab: NotRequired[nemo_rl.utils.logger.SwanlabConfig]#: None

mlflow: NotRequired[nemo_rl.utils.logger.MLflowConfig]#: None

monitor_gpus: bool#: None

gpu_monitoring: nemo_rl.utils.logger.GPUMonitoringConfig#: None

num_val_samples_to_print: NotRequired[int]#: None

class nemo_rl.utils.logger.LoggerInterface#

Bases: abc.ABC

Abstract base class for logger backends.

abstractmethod log_metrics( metrics: dict[str, Any], step: int, prefix: Optional[str] = '', step_metric: Optional[str] = None, step_finished: bool = False, ) → None#: Log a dictionary of metrics.

abstractmethod log_hyperparams(params: Mapping[str, Any]) → None#: Log dictionary of hyperparameters.

abstractmethod log_histogram( histogram: list[Any], step: int, name: str, ) → None#: Log histogram metrics.

abstractmethod log_plot( figure: matplotlib.pyplot.Figure, step: int, name: str, ) → None#: Log a matplotlib figure.

class nemo_rl.utils.logger.TensorboardLogger( cfg: nemo_rl.utils.logger.TensorboardConfig, log_dir: Optional[str] = None, )#

Bases: nemo_rl.utils.logger.LoggerInterface

Tensorboard logger backend.

Initialization

static _coerce_to_scalar( value: Any, ) → int | float | bool | str | None#

Coerce a value to a Python scalar for TensorBoard logging.

Returns the coerced value, or None if it can’t be converted to a scalar.

log_metrics( metrics: dict[str, Any], step: int, prefix: Optional[str] = '', step_metric: Optional[str] = None, step_finished: bool = False, ) → None#

Log metrics to Tensorboard.

Parameters:

metrics – Dict of metrics to log
step – Global step value
prefix – Optional prefix for metric names
step_metric – Optional step metric name (ignored in TensorBoard)

log_histogram( histogram: list[Any], step: int, name: str, ) → None#: Log histogram metrics to Tensorboard.

log_hyperparams(params: Mapping[str, Any]) → None#

Log hyperparameters to Tensorboard.

Parameters:: params – Dictionary of hyperparameters to log

log_plot( figure: matplotlib.pyplot.Figure, step: int, name: str, ) → None#

Log a plot to Tensorboard.

Parameters:

plot_data – Dictionary of plot data
step – Global step value

class nemo_rl.utils.logger.WandbLogger( cfg: nemo_rl.utils.logger.WandbConfig, log_dir: Optional[str] = None, )#

Bases: nemo_rl.utils.logger.LoggerInterface

Weights & Biases logger backend.

Initialization

_log_diffs()#

Log git diffs to wandb.

This function captures and logs two types of diffs:

Uncommitted changes (working tree diff against HEAD)
All changes (including uncommitted) against the main branch

Each diff is saved as a text file in a wandb artifact.

_log_code()#

Log code that is tracked by git to wandb.

This function gets a list of all files tracked by git in the project root and manually uploads them to the current wandb run as an artifact.

define_metric( name: str, step_metric: Optional[str] = None, ) → None#

Define a metric with custom step metric.

Parameters:

name – Name of the metric or pattern (e.g. ‘ray/*’)
step_metric – Optional name of the step metric to use

log_metrics( metrics: dict[str, Any], step: int, prefix: Optional[str] = '', step_metric: Optional[str] = None, step_finished: bool = False, ) → None#

Log metrics to wandb.

Parameters:

metrics – Dict of metrics to log
step – Global step value
prefix – Optional prefix for metric names
step_metric – Optional name of a field in metrics to use as step instead of the provided step value

log_hyperparams(params: Mapping[str, Any]) → None#

Log hyperparameters to wandb.

Parameters:: params – Dict of hyperparameters to log

log_plot( figure: matplotlib.pyplot.Figure, step: int, name: str, ) → None#

Log a plot to wandb.

Parameters:

figure – Matplotlib figure to log
step – Global step value

log_histogram( histogram: list[Any], step: int, name: str, ) → None#

Log histogram metrics to wandb.

Parameters:

histogram – List of histogram values
step – Global step value
name – Name of the metric

class nemo_rl.utils.logger.SwanlabLogger( cfg: nemo_rl.utils.logger.SwanlabConfig, log_dir: Optional[str] = None, )#

Bases: nemo_rl.utils.logger.LoggerInterface

SwanLab logger backend.

Initialization

Initialize the SwanlabLogger by starting a Swanlab run and storing the resulting run on self.run.

Parameters:

cfg (SwanlabConfig) – Configuration for the Swanlab run (e.g., project and name).
log_dir (Optional[str]) – Optional offline log directory passed to Swanlab’s init.

log_metrics( metrics: dict[str, Any], step: int, prefix: Optional[str] = '', step_metric: Optional[str] = None, step_finished: bool = False, ) → None#

Log metrics to the associated Swanlab run.

Parameters:

metrics (dict[str, Any]) – Mapping of metric names to metric values.
step (int) – Global step value to associate with all logged metrics.
prefix (Optional[str]) – Optional prefix applied to metric names; metric names equal to step_metric are not prefixed.
step_metric (Optional[str]) – Name of a metric that should be excluded from prefixing.

log_hyperparams(params: Mapping[str, Any]) → None#

Update the Swanlab run configuration with the provided hyperparameters.

Parameters:: params (Mapping[str, Any]) – Mapping of hyperparameter names to values to store in the run configuration.

log_plot( figure: matplotlib.pyplot.Figure, step: int, name: str, ) → None#

Log a plot to swanlab.

Parameters:

figure – Matplotlib figure to log
step – Global step value

log_histogram( histogram: list[Any], step: int, name: str, ) → None#: Log histogram metrics to swanlab.

class nemo_rl.utils.logger.GpuMetricSnapshot#

Bases: typing.TypedDict

step: int#: None

metrics: dict[str, Any]#: None

class nemo_rl.utils.logger.RayGpuMonitorLogger( collection_interval: int | float, flush_interval: int | float, metric_prefix: str, step_metric: str, parent_logger: Optional[nemo_rl.utils.logger.Logger] = None, )#

Monitor GPU utilization across a Ray cluster and log metrics to a parent logger.

Initialization

Initialize the GPU monitor.

Parameters:

collection_interval – Interval in seconds to collect GPU metrics
flush_interval – Interval in seconds to flush metrics to parent logger
step_metric – Name of the field to use as the step metric
parent_logger – Logger to receive the collected metrics

start() → None#: Start the GPU monitoring thread.

stop() → None#: Stop the GPU monitoring thread.

_collection_loop() → None#: Main collection loop that runs in a separate thread.

_parse_metric( sample: prometheus_client.samples.Sample, node_idx: int, ) → dict[str, Any]#

Parse a metric sample into a standardized format.

Parameters:

sample – Prometheus metric sample
node_idx – Index of the node

Returns:

Dictionary with metric name and value

_parse_gpu_sku( sample: prometheus_client.samples.Sample, node_idx: int, ) → dict[str, str]#

Parse a GPU metric sample into a standardized format.

Parameters:

sample – Prometheus metric sample
node_idx – Index of the node

Returns:

Dictionary with metric name and value

_collect_gpu_sku() → dict[str, str]#

Collect GPU SKU from all Ray nodes.

Note: This is an internal API and users are not expected to call this.

Returns:: Dictionary of SKU types on all Ray nodes

_collect_metrics() → dict[str, Any]#

Collect GPU metrics from all Ray nodes.

Returns:: Dictionary of collected metrics

_collect( metrics: bool = False, sku: bool = False, ) → dict[str, Any]#

Collect GPU metrics from all Ray nodes.

Returns:: Dictionary of collected metrics

_fetch_and_parse_metrics( node_idx: int, metric_address: str, parser_fn: Callable, )#

Fetch metrics from a node and parse GPU metrics.

Parameters:

node_idx – Index of the node
metric_address – Address of the metrics endpoint

Returns:

Dictionary of GPU metrics

flush() → None#: Flush collected metrics to the parent logger.

nemo_rl.utils.logger._merge_generation_logger_workers( metrics: dict[str, Any], ) → dict[str, Any]#

Merge per-worker async-vLLM generation-logger lists into one list per metric.

generation_logger_metrics is {metric: {worker_id: [samples]}}; merging across workers gives a cluster-level view and avoids one curve per (worker, statistic). Returns a shallow copy; the input dict is not mutated.

nemo_rl.utils.logger._summarize_list(values: list[Any]) → dict[str, float]#: Reduce a list to bounded summary stats for scalar-only backends (MLflow).

class nemo_rl.utils.logger.MLflowLogger( cfg: nemo_rl.utils.logger.MLflowConfig, log_dir: Optional[str] = None, )#

Bases: nemo_rl.utils.logger.LoggerInterface

MLflow logger backend.

Initialization

Initialize MLflow logger.

Parameters:

cfg – MLflow configuration
log_dir – Optional log directory (used as fallback if artifact_location not in cfg)

log_metrics( metrics: dict[str, Any], step: int, prefix: Optional[str] = '', step_metric: Optional[str] = None, step_finished: bool = False, ) → None#

Log metrics to MLflow.

Parameters:

metrics – Dict of metrics to log
step – Global step value
prefix – Optional prefix for metric names
step_metric – Optional step metric name (ignored in MLflow)

log_hyperparams(params: Mapping[str, Any]) → None#

Log hyperparameters to MLflow.

Parameters:: params – Dictionary of hyperparameters to log

log_plot( figure: matplotlib.pyplot.Figure, step: int, name: str, ) → None#

Log a plot to MLflow.

Parameters:

figure – Matplotlib figure to log
step – Global step value
name – Name of the plot

log_histogram( histogram: list[Any], step: int, name: str, ) → None#: Log histogram metrics to MLflow.

__del__() → None#: Clean up resources when the logger is destroyed.

class nemo_rl.utils.logger.Logger(cfg: nemo_rl.utils.logger.LoggerConfig)#

Bases: nemo_rl.utils.logger.LoggerInterface

Main logger class that delegates to multiple backend loggers.

Initialization

Create and configure enabled logging backends and optionally start GPU monitoring.

Parameters:

cfg (LoggerConfig) –

Configuration mapping. Expected keys include:

”log_dir”: base directory for backend logs.
”wandb_enabled”, “swanlab_enabled”, “tensorboard_enabled”, “mlflow_enabled”: booleans to enable backends.
”wandb”, “swanlab”, “tensorboard”, “mlflow”: per-backend configuration dicts.
”monitor_gpus”: boolean to enable Ray GPU monitoring.
”gpu_monitoring”: dict with “collection_interval” and “flush_interval” when GPU monitoring is enabled.

log_metrics( metrics: dict[str, Any], step: int, prefix: Optional[str] = '', step_metric: Optional[str] = None, step_finished: bool = False, ) → None#

Log metrics to all enabled backends.

Parameters:

metrics – Dict of metrics to log
step – Global step value
prefix – Optional prefix for metric names
step_metric – Optional name of a field in metrics to use as step instead of the provided step value (currently only needed for wandb)

log_hyperparams(params: Mapping[str, Any]) → None#

Log hyperparameters to all enabled backends.

Parameters:: params – Dict of hyperparameters to log

log_batched_dict_as_jsonl( to_log: nemo_rl.distributed.batched_data_dict.BatchedDataDict[Any] | dict[str, Any], filename: str, ) → None#

Log a list of dictionaries to a JSONL file.

Parameters:

to_log – BatchedDataDict to log
filename – Filename to log to (within the log directory)

log_string_list_as_jsonl(to_log: list[str], filename: str) → None#

Log a list of strings to a JSONL file.

Parameters:

to_log – list of strings to log
filename – Filename to log to (within the log directory)

log_plot_per_worker_timeline_metrics( metrics: dict[int, list[Any]], step: int, prefix: str, name: str, timeline_interval: float, ) → None#

Log a plot of per-worker timeline metrics.

Parameters:

metrics (-) – Dictionary of metrics to log, where the keys are the worker IDs and the values are the lists of metric values
metrics – dict[str, list[Any]] = {worker_id: [metric_value_1, metric_value_2, …]}
time (- metric values are time series values over)
timeline_interval (the timing gap between the values is the)
step – Global step value
name – Name of the plot
timeline_interval – Interval between timeline points (in seconds)

log_histogram( histogram: list[Any], step: int, name: str, ) → None#

Log histogram metrics to all backends if available.

Parameters:

histogram – List of histogram values
step – Global step value
name – Name of the metric

log_plot( figure: matplotlib.pyplot.Figure, step: int, name: str, ) → None#

Log a matplotlib figure to all backends.

Parameters:

figure – Matplotlib figure to log
step – Global step value
name – Name of the plot

log_plot_token_mult_prob_error( data: dict[str, Any], step: int, name: str, ) → None#

Log a plot of log probability errors in samples.

This function logs & plots the per-token log-probabilities and errors over the sequence for the sample with the highest multiplicative probability error in the batch.

Parameters:

log_data – Dictionary of log probability samples
step – Global step value
name – Name of the plot

__del__() → None#: Clean up resources when the logger is destroyed.

nemo_rl.utils.logger.flatten_dict( d: Mapping[str, Any], sep: str = '.', expand_lists: bool = True, ) → dict[str, Any]#

Flatten a nested dictionary.

Handles nested dictionaries and lists by creating keys with separators. For lists, the index is used as part of the key.

Parameters:

d – Dictionary to flatten
sep – Separator to use between nested keys
expand_lists – If True (default), expand list values into one key per index. If False, keep each list intact under its stable key.

Returns:

Flattened dictionary with compound keys

.. rubric:: Examples

>>> from nemo_rl.utils.logger import flatten_dict
>>> flatten_dict({"a": 1, "b": {"c": 2}})
{'a': 1, 'b.c': 2}

>>> flatten_dict({"a": [1, 2], "b": {"c": [3, 4]}})
{'a.0': 1, 'a.1': 2, 'b.c.0': 3, 'b.c.1': 4}

>>> flatten_dict({"a": [{"b": 1}, {"c": 2}]})
{'a.0.b': 1, 'a.1.c': 2}

nemo_rl.utils.logger.configure_rich_logging( level: str = 'INFO', show_time: bool = True, show_path: bool = True, ) → None#

Configure rich logging for more visually appealing log output.

Parameters:

level – The logging level to use
show_time – Whether to show timestamps in logs
show_path – Whether to show file paths in logs

nemo_rl.utils.logger.print_message_log_samples( message_logs: list[nemo_rl.data.interfaces.LLMMessageLogType], rewards: list[float], num_samples: int = 5, step: int = 0, ) → None#

Visualization for message logs and rewards using a more visual approach with emoji indicators and horizontal layout.

Parameters:

message_logs – List of message logs to sample from
rewards – List of rewards corresponding to each message log
num_samples – Number of samples to display (default: 5)
step – Current training step (for display purposes)

nemo_rl.utils.logger.get_next_experiment_dir(base_log_dir: str) → str#

Create a new experiment directory with an incremented ID.

Parameters:: base_log_dir (str) – The base log directory path
Returns:: Path to the new experiment directory with incremented ID
Return type:: str

nemo_rl.utils.logger.log_container_init_timing() → None#

Log pre-Python container timing from environment variables.

Reads epoch timestamps set by the launch script (ray.sub) and prints the Container Init breakdown. Missing variables are silently skipped so this is safe to call in any environment.

nemo_rl.utils.logger#

Module Contents#

Classes#

Functions#

Data#

API#

`nemo_rl.utils.logger`#