feat: Migrating client toward DiracX

Author: Robin-Van-de-Merghel

src/DIRAC/WorkloadManagementSystem/FutureClient/PilotManagerClient.py

@@ -0,0 +1,101 @@
+from diracx.client._generated.models import BodyPilotsAddJobsToPilot, BodyPilotsAddPilotStamps, ScalarSearchSpec, SearchParams
+from DIRAC.Core.Utilities.ReturnValues import convertToReturnValue
+
+from DIRAC.Core.Security.DiracX import DiracXClient
+
+
+class PilotManagerClient:
+    @convertToReturnValue
+    def addPilotReferences(self, pilot_stamps, VO, gridType="DIRAC", pilot_references={}):
+        with DiracXClient() as api:
+            # We will move toward a stamp as identifier for the pilot
+            return api.pilots.add_pilot_stamps(BodyPilotsAddPilotStamps(
+                pilot_stamps=pilot_stamps,
+                vo=VO,
+                grid_type=gridType,
+                pilot_references=pilot_references
+            ))
+
+
+    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
+            )
+
+    @convertToReturnValue
+    def setPilotBenchmark(self, pilotStamp, mark):
+        return self.set_pilot_field(pilotStamp, {"BenchMark": mark})
+
+    @convertToReturnValue
+    def setAccountingFlag(self, pilotStamp, flag):
+        return self.set_pilot_field(pilotStamp, {"AccountingSent": flag})
+
+    @convertToReturnValue
+    def setPilotStatus(self, pilot_stamp, status, destination=None, reason=None, grid_site=None, queue=None):
+        return self.set_pilot_field(pilot_stamp, {
+            "Status": status,
+            "DestinationSite": destination,
+            "StatusReason": reason,
+            "GridSite": grid_site,
+            "Queue": queue
+        })
+    
+    @convertToReturnValue
+    def clearPilots(self, interval=30, aborted_interval=7):
+        with DiracXClient() as api:
+            api.pilots.delete_pilots(
+                age_in_days=interval,
+                delete_only_aborted=False
+            )
+            api.pilots.delete_pilots(
+                age_in_days=aborted_interval,
+                delete_only_aborted=True
+            )
+
+    @convertToReturnValue
+    def deletePilots(self, pilot_stamps):
+        with DiracXClient() as api:
+            api.pilots.delete_pilots(
+                pilot_stamps=pilot_stamps 
+            )
+
+    @convertToReturnValue
+    def setJobForPilot(self, job_id, pilot_stamp, destination=None):
+        with DiracXClient() as api:
+            api.pilots.add_jobs_to_pilot(BodyPilotsAddJobsToPilot(
+                pilot_stamp=pilot_stamp, 
+                job_ids=[job_id]
+            ))
+
+            self.set_pilot_field(pilot_stamp, {
+                "DestinationSite": destination,
+            })
+    
+    @convertToReturnValue
+    def getPilots(self, job_id):
+        with DiracXClient() as api:
+            pilot_ids = api.pilots.get_pilot_jobs(
+                job_id=job_id
+            )
+
+            query = [{"parameter": "PilotID", "operator": "in", "value": pilot_ids}]
+
+            return api.pilots.search(
+                parameters=[],
+                search=query,
+                sort=[]
+            )
+
+    
+    @convertToReturnValue
+    def getPilotInfo(self, pilot_stamp):
+        with DiracXClient() as api:
+            query = [{"parameter": "PilotStamp", "operator": "eq", "value": pilot_stamp}]
+
+            return api.pilots.search(
+                parameters=[],
+                search=query,
+                sort=[]
+            )

src/DIRAC/WorkloadManagementSystem/Service/PilotManagerHandler.py

@@ -69,14 +69,6 @@ def export_getCurrentPilotCounters(cls, attrDict={}):
 
         return S_OK(resultDict)
 
-    ##########################################################################################
-    types_addPilotReferences = [list, str]
-
-    @classmethod
-    def export_addPilotReferences(cls, pilotRef, VO, gridType="DIRAC", pilotStampDict={}):
-        """Add a new pilot job reference"""
-        return cls.pilotAgentsDB.addPilotReferences(pilotRef, VO, gridType, pilotStampDict)
-
     ##############################################################################
     types_getPilotOutput = [str]
 
@@ -202,17 +194,11 @@ def _getRemotePilotOutput(self, pilotReference, pilotDict):
         # return res, correct or not
         return res
 
-    ##############################################################################
-    types_getPilotInfo = [[list, str]]
-
-    @classmethod
-    def export_getPilotInfo(cls, pilotReference):
-        """Get the info about a given pilot job reference"""
-        return cls.pilotAgentsDB.getPilotInfo(pilotReference)
 
     ##############################################################################
     types_selectPilots = [dict]
 
+    # Won't be moved to DiracX: not used at all anywhere.
     @classmethod
     def export_selectPilots(cls, condDict):
         """Select pilots given the selection conditions"""
@@ -288,21 +274,35 @@ def export_getGroupedPilotSummary(cls, columnList):
         """
         return cls.pilotAgentsDB.getGroupedPilotSummary(columnList)
 
-    ##############################################################################
-    types_getPilots = [[str, int]]
+
+    types_countPilots = [dict]
 
     @classmethod
-    def export_getPilots(cls, jobID):
-        """Get pilots executing/having executed the Job"""
-        result = cls.pilotAgentsDB.getPilotsForJobID(int(jobID))
-        if not result["OK"] or not result["Value"]:
-            return S_ERROR(f"Failed to get pilot for Job {int(jobID)}: {result.get('Message', '')}")
+    def export_countPilots(cls, condDict, older=None, newer=None, timeStamp="SubmissionTime"):
+        """Set the pilot agent status"""
 
-        return cls.pilotAgentsDB.getPilotInfo(pilotID=result["Value"])
+        return cls.pilotAgentsDB.countPilots(condDict, older, newer, timeStamp)
 
-    ##############################################################################
-    types_setJobForPilot = [[str, int], str]
 
+    # --------------- Moved to DiracX ---------------
+
+    #############################################
+    types_addPilotReferences = [list, str]
+
+    # Moved to DiracX
+    @classmethod
+    def export_addPilotReferences(cls, pilotStamps, VO, gridType="DIRAC", pilotRefDict={}):
+        """Add a new pilot job reference"""
+        pilot_references = pilotRefDict.values()
+        pilot_stamp_dict = dict(zip(pilotStamps, pilot_references))
+
+
+        return cls.pilotAgentsDB.addPilotReferences(pilot_references, VO, gridType, pilot_stamp_dict)
+
+
+    #############################################
+    types_setJobForPilot = [[str, int], str]
+    
     @classmethod
     def export_setJobForPilot(cls, jobID, pilotRef, destination=None):
         """Report the DIRAC job ID which is executed by the given pilot job"""
@@ -318,23 +318,24 @@ def export_setJobForPilot(cls, jobID, pilotRef, destination=None):
 
         return result
 
-    ##########################################################################################
+    #############################################
     types_setPilotBenchmark = [str, float]
 
     @classmethod
     def export_setPilotBenchmark(cls, pilotRef, mark):
         """Set the pilot agent benchmark"""
         return cls.pilotAgentsDB.setPilotBenchmark(pilotRef, mark)
 
-    ##########################################################################################
+    #############################################
     types_setAccountingFlag = [str]
 
     @classmethod
     def export_setAccountingFlag(cls, pilotRef, mark="True"):
         """Set the pilot AccountingSent flag"""
         return cls.pilotAgentsDB.setAccountingFlag(pilotRef, mark)
 
-    ##########################################################################################
+
+    #############################################
     types_setPilotStatus = [str, str]
 
     @classmethod
@@ -345,22 +346,16 @@ def export_setPilotStatus(cls, pilotRef, status, destination=None, reason=None,
             pilotRef, status, destination=destination, statusReason=reason, gridSite=gridSite, queue=queue
         )
 
-    ##########################################################################################
-    types_countPilots = [dict]
-
-    @classmethod
-    def export_countPilots(cls, condDict, older=None, newer=None, timeStamp="SubmissionTime"):
-        """Set the pilot agent status"""
-
-        return cls.pilotAgentsDB.countPilots(condDict, older, newer, timeStamp)
-
-    ##########################################################################################
+    #############################################
     types_deletePilots = [[list, str, int]]
-
+    
     @classmethod
     def export_deletePilots(cls, pilotIDs):
         if isinstance(pilotIDs, str):
             return cls.pilotAgentsDB.deletePilot(pilotIDs)
+        
+        # And list[str]????
+        # pilot_id>>>S<<<
 
         if isinstance(pilotIDs, int):
             pilotIDs = [
@@ -373,9 +368,29 @@ def export_deletePilots(cls, pilotIDs):
 
         return S_OK()
 
-    ##############################################################################
+    #############################################
     types_clearPilots = [int, int]
 
     @classmethod
     def export_clearPilots(cls, interval=30, aborted_interval=7):
         return cls.pilotAgentsDB.clearPilots(interval, aborted_interval)
+    
+    ##############################################################################
+    types_getPilots = [[str, int]]
+
+    @classmethod
+    def export_getPilots(cls, jobID):
+        """Get pilots executing/having executed the Job"""
+        result = cls.pilotAgentsDB.getPilotsForJobID(int(jobID))
+        if not result["OK"] or not result["Value"]:
+            return S_ERROR(f"Failed to get pilot for Job {int(jobID)}: {result.get('Message', '')}")
+
+        return cls.pilotAgentsDB.getPilotInfo(pilotID=result["Value"])
+    
+    ##############################################################################
+    types_getPilotInfo = [[list, str]]
+
+    @classmethod
+    def export_getPilotInfo(cls, pilotReference):
+        """Get the info about a given pilot job reference"""
+        return cls.pilotAgentsDB.getPilotInfo(pilotReference)

src/DIRAC/WorkloadManagementSystem/scripts/dirac_admin_add_pilot.py

@@ -74,7 +74,7 @@ def main():
         if not DErrno.cmpError(res, DErrno.EWMSNOPILOT):
             gLogger.error(res["Message"])
             DIRACExit(1)
-        res = pmc.addPilotReferences([pilotRef], VO, gridType, {pilotRef: pilotStamp})
+        res = pmc.addPilotReferences(pilotStamp, VO, gridType, {pilotStamp: pilotRef})
         if not res["OK"]:
             gLogger.error(res["Message"])
             DIRACExit(1)

tests/Integration/WorkloadManagementSystem/Test_PilotsClient.py

@@ -25,8 +25,8 @@ def test_PilotsDB():
     webapp = WebAppClient()
 
     # This will allow you to run the test again if necessary
-    for jobID in ["aPilot", "anotherPilot"]:
-        pilots.deletePilots(jobID)
+    for pilot_stamp in ["aPilot", "anotherPilot"]:
+        pilots.deletePilot(pilot_stamp)
 
     res = pilots.addPilotReferences(["aPilot"], "VO")
     assert res["OK"], res["Message"]