Merge pull request #8252 from fstagni/90_sbstore_fixes

[9.0] SandboxStore simplifications

Author: fstagni

src/DIRAC/WorkloadManagementSystem/Agent/JobCleaningAgent.py

@@ -36,9 +36,9 @@
 from DIRAC.RequestManagementSystem.Client.Request import Request
 from DIRAC.WorkloadManagementSystem.Client import JobStatus
 from DIRAC.WorkloadManagementSystem.Client.JobMonitoringClient import JobMonitoringClient
-from DIRAC.WorkloadManagementSystem.Client.SandboxStoreClient import SandboxStoreClient
 from DIRAC.WorkloadManagementSystem.Client.WMSClient import WMSClient
 from DIRAC.WorkloadManagementSystem.DB.JobDB import JobDB
+from DIRAC.WorkloadManagementSystem.DB.SandboxMetadataDB import SandboxMetadataDB
 
 
 class JobCleaningAgent(AgentModule):
@@ -152,8 +152,9 @@ def removeDeletedJobs(self):
             return S_OK()
 
         self.log.info("Unassigning sandboxes from soon to be deleted jobs", f"({len(jobList)})")
-        result = SandboxStoreClient(useCertificates=True).unassignJobs(jobList)
-        if not result["OK"]:
+
+        entitiesList = [f"Job:{jobId}" for jobId in jobList]
+        if not (result := SandboxMetadataDB().unassignEntities(entitiesList))["OK"]:
             self.log.error("Cannot unassign jobs to sandboxes", result["Message"])
             return result
 

src/DIRAC/WorkloadManagementSystem/Client/SandboxStoreClient.py

@@ -20,7 +20,6 @@
 
 class SandboxStoreClient:
     __validSandboxTypes = ("Input", "Output")
-    __smdb = None
 
     def __init__(self, rpcClient=None, transferClient=None, smdb=False, **kwargs):
         """Constructor
@@ -37,21 +36,8 @@ def __init__(self, rpcClient=None, transferClient=None, smdb=False, **kwargs):
         self.__transferClient = transferClient
         self.__kwargs = kwargs
         self.__vo = None
-        SandboxStoreClient.__smdb = smdb
         if "delegatedGroup" in kwargs:
             self.__vo = getVOForGroup(kwargs["delegatedGroup"])
-        if SandboxStoreClient.__smdb is True:
-            try:
-                from DIRAC.WorkloadManagementSystem.DB.SandboxMetadataDB import SandboxMetadataDB
-
-                SandboxStoreClient.__smdb = SandboxMetadataDB()
-                result = SandboxStoreClient.__smdb._getConnection()  # pylint: disable=protected-access
-                if not result["OK"]:
-                    SandboxStoreClient.__smdb = False
-                else:
-                    result["Value"].close()
-            except (ImportError, RuntimeError, AttributeError):
-                SandboxStoreClient.__smdb = False
 
     def __getRPCClient(self):
         """Get an RPC client for SB service"""
@@ -227,29 +213,6 @@ def downloadSandbox(self, sbLocation, destinationDir="", inMemory=False, unpack=
     ##############
     # Jobs
 
-    def assignSandboxesToJob(self, jobId, sbList, ownerName="", ownerGroup=""):
-        """
-        Assign sandboxes to a job.
-        sbList must be a list of sandboxes and relation types
-        sbList = [ ( "SB:SEName|SEPFN", "Input" ), ( "SB:SEName|SEPFN", "Output" ) ]
-        """
-        eId = f"Job:{jobId}"
-        for sbT in sbList:
-            if sbT[1] not in self.__validSandboxTypes:
-                return S_ERROR(f"Invalid Sandbox type {sbT[1]}")
-        if SandboxStoreClient.__smdb and ownerName and ownerGroup:
-            return SandboxStoreClient.__smdb.assignSandboxesToEntities({eId: sbList}, ownerName, ownerGroup)
-        return self.__getRPCClient().assignSandboxesToEntities({eId: sbList}, ownerName, ownerGroup)
-
-    def unassignJobs(self, jobIdList):
-        """Unassign SB to a job"""
-        if isinstance(jobIdList, int):
-            jobIdList = [jobIdList]
-        entitiesList = []
-        for jobId in jobIdList:
-            entitiesList.append(f"Job:{jobId}")
-        return self.__getRPCClient().unassignEntities(entitiesList)
-
     def downloadSandboxForJob(self, jobId, sbType, destinationPath="", inMemory=False, unpack=True):
         """Download SB for a job"""
         result = self.__getRPCClient().getSandboxesAssignedToEntity(f"Job:{jobId}")

src/DIRAC/WorkloadManagementSystem/DB/SandboxMetadataDB.py

@@ -1,6 +1,6 @@
 """ SandboxMetadataDB class is a front-end to the metadata for sandboxes
 """
-from DIRAC import S_ERROR, S_OK, gLogger
+from DIRAC import S_ERROR, S_OK
 from DIRAC.ConfigurationSystem.Client.Helpers import Registry
 from DIRAC.Core.Base.DB import DB
 from DIRAC.Core.Security import Properties
@@ -220,48 +220,22 @@ def assignSandboxesToEntities(self, enDict, requesterName, requesterGroup, owner
             return result
         return S_OK(assigned)
 
-    def __entitiesByRequesterCond(self, requesterName, requesterGroup):
-        sqlCond = []
-        requesterProps = Registry.getPropertiesForEntity(requesterGroup, name=requesterName)
-        if Properties.JOB_ADMINISTRATOR in requesterProps:
-            # Do nothing, just ensure it doesn't fit in the other cases
-            pass
-        elif Properties.JOB_SHARING in requesterProps:
-            sqlCond.append(f"o.OwnerGroup='{requesterGroup}'")
-        elif Properties.NORMAL_USER in requesterProps:
-            sqlCond.append(f"o.OwnerGroup='{requesterGroup}'")
-            sqlCond.append(f"o.Owner='{requesterName}'")
-        else:
-            return S_ERROR("Not authorized to access sandbox")
-        return sqlCond
-
     @convertToReturnValue
-    def unassignEntities(self, entities, requesterName, requesterGroup):
+    def unassignEntities(self, entities: list):
         """
-        Unassign jobs to sandboxes
+        Unassign entities to sandboxes. Entities are a list of strings, e.g. ['job:1234', 'job:5678'].
 
         :param list entities: list of entities to unassign
         """
         if not entities:
             return None
-        conds = self.__entitiesByRequesterCond(requesterName, requesterGroup)
 
-        sqlCmd = "CREATE TEMPORARY TABLE to_delete_EntityId (EntityId VARCHAR(128)  NOT NULL, PRIMARY KEY (EntityId)) ENGINE=MEMORY;"
+        sqlCmd = "CREATE TEMPORARY TABLE to_delete_EntityId (EntityId VARCHAR(128) NOT NULL, PRIMARY KEY (EntityId)) ENGINE=MEMORY;"
         returnValueOrRaise(self._update(sqlCmd))
         try:
             sqlCmd = "INSERT INTO to_delete_EntityId (EntityId) VALUES ( %s )"
             returnValueOrRaise(self._updatemany(sqlCmd, [(e,) for e in entities]))
             sqlCmd = "DELETE m from `sb_EntityMapping` m JOIN to_delete_EntityId t USING (EntityId)"
-            if conds:
-                sqlCmd = " ".join(
-                    [
-                        sqlCmd,
-                        "JOIN `sb_SandBoxes` s ON s.SBId = m.SBId",
-                        "JOIN `sb_Owners` o ON s.OwnerId = o.OwnerId",
-                        "WHERE",
-                        " AND ".join(conds),
-                    ]
-                )
             returnValueOrRaise(self._update(sqlCmd))
         finally:
             sqlCmd = "DROP TEMPORARY TABLE to_delete_EntityId"
@@ -313,16 +287,30 @@ def getUnusedSandboxes(self):
         sqlCmd = f"SELECT SBId, SEName, SEPFN FROM `sb_SandBoxes` WHERE SEPFN not like '/S3/%' AND (( {' ) OR ( '.join(sqlCond)} ))"
         return self._query(sqlCmd)
 
+    @convertToReturnValue
     def deleteSandboxes(self, SBIdList):
         """
-        Delete sandboxes
+        Delete sandboxes using a temporary table for efficiency and consistency.
         """
-        sqlSBList = ", ".join([str(sbid) for sbid in SBIdList])
-        for table in ("sb_SandBoxes", "sb_EntityMapping"):
-            sqlCmd = f"DELETE FROM `{table}` WHERE SBId IN ( {sqlSBList} )"
-            result = self._update(sqlCmd)
-            if not result["OK"]:
-                return result
+        if not SBIdList:
+            return S_OK()
+        # Create temporary table
+        sqlCmd = "CREATE TEMPORARY TABLE to_delete_SBId (SBId INTEGER(10) UNSIGNED NOT NULL, PRIMARY KEY (SBId)) ENGINE=MEMORY;"
+        returnValueOrRaise(self._update(sqlCmd))
+        try:
+            # Insert SBIds into temporary table
+            sqlCmd = "INSERT INTO to_delete_SBId (SBId) VALUES (%s)"
+            returnValueOrRaise(self._updatemany(sqlCmd, [(sbid,) for sbid in SBIdList]))
+            # Delete from sb_EntityMapping first (to respect FK constraints if any)
+            sqlCmd = "DELETE FROM `sb_EntityMapping` WHERE SBId IN (SELECT SBId FROM to_delete_SBId)"
+            returnValueOrRaise(self._update(sqlCmd))
+            # Delete from sb_SandBoxes
+            sqlCmd = "DELETE FROM `sb_SandBoxes` WHERE SBId IN (SELECT SBId FROM to_delete_SBId)"
+            returnValueOrRaise(self._update(sqlCmd))
+        finally:
+            # Drop temporary table
+            sqlCmd = "DROP TEMPORARY TABLE to_delete_SBId"
+            returnValueOrRaise(self._update(sqlCmd))
         return S_OK()
 
     def getSandboxId(self, SEName, SEPFN, requesterName, requesterGroup, field="SBId"):

src/DIRAC/WorkloadManagementSystem/Executor/JobSanity.py

@@ -1,6 +1,6 @@
 """ The Job Sanity executor assigns sandboxes to the job """
 from DIRAC import S_OK
-from DIRAC.WorkloadManagementSystem.Client.SandboxStoreClient import SandboxStoreClient
+from DIRAC.WorkloadManagementSystem.DB.SandboxMetadataDB import SandboxMetadataDB
 from DIRAC.WorkloadManagementSystem.Executor.Base.OptimizerExecutor import OptimizerExecutor
 
 
@@ -15,7 +15,6 @@ class JobSanity(OptimizerExecutor):
     @classmethod
     def initializeOptimizer(cls):
         """Initialize specific parameters for JobSanityAgent."""
-        cls.sandboxClient = SandboxStoreClient(useCertificates=True, smdb=True)
         return S_OK()
 
     def optimizeJob(self, jid, jobState):
@@ -57,7 +56,8 @@ def checkInputSandbox(self, jobState, manifest):
         if not numSBsToAssign:
             return S_OK(0)
         self.jobLog.info("Assigning sandboxes", f"({numSBsToAssign} on behalf of {ownerName}@{ownerGroup}@{vo})")
-        result = self.sandboxClient.assignSandboxesToJob(jobState.jid, sbsToAssign, ownerName, ownerGroup)
+        eId = f"Job:{jobState.jid}"
+        result = SandboxMetadataDB().assignSandboxesToEntities({eId: sbsToAssign}, ownerName, ownerGroup)
         if not result["OK"]:
             self.jobLog.error("Could not assign sandboxes in the SandboxStore")
             return result

src/DIRAC/WorkloadManagementSystem/Service/SandboxStoreHandler.py

@@ -22,7 +22,6 @@
 from DIRAC.Core.Security import Properties
 from DIRAC.Core.Utilities.File import getGlobbedTotalSize, mkDir
 from DIRAC.Core.Utilities.ObjectLoader import ObjectLoader
-from DIRAC.DataManagementSystem.Service.StorageElementHandler import getDiskSpace
 from DIRAC.FrameworkSystem.Utilities.diracx import TheImpersonator
 
 
@@ -139,7 +138,7 @@ def _getFromClient(self, fileId, token, fileSize, fileHelper=None, data=""):
                     gLogger.debug("Sandbox already exists in storage backend", res.pfn)
 
                 assignTo = {key: [(res.pfn, assignTo[key])] for key in assignTo}
-                result = self.export_assignSandboxesToEntities(assignTo)
+                result = self.sandboxDB.assignSandboxesToEntities(assignTo, credDict["username"], credDict["group"])
                 if not result["OK"]:
                     return result
                 return S_OK(res.pfn)
@@ -153,7 +152,7 @@ def _getFromClient(self, fileId, token, fileSize, fileHelper=None, data=""):
                 fileHelper.markAsTransferred()
             sbURL = f"SB:{self.__localSEName}|{sbPath}"
             assignTo = {key: [(sbURL, assignTo[key])] for key in assignTo}
-            result = self.export_assignSandboxesToEntities(assignTo)
+            result = self.sandboxDB.assignSandboxesToEntities(assignTo, credDict["username"], credDict["group"])
             if not result["OK"]:
                 return result
             return S_OK(sbURL)
@@ -210,7 +209,9 @@ def _getFromClient(self, fileId, token, fileSize, fileHelper=None, data=""):
 
         sbURL = f"SB:{self.__localSEName}|{sbPath}"
         assignTo = {key: [(sbURL, assignTo[key])] for key in assignTo}
-        if not (result := self.export_assignSandboxesToEntities(assignTo))["OK"]:
+        if not (result := self.sandboxDB.assignSandboxesToEntities(assignTo, credDict["username"], credDict["group"]))[
+            "OK"
+        ]:
             return result
         return S_OK(sbURL)
 
@@ -328,33 +329,6 @@ def __moveToFinalLocation(self, localFilePath, sbPath):
 
         return result
 
-    ##################
-    # Assigning sbs to jobs
-
-    types_assignSandboxesToEntities = [dict]
-
-    def export_assignSandboxesToEntities(self, enDict, ownerName="", ownerGroup=""):
-        """
-        Assign sandboxes to jobs.
-        Expects a dict of { entityId : [ ( SB, SBType ), ... ] }
-        """
-        credDict = self.getRemoteCredentials()
-        return self.sandboxDB.assignSandboxesToEntities(
-            enDict, credDict["username"], credDict["group"], ownerName, ownerGroup
-        )
-
-    ##################
-    # Unassign sbs to jobs
-
-    types_unassignEntities = [(list, tuple)]
-
-    def export_unassignEntities(self, entitiesList):
-        """
-        Unassign a list of jobs
-        """
-        credDict = self.getRemoteCredentials()
-        return self.sandboxDB.unassignEntities(entitiesList, credDict["username"], credDict["group"])
-
     ##################
     # Getting assigned sandboxes
 
@@ -376,26 +350,6 @@ def export_getSandboxesAssignedToEntity(self, entityId):
             sbDict[SBType].append(f"SB:{SEName}|{SEPFN}")
         return S_OK(sbDict)
 
-    ##################
-    # Disk space left management
-
-    types_getFreeDiskSpace = []
-
-    def export_getFreeDiskSpace(self):
-        """Get the free disk space of the storage element
-        If no size is specified, terabytes will be used by default.
-        """
-
-        return getDiskSpace(self.getCSOption("BasePath", "/opt/dirac/storage/sandboxes"))
-
-    types_getTotalDiskSpace = []
-
-    def export_getTotalDiskSpace(self):
-        """Get the total disk space of the storage element
-        If no size is specified, terabytes will be used by default.
-        """
-        return getDiskSpace(self.getCSOption("BasePath", "/opt/dirac/storage/sandboxes"), total=True)
-
     ##################
     # Download sandboxes
 

tests/Integration/WorkloadManagementSystem/Test_SandboxStoreClient.py

@@ -38,26 +38,15 @@
 from DIRAC.tests.Utilities.utils import find_all
 from DIRAC.WorkloadManagementSystem.Client.SandboxStoreClient import SandboxStoreClient
 
-
 gLogger.setLevel("DEBUG")
 
 
 def test_SSCChain():
     """full test of functionalities"""
     ssc = SandboxStoreClient()
 
-    jobId = 1
-
     exeScriptLocation = find_all("exe-script.py", "../..", "/DIRAC/tests/Integration")[0]
     fileList = [exeScriptLocation]
 
     res = ssc.uploadFilesAsSandbox(fileList)
     assert res["OK"], res["Message"]
-
-    # TODO : FIXME
-    # res = ssc.downloadSandboxForJob(jobId, "Input")  # to run this we need the RSS on
-    # print(res)  # for debug...
-    # assert res["OK"], res["Message"]
-
-    res = ssc.unassignJobs([jobId])
-    assert res["OK"], res["Message"]