We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN 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.

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `src/prefect/task_server.py`
Content:
```
1 import asyncio
2 import inspect
3 import os
4 import signal
5 import socket
6 import sys
7 from contextlib import AsyncExitStack
8 from typing import List
9 
10 import anyio
11 from websockets.exceptions import InvalidStatusCode
12 
13 from prefect import Task, get_client
14 from prefect._internal.concurrency.api import create_call, from_sync
15 from prefect.client.schemas.objects import TaskRun
16 from prefect.client.subscriptions import Subscription
17 from prefect.exceptions import Abort, PrefectHTTPStatusError
18 from prefect.logging.loggers import get_logger
19 from prefect.results import ResultFactory
20 from prefect.settings import (
21     PREFECT_API_URL,
22     PREFECT_EXPERIMENTAL_ENABLE_TASK_SCHEDULING,
23     PREFECT_TASK_SCHEDULING_DELETE_FAILED_SUBMISSIONS,
24 )
25 from prefect.states import Pending
26 from prefect.task_engine import run_task_async, run_task_sync
27 from prefect.utilities.asyncutils import asyncnullcontext, sync_compatible
28 from prefect.utilities.engine import emit_task_run_state_change_event, propose_state
29 from prefect.utilities.processutils import _register_signal
30 
31 logger = get_logger("task_server")
32 
33 
34 class StopTaskServer(Exception):
35     """Raised when the task server is stopped."""
36 
37     pass
38 
39 
40 def should_try_to_read_parameters(task: Task, task_run: TaskRun) -> bool:
41     """Determines whether a task run should read parameters from the result factory."""
42     new_enough_state_details = hasattr(
43         task_run.state.state_details, "task_parameters_id"
44     )
45     task_accepts_parameters = bool(inspect.signature(task.fn).parameters)
46 
47     return new_enough_state_details and task_accepts_parameters
48 
49 
50 class TaskServer:
51     """This class is responsible for serving tasks that may be executed in the background
52     by a task runner via the traditional engine machinery.
53 
54     When `start()` is called, the task server will open a websocket connection to a
55     server-side queue of scheduled task runs. When a scheduled task run is found, the
56     scheduled task run is submitted to the engine for execution with a minimal `EngineContext`
57     so that the task run can be governed by orchestration rules.
58 
59     Args:
60         - tasks: A list of tasks to serve. These tasks will be submitted to the engine
61             when a scheduled task run is found.
62     """
63 
64     def __init__(
65         self,
66         *tasks: Task,
67     ):
68         self.tasks: List[Task] = tasks
69 
70         self.started: bool = False
71         self.stopping: bool = False
72 
73         self._client = get_client()
74         self._exit_stack = AsyncExitStack()
75 
76         if not asyncio.get_event_loop().is_running():
77             raise RuntimeError(
78                 "TaskServer must be initialized within an async context."
79             )
80 
81         self._runs_task_group: anyio.abc.TaskGroup = anyio.create_task_group()
82 
83     @property
84     def _client_id(self) -> str:
85         return f"{socket.gethostname()}-{os.getpid()}"
86 
87     def handle_sigterm(self, signum, frame):
88         """
89         Shuts down the task server when a SIGTERM is received.
90         """
91         logger.info("SIGTERM received, initiating graceful shutdown...")
92         from_sync.call_in_loop_thread(create_call(self.stop))
93 
94         sys.exit(0)
95 
96     @sync_compatible
97     async def start(self) -> None:
98         """
99         Starts a task server, which runs the tasks provided in the constructor.
100         """
101         _register_signal(signal.SIGTERM, self.handle_sigterm)
102 
103         async with asyncnullcontext() if self.started else self:
104             logger.info("Starting task server...")
105             try:
106                 await self._subscribe_to_task_scheduling()
107             except InvalidStatusCode as exc:
108                 if exc.status_code == 403:
109                     logger.error(
110                         "Could not establish a connection to the `/task_runs/subscriptions/scheduled`"
111                         f" endpoint found at:\n\n {PREFECT_API_URL.value()}"
112                         "\n\nPlease double-check the values of your"
113                         " `PREFECT_API_URL` and `PREFECT_API_KEY` environment variables."
114                     )
115                 else:
116                     raise
117 
118     @sync_compatible
119     async def stop(self):
120         """Stops the task server's polling cycle."""
121         if not self.started:
122             raise RuntimeError(
123                 "Task server has not yet started. Please start the task server by"
124                 " calling .start()"
125             )
126 
127         self.started = False
128         self.stopping = True
129 
130         raise StopTaskServer
131 
132     async def _subscribe_to_task_scheduling(self):
133         logger.info(
134             f"Subscribing to tasks: {' | '.join(t.task_key.split('.')[-1] for t in self.tasks)}"
135         )
136         async for task_run in Subscription(
137             model=TaskRun,
138             path="/task_runs/subscriptions/scheduled",
139             keys=[task.task_key for task in self.tasks],
140             client_id=self._client_id,
141         ):
142             logger.info(f"Received task run: {task_run.id} - {task_run.name}")
143             await self._submit_scheduled_task_run(task_run)
144 
145     async def _submit_scheduled_task_run(self, task_run: TaskRun):
146         logger.debug(
147             f"Found task run: {task_run.name!r} in state: {task_run.state.name!r}"
148         )
149 
150         task = next((t for t in self.tasks if t.task_key == task_run.task_key), None)
151 
152         if not task:
153             if PREFECT_TASK_SCHEDULING_DELETE_FAILED_SUBMISSIONS.value():
154                 logger.warning(
155                     f"Task {task_run.name!r} not found in task server registry."
156                 )
157                 await self._client._client.delete(f"/task_runs/{task_run.id}")
158 
159             return
160 
161         # The ID of the parameters for this run are stored in the Scheduled state's
162         # state_details. If there is no parameters_id, then the task was created
163         # without parameters.
164         parameters = {}
165         if should_try_to_read_parameters(task, task_run):
166             parameters_id = task_run.state.state_details.task_parameters_id
167             task.persist_result = True
168             factory = await ResultFactory.from_autonomous_task(task)
169             try:
170                 parameters = await factory.read_parameters(parameters_id)
171             except Exception as exc:
172                 logger.exception(
173                     f"Failed to read parameters for task run {task_run.id!r}",
174                     exc_info=exc,
175                 )
176                 if PREFECT_TASK_SCHEDULING_DELETE_FAILED_SUBMISSIONS.value():
177                     logger.info(
178                         f"Deleting task run {task_run.id!r} because it failed to submit"
179                     )
180                     await self._client._client.delete(f"/task_runs/{task_run.id}")
181                 return
182 
183         logger.debug(
184             f"Submitting run {task_run.name!r} of task {task.name!r} to engine"
185         )
186 
187         try:
188             state = await propose_state(
189                 client=get_client(),  # TODO prove that we cannot use self._client here
190                 state=Pending(),
191                 task_run_id=task_run.id,
192             )
193         except Abort as exc:
194             logger.exception(
195                 f"Failed to submit task run {task_run.id!r} to engine", exc_info=exc
196             )
197             return
198         except PrefectHTTPStatusError as exc:
199             if exc.response.status_code == 404:
200                 logger.warning(
201                     f"Task run {task_run.id!r} not found. It may have been deleted."
202                 )
203                 return
204             raise
205 
206         if not state.is_pending():
207             logger.warning(
208                 f"Cancelling submission of task run {task_run.id!r} -"
209                 f" server returned a non-pending state {state.type.value!r}."
210             )
211             return
212 
213         emit_task_run_state_change_event(
214             task_run=task_run,
215             initial_state=task_run.state,
216             validated_state=state,
217         )
218 
219         if task.isasync:
220             await run_task_async(
221                 task=task,
222                 task_run_id=task_run.id,
223                 task_run=task_run,
224                 parameters=parameters,
225                 return_type="state",
226             )
227         else:
228             run_task_sync(
229                 task=task,
230                 task_run_id=task_run.id,
231                 task_run=task_run,
232                 parameters=parameters,
233                 return_type="state",
234             )
235 
236     async def execute_task_run(self, task_run: TaskRun):
237         """Execute a task run in the task server."""
238         async with self if not self.started else asyncnullcontext():
239             await self._submit_scheduled_task_run(task_run)
240 
241     async def __aenter__(self):
242         logger.debug("Starting task server...")
243 
244         if self._client._closed:
245             self._client = get_client()
246 
247         await self._exit_stack.enter_async_context(self._client)
248         await self._runs_task_group.__aenter__()
249 
250         self.started = True
251         return self
252 
253     async def __aexit__(self, *exc_info):
254         logger.debug("Stopping task server...")
255         self.started = False
256         await self._runs_task_group.__aexit__(*exc_info)
257         await self._exit_stack.__aexit__(*exc_info)
258 
259 
260 @sync_compatible
261 async def serve(*tasks: Task):
262     """Serve the provided tasks so that their runs may be submitted to and executed.
263     in the engine. Tasks do not need to be within a flow run context to be submitted.
264     You must `.submit` the same task object that you pass to `serve`.
265 
266     Args:
267         - tasks: A list of tasks to serve. When a scheduled task run is found for a
268             given task, the task run will be submitted to the engine for execution.
269 
270     Example:
271         ```python
272         from prefect import task
273         from prefect.task_server import serve
274 
275         @task(log_prints=True)
276         def say(message: str):
277             print(message)
278 
279         @task(log_prints=True)
280         def yell(message: str):
281             print(message.upper())
282 
283         # starts a long-lived process that listens for scheduled runs of these tasks
284         if __name__ == "__main__":
285             serve(say, yell)
286         ```
287     """
288     if not PREFECT_EXPERIMENTAL_ENABLE_TASK_SCHEDULING.value():
289         raise RuntimeError(
290             "To enable task scheduling, set PREFECT_EXPERIMENTAL_ENABLE_TASK_SCHEDULING"
291             " to True."
292         )
293 
294     task_server = TaskServer(*tasks)
295     try:
296         await task_server.start()
297 
298     except StopTaskServer:
299         logger.info("Task server stopped.")
300 
301     except asyncio.CancelledError:
302         logger.info("Task server interrupted, stopping...")
303
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

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 @@
             )
 
         self._runs_task_group: anyio.abc.TaskGroup = anyio.create_task_group()
+        self._executor = ThreadPoolExecutor()
 
     @property
     def _client_id(self) -> str:
@@ -140,7 +144,7 @@
             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 @@
                 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 @@
             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 __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)

{"golden_diff": "diff --git a/src/prefect/task_server.py b/src/prefect/task_server.py\n--- a/src/prefect/task_server.py\n+++ b/src/prefect/task_server.py\n@@ -4,10 +4,13 @@\n import signal\n import socket\n import sys\n+from concurrent.futures import ThreadPoolExecutor\n from contextlib import AsyncExitStack\n+from contextvars import copy_context\n from typing import List\n \n import anyio\n+import anyio.abc\n from websockets.exceptions import InvalidStatusCode\n \n from prefect import Task, get_client\n@@ -79,6 +82,7 @@\n             )\n \n         self._runs_task_group: anyio.abc.TaskGroup = anyio.create_task_group()\n+        self._executor = ThreadPoolExecutor()\n \n     @property\n     def _client_id(self) -> str:\n@@ -140,7 +144,7 @@\n             client_id=self._client_id,\n         ):\n             logger.info(f\"Received task run: {task_run.id} - {task_run.name}\")\n-            await self._submit_scheduled_task_run(task_run)\n+            self._runs_task_group.start_soon(self._submit_scheduled_task_run, task_run)\n \n     async def _submit_scheduled_task_run(self, task_run: TaskRun):\n         logger.debug(\n@@ -225,13 +229,17 @@\n                 return_type=\"state\",\n             )\n         else:\n-            run_task_sync(\n+            context = copy_context()\n+            future = self._executor.submit(\n+                context.run,\n+                run_task_sync,\n                 task=task,\n                 task_run_id=task_run.id,\n                 task_run=task_run,\n                 parameters=parameters,\n                 return_type=\"state\",\n             )\n+            await asyncio.wrap_future(future)\n \n     async def execute_task_run(self, task_run: TaskRun):\n         \"\"\"Execute a task run in the task server.\"\"\"\n@@ -245,7 +253,8 @@\n             self._client = get_client()\n \n         await self._exit_stack.enter_async_context(self._client)\n-        await self._runs_task_group.__aenter__()\n+        await self._exit_stack.enter_async_context(self._runs_task_group)\n+        self._exit_stack.enter_context(self._executor)\n \n         self.started = True\n         return self\n@@ -253,7 +262,6 @@\n     async def __aexit__(self, *exc_info):\n         logger.debug(\"Stopping task server...\")\n         self.started = False\n-        await self._runs_task_group.__aexit__(*exc_info)\n         await self._exit_stack.__aexit__(*exc_info)\n", "issue": "Task server should submit sync tasks to thread or event loop\nNow 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.\n\nSomething else to consider is, can we distribute work across multiple CPUs with a single task server? This could be a separate issue, too.\n", "before_files": [{"content": "import asyncio\nimport inspect\nimport os\nimport signal\nimport socket\nimport sys\nfrom contextlib import AsyncExitStack\nfrom typing import List\n\nimport anyio\nfrom websockets.exceptions import InvalidStatusCode\n\nfrom prefect import Task, get_client\nfrom prefect._internal.concurrency.api import create_call, from_sync\nfrom prefect.client.schemas.objects import TaskRun\nfrom prefect.client.subscriptions import Subscription\nfrom prefect.exceptions import Abort, PrefectHTTPStatusError\nfrom prefect.logging.loggers import get_logger\nfrom prefect.results import ResultFactory\nfrom prefect.settings import (\n    PREFECT_API_URL,\n    PREFECT_EXPERIMENTAL_ENABLE_TASK_SCHEDULING,\n    PREFECT_TASK_SCHEDULING_DELETE_FAILED_SUBMISSIONS,\n)\nfrom prefect.states import Pending\nfrom prefect.task_engine import run_task_async, run_task_sync\nfrom prefect.utilities.asyncutils import asyncnullcontext, sync_compatible\nfrom prefect.utilities.engine import emit_task_run_state_change_event, propose_state\nfrom prefect.utilities.processutils import _register_signal\n\nlogger = get_logger(\"task_server\")\n\n\nclass StopTaskServer(Exception):\n    \"\"\"Raised when the task server is stopped.\"\"\"\n\n    pass\n\n\ndef should_try_to_read_parameters(task: Task, task_run: TaskRun) -> bool:\n    \"\"\"Determines whether a task run should read parameters from the result factory.\"\"\"\n    new_enough_state_details = hasattr(\n        task_run.state.state_details, \"task_parameters_id\"\n    )\n    task_accepts_parameters = bool(inspect.signature(task.fn).parameters)\n\n    return new_enough_state_details and task_accepts_parameters\n\n\nclass TaskServer:\n    \"\"\"This class is responsible for serving tasks that may be executed in the background\n    by a task runner via the traditional engine machinery.\n\n    When `start()` is called, the task server will open a websocket connection to a\n    server-side queue of scheduled task runs. When a scheduled task run is found, the\n    scheduled task run is submitted to the engine for execution with a minimal `EngineContext`\n    so that the task run can be governed by orchestration rules.\n\n    Args:\n        - tasks: A list of tasks to serve. These tasks will be submitted to the engine\n            when a scheduled task run is found.\n    \"\"\"\n\n    def __init__(\n        self,\n        *tasks: Task,\n    ):\n        self.tasks: List[Task] = tasks\n\n        self.started: bool = False\n        self.stopping: bool = False\n\n        self._client = get_client()\n        self._exit_stack = AsyncExitStack()\n\n        if not asyncio.get_event_loop().is_running():\n            raise RuntimeError(\n                \"TaskServer must be initialized within an async context.\"\n            )\n\n        self._runs_task_group: anyio.abc.TaskGroup = anyio.create_task_group()\n\n    @property\n    def _client_id(self) -> str:\n        return f\"{socket.gethostname()}-{os.getpid()}\"\n\n    def handle_sigterm(self, signum, frame):\n        \"\"\"\n        Shuts down the task server when a SIGTERM is received.\n        \"\"\"\n        logger.info(\"SIGTERM received, initiating graceful shutdown...\")\n        from_sync.call_in_loop_thread(create_call(self.stop))\n\n        sys.exit(0)\n\n    @sync_compatible\n    async def start(self) -> None:\n        \"\"\"\n        Starts a task server, which runs the tasks provided in the constructor.\n        \"\"\"\n        _register_signal(signal.SIGTERM, self.handle_sigterm)\n\n        async with asyncnullcontext() if self.started else self:\n            logger.info(\"Starting task server...\")\n            try:\n                await self._subscribe_to_task_scheduling()\n            except InvalidStatusCode as exc:\n                if exc.status_code == 403:\n                    logger.error(\n                        \"Could not establish a connection to the `/task_runs/subscriptions/scheduled`\"\n                        f\" endpoint found at:\\n\\n {PREFECT_API_URL.value()}\"\n                        \"\\n\\nPlease double-check the values of your\"\n                        \" `PREFECT_API_URL` and `PREFECT_API_KEY` environment variables.\"\n                    )\n                else:\n                    raise\n\n    @sync_compatible\n    async def stop(self):\n        \"\"\"Stops the task server's polling cycle.\"\"\"\n        if not self.started:\n            raise RuntimeError(\n                \"Task server has not yet started. Please start the task server by\"\n                \" calling .start()\"\n            )\n\n        self.started = False\n        self.stopping = True\n\n        raise StopTaskServer\n\n    async def _subscribe_to_task_scheduling(self):\n        logger.info(\n            f\"Subscribing to tasks: {' | '.join(t.task_key.split('.')[-1] for t in self.tasks)}\"\n        )\n        async for task_run in Subscription(\n            model=TaskRun,\n            path=\"/task_runs/subscriptions/scheduled\",\n            keys=[task.task_key for task in self.tasks],\n            client_id=self._client_id,\n        ):\n            logger.info(f\"Received task run: {task_run.id} - {task_run.name}\")\n            await self._submit_scheduled_task_run(task_run)\n\n    async def _submit_scheduled_task_run(self, task_run: TaskRun):\n        logger.debug(\n            f\"Found task run: {task_run.name!r} in state: {task_run.state.name!r}\"\n        )\n\n        task = next((t for t in self.tasks if t.task_key == task_run.task_key), None)\n\n        if not task:\n            if PREFECT_TASK_SCHEDULING_DELETE_FAILED_SUBMISSIONS.value():\n                logger.warning(\n                    f\"Task {task_run.name!r} not found in task server registry.\"\n                )\n                await self._client._client.delete(f\"/task_runs/{task_run.id}\")\n\n            return\n\n        # The ID of the parameters for this run are stored in the Scheduled state's\n        # state_details. If there is no parameters_id, then the task was created\n        # without parameters.\n        parameters = {}\n        if should_try_to_read_parameters(task, task_run):\n            parameters_id = task_run.state.state_details.task_parameters_id\n            task.persist_result = True\n            factory = await ResultFactory.from_autonomous_task(task)\n            try:\n                parameters = await factory.read_parameters(parameters_id)\n            except Exception as exc:\n                logger.exception(\n                    f\"Failed to read parameters for task run {task_run.id!r}\",\n                    exc_info=exc,\n                )\n                if PREFECT_TASK_SCHEDULING_DELETE_FAILED_SUBMISSIONS.value():\n                    logger.info(\n                        f\"Deleting task run {task_run.id!r} because it failed to submit\"\n                    )\n                    await self._client._client.delete(f\"/task_runs/{task_run.id}\")\n                return\n\n        logger.debug(\n            f\"Submitting run {task_run.name!r} of task {task.name!r} to engine\"\n        )\n\n        try:\n            state = await propose_state(\n                client=get_client(),  # TODO prove that we cannot use self._client here\n                state=Pending(),\n                task_run_id=task_run.id,\n            )\n        except Abort as exc:\n            logger.exception(\n                f\"Failed to submit task run {task_run.id!r} to engine\", exc_info=exc\n            )\n            return\n        except PrefectHTTPStatusError as exc:\n            if exc.response.status_code == 404:\n                logger.warning(\n                    f\"Task run {task_run.id!r} not found. It may have been deleted.\"\n                )\n                return\n            raise\n\n        if not state.is_pending():\n            logger.warning(\n                f\"Cancelling submission of task run {task_run.id!r} -\"\n                f\" server returned a non-pending state {state.type.value!r}.\"\n            )\n            return\n\n        emit_task_run_state_change_event(\n            task_run=task_run,\n            initial_state=task_run.state,\n            validated_state=state,\n        )\n\n        if task.isasync:\n            await run_task_async(\n                task=task,\n                task_run_id=task_run.id,\n                task_run=task_run,\n                parameters=parameters,\n                return_type=\"state\",\n            )\n        else:\n            run_task_sync(\n                task=task,\n                task_run_id=task_run.id,\n                task_run=task_run,\n                parameters=parameters,\n                return_type=\"state\",\n            )\n\n    async def execute_task_run(self, task_run: TaskRun):\n        \"\"\"Execute a task run in the task server.\"\"\"\n        async with self if not self.started else asyncnullcontext():\n            await self._submit_scheduled_task_run(task_run)\n\n    async def __aenter__(self):\n        logger.debug(\"Starting task server...\")\n\n        if self._client._closed:\n            self._client = get_client()\n\n        await self._exit_stack.enter_async_context(self._client)\n        await self._runs_task_group.__aenter__()\n\n        self.started = True\n        return self\n\n    async def __aexit__(self, *exc_info):\n        logger.debug(\"Stopping task server...\")\n        self.started = False\n        await self._runs_task_group.__aexit__(*exc_info)\n        await self._exit_stack.__aexit__(*exc_info)\n\n\n@sync_compatible\nasync def serve(*tasks: Task):\n    \"\"\"Serve the provided tasks so that their runs may be submitted to and executed.\n    in the engine. Tasks do not need to be within a flow run context to be submitted.\n    You must `.submit` the same task object that you pass to `serve`.\n\n    Args:\n        - tasks: A list of tasks to serve. When a scheduled task run is found for a\n            given task, the task run will be submitted to the engine for execution.\n\n    Example:\n        ```python\n        from prefect import task\n        from prefect.task_server import serve\n\n        @task(log_prints=True)\n        def say(message: str):\n            print(message)\n\n        @task(log_prints=True)\n        def yell(message: str):\n            print(message.upper())\n\n        # starts a long-lived process that listens for scheduled runs of these tasks\n        if __name__ == \"__main__\":\n            serve(say, yell)\n        ```\n    \"\"\"\n    if not PREFECT_EXPERIMENTAL_ENABLE_TASK_SCHEDULING.value():\n        raise RuntimeError(\n            \"To enable task scheduling, set PREFECT_EXPERIMENTAL_ENABLE_TASK_SCHEDULING\"\n            \" to True.\"\n        )\n\n    task_server = TaskServer(*tasks)\n    try:\n        await task_server.start()\n\n    except StopTaskServer:\n        logger.info(\"Task server stopped.\")\n\n    except asyncio.CancelledError:\n        logger.info(\"Task server interrupted, stopping...\")\n", "path": "src/prefect/task_server.py"}], "after_files": [{"content": "import asyncio\nimport inspect\nimport os\nimport signal\nimport socket\nimport sys\nfrom concurrent.futures import ThreadPoolExecutor\nfrom contextlib import AsyncExitStack\nfrom contextvars import copy_context\nfrom typing import List\n\nimport anyio\nimport anyio.abc\nfrom websockets.exceptions import InvalidStatusCode\n\nfrom prefect import Task, get_client\nfrom prefect._internal.concurrency.api import create_call, from_sync\nfrom prefect.client.schemas.objects import TaskRun\nfrom prefect.client.subscriptions import Subscription\nfrom prefect.exceptions import Abort, PrefectHTTPStatusError\nfrom prefect.logging.loggers import get_logger\nfrom prefect.results import ResultFactory\nfrom prefect.settings import (\n    PREFECT_API_URL,\n    PREFECT_EXPERIMENTAL_ENABLE_TASK_SCHEDULING,\n    PREFECT_TASK_SCHEDULING_DELETE_FAILED_SUBMISSIONS,\n)\nfrom prefect.states import Pending\nfrom prefect.task_engine import run_task_async, run_task_sync\nfrom prefect.utilities.asyncutils import asyncnullcontext, sync_compatible\nfrom prefect.utilities.engine import emit_task_run_state_change_event, propose_state\nfrom prefect.utilities.processutils import _register_signal\n\nlogger = get_logger(\"task_server\")\n\n\nclass StopTaskServer(Exception):\n    \"\"\"Raised when the task server is stopped.\"\"\"\n\n    pass\n\n\ndef should_try_to_read_parameters(task: Task, task_run: TaskRun) -> bool:\n    \"\"\"Determines whether a task run should read parameters from the result factory.\"\"\"\n    new_enough_state_details = hasattr(\n        task_run.state.state_details, \"task_parameters_id\"\n    )\n    task_accepts_parameters = bool(inspect.signature(task.fn).parameters)\n\n    return new_enough_state_details and task_accepts_parameters\n\n\nclass TaskServer:\n    \"\"\"This class is responsible for serving tasks that may be executed in the background\n    by a task runner via the traditional engine machinery.\n\n    When `start()` is called, the task server will open a websocket connection to a\n    server-side queue of scheduled task runs. When a scheduled task run is found, the\n    scheduled task run is submitted to the engine for execution with a minimal `EngineContext`\n    so that the task run can be governed by orchestration rules.\n\n    Args:\n        - tasks: A list of tasks to serve. These tasks will be submitted to the engine\n            when a scheduled task run is found.\n    \"\"\"\n\n    def __init__(\n        self,\n        *tasks: Task,\n    ):\n        self.tasks: List[Task] = tasks\n\n        self.started: bool = False\n        self.stopping: bool = False\n\n        self._client = get_client()\n        self._exit_stack = AsyncExitStack()\n\n        if not asyncio.get_event_loop().is_running():\n            raise RuntimeError(\n                \"TaskServer must be initialized within an async context.\"\n            )\n\n        self._runs_task_group: anyio.abc.TaskGroup = anyio.create_task_group()\n        self._executor = ThreadPoolExecutor()\n\n    @property\n    def _client_id(self) -> str:\n        return f\"{socket.gethostname()}-{os.getpid()}\"\n\n    def handle_sigterm(self, signum, frame):\n        \"\"\"\n        Shuts down the task server when a SIGTERM is received.\n        \"\"\"\n        logger.info(\"SIGTERM received, initiating graceful shutdown...\")\n        from_sync.call_in_loop_thread(create_call(self.stop))\n\n        sys.exit(0)\n\n    @sync_compatible\n    async def start(self) -> None:\n        \"\"\"\n        Starts a task server, which runs the tasks provided in the constructor.\n        \"\"\"\n        _register_signal(signal.SIGTERM, self.handle_sigterm)\n\n        async with asyncnullcontext() if self.started else self:\n            logger.info(\"Starting task server...\")\n            try:\n                await self._subscribe_to_task_scheduling()\n            except InvalidStatusCode as exc:\n                if exc.status_code == 403:\n                    logger.error(\n                        \"Could not establish a connection to the `/task_runs/subscriptions/scheduled`\"\n                        f\" endpoint found at:\\n\\n {PREFECT_API_URL.value()}\"\n                        \"\\n\\nPlease double-check the values of your\"\n                        \" `PREFECT_API_URL` and `PREFECT_API_KEY` environment variables.\"\n                    )\n                else:\n                    raise\n\n    @sync_compatible\n    async def stop(self):\n        \"\"\"Stops the task server's polling cycle.\"\"\"\n        if not self.started:\n            raise RuntimeError(\n                \"Task server has not yet started. Please start the task server by\"\n                \" calling .start()\"\n            )\n\n        self.started = False\n        self.stopping = True\n\n        raise StopTaskServer\n\n    async def _subscribe_to_task_scheduling(self):\n        logger.info(\n            f\"Subscribing to tasks: {' | '.join(t.task_key.split('.')[-1] for t in self.tasks)}\"\n        )\n        async for task_run in Subscription(\n            model=TaskRun,\n            path=\"/task_runs/subscriptions/scheduled\",\n            keys=[task.task_key for task in self.tasks],\n            client_id=self._client_id,\n        ):\n            logger.info(f\"Received task run: {task_run.id} - {task_run.name}\")\n            self._runs_task_group.start_soon(self._submit_scheduled_task_run, task_run)\n\n    async def _submit_scheduled_task_run(self, task_run: TaskRun):\n        logger.debug(\n            f\"Found task run: {task_run.name!r} in state: {task_run.state.name!r}\"\n        )\n\n        task = next((t for t in self.tasks if t.task_key == task_run.task_key), None)\n\n        if not task:\n            if PREFECT_TASK_SCHEDULING_DELETE_FAILED_SUBMISSIONS.value():\n                logger.warning(\n                    f\"Task {task_run.name!r} not found in task server registry.\"\n                )\n                await self._client._client.delete(f\"/task_runs/{task_run.id}\")\n\n            return\n\n        # The ID of the parameters for this run are stored in the Scheduled state's\n        # state_details. If there is no parameters_id, then the task was created\n        # without parameters.\n        parameters = {}\n        if should_try_to_read_parameters(task, task_run):\n            parameters_id = task_run.state.state_details.task_parameters_id\n            task.persist_result = True\n            factory = await ResultFactory.from_autonomous_task(task)\n            try:\n                parameters = await factory.read_parameters(parameters_id)\n            except Exception as exc:\n                logger.exception(\n                    f\"Failed to read parameters for task run {task_run.id!r}\",\n                    exc_info=exc,\n                )\n                if PREFECT_TASK_SCHEDULING_DELETE_FAILED_SUBMISSIONS.value():\n                    logger.info(\n                        f\"Deleting task run {task_run.id!r} because it failed to submit\"\n                    )\n                    await self._client._client.delete(f\"/task_runs/{task_run.id}\")\n                return\n\n        logger.debug(\n            f\"Submitting run {task_run.name!r} of task {task.name!r} to engine\"\n        )\n\n        try:\n            state = await propose_state(\n                client=get_client(),  # TODO prove that we cannot use self._client here\n                state=Pending(),\n                task_run_id=task_run.id,\n            )\n        except Abort as exc:\n            logger.exception(\n                f\"Failed to submit task run {task_run.id!r} to engine\", exc_info=exc\n            )\n            return\n        except PrefectHTTPStatusError as exc:\n            if exc.response.status_code == 404:\n                logger.warning(\n                    f\"Task run {task_run.id!r} not found. It may have been deleted.\"\n                )\n                return\n            raise\n\n        if not state.is_pending():\n            logger.warning(\n                f\"Cancelling submission of task run {task_run.id!r} -\"\n                f\" server returned a non-pending state {state.type.value!r}.\"\n            )\n            return\n\n        emit_task_run_state_change_event(\n            task_run=task_run,\n            initial_state=task_run.state,\n            validated_state=state,\n        )\n\n        if task.isasync:\n            await run_task_async(\n                task=task,\n                task_run_id=task_run.id,\n                task_run=task_run,\n                parameters=parameters,\n                return_type=\"state\",\n            )\n        else:\n            context = copy_context()\n            future = self._executor.submit(\n                context.run,\n                run_task_sync,\n                task=task,\n                task_run_id=task_run.id,\n                task_run=task_run,\n                parameters=parameters,\n                return_type=\"state\",\n            )\n            await asyncio.wrap_future(future)\n\n    async def execute_task_run(self, task_run: TaskRun):\n        \"\"\"Execute a task run in the task server.\"\"\"\n        async with self if not self.started else asyncnullcontext():\n            await self._submit_scheduled_task_run(task_run)\n\n    async def __aenter__(self):\n        logger.debug(\"Starting task server...\")\n\n        if self._client._closed:\n            self._client = get_client()\n\n        await self._exit_stack.enter_async_context(self._client)\n        await self._exit_stack.enter_async_context(self._runs_task_group)\n        self._exit_stack.enter_context(self._executor)\n\n        self.started = True\n        return self\n\n    async def __aexit__(self, *exc_info):\n        logger.debug(\"Stopping task server...\")\n        self.started = False\n        await self._exit_stack.__aexit__(*exc_info)\n\n\n@sync_compatible\nasync def serve(*tasks: Task):\n    \"\"\"Serve the provided tasks so that their runs may be submitted to and executed.\n    in the engine. Tasks do not need to be within a flow run context to be submitted.\n    You must `.submit` the same task object that you pass to `serve`.\n\n    Args:\n        - tasks: A list of tasks to serve. When a scheduled task run is found for a\n            given task, the task run will be submitted to the engine for execution.\n\n    Example:\n        ```python\n        from prefect import task\n        from prefect.task_server import serve\n\n        @task(log_prints=True)\n        def say(message: str):\n            print(message)\n\n        @task(log_prints=True)\n        def yell(message: str):\n            print(message.upper())\n\n        # starts a long-lived process that listens for scheduled runs of these tasks\n        if __name__ == \"__main__\":\n            serve(say, yell)\n        ```\n    \"\"\"\n    if not PREFECT_EXPERIMENTAL_ENABLE_TASK_SCHEDULING.value():\n        raise RuntimeError(\n            \"To enable task scheduling, set PREFECT_EXPERIMENTAL_ENABLE_TASK_SCHEDULING\"\n            \" to True.\"\n        )\n\n    task_server = TaskServer(*tasks)\n    try:\n        await task_server.start()\n\n    except StopTaskServer:\n        logger.info(\"Task server stopped.\")\n\n    except asyncio.CancelledError:\n        logger.info(\"Task server interrupted, stopping...\")\n", "path": "src/prefect/task_server.py"}]}