Repository: dmwm/wmcore
Issue: DrainStatusPoller KeyError crash in 1.1.18

I've just seen this reported in wmstats for vocms0282. Traceback is in [1]

[1]
```
2019-04-06 20:00:16,208:140226378852096:INFO:DrainStatusPoller:Running agent drain algorithm...
2019-04-06 20:00:16,461:140226387244800:WARNING:ResourceControlUpdater:Site T2_TH_CUNSTDA has invalid IO thresholds in SSB. Taking no action
2019-04-06 20:00:16,462:140226387244800:WARNING:ResourceControlUpdater:Site: T3_US_Omaha has incomplete SSB metrics, see {'slotsCPU': 2001}
2019-04-06 20:00:16,462:140226387244800:WARNING:ResourceControlUpdater:Site: T2_TH_CUNSTDA has incomplete SSB metrics, see {'slotsCPU': 67}
2019-04-06 20:00:16,519:140226387244800:ERROR:GenericDataCache:Traceback (most recent call last):
  File "/data/srv/wmagent/v1.1.20.patch3/sw/slc7_amd64_gcc630/cms/wmagent/1.1.20.patch3/lib/python2.7/site-packages/WMCore/Cache/GenericDataCache.py", line 37, in getData
    self.data = self.func(**self.kwargs)
  File "/data/srv/wmagent/v1.1.20.patch3/sw/slc7_amd64_gcc630/cms/wmagent/1.1.20.patch3/lib/python2.7/site-packages/WMCore/Services/ReqMgrAux/ReqMgrAux.py", line 49, in _getResult
    verb=verb, contentType=contentType)
  File "/data/srv/wmagent/v1.1.20.patch3/sw/slc7_amd64_gcc630/cms/wmagent/1.1.20.patch3/lib/python2.7/site-packages/WMCore/Services/Service.py", line 211, in refreshCache
    return open(cachefile, 'r')
IOError: [Errno 2] No such file or directory: '/data/srv/wmagent/v1.1.20.patch3/install/wmagent/AgentStatusWatcher/.wmcore_cache/.wmcore_cache_31961/requests/cmsweb.cern.ch:8443/15744094587142094_GET_wmagentconfig_vocms0282.cern.ch'

2019-04-06 20:00:16,520:140226378852096:ERROR:GenericDataCache:Traceback (most recent call last):
  File "/data/srv/wmagent/v1.1.20.patch3/sw/slc7_amd64_gcc630/cms/wmagent/1.1.20.patch3/lib/python2.7/site-packages/WMCore/Cache/GenericDataCache.py", line 37, in getData
    self.data = self.func(**self.kwargs)
  File "/data/srv/wmagent/v1.1.20.patch3/sw/slc7_amd64_gcc630/cms/wmagent/1.1.20.patch3/lib/python2.7/site-packages/WMCore/Services/ReqMgrAux/ReqMgrAux.py", line 49, in _getResult
    verb=verb, contentType=contentType)
  File "/data/srv/wmagent/v1.1.20.patch3/sw/slc7_amd64_gcc630/cms/wmagent/1.1.20.patch3/lib/python2.7/site-packages/WMCore/Services/Service.py", line 211, in refreshCache
    return open(cachefile, 'r')
IOError: [Errno 2] No such file or directory: '/data/srv/wmagent/v1.1.20.patch3/install/wmagent/AgentStatusWatcher/.wmcore_cache/.wmcore_cache_31961/requests/cmsweb.cern.ch:8443/15744094587142094_GET_wmagentconfig_vocms0282.cern.ch'

2019-04-06 20:00:16,520:140226387244800:WARNING:ReqMgrAux:agent config is not correct: {}
2019-04-06 20:00:16,520:140226378852096:WARNING:ReqMgrAux:agent config is not correct: {}
2019-04-06 20:00:16,520:140226387244800:ERROR:ResourceControlUpdater:Error occurred, will retry later:
2019-04-06 20:00:16,521:140226387244800:ERROR:ResourceControlUpdater:'UserDrainMode'
2019-04-06 20:00:16,524:140226387244800:ERROR:ResourceControlUpdater:Trace back: 
Traceback (most recent call last):
  File "/data/srv/wmagent/v1.1.20.patch3/sw/slc7_amd64_gcc630/cms/wmagent/1.1.20.patch3/lib/python2.7/site-packages/WMComponent/AgentStatusWatcher/ResourceControlUpdater.py", line 104, in algorithm
    self.getAgentsByTeam()
  File "/data/srv/wmagent/v1.1.20.patch3/sw/slc7_amd64_gcc630/cms/wmagent/1.1.20.patch3/lib/python2.7/site-packages/WMComponent/AgentStatusWatcher/ResourceControlUpdater.py", line 120, in getAgentsByTeam
    if isDrainMode(self.config):
  File "/data/srv/wmagent/v1.1.20.patch3/sw/slc7_amd64_gcc630/cms/wmagent/1.1.20.patch3/lib/python2.7/site-packages/WMCore/Services/ReqMgrAux/ReqMgrAux.py", line 199, in isDrainMode
    return AUXDB_AGENT_CONFIG_CACHE["UserDrainMode"] or AUXDB_AGENT_CONFIG_CACHE["AgentDrainMode"]
KeyError: 'UserDrainMode'

2019-04-06 20:00:16,524:140226387244800:INFO:ResourceControlUpdater:Resource control cycle finished updating site state and thresholds.
2019-04-06 20:00:16,524:140226387244800:INFO:BaseWorkerThread:ResourceControlUpdater took 0.414 secs to execute
2019-04-06 20:00:16,545:140226378852096:ERROR:BaseWorkerThread:Error in worker algorithm (1):
Backtrace:
  <WMComponent.AgentStatusWatcher.DrainStatusPoller.DrainStatusPoller object at 0x7f89041435d0> 'UserDrainMode'  File "/data/srv/wmagent/v1.1.20.patch3/sw/slc7_amd64_gcc630/cms/wmagent/1.1.20.patch3/lib/python2.7/site-packages/WMCore/WorkerThreads/BaseWorkerThread.py", line 182, in __call__
    tSpent, results, _ = algorithmWithDBExceptionHandler(parameters)
  File "/data/srv/wmagent/v1.1.20.patch3/sw/slc7_amd64_gcc630/cms/wmagent/1.1.20.patch3/lib/python2.7/site-packages/WMCore/Database/DBExceptionHandler.py", line 35, in wrapper
    return f(*args, **kwargs)
  File "/data/srv/wmagent/v1.1.20.patch3/sw/slc7_amd64_gcc630/cms/wmagent/1.1.20.patch3/lib/python2.7/site-packages/Utils/Timers.py", line 24, in wrapper
    res = func(*arg, **kw)
  File "/data/srv/wmagent/v1.1.20.patch3/sw/slc7_amd64_gcc630/cms/wmagent/1.1.20.patch3/lib/python2.7/site-packages/WMComponent/AgentStatusWatcher/DrainStatusPoller.py", line 48, in algorithm
    if isDrainMode(self.config):
  File "/data/srv/wmagent/v1.1.20.patch3/sw/slc7_amd64_gcc630/cms/wmagent/1.1.20.patch3/lib/python2.7/site-packages/WMCore/Services/ReqMgrAux/ReqMgrAux.py", line 199, in isDrainMode
    return AUXDB_AGENT_CONFIG_CACHE["UserDrainMode"] or AUXDB_AGENT_CONFIG_CACHE["AgentDrainMode"]

2019-04-06 20:00:16,545:140226378852096:INFO:Harness:>>>Terminating worker threads
2019-04-06 20:00:16,545:140226378852096:ERROR:BaseWorkerThread:Error in event loop (2): <WMComponent.AgentStatusWatcher.DrainStatusPoller.DrainStatusPoller object at 0x7f89041435d0> 'UserDrainMode'
Backtrace:
  File "/data/srv/wmagent/v1.1.20.patch3/sw/slc7_amd64_gcc630/cms/wmagent/1.1.20.patch3/lib/python2.7/site-packages/WMCore/WorkerThreads/BaseWorkerThread.py", line 208, in __call__
    raise ex

2019-04-06 20:00:16,607:140226378852096:INFO:BaseWorkerThread:Worker thread <WMComponent.AgentStatusWatcher.DrainStatusPoller.DrainStatusPoller object at 0x7f89041435d0> terminated
```

Changed files:
src/python/WMComponent/AgentStatusWatcher/DrainStatusPoller.py

diff --git a/src/python/WMComponent/AgentStatusWatcher/DrainStatusPoller.py b/src/python/WMComponent/AgentStatusWatcher/DrainStatusPoller.py
index 2697de253..4de5fe7ec 100644
--- a/src/python/WMComponent/AgentStatusWatcher/DrainStatusPoller.py
+++ b/src/python/WMComponent/AgentStatusWatcher/DrainStatusPoller.py
@@ -10,10 +10,11 @@ __all__ = []
 import logging
 from Utils.Timers import timeFunction
 from WMCore.WorkerThreads.BaseWorkerThread import BaseWorkerThread
-from WMCore.Services.ReqMgrAux.ReqMgrAux import isDrainMode, ReqMgrAux
+from WMCore.Services.ReqMgrAux.ReqMgrAux import ReqMgrAux
 from WMComponent.AgentStatusWatcher.DrainStatusAPI import DrainStatusAPI
 from WMCore.Services.PyCondor.PyCondorAPI import PyCondorAPI
 
+
 class DrainStatusPoller(BaseWorkerThread):
     """
     Collects information related to the agent drain status
@@ -45,25 +46,25 @@ class DrainStatusPoller(BaseWorkerThread):
             logging.error("Failed to fetch agent configuration from the auxiliary DB")
             return
 
-        if isDrainMode(self.config):
-            # check to see if the agent hit any speed drain thresholds
-            thresholdsHit = self.checkSpeedDrainThresholds()
-            if thresholdsHit:
-                logging.info("Updating agent configuration for speed drain...")
-                self.updateAgentSpeedDrainConfig(thresholdsHit)
-            # now collect drain statistics
-            try:
+        try:
+            # check the agent aux db config to see if the agent is in drain mode
+            if self.agentConfig["UserDrainMode"] or self.agentConfig["AgentDrainMode"]:
+                # check to see if the agent hit any speed drain thresholds
+                thresholdsHit = self.checkSpeedDrainThresholds()
+                if thresholdsHit:
+                    logging.info("Updating agent configuration for speed drain...")
+                    self.updateAgentSpeedDrainConfig(thresholdsHit)
+                # now collect drain statistics
                 DrainStatusPoller.drainStats = self.drainAPI.collectDrainInfo()
                 logging.info("Finished collecting agent drain status.")
-                logging.info("Drain stats: " + str(DrainStatusPoller.drainStats))
-
-            except Exception as ex:
-                msg = "Error occurred, will retry later:\\n"
-                msg += str(ex)
-                logging.exception(msg)
-        else:
-            logging.info("Agent not in drain mode. Resetting flags and skipping drain check...")
-            self.resetAgentSpeedDrainConfig()
+                logging.info("Drain stats: %s", str(DrainStatusPoller.drainStats))
+            else:
+                logging.info("Agent not in drain mode. Resetting flags and skipping drain check...")
+                self.resetAgentSpeedDrainConfig()
+        except Exception as ex:
+            msg = "Error occurred, will retry later:\\n"
+            msg += str(ex)
+            logging.exception(msg)
 
     @classmethod
     def getDrainInfo(cls):
@@ -157,4 +158,4 @@ class DrainStatusPoller(BaseWorkerThread):
             else:
                 logging.warning("Speed drain configuration error for %s.  Please check aux db contents.", k)
 
-        return enableKeys
\\ No newline at end of file
+        return enableKeys