Merge pull request #6890 from atsareg/dev-ce-token

[8.0] Perform pilot operations with tokens

Author: fstagni

src/DIRAC/Resources/Computing/AREXComputingElement.py

@@ -356,11 +356,12 @@ def _killJob(self, arcJobList):
             return S_ERROR("REST interface not initialised. Cannot kill jobs.")
 
         # Get a proxy
-        result = self._prepareProxy()
-        if not result["OK"]:
-            self.log.error("Failed to set up proxy", result["Message"])
-            return result
-        self.session.cert = Locations.getProxyLocation()
+        if self.proxy:
+            result = self._prepareProxy()
+            if not result["OK"]:
+                self.log.error("Failed to set up proxy", result["Message"])
+                return result
+            self.session.cert = Locations.getProxyLocation()
 
         # List of jobs in json format for the REST query
         jobsJson = {"job": [{"id": job} for job in arcJobList]}
@@ -393,11 +394,12 @@ def submitJob(self, executableFile, proxy, numberOfJobs=1, inputs=None, outputs=
         self.log.verbose("Executable file path:", executableFile)
 
         # Get a proxy
-        result = self._prepareProxy()
-        if not result["OK"]:
-            self.log.error("Failed to set up proxy", result["Message"])
-            return result
-        self.session.cert = Locations.getProxyLocation()
+        if self.proxy:
+            result = self._prepareProxy()
+            if not result["OK"]:
+                self.log.error("Failed to set up proxy", result["Message"])
+                return result
+            self.session.cert = Locations.getProxyLocation()
 
         # Get a "delegation" and use the same delegation for all the jobs
         delegation = ""
@@ -476,11 +478,12 @@ def getCEStatus(self):
             return S_ERROR("REST interface not initialised. Cannot get CE status.")
 
         # Get a proxy
-        result = self._prepareProxy()
-        if not result["OK"]:
-            self.log.error("Failed to set up proxy", result["Message"])
-            return result
-        self.session.cert = Locations.getProxyLocation()
+        if self.proxy:
+            result = self._prepareProxy()
+            if not result["OK"]:
+                self.log.error("Failed to set up proxy", result["Message"])
+                return result
+            self.session.cert = Locations.getProxyLocation()
 
         # Try to find out which VO we are running for.
         # Essential now for REST interface.
@@ -597,11 +600,12 @@ def getJobStatus(self, jobIDList):
             return S_ERROR("REST interface not initialised. Cannot get job status.")
 
         # Get a proxy
-        result = self._prepareProxy()
-        if not result["OK"]:
-            self.log.error("AREXComputingElement: failed to set up proxy", result["Message"])
-            return result
-        self.session.cert = Locations.getProxyLocation()
+        if self.proxy:
+            result = self._prepareProxy()
+            if not result["OK"]:
+                self.log.error("AREXComputingElement: failed to set up proxy", result["Message"])
+                return result
+            self.session.cert = Locations.getProxyLocation()
 
         if not isinstance(jobIDList, list):
             jobIDList = [jobIDList]
@@ -688,11 +692,12 @@ def getJobLog(self, jobID):
             return S_ERROR("REST interface not initialised. Cannot get job output.")
 
         # Get a proxy
-        result = self._prepareProxy()
-        if not result["OK"]:
-            self.log.error("AREXComputingElement: failed to set up proxy", result["Message"])
-            return result
-        self.session.cert = Locations.getProxyLocation()
+        if self.proxy:
+            result = self._prepareProxy()
+            if not result["OK"]:
+                self.log.error("AREXComputingElement: failed to set up proxy", result["Message"])
+                return result
+            self.session.cert = Locations.getProxyLocation()
 
         # Extract stamp from the Job ID
         if ":::" in jobID:
@@ -749,11 +754,12 @@ def getJobOutput(self, jobID, workingDirectory=None):
             return S_ERROR("REST interface not initialised. Cannot get job output.")
 
         # Get a proxy
-        result = self._prepareProxy()
-        if not result["OK"]:
-            self.log.error("AREXComputingElement: failed to set up proxy", result["Message"])
-            return result
-        self.session.cert = Locations.getProxyLocation()
+        if self.proxy:
+            result = self._prepareProxy()
+            if not result["OK"]:
+                self.log.error("AREXComputingElement: failed to set up proxy", result["Message"])
+                return result
+            self.session.cert = Locations.getProxyLocation()
 
         # Extract stamp from the Job ID
         if ":::" in jobID:

src/DIRAC/Resources/Computing/ComputingElement.py

@@ -91,20 +91,13 @@ def setToken(self, token, valid=0):
 
     def _prepareProxy(self):
         """Set the environment variable X509_USER_PROXY"""
-        if not self.proxy:
-            result = getProxyInfo()
-            if not result["OK"]:
-                return S_ERROR("No proxy available")
-            if "path" in result["Value"]:
-                os.environ["X509_USER_PROXY"] = result["Value"]["path"]
-                return S_OK()
-        else:
+        if self.proxy:
             result = gProxyManager.dumpProxyToFile(self.proxy, requiredTimeLeft=self.minProxyTime)
             if not result["OK"]:
                 return result
             os.environ["X509_USER_PROXY"] = result["Value"]
 
-        self.log.debug(f"Set proxy variable X509_USER_PROXY to {os.environ['X509_USER_PROXY']}")
+            self.log.debug(f"Set proxy variable X509_USER_PROXY to {os.environ['X509_USER_PROXY']}")
         return S_OK()
 
     def isProxyValid(self, valid=1000):

src/DIRAC/WorkloadManagementSystem/Service/PilotManagerHandler.py

@@ -14,7 +14,7 @@
 from DIRAC.WorkloadManagementSystem.Client import PilotStatus
 from DIRAC.WorkloadManagementSystem.Service.WMSUtilities import (
     getPilotCE,
-    getPilotProxy,
+    setPilotCredentials,
     getPilotRef,
     killPilotsInQueues,
 )
@@ -135,13 +135,10 @@ def export_getPilotOutput(self, pilotReference):
         if not hasattr(ce, "getJobOutput"):
             return S_ERROR(f"Pilot output not available for {pilotDict['GridType']} CEs")
 
-        result = getPilotProxy(pilotDict)
+        result = setPilotCredentials(ce, pilotDict)
         if not result["OK"]:
             return result
 
-        proxy = result["Value"]
-        ce.setProxy(proxy)
-
         result = getPilotRef(pilotReference, pilotDict)
         if not result["OK"]:
             return result
@@ -213,13 +210,11 @@ def export_getPilotLoggingInfo(self, pilotReference):
             self.log.info("Pilot logging not available for", f"{pilotDict['GridType']} CEs")
             return S_ERROR(f"Pilot logging not available for {pilotDict['GridType']} CEs")
 
-        result = getPilotProxy(pilotDict)
+        # Set proxy or token for the CE
+        result = setPilotCredentials(ce, pilotDict)
         if not result["OK"]:
             return result
 
-        proxy = result["Value"]
-        ce.setProxy(proxy)
-
         result = getPilotRef(pilotReference, pilotDict)
         if not result["OK"]:
             return result

src/DIRAC/WorkloadManagementSystem/Service/WMSUtilities.py

@@ -6,8 +6,9 @@
 
 from DIRAC import S_OK, S_ERROR, gLogger, gConfig
 from DIRAC.ConfigurationSystem.Client.Helpers.Resources import getQueue
-from DIRAC.ConfigurationSystem.Client.Helpers.Registry import getGroupOption
+from DIRAC.ConfigurationSystem.Client.Helpers.Registry import getGroupOption, getUsernameForDN
 from DIRAC.FrameworkSystem.Client.ProxyManagerClient import gProxyManager
+from DIRAC.FrameworkSystem.Client.TokenManagerClient import gTokenManager
 from DIRAC.Resources.Computing.ComputingElementFactory import ComputingElementFactory
 
 
@@ -49,7 +50,11 @@ def getPilotCE(pilotDict):
 
 
 def getPilotProxy(pilotDict):
-    """Get a proxy bound to a pilot"""
+    """Get a proxy bound to a pilot
+
+    :param dict pilotDict: pilot parameters
+    :return: S_OK/S_ERROR with proxy as Value
+    """
     owner = pilotDict["OwnerDN"]
     group = pilotDict["OwnerGroup"]
 
@@ -62,6 +67,47 @@ def getPilotProxy(pilotDict):
     return S_OK(proxy)
 
 
+def getPilotToken(pilotDict):
+    """Get a token corresponding to the pilot
+
+    :param dict pilotDict: pilot parameters
+    :return: S_OK/S_ERROR with token as Value
+    """
+    ownerDN = pilotDict["OwnerDN"]
+    group = pilotDict["OwnerGroup"]
+
+    result = getUsernameForDN(ownerDN)
+    if not result["OK"]:
+        return result
+    username = result["Value"]
+    result = gTokenManager.getToken(
+        username=username,
+        userGroup=group,
+        requiredTimeLeft=3600,
+    )
+    return result
+
+
+def setPilotCredentials(ce, pilotDict):
+    """Instrument the given CE with proxy or token
+
+    :param obj ce:  CE object
+    :param pilotDict: pilot parameter dictionary
+    :return: S_OK/S_ERROR
+    """
+    if "Token" in ce.ceParameters.get("Tag", []):
+        result = getPilotToken(pilotDict)
+        if not result["OK"]:
+            return result
+        ce.setToken(result["Value"], 3500)
+    else:
+        result = getPilotProxy(pilotDict)
+        if not result["OK"]:
+            return result
+        ce.setProxy(result["Value"])
+    return S_OK()
+
+
 def getPilotRef(pilotReference, pilotDict):
     """Add the pilotStamp to the pilotReference, if the pilotStamp is in the dictionary,
     otherwise return unchanged pilotReference.