diff --git a/docs/guides/tasks.md b/docs/guides/tasks.md index 7c0cc8d..8c94484 100644 --- a/docs/guides/tasks.md +++ b/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: diff --git a/macfleet/compute/models.py b/macfleet/compute/models.py index abba120..f7cc6a2 100644 --- a/macfleet/compute/models.py +++ b/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() diff --git a/macfleet/compute/worker.py b/macfleet/compute/worker.py index 24c05a5..d8f0f14 100644 --- a/macfleet/compute/worker.py +++ b/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: diff --git a/macfleet/sdk/pool.py b/macfleet/sdk/pool.py index 8f5f460..e4067e7 100644 --- a/macfleet/sdk/pool.py +++ b/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(). diff --git a/tests/test_compute/test_dispatch.py b/tests/test_compute/test_dispatch.py index 7abb5db..8d36b46 100644 --- a/tests/test_compute/test_dispatch.py +++ b/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): diff --git a/tests/test_compute/test_models.py b/tests/test_compute/test_models.py index f02591f..fe46ba5 100644 --- a/tests/test_compute/test_models.py +++ b/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"]) diff --git a/tests/test_compute/test_pool_compute.py b/tests/test_compute/test_pool_compute.py index 684b418..7b756bd 100644 --- a/tests/test_compute/test_pool_compute.py +++ b/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):