qq_lib.batch.slurmlumi

 23class SlurmLumi(SlurmIT4I):
 24    """
 25    Implementation of BatchInterface for Slurm on the LUMI supercomputer.
 26    """
 27
 28    # all scratch directory types supported by SlurmLumi
 29    SUPPORTED_SCRATCHES = ["scratch", "flash"]
 30
 31    @classmethod
 32    def env_name(cls) -> str:
 33        return "SlurmLumi"
 34
 35    @classmethod
 36    def is_available(cls) -> bool:
 37        return shutil.which("lumi-allocations") is not None
 38
 39    @classmethod
 40    def job_submit(
 41        cls,
 42        res: Resources,
 43        queue: str,
 44        script: Path,
 45        job_name: str,
 46        depend: list[Depend],
 47        env_vars: dict[str, str],
 48        account: str | None = None,
 49        server: str | None = None,
 50        remote_host: str | None = None,
 51    ) -> str:
 52        # set the 'lumi_scratch_type' env var to be able to decide in get_scratch_dir
 53        # whether to create a scratch directory on /scratch or on /flash
 54        if res.uses_scratch():
 55            assert res.work_dir is not None
 56            env_vars[CFG.env_vars.lumi_scratch_type] = res.work_dir
 57
 58        return super().job_submit(
 59            res, queue, script, job_name, depend, env_vars, account, server, remote_host
 60        )
 61
 62    @classmethod
 63    def create_work_dir_on_scratch(cls, job_id: str) -> Path:
 64        if not (account := os.environ.get(CFG.env_vars.slurm_job_account)):
 65            raise QQError(f"No account is defined for job '{job_id}'.")
 66
 67        # get the storage type (scratch or flash)
 68        if not (storage_type := os.environ.get(CFG.env_vars.lumi_scratch_type)):
 69            raise QQError(
 70                f"Environment variable '{CFG.env_vars.lumi_scratch_type}' is not defined. This is a bug!"
 71            )
 72
 73        user = getpass.getuser()
 74
 75        # we attempt to create the scratch directory multiple times in different user directory;
 76        # if the user directory is already created but the user does not have permissions
 77        # to write into it, we append a number to the user's name and try creating a new directory
 78        last_exception = None
 79        for attempt in range(CFG.slurm_lumi_options.scratch_dir_attempts):
 80            user_component = (
 81                user if attempt == 0 else f"{user}{attempt + 1}"
 82            )  # appended number is 2 for the second attempt
 83
 84            scratch = Path(
 85                f"/{storage_type}/{account.lower()}/{user_component}/qq-jobs/job_{job_id}"
 86            )
 87            logger.debug(
 88                f"Creating directory '{str(scratch)}' on '{storage_type}' storage."
 89            )
 90
 91            try:
 92                scratch.mkdir(parents=True, exist_ok=True)
 93                return scratch
 94            except Exception as e:
 95                last_exception = e
 96
 97        # if all attempts failed
 98        raise QQError(
 99            f"Could not create a working directory on {storage_type} for job '{job_id}' after {CFG.slurm_lumi_options.scratch_dir_attempts} attempts: {last_exception}"
100        ) from last_exception
101
102    @classmethod
103    def get_nodes(cls, server: str | None = None) -> list[SlurmLumiNode]:  # ty: ignore[invalid-method-override]
104        nodes = Slurm.get_nodes(server)
105        for node in nodes:
106            node.__class__ = SlurmLumiNode
107
108        return cast("list[SlurmLumiNode]", nodes)
109
110    @classmethod
111    def get_supported_work_dir_types(cls) -> list[str]:
112        return cls.SUPPORTED_SCRATCHES + [
113            "input_dir",
114            "job_dir",  # same as input_dir
115        ]
116
117    @classmethod
118    def _get_default_resources(cls) -> Resources:
119        return Resources(
120            nnodes=1,
121            ncpus_per_node=128,
122            mem_per_cpu="500mb",
123            work_dir="scratch",
124            walltime="1d",
125        )