fix: propagate asynchronous pipeline failures

src/htrflow/cli.py

@@ -1,7 +1,7 @@
 import logging
 import socket
 import time
-from concurrent.futures import ThreadPoolExecutor
+from concurrent.futures import ThreadPoolExecutor, as_completed
 from enum import Enum
 from typing import Iterable
 
@@ -112,9 +112,13 @@ def pipeline(
     n_pages = 0
 
     with ThreadPoolExecutor(workers) as executor:
+        futures = []
         for document in auto_import(inputs):
-            executor.submit(pipeline.run, document)
+            futures.append(executor.submit(pipeline.run, document))
             n_pages += 1
+
+        for future in as_completed(futures):
+            future.result()
     toc = time.time()
 
     total_time = toc - tic

src/htrflow/pipeline/steps.py

@@ -109,9 +109,17 @@ def from_config(cls, config):
     def _process(self):
         while 1:
             batch = self._queue.get()
-            outputs = self.model([item.item for item in batch], **self.generation_kwargs)
-            for output, item in zip(outputs, batch):
-                item.future.set_result(output)
+            try:
+                outputs = list(self.model([item.item for item in batch], **self.generation_kwargs))
+                if len(outputs) != len(batch):
+                    raise ValueError(f"Model returned {len(outputs)} outputs for a batch of {len(batch)} inputs")
+                for output, item in zip(outputs, batch):
+                    if not item.future.done():
+                        item.future.set_result(output)
+            except Exception as exception:
+                for item in batch:
+                    if not item.future.done():
+                        item.future.set_exception(exception)
 
     def run(self, document: Document):
         images = ImageLoader(document)

tests/unit/test_cli.py

@@ -1,5 +1,6 @@
 import pytest
 import typer
+from typer.testing import CliRunner
 
 from htrflow import cli
 
@@ -23,3 +24,33 @@ def test_get_inputs_file(tmp_path):
     inputs_file.write_text("\n".join(INPUTS))
     inputs = list(cli.get_inputs(None, inputs_file))
     assert inputs == INPUTS
+
+
+def test_pipeline_propagates_worker_exception(monkeypatch):
+    from htrflow.pipeline import steps
+    from htrflow.pipeline.pipeline import Pipeline
+
+    class FailingPipeline:
+        steps = []
+
+        def run(self, document):
+            raise RuntimeError("page failed")
+
+    class RecordingLogger:
+        def __init__(self):
+            self.messages = []
+
+        def info(self, message, *args):
+            self.messages.append(message % args)
+
+    logger = RecordingLogger()
+    monkeypatch.setattr(cli, "setup_pipeline_logging", lambda *args: logger)
+    monkeypatch.setattr(Pipeline, "from_config", classmethod(lambda cls, path: FailingPipeline()))
+    monkeypatch.setattr(steps, "auto_import", lambda inputs: iter([object()]))
+
+    result = CliRunner().invoke(cli.app, ["pipeline", "pipeline.yaml", "page.jpg"])
+
+    assert result.exit_code == 1
+    assert isinstance(result.exception, RuntimeError)
+    assert str(result.exception) == "page failed"
+    assert not any("Processed" in message for message in logger.messages)

tests/unit/test_steps.py

@@ -0,0 +1,75 @@
+import pytest
+
+from htrflow.pipeline import steps
+
+
+class FailingModel:
+    metadata = {}
+
+    def __call__(self, images, **kwargs):
+        raise RuntimeError("model failed")
+
+
+def test_inference_propagates_model_exception(monkeypatch):
+    step = steps.Inference(FailingModel(), batch_size=1)
+    monkeypatch.setattr(steps, "ImageLoader", lambda document: iter([(object(), object())]))
+
+    with pytest.raises(RuntimeError, match="model failed"):
+        step.run(object())
+
+    assert step._thread.is_alive()
+
+
+def test_inference_thread_survives_failed_batch():
+    class FailOnceModel:
+        metadata = {}
+
+        def __init__(self):
+            self.calls = 0
+
+        def __call__(self, images, **kwargs):
+            self.calls += 1
+            if self.calls == 1:
+                raise RuntimeError("model failed")
+            return [f"result {index}" for index, _ in enumerate(images)]
+
+    step = steps.Inference(FailOnceModel(), batch_size=2)
+    failed_futures = [step._queue.put(item) for item in ("first", "second")]
+    for future in failed_futures:
+        with pytest.raises(RuntimeError, match="model failed"):
+            future.result(timeout=5)
+
+    successful_futures = [step._queue.put(item) for item in ("third", "fourth")]
+    assert [future.result(timeout=5) for future in successful_futures] == ["result 0", "result 1"]
+    assert step._thread.is_alive()
+
+
+def test_inference_rejects_wrong_number_of_outputs():
+    class ShortModel:
+        metadata = {}
+
+        def __call__(self, images, **kwargs):
+            return ["only one result"]
+
+    step = steps.Inference(ShortModel(), batch_size=2)
+    futures = [step._queue.put(item) for item in ("first", "second")]
+
+    for future in futures:
+        with pytest.raises(ValueError, match="returned 1 outputs for a batch of 2 inputs"):
+            future.result(timeout=5)
+
+
+def test_inference_propagates_output_iteration_exception():
+    class FailingGeneratorModel:
+        metadata = {}
+
+        def __call__(self, images, **kwargs):
+            yield "partial result"
+            raise RuntimeError("output iteration failed")
+
+    step = steps.Inference(FailingGeneratorModel(), batch_size=2)
+    futures = [step._queue.put(item) for item in ("first", "second")]
+
+    for future in futures:
+        with pytest.raises(RuntimeError, match="output iteration failed"):
+            future.result(timeout=5)