diff --git a/macfleet/compute/dispatch.py b/macfleet/compute/dispatch.py index d941cca..c01e2a5 100644 --- a/macfleet/compute/dispatch.py +++ b/macfleet/compute/dispatch.py @@ -125,6 +125,10 @@ async def submit( raise RuntimeError("TaskDispatcher must be started before submitting tasks") spec = TaskSpec.from_call(fn, args, kwargs, timeout=timeout) + try: + payload = spec.pack() + except (TypeError, ValueError) as e: + raise ValueError(f"TaskSpec could not be serialized: {e}") from e future = TaskFuture(task_id=spec.task_id) worker_id = self._pick_live_worker() @@ -141,7 +145,7 @@ async def submit( self._pending[spec.task_id] = future self._task_to_worker[spec.task_id] = worker_id try: - await self._transport.send(worker_id, spec.pack(), msg_type=MessageType.TASK) + await self._transport.send(worker_id, payload, msg_type=MessageType.TASK) except (ConnectionError, OSError) as e: self._pending.pop(spec.task_id, None) self._task_to_worker.pop(spec.task_id, None) diff --git a/macfleet/compute/worker.py b/macfleet/compute/worker.py index d8f0f14..a938614 100644 --- a/macfleet/compute/worker.py +++ b/macfleet/compute/worker.py @@ -199,10 +199,30 @@ async def _execute_and_reply(self, spec: TaskSpec) -> None: ) result = TaskResult.failure(spec.task_id, e) + try: + payload = result.pack() + except (TypeError, ValueError) as e: + logger.warning( + "Task %s produced an unserializable result: %s", + spec.task_id[:8], e, + ) + audit_event( + "task.result_serialization_failed", + task_name=spec.task_name, + coordinator=self._coordinator, + error_type=type(e).__name__, + ) + result = TaskResult( + task_id=spec.task_id, + ok=False, + error=f"Task result could not be serialized: {e}", + ) + payload = result.pack() + try: await self._transport.send( self._coordinator, - result.pack(), + payload, msg_type=MessageType.RESULT, ) except (ConnectionError, OSError) as e: diff --git a/tests/test_compute/test_dispatch.py b/tests/test_compute/test_dispatch.py index 8d36b46..c188249 100644 --- a/tests/test_compute/test_dispatch.py +++ b/tests/test_compute/test_dispatch.py @@ -61,6 +61,11 @@ def slow_identity(x: int, delay: float = 0.3) -> int: return x +@task +def returns_unserializable_object(): + return object() + + async def _setup_pair() -> tuple[PeerTransport, PeerTransport, int]: """Create a connected coordinator-worker transport pair over loopback.""" coordinator = PeerTransport(local_id="coordinator", config=CONFIG) @@ -382,6 +387,34 @@ async def test_worker_stop_returns_after_timed_out_task(self): finally: await _teardown(coordinator, worker) + @pytest.mark.asyncio + async def test_unserializable_worker_result_returns_remote_error(self): + """Worker return values that cannot cross msgpack fail the task, not the listener.""" + 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(returns_unserializable_object, timeout=5.0) + with pytest.raises(RemoteTaskError) as exc_info: + await future.result(timeout=5.0) + + assert "Task result could not be serialized" in exc_info.value.remote_traceback + assert dispatcher.pending_count == 0 + assert tw._listener_task is not None + assert not tw._listener_task.done() + + healthy = await dispatcher.submit(times_two, 21, timeout=5.0) + assert await healthy.result(timeout=5.0) == 42 + + await tw.stop() + await dispatcher.stop() + finally: + await _teardown(coordinator, worker) + class TestDispatcherEdgeCases: def test_no_workers_raises(self): @@ -408,6 +441,19 @@ async def run(): asyncio.run(run()) + def test_unserializable_task_args_do_not_leak_pending_future(self): + transport = PeerTransport(local_id="solo", config=CONFIG) + dispatcher = TaskDispatcher(transport, ["worker-0"]) + + async def run(): + await dispatcher.start() + with pytest.raises(ValueError, match="TaskSpec could not be serialized"): + await dispatcher.submit(times_two, object(), timeout=1.0) + assert dispatcher.pending_count == 0 + await dispatcher.stop() + + asyncio.run(run()) + def test_dispatcher_cannot_start_twice_and_can_restart_after_stop(self): transport = PeerTransport(local_id="solo", config=CONFIG) dispatcher = TaskDispatcher(transport, ["worker-0"])