From 81fc9af217a1ad8f2038759bd4d0b7890dfe24a2 Mon Sep 17 00:00:00 2001 From: Ben Lugassi Date: Tue, 14 Jul 2026 16:19:23 +0300 Subject: [PATCH 1/2] 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 --- .../slurm_command_gen_strategy.py | 1 + tests/ref_data/megatron-bridge.sbatch | 2 +- .../test_command_gen_strategy_slurm.py | 23 +++++++++++++++++++ 3 files changed, 25 insertions(+), 1 deletion(-) diff --git a/src/cloudai/workloads/megatron_bridge/slurm_command_gen_strategy.py b/src/cloudai/workloads/megatron_bridge/slurm_command_gen_strategy.py index d0f24d3e9..96cc2da71 100644 --- a/src/cloudai/workloads/megatron_bridge/slurm_command_gen_strategy.py +++ b/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", diff --git a/tests/ref_data/megatron-bridge.sbatch b/tests/ref_data/megatron-bridge.sbatch index 8a86f679e..d04877494 100644 --- a/tests/ref_data/megatron-bridge.sbatch +++ b/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="" diff --git a/tests/workloads/megatron_bridge/test_command_gen_strategy_slurm.py b/tests/workloads/megatron_bridge/test_command_gen_strategy_slurm.py index ebad2bb82..c824173a0 100644 --- a/tests/workloads/megatron_bridge/test_command_gen_strategy_slurm.py +++ b/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: From ab02db260da131c5b1f8f2b8b7a45bc8a4ab639e Mon Sep 17 00:00:00 2001 From: Ben Lugassi Date: Tue, 14 Jul 2026 16:22:07 +0300 Subject: [PATCH 2/2] fix(training_report): improve can_parse validation and logs Reject missing or invalid training artifacts before parsing and log why report generation is skipped. Signed-off-by: Ben Lugassi --- .../report_generator/training/parser.py | 25 +++++++++- .../training/test_training_parser.py | 47 ++++++++++++++++++- 2 files changed, 68 insertions(+), 4 deletions(-) diff --git a/src/cloudai/report_generator/training/parser.py b/src/cloudai/report_generator/training/parser.py index 3bb70a77a..ccff5ab55 100644 --- a/src/cloudai/report_generator/training/parser.py +++ b/src/cloudai/report_generator/training/parser.py @@ -72,13 +72,34 @@ 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 + if config_path.suffix.lower() in {".json", ".yaml", ".yml"}: + try: + config = self.get_model_config(tr) + except (json.JSONDecodeError, yaml.YAMLError) as exc: + logging.warning(f"{name}: invalid config artifact at '{config_path}' ({exc}); skipping training report") + return False + if not isinstance(config, dict) or not config: + logging.warning( + f"{name}: empty or invalid config artifact at '{config_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.""" diff --git a/tests/report_generator/training/test_training_parser.py b/tests/report_generator/training/test_training_parser.py index 2e8e112d3..f8d8ddc57 100644 --- a/tests/report_generator/training/test_training_parser.py +++ b/tests/report_generator/training/test_training_parser.py @@ -472,13 +472,56 @@ 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("{}") + (tmp_path / "nemo_config.json").write_text('{"model": {}}') assert parser.can_parse(_tr(output_path=tmp_path)) is True +@pytest.mark.parametrize( + ("parser", "tb_path", "config_path", "config_contents"), + ( + pytest.param(NeMoRunParser(), "tb_logs/default", "nemo_config.json", "", id="nemo-empty-file"), + pytest.param(NeMoRunParser(), "tb_logs/default", "nemo_config.json", "{}", id="nemo-empty-object"), + pytest.param(NeMoRunParser(), "tb_logs/default", "nemo_config.json", "{", id="nemo-malformed"), + pytest.param( + MegatronBridgeParser(), + "experiment/tb_logs", + "experiment/configs/ConfigContainer.yaml", + "", + id="bridge-empty-file", + ), + pytest.param( + MegatronBridgeParser(), + "experiment/tb_logs", + "experiment/configs/ConfigContainer.yaml", + "{}", + id="bridge-empty-object", + ), + pytest.param( + MegatronBridgeParser(), + "experiment/tb_logs", + "experiment/configs/ConfigContainer.yaml", + "key: [", + id="bridge-malformed", + ), + ), +) +def test_can_parse_rejects_invalid_file_config(tmp_path, parser, tb_path, config_path, config_contents): + tb_dir = tmp_path / tb_path + tb_dir.mkdir(parents=True) + (tb_dir / "events.out.tfevents.1").write_text("x") + artifact = tmp_path / config_path + artifact.parent.mkdir(parents=True, exist_ok=True) + artifact.write_text(config_contents) + + assert parser.can_parse(_tr(output_path=tmp_path)) is False + + def test_megatron_config_path_points_to_tb_event_file(tmp_path): tb_dir = tmp_path / "tensorboard" tb_dir.mkdir()