SMART-Lab /
smartdispatch
| @@ 7-27 (lines=21) @@ | ||
| 4 | from verify_slurm_cluster import VerifySlurmCluster, set_defaults |
|
| 5 | ||
| 6 | ||
| 7 | class VerifyCedarCluster(VerifySlurmCluster): |
|
| 8 | ||
| 9 | WALLTIME = 60 |
|
| 10 | CORES_PER_NODE = 24 |
|
| 11 | GPUS_PER_NODE = 4 |
|
| 12 | ||
| 13 | def get_arguments(self, **kwargs): |
|
| 14 | ||
| 15 | kwargs = super(VerifyCedarCluster, self).get_arguments(**kwargs) |
|
| 16 | ||
| 17 | if kwargs["gpusPerCommand"] == 0: |
|
| 18 | account = os.environ.get("CPU_SLURM_ACCOUNT") |
|
| 19 | else: |
|
| 20 | account = os.environ.get("GPU_SLURM_ACCOUNT") |
|
| 21 | ||
| 22 | if "sbatchFlags" not in kwargs or len(kwargs["sbatchFlags"]) == 0: |
|
| 23 | kwargs["sbatchFlags"] = "--account=" + account |
|
| 24 | else: |
|
| 25 | kwargs["sbatchFlags"] += " --account=" + account |
|
| 26 | ||
| 27 | return kwargs |
|
| 28 | ||
| 29 | ||
| 30 | if __name__ == "__main__": |
|
| @@ 6-26 (lines=21) @@ | ||
| 3 | from verify_slurm_cluster import VerifySlurmCluster, set_defaults |
|
| 4 | ||
| 5 | ||
| 6 | class VerifyGrahamCluster(VerifySlurmCluster): |
|
| 7 | ||
| 8 | WALLTIME = 60 |
|
| 9 | CORES_PER_NODE = 32 |
|
| 10 | GPUS_PER_NODE = 2 |
|
| 11 | ||
| 12 | def get_arguments(self, **kwargs): |
|
| 13 | ||
| 14 | kwargs = super(VerifyGrahamCluster, self).get_arguments(**kwargs) |
|
| 15 | ||
| 16 | if kwargs["gpusPerCommand"] == 0: |
|
| 17 | account = os.environ.get("CPU_SLURM_ACCOUNT") |
|
| 18 | else: |
|
| 19 | account = os.environ.get("GPU_SLURM_ACCOUNT") |
|
| 20 | ||
| 21 | if "sbatchFlags" not in kwargs or len(kwargs["sbatchFlags"]) == 0: |
|
| 22 | kwargs["sbatchFlags"] = "--account=" + account |
|
| 23 | else: |
|
| 24 | kwargs["sbatchFlags"] += " --account=" + account |
|
| 25 | ||
| 26 | return kwargs |
|
| 27 | ||
| 28 | ||
| 29 | if __name__ == "__main__": |
|
