Merge pull request #134 from TaekyungHeo/nemo-env-var-fix

Pass all env vars to final command in NeMo launcher test template

Author: TaekyungHeo

src/cloudai/schema/test_template/nemo_launcher/slurm_command_gen_strategy.py

@@ -20,17 +20,6 @@
 
 from .slurm_install_strategy import NeMoLauncherSlurmInstallStrategy
 
-REQUIRE_ENV_VARS = [
-    "NCCL_SOCKET_IFNAME",
-    "NCCL_IB_GID_INDEX",
-    "NCCL_IB_TC",
-    "NCCL_IB_QPS_PER_CONNECTION",
-    "UCX_IB_GID_INDEX",
-    "NCCL_IB_ADAPTIVE_ROUTING",
-    "NCCL_IB_SPLIT_DATA_ON_QPS",
-    "NCCL_IBEXT_DISABLE",
-]
-
 
 class NeMoLauncherSlurmCommandGenStrategy(SlurmCommandGenStrategy):
     """
@@ -50,10 +39,8 @@ def gen_exec_command(
         num_nodes: int,
         nodes: List[str],
     ) -> str:
-        # Ensure required environment variables are included
-        for key in REQUIRE_ENV_VARS:
-            if key not in extra_env_vars:
-                extra_env_vars[key] = self.slurm_system.global_env_vars[key]
+        final_env_vars = self._override_env_vars(self.default_env_vars, env_vars)
+        final_env_vars = self._override_env_vars(final_env_vars, extra_env_vars)
 
         launcher_path = os.path.join(
             self.install_path,
@@ -67,7 +54,7 @@ def gen_exec_command(
         self.final_cmd_args["base_results_dir"] = output_path
         self.final_cmd_args["training.model.data.index_mapping_dir"] = output_path
         self.final_cmd_args["launcher_scripts_path"] = os.path.join(launcher_path, "launcher_scripts")
-        for key, value in extra_env_vars.items():
+        for key, value in final_env_vars.items():
             self.final_cmd_args[f"env_vars.{key}"] = value
         self.final_cmd_args["cluster.partition"] = self.slurm_system.default_partition
         nodes = self.slurm_system.parse_nodes(nodes)
@@ -96,7 +83,7 @@ def gen_exec_command(
                 tokenizer_path = extra_cmd_args.split("training.model.tokenizer.model=")[1].split(" ")[0]
                 full_cmd += " " + f"container_mounts=[{tokenizer_path}:{tokenizer_path}]"
 
-        env_vars_str = " ".join(f"{key}={value}" for key, value in extra_env_vars.items())
+        env_vars_str = " ".join(f"{key}={value}" for key, value in final_env_vars.items())
         full_cmd = f"{env_vars_str} {full_cmd}" if env_vars_str else full_cmd
 
         return full_cmd.strip()
@@ -130,13 +117,19 @@ def _generate_cmd_args_str(self, args: Dict[str, str], nodes: List[str]) -> str:
         Returns:
             str: A string of command-line arguments.
         """
-        arg_str_parts = []
+        cmd_arg_str_parts = []
+        env_var_str_parts = []
+
         for key, value in args.items():
-            formatted_key = f"+{key}" if key.startswith("env_vars.") else key
-            arg_str_parts.append(f"{formatted_key}={value}")
+            if key.startswith("env_vars."):
+                if isinstance(value, str) and "," in value:
+                    value = f"\\'{value}\\'"
+                env_var_str_parts.append(f"+{key}={value}")
+            else:
+                cmd_arg_str_parts.append(f"{key}={value}")
 
         if nodes:
             nodes_str = ",".join(nodes)
-            arg_str_parts.append(f"+cluster.nodelist=\\'{nodes_str}\\'")
+            cmd_arg_str_parts.append(f"+cluster.nodelist=\\'{nodes_str}\\'")
 
-        return " ".join(arg_str_parts)
+        return " ".join(cmd_arg_str_parts + env_var_str_parts)

tests/test_slurm_command_gen_strategy.py

@@ -17,7 +17,6 @@
 
 import pytest
 from cloudai.schema.test_template.nemo_launcher.slurm_command_gen_strategy import (
-    REQUIRE_ENV_VARS,
     NeMoLauncherSlurmCommandGenStrategy,
 )
 from cloudai.systems import SlurmSystem
@@ -121,15 +120,8 @@ def nemo_cmd_gen(self, slurm_system: SlurmSystem) -> NeMoLauncherSlurmCommandGen
         strategy = NeMoLauncherSlurmCommandGenStrategy(slurm_system, env_vars, cmd_args)
         return strategy
 
-    def test_raises_if_required_env_var_missed(self, nemo_cmd_gen: NeMoLauncherSlurmCommandGenStrategy):
-        with pytest.raises(KeyError) as exc_info:
-            nemo_cmd_gen.gen_exec_command(
-                env_vars={}, cmd_args={}, extra_env_vars={}, extra_cmd_args="", output_path="", num_nodes=1, nodes=[]
-            )
-        assert REQUIRE_ENV_VARS[0] in str(exc_info.value)
-
     def test_extra_env_vars_added(self, nemo_cmd_gen: NeMoLauncherSlurmCommandGenStrategy):
-        extra_env_vars = {v: "fake" for v in REQUIRE_ENV_VARS}
+        extra_env_vars = {"TEST_VAR_1": "value1", "TEST_VAR_2": "value2"}
         cmd_args = {
             "docker_image_url": "fake",
             "repository_url": "fake",
@@ -148,8 +140,27 @@ def test_extra_env_vars_added(self, nemo_cmd_gen: NeMoLauncherSlurmCommandGenStr
         for k, v in extra_env_vars.items():
             assert f"{k}={v}" in cmd
 
+    def test_env_var_escaping(self, nemo_cmd_gen: NeMoLauncherSlurmCommandGenStrategy):
+        extra_env_vars = {"TEST_VAR": "value,with,commas"}
+        cmd_args = {
+            "docker_image_url": "fake",
+            "repository_url": "fake",
+            "repository_commit_hash": "fake",
+        }
+        cmd = nemo_cmd_gen.gen_exec_command(
+            env_vars={},
+            cmd_args=cmd_args,
+            extra_env_vars=extra_env_vars,
+            extra_cmd_args="",
+            output_path="",
+            num_nodes=1,
+            nodes=[],
+        )
+
+        assert "TEST_VAR=\\'value,with,commas\\'" in cmd
+
     def test_tokenizer_handled(self, nemo_cmd_gen: NeMoLauncherSlurmCommandGenStrategy):
-        extra_env_vars = {v: "fake" for v in REQUIRE_ENV_VARS}
+        extra_env_vars = {"TEST_VAR_1": "value1"}
         cmd_args = {
             "docker_image_url": "fake",
             "repository_url": "fake",