feat: set env variable OMP_NUM_THREADS

fstagni · Nov 27, 2023 · 2018021 · 2018021
1 parent df4e1ff
commit 2018021
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/src/DIRAC/Resources/Computing/InProcessComputingElement.py b/src/DIRAC/Resources/Computing/InProcessComputingElement.py
@@ -35,7 +35,7 @@ def submitJob(self, executableFile, proxy=None, inputs=None, **kwargs):
         """
         payloadEnv = dict(os.environ)
         if mp_threads := kwargs.get("numberOfProcessors"):
-            payloadEnv["OMP_NUM_THREADS"] = mp_threads
+            payloadEnv["OMP_NUM_THREADS"] = str(mp_threads)
         payloadProxy = ""
         if proxy:
             self.log.verbose("Setting up proxy for payload")

diff --git a/src/DIRAC/Resources/Computing/SingularityComputingElement.py b/src/DIRAC/Resources/Computing/SingularityComputingElement.py
@@ -94,6 +94,7 @@
     r"Xrd.*",
     r"DIRAC_.*",
     r"BEARER_TOKEN.*",
+    r"OMP_NUM_TREADS",
 ]
 ENV_VAR_WHITELIST = re.compile(r"^(" + r"|".join(ENV_VAR_WHITELIST) + r")$")
 
@@ -374,6 +375,9 @@ def submitJob(self, executableFile, proxy=None, **kwargs):
         baseDir = ret["baseDir"]
         tmpDir = ret["tmpDir"]
 
+        if mp_threads := kwargs.get("numberOfProcessors"):
+            os.environ["OMP_NUM_THREADS"] = str(mp_threads)
+
         if proxy:
             payloadProxyLoc = ret["proxyLocation"]