Skip to content

Commit

Permalink
sweep: DIRACGrid#7913 Add option to include proxy on AREX token submi…
Browse files Browse the repository at this point in the history
…ssion
  • Loading branch information
sfayer authored and fstagni committed Nov 29, 2024
1 parent 8d7e680 commit 8508d7d
Showing 1 changed file with 22 additions and 8 deletions.
30 changes: 22 additions & 8 deletions src/DIRAC/Resources/Computing/AREXComputingElement.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,24 +36,29 @@
The XRSLExtraString note about times also applies to this configuration option.
AlwaysIncludeProxy:
A boolean, set to true to include the proxy in job submission even
in cases where tokens are the primary authentication method.
(Recommended for ARC6 tokens, deprecated for ARC7)
**Code Documentation**
"""

import os
import json
import requests
import os
import shutil
import stat
import uuid

from DIRAC import S_OK, S_ERROR
import requests

from DIRAC import S_ERROR, S_OK
from DIRAC.Core.Security import Locations
from DIRAC.Core.Security.ProxyInfo import getVOfromProxyGroup
from DIRAC.Core.Security.X509Chain import X509Chain # pylint: disable=import-error
from DIRAC.Resources.Computing.ComputingElement import ComputingElement
from DIRAC.WorkloadManagementSystem.Client import PilotStatus
from DIRAC.Resources.Computing.PilotBundle import writeScript

from DIRAC.WorkloadManagementSystem.Client import PilotStatus

MANDATORY_PARAMETERS = ["Queue"]

Expand Down Expand Up @@ -115,6 +120,8 @@ def __init__(self, ceUniqueID):
}
# URL used to communicate with the REST interface
self.base_url = ""
# A flag to always include a proxy, even if a token is the primary auth method
self.alwaysIncludeProxy = False

#############################################################################

Expand Down Expand Up @@ -149,6 +156,10 @@ def _reset(self):
service_url = os.path.join("https://", f"{self.ceName}:{self.port}")
self.base_url = os.path.join(service_url, "arex", "rest", self.restVersion)

self.alwaysIncludeProxy = False
if self.ceParameters.get("AlwaysIncludeProxy", "false").lower() in ("true", "yes"):
self.alwaysIncludeProxy = True

# Set up the request framework
self.session = requests.Session()
self.session.verify = Locations.getCAsLocation()
Expand Down Expand Up @@ -247,13 +258,16 @@ def _checkSession(self):
if not (self.token or self.proxy):
self.log.error("Proxy or token not set")
return S_ERROR("Proxy or token not set")
if not self.proxy and self.alwaysIncludeProxy:
self.log.error("Proxy required but not set")
return S_ERROR("Proxy required but not set")

# If a token is set, we use it
if self.token:
# Attach the token to the headers if present
self.headers["Authorization"] = f"Bearer {self.token['access_token']}"
self.log.verbose("A token is attached to the header of the request(s)")
else:
if not self.token or self.alwaysIncludeProxy:
# Prepare the proxy in X509_USER_PROXY
if not (result := self._prepareProxy())["OK"]:
self.log.error("Failed to set up proxy", result["Message"])
Expand Down Expand Up @@ -573,7 +587,7 @@ def submitJob(self, executableFile, proxy, numberOfJobs=1, inputs=None, outputs=

# Delegation cannot be used with a token
delegation = ""
if not self.token:
if not self.token or self.alwaysIncludeProxy:
# Get existing delegations
result = self._getDelegationIDs()
if not result["OK"]:
Expand Down Expand Up @@ -918,7 +932,7 @@ def getJobStatus(self, jobIDList):
self.log.debug(f"Killing held job {jobReference}")

# Renew delegations to renew the proxies of the jobs
if not self.token:
if not self.token or self.alwaysIncludeProxy:
result = self._getDelegationIDs()
if not result["OK"]:
return result
Expand Down

0 comments on commit 8508d7d

Please sign in to comment.