diff --git a/src/psij/executors/batch/lsf/lsf.mustache b/src/psij/executors/batch/lsf/lsf.mustache index aaa5f319..a8b60777 100644 --- a/src/psij/executors/batch/lsf/lsf.mustache +++ b/src/psij/executors/batch/lsf/lsf.mustache @@ -32,6 +32,10 @@ #BSUB -gpu num={{.}}/task {{/gpu_cores_per_process}} + {{#memory}} +#BSUB -M {{memory_kb}}KB + {{/memory}} + {{/job.spec.resources}} diff --git a/src/psij/executors/batch/pbs/pbs_classic.mustache b/src/psij/executors/batch/pbs/pbs_classic.mustache index 03a3d015..90997b25 100644 --- a/src/psij/executors/batch/pbs/pbs_classic.mustache +++ b/src/psij/executors/batch/pbs/pbs_classic.mustache @@ -10,7 +10,7 @@ {{/job.spec.inherit_environment}} {{#job.spec.resources}} -#PBS -l nodes={{computed_node_count}}:ppn={{computed_processes_per_node}}{{#gpu_cores_per_process}}:gpus={{.}}{{/gpu_cores_per_process}} +#PBS -l nodes={{computed_node_count}}:ppn={{computed_processes_per_node}}{{#gpu_cores_per_process}}:gpus={{.}}{{/gpu_cores_per_process}}{{#memory}}:mem={{.}}{{/memory}} {{#exclusive_node_use}} #PBS -n {{/exclusive_node_use}} diff --git a/src/psij/executors/batch/pbs/pbspro.mustache b/src/psij/executors/batch/pbs/pbspro.mustache index 19176687..0617e2d0 100644 --- a/src/psij/executors/batch/pbs/pbspro.mustache +++ b/src/psij/executors/batch/pbs/pbspro.mustache @@ -10,7 +10,7 @@ {{/job.spec.inherit_environment}} {{#job.spec.resources}} -#PBS -l select={{computed_node_count}}:ncpus={{computed_processes_per_node}}:mpiprocs={{computed_processes_per_node}} +#PBS -l select={{computed_node_count}}:ncpus={{computed_processes_per_node}}:mpiprocs={{computed_processes_per_node}}{{#memory}}:mem={{.}}{{/memory}} {{#exclusive_node_use}} #PBS -l place=scatter:exclhost {{/exclusive_node_use}} diff --git a/src/psij/executors/batch/slurm/slurm.mustache b/src/psij/executors/batch/slurm/slurm.mustache index 68e86271..a45b1b55 100644 --- a/src/psij/executors/batch/slurm/slurm.mustache +++ b/src/psij/executors/batch/slurm/slurm.mustache @@ -37,6 +37,10 @@ hyperthreaded CPUs are used) or CPU core (for non-hyperthreaded CPUs).}} #SBATCH --cpus-per-task={{.}} {{/cpu_cores_per_process}} + + {{#memory}} +#SBATCH --mem={{memory_kb}}K + {{/memory}} {{/job.spec.resources}} {{#formatted_job_duration}} diff --git a/src/psij/resource_spec.py b/src/psij/resource_spec.py index 7b565fe4..753228c7 100644 --- a/src/psij/resource_spec.py +++ b/src/psij/resource_spec.py @@ -54,7 +54,8 @@ def __init__(self, node_count: Optional[int] = None, processes_per_node: Optional[int] = None, cpu_cores_per_process: Optional[int] = None, gpu_cores_per_process: Optional[int] = None, - exclusive_node_use: bool = False) -> None: + exclusive_node_use: bool = False, + memory: Optional[int] = None) -> None: """ Some of the properties of this class are constrained. Specifically, `process_count = node_count * processes_per_node`. Specifying all constrained properties @@ -79,6 +80,7 @@ def __init__(self, node_count: Optional[int] = None, set to `False`, which is the default, the LRM is free to co-schedule multiple jobs on a given node if the number of cores requested by those jobs total less than the amount available on the node. + :param memory: The total amount, in bytes, of memory requested for the job. All constructor parameters are accessible as properties. """ @@ -88,6 +90,7 @@ def __init__(self, node_count: Optional[int] = None, self.cpu_cores_per_process = cpu_cores_per_process self.gpu_cores_per_process = gpu_cores_per_process self.exclusive_node_use = exclusive_node_use + self.memory = memory self._check_constraints() def _check_constraints(self) -> None: @@ -184,6 +187,19 @@ def computed_processes_per_node(self) -> int: assert self._computed_ppn is not None return self._computed_ppn + @property + def memory_kb(self) -> Optional[int]: + """ + Returns the memory limit specified by the `memory` property, but in KB. + + :return: If the `memory` property is set on this object, returns `memory // 1024`. If the + `memory` property is `None`, this method returns `None`. + """ + if self.memory: + return self.memory // 1024 + else: + return None + @property def version(self) -> int: """Returns the version of this `ResourceSpec`, which is 1 for this class.""" diff --git a/tests/test_resources.py b/tests/test_resources.py index d44b9905..878d3677 100644 --- a/tests/test_resources.py +++ b/tests/test_resources.py @@ -94,3 +94,13 @@ def _supported(ep: ExecutorTestParams) -> bool: return ep.launcher in ['mpirun', 'srun', 'ibrun'] return False + + +def test_memory(execparams: ExecutorTestParams) -> None: + job = Job(JobSpec(executable='/bin/hostname', launcher=execparams.launcher)) + assert job.spec is not None + job.spec.resources = ResourceSpecV1(memory=1024 * 1024 * 100) + ex = _get_executor_instance(execparams, job) + ex.submit(job) + status = job.wait(timeout=_get_timeout(execparams)) + assert_completed(job, status)