adding proportional cpu and memory args when not asking for the full node (#1678)

Summary: Pull Request resolved: #1678

Differential Revision: D85610860

Pulled By: luccabb

Author: luccabb

python/monarch/_src/job/slurm.py

@@ -13,6 +13,8 @@
 import sys
 from typing import Any, cast, Dict, FrozenSet, List, Optional, Sequence
 
+import clusterscope
+
 from monarch._rust_bindings.monarch_hyperactor.channel import ChannelTransport
 from monarch._rust_bindings.monarch_hyperactor.config import configure
 
@@ -134,6 +136,20 @@ def _submit_slurm_job(self, num_nodes: int) -> str:
         if self._partition:
             sbatch_directives.append(f"#SBATCH --partition={self._partition}")
 
+        # add proportional cpu and memory args when not taking the full node
+        if not self._exclusive and self._partition and self._gpus_per_node:
+            gpus_per_task = self._gpus_per_node // self._ntasks_per_node
+            assert self._partition, "Partition must be set to get cpu and memory args"
+            slurm_args = clusterscope.job_gen_task_slurm(
+                partition=self._partition,
+                gpus_per_task=gpus_per_task,
+                tasks_per_node=self._ntasks_per_node,
+            )
+            sbatch_directives.append(
+                f"#SBATCH --cpus-per-task={slurm_args["cpus_per_task"]}"
+            )
+            sbatch_directives.append(f"#SBATCH --mem={slurm_args["memory"]}")
+
         # Add any additional slurm args as directives
         for arg in self._slurm_args:
             if arg.startswith("-"):

requirements.txt

@@ -10,3 +10,4 @@ torchx-nightly
 lark
 tabulate
 opentelemetry-api
+clusterscope