feat: Add secrets in Arex CE

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):
+    def _writeXRSL(self, executableFile, inputs, outputs, diracXSecret):
         """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):
 (inputFiles=({executable} "{executableFile}") {xrslInputAdditions})
 (stdout="{diracStamp}.out")
 (stderr="{diracStamp}.err")
-(environment=("DIRAC_PILOT_STAMP" "{diracStamp}"))
+(environment=("DIRAC_PILOT_STAMP" "{diracStamp}") ("DIRACX_SECRET" "{diracXSecret}"))
 (outputFiles={xrslOutputFiles})
 (queue={queue})
 {xrslMPAdditions}
@@ -476,6 +476,7 @@ def _writeXRSL(self, executableFile, inputs, outputs):
             executable=os.path.basename(executableFile),
             xrslInputAdditions=xrslInputs,
             diracStamp=diracStamp,
+            diracXSecret=diracXSecret,
             queue=self.queue,
             xrslOutputFiles=xrslOutputs,
             xrslMPAdditions=xrslMPAdditions,
@@ -501,7 +502,7 @@ def _bundlePreamble(self, executableFile):
             bundleFile.write(wrapperContent)
             return bundleFile.name
 
-    def _getArcJobID(self, executableFile, inputs, outputs, delegation):
+    def _getArcJobID(self, executableFile, inputs, outputs, delegation, diracXSecret):
         """Get an ARC JobID endpoint to upload executables and inputs.
 
         :param str executableFile: executable to submit
@@ -516,7 +517,7 @@ def _getArcJobID(self, executableFile, inputs, outputs, delegation):
         query = self._urlJoin("jobs")
 
         # Get the job into the ARC way
-        xrslString, diracStamp = self._writeXRSL(executableFile, inputs, outputs)
+        xrslString, diracStamp = self._writeXRSL(executableFile, inputs, outputs, diracXSecret)
         xrslString += delegation
         self.log.debug("XRSL string submitted", f"is {xrslString}")
         self.log.debug("DIRAC stamp for job", f"is {diracStamp}")
@@ -569,10 +570,12 @@ 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):
+    def submitJob(self, executableFile, proxy, numberOfJobs=1, inputs=None, outputs=None, diracXSecrets=[]):
         """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
+
+        For V9+: Will give back also a {"stamp": "secret"} dictionnary.
         """
         result = self._checkSession()
         if not result["OK"]:
@@ -650,8 +653,13 @@ def submitJob(self, executableFile, proxy, numberOfJobs=1, inputs=None, outputs=
         # Also : https://bugzilla.nordugrid.org/show_bug.cgi?id=4069
         batchIDList = []
         stampDict = {}
-        for _ in range(numberOfJobs):
-            result = self._getArcJobID(executableFile, inputs, outputs, delegation)
+        secretDict = {}
+        for i in range(numberOfJobs):
+            if i > len(diracXSecrets):
+                currentSecret = ""
+            else:
+                currentSecret = diracXSecrets[i]
+            result = self._getArcJobID(executableFile, inputs, outputs, delegation, currentSecret)
             if not result["OK"]:
                 break
             arcJobID, diracStamp = result["Value"]
@@ -665,6 +673,8 @@ 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
             self.log.debug(
                 "Successfully submitted job",
                 f"{jobReference} to CE {self.ceName}",
@@ -677,6 +687,7 @@ def submitJob(self, executableFile, proxy, numberOfJobs=1, inputs=None, outputs=
         if batchIDList:
             result = S_OK(batchIDList)
             result["PilotStampDict"] = stampDict
+            result["SecretDict"] = secretDict
         else:
             result = S_ERROR("No ID obtained from the ARC job submission")
         return result

src/DIRAC/Resources/Computing/CloudComputingElement.py

@@ -372,7 +372,7 @@ def __init__(self, *args, **kwargs):
         self.ceType = CE_NAME
         self._cloudDriver = None
 
-    def submitJob(self, executableFile, proxy, numberOfJobs=1):
+    def submitJob(self, executableFile, proxy, numberOfJobs=1, diracXSecrets=[]):
         """Creates VM instances
 
         :param str executableFile: Path to pilot job wrapper file to use

src/DIRAC/Resources/Computing/HTCondorCEComputingElement.py

@@ -308,7 +308,7 @@ def _executeCondorCommand(self, cmd, keepTokenFile=False):
         return S_OK(stdout.strip())
 
     #############################################################################
-    def submitJob(self, executableFile, proxy, numberOfJobs=1):
+    def submitJob(self, executableFile, proxy, numberOfJobs=1, diracXSecrets=[]):
         """Method to submit job"""
 
         self.log.verbose("Executable file path:", executableFile)

src/DIRAC/Resources/Computing/LocalComputingElement.py

@@ -148,7 +148,7 @@ def _prepareHost(self):
 
         return S_OK()
 
-    def submitJob(self, executableFile, proxy=None, numberOfJobs=1):
+    def submitJob(self, executableFile, proxy=None, numberOfJobs=1, diracXSecrets=[]):
         if not os.access(executableFile, 5):
             os.chmod(executableFile, stat.S_IRWXU | stat.S_IRGRP | stat.S_IXGRP | stat.S_IROTH | stat.S_IXOTH)
 

src/DIRAC/Resources/Computing/SSHBatchComputingElement.py

@@ -50,7 +50,7 @@ def _reset(self):
         return S_OK()
 
     #############################################################################
-    def submitJob(self, executableFile, proxy, numberOfJobs=1):
+    def submitJob(self, executableFile, proxy, numberOfJobs=1, diracXSecrets=[]):
         """Method to submit job"""
 
         # Choose eligible hosts, rank them by the number of available slots

src/DIRAC/Resources/Computing/SSHComputingElement.py

@@ -526,7 +526,7 @@ def __executeHostCommand(self, command, options, ssh=None, host=None):
         else:
             return S_ERROR("\n".join([sshStdout, sshStderr]))
 
-    def submitJob(self, executableFile, proxy, numberOfJobs=1):
+    def submitJob(self, executableFile, proxy, numberOfJobs=1, diracXSecrets=[]):
         #    self.log.verbose( "Executable file path: %s" % executableFile )
         if not os.access(executableFile, 5):
             os.chmod(executableFile, stat.S_IRWXU | stat.S_IRGRP | stat.S_IXGRP | stat.S_IROTH | stat.S_IXOTH)

src/DIRAC/WorkloadManagementSystem/Agent/SiteDirector.py

@@ -37,6 +37,7 @@
 from DIRAC.WorkloadManagementSystem.Client.PilotScopes import PILOT_SCOPES
 from DIRAC.WorkloadManagementSystem.Client.ServerUtils import getPilotAgentsDB
 from DIRAC.WorkloadManagementSystem.private.ConfigHelper import findGenericPilotCredentials
+from DIRAC.WorkloadManagementSystem.Client.PilotManagerClient import PilotManagerClient
 from DIRAC.WorkloadManagementSystem.Utilities.PilotWrapper import (
     _writePilotWrapperFile,
     getPilotFilesCompressedEncodedDict,
@@ -103,6 +104,7 @@ def initialize(self):
         self.rssClient = ResourceStatus()
         self.pilotAgentsDB = getPilotAgentsDB()
         self.matcherClient = MatcherClient()
+        self.pilotManagementClient = PilotManagerClient()
 
         return S_OK()
 
@@ -348,15 +350,15 @@ def _submitPilotsPerQueue(self, queueName: str):
         if not result["OK"]:
             self.log.info("Failed pilot submission", f"Queue: {queueName}")
             return result
-        pilotList, stampDict = result["Value"]
+        stampDict, secretDict = result["Value"]
 
-        # updating the pilotAgentsDB... done by default but maybe not strictly necessary
-        result = self._addPilotReferences(queueName, pilotList, stampDict)
+        submittedPilots = len(stampDict)
+        self.log.info("Total number of pilots submitted", f"to {queueName}: {submittedPilots}")
+
+        result = self._addPilotReferences(queueName, stampDict, secretDict)
         if not result["OK"]:
             return result
 
-        submittedPilots = len(pilotList)
-        self.log.info("Total number of pilots submitted", f"to {queueName}: {submittedPilots}")
         return S_OK(submittedPilots)
 
     def _getQueueSlots(self, queue: str):
@@ -460,8 +462,12 @@ 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)
+
         # Submit the job
-        submitResult = ce.submitJob(executable, "", pilotsToSubmit)
+        # NOTE FOR DIRACX /!\ : We need in each CE to create a secret
+        submitResult = ce.submitJob(executable, "", pilotsToSubmit, diracXSecrets=secrets)
+
         # 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
         if submitResult.get("ExecutableToKeep") != executable:
@@ -531,34 +537,56 @@ def _submitPilotsToQueue(self, pilotsToSubmit: int, ce: ComputingElement, queue:
             if not result["OK"]:
                 self.log.error("Failure submitting Monitoring report", result["Message"])
 
-        return S_OK((pilotList, stampDict))
+        secretDict = {}
+        if "SecretDict" in submitResult:
+            # TODO: Update this comment as we add DiracX support
+            # V9+, only for:
+            # 1. Arex
+
+            # Result body: {"secret": "PilotStamps": ["stamp"]}
+            secretDict = submitResult["SecretDict"]
+
+        references = stampDict.keys()
+        stamps = stampDict.values()
+        stampDict = dict(zip(stamps, references))
+
+        return S_OK((stampDict, secretDict))
 
-    def _addPilotReferences(self, queue: str, pilotList: list[str], stampDict: dict[str, str]):
+    def _addPilotReferences(self, queue: str, stampDict: dict[str, str], secretDict: dict[str, str]):
         """Add pilotReference to pilotAgentsDB
 
         :param queue: the queue name
         :param pilotList: list of pilots
-        :param stampDict: dictionary of pilots timestamps
+        :param refDict: dictionary {"pilotstamp":"pilotref"}
+        :param secretDict: dictionary {"pilotstamp":"secret"}
         """
-        result = self.pilotAgentsDB.addPilotReferences(
-            pilotList,
-            self.vo,
-            self.queueDict[queue]["CEType"],
-            stampDict,
+        # FIXME: Change for a client or at least request to DiracX
+
+        # First, create pilots
+        stamps = stampDict.keys()
+        result = self.pilotManagementClient.addPilotReferences(
+            stamps, self.vo, self.queueDict[queue]["CEType"], stampDict
         )
         if not result["OK"]:
-            self.log.error("Failed add pilots to the PilotAgentsDB", result["Message"])
             return result
 
-        for pilot in pilotList:
-            result = self.pilotAgentsDB.setPilotStatus(
-                pilot,
-                PilotStatus.SUBMITTED,
-                self.queueDict[queue]["CEName"],
-                "Successfully submitted by the SiteDirector",
-                self.queueDict[queue]["Site"],
-                self.queueDict[queue]["QueueName"],
+        # We associate all of the pilots with their secrets
+        if secretDict:
+            result = self.pilotManagementClient.associatePilotWithSecret(secretDict)
+            if not result["OK"]:
+                return result
+
+        for stamp in stamps:
+            result = self.pilotManagementClient.set_pilot_field(
+                stamp,
+                {
+                    "DestinationSite": self.queueDict[queue]["CEName"],
+                    "StatusReason": "Successfully submitted by the SiteDirector",
+                    "GridSite": self.queueDict[queue]["Site"],
+                    "Queue": self.queueDict[queue]["QueueName"],
+                },
             )
+
             if not result["OK"]:
                 self.log.error("Failed to set pilot status", result["Message"])
                 return result
@@ -591,14 +619,13 @@ def _getExecutable(self, queue: str, proxy: X509Chain, jobExecDir: str = "", env
         ce = self.queueCECache[queue]["CE"]
         workingDirectory = getattr(ce, "workingDirectory", self.workingDirectory)
 
-        executable = self._writePilotScript(
+        return self._writePilotScript(
             workingDirectory=workingDirectory,
             pilotOptions=pilotOptions,
             proxy=proxy,
             pilotExecDir=jobExecDir,
             envVariables=envVariables,
         )
-        return executable
 
     def _getPilotOptions(self, queue: str) -> list[str]:
         """Prepare pilot options
@@ -680,6 +707,13 @@ def _getPilotOptions(self, queue: str) -> list[str]:
         if "PipInstallOptions" in queueDict:
             pilotOptions.append(f"--pipInstallOptions={queueDict['PipInstallOptions']}")
 
+        # FIXME: Get secret
+        # if "secret" in queueDict:
+        #     pilotOptions.append(f"--pilotSecret={queueDict['...']}")
+        # FIXME: Get clientID
+        # pilotOptions.append(f"--clientID={opsHelper.getValue('TO CHANGE')})
+        pilotOptions.append(f"--diracx_URL={DIRAC.gConfig.getValue('/DiracX/URL')}")
+
         return pilotOptions
 
     def _writePilotScript(

src/DIRAC/WorkloadManagementSystem/FutureClient/PilotManagerClient.py

@@ -9,13 +9,13 @@ def addPilotReferences(self, pilot_stamps, VO, gridType="DIRAC", pilot_reference
         with DiracXClient() as api:
             # We will move toward a stamp as identifier for the pilot
             return api.pilots.add_pilot_stamps(
-                {"pilot_stamps": pilot_stamps, "vo": VO, "grid_type": gridType, "pilot_references": pilot_references}
-            )
+                {"pilot_stamps": pilot_stamps, "vo": VO, "grid_type": gridType, "pilot_references": pilot_references, "generate_secrets": False}  # type: ignore
+            )  # type: ignore
 
     def set_pilot_field(self, pilot_stamp, values_dict):
         with DiracXClient() as api:
             values_dict["PilotStamp"] = pilot_stamp
-            return api.pilots.update_pilot_fields(values_dict)
+            return api.pilots.update_pilot_fields({"pilot_stamps_to_fields_mapping": [values_dict]})  # type: ignore
 
     @convertToReturnValue
     def setPilotBenchmark(self, pilotStamp, mark):
@@ -48,7 +48,7 @@ def clearPilots(self, interval=30, aborted_interval=7):
     def deletePilots(self, pilot_stamps):
         with DiracXClient() as api:
             pilot_ids = None
-            if isinstance(pilot_stamps, list[int]):
+            if isinstance(pilot_stamps, list[int]):  # type: ignore
                 # Multiple elements (int)
                 pilot_ids = pilot_stamps  # Semantic
             elif isinstance(pilot_stamps, int):
@@ -66,12 +66,12 @@ def deletePilots(self, pilot_stamps):
                 pilots = api.pilots.search(parameters=["PilotStamp"], search=query, sort=[])
                 pilot_stamps = [pilot["PilotStamp"] for pilot in pilots]
 
-            api.pilots.delete_pilots(pilot_stamps=pilot_stamps)
+            api.pilots.delete_pilots(pilot_stamps=pilot_stamps)  # type: ignore
 
     @convertToReturnValue
     def setJobForPilot(self, job_id, pilot_stamp, destination=None):
         with DiracXClient() as api:
-            api.pilots.add_jobs_to_pilot({"pilot_stamp": pilot_stamp, "job_ids": [job_id]})
+            api.pilots.add_jobs_to_pilot({"pilot_stamp": pilot_stamp, "job_ids": [job_id]})  # type: ignore
 
             self.set_pilot_field(
                 pilot_stamp,
@@ -95,3 +95,21 @@ def getPilotInfo(self, pilot_stamp):
             query = [{"parameter": "PilotStamp", "operator": "eq", "value": pilot_stamp}]
 
             return api.pilots.search(parameters=[], search=query, sort=[])
+
+    @convertToReturnValue
+    def associatePilotWithSecret(self, secretDict):
+        # secretDict format: {"secret": ["stamp"]}
+        with DiracXClient() as api:
+            return api.pilots.update_secrets_constraints(secretDict)  # type: ignore
+
+    @convertToReturnValue
+    def createNSecrets(self, vo, n=100, expiration_minutes=120, pilot_secret_use_count_max=1):
+        with DiracXClient() as api:
+            return api.pilots.create_pilot_secrets(
+                {
+                    "n": n,
+                    "expiration_minutes": expiration_minutes,
+                    "pilot_secret_use_count_max": pilot_secret_use_count_max,
+                    "vo": vo,
+                }
+            )  # type: ignore

src/DIRAC/WorkloadManagementSystem/Utilities/PilotWrapper.py

@@ -182,6 +182,8 @@ def pilotWrapperScript(
 
     if envVariables is None:
         envVariables = {}
+    elif "DIRACX_SECRET" in envVariables:
+        pilotOptions += f" --pilotSecret={envVariables['DIRACX_SECRET']}"
 
     if not CVMFS_locations:
         # What is in this location is almost certainly incorrect, especially the pilot.json

src/DIRAC/WorkloadManagementSystem/scripts/dirac_admin_add_pilot.py

@@ -51,6 +51,8 @@ def main():
     """
     params = Params()
 
+    # TODO: Add also site here.
+    # Later deprecated in V9.
     Script.registerSwitches(params.switches)
     Script.registerArgument("pilotRef: pilot reference")
     Script.registerArgument("VO: VO, or pilot owner group")