diff --git a/roles/provision_slurm/templates/provision_slurm.py.j2 b/roles/provision_slurm/templates/provision_slurm.py.j2 index 0d30d93db6ca656cd2faa82418ae3b72b39dd640..bebbe6593fef844a3d40fbb51d5fd81c94c90ebb 100644 --- a/roles/provision_slurm/templates/provision_slurm.py.j2 +++ b/roles/provision_slurm/templates/provision_slurm.py.j2 @@ -5,6 +5,7 @@ import os import stat import subprocess + class ldapSearchConfig: def __init__(self): self.ldapserver="" @@ -14,12 +15,13 @@ class ldapSearchConfig: self.searchFilter="" self.cacertfile='' + class genericUser: def __init__(self): self.dn="" self.cn="" self.entry="" - self.uid="" + self.uid="" def get_users(server): @@ -52,14 +54,24 @@ def mk_slurmaccount(acct): else: subprocess.call(["{{ slurm_dir }}/bin/sacctmgr","-i","create","account",acct]) -def mk_slurmuser(user,acct): - p=subprocess.Popen(["{{ slurm_dir }}/bin/sacctmgr","--noheader","list","Association","user=%s"%user,"format=account"],stdout=subprocess.PIPE,stderr=subprocess.PIPE) +def mk_slurmuser_batch(userbatch,acct): + """ + Fetch a list of all accounts associated with users in userbatch, then add + those that aren't a member of acct to acct using sacctmgr. + """ + p=subprocess.Popen(["/opt/slurm-16.05.4/bin/sacctmgr","--noheader","list","Association","user=%s"%userbatch,"format=user,account"],stdout=subprocess.PIPE,stderr=subprocess.PIPE) + (output,error)=p.communicate() - if acct in output.splitlines(): - return - else: - subprocess.call(["{{ slurm_dir }}/bin/sacctmgr","-i","create","user",user,"account=%s"%acct,"DefaultAccount=%s"%acct]) + + userdict = defaultdict(list) + for line in output.splitlines(): + (u,a) = line.split() + userdict[u].append(a) + + newuserbatch = ",".join([ x[0] for x in userdict.items() if acct not in x[1] ]) + if newuserbatch is not None: + subprocess.call(["/opt/slurm-16.05.4/bin/sacctmgr","-i","create","user",newuserbatch,"account=%s"%acct,"DefaultAccount=%s"%acct]) s=ldapSearchConfig() @@ -81,11 +93,11 @@ for user in users: userlist=[] i=i+1 userlist.append(users[user].entry['uid'][0]) + usergrouplist.append(",".join(userlist)) for usergroup in usergrouplist: - - try: - mk_slurmuser(usergroup,"default") - except: - print traceback.format_exc() - pass + try: + mk_slurmuser_batch(usergroup,"default") + except: + print traceback.format_exc() + pass