provision_slurm.py.j2

#!/usr/bin/python
import ldap
import traceback
import os
import stat
import subprocess


class ldapSearchConfig:
    def __init__(self):
        self.ldapserver=""
        self.binddn=""
        self.bindpw=""
        self.baseDN=""
        self.searchFilter=""
        self.cacertfile=''


class genericUser:
    def __init__(self):
        self.dn=""
        self.cn=""
        self.entry=""
        self.uid=""


def get_users(server):
#    ldap.set_option(ldap.OPT_X_TLS_CACERTFILE,server.cacertfile)
    ldap.set_option( ldap.OPT_X_TLS_REQUIRE_CERT, ldap.OPT_X_TLS_NEVER )
    l=ldap.initialize(server.ldapserver)
    l.simple_bind_s(server.binddn,server.bindpw)
    retrieveAttributes = ["*"]
    searchScope=ldap.SCOPE_SUBTREE
    try:
        ldap_result_id = l.search(server.baseDN,searchScope,server.searchFilter,retrieveAttributes)
    except ldap.LDAPError, e:
        pass
    rtype,rdata = l.result(ldap_result_id,1)
    allusers={}
    for user in rdata:
        dn=user[0]
        attrs=user[1]
        allusers[dn]=genericUser()
        allusers[dn].dn=dn
        allusers[dn].entry=attrs
    return allusers


def mk_slurmaccount(acct):
    p=subprocess.Popen(["{{ slurm_dir }}/bin/sacctmgr","--noheader","list","account",acct],stdout=subprocess.PIPE,stderr=subprocess.PIPE)
    (output,error)=p.communicate()
    if acct in output:
        return
    else:
        subprocess.call(["{{ slurm_dir }}/bin/sacctmgr","-i","create","account",acct])


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()
    
    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()
s.ldapserver="{{ ldapURI }}"
s.binddn="{{ ldapBindDN }}"
s.bindpw="{{ ldapBindDNPassword }}"
s.baseDN="{{ ldapBase }}"
s.searchFilter = "{{ ldap_access_filter }}"

users=get_users(s)
mk_slurmaccount("default")
usergrouplist=[]
userlist=[]
i=0
for user in users:
    if i==200:
        i=0
        usergrouplist.append(",".join(userlist))
        userlist=[]
    i=i+1
    userlist.append(users[user].entry['uid'][0])
    
usergrouplist.append(",".join(userlist))
for usergroup in usergrouplist:
    try:
        mk_slurmuser_batch(usergroup,"default")
    except:
        print traceback.format_exc()
        pass