diff --git a/macfleet/compute/dispatch.py b/macfleet/compute/dispatch.py index 9350e16..d941cca 100644 --- a/macfleet/compute/dispatch.py +++ b/macfleet/compute/dispatch.py @@ -182,10 +182,9 @@ async def map( future = await self.submit(fn, item, timeout=timeout) futures.append(future) - results = [] - for future in futures: - results.append(await future.result(timeout=timeout)) - return results + return list(await asyncio.gather( + *(future.result(timeout=timeout) for future in futures) + )) async def _worker_listener(self, worker_id: str) -> None: """Listen for RESULT messages from one worker. diff --git a/macfleet/compute/models.py b/macfleet/compute/models.py index cc0c1c1..abba120 100644 --- a/macfleet/compute/models.py +++ b/macfleet/compute/models.py @@ -61,7 +61,7 @@ def _unpack_msgpack_dict(data: bytes, label: str) -> dict: """Decode msgpack and require a top-level mapping.""" try: payload = msgpack.unpackb(data, raw=False) - except msgpack.exceptions.UnpackException as e: + except (msgpack.exceptions.UnpackException, ValueError) as e: raise ValueError(f"{label} payload is not valid msgpack") from e if not isinstance(payload, dict): raise ValueError(f"{label} payload not a dict: {type(payload).__name__}") @@ -82,6 +82,25 @@ def _require_bool(payload: dict, key: str, label: str) -> bool: return value +def _validate_optional_timeout(timeout: Optional[float]) -> Optional[float]: + if timeout is None: + return None + if ( + isinstance(timeout, bool) + or not isinstance(timeout, (int, float)) + or not math.isfinite(float(timeout)) + or timeout <= 0 + ): + raise ValueError("timeout must be None or a positive finite number") + return float(timeout) + + +def _validate_kwargs(kwargs: dict, label: str) -> dict: + if any(not isinstance(key, str) or not key for key in kwargs): + raise ValueError(f"{label} kwargs keys must be non-empty strings") + return kwargs + + class RemoteTaskError(Exception): """Raised when a task fails on a remote worker. @@ -158,7 +177,7 @@ def from_call( # dump to msgpack-native types. schema = getattr(fn, "schema", None) arg_list: list = list(args) - kwarg_dict: dict = dict(kwargs or {}) + kwarg_dict: dict = _validate_kwargs(dict(kwargs or {}), "TaskSpec") if schema is not None: if len(arg_list) == 1 and isinstance(arg_list[0], schema): # Common shape: @task(schema=X) def f(args: X): ... @@ -210,6 +229,7 @@ def unpack(cls, data: bytes) -> TaskSpec: raise ValueError("TaskSpec field 'args' must be a list") if not isinstance(kwargs, dict): raise ValueError("TaskSpec field 'kwargs' must be a dict") + kwargs = _validate_kwargs(kwargs, "TaskSpec") if isinstance(timeout, bool) or not isinstance(timeout, (int, float)): raise ValueError("TaskSpec field 'timeout' must be a number") if not math.isfinite(float(timeout)) or float(timeout) <= 0: @@ -318,11 +338,14 @@ def unpack(cls, data: bytes) -> TaskResult: if len(data) > MAX_RESULT_BYTES: raise ValueError(f"TaskResult size {len(data)}B exceeds max {MAX_RESULT_BYTES}B") d = _unpack_msgpack_dict(data, "TaskResult") + error = d.get("error") + if error is not None and not isinstance(error, str): + raise ValueError("TaskResult field 'error' must be a string or null") return cls( task_id=_require_str(d, "task_id", "TaskResult"), ok=_require_bool(d, "ok", "TaskResult"), value=d.get("value"), - error=d.get("error"), + error=error, ) def unwrap(self) -> Any: @@ -366,7 +389,9 @@ async def result(self, timeout: Optional[float] = None) -> Any: Raises: RemoteTaskError: If the task failed on the worker. asyncio.TimeoutError: If timeout is exceeded. + ValueError: If timeout is not None or a positive finite number. """ - await asyncio.wait_for(self._event.wait(), timeout=timeout) + wait_timeout = _validate_optional_timeout(timeout) + await asyncio.wait_for(self._event.wait(), timeout=wait_timeout) assert self._result is not None return self._result.unwrap() diff --git a/tests/test_compute/test_dispatch.py b/tests/test_compute/test_dispatch.py index bdda881..7abb5db 100644 --- a/tests/test_compute/test_dispatch.py +++ b/tests/test_compute/test_dispatch.py @@ -216,6 +216,8 @@ async def test_worker_ignores_malformed_task_frame(self): @pytest.mark.asyncio async def test_dispatcher_ignores_malformed_result_frame(self): """A malformed RESULT frame must not kill the coordinator listener.""" + import msgpack + coordinator, worker, _ = await _setup_pair() try: dispatcher = TaskDispatcher(coordinator, ["worker-0"]) @@ -225,6 +227,16 @@ async def test_dispatcher_ignores_malformed_result_frame(self): await tw.start() await worker.send("coordinator", b"\x91\x01", msg_type=MessageType.RESULT) + await worker.send( + "coordinator", + msgpack.packb({ + "task_id": "bad-result-error", + "ok": False, + "value": None, + "error": ["not", "a", "string"], + }, use_bin_type=True), + msg_type=MessageType.RESULT, + ) await asyncio.sleep(0.05) listener = dispatcher._worker_listeners["worker-0"] assert not listener.done() @@ -267,6 +279,38 @@ async def test_map_empty_iterable(self): finally: await _teardown(coordinator, worker) + @pytest.mark.asyncio + async def test_map_awaits_results_concurrently(self, monkeypatch): + """map() should not multiply the per-call wait timeout by input size.""" + transport = PeerTransport(local_id="solo", config=CONFIG) + dispatcher = TaskDispatcher(transport, ["worker-0"]) + waiters = 0 + all_waiting = asyncio.Event() + + class FakeFuture: + def __init__(self, value: int): + self.value = value + + async def result(self, timeout: float = 300.0) -> int: + nonlocal waiters + waiters += 1 + if waiters == 3: + all_waiting.set() + await all_waiting.wait() + return self.value + + async def fake_submit(fn, item, timeout=300.0): + return FakeFuture(item) + + monkeypatch.setattr(dispatcher, "submit", fake_submit) + + results = await asyncio.wait_for( + dispatcher.map(times_two, [1, 2, 3], timeout=1.0), + timeout=0.5, + ) + + assert results == [1, 2, 3] + @pytest.mark.asyncio async def test_multiple_tasks_round_robin(self): """With 2 workers, tasks alternate between them.""" diff --git a/tests/test_compute/test_models.py b/tests/test_compute/test_models.py index f3d1d03..f02591f 100644 --- a/tests/test_compute/test_models.py +++ b/tests/test_compute/test_models.py @@ -72,6 +72,11 @@ def test_from_call_with_kwargs(self): assert spec.timeout_sec == 60.0 assert spec.args == [[1, 2, 3]] + @pytest.mark.parametrize("kwargs", [{1: "bad"}, {"": "bad"}]) + def test_from_call_rejects_bad_kwarg_keys(self, kwargs): + with pytest.raises(ValueError, match="kwargs"): + TaskSpec.from_call(square, kwargs=kwargs) + @pytest.mark.parametrize("timeout", [0, -1, True, float("nan"), float("inf")]) def test_from_call_rejects_invalid_timeout(self, timeout): with pytest.raises(ValueError, match="timeout"): @@ -160,6 +165,33 @@ def test_unpack_rejects_bad_field_shapes(self): with pytest.raises(ValueError, match="args"): TaskSpec.unpack(data) + @pytest.mark.parametrize("kwargs", [{b"bad": 1}, {"": 1}]) + def test_unpack_rejects_bad_kwarg_keys(self, kwargs): + import msgpack + + data = msgpack.packb({ + "task_id": "abc123", + "name": square.task_name, + "args": [], + "kwargs": kwargs, + "timeout": 300.0, + }, use_bin_type=True) + with pytest.raises(ValueError, match="kwargs"): + TaskSpec.unpack(data) + + def test_unpack_wraps_strict_msgpack_map_key_errors(self): + import msgpack + + data = msgpack.packb({ + "task_id": "abc123", + "name": square.task_name, + "args": [], + "kwargs": {1: "bad"}, + "timeout": 300.0, + }, use_bin_type=True) + with pytest.raises(ValueError, match="TaskSpec payload is not valid msgpack"): + TaskSpec.unpack(data) + # --------------------------------------------------------------------------- # # Pydantic schema validation # @@ -297,6 +329,18 @@ def test_unpack_rejects_bad_result_shape(self): with pytest.raises(ValueError, match="ok"): TaskResult.unpack(data) + def test_unpack_rejects_bad_error_shape(self): + import msgpack + + data = msgpack.packb({ + "task_id": "task-bad-error", + "ok": False, + "value": None, + "error": ["not", "a", "string"], + }, use_bin_type=True) + with pytest.raises(ValueError, match="error"): + TaskResult.unpack(data) + def test_pack_rejects_result_larger_than_limit(self, monkeypatch): monkeypatch.setattr(models, "MAX_RESULT_BYTES", 1) @@ -347,6 +391,13 @@ async def test_result_timeout(self): with pytest.raises(asyncio.TimeoutError): await future.result(timeout=0.05) + @pytest.mark.parametrize("timeout", [0, -1, True, float("nan"), float("inf"), "soon"]) + @pytest.mark.asyncio + async def test_result_rejects_invalid_timeout(self, timeout): + future = TaskFuture(task_id="f-3-invalid") + with pytest.raises(ValueError, match="timeout"): + await future.result(timeout=timeout) + @pytest.mark.asyncio async def test_result_failure(self): future = TaskFuture(task_id="f-4")