Revert "Update docker support for Nemo/Nccl testing"

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

@@ -60,17 +60,6 @@ def gen_exec_command(
         srun_command = self._generate_srun_command(slurm_args, final_env_vars, final_cmd_args, extra_cmd_args)
         return self._write_sbatch_script(slurm_args, env_vars_str, srun_command, output_path)
 
-    def get_docker_image_path(self, cmd_args: Dict[str, str]) -> str:
-        if os.path.isfile(cmd_args["docker_image_url"]):
-            image_path = cmd_args["docker_image_url"]
-        else:
-            image_path = os.path.join(
-                self.install_path,
-                NcclTestSlurmInstallStrategy.SUBDIR_PATH,
-                NcclTestSlurmInstallStrategy.DOCKER_IMAGE_FILENAME,
-            )
-        return image_path
-
     def _parse_slurm_args(
         self,
         job_name_prefix: str,
@@ -80,7 +69,11 @@ def _parse_slurm_args(
     ) -> Dict[str, Any]:
         base_args = super()._parse_slurm_args(job_name_prefix, env_vars, cmd_args, nodes)
 
-        image_path = self.get_docker_image_path(cmd_args)
+        image_path = os.path.join(
+            self.install_path,
+            NcclTestSlurmInstallStrategy.SUBDIR_PATH,
+            NcclTestSlurmInstallStrategy.DOCKER_IMAGE_FILENAME,
+        )
 
         container_mounts = ""
         if "NCCL_TOPO_FILE" in env_vars and "DOCKER_NCCL_TOPO_FILE" in env_vars:

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

@@ -85,9 +85,7 @@ def gen_exec_command(
         nodes = self.slurm_system.parse_nodes(nodes)
         if nodes:
             self.final_cmd_args["training.trainer.num_nodes"] = str(len(nodes))
-
-        self.set_container_arg()
-
+        self.final_cmd_args["container"] = self.final_cmd_args["docker_image_url"]
         del self.final_cmd_args["repository_url"]
         del self.final_cmd_args["repository_commit_hash"]
         del self.final_cmd_args["docker_image_url"]
@@ -104,16 +102,6 @@ def gen_exec_command(
 
         return full_cmd.strip()
 
-    def set_container_arg(self) -> None:
-        if os.path.isfile(self.final_cmd_args["docker_image_url"]):
-            self.final_cmd_args["container"] = self.final_cmd_args["docker_image_url"]
-        else:
-            self.final_cmd_args["container"] = os.path.join(
-                self.install_path,
-                NeMoLauncherSlurmInstallStrategy.SUBDIR_PATH,
-                NeMoLauncherSlurmInstallStrategy.DOCKER_IMAGE_FILENAME,
-            )
-
     def _handle_special_keys(self, key: str, value: Any, launcher_path: str, output_path: str) -> Any:
         """
         Handles special formatting for specific keys.

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

@@ -108,11 +108,7 @@ def is_installed(self) -> bool:
         docker_image_path = os.path.join(subdir_path, self.DOCKER_IMAGE_FILENAME)
         repo_path = os.path.join(subdir_path, self.REPOSITORY_NAME)
         repo_installed = os.path.isdir(repo_path)
-
-        if not os.path.isfile(self.docker_image_url):
-            docker_image_installed = os.path.isfile(docker_image_path)
-        else:
-            docker_image_installed = True
+        docker_image_installed = os.path.isfile(docker_image_path)
 
         data_dir_path = self.default_cmd_args["data_dir"]
         datasets_ready = self._check_datasets_on_nodes(data_dir_path)
@@ -146,8 +142,7 @@ def install(self) -> None:
             )
 
         self._clone_repository(subdir_path)
-        if not os.path.isfile(self.docker_image_url):
-            self._setup_docker_image(self.slurm_system, subdir_path)
+        self._setup_docker_image(self.slurm_system, subdir_path)
 
     def _check_install_path_access(self):
         """

tests/test_slurm_command_gen_strategy.py

@@ -4,26 +4,17 @@
 from cloudai.schema.system import SlurmSystem
 from cloudai.schema.system.slurm import SlurmNode, SlurmNodeState
 from cloudai.schema.system.slurm.strategy import SlurmCommandGenStrategy
-from cloudai.schema.test_template.nccl_test.slurm_command_gen_strategy import NcclTestSlurmCommandGenStrategy
-from cloudai.schema.test_template.nemo_launcher.slurm_command_gen_strategy import NeMoLauncherSlurmCommandGenStrategy
 
 
 @pytest.fixture
-def slurm_system(tmp_path: Path) -> SlurmSystem:
+def strategy_fixture() -> SlurmCommandGenStrategy:
     slurm_system = SlurmSystem(
         name="TestSystem",
-        install_path=str(tmp_path / "install"),
-        output_path=str(tmp_path / "output"),
+        install_path="/path/to/install",
+        output_path="/path/to/output",
         default_partition="main",
         partitions={"main": [SlurmNode(name="node1", partition="main", state=SlurmNodeState.IDLE)]},
     )
-    Path(slurm_system.install_path).mkdir()
-    Path(slurm_system.output_path).mkdir()
-    return slurm_system
-
-
-@pytest.fixture
-def strategy_fixture(slurm_system: SlurmSystem) -> SlurmCommandGenStrategy:
     env_vars = {"TEST_VAR": "VALUE"}
     cmd_args = {"test_arg": "test_value"}
     strategy = SlurmCommandGenStrategy(slurm_system, env_vars, cmd_args)
@@ -51,46 +42,3 @@ def test_filename_generation(strategy_fixture: SlurmCommandGenStrategy, tmp_path
 
     # Check the correctness of the sbatch command format
     assert sbatch_command == f"sbatch {filepath_from_command}"
-
-
-class TestNcclTestSlurmCommandGenStrategy__GetDockerImagePath:
-    @pytest.fixture
-    def nccl_slurm_cmd_gen_strategy_fixture(self, slurm_system: SlurmSystem) -> NcclTestSlurmCommandGenStrategy:
-        env_vars = {"TEST_VAR": "VALUE"}
-        cmd_args = {"test_arg": "test_value"}
-        strategy = NcclTestSlurmCommandGenStrategy(slurm_system, env_vars, cmd_args)
-        return strategy
-
-    def test_cmd_arg_file_doesnt_exist(self, nccl_slurm_cmd_gen_strategy_fixture: NcclTestSlurmCommandGenStrategy):
-        cmd_args = {"docker_image_url": f"{nccl_slurm_cmd_gen_strategy_fixture.install_path}/docker_image"}
-        image_path = nccl_slurm_cmd_gen_strategy_fixture.get_docker_image_path(cmd_args)
-        assert image_path == f"{nccl_slurm_cmd_gen_strategy_fixture.install_path}/nccl-test/nccl_test.sqsh"
-
-    def test_cmd_arg_file_exists(self, nccl_slurm_cmd_gen_strategy_fixture: NcclTestSlurmCommandGenStrategy):
-        cmd_args = {"docker_image_url": f"{nccl_slurm_cmd_gen_strategy_fixture.install_path}/docker_image"}
-        Path(cmd_args["docker_image_url"]).touch()
-        image_path = nccl_slurm_cmd_gen_strategy_fixture.get_docker_image_path(cmd_args)
-        assert image_path == cmd_args["docker_image_url"]
-
-
-class TestNeMoLauncherSlurmCommandGenStrategy__SetContainerArg:
-    @pytest.fixture
-    def nemo_cmd_gen(self, slurm_system: SlurmSystem) -> NeMoLauncherSlurmCommandGenStrategy:
-        env_vars = {"TEST_VAR": "VALUE"}
-        cmd_args = {"test_arg": "test_value"}
-        strategy = NeMoLauncherSlurmCommandGenStrategy(slurm_system, env_vars, cmd_args)
-        return strategy
-
-    def test_docker_image_url_is_not_file(self, nemo_cmd_gen: NeMoLauncherSlurmCommandGenStrategy):
-        nemo_cmd_gen.final_cmd_args["docker_image_url"] = f"{nemo_cmd_gen.install_path}/docker_image"
-        nemo_cmd_gen.set_container_arg()
-        assert (
-            nemo_cmd_gen.final_cmd_args["container"]
-            == f"{nemo_cmd_gen.install_path}/NeMo-Megatron-Launcher/nemo_megatron_launcher.sqsh"
-        )
-
-    def test_docker_image_url_is_file(self, nemo_cmd_gen: NeMoLauncherSlurmCommandGenStrategy):
-        nemo_cmd_gen.final_cmd_args["docker_image_url"] = f"{nemo_cmd_gen.install_path}/docker_image"
-        Path(nemo_cmd_gen.final_cmd_args["docker_image_url"]).touch()
-        nemo_cmd_gen.set_container_arg()
-        assert nemo_cmd_gen.final_cmd_args["container"] == nemo_cmd_gen.final_cmd_args["docker_image_url"]