Merge pull request #7049 from atsareg/tornado-sandbox

[8.0] Simple  TornadoSandboxStore handler

Author: fstagni

src/DIRAC/Core/Tornado/Client/TornadoClient.py

@@ -24,6 +24,7 @@
 
 from DIRAC.Core.Tornado.Client.private.TornadoBaseClient import TornadoBaseClient
 from DIRAC.Core.Utilities.JEncode import encode
+from DIRAC.Core.Utilities.File import getGlobbedTotalSize
 
 
 class TornadoClient(TornadoBaseClient):
@@ -78,6 +79,25 @@ def receiveFile(self, destFile, *args):
         retVal = self._request(outputFile=destFile, **rpcCall)
         return retVal
 
+    def sendFile(self, filename, fileID):
+        """
+        Equivalent of :py:meth:`~DIRAC.Core.DISET.TransferClient.TransferClient.sendFile`
+
+        In practice, it calls the remote method `streamFromClient` and transfers the file
+        as a string type value
+
+        :param str filename: file (or path) where to store the result
+        :param any fileID: tuple if file identifiers
+        :returns: S_OK/S_ERROR result of the remote RPC call
+        """
+        fileSize = getGlobbedTotalSize(filename)
+        token = ""
+
+        with open(filename, "br") as input:
+            result = self.executeRPC("streamFromClient", *[fileID, token, fileSize, input.read()])
+
+        return result
+
 
 def executeRPCStub(rpcStub):
     """

src/DIRAC/FrameworkSystem/Client/TokenManagerClient.py

@@ -57,7 +57,7 @@ def getToken(
         idpObj = result["Value"]
 
         # Search for an existing token in tokensCache
-        cachedKey = getCachedKey(idpObj.name, username, userGroup, scope, audience)
+        cachedKey = getCachedKey(idpObj, username, userGroup, scope, audience)
         result = getCachedToken(self.__tokensCache, cachedKey, requiredTimeLeft)
         if result["OK"]:
             # A valid token has been found and is returned

src/DIRAC/FrameworkSystem/Service/TokenManagerHandler.py

@@ -213,7 +213,7 @@ def export_getToken(
         idpObj = result["Value"]
 
         # Search for an existing token in tokensCache
-        cachedKey = getCachedKey(idpObj.name, username, userGroup, scope, audience)
+        cachedKey = getCachedKey(idpObj, username, userGroup, scope, audience)
         result = getCachedToken(self.__tokensCache, cachedKey, requiredTimeLeft)
         if result["OK"]:
             # A valid token has been found and is returned
@@ -259,7 +259,7 @@ def export_getToken(
                     result = self.__checkProperties(dn, userGroup)
                     if result["OK"]:
                         # refresh token with requested scope
-                        result = idpObj.refreshToken(tokens.get("refresh_token"))
+                        result = idpObj.refreshToken(tokens.get("refresh_token"), group=userGroup, scope=scope)
                         if result["OK"]:
                             # caching new tokens
                             self.__tokensCache.add(

src/DIRAC/Resources/IdProvider/Utilities.py

@@ -39,7 +39,8 @@ def getIdProviderIdentifierFromIssuerAndClientID(issuer, clientID):
 
     for identifier in result["Value"]:
         testedClientID = gConfig.getValue(f"/Resources/IdProviders/{identifier}/client_id")
-        if testedClientID and testedClientID == clientID:
+        # clientID is not always available, e.g. in case of a token of particular user
+        if not clientID or (testedClientID and testedClientID == clientID):
             # Found the client ID but need to check the issuer
             # 2 different issuers could theoretically have a same client ID
             testedIssuer = gConfig.getValue(f"/Resources/IdProviders/{identifier}/issuer")

src/DIRAC/Resources/Storage/DIPStorage.py

@@ -160,7 +160,7 @@ def getFile(self, path, localPath=False):
 
     def __getFile(self, src_url, dest_file):
         transferClient = TransferClient(self.url)
-        res = transferClient.receiveFile(dest_file, src_url, token=self.checkSum)
+        res = transferClient.receiveFile(dest_file, src_url, self.checkSum)
         if not res["OK"]:
             return res
         if not os.path.exists(dest_file):

src/DIRAC/WorkloadManagementSystem/ConfigTemplate.cfg

@@ -149,6 +149,28 @@ Services
     }
   }
   ##END
+  ##BEGIN TornadoSandboxStore
+  TornadoSandboxStore
+  {
+    Protocol = https
+    LocalSE = ProductionSandboxSE
+    MaxThreads = 200
+    toClientMaxThreads = 100
+    Backend = local
+    MaxSandboxSizeMiB = 10
+    SandboxPrefix = Sandbox
+    BasePath = /opt/dirac/storage/sandboxes
+    DelayedExternalDeletion = True
+    Authorization
+    {
+      Default = authenticated
+      FileTransfer
+      {
+        Default = authenticated
+      }
+    }
+  }
+  ##END
   OptimizationMind
   {
     Port = 9175

src/DIRAC/WorkloadManagementSystem/Service/SandboxStoreHandler.py

@@ -11,6 +11,7 @@
 import time
 import threading
 import tempfile
+import hashlib
 
 from DIRAC import gLogger, S_OK, S_ERROR
 from DIRAC.Core.DISET.RequestHandler import RequestHandler
@@ -24,9 +25,10 @@
 from DIRAC.RequestManagementSystem.Client.Operation import Operation
 from DIRAC.RequestManagementSystem.Client.File import File
 from DIRAC.Resources.Storage.StorageElement import StorageElement
+from DIRAC.Core.Utilities.File import getGlobbedTotalSize
 
 
-class SandboxStoreHandler(RequestHandler):
+class SandboxStoreHandlerMixin:
     __purgeCount = -1
     __purgeLock = threading.Lock()
     __purgeWorking = False
@@ -44,10 +46,10 @@ def initializeHandler(cls, serviceInfoDict):
             return S_ERROR(f"Can't connect to DB: {repr(excp)}")
         return S_OK()
 
-    def initialize(self):
+    def initializeRequest(self):
         self.__backend = self.getCSOption("Backend", "local")
         self.__localSEName = self.getCSOption("LocalSE", "SandboxSE")
-        self.__maxUploadBytes = self.getCSOption("MaxSandboxSizeMiB", 10) * 1048576
+        self._maxUploadBytes = self.getCSOption("MaxSandboxSizeMiB", 10) * 1048576
         if self.__backend.lower() == "local" or self.__backend == self.__localSEName:
             self.__useLocalStorage = True
             self.__seNameToUse = self.__localSEName
@@ -75,13 +77,14 @@ def __getSandboxPath(self, md5):
         pathItems.extend([md5[0:3], md5[3:6], md5])
         return os.path.join(*pathItems)
 
-    def transfer_fromClient(self, fileId, token, fileSize, fileHelper):
+    def _getFromClient(self, fileId, token, fileSize, fileHelper=None, data=""):
         """
         Receive a file as a sandbox
         """
 
-        if self.__maxUploadBytes and fileSize > self.__maxUploadBytes:
-            fileHelper.markAsTransferred()
+        if self._maxUploadBytes and fileSize > self._maxUploadBytes:
+            if fileHelper:
+                fileHelper.markAsTransferred()
             return S_ERROR("Sandbox is too big. Please upload it to a grid storage element")
 
         if isinstance(fileId, (list, tuple)):
@@ -113,7 +116,8 @@ def transfer_fromClient(self, fileId, token, fileSize, fileHelper):
         result = self.sandboxDB.getSandboxId(seName, sePFN, credDict["username"], credDict["group"])
         if result["OK"]:
             gLogger.info("Sandbox already exists. Skipping upload")
-            fileHelper.markAsTransferred()
+            if fileHelper:
+                fileHelper.markAsTransferred()
             sbURL = f"SB:{seName}|{sePFN}"
             assignTo = {key: [(sbURL, assignTo[key])] for key in assignTo}
             result = self.export_assignSandboxesToEntities(assignTo)
@@ -126,14 +130,35 @@ def transfer_fromClient(self, fileId, token, fileSize, fileHelper):
         else:
             hdPath = False
         # Write to local file
-        result = self.__networkToFile(fileHelper, hdPath)
+
+        if fileHelper:
+            result = self.__networkToFile(fileHelper, hdPath)
+        elif data:
+            hdPath = os.path.realpath(hdPath)
+            mkDir(os.path.dirname(hdPath))
+            with open(hdPath, "bw") as output:
+                output.write(data)
+            result = S_OK(hdPath)
+        else:
+            result = S_ERROR("No data provided")
+
         if not result["OK"]:
             gLogger.error("Error while receiving sandbox file", result["Message"])
             return result
         hdPath = result["Value"]
         gLogger.info("Wrote sandbox to file", hdPath)
         # Check hash!
-        if fileHelper.getHash() != aHash:
+        if fileHelper:
+            hdHash = fileHelper.getHash()
+        else:
+            oMD5 = hashlib.md5()
+            with open(hdPath, "rb") as fd:
+                bData = fd.read(10240)
+                while bData:
+                    oMD5.update(bData)
+                    bData = fd.read(10240)
+            hdHash = oMD5.hexdigest()
+        if hdHash != aHash:
             self.__secureUnlinkFile(hdPath)
             gLogger.error("Hashes don't match! Client defined hash is different with received data hash!")
             return S_ERROR("Hashes don't match!")
@@ -147,13 +172,14 @@ def transfer_fromClient(self, fileId, token, fileSize, fileHelper):
             sbPath = result["Value"][1]
         # Register!
         gLogger.info("Registering sandbox in the DB with", f"SB:{self.__seNameToUse}|{sbPath}")
+        fSize = getGlobbedTotalSize(hdPath)
         result = self.sandboxDB.registerAndGetSandbox(
             credDict["username"],
             credDict["DN"],
             credDict["group"],
             self.__seNameToUse,
             sbPath,
-            fileHelper.getTransferedBytes(),
+            fSize,
         )
         if not result["OK"]:
             self.__secureUnlinkFile(hdPath)
@@ -166,6 +192,13 @@ def transfer_fromClient(self, fileId, token, fileSize, fileHelper):
             return result
         return S_OK(sbURL)
 
+    def transfer_fromClient(self, fileId, token, fileSize, fileHelper):
+        """
+        Receive a file as a sandbox
+        """
+
+        return self._getFromClient(fileId, token, fileSize, fileHelper=fileHelper)
+
     def transfer_bulkFromClient(self, fileId, token, _fileSize, fileHelper):
         """Receive files packed into a tar archive by the fileHelper logic.
         token is used for access rights confirmation.
@@ -259,7 +292,7 @@ def __networkToFile(self, fileHelper, destFileName=False):
                 fd = tfd
             else:
                 fd = open(destFileName, "wb")
-            result = fileHelper.networkToDataSink(fd, maxFileSize=self.__maxUploadBytes)
+            result = fileHelper.networkToDataSink(fd, maxFileSize=self._maxUploadBytes)
             fd.close()
         except Exception as e:
             gLogger.error("Cannot open to write destination file", f"{destFileName}: {repr(e).replace(',)', ')')}")
@@ -328,7 +361,7 @@ def export_assignSandboxesToEntities(self, enDict, ownerName="", ownerGroup="",
         Expects a dict of { entityId : [ ( SB, SBType ), ... ] }
         """
         if not entitySetup:
-            entitySetup = self.serviceInfoDict["clientSetup"]
+            entitySetup = self.diracSetup
         credDict = self.getRemoteCredentials()
         return self.sandboxDB.assignSandboxesToEntities(
             enDict, credDict["username"], credDict["group"], entitySetup, ownerName, ownerGroup
@@ -344,7 +377,7 @@ def export_unassignEntities(self, entitiesList, entitiesSetup=False):
         Unassign a list of jobs
         """
         if not entitiesSetup:
-            entitiesSetup = self.serviceInfoDict["clientSetup"]
+            entitiesSetup = self.diracSetup
         credDict = self.getRemoteCredentials()
         return self.sandboxDB.unassignEntities({entitiesSetup: entitiesList}, credDict["username"], credDict["group"])
 
@@ -358,7 +391,7 @@ def export_getSandboxesAssignedToEntity(self, entityId, entitySetup=False):
         Get the sandboxes associated to a job and the association type
         """
         if not entitySetup:
-            entitySetup = self.serviceInfoDict["clientSetup"]
+            entitySetup = self.diracSetup
         credDict = self.getRemoteCredentials()
         result = self.sandboxDB.getSandboxesAssignedToEntity(
             entityId, entitySetup, credDict["username"], credDict["group"]
@@ -400,6 +433,10 @@ def transfer_toClient(self, fileID, token, fileHelper):
         fileID is the local file name in the SE.
         token is used for access rights confirmation.
         """
+
+        return self._sendToClient(fileID, token, fileHelper=fileHelper)
+
+    def _sendToClient(self, fileID, token, fileHelper=None, raw=False):
         credDict = self.getRemoteCredentials()
         serviceURL = self.serviceInfoDict["URL"]
         filePath = fileID.replace(serviceURL, "")
@@ -412,13 +449,20 @@ def transfer_toClient(self, fileID, token, fileHelper):
         hdPath = self.__sbToHDPath(filePath)
         if not os.path.isfile(hdPath):
             return S_ERROR("Sandbox does not exist")
-        result = fileHelper.getFileDescriptor(hdPath, "rb")
-        if not result["OK"]:
-            return S_ERROR(f"Failed to get file descriptor: {result['Message']}")
-        fd = result["Value"]
-        result = fileHelper.FDToNetwork(fd)
-        fileHelper.oFile.close()
-        return result
+
+        if fileHelper:
+            result = fileHelper.getFileDescriptor(hdPath, "rb")
+            if not result["OK"]:
+                return result
+            fd = result["Value"]
+            result = fileHelper.FDToNetwork(fd)
+            fileHelper.oFile.close()
+            return result
+
+        with open(hdPath, "rb") as fd:
+            if raw:
+                return fd.read()
+            return S_OK(fd.read())
 
     ##################
     # Purge sandboxes
@@ -531,3 +575,10 @@ def __deleteSandboxFromExternalBackend(self, SEName, SEPFN):
             except Exception:
                 gLogger.exception("RM raised an exception while trying to delete a remote sandbox")
                 return S_ERROR("RM raised an exception while trying to delete a remote sandbox")
+
+
+class SandboxStoreHandler(SandboxStoreHandlerMixin, RequestHandler):
+    def initialize(self):
+        # we need it still in 8.0
+        self.diracSetup = self.serviceInfoDict["clientSetup"]
+        return self.initializeRequest()

src/DIRAC/WorkloadManagementSystem/Service/TornadoSandboxStoreHandler.py

@@ -0,0 +1,26 @@
+""" Tornado-based HTTPs SandboxStore service.
+
+.. literalinclude:: ../ConfigTemplate.cfg
+  :start-after: ##BEGIN TornadoSandboxStore
+  :end-before: ##END
+  :dedent: 2
+  :caption: SandboxStore options
+
+"""
+from DIRAC import gLogger, S_OK, S_ERROR
+from DIRAC.Core.Tornado.Server.TornadoService import TornadoService
+from DIRAC.WorkloadManagementSystem.Service.SandboxStoreHandler import SandboxStoreHandlerMixin
+
+
+class TornadoSandboxStoreHandler(SandboxStoreHandlerMixin, TornadoService):
+    def initializeRequest(self):
+        self.diracSetup = self.get_argument("clientSetup")
+        # Ugly, but makes DIPS and HTTPS handlers compatible, TBD properly
+        self.serviceInfoDict = self._serviceInfoDict
+        return SandboxStoreHandlerMixin.initializeRequest(self)
+
+    def export_streamFromClient(self, fileId, token, fileSize, data):
+        return self._getFromClient(fileId, token, fileSize, data=data)
+
+    def export_streamToClient(self, fileId, token=""):
+        return self._sendToClient(fileId, token, raw=True)

tests/Jenkins/dirac-cfg-update-server.py

@@ -25,6 +25,8 @@
 
 csAPI.setOption("Systems/WorkloadManagement/Production/Services/SandboxStore/BasePath", f"{setupName}/sandboxes")
 csAPI.setOption("Systems/WorkloadManagement/Production/Services/SandboxStore/LogLevel", "DEBUG")
+csAPI.setOption("Systems/WorkloadManagement/Production/Services/TornadoSandboxStore/BasePath", f"{setupName}/sandboxes")
+csAPI.setOption("Systems/WorkloadManagement/Production/Services/TornadoSandboxStore/LogLevel", "DEBUG")
 
 # Now setting a SandboxSE as the following:
 #     ProductionSandboxSE