non sleep solution, with test

Author: khsrali

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

@@ -61,6 +61,7 @@ def __init__(self, authinfo: AuthInfo, transport_queue: 'TransportQueue', last_u
         self._job_update_requests: Dict[Hashable, asyncio.Future] = {}  # Mapping: {job_id: Future}
         self._last_updated = last_updated
         self._update_handle: Optional[asyncio.TimerHandle] = None
+        self._inspecting_jobs: List[str] = []
 
     @property
     def logger(self) -> logging.Logger:
@@ -101,10 +102,11 @@ async def _get_jobs_from_scheduler(self) -> Dict[Hashable, 'JobInfo']:
             scheduler.set_transport(transport)
 
             kwargs: Dict[str, Any] = {'as_dict': True}
+            self._inspecting_jobs = self._get_jobs_with_scheduler()
             if scheduler.get_feature('can_query_by_user'):
                 kwargs['user'] = '$USER'
             else:
-                kwargs['jobs'] = self._get_jobs_with_scheduler()
+                kwargs['jobs'] = self._inspecting_jobs
 
             scheduler_response = scheduler.get_jobs(**kwargs)
 
@@ -124,6 +126,7 @@ async def _update_job_info(self) -> None:
         This will set the futures for all pending update requests where the corresponding job has a new status compared
         to the last update.
         """
+        racing_requests = {}
         try:
             if not self._update_requests_outstanding():
                 return
@@ -146,9 +149,12 @@ async def _update_job_info(self) -> None:
         else:
             for job_id, future in self._job_update_requests.items():
                 if not future.done():
-                    future.set_result(self._jobs_cache.get(job_id, None))
+                    if str(job_id) in self._inspecting_jobs:
+                        future.set_result(self._jobs_cache.get(job_id, None))
+                    else:
+                        racing_requests[job_id] = future
         finally:
-            self._job_update_requests = {}
+            self._job_update_requests = racing_requests
 
     @contextlib.contextmanager
     def request_job_info_update(self, authinfo: AuthInfo, job_id: Hashable) -> Iterator['asyncio.Future[JobInfo]']:

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

@@ -209,15 +209,6 @@ async def do_update():
     try:
         logger.info(f'scheduled request to update CalcJob<{node.pk}>')
         ignore_exceptions = (plumpy.futures.CancelledError, plumpy.process_states.Interruption)
-
-        if not node.get_last_job_info():
-            # One can pass this value from the scheduler,
-            # so we sleep only for slurm. And only for the first time?
-            # scheduler = node.computer.get_scheduler()
-            # sleeptim e= scheduler._FIRST_FETCH_SLEEP
-            # Some Schedulers take some time to show the job status
-            await asyncio.sleep(5)
-
         job_done = await utils.exponential_backoff_retry(
             do_update, initial_interval, max_attempts, logger=node.logger, ignore_exceptions=ignore_exceptions
         )

tests/engine/processes/calcjobs/test_calc_job.py

@@ -1225,7 +1225,7 @@ def test_monitor_result_parse(get_calcjob_builder, entry_points):
     entry_points.add(monitor_skip_parse, group='aiida.calculations.monitors', name='core.skip_parse')
 
     builder = get_calcjob_builder()
-    builder.metadata.options.sleep = 8
+    builder.metadata.options.sleep = 3
     builder.monitors = {'monitor': orm.Dict({'entry_point': 'core.skip_parse'})}
     _, node = launch.run_get_node(builder)
     assert sorted(node.outputs) == ['remote_folder', 'retrieved']
@@ -1245,7 +1245,7 @@ def test_monitor_result_retrieve(get_calcjob_builder, entry_points):
     entry_points.add(monitor_skip_retrieve, group='aiida.calculations.monitors', name='core.skip_retrieval')
 
     builder = get_calcjob_builder()
-    builder.metadata.options.sleep = 8
+    builder.metadata.options.sleep = 3
     builder.monitors = {'monitor': orm.Dict({'entry_point': 'core.skip_retrieval'})}
     _, node = launch.run_get_node(builder)
     assert 'retrieved' not in node.outputs
@@ -1265,7 +1265,7 @@ def test_monitor_result_override_exit_code(get_calcjob_builder, entry_points):
     entry_points.add(monitor_override_exit_code, group='aiida.calculations.monitors', name='core.override_exit_code')
 
     builder = get_calcjob_builder()
-    builder.metadata.options.sleep = 8
+    builder.metadata.options.sleep = 3
     builder.monitors = {'monitor': orm.Dict({'entry_point': 'core.override_exit_code'})}
     _, node = launch.run_get_node(builder)
     assert sorted(node.outputs) == ['remote_folder', 'retrieved']
@@ -1315,7 +1315,7 @@ def test_monitor_result_action_disable_self(get_calcjob_builder, entry_points, c
     entry_points.add(monitor_disable_self, group='aiida.calculations.monitors', name='core.disable_self')
 
     builder = get_calcjob_builder()
-    builder.metadata.options.sleep = 6
+    builder.metadata.options.sleep = 1
     builder.monitors = {'disable_self': orm.Dict({'entry_point': 'core.disable_self'})}
     _, node = launch.run_get_node(builder)
     assert node.is_finished_ok

tests/engine/processes/calcjobs/test_monitors.py

@@ -198,7 +198,7 @@ def test_calc_job_monitors_process_poll_interval_integrated(entry_points, aiida_
     builder.x = Int(1)
     builder.y = Int(1)
     builder.monitors = {'always_kill': Dict({'entry_point': 'core.emit_warning', 'minimum_poll_interval': 5})}
-    builder.metadata = {'options': {'sleep': 6, 'resources': {'num_machines': 1}}}
+    builder.metadata = {'options': {'sleep': 1, 'resources': {'num_machines': 1}}}
 
     _, node = run_get_node(builder)
     assert node.is_finished_ok
@@ -218,7 +218,7 @@ def test_calc_job_monitors_outputs(entry_points, aiida_code_installed):
     builder.x = Int(1)
     builder.y = Int(1)
     builder.monitors = {'store_message': Dict({'entry_point': 'core.store_message', 'minimum_poll_interval': 1})}
-    builder.metadata = {'options': {'sleep': 8, 'resources': {'num_machines': 1}}}
+    builder.metadata = {'options': {'sleep': 3, 'resources': {'num_machines': 1}}}
 
     _, node = run_get_node(builder)
     assert node.is_finished_ok

tests/engine/test_manager.py

@@ -71,3 +71,47 @@ def test_last_updated(self):
         last_updated = time.time()
         jobs_list = JobsList(self.auth_info, self.transport_queue, last_updated=last_updated)
         assert jobs_list.last_updated == last_updated
+
+    def test_prevent_racing_condition(self):
+        """Test that the `JobsList` prevents racing condition when updating job info.
+
+        This test simulates a race condition where:
+        1. Job 'job1' requests an update
+        2. During the scheduler query, a new job 'job2' also requests an update
+        3. JobList must only update about 'job1'
+        4. 'job2' future should be kept pending for the next update cycle
+        """
+        from unittest.mock import patch
+
+        from aiida.schedulers.datastructures import JobInfo, JobState
+
+        jobs_list = self.jobs_list
+
+        mock_job_info = JobInfo()
+        mock_job_info.job_id = 'job1'
+        mock_job_info.job_state = JobState.RUNNING
+
+        def mock_get_jobs(**kwargs):
+            # Simulate the race: job2 is added to _job_update_requests while we're querying the scheduler
+            jobs_list._job_update_requests.setdefault('job2', asyncio.Future())
+
+            # Return only job1 (scheduler was queried with only job1)
+            return {'job1': mock_job_info}
+
+        # Request update for job1
+        future1 = jobs_list._job_update_requests.setdefault('job1', asyncio.Future())
+
+        # Patch the scheduler's get_jobs
+        scheduler = self.auth_info.computer.get_scheduler()
+        with patch.object(scheduler.__class__, 'get_jobs', side_effect=mock_get_jobs):
+            self.loop.run_until_complete(jobs_list._update_job_info())
+
+        # Verify job1 was resolved correctly
+        assert future1.done(), 'job1 future should be resolved'
+        assert future1.result() == mock_job_info, 'job1 should have the correct JobInfo'
+
+        # Verify job2 was NOT resolved and it has remained in _job_update_requests for the next cycle
+        assert 'job2' in jobs_list._job_update_requests, 'job2 should still be in update requests'
+        future2 = jobs_list._job_update_requests['job2']
+        assert not future2.done(), 'job2 future should NOT be resolved yet (prevented racing bug)'
+        assert len(jobs_list._job_update_requests) == 1, 'Only job2 should remain in update requests'