🚧 Some SlurmSite logic

xtl.jobs.sites:SlurmLocalSite/SlurmModuleSite
- Implementation of .prepare_preamble()

xtl.jobs.config2
- New ResourcesConfig class for holding compute resources related
  options
- Added .resources attribute to BatchJobConfig

Author: dtriand

src/xtl/jobs/config2.py

@@ -1,10 +1,11 @@
+from datetime import timedelta
 from pathlib import Path
 from typing import Optional
 
 from pydantic import PrivateAttr, model_validator, computed_field, field_serializer
 
 from xtl import settings, Logger
-from xtl.common.options import Option
+from xtl.common.options import Option, Options
 from xtl.common.os import FilePermissions
 from xtl.common.serializers import PermissionOctal
 from xtl.common.validators import cast_as_temp_dir_if_none
@@ -16,8 +17,66 @@
 logger = Logger(__name__)
 
 
+class ResourcesConfig(Options):
+
+    # TODO: Custom formatters & aliases for SLURM
+    # TODO: Custom validators for SLURM-like input
+    cpus: int = \
+        Option(
+            default=1, ge=1,
+            desc='Number of CPU cores required for the job',
+            alias='cpus-per-task'
+        )
+
+    memory: float = \
+        Option(
+            default=1.0, ge=0.0,
+            desc='Amount of memory (in GB) required for the job',
+            alias='mem',
+            formatter=lambda x: f'{x}G'
+        )
+
+    timeout: float | str | timedelta | None = \
+        Option(
+            default=None,
+            desc='Maximum runtime for the job (in minutes or D-HH:MM:SS format)',
+            alias='time',
+            # cast_as=...,
+            # formatter=...
+        )
+
+    gpus: int = \
+        Option(
+            default=0, ge=0,
+            desc='Number of GPUs required for the job',
+            alias='gpus'
+        )
+
+    no_tasks: int = \
+        Option(
+            default=1, ge=1,
+            desc='Number of tasks required for the job (only used in MPI jobs)',
+            alias='ntasks'
+        )
+
+    no_nodes: int = \
+        Option(
+            default=1, ge=1,
+            desc='Number of nodes required for the job (only used in MPI jobs)',
+            alias='nodes'
+        )
+
+    def to_slurm(self) -> list[str]:
+        args = []
+        for field, value in self.to_dict(by_alias=True).items():
+            if value:
+                args.append(f'--{field}={value}')
+        return args
+
+
 class BatchJobConfig(JobConfig):
 
+    # Generate a temporary directory if not provided
     job_directory: Optional[Path] = Option(
         default_factory=lambda: cast_as_temp_dir_if_none(None, prefix='xtl_batch_'),
         desc='Directory for job execution and results',
@@ -40,7 +99,7 @@ class BatchJobConfig(JobConfig):
         )  # for slurm --comment
     permissions: FilePermissions | str | int = \
         Option(
-            default=FilePermissions(0o700),
+            default=settings.jobs.batch.permissions,
             desc='Permissions for the batch file in octal format (e.g., 700)',
             cast_as=FilePermissions,
             formatter=PermissionOctal
@@ -66,6 +125,11 @@ class BatchJobConfig(JobConfig):
             default_factory=dict,
             desc='Templates for the content of the batch file for different shells'
         )
+    resources: ResourcesConfig = \
+        Option(
+            default_factory=ResourcesConfig,
+            desc='Resources required for the batch job'
+        )
 
     _shell: Shell | None = PrivateAttr(None)
     """The shell that will be used to execute the batch file"""
@@ -171,3 +235,17 @@ def get_template(self) -> Optional[str]:
         Get the template for the content of the batch file for the selected shell.
         """
         return self.templates.get(self.shell, None)
+
+    def to_slurm(self) -> list[str]:
+        """
+        Convert the batch job configuration to a list of SLURM command-line arguments.
+        """
+        args = []
+        if self.name:
+            args.append(f'--job-name={self.name}')
+        if self.description:
+            args.append(f'--comment={self.description}')
+        args.extend(self.resources.to_slurm())
+        args.append(f'--output={self.stdout}')
+        args.append(f'--error={self.stderr}')
+        return args

src/xtl/jobs/sites.py

@@ -11,6 +11,7 @@
 import aiofiles
 
 if TYPE_CHECKING:
+    from xtl.jobs.config2 import BatchJobConfig
     from xtl.config.settings import DependencySettings
 from xtl.common.compatibility import PY310_OR_LESS, XTL_COMPUTE_SITE
 from xtl.jobs.batchfiles import BatchFile, BatchFileStatus
@@ -191,6 +192,12 @@ def execute_command(command: Sequence[str], timeout: int = 20) \
     async def execute_batch(self, batch: BatchFile, **kwargs):
         ...
 
+    @staticmethod
+    def _prepare_shebang(shell: ShellType) -> str:
+        if shell.shebang:
+            return f'{shell.shebang}{shell.new_line_char}'
+        return ''
+
 
 class LocalSite(BaseComputeSite):
     """
@@ -202,7 +209,8 @@ def prepare_preamble(self, dependencies: Iterable[DependencySettings | str]
                                              | DependencySettings | str
                                              | None,
                          shell: ShellType) -> str:
-        return self.policy.intercept_preamble('') if self.policy else ''
+        preamble = self._prepare_shebang(shell)
+        return self.policy.intercept_preamble(preamble) if self.policy else preamble
 
     def prepare_command(self, command: str) -> str:
         return self.policy.intercept_command(command) if self.policy else command
@@ -361,7 +369,6 @@ async def _log_stream_to_file(stream: asyncio.StreamReader,
                 await f.close()
 
 
-
 class ModulesSite(LocalSite):
     """
     A compute site that uses
@@ -400,6 +407,16 @@ def _load_modules(modules: str | Iterable[str], shell: ShellType) -> str:
             cmd = f'call {cmd}'
         return cmd
 
+    def _prepare_modules_preamble(self, dependencies: Iterable[DependencySettings | str]
+                                                 | DependencySettings | str
+                                                 | None,
+                                 shell: ShellType) -> str:
+        preamble = self._purge_modules(shell=shell)
+        for dep in self.resolve_dependencies(dependencies):
+            if dep.modules:
+                preamble += self._load_modules(modules=dep.modules, shell=shell)
+        return preamble
+
     def prepare_preamble(self, dependencies: Iterable[DependencySettings | str]
                                              | DependencySettings | str
                                              | None,
@@ -417,10 +434,9 @@ def prepare_preamble(self, dependencies: Iterable[DependencySettings | str]
         :param shell: The shell type to generate the commands for.
         :return: The preamble string.
         """
-        preamble = self._purge_modules(shell=shell)
-        for dep in self.resolve_dependencies(dependencies):
-            if dep.modules:
-                preamble += self._load_modules(modules=dep.modules, shell=shell)
+        preamble = self._prepare_shebang(shell)
+        preamble += self._prepare_modules_preamble(dependencies=dependencies,
+                                                   shell=shell)
         return self.policy.intercept_preamble(preamble) if self.policy else preamble
 
     def check_dependencies(self, dependencies: Iterable[DependencySettings | str]
@@ -506,12 +522,28 @@ class SlurmSite(SchedulerSite, ABC):
     Abstract base class for compute sites that utilize the
     `SLURM <https://slurm.schedmd.com/>`_ job scheduler.
     """
+    _default_shell = Shell.BASH
+    _supported_shells = frozenset([Shell.BASH])
 
-    def prepare_preamble(self, dependencies: Iterable[DependencySettings | str]
-                                             | DependencySettings | str
-                                             | None,
-                         shell: ShellType) -> str:
-        # SBATCH preamble
+    @staticmethod
+    def _prepare_slurm_preamble(shell: ShellType,
+                                config: BatchJobConfig | dict = None) -> str:
+        from xtl.jobs.config2 import BatchJobConfig
+
+        args = []
+        if config:
+            if not isinstance(config, BatchJobConfig):
+                config = BatchJobConfig(**config)
+            args.extend(config.to_slurm())
+
+        if args:
+            nl = shell.new_line_char
+            return nl.join(f'#SBATCH {arg}' for arg in args) + nl
+        return ''
+
+    async def validate_submission(self, batch: BatchFile, **kwargs) -> bool:
+        # Run `sbatch --test-only <batchfile>` to check if allocation is possible
+        #  and check the exit code
         raise NotImplementedError()
 
     async def schedule_batch(self, batch: BatchFile, **kwargs) -> str:
@@ -536,12 +568,14 @@ class SlurmLocalSite(SlurmSite, LocalSite):
     def prepare_preamble(self, dependencies: Iterable[DependencySettings | str]
                                              | DependencySettings | str
                                              | None,
-                         shell: ShellType) -> str:
-        slurm_preamble = SlurmSite.prepare_preamble(self, dependencies=dependencies,
-                                                    shell=shell)
-        local_preamble = LocalSite.prepare_preamble(self, dependencies=dependencies,
-                                                    shell=shell)
-        raise NotImplementedError()
+                         shell: ShellType,
+                         config: BatchJobConfig | dict = None,
+                         **kwargs) -> str:
+        preamble = self._prepare_shebang(shell=shell)
+        # SBATCH preamble
+        if config:
+            preamble += self._prepare_slurm_preamble(shell=shell, config=config)
+        return self.policy.intercept_preamble(preamble) if self.policy else preamble
 
 
 class SlurmModulesSite(SlurmSite, ModulesSite):
@@ -554,12 +588,17 @@ class SlurmModulesSite(SlurmSite, ModulesSite):
     def prepare_preamble(self, dependencies: Iterable[DependencySettings | str]
                                              | DependencySettings | str
                                              | None,
-                         shell: ShellType) -> str:
-        slurm_preamble = SlurmSite.prepare_preamble(self, dependencies=dependencies,
-                                                    shell=shell)
-        modules_preamble = ModulesSite.prepare_preamble(self, dependencies=dependencies,
-                                                        shell=shell)
-        raise NotImplementedError()
+                         shell: ShellType,
+                         config: BatchJobConfig | dict = None,
+                         **kwargs) -> str:
+        preamble = self._prepare_shebang(shell=shell)
+        # SBATCH preamble
+        if config:
+            preamble += self._prepare_slurm_preamble(shell=shell, config=config)
+        # Modules preamble
+        preamble += self._prepare_modules_preamble(dependencies=dependencies,
+                                                   shell=shell)
+        return self.policy.intercept_preamble(preamble) if self.policy else preamble
 
 
 class VirtualizationSite(BaseComputeSite, ABC): ...