feat: Modified legacy adaptor for PilotManager to make it work.

Author: Robin VAN DE MERGHEL

src/DIRAC/WorkloadManagementSystem/Client/PilotManagerClient.py

@@ -4,10 +4,17 @@
 from DIRAC.Core.Base.Client import Client, createClient
 
 
+from DIRAC.WorkloadManagementSystem.FutureClient.PilotManagerClient import (
+    PilotManagerClient as futurePilotManagerClient,
+)
+
+
 @createClient("WorkloadManagement/PilotManager")
 class PilotManagerClient(Client):
     """PilotManagerClient sets url for the PilotManagerHandler."""
 
+    diracxClient = futurePilotManagerClient
+
     def __init__(self, url=None, **kwargs):
         """
         Sets URL for PilotManager handler

src/DIRAC/WorkloadManagementSystem/FutureClient/PilotManagerClient.py

@@ -1,29 +1,20 @@
 from DIRAC.Core.Utilities.ReturnValues import convertToReturnValue
+from DIRAC.Core.Security.DiracX import DiracXClient, FutureClient
 
-from DIRAC.Core.Security.DiracX import DiracXClient
 
-
-class PilotManagerClient:
+class PilotManagerClient(FutureClient):
     @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(
-                {"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}  # 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)
-
-    @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})
+            return api.pilots.update_pilot_fields({"pilot_stamps_to_fields_mapping": [values_dict]})  # type: ignore
 
     @convertToReturnValue
     def setPilotStatus(self, pilot_stamp, status, destination=None, reason=None, grid_site=None, queue=None):
@@ -44,13 +35,74 @@ def clearPilots(self, interval=30, aborted_interval=7):
             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 deletePilot(self, pilot_stamp):
+        with DiracXClient() as api:
+            pilot_stamps = [pilot_stamp]
+            return api.pilots.delete_pilots(pilot_stamps=pilot_stamps)
+
+    @convertToReturnValue
+    def getJobsForPilotByStamp(self, pilotStamp):
+        with DiracXClient() as api:
+            return api.pilots.get_pilot_jobs(pilot_stamp=pilotStamp)
+
+    @convertToReturnValue
+    def getPilots(self, job_id):
+        with DiracXClient() as api:
+            pilot_ids = api.pilots.get_pilot_jobs(job_id=job_id)
+            search = [{"parameter": "PilotID", "operator": "in", "value": pilot_ids}]
+            return api.pilots.search(parameters=[], search=search, sort=[])  # type: ignore
+
+    @convertToReturnValue
+    def getPilotInfo(self, pilot_reference):
+        """Important: We assume that to one stamp is mapped one pilot."""
+        with DiracXClient() as api:
+            search = [{"parameter": "PilotJobReference", "operator": "eq", "value": pilot_reference}]
+            pilot = api.pilots.search(parameters=[], search=search, sort=[])[0]  # type: ignore
+
+            if not pilot:
+                # Return an error as in the legacy code
+                return []
+
+            # Convert all bools in pilot to str
+            for k, v in pilot.items():
+                if isinstance(v, bool):
+                    pilot[k] = str(v)
+
+            # Transform the list of pilots into a dict keyed by PilotJobReference
+            resDict = {}
+
+            pilotRef = pilot.get("PilotJobReference", None)
+            assert pilot_reference == pilotRef
+            pilotStamp = pilot.get("PilotStamp", None)
+
+            if pilotRef is not None:
+                resDict[pilotRef] = pilot
+            else:
+                # Fallback: use PilotStamp or another key if PilotJobReference is missing
+                resDict[pilotStamp] = pilot
+
+            jobIDs = self.getJobsForPilotByStamp(pilotStamp)
+            if jobIDs:  # Only add if jobs exist
+                for pilotRef, pilotInfo in resDict.items():
+                    pilotInfo["Jobs"] = jobIDs  # Attach the entire list
+
+            return resDict
+
+    @convertToReturnValue
+    def getGroupedPilotSummary(self, column_list):
+        with DiracXClient() as api:
+            return api.pilots.summary(grouping=column_list)
+
     @convertToReturnValue
     def deletePilots(self, pilot_stamps):
+        # Used by no one, but we won't raise `UnimplementedError` because we still use it in tests.
         with DiracXClient() as api:
             pilot_ids = None
-            if isinstance(pilot_stamps, list[int]):
-                # Multiple elements (int)
-                pilot_ids = pilot_stamps  # Semantic
+            if pilot_stamps and isinstance(pilot_stamps, list):
+                if isinstance(pilot_stamps[0], int):
+                    # Multiple elements (int)
+                    pilot_ids = pilot_stamps  # Semantic
             elif isinstance(pilot_stamps, int):
                 # Only one element (int)
                 pilot_ids = [pilot_stamps]
@@ -61,37 +113,45 @@ def deletePilots(self, pilot_stamps):
 
             if pilot_ids:
                 # If we have defined pilot_ids, then we have to change them to pilot_stamps
-                query = [{"parameter": "PilotID", "operator": "in", "value": pilot_ids}]
+                search = [{"parameter": "PilotID", "operator": "in", "value": pilot_ids}]
 
-                pilots = api.pilots.search(parameters=["PilotStamp"], search=query, sort=[])
+                pilots = api.pilots.search(parameters=["PilotStamp"], search=search, sort=[])  # type: ignore
                 pilot_stamps = [pilot["PilotStamp"] for pilot in pilots]
 
-            api.pilots.delete_pilots(pilot_stamps=pilot_stamps)
+            return 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]})
-
-            self.set_pilot_field(
-                pilot_stamp,
-                {
-                    "DestinationSite": destination,
-                },
-            )
+        raise NotImplementedError(
+            "This function is used by no one. I won't be adapted into DiracX via a legacy adaptor."
+        )
 
     @convertToReturnValue
-    def getPilots(self, job_id):
-        with DiracXClient() as api:
-            pilot_ids = api.pilots.get_pilot_jobs(job_id=job_id)
+    def countPilots(self, condDict, older=None, newer=None):
+        raise NotImplementedError(
+            "This function is used by no one. I won't be adapted into DiracX via a legacy adaptor."
+        )
 
-            query = [{"parameter": "PilotID", "operator": "in", "value": pilot_ids}]
+    @convertToReturnValue
+    def selectPilots(self, condDict):
+        raise NotImplementedError(
+            "This function is used by no one. I won't be adapted into DiracX via a legacy adaptor."
+        )
 
-            return api.pilots.search(parameters=[], search=query, sort=[])
+    @convertToReturnValue
+    def getCurrentPilotCounters(self, attrDict={}):
+        raise NotImplementedError(
+            "This function is used by no one. I won't be adapted into DiracX via a legacy adaptor."
+        )
 
     @convertToReturnValue
-    def getPilotInfo(self, pilot_stamp):
-        with DiracXClient() as api:
-            query = [{"parameter": "PilotStamp", "operator": "eq", "value": pilot_stamp}]
+    def setPilotBenchmark(self, pilotRef, mark):
+        raise NotImplementedError(
+            "This function is used by no one. I won't be adapted into DiracX via a legacy adaptor."
+        )
 
-            return api.pilots.search(parameters=[], search=query, sort=[])
+    @convertToReturnValue
+    def setAccountingFlag(self, pilotRef, flag="True"):
+        raise NotImplementedError(
+            "This function is used by no one. I won't be adapted into DiracX via a legacy adaptor."
+        )

src/DIRAC/WorkloadManagementSystem/Service/PilotManagerHandler.py

@@ -38,38 +38,7 @@ def initializeHandler(cls, serviceInfoDict):
         return S_OK()
 
     ##############################################################################
-    types_getCurrentPilotCounters = [dict]
-
-    @classmethod
-    def export_getCurrentPilotCounters(cls, attrDict={}):
-        """Get pilot counters per Status with attrDict selection. Final statuses are given for
-        the last day.
-        """
-
-        result = cls.pilotAgentsDB.getCounters("PilotAgents", ["Status"], attrDict, timeStamp="LastUpdateTime")
-        if not result["OK"]:
-            return result
-        last_update = datetime.datetime.utcnow() - TimeUtilities.day
-        resultDay = cls.pilotAgentsDB.getCounters(
-            "PilotAgents", ["Status"], attrDict, newer=last_update, timeStamp="LastUpdateTime"
-        )
-        if not resultDay["OK"]:
-            return resultDay
-
-        resultDict = {}
-        for statusDict, count in result["Value"]:
-            status = statusDict["Status"]
-            resultDict[status] = count
-            if status in PilotStatus.PILOT_FINAL_STATES:
-                resultDict[status] = 0
-                for statusDayDict, ccount in resultDay["Value"]:
-                    if status == statusDayDict["Status"]:
-                        resultDict[status] = ccount
-                    break
-
-        return S_OK(resultDict)
 
-    ##############################################################################
     types_getPilotOutput = [str]
 
     def export_getPilotOutput(self, pilotReference):
@@ -198,15 +167,7 @@ def _getRemotePilotOutput(self, pilotReference, pilotDict):
         return res
 
     ##############################################################################
-    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"""
-        return cls.pilotAgentsDB.selectPilots(condDict)
-
-    ##############################################################################
     types_storePilotOutput = [str, str, str]
 
     @classmethod
@@ -263,6 +224,19 @@ def export_getPilotSummary(cls, startdate="", enddate=""):
         return cls.pilotAgentsDB.getPilotSummary(startdate, enddate)
 
     ##############################################################################
+
+    # --------------- Moved to DiracX ---------------
+
+    types_selectPilots = [dict]
+
+    @classmethod
+    def export_selectPilots(cls, condDict):
+        """Select pilots given the selection conditions"""
+        # Used by no one
+        return cls.pilotAgentsDB.selectPilots(condDict)
+
+    #############################################
+
     types_getGroupedPilotSummary = [list]
 
     @classmethod
@@ -280,11 +254,41 @@ def export_getGroupedPilotSummary(cls, columnList):
 
     @classmethod
     def export_countPilots(cls, condDict, older=None, newer=None, timeStamp="SubmissionTime"):
-        """Set the pilot agent status"""
-
+        """Count pilots"""
+        # Used by no one
         return cls.pilotAgentsDB.countPilots(condDict, older, newer, timeStamp)
 
-    # --------------- Moved to DiracX ---------------
+    types_getCurrentPilotCounters = [dict]
+
+    @classmethod
+    def export_getCurrentPilotCounters(cls, attrDict={}):
+        """Get pilot counters per Status with attrDict selection. Final statuses are given for
+        the last day.
+        """
+        # Used by no one
+
+        result = cls.pilotAgentsDB.getCounters("PilotAgents", ["Status"], attrDict, timeStamp="LastUpdateTime")
+        if not result["OK"]:
+            return result
+        last_update = datetime.datetime.utcnow() - TimeUtilities.day
+        resultDay = cls.pilotAgentsDB.getCounters(
+            "PilotAgents", ["Status"], attrDict, newer=last_update, timeStamp="LastUpdateTime"
+        )
+        if not resultDay["OK"]:
+            return resultDay
+
+        resultDict = {}
+        for statusDict, count in result["Value"]:
+            status = statusDict["Status"]
+            resultDict[status] = count
+            if status in PilotStatus.PILOT_FINAL_STATES:
+                resultDict[status] = 0
+                for statusDayDict, ccount in resultDay["Value"]:
+                    if status == statusDayDict["Status"]:
+                        resultDict[status] = ccount
+                    break
+
+        return S_OK(resultDict)
 
     #############################################
     types_addPilotReferences = [list, str]
@@ -304,6 +308,7 @@ def export_addPilotReferences(cls, pilotStamps, VO, gridType="DIRAC", pilotRefDi
     @classmethod
     def export_setJobForPilot(cls, jobID, pilotRef, destination=None):
         """Report the DIRAC job ID which is executed by the given pilot job"""
+        # Used by no one.
 
         result = cls.pilotAgentsDB.setJobForPilot(int(jobID), pilotRef)
         if not result["OK"]:
@@ -322,14 +327,16 @@ def export_setJobForPilot(cls, jobID, pilotRef, destination=None):
     @classmethod
     def export_setPilotBenchmark(cls, pilotRef, mark):
         """Set the pilot agent benchmark"""
+        # Used by no one.
         return cls.pilotAgentsDB.setPilotBenchmark(pilotRef, mark)
 
     #############################################
     types_setAccountingFlag = [str]
 
     @classmethod
-    def export_setAccountingFlag(cls, pilotRef, mark="True"):
+    def export_setAccountingFlag(cls, pilotRef, flag="True"):
         """Set the pilot AccountingSent flag"""
+        # Used by no one
         return cls.pilotAgentsDB.setAccountingFlag(pilotRef, mark)
 
     #############################################
@@ -348,6 +355,7 @@ def export_setPilotStatus(cls, pilotRef, status, destination=None, reason=None,
 
     @classmethod
     def export_deletePilots(cls, pilotIDs):
+        # Used by no one.
         if isinstance(pilotIDs, str):
             return cls.pilotAgentsDB.deletePilot(pilotIDs)
 

tests/.dirac-ci-config.yaml

@@ -39,6 +39,8 @@ config:
 }"
   PILOT_DOWNLOAD_COMMAND: "git clone --single-branch --branch master https://github.com/DIRACGrid/Pilot.git && mv Pilot/Pilot/*.py . && rm -rf Pilot"
 
+# ^^^^^^^ TO MODIFY AHHAAHAH
+
 # List of feature variables which must be passed when preparing
 required-feature-flags: []