diff --git a/src/htrflow/cli.py b/src/htrflow/cli.py index 97fbc18..2f331ae 100644 --- a/src/htrflow/cli.py +++ b/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 diff --git a/src/htrflow/pipeline/steps.py b/src/htrflow/pipeline/steps.py index 6800b28..7a564c9 100644 --- a/src/htrflow/pipeline/steps.py +++ b/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) diff --git a/tests/unit/test_cli.py b/tests/unit/test_cli.py index 5d63bbe..262da43 100644 --- a/tests/unit/test_cli.py +++ b/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) diff --git a/tests/unit/test_steps.py b/tests/unit/test_steps.py new file mode 100644 index 0000000..a12d6ff --- /dev/null +++ b/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)