workloads/megatron_run: restrict bool sentinel to bool fields, add tests

Address two code review comments:
- Restrict "true"/"false"/"" sentinel logic to fields declared with
  default="" (known boolean flags) and model_extra fields (boolean by
  convention). Typed string fields like tensorboard_dir pass their
  value through unchanged.
- Add parameterized tests covering all sentinel cases, extra-field
  behavior, and the non-sentinel pass-through for string-valued fields.

Author: saivishal1999

src/cloudai/workloads/megatron_run/megatron_run.py

@@ -80,14 +80,20 @@ def no_dashed_args(self):
 
     @property
     def cmd_args(self) -> dict[str, Union[str, list[str]]]:
+        # Fields with default="" are boolean flags; extra fields (model_extra) are also boolean by convention.
+        bool_flags = {k for k, f in MegatronRunCmdArgs.model_fields.items() if f.default == ""}
+        extra_keys = set(self.model_extra or {})
         args = self.model_dump(exclude_none=True, exclude={"docker_image_url", "run_script"})
         result: dict[str, Union[str, list[str]]] = {}
         for k, v in args.items():
             flag = f"--{k.replace('_', '-')}"
-            if v == "false":
-                continue
-            elif v in ("true", ""):
-                result[flag] = ""
+            if k in bool_flags or k in extra_keys:
+                if v == "false":
+                    continue
+                elif v in ("true", ""):
+                    result[flag] = ""
+                else:
+                    result[flag] = v
             else:
                 result[flag] = v
         return result

tests/workloads/megatron_run/test_megatron_run.py

@@ -24,6 +24,51 @@
 from cloudai.workloads.megatron_run import MegatronRunCmdArgs, MegatronRunTestDefinition
 
 
+class TestMegatronRunCmdArgs:
+    BASE_ARGS = dict(docker_image_url="http://url", run_script=Path(__file__))
+
+    @pytest.mark.parametrize(
+        ("field", "value", "expected_flag", "expected_value"),
+        [
+            ("log_timers_to_tensorboard", "", "--log-timers-to-tensorboard", ""),
+            ("log_timers_to_tensorboard", "true", "--log-timers-to-tensorboard", ""),
+            ("log_throughput", "true", "--log-throughput", ""),
+            ("log_memory_to_tensorboard", "true", "--log-memory-to-tensorboard", ""),
+        ],
+    )
+    def test_bool_flag_true_and_empty_emit_bare_flag(self, field, value, expected_flag, expected_value):
+        args = MegatronRunCmdArgs(**self.BASE_ARGS, **{field: value})
+        assert args.cmd_args[expected_flag] == expected_value
+
+    @pytest.mark.parametrize("field", ["log_timers_to_tensorboard", "log_throughput", "log_memory_to_tensorboard"])
+    def test_bool_flag_false_omits_flag(self, field):
+        args = MegatronRunCmdArgs(**self.BASE_ARGS, **{field: "false"})
+        flag = f"--{field.replace('_', '-')}"
+        assert flag not in args.cmd_args
+
+    def test_extra_bool_flag_true_emits_bare_flag(self):
+        args = MegatronRunCmdArgs(**self.BASE_ARGS, mock_data="true")
+        assert args.cmd_args["--mock-data"] == ""
+
+    def test_extra_bool_flag_false_omits_flag(self):
+        args = MegatronRunCmdArgs(**self.BASE_ARGS, mock_data="false")
+        assert "--mock-data" not in args.cmd_args
+
+    def test_ordinary_value_passed_through(self):
+        args = MegatronRunCmdArgs(**self.BASE_ARGS, global_batch_size=32)
+        assert args.cmd_args["--global-batch-size"] == 32
+
+    def test_string_field_with_non_sentinel_value_passed_through(self):
+        args = MegatronRunCmdArgs(**self.BASE_ARGS, tensorboard_dir="/some/path")
+        assert args.cmd_args["--tensorboard-dir"] == "/some/path"
+
+    def test_string_field_literal_true_not_treated_as_sentinel(self):
+        # tensorboard_dir is a typed string field (default is a path, not ""),
+        # so "true" must be passed as a literal value, not a bare flag.
+        args = MegatronRunCmdArgs(**self.BASE_ARGS, tensorboard_dir="true")
+        assert args.cmd_args["--tensorboard-dir"] == "true"
+
+
 class TestMegatronRunSuccessCheck:
     def setup_method(self) -> None:
         self.megatron_tdef = MegatronRunTestDefinition(