From efa723a2edd9cffa5aacfa0ec2a5f572589327a3 Mon Sep 17 00:00:00 2001 From: Federico Stagni Date: Tue, 23 Apr 2024 17:10:00 +0200 Subject: [PATCH] feat: PilotManager can interact with ElasticPilotParameters --- .../DB/ElasticPilotParametersDB.py | 7 +++--- .../Service/PilotManagerHandler.py | 23 +++++++++++++++++++ .../Test_ElasticPilotParametersDB.py | 3 --- 3 files changed, 26 insertions(+), 7 deletions(-) diff --git a/src/DIRAC/WorkloadManagementSystem/DB/ElasticPilotParametersDB.py b/src/DIRAC/WorkloadManagementSystem/DB/ElasticPilotParametersDB.py index d7dd95724e9..b8551e125cc 100644 --- a/src/DIRAC/WorkloadManagementSystem/DB/ElasticPilotParametersDB.py +++ b/src/DIRAC/WorkloadManagementSystem/DB/ElasticPilotParametersDB.py @@ -5,12 +5,11 @@ - setPilotParameter() - deletePilotParameters() """ -from DIRAC import S_OK, S_ERROR, gConfig -from DIRAC.Core.Utilities import TimeUtilities -from DIRAC.ConfigurationSystem.Client.PathFinder import getDatabaseSection +from DIRAC import S_ERROR, S_OK, gConfig from DIRAC.ConfigurationSystem.Client.Helpers import CSGlobals +from DIRAC.ConfigurationSystem.Client.PathFinder import getDatabaseSection from DIRAC.Core.Base.ElasticDB import ElasticDB - +from DIRAC.Core.Utilities import TimeUtilities mapping = { "properties": { 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() diff --git a/tests/Integration/WorkloadManagementSystem/Test_ElasticPilotParametersDB.py b/tests/Integration/WorkloadManagementSystem/Test_ElasticPilotParametersDB.py index 2cb6a1bd1d3..5d703811a4e 100644 --- a/tests/Integration/WorkloadManagementSystem/Test_ElasticPilotParametersDB.py +++ b/tests/Integration/WorkloadManagementSystem/Test_ElasticPilotParametersDB.py @@ -150,9 +150,6 @@ def test_setAndGetPilotFromDB(): assert len(res["Value"][1010000]) == 0 # delete the indexes - res = elasticPilotParametersDB.deleteIndex(elasticPilotParametersDB.indexName_base) - assert res["OK"] - assert res["Value"] == "Nothing to delete" res = elasticPilotParametersDB.deleteIndex(elasticPilotParametersDB._indexName(100)) assert res["OK"] res = elasticPilotParametersDB.deleteIndex(elasticPilotParametersDB._indexName(1010000))