def loookup_replicas(files, protocol=['xroot', 'root']):
from DIRAC.DataManagementSystem.Client.DataManager import DataManager
from DIRAC.Resources.Storage.StorageElement import StorageElement
from LHCbDIRAC.BookkeepingSystem.Client.BookkeepingClient import BookkeepingClient
dm = DataManager()
bk = BookkeepingClient()
files_map = {f.lfn: f for f in files}
res = dm.getReplicas([f.lfn for f in files], getUrl=False)
replicas = res.get('Value', {}).get('Successful', {})
seList = sorted(set(se for f in files for se in replicas.get(f.lfn, {})))
# banned_SE_list = [se for se in seList if 'CNAF' in se]
banned_SE_list = []
print('Found SE list of', seList)
# Check if files are MDF
bkRes = bk.getFileTypeVersion([f.lfn for f in files])
assert not set(lfn for lfn, fileType in bkRes.get('Value', {}).iteritems()
if fileType == 'MDF')
for se in seList:
# TODO Check if SEs are available
lfns = [f.lfn for f in files if se in replicas.get(f.lfn, [])]
if se in banned_SE_list:
print('Skipping banned SE', se)
for lfn in lfns:
files_map[lfn].replicas.append(Replica(lfn, se, banned=True))
continue
else:
print('Looking up replicas for', len(lfns), 'files at', se)
if lfns:
res = StorageElement(se).getURL(lfns, protocol=protocol)
if res['OK']:
for lfn, pfn in res['Value']['Successful'].items():
files_map[lfn].replicas.append(Replica(lfn, se, pfn=pfn))
for lfn in res['Value']['Failed']:
files_map[lfn].replicas.append(Replica(lfn, se, error=res))
else:
print('LFN -> PFN lookup failed for', se, 'with error:',
res['Message'])
for lfn in lfns:
files_map[lfn].replicas.append(
Replica(lfn, se, error=res['Message']))
def setInputData(self,
lfns,
bkClient=None,
runNumber=None,
persistencyType=None):
""" Add the input data and the run number, if available
"""
if not lfns:
self.log.warn(
"no lfns passed in setInputData, was that intentional?")
return S_OK("Nothing to do")
res = Job.setInputData(self, lfns)
if not res['OK']:
return res
if not runNumber or not persistencyType:
if not bkClient:
bkClient = BookkeepingClient()
if not runNumber:
res = bkClient.getFileMetadata(lfns)
if not res['OK']:
return res
runNumbers = []
for fileMeta in res['Value']['Successful'].itervalues():
try:
if fileMeta['RunNumber'] not in runNumbers and fileMeta[
'RunNumber'] is not None:
runNumbers.append(fileMeta['RunNumber'])
except KeyError:
continue
if len(runNumbers) > 1:
runNumber = 'Multiple'
elif len(runNumbers) == 1:
runNumber = str(runNumbers[0])
else:
runNumber = 'Unknown'
if not persistencyType:
res = bkClient.getFileTypeVersion(lfns)
if not res['OK']:
return res
typeVersions = res['Value']
if not typeVersions:
self.log.verbose('The requested files do not exist in the BKK')
typeVersion = ''
else:
self.log.verbose('Found file types %s for LFNs: %s' %
(typeVersions.values(), typeVersions.keys()))
typeVersionsList = list(set(typeVersions.values()))
if len(typeVersionsList) == 1:
typeVersion = typeVersionsList[0]
else:
typeVersion = ''
self._addParameter(self.workflow, 'runNumber', 'JDL', runNumber,
'Input run number')
self._addParameter(self.workflow, 'persistency', 'String', typeVersion,
'Persistency type of the inputs')
return S_OK()
class InputDataResolution(DIRACInputDataResolution):
""" Define the Input Data Policy
"""
def __init__(self, argumentsDict, bkkClient=None):
""" Standard constructor
"""
super(InputDataResolution, self).__init__(argumentsDict)
if not bkkClient:
from LHCbDIRAC.BookkeepingSystem.Client.BookkeepingClient import BookkeepingClient
self.bkkClient = BookkeepingClient()
else:
self.bkkClient = bkkClient
#############################################################################
def execute(self):
"""Given the arguments from the Job Wrapper, this function calls existing
utilities in DIRAC to resolve input data according to LHCb VO policy.
"""
result = super(InputDataResolution, self).execute()
if not result['OK'] or not result['Value'].get('Successful', {}):
return result
resolvedData = result['Value']['Successful']
resolvedData = self._addPfnType(resolvedData)
if not resolvedData['OK']:
return resolvedData
resolvedData = resolvedData['Value']
# TODO: Below is temporary behaviour to prepend mdf: to resolved TURL(s) for case when not a ROOT file
# This instructs the Gaudi applications to use root to access different file types e.g. for MDF.
# In the longer term this should be derived from file catalog metadata information.
# 24/08/2010 - updated hack to use "mdf:" after udpates from Markus
for lfn, mdataList in resolvedData.items():
if not isinstance(mdataList, list):
mdataList = [mdataList]
for mdata in mdataList:
if mdata['pfntype'] == 'MDF':
mdata['turl'] = 'mdf:' + mdata['turl']
self.log.info('Prepending mdf: to TURL for %s' % lfn)
catalogName = self.arguments['Configuration'].get(
'CatalogName', 'pool_xml_catalog.xml')
self.log.verbose('Catalog name will be: %s' % catalogName)
appCatalog = PoolXMLSlice(catalogName)
check = appCatalog.execute(resolvedData)
if not check['OK']:
return check
return result
#############################################################################
def _addPfnType(self, resolvedData):
""" Add the pfn type to the lfn list in input
"""
typeVersions = self.bkkClient.getFileTypeVersion(resolvedData.keys())
if not typeVersions['OK']:
return typeVersions
typeVersions = typeVersions['Value']
for lfn, mdataList in resolvedData.items():
if not isinstance(mdataList, list):
mdataList = [mdataList]
if lfn not in typeVersions:
self.log.warn(
'The file %s do not exist in the BKK, assuming ROOT, unless it is a RAW (MDF)'
% lfn)
if lfn.split('.')[-1].lower() == 'raw':
lfnType = 'MDF'
else:
lfnType = 'ROOT'
else:
self.log.verbose('Adding PFN file type %s for %s' %
(typeVersions[lfn], lfn))
lfnType = typeVersions[lfn]
for mdata in mdataList:
mdata['pfntype'] = lfnType
return S_OK(resolvedData)
