Merge pull request #8242 from chaen/v9.0_feat_cleanTransfoCache

[9.0] Clean TransformationAgentCache and speedups

Author: fstagni

src/DIRAC/ConfigurationSystem/Client/Helpers/Registry.py

@@ -1,24 +1,53 @@
-""" Helper for /Registry section
-"""
+"""Helper for /Registry section"""
 
 import errno
+import inspect
+import sys
 
 from threading import Lock
+from collections.abc import Iterable
 
 from cachetools import TTLCache, cached
+from cachetools.keys import hashkey
 
 
+from typing import Optional
+from collections.abc import Iterable
+
 from DIRAC import S_OK, S_ERROR
 from DIRAC.ConfigurationSystem.Client.Config import gConfig
 from DIRAC.ConfigurationSystem.Client.Helpers.CSGlobals import getVO
 
 ID_DN_PREFIX = "/O=DIRAC/CN="
 
+# 300 is the default CS refresh time
+
+CACHE_REFRESH_TIME = 300
 # pylint: disable=missing-docstring
 
 gBaseRegistrySection = "/Registry"
 
 
+def reset_all_caches():
+    """This method is called to clear all caches.
+    It is necessary to reinitialize them after the central CS
+    has been loaded
+    """
+    for cache in [
+        obj
+        for name, obj in inspect.getmembers(sys.modules[__name__])
+        if (inspect.isfunction(obj) and hasattr(obj, "cache_clear"))
+    ]:
+        cache.cache_clear()
+
+
+def get_username_for_dn_key(dn: str, userList: Optional[Iterable[str]] = None):
+    if userList:
+        return hashkey(dn, *sorted(userList))
+    return hashkey(dn)
+
+
+@cached(TTLCache(maxsize=1000, ttl=CACHE_REFRESH_TIME), lock=Lock(), key=get_username_for_dn_key)
 def getUsernameForDN(dn, usersList=None):
     """Find DIRAC user for DN
 
@@ -39,6 +68,7 @@ def getUsernameForDN(dn, usersList=None):
     return S_ERROR(f"No username found for dn {dn}")
 
 
+@cached(TTLCache(maxsize=1000, ttl=CACHE_REFRESH_TIME), lock=Lock())
 def getDNForUsername(username):
     """Get user DN for user
 
@@ -419,6 +449,7 @@ def getBannedIPs():
     return gConfig.getValue(f"{gBaseRegistrySection}/BannedIPs", [])
 
 
+@cached(TTLCache(maxsize=1000, ttl=CACHE_REFRESH_TIME), lock=Lock())
 def getVOForGroup(group):
     """Search VO name for group
 
@@ -634,10 +665,7 @@ def getDNProperty(userDN, value, defaultValue=None):
     return S_OK(defaultValue)
 
 
-_cache_getProxyProvidersForDN = TTLCache(maxsize=1000, ttl=60)
-
-
-@cached(_cache_getProxyProvidersForDN, lock=Lock())
+@cached(TTLCache(maxsize=1000, ttl=CACHE_REFRESH_TIME), lock=Lock())
 def getProxyProvidersForDN(userDN):
     """Get proxy providers by user DN
 

src/DIRAC/ConfigurationSystem/Client/LocalConfiguration.py

@@ -568,6 +568,9 @@ def enableCS(self):
             objLoader = ObjectLoader()
             objLoader.reloadRootModules()
             self.__initLogger(self.componentName, self.loggingSection, forceInit=True)
+            from DIRAC.ConfigurationSystem.Client.Helpers.Registry import reset_all_caches
+
+            reset_all_caches()
         return res
 
     def isCSEnabled(self):

src/DIRAC/Core/Security/ProxyInfo.py

@@ -1,10 +1,13 @@
 """
- Set of utilities to retrieve Information from proxy
+Set of utilities to retrieve Information from proxy
 """
+
 import base64
 
 from DIRAC import S_ERROR, S_OK, gLogger
 from DIRAC.ConfigurationSystem.Client.Helpers import Registry
+from DIRAC.ConfigurationSystem.Client.Helpers.CSGlobals import getVO
+
 from DIRAC.Core.Security import Locations
 from DIRAC.Core.Security.DiracX import diracxTokenFromPEM
 from DIRAC.Core.Security.VOMS import VOMS
@@ -207,10 +210,11 @@ def getVOfromProxyGroup():
     """
     Return the VO associated to the group in the proxy
     """
-    voName = Registry.getVOForGroup("NoneExistingGroup")
+
     ret = getProxyInfo(disableVOMS=True)
-    if not ret["OK"]:
-        return S_OK(voName)
-    if "group" in ret["Value"]:
+    if not ret["OK"] or "group" not in ret["Value"]:
+        voName = getVO()
+    else:
         voName = Registry.getVOForGroup(ret["Value"]["group"])
+
     return S_OK(voName)

src/DIRAC/TransformationSystem/Agent/TransformationAgent.py

@@ -1,4 +1,4 @@
-"""  TransformationAgent processes transformations found in the transformation database.
+"""TransformationAgent processes transformations found in the transformation database.
 
 The following options can be set for the TransformationAgent.
 
@@ -8,13 +8,15 @@
   :dedent: 2
   :caption: TransformationAgent options
 """
+
 from importlib import import_module
 
 import time
 import os
 import datetime
 import pickle
 import concurrent.futures
+from pathlib import Path
 
 from DIRAC import S_OK, S_ERROR
 from DIRAC.ConfigurationSystem.Client.Helpers.Operations import Operations
@@ -127,6 +129,9 @@ def execute(self):
         if not res["OK"]:
             self._logError("Failed to obtain transformations:", res["Message"])
             return S_OK()
+
+        active_trans_ids = [t["TransformationID"] for t in res["Value"]]
+        self.cleanOldTransformationCache(active_trans_ids)
         # Process the transformations
         count = 0
         future_to_transID = {}
@@ -164,6 +169,22 @@ def execute(self):
 
         return S_OK()
 
+    def cleanOldTransformationCache(self, active_trans_ids: list[int]):
+        cache_filenames = {Path(self.__cacheFile(tid)) for tid in active_trans_ids}
+        existing_caches = set(Path(self.workDirectory).glob("*.pkl"))
+        useless_cache_files = existing_caches - cache_filenames
+
+        if useless_cache_files:
+            self._logInfo(f"Found potentially {len(useless_cache_files)} useless cache files")
+
+        # Since idle transformations aren't in active_trans_ids, let's filter it more
+        # and take only files that haven't been touched for 2 month
+        last_update_threshold = (datetime.datetime.utcnow() - datetime.timedelta(days=60)).timestamp()
+
+        for cache_file in useless_cache_files:
+            if Path(cache_file).stat().st_mtime < last_update_threshold:
+                cache_file.unlink()
+
     def getTransformations(self):
         """Obtain the transformations to be executed - this is executed at the start of every loop (it's really the
         only real thing in the execute()

src/DIRAC/TransformationSystem/Client/Utilities.py

@@ -9,6 +9,8 @@
 import ast
 import random
 
+from cachetools import LRUCache, cached
+from cachetools.keys import hashkey
 from DIRAC import S_OK, S_ERROR, gLogger
 
 from DIRAC.Core.Utilities.List import breakListIntoChunks
@@ -400,6 +402,10 @@ def isSameSE(self, se1, se2):
 
         return StorageElement(se1).isSameSE(StorageElement(se2))
 
+    @cached(
+        LRUCache(maxsize=1024),
+        key=lambda _, a, b: hashkey(a, *sorted(b)),
+    )
     def isSameSEInList(self, se1, seList):
         """Check if an SE is the same as any in a list"""
         if se1 in seList: