[NeMo-UX] Wrap task config save in a try/except (NVIDIA#9956) (NVIDIA…

…#9984) * wrap task config save in a try/except * move fiddle import --------- Signed-off-by: ashors1 <[email protected]> Co-authored-by: Anna Shors <[email protected]>
paarthneekhara · Jan 18, 2025 · f3453ea · f3453ea
1 parent 585ace4
commit f3453ea
Showing 1 changed file with 9 additions and 5 deletions.
diff --git a/nemo/lightning/nemo_logger.py b/nemo/lightning/nemo_logger.py
@@ -7,7 +7,6 @@
 
 import lightning_fabric as fl
 import pytorch_lightning as pl
-from fiddle._src.experimental import serialization
 from pytorch_lightning.callbacks.model_checkpoint import ModelCheckpoint as PTLModelCheckpoint
 from pytorch_lightning.loggers import Logger, TensorBoardLogger, WandbLogger
 
@@ -211,10 +210,15 @@ def _setup_trainer_model_checkpoint(self, trainer, log_dir, ckpt=None):
                 ModelCheckpoint.CHECKPOINT_NAME_LAST = callback.filename + '-last'
 
     def _handle_task_config(self, task_config, log_dir):
-        task_config.save_config_img(log_dir / "task.png")
-        task_json = serialization.dump_json(task_config)
-        with open(log_dir / "task.json", "w") as f:
-            f.write(task_json)
+        try:
+            from fiddle._src.experimental import serialization
+
+            task_config.save_config_img(log_dir / "task.png")
+            task_json = serialization.dump_json(task_config)
+            with open(log_dir / "task.json", "w") as f:
+                f.write(task_json)
+        except Exception as e:
+            logging.warning(f'Saving task config failed: {e}. Skipping saving')
 
     def _setup_file_logging(self, log_dir):
         """Set up file logging based on rank settings."""