CRAB/python/BdiiLdap.py

#!/usr/bin/env python
import re
import sys
import ldap

DEBUG = 0
map_source = {'ceList': [], 'bdii': ''}
ce_to_cluster_map = {}
cluster_to_site_map = {}

def runldapquery(filter, attribute, bdii):
    if DEBUG:
        print "runldapquery ["+bdii+"]", filter[0:100], attribute[0:100]
    attribute = attribute.split(' ')
    filter = filter.strip()
    filter = filter.lstrip("'").rstrip("'")

    bdiiuri = 'ldap://' + bdii + ':2170'
    l = ldap.initialize(bdiiuri)

    l.simple_bind_s('', '')

    base = "o=grid"
    scope = ldap.SCOPE_SUBTREE
    timeout = 0
    result_set = []
    filter = filter.strip("'")

    try:
        result_id = l.search(base, scope, filter, attribute)
        while 1:
            result_type, result_data = l.result(result_id, timeout)
            if (result_data == []):
                break
            else:
                if result_type == ldap.RES_SEARCH_ENTRY:
                    result_set.append(result_data)

    except ldap.LDAPError, error_message:
        print error_message

    return result_set

def getJMListFromSEList(selist, bdii='exp-bdii.cern.ch'):
    """
    Given a list of SE FQDNs, return list of CEUniqueIDs that advertise CMS
    support and are in Production,
    sorted by number of waiting jobs in descending order
    """
    jmlist = []

    query = buildOrQuery('GlueCESEBindGroupSEUniqueID', selist)
    pout = runldapquery(query, 'GlueCESEBindGroupCEUniqueID', bdii)

    query = "(&(GlueCEAccessControlBaseRule=VO:cms)(GlueCEStateStatus=Production)"
    query += buildOrQuery('GlueCEUniqueID', [l[0][1]['GlueCESEBindGroupCEUniqueID'][0] for l in pout])
    query += ")"
    pout = runldapquery(query, 'GlueCEUniqueID GlueCEStateWaitingJobs', bdii)

    jminfo_list = []
    for x in pout:
        jminfo = {}
        ce = x[0][1]['GlueCEUniqueID'][0]
        waiting_jobs = x[0][1]['GlueCEStateWaitingJobs'][0]

        jminfo['ce'] = ce
        jminfo['waiting_jobs'] = waiting_jobs
        jminfo_list.append(jminfo)

    def compare_by (fieldname):
        """ Comparison function for sorting dicts """
        def compare_two_dicts (a, b):
            return cmp(int(a[fieldname]), int(b[fieldname]))
        return compare_two_dicts

    jminfo_list.sort(compare_by('waiting_jobs'))
    jmlist = [x['ce'] for x in jminfo_list]
    return jmlist

def generateMaps(ceList, bdii='exp-bdii.cern.ch'):
    """
    Generate maps of CE to Cluster and Cluster to Site as the globals
    ce_to_cluster_map, cluster_to_site_map

    ceList: list of GlueCEUniqueIDs
    bdii: BDII instance to query
    """
    if (ceList == map_source['ceList']
        and bdii == map_source['bdii']): return

    query = buildOrQuery('GlueCEUniqueID', ceList)

    pout = runldapquery(query, 'GlueCEUniqueID GlueForeignKey', bdii)

    r = re.compile('^GlueClusterUniqueID\s*=\s*(.*)')
    for x in pout:
        host = x[0][1]['GlueCEUniqueID'][0]
        clusterid = x[0][1]['GlueForeignKey'][0]
        m = r.match(clusterid)
        if m: ce_to_cluster_map[host] = m.groups()[0]

    query = "(&(objectClass=GlueCluster)"
    query += buildOrQuery('GlueClusterUniqueID', ce_to_cluster_map.values())
    query += ")"

    pout = runldapquery(query, 'GlueClusterUniqueID GlueForeignKey', bdii)
    r = re.compile('^GlueSiteUniqueID=(.*)')
    for x in pout:
        cluster = x[0][1]['GlueClusterUniqueID'][0]
        foreign_keys = x[0][1]['GlueForeignKey']
        for foreign_key in foreign_keys:
            m = r.match(foreign_key)
            if m:
                site = m.groups()[0]
                cluster_to_site_map[cluster] = site

    # cache the list sources
    map_source['ceList'] = ceList
    map_source['bdii'] = bdii

    if (DEBUG): print 40*'*', 'exit generateMaps', 40*'*'
def buildOrQuery(gluekey, list):
    """
    Returns a nugget of LDAP requesting the OR of all items
    of the list equal to the gluekey
    """

    query = "(|"
    for x in list:
        query += "(%s=%s)" % (gluekey, x)
    query += ")"
    return query

def isOSGSite(host_list, bdii='exp-bdii.cern.ch'):
    """
    Given a list of CEs, return only the ones which belong to OSG sites
    """
    generateMaps(host_list, bdii)

    query = buildOrQuery('GlueSiteUniqueID', cluster_to_site_map.values())
    pout = runldapquery(query, 'GlueSiteUniqueID GlueSiteDescription', bdii)
    osg_site_list = []
    for x in pout:
        site_descr = x[0][1]['GlueSiteDescription'][0]

        if (site_descr.find('OSG') != -1):
            osg_site_list.append(x[0][1]['GlueSiteUniqueID'][0])

    osg_host_list = []

    for host in host_list:
        cluster = ce_to_cluster_map[host]
        site = cluster_to_site_map[cluster]

        if (osg_site_list.count(site)):
            osg_host_list.append(host)

    return osg_host_list

def getSoftwareAndArch(host_list, software, arch, bdii='exp-bdii.cern.ch'):
    """
    Given a list of CEs, return only those that match a given software
    and architecture tag
    """
    generateMaps(host_list, bdii)

    results_list = []
    software = 'VO-cms-' + software
    arch = 'VO-cms-' + arch
    query = "'(&(GlueHostApplicationSoftwareRunTimeEnvironment="+software+ ")"
    query +=   "(GlueHostApplicationSoftwareRunTimeEnvironment="+arch+")"

    query += buildOrQuery('GlueChunkKey=GlueClusterUniqueID', [ce_to_cluster_map[h] for h in host_list])
    query += ")"

    pout = runldapquery(query, 'GlueHostApplicationSoftwareRunTimeEnvironment GlueChunkKey', bdii)
    clusterlist = [x[0][1]['GlueChunkKey'][0] for x in pout]

    results_list = []
    for jm in host_list:
        cluster = "GlueClusterUniqueID=" + ce_to_cluster_map[jm]
        if (clusterlist.count(cluster) != 0):
            results_list.append(jm)

    return results_list


def getJobManagerList(selist, software, arch, bdii='exp-bdii.cern.ch', onlyOSG=True):
    """
    Given a list of SE FQDNs, return list of CEUniqueIDs that advertise CMS
    support and are in Production, sorted by number of waiting jobs in
    descending order that have a given software and arch.

    If OnlyOSG is True, return only OSG Sites.
    """

    jmlist = getJMListFromSEList(selist, bdii)
    jmlist = filterCE(jmlist, software, arch, bdii, onlyOSG)

def filterCE(ceList, software, arch, bdii, onlyOSG):
    """
    Given a list of CEUniqueIDs, filter out only the ones with given
    software, arch, and if it belongs to an OSG site.
    """
    generateMaps(ceList, bdii)
    if (onlyOSG): ceList = isOSGSite(ceList, bdii)
    ceList = getSoftwareAndArch(ceList, software, arch, bdii)
    res = removeQueues(ceList)
    return res

def removeQueues(celist):
    """
    Given a list of CEUniqueIDs, return a list of jobmanager contact
    strings.
    """
    r = re.compile('^(.*:.*/jobmanager-.*?)-(.*)')
    jmlist = []
    for x in celist:
        m = r.match(x)
        if m:
            item = m.groups()[0]
            if (jmlist.count(item) == 0):
                jmlist.append(item)
    return jmlist

def listAllCEs(software='', arch='', onlyOSG=False, bdii='exp-bdii.cern.ch'):
    ''' List all GlueCEUniqueIDs that advertise support for CMS '''

    RE_cename = re.compile('^GlueCEUniqueID: (.*)', re.IGNORECASE)
    filt = "'(&(GlueCEUniqueID=*)(GlueCEAccessControlBaseRule=VO:cms))'"
    res = runldapquery(filt, 'GlueCEUniqueID', bdii)
    ceList = [x[0][1]['GlueCEUniqueID'][0] for x in res]

    if (software or arch or onlyOSG):
        ceList = filterCE(ceList, software, arch, bdii, onlyOSG)

    return ceList

def listAllSEs(bdii='exp-bdii.cern.ch'):
    ''' List all SEs that are bound to CEs that advertise support for CMS '''

    RE_sename = re.compile('^GlueCESEBindGroupSEUniqueID: (.*)', re.IGNORECASE)
    seList = []
    ceList = listAllCEs(bdii=bdii)

    query = buildOrQuery('GlueCESEBindGroupCEUniqueID', ceList)
    res = runldapquery(query, 'GlueCESEBindGroupSEUniqueID', bdii)

    for x in res:
        try:
            item = x[0][1]['GlueCESEBindGroupSEUniqueID'][0]
        except KeyError:
            # Sometimes we publish a CESE BindGroup without an SE attached
            pass

        if (seList.count(item) == 0): seList.append(item)
    return seList
Revision:	1.2
Committed:	Thu Aug 14 19:24:00 2008 UTC (16 years, 8 months ago) by ewv
Content type:	text/x-python
Branch:	MAIN
Changes since 1.1:	+48 -29 lines
Log Message:	Consistent API with BdiiLdapsearch
#	User	Rev	Content
1	ewv	1.1	#!/usr/bin/env python
2	ewv	1.2	import re
3			import sys
4	ewv	1.1	import ldap
5
6			DEBUG = 0
7	ewv	1.2	map_source = {'ceList': [], 'bdii': ''}
8	ewv	1.1	ce_to_cluster_map = {}
9			cluster_to_site_map = {}
10
11			def runldapquery(filter, attribute, bdii):
12			if DEBUG:
13			print "runldapquery ["+bdii+"]", filter[0:100], attribute[0:100]
14			attribute = attribute.split(' ')
15			filter = filter.strip()
16			filter = filter.lstrip("'").rstrip("'")
17
18			bdiiuri = 'ldap://' + bdii + ':2170'
19			l = ldap.initialize(bdiiuri)
20	ewv	1.2
21	ewv	1.1	l.simple_bind_s('', '')
22
23			base = "o=grid"
24			scope = ldap.SCOPE_SUBTREE
25			timeout = 0
26			result_set = []
27			filter = filter.strip("'")
28	ewv	1.2
29	ewv	1.1	try:
30			result_id = l.search(base, scope, filter, attribute)
31			while 1:
32			result_type, result_data = l.result(result_id, timeout)
33			if (result_data == []):
34			break
35			else:
36			if result_type == ldap.RES_SEARCH_ENTRY:
37			result_set.append(result_data)
38
39			except ldap.LDAPError, error_message:
40			print error_message
41
42			return result_set
43
44			def getJMListFromSEList(selist, bdii='exp-bdii.cern.ch'):
45			"""
46			Given a list of SE FQDNs, return list of CEUniqueIDs that advertise CMS
47			support and are in Production,
48	ewv	1.2	sorted by number of waiting jobs in descending order
49	ewv	1.1	"""
50			jmlist = []
51
52			query = buildOrQuery('GlueCESEBindGroupSEUniqueID', selist)
53			pout = runldapquery(query, 'GlueCESEBindGroupCEUniqueID', bdii)
54
55			query = "(&(GlueCEAccessControlBaseRule=VO:cms)(GlueCEStateStatus=Production)"
56			query += buildOrQuery('GlueCEUniqueID', [l[0][1]['GlueCESEBindGroupCEUniqueID'][0] for l in pout])
57			query += ")"
58			pout = runldapquery(query, 'GlueCEUniqueID GlueCEStateWaitingJobs', bdii)
59
60			jminfo_list = []
61			for x in pout:
62			jminfo = {}
63			ce = x[0][1]['GlueCEUniqueID'][0]
64			waiting_jobs = x[0][1]['GlueCEStateWaitingJobs'][0]
65
66			jminfo['ce'] = ce
67			jminfo['waiting_jobs'] = waiting_jobs
68			jminfo_list.append(jminfo)
69
70			def compare_by (fieldname):
71			""" Comparison function for sorting dicts """
72			def compare_two_dicts (a, b):
73			return cmp(int(a[fieldname]), int(b[fieldname]))
74			return compare_two_dicts
75
76			jminfo_list.sort(compare_by('waiting_jobs'))
77			jmlist = [x['ce'] for x in jminfo_list]
78			return jmlist
79
80	ewv	1.2	def generateMaps(ceList, bdii='exp-bdii.cern.ch'):
81	ewv	1.1	"""
82			Generate maps of CE to Cluster and Cluster to Site as the globals
83			ce_to_cluster_map, cluster_to_site_map
84	ewv	1.2
85			ceList: list of GlueCEUniqueIDs
86	ewv	1.1	bdii: BDII instance to query
87			"""
88	ewv	1.2	if (ceList == map_source['ceList']
89			and bdii == map_source['bdii']): return
90
91			query = buildOrQuery('GlueCEUniqueID', ceList)
92	ewv	1.1
93			pout = runldapquery(query, 'GlueCEUniqueID GlueForeignKey', bdii)
94
95			r = re.compile('^GlueClusterUniqueID\s=\s(.*)')
96			for x in pout:
97			host = x[0][1]['GlueCEUniqueID'][0]
98			clusterid = x[0][1]['GlueForeignKey'][0]
99			m = r.match(clusterid)
100			if m: ce_to_cluster_map[host] = m.groups()[0]
101	ewv	1.2
102	ewv	1.1	query = "(&(objectClass=GlueCluster)"
103			query += buildOrQuery('GlueClusterUniqueID', ce_to_cluster_map.values())
104			query += ")"
105	ewv	1.2
106	ewv	1.1	pout = runldapquery(query, 'GlueClusterUniqueID GlueForeignKey', bdii)
107			r = re.compile('^GlueSiteUniqueID=(.*)')
108			for x in pout:
109			cluster = x[0][1]['GlueClusterUniqueID'][0]
110			foreign_keys = x[0][1]['GlueForeignKey']
111			for foreign_key in foreign_keys:
112			m = r.match(foreign_key)
113			if m:
114			site = m.groups()[0]
115			cluster_to_site_map[cluster] = site
116
117	ewv	1.2	# cache the list sources
118			map_source['ceList'] = ceList
119			map_source['bdii'] = bdii
120
121			if (DEBUG): print 40'', 'exit generateMaps', 40''
122	ewv	1.1	def buildOrQuery(gluekey, list):
123			"""
124			Returns a nugget of LDAP requesting the OR of all items
125			of the list equal to the gluekey
126			"""
127
128			query = "(\|"
129			for x in list:
130			query += "(%s=%s)" % (gluekey, x)
131	ewv	1.2	query += ")"
132	ewv	1.1	return query
133
134			def isOSGSite(host_list, bdii='exp-bdii.cern.ch'):
135			"""
136			Given a list of CEs, return only the ones which belong to OSG sites
137			"""
138	ewv	1.2	generateMaps(host_list, bdii)
139	ewv	1.1
140			query = buildOrQuery('GlueSiteUniqueID', cluster_to_site_map.values())
141			pout = runldapquery(query, 'GlueSiteUniqueID GlueSiteDescription', bdii)
142			osg_site_list = []
143			for x in pout:
144			site_descr = x[0][1]['GlueSiteDescription'][0]
145
146			if (site_descr.find('OSG') != -1):
147			osg_site_list.append(x[0][1]['GlueSiteUniqueID'][0])
148
149			osg_host_list = []
150
151			for host in host_list:
152			cluster = ce_to_cluster_map[host]
153			site = cluster_to_site_map[cluster]
154
155			if (osg_site_list.count(site)):
156			osg_host_list.append(host)
157
158			return osg_host_list
159
160			def getSoftwareAndArch(host_list, software, arch, bdii='exp-bdii.cern.ch'):
161			"""
162			Given a list of CEs, return only those that match a given software
163			and architecture tag
164			"""
165	ewv	1.2	generateMaps(host_list, bdii)
166	ewv	1.1
167			results_list = []
168			software = 'VO-cms-' + software
169			arch = 'VO-cms-' + arch
170			query = "'(&(GlueHostApplicationSoftwareRunTimeEnvironment="+software+ ")"
171			query += "(GlueHostApplicationSoftwareRunTimeEnvironment="+arch+")"
172
173			query += buildOrQuery('GlueChunkKey=GlueClusterUniqueID', [ce_to_cluster_map[h] for h in host_list])
174			query += ")"
175
176			pout = runldapquery(query, 'GlueHostApplicationSoftwareRunTimeEnvironment GlueChunkKey', bdii)
177			clusterlist = [x[0][1]['GlueChunkKey'][0] for x in pout]
178
179			results_list = []
180			for jm in host_list:
181			cluster = "GlueClusterUniqueID=" + ce_to_cluster_map[jm]
182			if (clusterlist.count(cluster) != 0):
183			results_list.append(jm)
184
185			return results_list
186
187	ewv	1.2
188	ewv	1.1	def getJobManagerList(selist, software, arch, bdii='exp-bdii.cern.ch', onlyOSG=True):
189			"""
190			Given a list of SE FQDNs, return list of CEUniqueIDs that advertise CMS
191			support and are in Production, sorted by number of waiting jobs in
192			descending order that have a given software and arch.
193
194			If OnlyOSG is True, return only OSG Sites.
195			"""
196	ewv	1.2
197	ewv	1.1	jmlist = getJMListFromSEList(selist, bdii)
198	ewv	1.2	jmlist = filterCE(jmlist, software, arch, bdii, onlyOSG)
199
200			def filterCE(ceList, software, arch, bdii, onlyOSG):
201			"""
202			Given a list of CEUniqueIDs, filter out only the ones with given
203			software, arch, and if it belongs to an OSG site.
204			"""
205			generateMaps(ceList, bdii)
206			if (onlyOSG): ceList = isOSGSite(ceList, bdii)
207			ceList = getSoftwareAndArch(ceList, software, arch, bdii)
208			res = removeQueues(ceList)
209	ewv	1.1	return res
210	ewv	1.2
211	ewv	1.1	def removeQueues(celist):
212			"""
213			Given a list of CEUniqueIDs, return a list of jobmanager contact
214			strings.
215			"""
216			r = re.compile('^(.:./jobmanager-.?)-(.)')
217			jmlist = []
218			for x in celist:
219			m = r.match(x)
220			if m:
221			item = m.groups()[0]
222			if (jmlist.count(item) == 0):
223			jmlist.append(item)
224			return jmlist
225	ewv	1.2
226			def listAllCEs(software='', arch='', onlyOSG=False, bdii='exp-bdii.cern.ch'):
227	ewv	1.1	''' List all GlueCEUniqueIDs that advertise support for CMS '''
228	ewv	1.2
229	ewv	1.1	RE_cename = re.compile('^GlueCEUniqueID: (.*)', re.IGNORECASE)
230			filt = "'(&(GlueCEUniqueID=*)(GlueCEAccessControlBaseRule=VO:cms))'"
231			res = runldapquery(filt, 'GlueCEUniqueID', bdii)
232			ceList = [x[0][1]['GlueCEUniqueID'][0] for x in res]
233	ewv	1.2
234			if (software or arch or onlyOSG):
235			ceList = filterCE(ceList, software, arch, bdii, onlyOSG)
236
237	ewv	1.1	return ceList
238	ewv	1.2
239	ewv	1.1	def listAllSEs(bdii='exp-bdii.cern.ch'):
240			''' List all SEs that are bound to CEs that advertise support for CMS '''
241	ewv	1.2
242	ewv	1.1	RE_sename = re.compile('^GlueCESEBindGroupSEUniqueID: (.*)', re.IGNORECASE)
243			seList = []
244	ewv	1.2	ceList = listAllCEs(bdii=bdii)
245	ewv	1.1
246			query = buildOrQuery('GlueCESEBindGroupCEUniqueID', ceList)
247			res = runldapquery(query, 'GlueCESEBindGroupSEUniqueID', bdii)
248
249			for x in res:
250			try:
251			item = x[0][1]['GlueCESEBindGroupSEUniqueID'][0]
252			except KeyError:
253			# Sometimes we publish a CESE BindGroup without an SE attached
254			pass
255
256			if (seList.count(item) == 0): seList.append(item)
257	ewv	1.2	return seList