Repository: PrefectHQ/prefect

Issue:
Task server should submit sync tasks to thread or event loop
Now that the task server uses the new engine, it should handle sync tasks differently than it did with the old engine. The task server is async but the new engine is sync, so if the task server runs sync tasks as they come in, each of those tasks will block the event loop. Instead, the task server can run async tasks normally but should submit sync tasks to a thread pool or other mechanism that won't block the event loop.

Something else to consider is, can we distribute work across multiple CPUs with a single task server? This could be a separate issue, too.


Discussion:

Resolving patch:
diff --git a/src/prefect/task_server.py b/src/prefect/task_server.py
--- a/src/prefect/task_server.py
+++ b/src/prefect/task_server.py
@@ -4,10 +4,13 @@
 import signal
 import socket
 import sys
+from concurrent.futures import ThreadPoolExecutor
 from contextlib import AsyncExitStack
+from contextvars import copy_context
 from typing import List
 
 import anyio
+import anyio.abc
 from websockets.exceptions import InvalidStatusCode
 
 from prefect import Task, get_client
@@ -79,6 +82,7 @@ def __init__(
             )
 
         self._runs_task_group: anyio.abc.TaskGroup = anyio.create_task_group()
+        self._executor = ThreadPoolExecutor()
 
     @property
     def _client_id(self) -> str:
@@ -140,7 +144,7 @@ async def _subscribe_to_task_scheduling(self):
             client_id=self._client_id,
         ):
             logger.info(f"Received task run: {task_run.id} - {task_run.name}")
-            await self._submit_scheduled_task_run(task_run)
+            self._runs_task_group.start_soon(self._submit_scheduled_task_run, task_run)
 
     async def _submit_scheduled_task_run(self, task_run: TaskRun):
         logger.debug(
@@ -225,13 +229,17 @@ async def _submit_scheduled_task_run(self, task_run: TaskRun):
                 return_type="state",
             )
         else:
-            run_task_sync(
+            context = copy_context()
+            future = self._executor.submit(
+                context.run,
+                run_task_sync,
                 task=task,
                 task_run_id=task_run.id,
                 task_run=task_run,
                 parameters=parameters,
                 return_type="state",
             )
+            await asyncio.wrap_future(future)
 
     async def execute_task_run(self, task_run: TaskRun):
         """Execute a task run in the task server."""
@@ -245,7 +253,8 @@ async def __aenter__(self):
             self._client = get_client()
 
         await self._exit_stack.enter_async_context(self._client)
-        await self._runs_task_group.__aenter__()
+        await self._exit_stack.enter_async_context(self._runs_task_group)
+        self._exit_stack.enter_context(self._executor)
 
         self.started = True
         return self
@@ -253,7 +262,6 @@ async def __aenter__(self):
     async def __aexit__(self, *exc_info):
         logger.debug("Stopping task server...")
         self.started = False
-        await self._runs_task_group.__aexit__(*exc_info)
         await self._exit_stack.__aexit__(*exc_info)
 
 


Test patch: