fix(megatron_bridge): save TensorBoard logs by default

Set the default log directory to the saved run output so training reports can consume TensorBoard metrics while preserving user overrides.

Signed-off-by: Ben Lugassi <blugassi@nvidia.com>

Author: blugassi

src/cloudai/report_generator/training/parser.py

@@ -72,13 +72,23 @@ def get_model_config(self, tr: TestRun) -> dict:
     def get_model_name(self, tr: TestRun) -> str:
         """CloudAI-sourced model identity for the run."""
 
+    @staticmethod
+    def _has_tb_event_files(tb_dir: Path) -> bool:
+        # TensorBoard documents event files as having "tfevents" in the filename.
+        return any(path.is_file() for path in tb_dir.rglob("*tfevents*"))
+
     def can_parse(self, tr: TestRun) -> bool:
         """Return True when the run produced the TB events and config artifact this parser needs."""
+        name = type(self).__name__
         tb_dir = self.get_tb_dir(tr)
-        if not (tb_dir.is_dir() and any(tb_dir.iterdir())):
+        if not (tb_dir.is_dir() and self._has_tb_event_files(tb_dir)):
+            logging.warning(f"{name}: no TensorBoard events at '{tb_dir}'; skipping training report")
             return False
         config_path = self.get_config_path(tr)
-        return config_path is not None and config_path.is_file()
+        if config_path is None or not config_path.is_file():
+            logging.warning(f"{name}: config artifact not found under '{tr.output_path}'; skipping training report")
+            return False
+        return True
 
     def parse(self, tr: TestRun, system: System) -> TrainingResults:
         """Read TB scalars + the config artifact and assemble TrainingResults."""

src/cloudai/workloads/megatron_bridge/slurm_command_gen_strategy.py

@@ -430,6 +430,7 @@ def _list_or_comma_str(self, val: str | list[str] | None) -> Optional[str]:
     def _add_extra_cmd_args(self, extra_cmd_args: dict[str, str]) -> list[str]:
         """Hydra overrides: defaults merged with the user's extra_cmd_args, which take precedence."""
         overrides = {
+            "logger.tensorboard_dir": "/nemo_run/tb_logs",
             "logger.log_timers_to_tensorboard": "true",
             "logger.log_throughput_to_tensorboard": "true",
             "logger.log_memory_to_tensorboard": "true",

tests/ref_data/megatron-bridge.sbatch

@@ -19,7 +19,7 @@ if [ "${WANDB_INSTALL_RC}" -ne 0 ]; then
 fi
 
 LAUNCH_RC=0
-NEMORUN_HOME="__OUTPUT_DIR__/output" __INSTALL_DIR__/Run__main-venv/bin/python __INSTALL_DIR__/Megatron-Bridge__main/scripts/performance/setup_experiment.py -p main -t 00:20:00 -i __OUTPUT_DIR__/output/megatron_bridge_image.sqsh -hf dummy_token -ng 8 -gn 8 -cm __INSTALL_DIR__/Megatron-Bridge__main:/opt/Megatron-Bridge -cb 'export CUDA_VISIBLE_DEVICES=0,1,2,3' -cb 'export NCCL_DEBUG=INFO' -m qwen3 -mr 30b_a3b --detach false --save_config_filepath /nemo_run/configs/ConfigContainer.yaml --additional_slurm_params 'gpus-per-node=8;gres=gpu:8' logger.log_timers_to_tensorboard=true logger.log_throughput_to_tensorboard=true logger.log_memory_to_tensorboard=true >>"$LOG" 2>&1 || LAUNCH_RC=$?
+NEMORUN_HOME="__OUTPUT_DIR__/output" __INSTALL_DIR__/Run__main-venv/bin/python __INSTALL_DIR__/Megatron-Bridge__main/scripts/performance/setup_experiment.py -p main -t 00:20:00 -i __OUTPUT_DIR__/output/megatron_bridge_image.sqsh -hf dummy_token -ng 8 -gn 8 -cm __INSTALL_DIR__/Megatron-Bridge__main:/opt/Megatron-Bridge -cb 'export CUDA_VISIBLE_DEVICES=0,1,2,3' -cb 'export NCCL_DEBUG=INFO' -m qwen3 -mr 30b_a3b --detach false --save_config_filepath /nemo_run/configs/ConfigContainer.yaml --additional_slurm_params 'gpus-per-node=8;gres=gpu:8' logger.tensorboard_dir=/nemo_run/tb_logs logger.log_timers_to_tensorboard=true logger.log_throughput_to_tensorboard=true logger.log_memory_to_tensorboard=true >>"$LOG" 2>&1 || LAUNCH_RC=$?
 
 
 JOB_ID=""

tests/report_generator/training/test_training_parser.py

@@ -472,7 +472,10 @@ def test_can_parse_requires_nonempty_tb_dir(tmp_path):
 
     tb_dir = tmp_path / "tb_logs" / "default"
     tb_dir.mkdir(parents=True)
-    (tb_dir / "events.out.tfevents").write_text("x")
+    (tb_dir / "placeholder.txt").write_text("x")
+    assert parser.can_parse(_tr(output_path=tmp_path)) is False  # non-empty dir without event files
+
+    (tb_dir / "custom.tfevents.log").write_text("x")
     assert parser.can_parse(_tr(output_path=tmp_path)) is False  # file-backed parsers also need config
 
     (tmp_path / "nemo_config.json").write_text("{}")

tests/workloads/megatron_bridge/test_command_gen_strategy_slurm.py

@@ -185,6 +185,29 @@ def test_defaults_not_emitted_when_not_set_in_toml(
         assert "--cuda_graph_impl" not in cmd
         assert " -ms " not in cmd
 
+    def test_tensorboard_dir_defaults_to_persisted_run_path(
+        self, configured_slurm_system: SlurmSystem, make_test_run: Callable[..., TestRun]
+    ) -> None:
+        tr = make_test_run(output_subdir="out_default_tensorboard")
+        cmd_gen = MegatronBridgeSlurmCommandGenStrategy(configured_slurm_system, tr)
+
+        wrapper_content = self._wrapper_content(cmd_gen)
+
+        assert "logger.tensorboard_dir=/nemo_run/tb_logs" in wrapper_content
+
+    def test_tensorboard_dir_can_be_overridden(
+        self, configured_slurm_system: SlurmSystem, make_test_run: Callable[..., TestRun]
+    ) -> None:
+        tr = make_test_run(output_subdir="out_custom_tensorboard")
+        tdef = cast(MegatronBridgeTestDefinition, tr.test)
+        tdef.extra_cmd_args["logger.tensorboard_dir"] = "/nemo_run/custom_tb"
+        cmd_gen = MegatronBridgeSlurmCommandGenStrategy(configured_slurm_system, tr)
+
+        wrapper_content = self._wrapper_content(cmd_gen)
+
+        assert "logger.tensorboard_dir=/nemo_run/custom_tb" in wrapper_content
+        assert "logger.tensorboard_dir=/nemo_run/tb_logs" not in wrapper_content
+
     def test_container_image_local_path_passed_verbatim(
         self, cmd_gen: MegatronBridgeSlurmCommandGenStrategy, test_run: TestRun
     ) -> None: