Support different architectures during job submission and runtime

etc/submit.sh

@@ -121,12 +121,17 @@ echo -e "======== WMAgent CMS environment load finished at $(TZ=GMT date) ======
 
 
 echo "======== WMAgent COMP Python bootstrap starting at $(TZ=GMT date) ========"
-# First, decide which COMP ScramArch to use based on the required OS
+# First, decide which COMP ScramArch to use based on the required OS and Architecture
+THIS_ARCH=`uname -m`  # if it's PowerPC, it returns `ppc64le`
+if [ "$THIS_ARCH" = "x86_64" ]
+then
+    THIS_ARCH="amd64"
+fi
 if [ "$REQUIRED_OS" = "rhel7" ];
 then
-    WMA_SCRAM_ARCH=slc7_amd64_gcc630
+    WMA_SCRAM_ARCH=slc7_${THIS_ARCH}_gcc630
 else
-    WMA_SCRAM_ARCH=slc6_amd64_gcc493
+    WMA_SCRAM_ARCH=slc6_${THIS_ARCH}_gcc700
 fi
 echo "Job requires OS: $REQUIRED_OS, thus setting ScramArch to: $WMA_SCRAM_ARCH"
 

etc/submit_py3.sh

@@ -122,12 +122,17 @@ echo -e "======== WMAgent CMS environment load finished at $(TZ=GMT date) ======
 
 
 echo "======== WMAgent COMP Python bootstrap starting at $(TZ=GMT date) ========"
-# First, decide which COMP ScramArch to use based on the required OS
+# First, decide which COMP ScramArch to use based on the required OS and Architecture
+THIS_ARCH=`uname -m`  # if it's PowerPC, it returns `ppc64le`
+if [ "$THIS_ARCH" = "x86_64" ]
+then
+    THIS_ARCH="amd64"
+fi
 if [ "$REQUIRED_OS" = "rhel7" ];
 then
-    WMA_SCRAM_ARCH=slc7_amd64_gcc630
+    WMA_SCRAM_ARCH=slc7_${THIS_ARCH}_gcc630
 else
-    WMA_SCRAM_ARCH=slc6_amd64_gcc700
+    WMA_SCRAM_ARCH=slc6_${THIS_ARCH}_gcc700
 fi
 echo "Job requires OS: $REQUIRED_OS, thus setting ScramArch to: $WMA_SCRAM_ARCH"
 

src/python/WMCore/BossAir/Plugins/BasePlugin.py

@@ -8,8 +8,9 @@
 from builtins import object, str, bytes
 from future.utils import viewitems, viewvalues
 
+from Utils.Utilities import decodeBytesToUnicode
 from WMCore.WMException import WMException
-from WMCore.WMRuntime.Tools.Scram import ARCH_TO_OS
+from WMCore.WMRuntime.Tools.Scram import ARCH_TO_OS, SCRAM_TO_ARCH
 
 
 
@@ -152,3 +153,40 @@ def scramArchtoRequiredOS(scramArch=None):
             requiredOSes.add('any')
 
         return ','.join(sorted(requiredOSes))
+
+    @staticmethod
+    def scramArchtoRequiredArch(scramArch=None):
+        """
+        Converts a given ScramArch to a unique target CPU architecture.
+        Note that an architecture precedence is enforced in case there are
+        multiple matches.
+        In case no scramArch is defined, leave the architecture undefined.
+        :param scramArch: can be either a string or a list of ScramArchs
+        :return: a string with the matched architecture
+        """
+        defaultArch = "X86_64"
+        requiredArchs = set()
+        if scramArch is None:
+            return None
+        elif isinstance(scramArch, (str, bytes)):
+            scramArch = [scramArch]
+
+        for item in scramArch:
+            item = decodeBytesToUnicode(item)
+            arch = item.split("_")[1]
+            if arch not in SCRAM_TO_ARCH:
+                msg = "Job configured to a ScramArch: '{}' not supported in BossAir".format(item)
+                raise BossAirPluginException(msg)
+            requiredArchs.add(SCRAM_TO_ARCH.get(arch))
+
+        # now we have the final list of architectures, return only 1 of them
+        if len(requiredArchs) == 1:
+            return requiredArchs.pop()
+        elif "X86_64" in requiredArchs:
+            return "X86_64"
+        elif "ppc64le" in requiredArchs:
+            return "ppc64le"
+        elif "aarch64" in requiredArchs:
+            return "aarch64"
+        else:  # should never get here!
+            return defaultArch

src/python/WMCore/BossAir/Plugins/SimpleCondorPlugin.py

@@ -625,7 +625,12 @@ def getJobParameters(self, jobList):
             ad['My.REQUIRED_OS'] = classad.quote(encodeUnicodeToBytesConditional(requiredOSes, condition=PY2))
             cmsswVersions = ','.join(job.get('swVersion'))
             ad['My.CMSSW_Versions'] = classad.quote(encodeUnicodeToBytesConditional(cmsswVersions, condition=PY2))
-
+            requiredArch = self.scramArchtoRequiredArch(job.get('scramArch'))
+            if not requiredArch:  # only Cleanup jobs should not have ScramArch defined
+                ad['Requirements'] = '(TARGET.Arch =!= Undefined)'
+            else:
+                ad['Requirements'] = '(TARGET.Arch =?= "{}")'.format(requiredArch)
+
             jobParameters.append(ad)    
 
         return jobParameters

src/python/WMCore/WMRuntime/Tools/Scram.py

@@ -38,8 +38,8 @@
 from Utils.PythonVersion import PY3
 from Utils.Utilities import encodeUnicodeToBytesConditional, decodeBytesToUnicodeConditional
 
+SCRAM_TO_ARCH = {'amd64': 'X86_64', 'aarch64': 'aarch64', 'ppc64le': 'ppc64le'}
 ARCH_TO_OS = {'slc5': ['rhel6'], 'slc6': ['rhel6'], 'slc7': ['rhel7']}
-
 OS_TO_ARCH = {}
 for arch, oses in viewitems(ARCH_TO_OS):
     for osName in oses:

test/python/WMCore_t/BossAir_t/BasePlugin_t.py

@@ -12,7 +12,7 @@
 
 from WMCore_t.BossAir_t.BossAir_t import BossAirTest
 
-from WMCore.BossAir.Plugins.BasePlugin import BasePlugin
+from WMCore.BossAir.Plugins.BasePlugin import BasePlugin, BossAirPluginException
 
 
 class BasePluginTest(BossAirTest):
@@ -40,6 +40,29 @@ def testScramArchToOS(self):
 
         return
 
+    def testScramArchtoRequiredArch(self):
+        """
+        Test mapping of ScramArch to a given architecture
+        """
+        bp = BasePlugin(config=None)
+
+        self.assertEqual(bp.scramArchtoRequiredArch('slc5_amd64_gcc481'), 'X86_64')
+        self.assertEqual(bp.scramArchtoRequiredArch('slc6_amd64_gcc630'), 'X86_64')
+        self.assertEqual(bp.scramArchtoRequiredArch('slc7_amd64_gcc10'), 'X86_64')
+        self.assertEqual(bp.scramArchtoRequiredArch('slc7_aarch64_gcc700'), 'aarch64')
+        self.assertEqual(bp.scramArchtoRequiredArch('slc7_ppc64le_gcc9'), 'ppc64le')
+        self.assertIsNone(bp.scramArchtoRequiredArch(None))
+        self.assertIsNone(bp.scramArchtoRequiredArch(None))
+        with self.assertRaises(BossAirPluginException):
+            bp.scramArchtoRequiredArch("slc7_BLAH_gcc700")
+
+        self.assertEqual(bp.scramArchtoRequiredArch(['slc5_amd64_gcc481', 'slc6_amd64_gcc630']), 'X86_64')
+        self.assertEqual(bp.scramArchtoRequiredArch(['slc7_amd64_gcc10', 'slc7_aarch64_gcc700']), 'X86_64')
+        self.assertEqual(bp.scramArchtoRequiredArch(
+                ['slc7_amd64_gcc10', 'slc7_aarch64_gcc700', 'slc7_ppc64le_gcc9']), 'X86_64')
+        self.assertEqual(bp.scramArchtoRequiredArch(['slc7_aarch64_gcc700', 'slc7_ppc64le_gcc9']), 'ppc64le')
+
+        return
 
 if __name__ == '__main__':
     unittest.main()