Improve task timeout and wire validation

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.

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()

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."""

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")