From 0d8b501c440372f80d84779528b1f8cc22ef1907 Mon Sep 17 00:00:00 2001 From: Federico Stagni Date: Thu, 18 Apr 2024 15:30:51 +0200 Subject: [PATCH] feat: PilotManager can interact with ElasticPilotParameters --- .../Service/PilotManagerHandler.py | 23 +++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/src/DIRAC/WorkloadManagementSystem/Service/PilotManagerHandler.py b/src/DIRAC/WorkloadManagementSystem/Service/PilotManagerHandler.py index de322394c81..edbd46bfb4d 100644 --- a/src/DIRAC/WorkloadManagementSystem/Service/PilotManagerHandler.py +++ b/src/DIRAC/WorkloadManagementSystem/Service/PilotManagerHandler.py @@ -36,6 +36,17 @@ def initializeHandler(cls, serviceInfoDict): defaultOption, defaultClass = "DownloadPlugin", "FileCacheDownloadPlugin" cls.configValue = getServiceOption(serviceInfoDict, defaultOption, defaultClass) cls.loggingPlugin = None + cls.elasticPilotParametersDB = None + try: + result = ObjectLoader().loadObject( + "WorkloadManagementSystem.DB.ElasticPilotParametersDB", "ElasticPilotParametersDB" + ) + if not result["OK"]: + return result + cls.elasticPilotParametersDB = result["Value"]() + except RuntimeError as excp: + return S_ERROR(f"Can't connect to DB: {excp}") + return S_OK() ############################################################################## @@ -479,3 +490,15 @@ def export_deletePilots(cls, pilotIDs): @classmethod def export_clearPilots(cls, interval=30, aborted_interval=7): return cls.pilotAgentsDB.clearPilots(interval, aborted_interval) + + #### ElasticPilotParameters + + types_setPilotParameters = [int, str, str] + + @classmethod + def export_setPilotParameter(cls, pilotID, key, value): + """Set Pilot parameters""" + if cls.elasticPilotParametersDB: + return cls.elasticPilotParametersDB.setPilotParameter(pilotID, key, value) + + return S_OK()