Improve task timeout handling

docs/guides/tasks.md

@@ -28,6 +28,22 @@ with macfleet.Pool() as pool:
 `pool.submit` and `pool.map` detect the `@task` decorator and route the
 call through the registry. Args/kwargs are serialized via msgpack.
 
+## Timeouts and local parallelism
+
+`timeout` is enforced for registered tasks in both single-Mac and
+fleet-backed execution:
+
+```python
+with macfleet.Pool() as pool:
+    result = pool.submit(resize, "/tmp/a.jpg", target_w=512, timeout=10)
+    results = pool.map(resize, image_paths, timeout=10, max_workers=4)
+```
+
+If a task misses its timeout, `Pool.submit`/`Pool.map` raise
+`TimeoutError` with the task id and timeout value. For a single-Mac
+pool, registered `pool.map` calls run in a local thread pool and preserve
+input order while honoring `max_workers`.
+
 ## Why not just `pool.submit(lambda x: ..., x)`?
 
 Undecorated functions are rejected by default:

macfleet/compute/models.py

@@ -392,6 +392,11 @@ async def result(self, timeout: Optional[float] = None) -> Any:
             ValueError: If timeout is not None or a positive finite number.
         """
         wait_timeout = _validate_optional_timeout(timeout)
-        await asyncio.wait_for(self._event.wait(), timeout=wait_timeout)
+        try:
+            await asyncio.wait_for(self._event.wait(), timeout=wait_timeout)
+        except asyncio.TimeoutError as e:
+            raise asyncio.TimeoutError(
+                f"Task {self.task_id!r} timed out after {wait_timeout}s"
+            ) from e
         assert self._result is not None
         return self._result.unwrap()

macfleet/compute/worker.py

@@ -130,7 +130,7 @@ async def stop(self) -> None:
             await asyncio.gather(*self._inflight, return_exceptions=True)
             self._inflight.clear()
         if self._executor:
-            self._executor.shutdown(wait=True)
+            self._executor.shutdown(wait=False, cancel_futures=True)
             self._executor = None
 
     async def _listen_tasks(self) -> None:

macfleet/sdk/pool.py

@@ -25,7 +25,13 @@
 import logging
 import threading
 import time
-from concurrent.futures import ProcessPoolExecutor
+from concurrent.futures import (
+    ProcessPoolExecutor,
+    ThreadPoolExecutor,
+)
+from concurrent.futures import (
+    TimeoutError as FutureTimeoutError,
+)
 from typing import TYPE_CHECKING, Any, Callable, Iterable, Optional
 
 if TYPE_CHECKING:
@@ -975,7 +981,12 @@ def process(img): ...
             return []
 
         if self._is_registered_task(fn):
-            return [self._run_registered_task(fn, item, timeout=timeout) for item in items]
+            return self._run_registered_task_map(
+                fn,
+                items,
+                timeout=timeout,
+                max_workers=max_workers,
+            )
 
         self._ensure_legacy_pickle_allowed("map")
 
@@ -1103,9 +1114,67 @@ def _run_registered_task(
         spec = TaskSpec.from_call(fn, args=args, kwargs=kwargs, timeout=timeout)
         entry = spec.resolve()
         resolved_args, resolved_kwargs = spec.validated_args(entry)
-        # Invoke the registered callable in-process. A future PR wires
-        # this to TaskDispatcher when pool.world_size > 1.
-        return entry.fn(*resolved_args, **resolved_kwargs)
+        executor = ThreadPoolExecutor(max_workers=1)
+        future = executor.submit(entry.fn, *resolved_args, **resolved_kwargs)
+        try:
+            return future.result(timeout=spec.timeout_sec)
+        except FutureTimeoutError as e:
+            future.cancel()
+            raise TimeoutError(
+                f"Pool task {spec.task_id!r} timed out after {spec.timeout_sec}s"
+            ) from e
+        finally:
+            executor.shutdown(wait=future.done(), cancel_futures=True)
+
+    def _run_registered_task_map(
+        self,
+        fn: Any,
+        items: list[Any],
+        timeout: float = 300.0,
+        max_workers: Optional[int] = None,
+    ) -> list[Any]:
+        """Run registered tasks locally with bounded parallelism and ordered results."""
+        if max_workers is not None and (
+            not isinstance(max_workers, int)
+            or isinstance(max_workers, bool)
+            or max_workers < 1
+        ):
+            raise ValueError("max_workers must be a positive integer when provided")
+
+        import os
+        from concurrent.futures import wait
+
+        from macfleet.compute.models import TaskSpec
+
+        workers = max_workers or min(os.cpu_count() or 1, 4)
+        calls = []
+        for item in items:
+            spec = TaskSpec.from_call(fn, args=(item,), kwargs={}, timeout=timeout)
+            entry = spec.resolve()
+            resolved_args, resolved_kwargs = spec.validated_args(entry)
+            calls.append((spec, entry, resolved_args, resolved_kwargs))
+
+        executor = ThreadPoolExecutor(max_workers=workers)
+        futures = [
+            executor.submit(entry.fn, *resolved_args, **resolved_kwargs)
+            for _, entry, resolved_args, resolved_kwargs in calls
+        ]
+        try:
+            _, not_done = wait(futures, timeout=timeout)
+            if not_done:
+                for future in not_done:
+                    future.cancel()
+                first_timeout = futures.index(next(iter(not_done)))
+                timed_out_spec = calls[first_timeout][0]
+                raise TimeoutError(
+                    f"Pool.map task {timed_out_spec.task_id!r} timed out after {timeout}s"
+                )
+            return [future.result() for future in futures]
+        finally:
+            executor.shutdown(
+                wait=all(future.done() for future in futures),
+                cancel_futures=True,
+            )
 
     def run(self, fn: Callable, *args: Any, **kwargs: Any) -> Any:
         """Run a function on the pool. Shorthand for submit().

tests/test_compute/test_dispatch.py

@@ -10,6 +10,7 @@
 """
 
 import asyncio
+import time
 
 import pytest
 
@@ -54,6 +55,12 @@ def local_only_secret(x: int) -> int:
     return x + 999
 
 
+@task
+def slow_identity(x: int, delay: float = 0.3) -> int:
+    time.sleep(delay)
+    return x
+
+
 async def _setup_pair() -> tuple[PeerTransport, PeerTransport, int]:
     """Create a connected coordinator-worker transport pair over loopback."""
     coordinator = PeerTransport(local_id="coordinator", config=CONFIG)
@@ -351,6 +358,30 @@ async def test_multiple_tasks_round_robin(self):
             await w0.disconnect_all()
             await w1.disconnect_all()
 
+    @pytest.mark.asyncio
+    async def test_worker_stop_returns_after_timed_out_task(self):
+        """A timed-out user task must not make worker shutdown wait for the thread."""
+        coordinator, worker, _ = await _setup_pair()
+        try:
+            dispatcher = TaskDispatcher(coordinator, ["worker-0"])
+            tw = TaskWorker(worker, "coordinator", max_workers=1)
+
+            await dispatcher.start()
+            await tw.start()
+
+            future = await dispatcher.submit(slow_identity, 7, timeout=0.05, delay=0.5)
+            with pytest.raises(RemoteTaskError, match="timed out"):
+                await future.result(timeout=1.0)
+
+            start = time.monotonic()
+            await tw.stop()
+            elapsed = time.monotonic() - start
+            assert elapsed < 0.25
+
+            await dispatcher.stop()
+        finally:
+            await _teardown(coordinator, worker)
+
 
 class TestDispatcherEdgeCases:
     def test_no_workers_raises(self):

tests/test_compute/test_models.py

@@ -388,7 +388,7 @@ async def set_later():
     @pytest.mark.asyncio
     async def test_result_timeout(self):
         future = TaskFuture(task_id="f-3")
-        with pytest.raises(asyncio.TimeoutError):
+        with pytest.raises(asyncio.TimeoutError, match="f-3.*0.05s"):
             await future.result(timeout=0.05)
 
     @pytest.mark.parametrize("timeout", [0, -1, True, float("nan"), float("inf"), "soon"])

tests/test_compute/test_pool_compute.py

@@ -1,5 +1,7 @@
 """Tests for Pool.map(), Pool.submit(), Pool.run() — single-node fast path."""
 
+import time
+
 import pytest
 
 from macfleet import task
@@ -31,6 +33,19 @@ def _get_value():
     return 42
 
 
+@task
+def _slow_value(delay):
+    time.sleep(delay)
+    return "done"
+
+
+@task
+def _slow_square(payload):
+    delay, value = payload
+    time.sleep(delay)
+    return value * value
+
+
 class TestPoolMap:
     def test_map_basic(self):
         with Pool(open=True) as pool:
@@ -62,6 +77,35 @@ def test_map_not_joined_raises(self):
         with pytest.raises(RuntimeError, match="Must join"):
             pool.map(_square, [1])
 
+    def test_map_honors_max_workers_for_registered_tasks(self):
+        with Pool(open=True) as pool:
+            start = time.monotonic()
+            results = pool.map(
+                _slow_square,
+                [(0.2, 1), (0.2, 2), (0.2, 3)],
+                timeout=1.0,
+                max_workers=3,
+            )
+            elapsed = time.monotonic() - start
+
+        assert results == [1, 4, 9]
+        assert elapsed < 0.4
+
+    def test_map_timeout_is_enforced_for_registered_tasks(self):
+        with Pool(open=True) as pool:
+            start = time.monotonic()
+            with pytest.raises(TimeoutError, match="Pool.map task .* timed out after 0.05s"):
+                pool.map(_slow_value, [0.5], timeout=0.05)
+            elapsed = time.monotonic() - start
+
+        assert elapsed < 0.25
+
+    @pytest.mark.parametrize("max_workers", [0, -1, True, 1.5])
+    def test_map_rejects_invalid_max_workers_for_registered_tasks(self, max_workers):
+        with Pool(open=True) as pool:
+            with pytest.raises(ValueError, match="max_workers"):
+                pool.map(_square, [1], max_workers=max_workers)
+
 
 class TestPoolSubmit:
     def test_submit_basic(self):
@@ -79,6 +123,15 @@ def test_submit_not_joined_raises(self):
         with pytest.raises(RuntimeError, match="Must join"):
             pool.submit(_square, 5)
 
+    def test_submit_timeout_is_enforced_for_registered_task(self):
+        with Pool(open=True) as pool:
+            start = time.monotonic()
+            with pytest.raises(TimeoutError, match="timed out after 0.05s"):
+                pool.submit(_slow_value, 0.5, timeout=0.05)
+            elapsed = time.monotonic() - start
+
+        assert elapsed < 0.25
+
 
 class TestPoolRun:
     def test_run_basic(self):