Keep transport open for quick tasks

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

@@ -72,7 +72,7 @@ async def task_upload_job(process: 'CalcJob', transport_queue: TransportQueue, c
     node = process.node
 
     if node.get_state() == CalcJobState.SUBMITTING:
-        logger.warning(f'CalcJob<{node.pk}> already marked as SUBMITTING, skipping task_update_job')
+        logger.warning(f'CalcJob<{node.pk}> already marked as SUBMITTING, skipping task_upload_job')
         return
 
     initial_interval = get_config_option(RETRY_INTERVAL_OPTION)

src/aiida/engine/transports.py

@@ -13,8 +13,10 @@
 import contextvars
 import logging
 import traceback
+from datetime import datetime
 from typing import TYPE_CHECKING, Awaitable, Dict, Hashable, Iterator, Optional
 
+from aiida.common import timezone
 from aiida.orm import AuthInfo
 
 if TYPE_CHECKING:
@@ -27,7 +29,6 @@ class TransportRequest:
     """Information kept about request for a transport object"""
 
     def __init__(self):
-        super().__init__()
         self.future: asyncio.Future = asyncio.Future()
         self.count = 0
 
@@ -47,6 +48,7 @@ def __init__(self, loop: Optional[asyncio.AbstractEventLoop] = None):
         """:param loop: An asyncio event, will use `asyncio.get_event_loop()` if not supplied"""
         self._loop = loop if loop is not None else asyncio.get_event_loop()
         self._transport_requests: Dict[Hashable, TransportRequest] = {}
+        self._last_close_time: Optional[datetime] = None
 
     @property
     def loop(self) -> asyncio.AbstractEventLoop:
@@ -99,7 +101,23 @@ def do_open():
             # 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())
+
+            # First request, submit immediately
+            if self._last_close_time is None:
+                open_callback_handle = self._loop.call_soon(do_open, context=contextvars.Context())
+
+            else:
+                close_timedelta = (timezone.localtime(timezone.now()) - self._last_close_time).total_seconds()
+
+                if close_timedelta > safe_open_interval:
+                    # If time since last close > `safe_open_interval`, open immediately
+                    open_callback_handle = self._loop.call_soon(do_open, context=contextvars.Context())
+
+                else:
+                    # Otherwise, wait only the difference required until the `safe_open_interval` is reached
+                    open_callback_handle = self._loop.call_later(
+                        safe_open_interval - close_timedelta, do_open, context=contextvars.Context()
+                    )
 
         try:
             transport_request.count += 1
@@ -120,7 +138,9 @@ def do_open():
                 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._last_close_time = timezone.localtime(timezone.now())
                 self._transport_requests.pop(authinfo.pk, None)