Fix how test-in-scenario is merge with test-in-toml

Author: amaslenn

src/cloudai/_core/test_scenario_parser.py

@@ -275,16 +275,30 @@ def _prepare_tdef(self, test_info: TestRunModel) -> Tuple[Test, TestDefinition]:
             if test_info.test_name not in self.test_mapping:
                 raise ValueError(f"Test '{test_info.test_name}' is not defined. Was tests directory correctly set?")
             test = self.test_mapping[test_info.test_name]
-        elif test_info.test_template_name:
+
+            test_defined = test.test_definition.model_dump()
+            tc_defined = test_info.tdef_model_dump()
+            merged_data = deep_merge(test_defined, tc_defined)
+            test.test_definition = tp.load_test_definition(merged_data, self.strict)
+        elif test_info.test_template_name:  # test fully defined in the scenario
             test = tp._parse_data(test_info.tdef_model_dump(), self.strict)
         else:
             # this should never happen, because we check for this in the modelvalidator
             raise ValueError(
                 f"Cannot configure test case '{test_info.id}' with both 'test_name' and 'test_template_name'."
             )
 
-        merged_data = test.test_definition.model_dump()
-        merged_data.update(test_info.tdef_model_dump())
-        tdef = tp.load_test_definition(merged_data, self.strict)
+        return test, test.test_definition
+
 
-        return test, tdef
+def deep_merge(a: dict, b: dict):
+    result = a.copy()
+    for key in b:
+        if key in result:
+            if isinstance(result[key], dict) and isinstance(b[key], dict):
+                result[key] = deep_merge(result[key], b[key])
+            else:
+                result[key] = b[key]
+        else:
+            result[key] = b[key]
+    return result

src/cloudai/models/scenario.py

@@ -75,8 +75,8 @@ def tdef_model_dump(self) -> dict:
             "agent_steps": self.agent_steps,
             "agent_metric": self.agent_metric,
             "extra_container_mounts": self.extra_container_mounts,
-            "cmd_args": self.cmd_args.model_dump() if self.cmd_args else None,
             "extra_env_vars": self.extra_env_vars if self.extra_env_vars else None,
+            "cmd_args": self.cmd_args.model_dump() if self.cmd_args else None,
             "git_repos": [repo.model_dump() for repo in self.git_repos] if self.git_repos else None,
             "nsys": self.nsys.model_dump() if self.nsys else None,
         }

tests/test_test_scenario.py

@@ -46,7 +46,11 @@
     JaxToolboxReportGenerationStrategy,
     NemotronTestDefinition,
 )
-from cloudai.workloads.megatron_run import CheckpointTimingReportGenerationStrategy, MegatronRunTestDefinition
+from cloudai.workloads.megatron_run import (
+    CheckpointTimingReportGenerationStrategy,
+    MegatronRunCmdArgs,
+    MegatronRunTestDefinition,
+)
 from cloudai.workloads.nccl_test import (
     NCCLCmdArgs,
     NCCLTestDefinition,
@@ -280,7 +284,7 @@ def test_total_time_limit_with_empty_hooks():
     assert result == "01:00:00"
 
 
-class TestSpec:
+class TestInScenario:
     @pytest.mark.parametrize("missing_arg", ["test_template_name", "name", "description"])
     def test_without_base(self, missing_arg: str):
         spec = {
@@ -414,6 +418,34 @@ def test_spec_can_set_unknown_args_no_base(self, test_scenario_parser: TestScena
         assert tdef.cmd_args_dict["unknown"] == 42
         assert isinstance(tdef.cmd_args, NCCLCmdArgs)
 
+    def test_data_is_merge_correctly(self, test_scenario_parser: TestScenarioParser, slurm_system: SlurmSystem):
+        test_scenario_parser.test_mapping = {
+            "megatron": Test(
+                test_definition=MegatronRunTestDefinition(
+                    name="megatron",
+                    description="desc",
+                    test_template_name="MegatronRun",
+                    cmd_args=MegatronRunCmdArgs(docker_image_url="docker://megatron", run_script=Path("run.sh")),
+                ),
+                test_template=TestTemplate(system=slurm_system),
+            )
+        }
+        model = TestScenarioModel.model_validate(
+            toml.loads(
+                """
+            name = "test"
+
+            [[Tests]]
+            id = "1"
+            test_name = "megatron"
+            cmd_args = { any = 42 }
+            """
+            )
+        )
+        _, tdef = test_scenario_parser._prepare_tdef(model.tests[0])
+        assert isinstance(tdef.cmd_args, MegatronRunCmdArgs)
+        assert tdef.cmd_args.run_script == Path("run.sh")
+
 
 class TestReporters:
     def test_default(self):