Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add the SshFabricTransport plugin #6154

Merged
merged 1 commit into from
Jul 23, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -174,6 +174,7 @@ requires-python = '>=3.9'
[project.entry-points.'aiida.transports']
'core.local' = 'aiida.transports.plugins.local:LocalTransport'
'core.ssh' = 'aiida.transports.plugins.ssh:SshTransport'
'core.ssh_auto' = 'aiida.transports.plugins.ssh_auto:SshAutoTransport'

[project.entry-points.'aiida.workflows']
'core.arithmetic.add_multiply' = 'aiida.workflows.arithmetic.add_multiply:add_multiply'
Expand Down
61 changes: 61 additions & 0 deletions src/aiida/transports/plugins/ssh_auto.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,61 @@
###########################################################################
# Copyright (c), The AiiDA team. All rights reserved. #
# This file is part of the AiiDA code. #
# #
# The code is hosted on GitHub at https://github.com/aiidateam/aiida-core #
# For further information on the license, see the LICENSE.txt file #
# For further information please visit http://www.aiida.net #
###########################################################################
"""Plugin for transport over SSH (and SFTP for file transfer)."""

import pathlib

import paramiko

from .ssh import SshTransport

__all__ = ('SshAutoTransport',)


class SshAutoTransport(SshTransport):
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

As mentioned in one of my comments above, I'm wondering if we should name the plugin SshConfigTransport (entry point ssh_config), since it takes the information from the ~/.ssh/config.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I see your point of it corresponding to the ~/.ssh/config file, but I don't think it is as obvious to a user. If a user knows nothing about the way these plugins work, but they have the option between core.ssh and core.ssh_auto, I think it is safe to say that they will pick the latter if they want the easy experience. In contrast, if the options are core.ssh and core.ssh_config I think they would guess that core.ssh_config would require more configuration. That is why I would vote for core.ssh_auto and just have the docstring explain that it automatically parses the config file. This also makes it less tied to exactly what it is parsing automatically if in the future we want to add more sources of information.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hmm, fair enough! Happy to keep auto then. :)

"""Support connection, command execution and data transfer to remote computers via SSH+SFTP."""

_valid_connect_params = []
_valid_auth_options = []

FILEPATH_CONFIG: pathlib.Path = pathlib.Path('~').expanduser() / '.ssh' / 'config'

def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs, key_policy='AutoAddPolicy')

config_ssh = paramiko.SSHConfig()

try:
with self.FILEPATH_CONFIG.open() as handle:
config_ssh.parse(handle)
except FileNotFoundError as exception:
raise RuntimeError(

Check warning on line 37 in src/aiida/transports/plugins/ssh_auto.py

View check run for this annotation

Codecov / codecov/patch

src/aiida/transports/plugins/ssh_auto.py#L36-L37

Added lines #L36 - L37 were not covered by tests
f'Could not determine connection configuration as the `{self.FILEPATH_CONFIG}` does not exist.'
) from exception
except PermissionError as exception:
raise RuntimeError(

Check warning on line 41 in src/aiida/transports/plugins/ssh_auto.py

View check run for this annotation

Codecov / codecov/patch

src/aiida/transports/plugins/ssh_auto.py#L40-L41

Added lines #L40 - L41 were not covered by tests
f'Could not determine connection configuration as the `{self.FILEPATH_CONFIG}` is not readable.'
) from exception

if self._machine not in config_ssh.get_hostnames():
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is there any case where the connection would still work when the host is not in the ~/.ssh/config? Is there a reason to not outright raise here?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

My thinking: it might. And since it will raise in that case anyway, why not give it a shot?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

One consideration might be that the user might focus more on the (somewhat unclear) error and miss the very clear warning.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fair enough. One more reason why I decided to use a warning is that the ~/.ssh/config is not the only place that SSH connections can be configured. There are also system wide configs like etc/sshd/* and others. These are also checked by fabric which is why it failed silently if the specified hostname is not part of the ~/.ssh/config file, because it cannot conclude that there is a problem if that is the case. We are adding this because for our purposes in most cases the user's config file would be the main place and we want to warn explicitly if we don't find the host there, but it is not necessarily a problem. Anyway, if others think it is better to make this an exception for now, happy to change it. We can always remove it later and turn it into a warning if there are cases where the plugin should rely on other configuration files as well.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's also not critical for me, and I see your point. In fact, now that I think of it, I might already know of a user that had a setup I wasn't familiar with, which might work fine even though the remote isn't specified in the ~/.ssh/config. So let's keep the warning.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@mbercx did you already tried the core.ssh_auto with your various real-life test cases? Did it work?

Copy link
Member

@mbercx mbercx Jul 22, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah, yes! I tried 3 different clusters, with and without proxy jump. Also messed with the configuration to make it fail and then fixed it. All seems to work as expected.

So I think once we remove ssh_fabric we can go ahead and merge this, let it into the wild and see what happens. 😁

self.logger.warning(

Check warning on line 46 in src/aiida/transports/plugins/ssh_auto.py

View check run for this annotation

Codecov / codecov/patch

src/aiida/transports/plugins/ssh_auto.py#L46

Added line #L46 was not covered by tests
f'The host `{self._machine}` is not defined in SSH config, connection will most likely fail to open.'
)

config_host = config_ssh.lookup(self._machine)

self._connect_args = {
'port': config_host.get('port', 22),
'username': config_host.get('user'),
'key_filename': config_host.get('identityfile', []),
'timeout': config_host.get('connecttimeout', 60),
'proxy_command': config_host.get('proxycommand', None),
'proxy_jump': config_host.get('proxyjump', None),
}

self._machine = config_host['hostname']
14 changes: 14 additions & 0 deletions tests/cmdline/commands/test_computer.py
Original file line number Diff line number Diff line change
Expand Up @@ -971,3 +971,17 @@ def time_use_login_shell(authinfo, auth_params, use_login_shell, iterations) ->
result = run_cli_command(computer_test, [aiida_localhost.label], use_subprocess=False)
assert 'Success: all 6 tests succeeded' in result.output
assert 'computer is configured to use a login shell, which is slower compared to a normal shell' in result.output


def test_computer_ssh_auto(run_cli_command, aiida_computer):
"""Test setup of computer with ``core.ssh_auto`` entry point.

The configure step should only require the common shared options ``safe_interval`` and ``use_login_shell``.
"""
computer = aiida_computer(transport_type='core.ssh_auto').store()
assert not computer.is_configured

# It is important that no other options (except for `--safe-interval`) have to be specified for this transport type.
options = ['core.ssh_auto', computer.uuid, '--non-interactive', '--safe-interval', '0']
run_cli_command(computer_configure, options, use_subprocess=False)
assert computer.is_configured
3 changes: 1 addition & 2 deletions tests/engine/daemon/test_execmanager.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,6 @@
from aiida.common.folders import SandboxFolder
from aiida.engine.daemon import execmanager
from aiida.orm import CalcJobNode, FolderData, PortableCode, RemoteData, SinglefileData
from aiida.plugins import entry_point
from aiida.transports.plugins.local import LocalTransport


Expand All @@ -40,7 +39,7 @@ def file_hierarchy_simple():
}


@pytest.fixture(params=entry_point.get_entry_point_names('aiida.transports'))
@pytest.fixture(params=('core.local', 'core.ssh'))
def node_and_calc_info(aiida_localhost, aiida_computer_ssh, aiida_code_installed, request):
"""Return a ``CalcJobNode`` and associated ``CalcInfo`` instance."""

Expand Down
12 changes: 10 additions & 2 deletions tests/transports/test_all_plugins.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,14 +34,22 @@


@pytest.fixture(scope='function', params=entry_point.get_entry_point_names('aiida.transports'))
def custom_transport(request) -> Transport:
def custom_transport(request, tmp_path, monkeypatch) -> Transport:
"""Fixture that parametrizes over all the registered implementations of the ``CommonRelaxWorkChain``."""
plugin = TransportFactory(request.param)

if request.param == 'core.ssh':
kwargs = {'machine': 'localhost', 'timeout': 30, 'load_system_host_keys': True, 'key_policy': 'AutoAddPolicy'}
elif request.param == 'core.ssh_auto':
kwargs = {'machine': 'localhost'}
filepath_config = tmp_path / 'config'
monkeypatch.setattr(plugin, 'FILEPATH_CONFIG', filepath_config)
if not filepath_config.exists():
filepath_config.write_text('Host localhost')
else:
kwargs = {}

return TransportFactory(request.param)(**kwargs)
return plugin(**kwargs)


def test_is_open(custom_transport):
Expand Down