Skip to content

Commit

Permalink
Remove outer usage of runner.coordinator
Browse files Browse the repository at this point in the history
  • Loading branch information
unkcpz committed Dec 28, 2024
1 parent 1857656 commit 01f92f7
Show file tree
Hide file tree
Showing 2 changed files with 8 additions and 12 deletions.
5 changes: 3 additions & 2 deletions src/aiida/engine/daemon/worker.py
Original file line number Diff line number Diff line change
Expand Up @@ -63,8 +63,9 @@ def create_daemon_runner(manager: Manager) -> 'Runner':
loader=persistence.get_object_loader(),
)

assert runner.coordinator is not None, 'coordinator not set for runner'
runner.coordinator.add_task_subscriber(task_receiver)
coordinator = manager.get_coordinator()
assert coordinator is not None, 'coordinator not set for runner'
coordinator.add_task_subscriber(task_receiver)

return runner

Expand Down
15 changes: 5 additions & 10 deletions src/aiida/engine/runners.py
Original file line number Diff line number Diff line change
Expand Up @@ -124,11 +124,6 @@ def persister(self) -> Optional[Persister]:
"""Get the persister used by this runner."""
return self._persister

@property
def coordinator(self) -> Optional[Coordinator]:
"""Get the coordinator used by this runner."""
return self._coordinator

@property
def plugin_version_provider(self) -> PluginVersionProvider:
return self._plugin_version_provider
Expand Down Expand Up @@ -337,16 +332,16 @@ def inline_callback(event, *args, **kwargs):
callback()
finally:
event.set()
if self.coordinator:
self.coordinator.remove_broadcast_subscriber(subscriber_identifier)
if self._coordinator:
self._coordinator.remove_broadcast_subscriber(subscriber_identifier)

broadcast_filter = kiwipy.BroadcastFilter(functools.partial(inline_callback, event), sender=pk)
for state in [ProcessState.FINISHED, ProcessState.KILLED, ProcessState.EXCEPTED]:
broadcast_filter.add_subject_filter(f'state_changed.*.{state.value}')

if self.coordinator:
if self._coordinator:
LOGGER.info('adding subscriber for broadcasts of %d', pk)
self.coordinator.add_broadcast_subscriber(broadcast_filter, subscriber_identifier)
self._coordinator.add_broadcast_subscriber(broadcast_filter, subscriber_identifier)
self._poll_process(node, functools.partial(inline_callback, event))

def get_process_future(self, pk: int) -> futures.ProcessFuture:
Expand All @@ -356,7 +351,7 @@ def get_process_future(self, pk: int) -> futures.ProcessFuture:
:return: A future representing the completion of the process node
"""
return futures.ProcessFuture(pk, self._loop, self._poll_interval, self.coordinator)
return futures.ProcessFuture(pk, self._loop, self._poll_interval, self._coordinator)

def _poll_process(self, node, callback):
"""Check whether the process state of the node is terminated and call the callback or reschedule it.
Expand Down

0 comments on commit 01f92f7

Please sign in to comment.