diff --git a/openhexa/sdk/pipelines/pipeline.py b/openhexa/sdk/pipelines/pipeline.py index 34795053..f37233a4 100644 --- a/openhexa/sdk/pipelines/pipeline.py +++ b/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]: diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index f586e652..f41fce03 100644 --- a/tests/test_pipeline.py +++ b/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 = [