diff --git a/src/DIRAC/WorkloadManagementSystem/Utilities/PilotWrapper.py b/src/DIRAC/WorkloadManagementSystem/Utilities/PilotWrapper.py index a2fe5d0a372..95c9e69395f 100644 --- a/src/DIRAC/WorkloadManagementSystem/Utilities/PilotWrapper.py +++ b/src/DIRAC/WorkloadManagementSystem/Utilities/PilotWrapper.py @@ -78,7 +78,12 @@ def pilotWrapperScript( - pilotFilesCompressedEncodedDict=None, pilotOptions="", pilotExecDir="", envVariables=None, location="" + pilotFilesCompressedEncodedDict=None, + pilotOptions="", + pilotExecDir="", + envVariables=None, + location="", + CVMFS_locations=[], ): """Returns the content of the pilot wrapper script. @@ -95,6 +100,8 @@ def pilotWrapperScript( :type envVariables: dict :param location: location where to get the pilot files :type location: string + :param location: optional CVMFS locations of where to get the pilot files + :type location: list :returns: content of the pilot wrapper :rtype: string @@ -106,6 +113,19 @@ def pilotWrapperScript( if envVariables is None: envVariables = {} + if CVMFS_locations: + # Here we are making the assumption that, if CVMFS_locations is, e.g., ['/cvmfs/somewhere', '/cvmfs/elsewhere'] + # and the project is 'LHCb', + # then the pilot can maybe be found at locations + # - file:/cvmfs/somewhere/lhcbdirac/pilot + # - file:/cvmfs/elsewhere/lhcbdirac/pilot + project = "" + if "-l" in pilotOptions: + project = pilotOptions.split(" ")[pilotOptions.split(" ").index("-l") + 1].lower() + CVMFS_locs = "[" + ",".join('"file:' + x + "/" + project + '/dirac/pilot"' for x in CVMFS_locations) + "]" + else: + CVMFS_locs = '["file:/cvmfs/dirac.egi.eu/pilot"]' + compressedString = "" # are there some pilot files to unpack? Then we create the unpacking string for pfName, encodedPf in pilotFilesCompressedEncodedDict.items(): @@ -178,8 +198,8 @@ def pilotWrapperScript( # we try from the available locations locs = [os.path.join('https://', loc) for loc in location] locations = locs + [os.path.join(loc, 'pilot') for loc in locs] -# adding also, as last the cvmfs location dirac.egi.eu, but this won't contain a valid JSON -locations += ['file:/cvmfs/dirac.egi.eu/pilot/'] +# adding also the cvmfs locations +locations += %(CVMFS_locs)s for loc in locations: print('Trying %%s' %% loc) @@ -262,7 +282,8 @@ def pilotWrapperScript( logger.debug('Checksum matched') """ % { - "location": location + "location": location, + "CVMFS_locs": CVMFS_locs, } localPilot += ( diff --git a/src/DIRAC/WorkloadManagementSystem/Utilities/test/Test_PilotWrapper.py b/src/DIRAC/WorkloadManagementSystem/Utilities/test/Test_PilotWrapper.py index 0b7f33f3c3d..96cec570f8b 100644 --- a/src/DIRAC/WorkloadManagementSystem/Utilities/test/Test_PilotWrapper.py +++ b/src/DIRAC/WorkloadManagementSystem/Utilities/test/Test_PilotWrapper.py @@ -1,9 +1,9 @@ """ This is a test of the creation of the pilot wrapper """ -import os import base64 import bz2 +import os from DIRAC.WorkloadManagementSystem.Utilities.PilotWrapper import pilotWrapperScript @@ -87,3 +87,37 @@ def test_scriptPilot3(): assert 'os.environ["someName"]="someValue"' in res assert "lhcb-portal.cern.ch" in res + assert """locations += ["file:/cvmfs/dirac.egi.eu/pilot"]""" in res + + +def test_scriptPilot3_2(): + """test script creation""" + res = pilotWrapperScript( + pilotFilesCompressedEncodedDict={"proxy": "thisIsSomeProxy"}, + pilotOptions="-c 123 --foo bar", + envVariables={"someName": "someValue", "someMore": "oneMore"}, + location="lhcb-portal.cern.ch", + CVMFS_locations=["/cvmfs/lhcb.cern.ch", "/cvmfs/dirac.egi.eu"], + ) + + assert 'os.environ["someName"]="someValue"' in res + assert "lhcb-portal.cern.ch" in res + assert """locations += ["file:/cvmfs/lhcb.cern.ch/dirac/pilot","file:/cvmfs/dirac.egi.eu/dirac/pilot"]""" in res + + +def test_scriptPilot3_3(): + """test script creation""" + res = pilotWrapperScript( + pilotFilesCompressedEncodedDict={"proxy": "thisIsSomeProxy"}, + pilotOptions="-c 123 --foo bar -l LHCb -h pippo", + envVariables={"someName": "someValue", "someMore": "oneMore"}, + location="lhcb-portal.cern.ch", + CVMFS_locations=["/cvmfs/lhcb.cern.ch", "/cvmfs/dirac.egi.eu"], + ) + + assert 'os.environ["someName"]="someValue"' in res + assert "lhcb-portal.cern.ch" in res + assert ( + """locations += ["file:/cvmfs/lhcb.cern.ch/lhcbdirac/pilot","file:/cvmfs/dirac.egi.eu/lhcbdirac/pilot"]""" + in res + )