Adding Allocation and Partition - New Approach (#117)

* added support for allocation and partition specification from slurm params

* add partition/allocation to hpc configs

* new format for hpc config to specify allocation and partition

* update types.ts will allocation parameter to fix error

* Tweak to adding allocation

* lint issues and . notation for hpcconfig

* added allocation/partition to hpcconfig in types.ts

* removed maintainer ability to be null

* adding back maintainer nullability and add check in slurmconnector

configs/config.ts

@@ -34,6 +34,8 @@ for (const hpc in rawHpc) {
       globus: undefined,
       mount: {},
       slurm_input_rules: {},
+      allocation: undefined,
+      partition: undefined
     },
     JSON.parse(JSON.stringify((rawHpc as Record<string, unknown>)[hpc]))
   ) as hpcConfig;

configs/hpc.example.json

@@ -26,7 +26,9 @@
         "/data/cigi/scratch/cigi-gisolve/compute_shared": "/compute_shared"
       },
       "allowlist": [],
-      "denylist": []
+      "denylist": [],
+      "allocation": null,
+      "partition": null
     },
     "expanse_community": {
       "ip": "login.expanse.sdsc.edu",
@@ -56,16 +58,16 @@
       ],
       "init_sbatch_options": [
         "#SBATCH --constraint=lustre",
-        "#SBATCH --partition=shared",
-        "#SBATCH --nodes=1",
-        "#SBATCH --account=TG-EAR190007"
+        "#SBATCH --nodes=1"
       ],
       "xsede_job_log_credential": {
         "xsederesourcename": "expanse.sdsc.xsede.org",
         "apikey": "XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX"
       },
       "allowlist": [],
-      "denylist": []
+      "denylist": [],
+      "allocation": "TG-EAR190007",
+      "partition": "shared"
     },
     "bridges_community_gpu": {
       "ip": "bridges2.psc.edu",
@@ -114,9 +116,11 @@
         "/anvil/scratch/x-cybergis/compute": "/compute_scratch"
       },
       "init_sbatch_script": ["module load gcc", "module load openmpi"],
-      "init_sbatch_options": ["#SBATCH --partition=shared", "#SBATCH --nodes=1"],
+      "init_sbatch_options": ["#SBATCH --nodes=1"],
       "allowlist": [],
-      "denylist": []
+      "denylist": [],
+      "allocation": null,
+      "partition": "shared"
     },
     "anvil_community_highmem": {
       "ip": "anvil.rcac.purdue.edu",
@@ -142,9 +146,11 @@
         "/anvil/scratch/x-cybergis/compute": "/compute_scratch"
       },
       "init_sbatch_script": ["module load gcc", "module load openmpi"],
-      "init_sbatch_options": ["#SBATCH --partition=highmem"],
+      "init_sbatch_options": [],
       "allowlist": [],
-      "denylist": []
+      "denylist": [],
+      "allocation": null,
+      "partition": "highmem"
     },
     "aces_community": {
       "ip": "aces-login1.hprc.tamu.edu",
@@ -173,7 +179,9 @@
         "[email protected]",
         "[email protected]"
       ],
-      "denylist": []
+      "denylist": [],
+      "allocation": null,
+      "partition": null
     }
 
   }

src/connectors/SlurmConnector.ts

@@ -1,6 +1,6 @@
 import * as path from "path";
 
-import { config } from "../../configs/config";
+import { config, hpcConfigMap } from "../../configs/config";
 import * as Helper from "../helpers/Helper";
 import { ConnectorError } from "../utils/errors";
 import { slurm } from "../utils/types";
@@ -35,6 +35,8 @@ class SlurmConnector extends BaseConnector {
    */
   prepare(cmd: string, config: slurm) {
     // prepare sbatch script
+    Helper.nullGuard(this.maintainer);
+    const hpc = hpcConfigMap[this.maintainer.job.hpc];
     config = Object.assign(
       {
         time: "01:00:00",
@@ -44,6 +46,14 @@ class SlurmConnector extends BaseConnector {
       config
     );
 
+    if (config.allocation == null) {
+      config.allocation = hpc.allocation;
+    }
+
+    if (config.partition == null) {
+      config.partition = hpc.partition;
+    }
+
     let modules = "";
     if (config.modules) {
       for (const module of config.modules)
@@ -88,6 +98,7 @@ ${
 }
 ${config.gpus_per_task ? `#SBATCH --gpus-per-task=${config.gpus_per_task}` : ""}
 ${config.partition ? `#SBATCH --partition=${config.partition}` : ""}
+${config.allocation ? `#SBATCH -A ${config.allocation}` : ""}
 ${this.getSBatchTagsFromArray("mail-type", config.mail_type)}
 ${this.getSBatchTagsFromArray("mail-user", config.mail_user)}
 module purge

src/utils/types.ts

@@ -120,6 +120,7 @@ export interface slurm {
   gpus_per_socket?: number | string;
   gpus_per_task?: number | string;
   partition?: string;
+  allocation?: string;
   mail_type?: string[];
   mail_user?: string[];
   modules?: string[];
@@ -210,6 +211,8 @@ export interface hpcConfig {
   xsede_job_log_credential: XSEDEJobLogCredential;
   allowlist: string[];
   denylist: string[];
+  allocation?: string;
+  partition?: string;
 }
 
 export interface XSEDEJobLogCredential {