chore: resolve some github comments.

Author: liquor233

nemo/collections/llm/fn/mixin.py

@@ -14,6 +14,7 @@
 
 from torch import nn
 from typing_extensions import Self
+import lightning.pytorch as pl
 
 from nemo.collections.llm.fn import base as fn
 from nemo.utils import logging
@@ -52,18 +53,12 @@ class FNMixin:
 
     def __init_subclass__(cls, **kwargs):
         # Add OneLogger timing hooks for LightningModule subclasses to enable telemetry tracking
-        try:
-            import lightning.pytorch as pl
-
-            if issubclass(cls, pl.LightningModule):
-                from nemo.lightning.one_logger_callback import hook_class_init_with_callbacks
-
-                hook_class_init_with_callbacks(cls, "on_model_init_start", "on_model_init_end")
-        except Exception:
-            # Continue gracefully if OneLogger hooks cannot be applied
-            pass
-        finally:
-            super().__init_subclass__(**kwargs)
+        if issubclass(cls, pl.LightningModule):
+            from nemo.lightning.one_logger_callback import hook_class_init_with_callbacks
+
+            hook_class_init_with_callbacks(cls, "on_model_init_start", "on_model_init_end")
+
+        super().__init_subclass__(**kwargs)
 
     def forall(self, func: fn.ModulePredicate, recurse: bool = False) -> bool:
         """

nemo/collections/vlm/neva/data/preloaded.py

@@ -516,10 +516,9 @@ def __init__(
         num_image_embeddings_per_tile: int = 576,
         seed: int = 1234,
     ) -> None:
-        if not hasattr(self, "_one_logger_init_started"):
-            from nemo.lightning.one_logger_callback import call_one_logger_callback
+        from nemo.lightning.one_logger_callback import call_one_logger_callback
 
-            call_one_logger_callback("on_dataloader_init_start")
+        call_one_logger_callback("on_dataloader_init_start")
 
         super().__init__()
         if not isinstance(paths, (list, tuple)):
@@ -581,8 +580,7 @@ def custom_on_megatron_step_start(self, step):
             dataloader_type="cyclic",
         )
 
-        if not hasattr(self, "_one_logger_init_started"):
-            call_one_logger_callback("on_dataloader_init_end")
+        call_one_logger_callback("on_dataloader_init_end")
 
     def setup(self, stage: str = "") -> None:
         assert len(self.paths) == 1, "not yet support blend dataset in Neva 2.0!"

nemo/lightning/io/mixin.py

@@ -33,6 +33,7 @@
 from fiddle._src import partial
 from fiddle._src.experimental import serialization
 from typing_extensions import Self
+import lightning.pytorch as pl
 
 from nemo.lightning.io.artifact.base import Artifact
 from nemo.lightning.io.capture import IOProtocol
@@ -190,18 +191,11 @@ def __init_subclass__(cls):
         _io_register_serialization(cls)
 
         # Add OneLogger timing hooks for data modules to enable telemetry tracking
-        try:
-            import lightning.pytorch as pl
-
-            if issubclass(cls, pl.LightningDataModule):
-                from nemo.lightning.one_logger_callback import hook_class_init_with_callbacks
+        if issubclass(cls, pl.LightningDataModule):
+            from nemo.lightning.one_logger_callback import hook_class_init_with_callbacks
 
-                hook_class_init_with_callbacks(cls, "on_dataloader_init_start", "on_dataloader_init_end")
-        except Exception:
-            # Continue gracefully if OneLogger hooks cannot be applied
-            pass
-        finally:
-            super().__init_subclass__()
+            hook_class_init_with_callbacks(cls, "on_dataloader_init_start", "on_dataloader_init_end")
+        super().__init_subclass__()
 
     def io_transform_args(self, init_fn, *args, **kwargs) -> Dict[str, Any]:
         """

nemo/lightning/one_logger_callback.py

@@ -153,22 +153,15 @@ def init_one_logger() -> None:
     if not enable_one_logger:
         return
 
-    try:
-        # Check if OneLogger is already configured
-        if TrainingTelemetryProvider.instance().one_logger_ready:
-            return
-
-        # Get initialization configuration
-        init_config = get_one_logger_init_config()
-        one_logger_config = OneLoggerConfig(**init_config)
-
-        # Configure the provider with entry-point exporters (automatically calls on_app_start)
-        TrainingTelemetryProvider.instance().with_base_config(
-            one_logger_config
-        ).with_export_config().configure_provider()
-        _ONE_LOGGER_CALLBACK = OneLoggerNeMoCallback(TrainingTelemetryProvider.instance())
-    except Exception:
-        _HAVE_ONE_LOGGER = False
+    # Get initialization configuration
+    init_config = get_one_logger_init_config()
+    one_logger_config = OneLoggerConfig(**init_config)
+
+    # Configure the provider with entry-point exporters (automatically calls on_app_start)
+    TrainingTelemetryProvider.instance().with_base_config(
+        one_logger_config
+    ).with_export_config().configure_provider()
+    _ONE_LOGGER_CALLBACK = OneLoggerNeMoCallback(TrainingTelemetryProvider.instance())
 
 
 def update_one_logger_config(

nemo/utils/meta_info_manager.py

@@ -18,6 +18,8 @@
 import os
 from typing import Any, Dict
 
+from lightning.pytorch.callbacks.model_checkpoint import ModelCheckpoint
+
 enable_one_logger = True
 
 
@@ -41,7 +43,7 @@ def get_one_logger_init_config() -> Dict[str, Any]:
     # Minimal configuration - required fields only
     init_config = {
         # Required fields (from OneLoggerConfig) - no defaults
-        "application_name": "nemo-application",
+        "application_name": "nemo",
         "session_tag_or_fn": session_tag,
         # Important fields with defaults - provide if available from config
         "enable_for_current_rank": _should_enable_for_current_rank(),
@@ -82,7 +84,6 @@ def _get_base_callback_config(
 
     world_size = int(os.environ.get('WORLD_SIZE', 1))
     max_steps = getattr(trainer, 'max_steps', 1)
-    # Use hardcoded value for log_every_n_steps instead of getting from trainer
     log_every_n_steps = getattr(trainer, 'log_every_n_steps', 10)
     micro_batch_size = global_batch_size // world_size
     # Get PERF_VERSION_TAG from environment
@@ -99,7 +100,7 @@ def _get_base_callback_config(
     is_validation_iterations_enabled = False
     save_checkpoint_strategy = "sync"
 
-    checkpoint_callbacks = [cb for cb in trainer.callbacks if 'Checkpoint' in type(cb).__name__]
+    checkpoint_callbacks = [cb for cb in trainer.callbacks if isinstance(cb, ModelCheckpoint)]
     is_save_checkpoint_enabled = len(checkpoint_callbacks) > 0
 
     val_check_interval = getattr(trainer, 'val_check_interval', -1)
@@ -166,6 +167,7 @@ def get_nemo_v1_callback_config(trainer: Any) -> Dict[str, Any]:
             global_batch_size = micro_batch_size * int(os.environ.get('WORLD_SIZE', 1))
         elif hasattr(model_cfg, 'train_ds') and hasattr(model_cfg.train_ds, 'bucket_batch_size'):
             # For ASR with bucketing, use the average batch size
+            # This is a temporary fix to support the bucketing 
             bucket_batch_sizes = model_cfg.train_ds.bucket_batch_size
             # Handle both ListConfig and regular list types
             if hasattr(bucket_batch_sizes, '__len__') and len(bucket_batch_sizes) > 0:
@@ -232,12 +234,6 @@ def _should_enable_for_current_rank() -> bool:
     Returns:
         True if OneLogger should be enabled for the current rank, False otherwise
     """
-    try:
-        rank = int(os.environ.get('RANK', 0))
-        world_size = int(os.environ.get('WORLD_SIZE', 1))
-
-        # Enable for rank 0 or the last rank (common pattern)
-        return rank == 0 or rank == world_size - 1
-    except (ValueError, TypeError):
-        # Default to True on invalid values (as expected by tests)
-        return True
+    rank = int(os.environ.get('RANK', 0))
+    # Enable for rank 0 or the last rank (common pattern)
+    return rank == 0