feat: PilotManager can interact with ElasticPilotParameters

DIRACGrid · Aug 8, 2023 · 7e13f8e · 7e13f8e
1 parent 40997c0
commit 7e13f8e
Showing 1 changed file with 26 additions and 0 deletions.
diff --git a/src/DIRAC/WorkloadManagementSystem/Service/PilotManagerHandler.py b/src/DIRAC/WorkloadManagementSystem/Service/PilotManagerHandler.py
@@ -9,7 +9,9 @@
 
 from DIRAC.Core.DISET.RequestHandler import RequestHandler
 from DIRAC.Core.Utilities.ObjectLoader import ObjectLoader
+from DIRAC.ConfigurationSystem.Client.Helpers.Operations import Operations
 from DIRAC.ConfigurationSystem.Client.Helpers.Registry import getUsernameForDN, getDNForUsername
+from DIRAC.WorkloadManagementSystem.DB.ElasticPilotParametersDB import ElasticPilotParametersDB
 from DIRAC.WorkloadManagementSystem.Client import PilotStatus
 from DIRAC.WorkloadManagementSystem.Service.WMSUtilities import (
     getPilotCE,
@@ -33,6 +35,18 @@ def initializeHandler(cls, serviceInfoDict):
         except RuntimeError as excp:
             return S_ERROR(f"Can't connect to DB: {excp}")
 
+        cls.elasticPilotParametersDB = None
+        if Operations().getValue("/Services/JobMonitoring/useESForPilotParametersFlag", False):
+            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()
 
     ##############################################################################
@@ -483,3 +497,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()