diff --git a/bank/system/slurm.py b/bank/system/slurm.py index cbaefed0..e8b56489 100644 --- a/bank/system/slurm.py +++ b/bank/system/slurm.py @@ -128,7 +128,7 @@ def get_locked_state(self, cluster: str) -> bool: raise ClusterNotFoundError(f'Cluster {cluster} is not configured with Slurm') cmd = f'sacctmgr -n -P show assoc account={self.account_name} format=GrpTresRunMins clusters={cluster}' - return 'cpu=0' in ShellCmd(cmd).out + return 'billing=0' in ShellCmd(cmd).out def set_locked_state(self, lock_state: bool, cluster: str) -> None: """Lock or unlock the current slurm account @@ -146,10 +146,7 @@ def set_locked_state(self, lock_state: bool, cluster: str) -> None: raise ClusterNotFoundError(f'Cluster {cluster} is not configured with Slurm') lock_state_int = 0 if lock_state else -1 - ShellCmd(f'sacctmgr -i modify account where account={self.account_name} cluster={cluster} ' - f'set GrpTresRunMins=cpu={lock_state_int}').raise_if_err() - ShellCmd(f'sacctmgr -i modify account where account={self.account_name} cluster={cluster} ' - f'set GrpTresRunMins=gres/gpu={lock_state_int}').raise_if_err() + ShellCmd(f'sacctmgr -i modify account where account={self.account_name} cluster={cluster} set GrpTresRunMins=billing={lock_state_int}').raise_if_err() def get_cluster_usage_per_user(self, cluster: str, start: date, end: date, in_hours: bool = True) -> Dict[str, int]: """Return the raw account usage per user on a given cluster