def make_stage_command_groups()

in launcher/nemo/stages.py [0:0]

• 14 lines of code
• 3 McCabe index (conditional complexity)

    def make_stage_command_groups(self, stage_cfg_path: Path) -> List[List[str]]:
        """
        Custom run stage which will invoke the entry script only
        [TODO] Make this compatiable with NeMo flow as well
        """
        if get_container_type(self.cfg.get("container", None)) == "docker":
            logger.warning(
                f"![WARNING] You're using docker container directly for slurm workload, we'd highly recommend using enroot instead"
            )
            command_groups = [
                [
                    # Launch container first
                    f"srun -l bash {stage_cfg_path.parents[0]}/launch_docker_container.sh",
                    f"srun -l bash {stage_cfg_path.parents[0]}/docker_exec_script.sh",
                ]
            ]
        # There will be only a single command group
        # enroot or conda/venv, no need to launch docker container
        else:
            command_groups = [[f"bash {stage_cfg_path.parents[0]}/train_script.sh"]]

        return command_groups