Merge pull request #111 from TaekyungHeo/nemo-bugfix

Add pip install requirements.txt step to NeMoLauncherSlurmInstallStrategy

Author: TaekyungHeo

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

@@ -157,6 +157,7 @@ def install(self) -> InstallStatusResult:
 
         try:
             self._clone_repository(subdir_path)
+            self._install_requirements(subdir_path)
             docker_image_result = self.docker_image_cache_manager.ensure_docker_image(
                 self.docker_image_url, self.SUBDIR_PATH, self.DOCKER_IMAGE_FILENAME
             )
@@ -293,3 +294,22 @@ def _clone_repository(self, subdir_path: str) -> None:
         result = subprocess.run(checkout_cmd, cwd=repo_path, capture_output=True, text=True)
         if result.returncode != 0:
             raise RuntimeError(f"Failed to checkout commit: {result.stderr}")
+
+    def _install_requirements(self, subdir_path: str) -> None:
+        """
+        Installs the required Python packages from the requirements.txt file in the cloned repository.
+
+        Args:
+            subdir_path (str): Subdirectory path for installation.
+        """
+        repo_path = os.path.join(subdir_path, self.REPOSITORY_NAME)
+        requirements_file = os.path.join(repo_path, "requirements.txt")
+
+        if os.path.isfile(requirements_file):
+            logging.debug("Installing requirements from %s", requirements_file)
+            install_cmd = ["pip", "install", "-r", requirements_file]
+            result = subprocess.run(install_cmd, capture_output=True, text=True)
+            if result.returncode != 0:
+                raise RuntimeError(f"Failed to install requirements: {result.stderr}")
+        else:
+            logging.warning("requirements.txt not found in %s", repo_path)

tests/test_slurm_install_strategy.py

@@ -155,6 +155,7 @@ def test_clone_repository_when_path_does_not_exist(self, strategy: NeMoLauncherS
             mock_run.return_value.returncode = 0
             mock_exists.side_effect = lambda path: path == str(subdir_path)
             strategy._clone_repository(str(subdir_path))
+            strategy._install_requirements(str(subdir_path))
 
             mock_run.assert_any_call(
                 ["git", "clone", strategy.repository_url, str(repo_path)], capture_output=True, text=True
@@ -166,6 +167,20 @@ def test_clone_repository_when_path_does_not_exist(self, strategy: NeMoLauncherS
                 text=True,
             )
 
+    def test_install_requirements(self, strategy: NeMoLauncherSlurmInstallStrategy):
+        subdir_path = Path(strategy.slurm_system.install_path) / strategy.SUBDIR_PATH
+        repo_path = subdir_path / strategy.REPOSITORY_NAME
+        requirements_file = repo_path / "requirements.txt"
+        repo_path.mkdir(parents=True, exist_ok=True)
+        requirements_file.touch()
+
+        with patch("subprocess.run") as mock_run:
+            mock_run.return_value.returncode = 0
+            strategy._install_requirements(str(subdir_path))
+            mock_run.assert_called_with(
+                ["pip", "install", "-r", str(requirements_file)], capture_output=True, text=True
+            )
+
     def test_clone_repository_when_path_exists(self, strategy: NeMoLauncherSlurmInstallStrategy):
         subdir_path = Path(strategy.slurm_system.install_path) / strategy.SUBDIR_PATH
         repo_path = subdir_path / strategy.REPOSITORY_NAME