fix: make SDK tasks debuggable in same process (HEXA-1833)

openhexa/sdk/pipelines/pipeline.py

@@ -98,10 +98,14 @@ def run(self, config: dict[str, typing.Any]):
         with heartbeat_manager(current_run, interval=30):
             # Execute pipeline function
             self.function(**validated_config)
-            # Execute tasks using Pool's built-in context manager
-            context = get_context("spawn")
-            with context.Pool() as pool:  # FIXME: set max size of pool
-                self._execute_tasks(pool)
+            if _should_run_tasks_in_process():
+                # Debuggers can't follow tasks into spawned subprocesses
+                self._execute_tasks_in_process()
+            else:
+                # Execute tasks using Pool's built-in context manager
+                context = get_context("spawn")
+                with context.Pool() as pool:  # FIXME: set max size of pool
+                    self._execute_tasks(pool)
 
         print(f'{get_timestamp()} Successfully completed pipeline "{self.name}"')
 
@@ -218,6 +222,29 @@ def _execute_tasks(self, pool):
                     # busy loop
                     time.sleep(0.3)
 
+    def _execute_tasks_in_process(self):
+        """Execute all tasks sequentially in the current process, in dependency order.
+
+        Raises
+        ------
+        PipelineRunError
+            If any task fails during execution.
+        """
+        total = len(self.tasks)
+        completed = 0
+
+        while tasks := self._get_available_tasks():
+            for task in tasks:
+                print(f'{get_timestamp()} Started task "{task.compute.__name__}"')
+                try:
+                    task.run()
+                except Exception as e:  # NOQA
+                    raise PipelineRunError(f"Pipeline {self.name} failed: {e}") from e
+
+                completed += 1
+                print(f'{get_timestamp()} Finished task "{task.compute.__name__}"')
+                self._update_progress(int(completed / total * 100))
+
     def to_dict(self):
         """Return a dictionary representation of the pipeline."""
         return {
@@ -301,6 +328,19 @@ def __call__(self, config: dict[str | typing.Any] | None = None):
         self.run(config)
 
 
+def _should_run_tasks_in_process() -> bool:
+    """Determine whether tasks should run in the current process instead of a multiprocessing pool.
+
+    This is the case when a debugger is in use, as debuggers do not follow tasks into the spawned worker processes:
+    either when running with "openhexa pipelines run --debug", or when a debugger (VS Code / debugpy, PyCharm /
+    pydevd) is attached to the current process.
+    """
+    if os.environ.get("REMOTE_DEBUGGER", "").strip().lower() == "true":
+        return True
+
+    return any(module in sys.modules for module in ("debugpy", "pydevd"))
+
+
 def pipeline(
     code: str = None, name: str = None, timeout: int = None, functional_type: str = None
 ) -> typing.Callable[[typing.Callable[..., typing.Any]], Pipeline]:

tests/test_pipeline.py

@@ -1,6 +1,7 @@
 """Pipeline test module."""
 
 import os
+import sys
 import time
 from unittest import TestCase
 from unittest.mock import Mock, patch
@@ -17,7 +18,7 @@
 from openhexa.sdk.pipelines.heartbeat import HeartbeatThread
 from openhexa.sdk.pipelines.log_level import LogLevel
 from openhexa.sdk.pipelines.parameter import Parameter, ParameterValueError
-from openhexa.sdk.pipelines.pipeline import Pipeline
+from openhexa.sdk.pipelines.pipeline import Pipeline, PipelineRunError, _should_run_tasks_in_process
 from openhexa.sdk.utils import Environment
 
 
@@ -335,6 +336,97 @@ def fast_init(self, run_context, interval=30):
     ), "Heartbeat should have been attempted multiple times despite failures"
 
 
+@patch.dict(os.environ, {"REMOTE_DEBUGGER": "true"})
+def test_pipeline_run_tasks_in_process():
+    """In debug mode, tasks are executed in the current process, in dependency order."""
+    calls = []
+
+    def pipeline_func():
+        task_3(task_1(), task_2())
+
+    pipeline = Pipeline("pipeline", pipeline_func, [])
+
+    @pipeline.task
+    def task_1():
+        calls.append("task_1")
+        return 1
+
+    @pipeline.task
+    def task_2():
+        calls.append("task_2")
+        return 2
+
+    @pipeline.task
+    def task_3(a, b):
+        calls.append("task_3")
+        return a + b
+
+    pipeline.run({})
+
+    assert calls[-1] == "task_3"
+    assert sorted(calls[:2]) == ["task_1", "task_2"]
+    assert pipeline.tasks[-1].result == 3
+
+
+@patch.dict(os.environ, {"REMOTE_DEBUGGER": "true"})
+def test_pipeline_run_tasks_in_process_failure():
+    """In debug mode, a failing task raises a PipelineRunError chained to the original exception."""
+
+    def pipeline_func():
+        failing_task()
+
+    pipeline = Pipeline("pipeline", pipeline_func, [])
+
+    @pipeline.task
+    def failing_task():
+        raise ValueError("boom")
+
+    with pytest.raises(PipelineRunError, match="boom") as exc_info:
+        pipeline.run({})
+    assert isinstance(exc_info.value.__cause__, ValueError)
+
+
+@pytest.mark.parametrize(
+    "env, debugger_module, expected",
+    [
+        ({}, None, False),
+        ({"REMOTE_DEBUGGER": "false"}, None, False),
+        ({"REMOTE_DEBUGGER": "true"}, None, True),
+        ({}, "debugpy", True),
+        ({}, "pydevd", True),
+    ],
+)
+def test_should_run_tasks_in_process(env, debugger_module, expected):
+    """Tasks run in process only when a debugger is in use."""
+    with patch.dict(os.environ, env, clear=True), patch.dict(sys.modules):
+        sys.modules.pop("debugpy", None)
+        sys.modules.pop("pydevd", None)
+        if debugger_module:
+            sys.modules[debugger_module] = Mock()
+        assert _should_run_tasks_in_process() is expected
+
+
+@patch("openhexa.sdk.pipelines.pipeline.get_context")
+def test_pipeline_run_tasks_in_process_when_debugger_attached(mock_get_context):
+    """When a debugger is attached, tasks run in the current process instead of a multiprocessing pool."""
+    task_pids = []
+
+    def pipeline_func():
+        task_1()
+
+    pipeline = Pipeline("pipeline", pipeline_func, [])
+
+    @pipeline.task
+    def task_1():
+        task_pids.append(os.getpid())
+
+    with patch.dict(os.environ, {}, clear=True), patch.dict(sys.modules, {"debugpy": Mock()}):
+        pipeline.run({})
+
+    mock_get_context.assert_not_called()
+    assert task_pids == [os.getpid()]
+
+
 class TestLogLevel(TestCase):
     def test_parse_log_level(self):
         test_cases = [