diff --git a/plexe/execution/training/local_runner.py b/plexe/execution/training/local_runner.py index de7944eb..bb584f10 100644 --- a/plexe/execution/training/local_runner.py +++ b/plexe/execution/training/local_runner.py @@ -8,7 +8,7 @@ import os import subprocess import sys -import time +import threading import uuid from pathlib import Path from typing import Any @@ -284,9 +284,20 @@ def run_training( env=env, ) - # Stream output in real-time while capturing it + # Stream output in real-time while capturing it. + # A watchdog enforces the timeout even when the subprocess is silent + # (data loading, quiet training loops), which reading the output pipe + # alone cannot do: readline() blocks until the next line arrives. stdout_lines = [] - start_time = time.time() + timed_out = threading.Event() + + def _kill_on_timeout(): + timed_out.set() + logger.warning(f"Training exceeded {timeout}s - killing process") + process.kill() + + watchdog = threading.Timer(timeout, _kill_on_timeout) + watchdog.start() try: for line in iter(process.stdout.readline, ""): @@ -296,17 +307,14 @@ def run_training( # Capture for error logging stdout_lines.append(line) - # Check timeout manually - elapsed = time.time() - start_time - if elapsed > timeout: - process.kill() - process.wait() - raise subprocess.TimeoutExpired(cmd, timeout) + if timed_out.is_set(): + raise subprocess.TimeoutExpired(cmd, timeout) # Wait for process to complete return_code = process.wait(timeout=5) # Short wait since process already finished finally: + watchdog.cancel() if process.stdout: process.stdout.close() diff --git a/tests/unit/execution/training/test_local_runner_timeout.py b/tests/unit/execution/training/test_local_runner_timeout.py new file mode 100644 index 00000000..7d6959b5 --- /dev/null +++ b/tests/unit/execution/training/test_local_runner_timeout.py @@ -0,0 +1,52 @@ +"""Tests for LocalProcessRunner training timeout enforcement.""" + +import threading +from unittest.mock import MagicMock, patch + +import pytest + +from plexe.execution.training.local_runner import LocalProcessRunner +from plexe.models import TrainingError + + +class _SilentProcess: + """Fake subprocess that produces no output until killed, like a quiet training run.""" + + def __init__(self): + self._killed = threading.Event() + self.stdout = self + self.kill_count = 0 + + def readline(self): + # Blocks like a pipe read on a silent child; returns EOF once killed. + self._killed.wait(timeout=60) + return "" + + def kill(self): + self.kill_count += 1 + self._killed.set() + + def wait(self, timeout=None): + return -9 + + def close(self): + pass + + +def test_timeout_enforced_when_process_is_silent(tmp_path): + """run_training must raise TrainingError after timeout even with no subprocess output.""" + proc = _SilentProcess() + runner = LocalProcessRunner(work_dir=str(tmp_path / "runs")) + + with patch("subprocess.Popen", return_value=proc), pytest.raises(TrainingError, match="timed out"): + runner.run_training( + template="train_xgboost", + model=object(), + feature_pipeline=MagicMock(), + train_uri=str(tmp_path / "train.parquet"), + val_uri=str(tmp_path / "val.parquet"), + timeout=2, + target_columns=["target"], + ) + + assert proc.kill_count == 1