Merge branch 'main' into ci_installer

Author: hategan

requirements-tests.txt

@@ -7,7 +7,7 @@ pytest >= 6.2.0
 requests >= 2.25.1
 pytest-cov
 pytest-timeout
-filelock >= 3.4, < 3.18
+filelock >= 3.4, < 3.19
 textual
 tree-sitter
 tree-sitter-bash

requirements.txt

@@ -1,4 +1,4 @@
 psutil >=5.9, <=7.0.0
 pystache>=0.6.0
 typeguard>=3.0.1
-packaging >= 24.0, <= 24.2
+packaging >= 24.0, <= 25.0

src/psij/executors/batch/batch_scheduler_executor.py

@@ -105,14 +105,21 @@ def __init__(self, launcher_log_file: Optional[Path] = None,
             Whether to keep submit files and auxiliary job files (exit code and output files) after
             a job has completed.
         """
-        super().__init__(work_directory, launcher_log_file)
+        super().__init__(launcher_log_file, work_directory)
         self.queue_polling_interval = queue_polling_interval
         self.initial_queue_polling_delay = initial_queue_polling_delay
         self.queue_polling_error_threshold = queue_polling_error_threshold
         self.keep_files = keep_files
         if 'PSIJ_BATCH_KEEP_FILES' in os.environ:
             self.keep_files = True
 
+    @classmethod
+    def _from_config(cls, config: JobExecutorConfig) -> 'BatchSchedulerExecutorConfig':
+        new = cls()
+        new.work_directory = config.work_directory
+        new.launcher_log_file = config.launcher_log_file
+        return new
+
 
 class InvalidJobStateError(Exception):
     """An exception that signals that a job cannot be cancelled due to it being already done."""
@@ -199,14 +206,21 @@ def __init__(self, url: Optional[str] = None,
             An configuration for this executor instance; if none is specified, a default
             configuration is used.
         """
-        super().__init__(url=url, config=config if config else BatchSchedulerExecutorConfig())
+        super().__init__(url=url, config=self._get_config(config))
         assert config
         self.work_directory = config.work_directory / self.name
         self._queue_poll_thread = self._start_queue_poll_thread()
 
     def _ensure_work_dir(self) -> None:
         self.work_directory.mkdir(parents=True, exist_ok=True)
 
+    def _get_config(self, config: Optional[JobExecutorConfig]) -> BatchSchedulerExecutorConfig:
+        if config is None:
+            return BatchSchedulerExecutorConfig()
+        if isinstance(config, BatchSchedulerExecutorConfig):
+            return config
+        return BatchSchedulerExecutorConfig._from_config(config)
+
     def submit(self, job: Job) -> None:
         """See :func:`~psij.JobExecutor.submit`."""
         logger.info('Job %s: submitting', job.id)

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/job_attributes.py

@@ -53,7 +53,8 @@ def __init__(self, duration: timedelta = timedelta(minutes=10),
         self.account = account
         self.duration = duration
         self.queue_name = queue_name
-        self.project_name = project_name
+        if project_name is not None:
+            self.project_name = project_name
         self.reservation_id = reservation_id
         self._custom_attributes = custom_attributes
 

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)