Merge pull request #521 from NVIDIA/am/fix-construction-megatron

Update MegatronRun model dump logic

Author: amaslenn

src/cloudai/workloads/megatron_run/megatron_run.py

@@ -28,8 +28,8 @@
 class MegatronRunCmdArgs(CmdArgs):
     """MegatronRun test command arguments."""
 
-    docker_image_url: str = Field(exclude=True)
-    run_script: Path = Field(exclude=True)
+    docker_image_url: str = Field()
+    run_script: Path = Field()
 
     global_batch_size: Optional[int] = 16
     hidden_size: Optional[int] = 4096
@@ -65,7 +65,7 @@ def no_dashed_args(self):
 
     @property
     def cmd_args(self) -> dict[str, Union[str, list[str]]]:
-        args = self.model_dump(exclude_none=True)
+        args = self.model_dump(exclude_none=True, exclude={"docker_image_url", "run_script"})
         args = {f"--{k.replace('_', '-')}": v for k, v in args.items()}
         return args
 

tests/test_test_definitions.py

@@ -291,11 +291,6 @@ def test_tokenizer_model(self, megatron_run: MegatronRunTestDefinition):
         assert megatron_run.cmd_args_dict["--tokenizer-model"] == Path("/path/to/tokenizer")
         assert megatron_run.cmd_args.tokenizer_model == Path("/path/to/tokenizer")
 
-    def test_auxiliary_fields_not_in_model_dump(self, megatron_run: MegatronRunTestDefinition):
-        d = megatron_run.cmd_args.model_dump()
-        assert "docker_image_url" not in d
-        assert "run_script" not in d
-
     @pytest.mark.parametrize("field", ["load", "save"])
     def test_load_is_set_but_not_mounted(self, field: str):
         with pytest.raises(ValueError) as exc_info:

tests/test_test_scenario.py

@@ -412,6 +412,7 @@ def test_spec_can_set_unknown_args_no_base(self, test_scenario_parser: TestScena
         )
         _, tdef = test_scenario_parser._prepare_tdef(model.tests[0])
         assert tdef.cmd_args_dict["unknown"] == 42
+        assert isinstance(tdef.cmd_args, NCCLCmdArgs)
 
 
 class TestReporters: