CRAB/python/DataLocation.py

#!/usr/bin/env python
import os, string, re
import common
from DLSInfo import *

# ####################################
class DataLocationError(exceptions.Exception):
    def __init__(self, errorMessage):
        self.args=errorMessage
        exceptions.Exception.__init__(self, self.args)
        pass
                                                                                      
    def getErrorMessage(self):
        """ Return exception error """
        return "%s" % (self.args)

# ####################################
# class to extact data location
class DataLocation:
    def __init__(self, Listfileblocks, cfg_params):

#       Attributes
        self.Listfileblocks = Listfileblocks  # DLS input: list of fileblocks lists

        self.cfg_params = cfg_params

        self.SelectedSites = {}        # DLS output: list of sites hosting fileblocks
                                       #  retrieved using method getSites

        SEBlackList = []
        try:
            tmpBad = string.split(self.cfg_params['EDG.se_black_list'],',')
            for tmp in tmpBad:
                tmp=string.strip(tmp)
                SEBlackList.append(tmp)
        except KeyError:
            pass
        common.logger.debug(5,'SEBlackList: '+str(SEBlackList))
        self.reSEBlackList=[]
        for bad in SEBlackList:
            self.reSEBlackList.append(re.compile( string.lower(bad) ))

        SEWhiteList = []
        try:
            tmpGood = string.split(self.cfg_params['EDG.se_white_list'],',')
            for tmp in tmpGood:
                tmp=string.strip(tmp)
                SEWhiteList.append(tmp)
        except KeyError:
            pass
        common.logger.debug(5,'SEWhiteList: '+str(SEWhiteList))
        self.reSEWhiteList=[]
        for good in SEWhiteList:
            self.reSEWhiteList.append(re.compile( string.lower(good) ))

        self.osgSitesDictionary = {"Wisconsin":"cmssrm.hep.wisc.edu", \
                                   "Purdue":"dcache.rcac.purdue.edu", \
                                   "Florida":"ufdcache.phys.ufl.edu", \
                                   "Nebraska":"thpc-1.unl.edu", \
                                   "Caltech":"cithep59.ultralight.org", \
                                   "UCSD":"t2data2.t2.ucsd.edu", \
                                   "fnal":"cmssrm.fnal.gov", \
                                   "MIT":"se01.cmsaf.mit.edu"}

# #######################################################################
    def fetchDLSInfo(self):
        """
        Contact DLS
        """
        countblock=0

        # check if dbs-2 is activated, set default to dbs
        # else, set default to dli
        try:
            use_dbs_1 = int(self.cfg_params['CMSSW.use_dbs_1'])
        except:
            use_dbs_1 = 0

        try:
            dlstype=self.cfg_params['CMSSW.dls_type']
        except KeyError:
            if use_dbs_1 == 1:
                dlstype='dli'
            else:
                dlstype='dbs'

        DLS_type="DLS_TYPE_%s"%dlstype.upper()

        ## find the replicas for each block
        blockSites = {}
        failCount = 0
        for fileblocks in self.Listfileblocks:
            countblock=countblock+1
            #dbspath=string.split(afileblock,'#')[0]
            #(null,ds,tier,ow)=string.split(dbspath,'/')
            #datablock=ow+"/"+ds
            #
            dls=DLSInfo(DLS_type,self.cfg_params)
            try:
                replicas=dls.getReplicas(fileblocks)
                common.logger.debug(5,"sites are %s"%replicas)
                replicas = self.checkOSGt2(replicas)
                replicas = self.checkBlackList(replicas, fileblocks)
                if len(replicas)!=0:
                    replicas = self.checkWhiteList(replicas, fileblocks)
                if len(replicas)!=0:
                    blockSites[fileblocks] = replicas
            except DLSNoReplicas, ex:
                common.logger.debug(5,str(ex.getErrorMessage()))
                common.logger.debug(5,"Proceeding without this block.\n")
                failCount = failCount + 1
            except:
                raise DataLocationError('')

        if countblock == failCount:
            msg = "All data blocks encountered a DLS error.  Quitting."
            raise DataLocationError(msg)

        if len(blockSites)==0:
            msg = 'No sites remaining that host any part of the requested data! Exiting... '
            raise CrabException(msg)

        self.SelectedSites = blockSites

# #######################################################################
    def getSites(self):
        """
        get the sites hosting all the needed data 
        """
        return self.SelectedSites

# #######################################################################
    def checkBlackList(self, Sites, fileblocks):
        """
        select sites that are not excluded by the user (via SE black list)
        """
        goodSites = []
        for aSite in Sites:
            common.logger.debug(10,'Site '+aSite)
            good=1
            for re in self.reSEBlackList:
                if re.search(string.lower(aSite)):
                    common.logger.debug(5,'SE in black list, skipping site '+aSite)
                    good=0
                pass
            if good: goodSites.append(aSite)
        if len(goodSites) == 0:
            msg = "No sites hosting the block %s after BlackList" % fileblocks
            common.logger.debug(5,msg)
            common.logger.debug(5,"Proceeding without this block.\n")
        else:
            common.logger.debug(5,"Selected sites for block "+str(fileblocks)+" via BlackList are "+str(goodSites)+"\n")
        return goodSites

# #######################################################################
    def checkWhiteList(self, Sites, fileblocks):
        """
        select sites that are defined by the user (via SE white list)
        """
        if len(self.reSEWhiteList)==0: return Sites
        goodSites = []
        for aSite in Sites:
            good=0
            for re in self.reSEWhiteList:
                if re.search(string.lower(aSite)):
                    common.logger.debug(5,'SE in white list, adding site '+aSite)
                    good=1
                pass
            if good: goodSites.append(aSite)
        if len(goodSites) == 0:
            msg = "No sites hosting the block %s after WhiteList" % fileblocks
            common.logger.debug(5,msg)
            common.logger.debug(5,"Proceeding without this block.\n")
        else:
            common.logger.debug(5,"Selected sites for block "+str(fileblocks)+" via WhiteList are "+str(goodSites)+"\n")
        return goodSites 

# #######################################################################
    def checkOSGt2(self, sites):
        fixedSites = []
        osgKeys = self.osgSitesDictionary.keys()
        for site in sites:
            fix = 0
            for osgSite in osgKeys:
                if string.lower(site) == string.lower(osgSite):
                    fixedSites.append(self.osgSitesDictionary[osgSite])
                    fix = 1
            if (fix == 0):
                fixedSites.append(site)

        return fixedSites

#######################################################################
    def uniquelist(self, old):
        """
        remove duplicates from a list
        """
        nd={}
        for e in old:
            nd[e]=0
        return nd.keys()
Revision:	1.22
Committed:	Thu May 10 14:34:13 2007 UTC (17 years, 11 months ago) by slacapra
Content type:	text/x-python
Branch:	MAIN
CVS Tags:	CRAB_1_5_3_pre1, CRAB_2_0_0_pre3, CRAB_1_5_2, CRAB_2_0_0_pre2, CRAB_2_0_0_pre1, CRAB_1_5_1, CRAB_1_5_1_pre4, CRAB_1_5_1_pre3, CRAB_1_5_1_pre2
Changes since 1.21:	+5 -5 lines
Log Message:	DBS2 is default
#	User	Rev	Content
1	gutsche	1.7	#!/usr/bin/env python
2	slacapra	1.20	import os, string, re
3	afanfani	1.1	import common
4			from DLSInfo import *
5
6			# ####################################
7	afanfani	1.3	class DataLocationError(exceptions.Exception):
8	slacapra	1.5	def __init__(self, errorMessage):
9	slacapra	1.20	self.args=errorMessage
10			exceptions.Exception.__init__(self, self.args)
11	slacapra	1.5	pass
12	afanfani	1.3
13	slacapra	1.5	def getErrorMessage(self):
14			""" Return exception error """
15			return "%s" % (self.args)
16	afanfani	1.1
17			# ####################################
18			# class to extact data location
19			class DataLocation:
20			def __init__(self, Listfileblocks, cfg_params):
21
22			# Attributes
23			self.Listfileblocks = Listfileblocks # DLS input: list of fileblocks lists
24	slacapra	1.5
25	afanfani	1.1	self.cfg_params = cfg_params
26
27	slacapra	1.15	self.SelectedSites = {} # DLS output: list of sites hosting fileblocks
28			# retrieved using method getSites
29	afanfani	1.1
30	gutsche	1.19	SEBlackList = []
31	afanfani	1.1	try:
32	gutsche	1.19	tmpBad = string.split(self.cfg_params['EDG.se_black_list'],',')
33	afanfani	1.1	for tmp in tmpBad:
34			tmp=string.strip(tmp)
35	gutsche	1.19	SEBlackList.append(tmp)
36	afanfani	1.1	except KeyError:
37			pass
38	gutsche	1.19	common.logger.debug(5,'SEBlackList: '+str(SEBlackList))
39			self.reSEBlackList=[]
40			for bad in SEBlackList:
41			self.reSEBlackList.append(re.compile( string.lower(bad) ))
42	afanfani	1.1
43	gutsche	1.19	SEWhiteList = []
44	afanfani	1.4	try:
45	gutsche	1.19	tmpGood = string.split(self.cfg_params['EDG.se_white_list'],',')
46	afanfani	1.4	for tmp in tmpGood:
47			tmp=string.strip(tmp)
48	gutsche	1.19	SEWhiteList.append(tmp)
49	afanfani	1.4	except KeyError:
50			pass
51	gutsche	1.19	common.logger.debug(5,'SEWhiteList: '+str(SEWhiteList))
52			self.reSEWhiteList=[]
53			for good in SEWhiteList:
54			self.reSEWhiteList.append(re.compile( string.lower(good) ))
55	afanfani	1.4
56	slacapra	1.15	self.osgSitesDictionary = {"Wisconsin":"cmssrm.hep.wisc.edu", \
57			"Purdue":"dcache.rcac.purdue.edu", \
58			"Florida":"ufdcache.phys.ufl.edu", \
59			"Nebraska":"thpc-1.unl.edu", \
60			"Caltech":"cithep59.ultralight.org", \
61			"UCSD":"t2data2.t2.ucsd.edu", \
62			"fnal":"cmssrm.fnal.gov", \
63			"MIT":"se01.cmsaf.mit.edu"}
64	afanfani	1.4
65	afanfani	1.1	# #######################################################################
66			def fetchDLSInfo(self):
67			"""
68			Contact DLS
69			"""
70			countblock=0
71
72	gutsche	1.21	# check if dbs-2 is activated, set default to dbs
73			# else, set default to dli
74			try:
75	slacapra	1.22	use_dbs_1 = int(self.cfg_params['CMSSW.use_dbs_1'])
76	gutsche	1.21	except:
77	slacapra	1.22	use_dbs_1 = 0
78	gutsche	1.21
79	afanfani	1.8	try:
80	slacapra	1.15	dlstype=self.cfg_params['CMSSW.dls_type']
81	afanfani	1.8	except KeyError:
82	slacapra	1.22	if use_dbs_1 == 1:
83			dlstype='dli'
84			else:
85	gutsche	1.21	dlstype='dbs'
86
87	afanfani	1.8	DLS_type="DLS_TYPE_%s"%dlstype.upper()
88	afanfani	1.6
89	mkirn	1.11	## find the replicas for each block
90	slacapra	1.15	blockSites = {}
91			failCount = 0
92	afanfani	1.1	for fileblocks in self.Listfileblocks:
93	slacapra	1.15	countblock=countblock+1
94			#dbspath=string.split(afileblock,'#')[0]
95			#(null,ds,tier,ow)=string.split(dbspath,'/')
96			#datablock=ow+"/"+ds
97			#
98			dls=DLSInfo(DLS_type,self.cfg_params)
99			try:
100			replicas=dls.getReplicas(fileblocks)
101			common.logger.debug(5,"sites are %s"%replicas)
102			replicas = self.checkOSGt2(replicas)
103			replicas = self.checkBlackList(replicas, fileblocks)
104			if len(replicas)!=0:
105			replicas = self.checkWhiteList(replicas, fileblocks)
106			if len(replicas)!=0:
107			blockSites[fileblocks] = replicas
108			except DLSNoReplicas, ex:
109	gutsche	1.17	common.logger.debug(5,str(ex.getErrorMessage()))
110			common.logger.debug(5,"Proceeding without this block.\n")
111	slacapra	1.15	failCount = failCount + 1
112			except:
113			raise DataLocationError('')
114
115			if countblock == failCount:
116			msg = "All data blocks encountered a DLS error. Quitting."
117			raise DataLocationError(msg)
118
119			if len(blockSites)==0:
120			msg = 'No sites remaining that host any part of the requested data! Exiting... '
121			raise CrabException(msg)
122
123			self.SelectedSites = blockSites
124	afanfani	1.1
125			# #######################################################################
126			def getSites(self):
127			"""
128	slacapra	1.5	get the sites hosting all the needed data
129	afanfani	1.1	"""
130			return self.SelectedSites
131
132			# #######################################################################
133	slacapra	1.15	def checkBlackList(self, Sites, fileblocks):
134	afanfani	1.1	"""
135	gutsche	1.19	select sites that are not excluded by the user (via SE black list)
136	afanfani	1.1	"""
137			goodSites = []
138			for aSite in Sites:
139	slacapra	1.15	common.logger.debug(10,'Site '+aSite)
140	afanfani	1.1	good=1
141	gutsche	1.19	for re in self.reSEBlackList:
142	slacapra	1.15	if re.search(string.lower(aSite)):
143	gutsche	1.19	common.logger.debug(5,'SE in black list, skipping site '+aSite)
144	afanfani	1.1	good=0
145			pass
146			if good: goodSites.append(aSite)
147			if len(goodSites) == 0:
148	slacapra	1.15	msg = "No sites hosting the block %s after BlackList" % fileblocks
149	gutsche	1.17	common.logger.debug(5,msg)
150			common.logger.debug(5,"Proceeding without this block.\n")
151	slacapra	1.15	else:
152			common.logger.debug(5,"Selected sites for block "+str(fileblocks)+" via BlackList are "+str(goodSites)+"\n")
153	afanfani	1.1	return goodSites
154
155	afanfani	1.4	# #######################################################################
156	slacapra	1.15	def checkWhiteList(self, Sites, fileblocks):
157	afanfani	1.4	"""
158	gutsche	1.19	select sites that are defined by the user (via SE white list)
159	afanfani	1.4	"""
160	gutsche	1.19	if len(self.reSEWhiteList)==0: return Sites
161	afanfani	1.4	goodSites = []
162			for aSite in Sites:
163			good=0
164	gutsche	1.19	for re in self.reSEWhiteList:
165	slacapra	1.15	if re.search(string.lower(aSite)):
166	gutsche	1.19	common.logger.debug(5,'SE in white list, adding site '+aSite)
167	afanfani	1.4	good=1
168			pass
169			if good: goodSites.append(aSite)
170			if len(goodSites) == 0:
171	slacapra	1.15	msg = "No sites hosting the block %s after WhiteList" % fileblocks
172	gutsche	1.17	common.logger.debug(5,msg)
173			common.logger.debug(5,"Proceeding without this block.\n")
174	slacapra	1.15	else:
175			common.logger.debug(5,"Selected sites for block "+str(fileblocks)+" via WhiteList are "+str(goodSites)+"\n")
176	afanfani	1.4	return goodSites
177
178	slacapra	1.15	# #######################################################################
179			def checkOSGt2(self, sites):
180			fixedSites = []
181			osgKeys = self.osgSitesDictionary.keys()
182			for site in sites:
183			fix = 0
184			for osgSite in osgKeys:
185			if string.lower(site) == string.lower(osgSite):
186			fixedSites.append(self.osgSitesDictionary[osgSite])
187			fix = 1
188			if (fix == 0):
189			fixedSites.append(site)
190
191			return fixedSites
192
193	afanfani	1.1	#######################################################################
194			def uniquelist(self, old):
195			"""
196	slacapra	1.5	remove duplicates from a list
197	afanfani	1.1	"""
198			nd={}
199			for e in old:
200			nd[e]=0
201			return nd.keys()