Merge pull request #6918 from fstagni/80_fixSingleton

[8.0] make singletons of VirtualMachineDB and PilotAgentsDB

Author: fstagni

src/DIRAC/Core/Utilities/ServerUtils.py

@@ -10,15 +10,14 @@
 from DIRAC.ConfigurationSystem.Client.Helpers import CSGlobals, Registry, Resources
 from DIRAC.WorkloadManagementSystem.Client.MatcherClient import MatcherClient
 from DIRAC.Core.Utilities.List import fromChar
-from DIRAC.WorkloadManagementSystem.Client.ServerUtils import pilotAgentsDB
+from DIRAC.WorkloadManagementSystem.Client.ServerUtils import getPilotAgentsDB, getVirtualMachineDB
 from DIRAC.ResourceStatusSystem.Client.SiteStatus import SiteStatus
 from DIRAC.Resources.Cloud.EndpointFactory import EndpointFactory
 from DIRAC.ConfigurationSystem.Client.Helpers.Resources import (
     findGenericCloudCredentials,
     getVMTypes,
     getPilotBootstrapParameters,
 )
-from DIRAC.WorkloadManagementSystem.Client.ServerUtils import virtualMachineDB
 from DIRAC.WorkloadManagementSystem.Utilities.Utils import getProxyFileForCloud
 
 
@@ -54,6 +53,8 @@ def __init__(self, *args, **kwargs):
 
     def initialize(self):
         self.siteClient = SiteStatus()
+        self.pilotAgentsDB = getPilotAgentsDB()
+        self.virtualMachineDB = getVirtualMachineDB()
         return S_OK()
 
     def beginExecution(self):
@@ -324,7 +325,7 @@ def createVMs(self):
 
         tqIDList = list(result["Value"].keys())
 
-        result = virtualMachineDB.getInstanceCounters("Status", {})
+        result = self.virtualMachineDB.getInstanceCounters("Status", {})
         totalVMs = 0
         if result["OK"]:
             for status in result["Value"]:
@@ -410,7 +411,7 @@ def createVMs(self):
 
             # Get the number of already instantiated VMs for these task queues
             totalWaitingVMs = 0
-            result = virtualMachineDB.getInstanceCounters("Status", {"Endpoint": endpoint})
+            result = self.virtualMachineDB.getInstanceCounters("Status", {"Endpoint": endpoint})
             if result["OK"]:
                 for status in result["Value"]:
                     if status in ["New", "Submitted"]:
@@ -465,7 +466,7 @@ def createVMs(self):
             for uuID in vmDict:
                 diracUUID = vmDict[uuID]["InstanceID"]
                 endpoint = f"{self.vmTypeDict[vmType]['Site']}::{ceName}"
-                result = virtualMachineDB.insertInstance(uuID, vmTypeName, diracUUID, endpoint, self.vo)
+                result = self.virtualMachineDB.insertInstance(uuID, vmTypeName, diracUUID, endpoint, self.vo)
                 if not result["OK"]:
                     continue
                 pRef = "vm://" + ceName + "/" + diracUUID + ":00"
@@ -489,7 +490,9 @@ def createVMs(self):
                 tqDict[tqID].append(pilotID)
 
             for tqID, pilotList in tqDict.items():
-                result = pilotAgentsDB.addPilotTQReference(pilotList, tqID, "", "", self.localhost, "Cloud", stampDict)
+                result = self.pilotAgentsDB.addPilotTQReference(
+                    pilotList, tqID, "", "", self.localhost, "Cloud", stampDict
+                )
                 if not result["OK"]:
                     self.log.error(f"Failed to insert pilots into the PilotAgentsDB: {result['Message']}")
 
@@ -499,7 +502,7 @@ def createVMs(self):
         return S_OK()
 
     def getVMInstances(self, endpoint, maxInstances):
-        result = virtualMachineDB.getInstanceCounters("Status", {"Endpoint": endpoint})
+        result = self.virtualMachineDB.getInstanceCounters("Status", {"Endpoint": endpoint})
         if not result["OK"]:
             return result
 

src/DIRAC/WorkloadManagementSystem/Agent/CloudDirector.py

@@ -33,7 +33,7 @@
 from DIRAC.ResourceStatusSystem.Client.SiteStatus import SiteStatus
 from DIRAC.WorkloadManagementSystem.Client import PilotStatus
 from DIRAC.WorkloadManagementSystem.Client.MatcherClient import MatcherClient
-from DIRAC.WorkloadManagementSystem.Client.ServerUtils import pilotAgentsDB
+from DIRAC.WorkloadManagementSystem.Client.ServerUtils import getPilotAgentsDB
 from DIRAC.WorkloadManagementSystem.private.ConfigHelper import findGenericPilotCredentials
 from DIRAC.WorkloadManagementSystem.Service.WMSUtilities import getGridEnv
 from DIRAC.WorkloadManagementSystem.Utilities.PilotWrapper import (
@@ -93,6 +93,7 @@ def __init__(self, *args, **kwargs):
         self.sendSubmissionMonitoring = False
         self.siteClient = None
         self.rssClient = None
+        self.pilotAgentsDB = None
         self.rssFlag = None
 
         self.globalParameters = {"NumberOfProcessors": 1, "MaxRAM": 2048}
@@ -148,6 +149,7 @@ def initialize(self):
         self.siteClient = SiteStatus()
         self.rssClient = ResourceStatus()
         self.matcherClient = MatcherClient()
+        self.pilotAgentsDB = getPilotAgentsDB()
 
         return S_OK()
 
@@ -390,7 +392,7 @@ def submitPilots(self):
             manyWaitingPilotsFlag = False
             if self.pilotWaitingFlag:
                 tqIDList = list(additionalInfo)
-                result = pilotAgentsDB.countPilots(
+                result = self.pilotAgentsDB.countPilots(
                     {"TaskQueueID": tqIDList, "Status": PilotStatus.PILOT_WAITING_STATES}, None
                 )
                 if not result["OK"]:
@@ -533,7 +535,9 @@ def _ifAndWhereToSubmit(self):
     def monitorJobsQueuesPilots(self, matchingTQs):
         """Just printout of jobs queues and pilots status in TQ"""
         tqIDList = list(matchingTQs)
-        result = pilotAgentsDB.countPilots({"TaskQueueID": tqIDList, "Status": PilotStatus.PILOT_WAITING_STATES}, None)
+        result = self.pilotAgentsDB.countPilots(
+            {"TaskQueueID": tqIDList, "Status": PilotStatus.PILOT_WAITING_STATES}, None
+        )
 
         totalWaitingJobs = 0
         for tqDescription in matchingTQs.values():
@@ -834,7 +838,7 @@ def _addPilotTQReference(self, queue, taskQueueDict, pilotList, stampDict):
             tqDict[tqID].append(pilotID)
 
         for tqID, pilotsList in tqDict.items():
-            result = pilotAgentsDB.addPilotTQReference(
+            result = self.pilotAgentsDB.addPilotTQReference(
                 pilotsList,
                 tqID,
                 self.pilotDN,
@@ -847,7 +851,7 @@ def _addPilotTQReference(self, queue, taskQueueDict, pilotList, stampDict):
                 self.log.error("Failed add pilots to the PilotAgentsDB", result["Message"])
                 continue
             for pilot in pilotsList:
-                result = pilotAgentsDB.setPilotStatus(
+                result = self.pilotAgentsDB.setPilotStatus(
                     pilot,
                     PilotStatus.SUBMITTED,
                     self.queueDict[queue]["CEName"],
@@ -871,7 +875,7 @@ def getQueueSlots(self, queue, manyWaitingPilotsFlag):
 
         # See if there are waiting pilots for this queue. If not, allow submission
         if totalSlots and manyWaitingPilotsFlag:
-            result = pilotAgentsDB.selectPilots(
+            result = self.pilotAgentsDB.selectPilots(
                 {"DestinationSite": ceName, "Queue": queueName, "Status": PilotStatus.PILOT_WAITING_STATES}
             )
             if result["OK"]:
@@ -886,7 +890,7 @@ def getQueueSlots(self, queue, manyWaitingPilotsFlag):
             if availableSlotsCount % self.availableSlotsUpdateCycleFactor == 0:
                 # Get the list of already existing pilots for this queue
                 jobIDList = None
-                result = pilotAgentsDB.selectPilots(
+                result = self.pilotAgentsDB.selectPilots(
                     {"DestinationSite": ceName, "Queue": queueName, "Status": PilotStatus.PILOT_TRANSIENT_STATES}
                 )
 
@@ -916,7 +920,7 @@ def getQueueSlots(self, queue, manyWaitingPilotsFlag):
                     waitingJobs = 0
                     totalJobs = 0
                     if jobIDList:
-                        result = pilotAgentsDB.getPilotInfo(jobIDList)
+                        result = self.pilotAgentsDB.getPilotInfo(jobIDList)
                         if not result["OK"]:
                             self.log.warn("Failed to check PilotAgentsDB", f"for queue {queue}: \n{result['Message']}")
                             self.failedQueues[queue] += 1
@@ -1127,7 +1131,7 @@ def updatePilotStatus(self):
             queueName = self.queueDict[queue]["QueueName"]
             ceType = self.queueDict[queue]["CEType"]
             siteName = self.queueDict[queue]["Site"]
-            result = pilotAgentsDB.selectPilots(
+            result = self.pilotAgentsDB.selectPilots(
                 {
                     "DestinationSite": ceName,
                     "Queue": queueName,
@@ -1144,7 +1148,7 @@ def updatePilotStatus(self):
             pilotRefs = result["Value"]
             if not pilotRefs:
                 continue
-            result = pilotAgentsDB.getPilotInfo(pilotRefs)
+            result = self.pilotAgentsDB.getPilotInfo(pilotRefs)
             if not result["OK"]:
                 self.log.error("Failed to get pilots info from DB", result["Message"])
                 continue
@@ -1155,7 +1159,7 @@ def updatePilotStatus(self):
 
             # Check if the accounting is to be sent
             if self.sendAccounting:
-                result = pilotAgentsDB.selectPilots(
+                result = self.pilotAgentsDB.selectPilots(
                     {
                         "DestinationSite": ceName,
                         "Queue": queueName,
@@ -1172,7 +1176,7 @@ def updatePilotStatus(self):
                 pilotRefs = result["Value"]
                 if not pilotRefs:
                     continue
-                result = pilotAgentsDB.getPilotInfo(pilotRefs)
+                result = self.pilotAgentsDB.getPilotInfo(pilotRefs)
                 if not result["OK"]:
                     self.log.error("Failed to get pilots info from DB", result["Message"])
                     continue
@@ -1195,7 +1199,7 @@ def _updatePilotStatusPerQueue(self, queue, proxy):
         ceType = self.queueDict[queue]["CEType"]
         siteName = self.queueDict[queue]["Site"]
 
-        result = pilotAgentsDB.selectPilots(
+        result = self.pilotAgentsDB.selectPilots(
             {
                 "DestinationSite": ceName,
                 "Queue": queueName,
@@ -1213,7 +1217,7 @@ def _updatePilotStatusPerQueue(self, queue, proxy):
         if not pilotRefs:
             return
 
-        result = pilotAgentsDB.getPilotInfo(pilotRefs)
+        result = self.pilotAgentsDB.getPilotInfo(pilotRefs)
         if not result["OK"]:
             self.log.error("Failed to get pilots info from DB", result["Message"])
             return
@@ -1289,7 +1293,7 @@ def _updatePilotStatus(self, pilotRefs, pilotDict, pilotCEDict):
 
             if newStatus:
                 self.log.info("Updating status", f"to {newStatus} for pilot {pRef}")
-                result = pilotAgentsDB.setPilotStatus(pRef, newStatus, "", "Updated by SiteDirector")
+                result = self.pilotAgentsDB.setPilotStatus(pRef, newStatus, "", "Updated by SiteDirector")
                 if not result["OK"]:
                     self.log.error(result["Message"])
                 if newStatus == "Aborted":
@@ -1328,7 +1332,7 @@ def _getPilotOutput(self, pRef, pilotDict, ce, ceName):
             output, error = result["Value"]
 
         if output:
-            result = pilotAgentsDB.storePilotOutput(pRef, output, error)
+            result = self.pilotAgentsDB.storePilotOutput(pRef, output, error)
             if not result["OK"]:
                 self.log.error("Failed to store pilot output", result["Message"])
         else:
@@ -1368,7 +1372,7 @@ def sendPilotAccounting(self, pilotDict):
                 self.log.error("Failed to send accounting info for pilot ", pRef)
             else:
                 # Set up AccountingSent flag
-                result = pilotAgentsDB.setAccountingFlag(pRef)
+                result = self.pilotAgentsDB.setAccountingFlag(pRef)
                 if not result["OK"]:
                     self.log.error("Failed to set accounting flag for pilot ", pRef)
 
@@ -1377,7 +1381,7 @@ def sendPilotAccounting(self, pilotDict):
         if result["OK"]:
             for pRef in pilotDict:
                 self.log.verbose("Setting AccountingSent flag", f"for pilot {pRef}")
-                result = pilotAgentsDB.setAccountingFlag(pRef)
+                result = self.pilotAgentsDB.setAccountingFlag(pRef)
                 if not result["OK"]:
                     self.log.error("Failed to set accounting flag for pilot ", pRef)
         else:

src/DIRAC/WorkloadManagementSystem/Agent/SiteDirector.py

@@ -10,7 +10,7 @@
 from DIRAC.ConfigurationSystem.Client.Helpers import Operations
 from DIRAC.Core.Base.AgentModule import AgentModule
 from DIRAC.Core.Utilities import List, Network
-from DIRAC.WorkloadManagementSystem.Client.ServerUtils import virtualMachineDB
+from DIRAC.WorkloadManagementSystem.Client.ServerUtils import getVirtualMachineDB
 
 
 class VirtualMachineMonitorAgent(AgentModule):
@@ -61,7 +61,7 @@ def __declareInstanceRunning(self):
         retries = 3
         sleepTime = 30
         for i in range(retries):
-            result = virtualMachineDB.declareInstanceRunning(self.uniqueID, self.ipAddress)
+            result = self.virtualMachineDB.declareInstanceRunning(self.uniqueID, self.ipAddress)
             if result["OK"]:
                 self.log.info("Declared instance running")
                 return result
@@ -87,6 +87,7 @@ def initialize(self):
         self.heartBeatPeriod = None
         self.am_setOption("MaxCycles", 0)
         self.am_setOption("PollingTime", 60)
+        self.virtualMachineDB = getVirtualMachineDB()
 
         # Discover net address
         self.ipAddress = None
@@ -100,7 +101,7 @@ def initialize(self):
 
         # Declare instance running
         self.uniqueID = ""
-        result = virtualMachineDB.getUniqueIDByName(self.vmID)
+        result = self.virtualMachineDB.getUniqueIDByName(self.vmID)
         if result["OK"]:
             self.uniqueID = result["Value"]
         result = self.__declareInstanceRunning()
@@ -169,7 +170,7 @@ def execute(self):
         if uptime % self.heartBeatPeriod <= self.am_getPollingTime():
             # Heartbeat time!
             self.log.info("Sending hearbeat...")
-            result = virtualMachineDB.instanceIDHeartBeat(
+            result = self.virtualMachineDB.instanceIDHeartBeat(
                 self.uniqueID,
                 avgLoad,
                 numJobs,
@@ -223,7 +224,7 @@ def __haltInstance(self, avgLoad=0.0):
         retries = 3
         sleepTime = 10
         for i in range(retries):
-            result = virtualMachineDB.declareInstanceHalting(self.uniqueID, avgLoad)
+            result = self.virtualMachineDB.declareInstanceHalting(self.uniqueID, avgLoad)
             if result["OK"]:
                 self.log.info("Declared instance halting")
                 break

src/DIRAC/WorkloadManagementSystem/Agent/VirtualMachineMonitorAgent.py

@@ -55,6 +55,7 @@ def sd(mocker):
     sd.log.setLevel("DEBUG")
     sd.rpcMatcher = MagicMock()
     sd.rssClient = MagicMock()
+    sd.pilotAgentsDB = MagicMock()
     sd.workingDirectory = ""
     sd.queueDict = {
         "aQueue": {

src/DIRAC/WorkloadManagementSystem/Agent/test/Test_Agent_SiteDirector.py

@@ -3,30 +3,24 @@
 It always try to insert the records directly. In case of failure a WMS client is used...
 """
 
-from DIRAC.Core.Utilities.ServerUtils import getDBOrClient
+from DIRAC.Core.Base.Client import Client
+from DIRAC.WorkloadManagementSystem.DB.PilotAgentsDB import gPilotAgentsDB
+from DIRAC.WorkloadManagementSystem.DB.VirtualMachineDB import gVirtualMachineDB
 
 
 def getPilotAgentsDB():
-    serverName = "WorkloadManagement/PilotManager"
-    PilotAgentsDB = None
     try:
-        from DIRAC.WorkloadManagementSystem.DB.PilotAgentsDB import PilotAgentsDB
+        if gPilotAgentsDB and gPilotAgentsDB._connected:
+            return gPilotAgentsDB
     except Exception:
         pass
-    return getDBOrClient(PilotAgentsDB, serverName)
-
-
-pilotAgentsDB = getPilotAgentsDB()
+    return Client(url="WorkloadManagement/PilotManager")
 
 
 def getVirtualMachineDB():
-    serverName = "WorkloadManagement/VirtualMachineManager"
-    VirtualMachineDB = None
     try:
-        from DIRAC.WorkloadManagementSystem.DB.VirtualMachineDB import VirtualMachineDB
+        if gVirtualMachineDB and gVirtualMachineDB._connected:
+            return gVirtualMachineDB
     except Exception:
         pass
-    return getDBOrClient(VirtualMachineDB, serverName)
-
-
-virtualMachineDB = getVirtualMachineDB()
+    return Client(url="WorkloadManagement/VirtualMachineManager")

src/DIRAC/WorkloadManagementSystem/Client/ServerUtils.py

@@ -1272,3 +1272,9 @@ def buildSQL(self, selectDict=None):
 
     def getColumnList(self):
         return self._columns
+
+
+try:
+    gPilotAgentsDB = PilotAgentsDB()
+except:
+    gPilotAgentsDB = None

src/DIRAC/WorkloadManagementSystem/DB/PilotAgentsDB.py

@@ -1165,3 +1165,9 @@ def __setError(self, element, iD, reason):
             return ret
 
         return S_ERROR(reason)
+
+
+try:
+    gVirtualMachineDB = VirtualMachineDB()
+except:
+    gVirtualMachineDB = None