Replaced hardcoded print() with logging (#754)

Replaced hardcoded runtime print() calls in aiida-workgraph with logging/process reporting so embedding apps can control output without stdout pollution.

Author: superstar54

src/aiida_workgraph/engine/awaitable_manager.py

@@ -131,7 +131,7 @@ def on_awaitable_finished(self, awaitable: Awaitable) -> None:
         try:
             self.process.resume()
         except Exception as e:
-            print(e)
+            self.logger.exception('Failed to resume process after awaitable completion: %s', e)
 
     def to_context(self, **kwargs: Awaitable | ProcessNode) -> None:
         """Add a dictionary of awaitables to the context.

src/aiida_workgraph/engine/task_manager.py

@@ -134,7 +134,7 @@ def should_run_task(self, task: 'Task') -> bool:
         # skip if the max number of awaitables is reached
         if task.task_type.upper() in process_task_types:
             if len(self.process._awaitables) >= self.process.wg.max_number_jobs:
-                print(MAX_NUMBER_AWAITABLES_MSG.format(self.process.wg.max_number_jobs, name))
+                self.process.report(MAX_NUMBER_AWAITABLES_MSG.format(self.process.wg.max_number_jobs, name))
                 return False
         # skip if the task is already executed or if the task is in a skippped state
         if name in self.ctx._executed_tasks or self.state_manager.get_task_runtime_info(name, 'state') in ['SKIPPED']:

src/aiida_workgraph/executors/builtins.py

@@ -1,6 +1,9 @@
+import logging
 from typing import Any
 from aiida import orm
 
+LOGGER = logging.getLogger(__name__)
+
 
 def UnavailableExecutor(*args, **kwargs):
     raise RuntimeError('This executor was defined dynamically and is not available from the database snapshot.')
@@ -53,5 +56,5 @@ def load_code(pk: int = None, uuid: str = None, label: str = None) -> orm.Code:
         pk = label.value if isinstance(label, orm.Str) else label
     else:
         pk = pk.value if isinstance(pk, orm.Int) else pk
-    print(f'Loading code with pk: {pk}')
+    LOGGER.info('Loading code with pk: %s', pk)
     return orm.load_code(pk)

src/aiida_workgraph/tasks/monitors.py

@@ -1,8 +1,11 @@
 import datetime
 import typing as t
+import logging
 
 from aiida_workgraph import task
 
+LOGGER = logging.getLogger(__name__)
+
 
 @task.monitor
 def monitor_file(filepath: str):
@@ -49,8 +52,8 @@ def monitor_task(task_name: str, workgraph_pk: int = None, workgraph_name: str =
         )
         if builder.count() == 0:
             return False
-    print('Found workgraph')
+    LOGGER.debug('Found workgraph')
     node = builder.first()[0]
     state = node.task_states.get(task_name, '')
-    print(f'Task state: {state}')
+    LOGGER.debug('Task state: %s', state)
     return state in ['FINISHED', 'FAILED', 'SKIPPED']

src/aiida_workgraph/utils/__init__.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import logging
 from typing import Any, Dict, Optional, Union, Callable, List
 from aiida.engine.processes import Process
 from aiida import orm
@@ -18,6 +19,8 @@
 from aiida_workgraph.orm.utils import deserialize_safe
 from copy import deepcopy
 
+LOGGER = logging.getLogger(__name__)
+
 
 def inspect_aiida_component_type(executor: Callable) -> str:
     task_type = None
@@ -261,9 +264,7 @@ def load_workgraph_data(node: Union[int, orm.Node]) -> Optional[Dict[str, Any]]:
     try:
         task_inputs = deserialize_safe(node.task_inputs or '')
     except (yaml.constructor.ConstructorError, yaml.YAMLError):
-        print(
-            'Info: could not deserialize inputs.The workgraph is still loaded and you can inspect tasks and outputs. '
-        )
+        LOGGER.info('Could not deserialize inputs. The workgraph is still loaded and tasks/outputs remain inspectable.')
         task_inputs = {}
 
     for name, data in task_inputs.items():

src/aiida_workgraph/utils/control.py

@@ -1,9 +1,13 @@
 from __future__ import annotations
 
+import logging
+
 from aiida.manage import get_manager
 from aiida import orm
 from aiida.engine.processes import control
 
+LOGGER = logging.getLogger(__name__)
+
 
 def create_task_action(
     pk: int,
@@ -35,7 +39,7 @@ def pause_tasks(pk: int, tasks: list[str], timeout: int = 5):
     node = orm.load_node(pk)
     if node.is_finished:
         message = 'WorkGraph is finished. Cannot pause tasks.'
-        print(message)
+        LOGGER.warning(message)
         return False, message
     elif node.process_state.value.upper() in [
         'CREATED',
@@ -54,15 +58,15 @@ def pause_tasks(pk: int, tasks: list[str], timeout: int = 5):
                         timeout=timeout,
                     )
                 except Exception as e:
-                    print(f'Pause task {name} failed: {e}')
+                    LOGGER.exception('Pause task %s failed: %s', name, e)
     return True, ''
 
 
 def play_tasks(pk: int, tasks: list, timeout: int = 5):
     node = orm.load_node(pk)
     if node.is_finished:
         message = 'WorkGraph is finished. Cannot kill tasks.'
-        print(message)
+        LOGGER.warning(message)
         return False, message
     elif node.process_state.value.upper() in [
         'CREATED',
@@ -86,15 +90,15 @@ def play_tasks(pk: int, tasks: list, timeout: int = 5):
                         timeout=timeout,
                     )
                 except Exception as e:
-                    print(f'Play task {name} failed: {e}')
+                    LOGGER.exception('Play task %s failed: %s', name, e)
     return True, ''
 
 
 def kill_tasks(pk: int, tasks: list, timeout: int = 5):
     node = orm.load_node(pk)
     if node.is_finished:
         message = 'WorkGraph is finished. Cannot kill tasks.'
-        print(message)
+        LOGGER.warning(message)
         return False, message
     elif node.process_state.value.upper() in [
         'CREATED',
@@ -114,7 +118,7 @@ def kill_tasks(pk: int, tasks: list, timeout: int = 5):
                 'PAUSED',
             ]:
                 if process is None:
-                    print(f'Task {name} is not a AiiDA process.')
+                    LOGGER.warning('Task %s is not an AiiDA process.', name)
                     create_task_action(pk, tasks, action='kill')
                 else:
                     try:
@@ -124,7 +128,7 @@ def kill_tasks(pk: int, tasks: list, timeout: int = 5):
                             timeout=timeout,
                         )
                     except Exception as e:
-                        print(f'Kill task {name} failed: {e}')
+                        LOGGER.exception('Kill task %s failed: %s', name, e)
     return True, ''
 
 
@@ -136,7 +140,7 @@ def reset_tasks(pk: int, tasks: list) -> None:
     node = orm.load_node(pk)
     if node.is_finished:
         message = 'WorkGraph is finished. Cannot kill tasks.'
-        print(message)
+        LOGGER.warning(message)
         return False, message
     elif node.process_state.value.upper() in [
         'CREATED',

src/aiida_workgraph/workgraph.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import logging
 import node_graph
 import aiida
 from aiida_workgraph.task import Task
@@ -12,6 +13,8 @@
 from aiida_workgraph.socket_spec import SocketSpecAPI
 from node_graph.error_handler import ErrorHandlerSpec
 
+LOGGER = logging.getLogger(__name__)
+
 
 class WorkGraph(node_graph.Graph):
     """Build flexible workflows with AiiDA.
@@ -207,7 +210,7 @@ def save(self, metadata: Optional[Dict[str, Any]] = None) -> None:
             self.process = process_inited.node
             self.process_inited = process_inited
             process_inited.close()
-            print(f'WorkGraph process created, PK: {self.process.pk}')
+            LOGGER.info('WorkGraph process created, PK: %s', self.process.pk)
         else:
             self.save_to_base(inputs)
         self.update()
@@ -290,7 +293,7 @@ def wait(self, timeout: int = 600, tasks: dict = None, interval: int = 5) -> Non
                 finished = self.state in terminating_states
 
             if finished:
-                print(f'Process {self.process.pk} finished with state: {self.state}')
+                LOGGER.info('Process %s finished with state: %s', self.process.pk, self.state)
                 return
 
             time.sleep(interval)
@@ -470,7 +473,7 @@ def kill_tasks(self, tasks: List[str]) -> None:
     def reset_tasks(self, tasks: List[str]) -> None:
         from aiida_workgraph.utils.control import reset_tasks
 
-        print(f'Reset tasks: {tasks}')
+        LOGGER.info('Reset tasks: %s', tasks)
 
         if self.process is None:
             for name in tasks: