Merge pull request #517 from ExaWorks/add_mem

Add `memory` property to `ResourceSpecV1`

Author: hategan

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}}
 
 

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}}

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}}

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}}

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."""

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)