Make the scheduler interface asynchronous

docs/source/howto/cookbook.rst

@@ -7,88 +7,6 @@ Cookbook
 This how-to page collects useful short scripts and code snippets that may be useful in the everyday usage of AiiDA.
 
 
-Checking the queued jobs on a scheduler
-=======================================
-
-If you want to know if which jobs are currently on the scheduler (e.g. to dynamically decide on which computer to submit, or to delay submission, etc.) you can use the following script as an example:
-
-.. code-block:: python
-
-    def get_scheduler_jobs(computer_label='localhost', only_current_user=True):
-        """Return a list of all current jobs in the scheduler.
-
-        .. note:: an SSH connection is open and closed at every launch of this function.
-
-        :param computer_label: the label of the computer.
-        :param only_current_user: if True, only retrieve jobs of the current default user.
-            (if this feature is supported by the scheduler plugin). Otherwise show all jobs.
-        """
-        from aiida import orm
-
-        computer = Computer.collection.get(label=computer_label)
-        transport = computer.get_transport()
-        scheduler = computer.get_scheduler()
-        scheduler.set_transport(transport)
-
-        # This opens the SSH connection, for SSH transports
-        with transport:
-            if only_current_user:
-                remote_username = transport.whoami()
-                all_jobs = scheduler.get_jobs(user=remote_username, as_dict=True)
-            else:
-                all_jobs = scheduler.get_jobs(as_dict=True)
-
-        return all_jobs
-
-    if __name__ == '__main__':
-        all_jobs = get_scheduler_jobs(only_current_user=False)
-        user_jobs = get_scheduler_jobs(only_current_user=True)
-
-        print(f'Current user has {len(user_jobs)} jobs out of {len(all_jobs)} in the scheduler'
-        print('Detailed job view:')
-
-        for job_id, job_info in user_jobs.items():
-            print(f'Job ID: {job_id}')
-            for k, v in job_info.items():
-                if k == 'raw_data':
-                    continue
-                print(f'  {k}: {v}')
-            print('')
-
-Use ``verdi run`` to execute it:
-
-.. code-block:: console
-
-    verdi run file_with_script.py
-
-.. important::
-
-    Every time you call the function, two SSH connections are opened!
-    So be careful and run this function sparsely, or your supercomputer center might block your account.
-    A possible work around to this limitation is to pass the transport as a parameter, and pass it in so that it can be reused.
-
-An example output would be::
-
-    Current user has 5 jobs out of 1425 in the scheduler
-    Detailed job view:
-    Job ID: 1658497
-        job_id: 1658497
-        wallclock_time_seconds: 38052
-        title: aiida-2324985
-        num_machines: 4
-        job_state: RUNNING
-        queue_name: parallel
-        num_mpiprocs: 64
-        allocated_machines_raw: r02-node[17-18,53-54]
-        submission_time: 2018-03-28 09:21:35
-        job_owner: some_remote_username
-        dispatch_time: 2018-03-28 09:21:35
-        annotation: None
-        requested_wallclock_time_seconds: 82800
-
-    (...)
-
-
 Getting an ``AuthInfo`` knowing the computer and the user
 =========================================================
 

src/aiida/cmdline/commands/cmd_computer.py

@@ -57,7 +57,14 @@ def _computer_test_get_jobs(transport, scheduler, authinfo, computer):
     :param authinfo: the AuthInfo object (from which one can get computer and aiidauser)
     :return: tuple of boolean indicating success or failure and an optional string message
     """
-    found_jobs = scheduler.get_jobs(as_dict=True)
+
+    from plumpy import run_until_complete
+
+    from aiida.manage import get_manager
+
+    loop = get_manager().get_runner().loop
+    found_jobs = run_until_complete(loop, scheduler.get_jobs_async(as_dict=True))
+
     return True, f'{len(found_jobs)} jobs found in the queue'
 
 

src/aiida/engine/daemon/execmanager.py

@@ -398,7 +398,7 @@ async def _copy_sandbox_files(logger, node, transport, folder, workdir: Path):
         await transport.put_async(folder.get_abs_path(filename), workdir.joinpath(filename))
 
 
-def submit_calculation(calculation: CalcJobNode, transport: Transport) -> str | ExitCode:
+async def submit_calculation(calculation: CalcJobNode, transport: Transport) -> str | ExitCode:
     """Submit a previously uploaded `CalcJob` to the scheduler.
 
     :param calculation: the instance of CalcJobNode to submit.
@@ -420,7 +420,7 @@ def submit_calculation(calculation: CalcJobNode, transport: Transport) -> str |
 
     submit_script_filename = calculation.get_option('submit_script_filename')
     workdir = calculation.get_remote_workdir()
-    result = scheduler.submit_job(workdir, submit_script_filename)
+    result = await scheduler.submit_job_async(workdir, submit_script_filename)
 
     if isinstance(result, str):
         calculation.set_job_id(result)
@@ -791,7 +791,7 @@ async def retrieve_calculation(
     return retrieved_files
 
 
-def kill_calculation(calculation: CalcJobNode, transport: Transport) -> None:
+async def kill_calculation(calculation: CalcJobNode, transport: Transport) -> None:
     """Kill the calculation through the scheduler
 
     :param calculation: the instance of CalcJobNode to kill.
@@ -808,19 +808,19 @@ def kill_calculation(calculation: CalcJobNode, transport: Transport) -> None:
     scheduler.set_transport(transport)
 
     # Call the proper kill method for the job ID of this calculation
-    result = scheduler.kill_job(job_id)
+    result = await scheduler.kill_job_async(job_id)
 
     if result is not True:
         # Failed to kill because the job might have already been completed
-        running_jobs = scheduler.get_jobs(jobs=[job_id], as_dict=True)
+        running_jobs = await scheduler.get_jobs_async(jobs=[job_id], as_dict=True)
         job = running_jobs.get(job_id, None)
 
         # If the job is returned it is still running and the kill really failed, so we raise
         if job is not None and job.job_state != JobState.DONE:
-            raise exceptions.RemoteOperationError(f'scheduler.kill_job({job_id}) was unsuccessful')
+            raise exceptions.RemoteOperationError(f'scheduler.kill_job_async({job_id}) was unsuccessful')
         else:
             EXEC_LOGGER.warning(
-                'scheduler.kill_job() failed but job<{%s}> no longer seems to be running regardless', job_id
+                'scheduler.kill_job_async() failed but job<{%s}> no longer seems to be running regardless', job_id
             )
 
 

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

@@ -106,9 +106,9 @@ async def _get_jobs_from_scheduler(self) -> Dict[str, 'JobInfo']:
             self._polling_jobs = frozenset([str(job_id) for job_id in self._job_update_requests.keys()])
 
             if scheduler.get_feature('can_query_by_user'):
-                scheduler_response = scheduler.get_jobs(user='$USER', as_dict=True)
+                scheduler_response = await scheduler.get_jobs_async(user='$USER', as_dict=True)
             else:
-                scheduler_response = scheduler.get_jobs(jobs=list(self._polling_jobs), as_dict=True)
+                scheduler_response = await scheduler.get_jobs_async(jobs=list(self._polling_jobs), as_dict=True)
 
             # Update the last update time and clear the jobs cache
             self._last_updated = time.time()

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

@@ -146,7 +146,7 @@ async def task_submit_job(node: CalcJobNode, transport_queue: TransportQueue, ca
     async def do_submit():
         async with transport_queue.request_transport(authinfo) as request:
             transport = await cancellable.with_interrupt(request)
-            return execmanager.submit_calculation(node, transport)
+            return await execmanager.submit_calculation(node, transport)
 
     try:
         logger.info(f'scheduled request to submit CalcJob<{node.pk}>')
@@ -312,7 +312,7 @@ async def do_retrieve():
                 retrieved = await execmanager.retrieve_calculation(node, transport, retrieved_temporary_folder)
             else:
                 try:
-                    detailed_job_info = scheduler.get_detailed_job_info(job_id)
+                    detailed_job_info = await scheduler.get_detailed_job_info_async(job_id)
                 except FeatureNotAvailable:
                     logger.info(f'detailed job info not available for scheduler of CalcJob<{node.pk}>')
                     node.set_detailed_job_info(None)
@@ -456,7 +456,7 @@ async def task_kill_job(node: CalcJobNode, transport_queue: TransportQueue, canc
     async def do_kill():
         async with transport_queue.request_transport(authinfo) as request:
             transport = await cancellable.with_interrupt(request)
-            return execmanager.kill_calculation(node, transport)
+            return await execmanager.kill_calculation(node, transport)
 
     try:
         logger.info(f'scheduled request to kill CalcJob<{node.pk}>')

src/aiida/engine/transports.py

@@ -69,15 +69,6 @@ async def transport_task(transport_queue, authinfo):
         :param authinfo: The authinfo to be used to get transport
         :return: A future that can be yielded to give the transport
         """
-
-        from plumpy import ensure_portal
-
-        # NOTE: We need to ensure the portal here only because
-        # our scheduler has only a sync interface and _get_jobs_from_scheduler is using that
-        # if we ever provide a fully async scheduler interface then we can remove this here
-        # An issue is opened to reference this https://github.com/aiidateam/aiida-core/issues/7222
-        await ensure_portal()
-
         open_callback_handle = None
         transport_request = self._transport_requests.get(authinfo.pk, None)
 

src/aiida/schedulers/plugins/bash.py

@@ -24,34 +24,34 @@
 class BashCliScheduler(Scheduler, metaclass=abc.ABCMeta):
     """Job scheduler that is interacted with through a CLI in bash."""
 
-    def submit_job(self, working_directory: str, filename: str) -> str | ExitCode:
+    async def submit_job_async(self, working_directory: str, filename: str) -> str | ExitCode:
         """Submit a job.
 
         :param working_directory: The absolute filepath to the working directory where the job is to be executed.
         :param filename: The filename of the submission script relative to the working directory.
         """
-        result = self.transport.exec_command_wait(
+        result = await self.transport.exec_command_wait_async(
             self._get_submit_command(escape_for_bash(filename)), workdir=working_directory
         )
         return self._parse_submit_output(*result)
 
     @t.overload
-    def get_jobs(
+    async def get_jobs_async(
         self,
         jobs: list[str] | None = None,
         user: str | None = None,
         as_dict: t.Literal[False] = False,
     ) -> list[JobInfo]: ...
 
     @t.overload
-    def get_jobs(
+    async def get_jobs_async(
         self,
         jobs: list[str] | None = None,
         user: str | None = None,
         as_dict: t.Literal[True] = ...,
     ) -> dict[str, JobInfo]: ...
 
-    def get_jobs(
+    async def get_jobs_async(
         self,
         jobs: list[str] | None = None,
         user: str | None = None,
@@ -65,8 +65,10 @@ def get_jobs(
             returned, where the ``job_id`` is the key and the values are the ``JobInfo`` objects.
         :returns: List of active jobs.
         """
-        with self.transport:
-            retval, stdout, stderr = self.transport.exec_command_wait(self._get_joblist_command(jobs=jobs, user=user))
+        async with self.transport:
+            retval, stdout, stderr = await self.transport.exec_command_wait_async(
+                self._get_joblist_command(jobs=jobs, user=user)
+            )
 
         joblist = self._parse_joblist_output(retval, stdout, stderr)
         if as_dict:
@@ -77,7 +79,7 @@ def get_jobs(
 
         return joblist
 
-    def kill_job(self, jobid: str) -> bool:
+    async def kill_job_async(self, jobid: str) -> bool:
         """Kill a remote job and parse the return value of the scheduler to check if the command succeeded.
 
         ..note::
@@ -88,7 +90,7 @@ def kill_job(self, jobid: str) -> bool:
         :param jobid: the job ID to be killed
         :returns: True if everything seems ok, False otherwise.
         """
-        retval, stdout, stderr = self.transport.exec_command_wait(self._get_kill_command(jobid))
+        retval, stdout, stderr = await self.transport.exec_command_wait_async(self._get_kill_command(jobid))
         return self._parse_kill_output(retval, stdout, stderr)
 
     @abc.abstractmethod

src/aiida/schedulers/plugins/direct.py

@@ -270,23 +270,23 @@ def _parse_joblist_output(self, retval: int, stdout: str, stderr: str) -> list[J
         return job_list
 
     @t.overload
-    def get_jobs(
+    async def get_jobs_async(
         self,
         jobs: list[str] | None = None,
         user: str | None = None,
         as_dict: t.Literal[False] = False,
     ) -> list[JobInfo]: ...
 
     @t.overload
-    def get_jobs(
+    async def get_jobs_async(
         self,
         jobs: list[str] | None = None,
         user: str | None = None,
         as_dict: t.Literal[True] = ...,
     ) -> dict[str, JobInfo]: ...
 
     @override
-    def get_jobs(
+    async def get_jobs_async(
         self,
         jobs: list[str] | None = None,
         user: str | None = None,
@@ -295,7 +295,7 @@ def get_jobs(
         """Overrides original method from BashScheduler in order to list
         missing processes as DONE.
         """
-        job_stats = super().get_jobs(jobs=jobs, user=user, as_dict=True)
+        job_stats = await super().get_jobs_async(jobs=jobs, user=user, as_dict=True)
 
         # Get the list of known jobs
         found_jobs = job_stats.keys()

src/aiida/schedulers/scheduler.py

@@ -126,32 +126,31 @@ def create_job_resource(cls, **kwargs: t.Any) -> JobResource:
         return cls._job_resource_class(**kwargs)
 
     @abc.abstractmethod
-    def submit_job(self, working_directory: str, filename: str) -> str | ExitCode:
+    async def submit_job_async(self, working_directory: str, filename: str) -> str | ExitCode:
         """Submit a job.
 
         :param working_directory: The absolute filepath to the working directory where the job is to be executed.
         :param filename: The filename of the submission script relative to the working directory.
-        :returns:
         """
 
     @t.overload
-    def get_jobs(
+    async def get_jobs_async(
         self,
         jobs: list[str] | None = None,
         user: str | None = None,
         as_dict: t.Literal[False] = False,
     ) -> list[JobInfo]: ...
 
     @t.overload
-    def get_jobs(
+    async def get_jobs_async(
         self,
         jobs: list[str] | None = None,
         user: str | None = None,
         as_dict: t.Literal[True] = ...,
     ) -> dict[str, JobInfo]: ...
 
     @abc.abstractmethod
-    def get_jobs(
+    async def get_jobs_async(
         self,
         jobs: list[str] | None = None,
         user: str | None = None,
@@ -167,7 +166,7 @@ def get_jobs(
         """
 
     @abc.abstractmethod
-    def kill_job(self, jobid: str) -> bool:
+    async def kill_job_async(self, jobid: str) -> bool:
         """Kill a remote job and parse the return value of the scheduler to check if the command succeeded.
 
         ..note::
@@ -352,7 +351,7 @@ def _get_detailed_job_info_command(self, job_id: str) -> str:
         """
         raise exceptions.FeatureNotAvailable('Cannot get detailed job info')
 
-    def get_detailed_job_info(self, job_id: str) -> dict[str, str | int]:
+    async def get_detailed_job_info_async(self, job_id: str) -> dict[str, str | int]:
         """Return the detailed job info.
 
         This will be a dictionary with the return value, stderr and stdout content returned by calling the command that
@@ -362,7 +361,7 @@ def get_detailed_job_info(self, job_id: str) -> dict[str, str | int]:
         :return: dictionary with `retval`, `stdout` and `stderr`.
         """
         command = self._get_detailed_job_info_command(job_id)
-        retval, stdout, stderr = self.transport.exec_command_wait(command)
+        retval, stdout, stderr = await self.transport.exec_command_wait_async(command)
 
         detailed_job_info = {
             'retval': retval,

tests/cmdline/commands/test_computer.py

@@ -968,7 +968,7 @@ def test_computer_test_stderr(run_cli_command, aiida_localhost, monkeypatch):
     aiida_localhost.configure()
     stderr = 'spurious output in standard error'
 
-    def exec_command_wait(self, command, **kwargs):
+    def exec_command_wait(self, command, *args, **kwargs):
         return 0, '', stderr
 
     monkeypatch.setattr(LocalTransport, 'exec_command_wait', exec_command_wait)
@@ -985,7 +985,7 @@ def test_computer_test_stdout(run_cli_command, aiida_localhost, monkeypatch):
     aiida_localhost.configure()
     stdout = 'spurious output in standard output'
 
-    def exec_command_wait(self, command, **kwargs):
+    def exec_command_wait(self, command, *args, **kwargs):
         return 0, stdout, ''
 
     monkeypatch.setattr(LocalTransport, 'exec_command_wait', exec_command_wait)