Set properly account in verify_graham

The tests were failing because the account was not specified.

Author: bouthilx

smartdispatch/tests/verify_graham.py

@@ -11,14 +11,17 @@ class VerifyGrahamCluster(VerifySlurmCluster):
 
     def get_arguments(self, **kwargs):
 
-        set_defaults(
-            kwargs,
-            coresPerCommand=1,
-            gpusPerCommand=0,
-            walltime=self.WALLTIME,
-            coresPerNode=self.CORES_PER_NODE,
-            gpusPerNode=self.GPUS_PER_NODE,
-            sbatchFlags="--account=rpp-bengioy")
+        kwargs = super(VerifyGrahamCluster, self).get_arguments(**kwargs)
+
+        if kwargs["gpusPerCommand"] == 0:
+            account = os.environ.get("CPU_SLURM_ACCOUNT")
+        else:
+            account = os.environ.get("GPU_SLURM_ACCOUNT")
+
+        if "sbatchFlags" not in kwargs or len(kwargs["sbatchFlags"]) == 0:
+            kwargs["sbatchFlags"] = "--account=" + account
+        else:
+            kwargs["sbatchFlags"] += " --account=" + account
 
         return kwargs