Transport & Engine: AsyncTransport plugin

.readthedocs.yml

@@ -20,7 +20,7 @@ build:
     - asdf install uv 0.2.9
     - asdf global uv 0.2.9
     post_install:
-    - VIRTUAL_ENV=$READTHEDOCS_VIRTUALENV_PATH uv pip install .[docs,tests,rest,atomic_tools]
+    - VIRTUAL_ENV=$READTHEDOCS_VIRTUALENV_PATH uv pip install .[docs,tests,rest,atomic_tools] --preview
 
 # Let the build fail if there are any warnings
 sphinx:

docs/source/topics/transport.rst

@@ -24,15 +24,15 @@ The generic transport class contains a set of minimal methods that an implementa
 If not, a ``NotImplementedError`` will be raised, interrupting the managing of the calculation or whatever is using the transport plugin.
 
 As for the general functioning of the plugin, the :py:meth:`~aiida.transports.transport.Transport.__init__` method is used only to initialize the class instance, without actually opening the transport channel.
-The connection must be opened only by the :py:meth:`~aiida.transports.transport.Transport.__enter__` method, (and closed by :py:meth:`~aiida.transports.transport.Transport.__exit__`).
-The :py:meth:`~aiida.transports.transport.Transport.__enter__` method lets you use the transport class using the ``with`` statement (see `python docs <https://docs.python.org/3/reference/compound_stmts.html#with>`_), in a way similar to the following:
+The connection must be opened only by the :py:meth:`~aiida.transports.transport._BaseTransport.__enter__` method, (and closed by :py:meth:`~aiida.transports.transport._BaseTransport.__exit__`).
+The :py:meth:`~aiida.transports.transport._BaseTransport.__enter__` method lets you use the transport class using the ``with`` statement (see `python docs <https://docs.python.org/3/reference/compound_stmts.html#with>`_), in a way similar to the following:
 
 .. code-block:: python
 
     with TransportPlugin() as transport:
         transport.some_method()
 
-To ensure this, for example, the local plugin uses a hidden boolean variable ``_is_open`` that is set when the :py:meth:`~aiida.transports.transport.Transport.__enter__` and :py:meth:`~aiida.transports.transport.Transport.__exit__` methods are called.
+To ensure this, for example, the local plugin uses a hidden boolean variable ``_is_open`` that is set when the :py:meth:`~aiida.transports.transport._BaseTransport.__enter__` and :py:meth:`~aiida.transports.transport._BaseTransport.__exit__` methods are called.
 The ``ssh`` logic is instead given by the property sftp.
 
 The other functions that require some care are the copying functions, called using the following terminology:

environment.yml

@@ -8,6 +8,7 @@ dependencies:
 - python~=3.9
 - alembic~=1.2
 - archive-path~=0.4.2
+- asyncssh~=2.19.0
 - circus~=0.18.0
 - click-spinner~=0.1.8
 - click~=8.1
@@ -22,7 +23,7 @@ dependencies:
 - importlib-metadata~=6.0
 - numpy~=1.21
 - paramiko~=3.0
-- plumpy~=0.22.3
+- plumpy~=0.23.0
 - pgsu~=0.3.0
 - psutil~=5.6
 - psycopg[binary]~=3.0

pyproject.toml

@@ -20,6 +20,7 @@ classifiers = [
 dependencies = [
   'alembic~=1.2',
   'archive-path~=0.4.2',
+  "asyncssh~=2.19.0",
   'circus~=0.18.0',
   'click-spinner~=0.1.8',
   'click~=8.1',
@@ -34,7 +35,7 @@ dependencies = [
   'importlib-metadata~=6.0',
   'numpy~=1.21',
   'paramiko~=3.0',
-  'plumpy~=0.22.3',
+  'plumpy~=0.23.0',
   'pgsu~=0.3.0',
   'psutil~=5.6',
   'psycopg[binary]~=3.0',
@@ -175,6 +176,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_async' = 'aiida.transports.plugins.ssh_async:AsyncSshTransport'
 'core.ssh_auto' = 'aiida.transports.plugins.ssh_auto:SshAutoTransport'
 
 [project.entry-points.'aiida.workflows']
@@ -308,6 +310,7 @@ module = 'tests.*'
 ignore_missing_imports = true
 module = [
   'ase.*',
+  'asyncssh.*',
   'bpython.*',
   'bs4.*',
   'CifFile.*',

src/aiida/calculations/monitors/base.py

@@ -4,12 +4,13 @@
 
 import tempfile
 from pathlib import Path
+from typing import Union
 
 from aiida.orm import CalcJobNode
-from aiida.transports import Transport
+from aiida.transports import AsyncTransport, Transport
 
 
-def always_kill(node: CalcJobNode, transport: Transport) -> str | None:
+def always_kill(node: CalcJobNode, transport: Union['Transport', 'AsyncTransport']) -> str | None:
     """Retrieve and inspect files in working directory of job to determine whether the job should be killed.
 
     This particular implementation is just for demonstration purposes and will kill the job as long as there is a

src/aiida/engine/processes/calcjobs/calcjob.py

@@ -524,7 +524,7 @@ def on_terminated(self) -> None:
         super().on_terminated()
 
     @override
-    def run(self) -> Union[plumpy.process_states.Stop, int, plumpy.process_states.Wait]:
+    async def run(self) -> Union[plumpy.process_states.Stop, int, plumpy.process_states.Wait]:
         """Run the calculation job.
 
         This means invoking the `presubmit` and storing the temporary folder in the node's repository. Then we move the
@@ -535,11 +535,11 @@ def run(self) -> Union[plumpy.process_states.Stop, int, plumpy.process_states.Wa
 
         """
         if self.inputs.metadata.dry_run:
-            self._perform_dry_run()
+            await self._perform_dry_run()
             return plumpy.process_states.Stop(None, True)
 
         if 'remote_folder' in self.inputs:
-            exit_code = self._perform_import()
+            exit_code = await self._perform_import()
             return exit_code
 
         # The following conditional is required for the caching to properly work. Even if the source node has a process
@@ -627,7 +627,7 @@ def _setup_inputs(self) -> None:
         if not self.node.computer:
             self.node.computer = self.inputs.code.computer
 
-    def _perform_dry_run(self):
+    async def _perform_dry_run(self):
         """Perform a dry run.
 
         Instead of performing the normal sequence of steps, just the `presubmit` is called, which will call the method
@@ -643,13 +643,13 @@ def _perform_dry_run(self):
         with LocalTransport() as transport:
             with SubmitTestFolder() as folder:
                 calc_info = self.presubmit(folder)
-                upload_calculation(self.node, transport, calc_info, folder, inputs=self.inputs, dry_run=True)
+                await upload_calculation(self.node, transport, calc_info, folder, inputs=self.inputs, dry_run=True)
                 self.node.dry_run_info = {  # type: ignore[attr-defined]
                     'folder': folder.abspath,
                     'script_filename': self.node.get_option('submit_script_filename'),
                 }
 
-    def _perform_import(self):
+    async def _perform_import(self):
         """Perform the import of an already completed calculation.
 
         The inputs contained a `RemoteData` under the key `remote_folder` signalling that this is not supposed to be run
@@ -669,7 +669,7 @@ def _perform_import(self):
                 with SandboxFolder(filepath_sandbox) as retrieved_temporary_folder:
                     self.presubmit(folder)
                     self.node.set_remote_workdir(self.inputs.remote_folder.get_remote_path())
-                    retrieved = retrieve_calculation(self.node, transport, retrieved_temporary_folder.abspath)
+                    retrieved = await retrieve_calculation(self.node, transport, retrieved_temporary_folder.abspath)
                     if retrieved is not None:
                         self.out(self.node.link_label_retrieved, retrieved)
                         self.update_outputs()

src/aiida/engine/processes/calcjobs/monitors.py

@@ -8,14 +8,15 @@
 import inspect
 import typing as t
 from datetime import datetime, timedelta
+from typing import Union
 
 from aiida.common.lang import type_check
 from aiida.common.log import AIIDA_LOGGER
 from aiida.orm import CalcJobNode, Dict
 from aiida.plugins import BaseFactory
 
 if t.TYPE_CHECKING:
-    from aiida.transports import Transport
+    from aiida.transports import AsyncTransport, Transport
 
 LOGGER = AIIDA_LOGGER.getChild(__name__)
 
@@ -122,7 +123,9 @@
         parameters = list(signature.parameters.keys())
 
         if any(required_parameter not in parameters for required_parameter in ('node', 'transport')):
-            correct_signature = '(node: CalcJobNode, transport: Transport, **kwargs) str | None:'
+            correct_signature = (
+                "(node: CalcJobNode, transport: Union['Transport', 'AsyncTransport'], **kwargs) str | None:"
+            )
             raise ValueError(
                 f'The monitor `{self.entry_point}` has an invalid function signature, it should be: {correct_signature}'
             )
@@ -176,7 +179,7 @@
     def process(
         self,
         node: CalcJobNode,
-        transport: Transport,
+        transport: Union['Transport', 'AsyncTransport'],
     ) -> CalcJobMonitorResult | None:
         """Call all monitors in order and return the result as one returns anything other than ``None``.
 

src/aiida/engine/processes/calcjobs/tasks.py

@@ -92,7 +92,7 @@
                 except Exception as exception:
                     raise PreSubmitException('exception occurred in presubmit call') from exception
                 else:
-                    remote_folder = execmanager.upload_calculation(node, transport, calc_info, folder)
+                    remote_folder = await execmanager.upload_calculation(node, transport, calc_info, folder)
                     if remote_folder is not None:
                         process.out('remote_folder', remote_folder)
                     skip_submit = calc_info.skip_submit or False
@@ -314,7 +314,7 @@
 
             if node.get_job_id() is None:
                 logger.warning(f'there is no job id for CalcJobNoe<{node.pk}>: skipping `get_detailed_job_info`')
-                retrieved = execmanager.retrieve_calculation(node, transport, retrieved_temporary_folder)
+                retrieved = await execmanager.retrieve_calculation(node, transport, retrieved_temporary_folder)
             else:
                 try:
                     detailed_job_info = scheduler.get_detailed_job_info(node.get_job_id())
@@ -324,7 +324,7 @@
                 else:
                     node.set_detailed_job_info(detailed_job_info)
 
-                retrieved = execmanager.retrieve_calculation(node, transport, retrieved_temporary_folder)
+                retrieved = await execmanager.retrieve_calculation(node, transport, retrieved_temporary_folder)
 
             if retrieved is not None:
                 process.out(node.link_label_retrieved, retrieved)
@@ -376,7 +376,7 @@
             transport = await cancellable.with_interrupt(request)
 
             logger.info(f'stashing calculation<{node.pk}>')
-            return execmanager.stash_calculation(node, transport)
+            return await execmanager.stash_calculation(node, transport)
 
     try:
         await exponential_backoff_retry(

src/aiida/engine/processes/functions.py

@@ -567,7 +567,7 @@ def _setup_db_record(self) -> None:
         self.node.store_source_info(self._func)
 
     @override
-    def run(self) -> 'ExitCode' | None:
+    async def run(self) -> 'ExitCode' | None:
         """Run the process."""
         from .exit_code import ExitCode
 

src/aiida/engine/processes/workchains/workchain.py

@@ -297,7 +297,7 @@ def _update_process_status(self) -> None:
 
     @override
     @Protect.final
-    def run(self) -> t.Any:
+    async def run(self) -> t.Any:
         self._stepper = self.spec().get_outline().create_stepper(self)  # type: ignore[arg-type]
         return self._do_step()
 

src/aiida/engine/transports.py

@@ -13,12 +13,12 @@
 import contextvars
 import logging
 import traceback
-from typing import TYPE_CHECKING, Awaitable, Dict, Hashable, Iterator, Optional
+from typing import TYPE_CHECKING, Awaitable, Dict, Hashable, Iterator, Optional, Union
 
 from aiida.orm import AuthInfo
 
 if TYPE_CHECKING:
-    from aiida.transports import Transport
+    from aiida.transports import AsyncTransport, Transport
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -54,7 +54,7 @@
         return self._loop
 
     @contextlib.contextmanager
-    def request_transport(self, authinfo: AuthInfo) -> Iterator[Awaitable['Transport']]:
+    def request_transport(self, authinfo: AuthInfo) -> Iterator[Awaitable[Union['Transport', 'AsyncTransport']]]:
         """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::

src/aiida/orm/authinfos.py

@@ -8,7 +8,7 @@
 ###########################################################################
 """Module for the `AuthInfo` ORM class."""
 
-from typing import TYPE_CHECKING, Any, Dict, Optional, Type
+from typing import TYPE_CHECKING, Any, Dict, Optional, Type, Union
 
 from aiida.common import exceptions
 from aiida.manage import get_manager
@@ -21,7 +21,7 @@
     from aiida.orm import Computer, User
     from aiida.orm.implementation import StorageBackend
     from aiida.orm.implementation.authinfos import BackendAuthInfo  # noqa: F401
-    from aiida.transports import Transport
+    from aiida.transports import AsyncTransport, Transport
 
 __all__ = ('AuthInfo',)
 
@@ -166,7 +166,7 @@
         except KeyError:
             return self.computer.get_workdir()
 
-    def get_transport(self) -> 'Transport':
+    def get_transport(self) -> Union['Transport', 'AsyncTransport']:
         """Return a fully configured transport that can be used to connect to the computer set for this instance."""
         computer = self.computer
         transport_type = computer.transport_type

src/aiida/orm/computers.py

@@ -23,7 +23,7 @@
     from aiida.orm import AuthInfo, User
     from aiida.orm.implementation import StorageBackend
     from aiida.schedulers import Scheduler
-    from aiida.transports import Transport
+    from aiida.transports import AsyncTransport, Transport
 
 __all__ = ('Computer',)
 
@@ -622,16 +622,16 @@
             # Return False if the user is not configured (in a sense, it is disabled for that user)
             return False
 
-    def get_transport(self, user: Optional['User'] = None) -> 'Transport':
+    def get_transport(self, user: Optional['User'] = None) -> Union['Transport', 'AsyncTransport']:
         """Return a Transport class, configured with all correct parameters.
         The Transport is closed (meaning that if you want to run any operation with
         it, you have to open it first (i.e., e.g. for a SSH transport, you have
-        to open a connection). To do this you can call ``transports.open()``, or simply
+        to open a connection). To do this you can call ``transport.open()``, or simply
         run within a ``with`` statement::
 
            transport = Computer.get_transport()
            with transport:
-               print(transports.whoami())
+               print(transport.whoami())
 
         :param user: if None, try to obtain a transport for the default user.
             Otherwise, pass a valid User.
@@ -646,7 +646,7 @@
         authinfo = authinfos.AuthInfo.get_collection(self.backend).get(dbcomputer=self, aiidauser=user)
         return authinfo.get_transport()
 
-    def get_transport_class(self) -> Type['Transport']:
+    def get_transport_class(self) -> Union[Type['Transport'], Type['AsyncTransport']]:
         """Get the transport class for this computer.  Can be used to instantiate a transport instance."""
         try:
             return TransportFactory(self.transport_type)

src/aiida/orm/nodes/data/remote/base.py

@@ -125,7 +125,8 @@ def listdir_withattributes(self, path='.'):
         """Connects to the remote folder and lists the directory content.
 
         :param relpath: If 'relpath' is specified, lists the content of the given subfolder.
-        :return: a list of dictionaries, where the documentation is in :py:class:Transport.listdir_withattributes.
+        :return: a list of dictionaries, where the documentation
+            is in :py:class:Transport.listdir_withattributes.
         """
         authinfo = self.get_authinfo()
 

src/aiida/orm/nodes/process/calculation/calcjob.py

@@ -26,7 +26,7 @@
     from aiida.parsers import Parser
     from aiida.schedulers.datastructures import JobInfo, JobState
     from aiida.tools.calculations import CalculationTools
-    from aiida.transports import Transport
+    from aiida.transports import AsyncTransport, Transport
 
 __all__ = ('CalcJobNode',)
 
@@ -450,10 +450,11 @@
 
         return computer.get_authinfo(self.user)
 
-    def get_transport(self) -> 'Transport':
+    def get_transport(self) -> Union['Transport', 'AsyncTransport']:
         """Return the transport for this calculation.
 
-        :return: `Transport` configured with the `AuthInfo` associated to the computer of this node
+        :return: Union['Transport', 'AsyncTransport'] configured
+            with the `AuthInfo` associated to the computer of this node
         """
         return self.get_authinfo().get_transport()
 

src/aiida/orm/utils/remote.py

@@ -12,6 +12,7 @@
 
 import os
 import typing as t
+from typing import Union
 
 from aiida.orm.nodes.data.remote.base import RemoteData
 
@@ -20,14 +21,14 @@
 
     from aiida import orm
     from aiida.orm.implementation import StorageBackend
-    from aiida.transports import Transport
+    from aiida.transports import AsyncTransport, Transport
 
 
-def clean_remote(transport: Transport, path: str) -> None:
+def clean_remote(transport: Union['Transport', 'AsyncTransport'], path: str) -> None:
     """Recursively remove a remote folder, with the given absolute path, and all its contents. The path should be
     made accessible through the transport channel, which should already be open
 
-    :param transport: an open Transport channel
+    :param transport: an open Union['Transport', 'AsyncTransport'] channel
     :param path: an absolute path on the remote made available through the transport
     """
     if not isinstance(path, str):