feat: Allow to add envMapping into a wrapper

Author: Robin VAN DE MERGHEL

src/DIRAC/Resources/Computing/AREXComputingElement.py

@@ -417,7 +417,7 @@ def _getProxyFromDelegationID(self, delegationID):
 
     #############################################################################
 
-    def _writeXRSL(self, executableFile, inputs, outputs, diracXSecret):
+    def _writeXRSL(self, executableFile, inputs, outputs, additionalEnv):
         """Create the JDL for submission
 
         :param str executableFile: executable to wrap in a XRSL file
@@ -465,7 +465,7 @@ def _writeXRSL(self, executableFile, inputs, outputs, diracXSecret):
 (inputFiles=({executable} "{executableFile}") {xrslInputAdditions})
 (stdout="{diracStamp}.out")
 (stderr="{diracStamp}.err")
-(environment=("DIRAC_PILOT_STAMP" "{diracStamp}") ("DIRACX_SECRET" "{diracXSecret}"))
+(environment=("DIRAC_PILOT_STAMP" "{diracStamp}") {additionalEnv})
 (outputFiles={xrslOutputFiles})
 (queue={queue})
 {xrslMPAdditions}
@@ -476,7 +476,7 @@ def _writeXRSL(self, executableFile, inputs, outputs, diracXSecret):
             executable=os.path.basename(executableFile),
             xrslInputAdditions=xrslInputs,
             diracStamp=diracStamp,
-            diracXSecret=diracXSecret,
+            additionalEnv=additionalEnv,
             queue=self.queue,
             xrslOutputFiles=xrslOutputs,
             xrslMPAdditions=xrslMPAdditions,
@@ -502,7 +502,7 @@ def _bundlePreamble(self, executableFile):
             bundleFile.write(wrapperContent)
             return bundleFile.name
 
-    def _getArcJobID(self, executableFile, inputs, outputs, delegation, diracXSecret):
+    def _getArcJobID(self, executableFile, inputs, outputs, delegation, additionalEnv):
         """Get an ARC JobID endpoint to upload executables and inputs.
 
         :param str executableFile: executable to submit
@@ -517,7 +517,7 @@ def _getArcJobID(self, executableFile, inputs, outputs, delegation, diracXSecret
         query = self._urlJoin("jobs")
 
         # Get the job into the ARC way
-        xrslString, diracStamp = self._writeXRSL(executableFile, inputs, outputs, diracXSecret)
+        xrslString, diracStamp = self._writeXRSL(executableFile, inputs, outputs, additionalEnv)
         xrslString += delegation
         self.log.debug("XRSL string submitted", f"is {xrslString}")
         self.log.debug("DIRAC stamp for job", f"is {diracStamp}")
@@ -570,7 +570,7 @@ def _uploadJobDependencies(self, arcJobID, executableFile, inputs):
             self.log.verbose("Input correctly uploaded", fileToSubmit)
         return S_OK()
 
-    def submitJob(self, executableFile, proxy, numberOfJobs=1, inputs=None, outputs=None, diracXSecrets=[]):
+    def submitJob(self, executableFile, proxy, numberOfJobs=1, inputs=None, outputs=None, additionalEnv=[]):
         """Method to submit job
         Assume that the ARC queues are always of the format nordugrid-<batchSystem>-<queue>
         And none of our supported batch systems have a "-" in their name
@@ -653,13 +653,16 @@ def submitJob(self, executableFile, proxy, numberOfJobs=1, inputs=None, outputs=
         # Also : https://bugzilla.nordugrid.org/show_bug.cgi?id=4069
         batchIDList = []
         stampDict = {}
-        secretDict = {}
+        additionalEnvMappingResponse = {}
         for i in range(numberOfJobs):
-            if i > len(diracXSecrets):
-                currentSecret = ""
+            if i > len(additionalEnv):
+                currentEnv = ""
             else:
-                currentSecret = diracXSecrets[i]
-            result = self._getArcJobID(executableFile, inputs, outputs, delegation, currentSecret)
+                # AdditionalEnv[i] format:
+                # {"secret": "1_l0v3_1c3cr34m", ...}
+                # We merge them to have the right format: '("key" "value") (...)'
+                currentEnv = " ".join([f"({key} {value})" for key, value in additionalEnv[i]])
+            result = self._getArcJobID(executableFile, inputs, outputs, delegation, currentEnv)
             if not result["OK"]:
                 break
             arcJobID, diracStamp = result["Value"]
@@ -673,8 +676,11 @@ def submitJob(self, executableFile, proxy, numberOfJobs=1, inputs=None, outputs=
             jobReference = self._arcIDToJobReference(arcJobID)
             batchIDList.append(jobReference)
             stampDict[jobReference] = diracStamp
-            secretDict[currentSecret] = {}
-            secretDict[currentSecret]["PilotStamps"] = [diracStamp]  # Used by DiracX to associate secrets and pilots
+
+            # Add all env variables we added into additionalEnvMappingResponse so we have:
+            # { "Stamp1": { "SECRET": "I_luv_strawberries", "...": "..." }, "Stamp2": {...} }
+            additionalEnvMappingResponse[diracStamp] = additionalEnv[i]
+
             self.log.debug(
                 "Successfully submitted job",
                 f"{jobReference} to CE {self.ceName}",
@@ -687,7 +693,7 @@ def submitJob(self, executableFile, proxy, numberOfJobs=1, inputs=None, outputs=
         if batchIDList:
             result = S_OK(batchIDList)
             result["PilotStampDict"] = stampDict
-            result["SecretDict"] = secretDict
+            result["EnvMapping"] = additionalEnvMappingResponse
         else:
             result = S_ERROR("No ID obtained from the ARC job submission")
         return result

src/DIRAC/WorkloadManagementSystem/Agent/SiteDirector.py

@@ -462,11 +462,31 @@ def _submitPilotsToQueue(self, pilotsToSubmit: int, ce: ComputingElement, queue:
         jobProxy = result["Value"]
         executable = self._getExecutable(queue, proxy=jobProxy, jobExecDir=jobExecDir, envVariables=envVariables)
 
-        secrets = self.pilotManagementClient.createNSecrets(vo=self.vo, n=pilotsToSubmit)
+        # ----------- Generate additionnal env variables -----------
+
+        additionalEnv = []
+
+        # First, create secrets
+        resultSecrets = self.pilotManagementClient.createNSecrets(vo=self.vo, n=pilotsToSubmit)
+        if not resultSecrets["OK"]:
+            self.log.error("Couldn't fetch secrets.")
+            secrets = []
+        else:
+            secrets = [{"DIRACX_SECRET": secret["pilot_secret"]} for secret in resultSecrets["Value"]]
+
+        # ---
+        # More env?
+        # ---
+
+        # Could be secrets, but also other things, increase in the for loop...
+        for i, el in enumerate(secrets):
+            additionalEnv[i].update(el)
+
+        # ----------------------------------------------------------
 
         # Submit the job
         # NOTE FOR DIRACX /!\ : We need in each CE to create a secret
-        submitResult = ce.submitJob(executable, "", pilotsToSubmit, diracXSecrets=secrets)
+        submitResult = ce.submitJob(executable, "", pilotsToSubmit, additionalEnv=additionalEnv)
 
         # In case the CE does not need the executable after the submission, we delete it
         # Else, we keep it, the CE will delete it after the end of the pilot execution
@@ -537,20 +557,30 @@ def _submitPilotsToQueue(self, pilotsToSubmit: int, ce: ComputingElement, queue:
             if not result["OK"]:
                 self.log.error("Failure submitting Monitoring report", result["Message"])
 
-        secretDict = {}
-        if "SecretDict" in submitResult:
+        secrets_request_body = {}
+        if "EnvMapping" in submitResult:
             # TODO: Update this comment as we add DiracX support
             # V9+, only for:
             # 1. Arex
 
-            # Result body: {"secret": "PilotStamps": ["stamp"]}
-            secretDict = submitResult["SecretDict"]
+            # EnvMapping body: { "Stamp1": { "DIRACX_SECRET": "I_luv_strawberries", "...": "..." }, "Stamp2": {...} }
+            # Request body for secrets: {"I_luv_stawberries": ["Stamp1", "..."]}
+            envMapping = submitResult["SecretDict"]
+            secrets_request_body = {}
+
+            for stamp, envVariables in envMapping.items():
+                diracx_secret = envVariables.get("DIRACX_SECRET")
+                if diracx_secret:
+                    if not diracx_secret in secrets_request_body:
+                        secrets_request_body[diracx_secret] = []
+                    secrets_request_body[diracx_secret].append(stamp)
 
+        # We reverse the ref-stamp dictionnary (DiracX behaviour)
         references = stampDict.keys()
         stamps = stampDict.values()
         stampDict = dict(zip(stamps, references))
 
-        return S_OK((stampDict, secretDict))
+        return S_OK((stampDict, secrets_request_body))
 
     def _addPilotReferences(self, queue: str, stampDict: dict[str, str], secretDict: dict[str, str]):
         """Add pilotReference to pilotAgentsDB