From e915976e512a2f937e400cf453d54733bd0f3276 Mon Sep 17 00:00:00 2001 From: Federico Stagni Date: Fri, 10 Jul 2026 17:48:10 +0200 Subject: [PATCH 1/2] fix: allow using PilotStamp when setting the pilot status --- .../WorkloadManagementSystem/Agent/JobAgent.py | 10 +++++++++- .../WorkloadManagementSystem/DB/PilotAgentsDB.py | 15 +++++++++++---- .../Service/PilotManagerHandler.py | 12 ++++++++++-- 3 files changed, 30 insertions(+), 7 deletions(-) diff --git a/src/DIRAC/WorkloadManagementSystem/Agent/JobAgent.py b/src/DIRAC/WorkloadManagementSystem/Agent/JobAgent.py index 3860b6e3ea1..c3388a0a6ba 100755 --- a/src/DIRAC/WorkloadManagementSystem/Agent/JobAgent.py +++ b/src/DIRAC/WorkloadManagementSystem/Agent/JobAgent.py @@ -62,6 +62,7 @@ def __init__(self, agentName, loadName, baseAgentName=False, properties=None): # Localsite options self.siteName = "Unknown" self.pilotReference = "Unknown" + self.pilotStamp = self.pilotReference self.defaultProxyLength = 86400 * 5 # Agent options @@ -116,6 +117,7 @@ def initialize(self): # Localsite options self.siteName = siteName() self.pilotReference = gConfig.getValue("/LocalSite/PilotReference", self.pilotReference) + self.pilotStamp = os.environ.get("DIRAC_PILOT_STAMP", self.pilotReference) self.defaultProxyLength = gConfig.getValue("/Registry/DefaultProxyLifeTime", self.defaultProxyLength) # Agent options # This is the factor to convert raw CPU to Normalized units (based on the CPU Model) @@ -895,7 +897,13 @@ def finalize(self): gridCE = gConfig.getValue("/LocalSite/GridCE", "") queue = gConfig.getValue("/LocalSite/CEQueue", "") result = PilotManagerClient().setPilotStatus( - str(self.pilotReference), PilotStatus.DONE, gridCE, "Report from JobAgent", self.siteName, queue + str(self.pilotReference), + PilotStatus.DONE, + gridCE, + "Report from JobAgent", + self.siteName, + queue, + self.pilotStamp, ) if not result["OK"]: self.log.warn("Issue setting the pilot status", result["Message"]) diff --git a/src/DIRAC/WorkloadManagementSystem/DB/PilotAgentsDB.py b/src/DIRAC/WorkloadManagementSystem/DB/PilotAgentsDB.py index 1097613241b..e72c3b46c2e 100755 --- a/src/DIRAC/WorkloadManagementSystem/DB/PilotAgentsDB.py +++ b/src/DIRAC/WorkloadManagementSystem/DB/PilotAgentsDB.py @@ -93,6 +93,7 @@ def setPilotStatus( benchmark=None, currentJob=None, updateTime=None, + pilotStamp=None, conn=False, ): """Set pilot job status""" @@ -123,17 +124,23 @@ def setPilotStatus( setList.append("CurrentJobID=%s") args.append(str(int(currentJob))) if destination: - setList.append(f"DestinationSite=%s") + setList.append("DestinationSite=%s") args.append(destination) if not gridSite: res = getCESiteMapping(destination) if res["OK"] and res["Value"]: - setList.append(f"GridSite=%s") + setList.append("GridSite=%s") args.append(res["Value"][destination]) set_string = ",".join(setList) - req = f"UPDATE PilotAgents SET {set_string} WHERE PilotJobReference=%s" # nosec - args.append(pilotRef) + req = f"UPDATE PilotAgents SET {set_string}" # nosec + + if pilotStamp: + req += " WHERE PilotStamp=%s" + args.append(pilotStamp) + else: + req += " WHERE PilotJobReference=%s" + args.append(pilotRef) return self._update(req, args=args, conn=conn) def selectPilots( diff --git a/src/DIRAC/WorkloadManagementSystem/Service/PilotManagerHandler.py b/src/DIRAC/WorkloadManagementSystem/Service/PilotManagerHandler.py index 1ad2c61a24d..cce439de2c5 100644 --- a/src/DIRAC/WorkloadManagementSystem/Service/PilotManagerHandler.py +++ b/src/DIRAC/WorkloadManagementSystem/Service/PilotManagerHandler.py @@ -73,9 +73,17 @@ def export_getPilots(cls, jobID): types_setPilotStatus = [str, str] @classmethod - def export_setPilotStatus(cls, pilotRef, status, destination=None, reason=None, gridSite=None, queue=None): + def export_setPilotStatus( + cls, pilotRef, status, destination=None, reason=None, gridSite=None, queue=None, pilotStamp=None + ): """Set the pilot agent status""" return cls.pilotAgentsDB.setPilotStatus( - pilotRef, status, destination=destination, statusReason=reason, gridSite=gridSite, queue=queue + pilotRef, + status, + destination=destination, + statusReason=reason, + gridSite=gridSite, + queue=queue, + pilotStamp=pilotStamp, ) From 1f9f94cca092912ca547c9f6f7e156ea702ce5b1 Mon Sep 17 00:00:00 2001 From: Federico Stagni Date: Tue, 14 Jul 2026 15:20:58 +0200 Subject: [PATCH 2/2] feat: added PilotManager legacy adaptor --- .../Client/PilotManagerClient.py | 8 +- .../FutureClient/PilotManagerClient.py | 137 ++++++++++++++++++ 2 files changed, 143 insertions(+), 2 deletions(-) create mode 100644 src/DIRAC/WorkloadManagementSystem/FutureClient/PilotManagerClient.py diff --git a/src/DIRAC/WorkloadManagementSystem/Client/PilotManagerClient.py b/src/DIRAC/WorkloadManagementSystem/Client/PilotManagerClient.py index 5cd731f2bea..800d9575892 100644 --- a/src/DIRAC/WorkloadManagementSystem/Client/PilotManagerClient.py +++ b/src/DIRAC/WorkloadManagementSystem/Client/PilotManagerClient.py @@ -1,13 +1,17 @@ -""" Module that contains client access to the Pilots handler. -""" +"""Module that contains client access to the Pilots handler.""" 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 diff --git a/src/DIRAC/WorkloadManagementSystem/FutureClient/PilotManagerClient.py b/src/DIRAC/WorkloadManagementSystem/FutureClient/PilotManagerClient.py new file mode 100644 index 00000000000..c6d1bed5ece --- /dev/null +++ b/src/DIRAC/WorkloadManagementSystem/FutureClient/PilotManagerClient.py @@ -0,0 +1,137 @@ +from DIRAC.Core.Security.DiracX import DiracXClient, FutureClient +from DIRAC.Core.Utilities.ReturnValues import convertToReturnValue + + +class PilotManagerClient(FutureClient): + @convertToReturnValue + def addPilotReferences( + self, + pilot_reference: str, + VO: str, + grid_type: str = "DIRAC", + pilot_stamp_dict: dict = {}, + grid_site: str = "Unknown", + destination_site: str = "NotAssigned", + ): + """Add a new pilot to the database. + Uses pilots/management/register_pilot + """ + + with DiracXClient() as api: + api.pilots.register_pilot( + pilot_stamp=pilot_stamp_dict[pilot_reference], + vo=VO, + grid_type=grid_type, + grid_site=grid_site, + destination_site=destination_site, + ) + + @convertToReturnValue + def getPilotInfo(self, pilot_reference: str): + """Get the info about a given pilot job reference""" + + 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 selectPilots(self, conditions_dict: dict): + """Select pilots given the selection conditions""" + + with DiracXClient() as api: + search = [{}] # FIXME + return api.pilots.search(parameters=[], search=search) + + @convertToReturnValue + def getPilotSummary(self, start_date: str, end_date: str): + """Get summary of the status of the Pilot Jobs""" + + with DiracXClient() as api: + search_filters = [] + if start_date: + search_filters.append({"parameter": "SubmissionTime", "operator": "gt", "value": start_date}) + if end_date: + search_filters.append({"parameter": "SubmissionTime", "operator": "lt", "value": end_date}) + + rows = api.pilots.summary(grouping=["DestinationSite", "Status"], search=search_filters) + + # Build nested result: { site: { status: count }, Total: { status: total_count } } + summary_dict = {"Total": {}} + for row in rows: + site = row["DestinationSite"] + status = row["Status"] + count = row["count"] + + if site not in summary_dict: + summary_dict[site] = {} + + summary_dict[site][status] = count + summary_dict["Total"].setdefault(status, 0) + summary_dict["Total"][status] += count + + return summary_dict + + @convertToReturnValue + def getPilots(self, job_id: str | int): + """Get pilots executing/having executed a Job""" + + with DiracXClient() as api: + pilot_ids = api.pilots.search(job_id=job_id) + search = [{"parameter": "PilotID", "operator": "in", "value": pilot_ids}] + return api.pilots.search(parameters=[], search=search, sort=[]) # type: ignore + + @convertToReturnValue + def setPilotStatus( + self, + pilot_reference: str, + status: str, + destination: str | None = None, + reason: str | None = None, + grid_site: str | None = None, + queue: str | None = None, + pilot_stamp: str | None = None, + ): + """Set the pilot status""" + + with DiracXClient() as api: + values_dict = ( + { + "PilotStamp": pilot_stamp, + "Status": status, + "DestinationSite": destination, + "StatusReason": reason, + "GridSite": grid_site, + "Queue": queue, + }, + ) + + return api.pilots.update_pilot_metadata(pilot_stamps_to_fields_mapping=[values_dict])