def create_authinfo(computer, store=False):
"""Allow the current user to use the given computer."""
from aiida.orm import AuthInfo
authinfo = AuthInfo(computer=computer, user=get_current_user())
if store:
authinfo.store()
return authinfo
def get_authinfo(computer):
"""Get existing authinfo or create one if not in place"""
try:
authinfo = computer.get_authinfo(get_current_user())
except NotExistent:
authinfo = AuthInfo(computer=computer, user=get_current_user())
authinfo.store()
return authinfo
def setUp(self):
super().setUp()
self.loop = asyncio.get_event_loop()
self.transport_queue = TransportQueue(self.loop)
self.user = User.objects.get_default()
self.auth_info = AuthInfo(self.computer, self.user).store()
self.jobs_list = JobsList(self.auth_info, self.transport_queue)
def setUp(self):
super(TestJobsList, self).setUp()
self.loop = tornado.ioloop.IOLoop()
self.transport_queue = TransportQueue(self.loop)
self.user = User.objects.get_default()
self.auth_info = AuthInfo(self.computer, self.user).store()
self.jobs_list = JobsList(self.auth_info, self.transport_queue)
def setUp(self):
super().setUp()
self.loop = tornado.ioloop.IOLoop()
self.transport_queue = TransportQueue(self.loop)
self.user = User.objects.get_default()
self.auth_info = AuthInfo(self.computer, self.user).store()
self.manager = JobManager(self.transport_queue)
def _configure_computer(self):
"""create AuthInfo"""
print("Configuring '{}'".format(self.name))
sshcfg = parse_sshconfig(self.hostname)
authparams = {
'compress': True,
'gss_auth': False,
'gss_deleg_creds': False,
'gss_host': self.hostname,
'gss_kex': False,
'key_policy': 'WarningPolicy',
'load_system_host_keys': True,
'port': 22,
'timeout': 60,
}
if 'user' in sshcfg:
authparams['username'] = sshcfg['user']
else:
print(
"SSH username is not provided, please run `verdi computer configure {}` "
"from the command line".format(self.name))
return
if 'proxycommand' in sshcfg:
authparams['proxy_command'] = sshcfg['proxycommand']
aiidauser = User.objects.get_default()
from aiida.orm import AuthInfo
authinfo = AuthInfo(computer=Computer.objects.get(name=self.name),
user=aiidauser)
authinfo.set_auth_params(authparams)
authinfo.store()
print(
check_output(['verdi', 'computer', 'show',
self.name]).decode('utf-8'))
def _configure_computer(self):
"""Create AuthInfo."""
print("Configuring '{}'".format(self.label))
sshcfg = parse_sshconfig(self.hostname)
authparams = {
'compress':
True,
'key_filename':
os.path.expanduser(
sshcfg.get('identityfile', ['~/.ssh/id_rsa'])[0]),
'gss_auth':
False,
'gss_deleg_creds':
False,
'gss_host':
self.hostname,
'gss_kex':
False,
'key_policy':
'WarningPolicy',
'load_system_host_keys':
True,
'port':
sshcfg.get('port', 22),
'timeout':
60,
'use_login_shell':
self._use_login_shell.value,
}
if 'user' in sshcfg:
authparams['username'] = sshcfg['user']
else:
print(
f"SSH username is not provided, please run `verdi computer configure {self.label}` "
"from the command line.")
return
if 'proxycommand' in sshcfg:
authparams['proxy_command'] = sshcfg['proxycommand']
aiidauser = User.objects.get_default()
from aiida.orm import AuthInfo
authinfo = AuthInfo(computer=Computer.objects.get(name=self.label),
user=aiidauser)
authinfo.set_auth_params(authparams)
authinfo.store()
print(
check_output(['verdi', 'computer', 'show',
self.label]).decode('utf-8'))
def _configure_computer(self):
"""Create AuthInfo."""
print(f"Configuring '{self.label}'")
sshcfg = parse_sshconfig(self.hostname)
authparams = {
"compress": True,
"key_filename": os.path.expanduser(
sshcfg.get("identityfile", ["~/.ssh/id_rsa"])[0]
),
"gss_auth": False,
"gss_deleg_creds": False,
"gss_host": self.hostname,
"gss_kex": False,
"key_policy": "WarningPolicy",
"load_system_host_keys": True,
"port": sshcfg.get("port", 22),
"timeout": 60,
"use_login_shell": self._use_login_shell.value,
"safe_interval": self.safe_interval,
}
if "user" in sshcfg:
authparams["username"] = sshcfg["user"]
else:
print(
f"SSH username is not provided, please run `verdi computer configure {self.label}` "
"from the command line."
)
return
if "proxycommand" in sshcfg:
authparams["proxy_command"] = sshcfg["proxycommand"]
aiidauser = User.objects.get_default()
from aiida.orm import AuthInfo
authinfo = AuthInfo(
computer=Computer.objects.get(label=self.label), user=aiidauser
)
authinfo.set_auth_params(authparams)
authinfo.store()
print(check_output(["verdi", "computer", "show", self.label]).decode("utf-8"))
def request_transport(
self, authinfo: AuthInfo) -> Iterator[Awaitable[Transport]]:
"""
Request a transport from an authinfo. Because the client is not allowed to
request a transport immediately they will instead be given back a future
that can be awaited to get the transport::
async def transport_task(transport_queue, authinfo):
with transport_queue.request_transport(authinfo) as request:
transport = await request
# Do some work with the transport
:param authinfo: The authinfo to be used to get transport
:return: A future that can be yielded to give the transport
"""
open_callback_handle = None
transport_request = self._transport_requests.get(authinfo.id, None)
if transport_request is None:
# There is no existing request for this transport (i.e. on this authinfo)
transport_request = TransportRequest()
self._transport_requests[authinfo.id] = transport_request
transport = authinfo.get_transport()
safe_open_interval = transport.get_safe_open_interval()
def do_open():
""" Actually open the transport """
if transport_request and transport_request.count > 0:
# The user still wants the transport so open it
_LOGGER.debug('Transport request opening transport for %s',
authinfo)
try:
transport.open()
except Exception as exception: # pylint: disable=broad-except
_LOGGER.error(
'exception occurred while trying to open transport:\n %s',
exception)
transport_request.future.set_exception(exception)
# Cleanup of the stale TransportRequest with the excepted transport future
self._transport_requests.pop(authinfo.id, None)
else:
transport_request.future.set_result(transport)
# Save the handle so that we can cancel the callback if the user no longer wants it
# Note: Don't pass the Process context, since (a) it is not needed by `do_open` and (b) the transport is
# passed around to many places, including outside aiida-core (e.g. paramiko). Anyone keeping a reference
# to this handle would otherwise keep the Process context (and thus the process itself) in memory.
# See https://github.com/aiidateam/aiida-core/issues/4698
open_callback_handle = self._loop.call_later(
safe_open_interval, do_open,
context=contextvars.Context()) # type: ignore[call-arg]
try:
transport_request.count += 1
yield transport_request.future
except asyncio.CancelledError: # pylint: disable=try-except-raise
# note this is only required in python<=3.7,
# where asyncio.CancelledError inherits from Exception
_LOGGER.debug('Transport task cancelled')
raise
except Exception:
_LOGGER.error('Exception whilst using transport:\n%s',
traceback.format_exc())
raise
finally:
transport_request.count -= 1
assert transport_request.count >= 0, 'Transport request count dropped below 0!'
# Check if there are no longer any users that want the transport
if transport_request.count == 0:
if transport_request.future.done():
_LOGGER.debug('Transport request closing transport for %s',
authinfo)
transport_request.future.result().close()
elif open_callback_handle is not None:
open_callback_handle.cancel()
self._transport_requests.pop(authinfo.id, None)
def setUpClass(cls):
super(TestRemoteData, cls).setUpClass()
user = User.objects.get_default()
authinfo = AuthInfo(cls.computer, user)
authinfo.store()
def setUpClass(cls): # pylint: disable=arguments-differ
super().setUpClass()
user = User.objects.get_default()
authinfo = AuthInfo(cls.computer, user)
authinfo.store()