Repository: PrefectHQ/prefect
Pull request: make task worker resilient to errors
Previously any unhandled exceptions inside of `_submit_scheduled_task_run` would be raised and could lead to the worker shutting down. This particularly problematic because we re-raise a lot of errors in the task run engine. 

https://github.com/PrefectHQ/prefect/blob/main/src/prefect/task_engine.py#L472-L486

This PR adds handling around task submission so the worker can keep going even if errors are raised.

closes: https://github.com/PrefectHQ/prefect/issues/13911

Problem statement:
task worker shuts down on `Abort`
### First check

- [X] I added a descriptive title to this issue.
- [X] I used the GitHub search to find a similar issue and didn't find it.
- [X] I searched the Prefect documentation for this issue.
- [X] I checked that this issue is related to Prefect and not one of its dependencies.

### Bug summary

When using a task worker, the worker shuts down unexpectedly.

### Reproduction

```python3
Run this flow

# task_submitter.py
from prefect import flow
from task_worker import my_background_task


@flow
def my_flow():
    my_background_task.delay()


if __name__ == "__main__":
    my_flow()
```
then after the flow is done, run the task worker:
```
# task_worker.py
from prefect import task
from prefect.task_worker import serve


@task
def my_background_task():
    print("In a background task!")


if __name__ == "__main__":
    serve(my_background_task)
```
```


### Error

```python3
(demo-flows311) ➜  task_scheduling git:(main) ✗ python task_worker.py   
15:38:31.274 | INFO    | prefect.task_worker - Starting task worker...
15:38:31.275 | INFO    | prefect.task_worker - Subscribing to tasks: my_background_task-7c388d0bf59fa4394379971e841498bc
15:38:31.553 | INFO    | prefect.task_worker - Received task run: c7b3dcfc-3986-4c67-beae-c45146528688 - my_background_task-0
15:38:32.352 | ERROR   | Task run 'my_background_task-0' - Task run was aborted: The enclosing flow must be running to begin task execution.
15:38:32.355 | ERROR   | Task run 'my_background_task-0' - Finished in state Pending()
Please wait for all submitted tasks to complete before exiting your flow by calling `.wait()` on the `PrefectFuture` returned from your `.submit()` calls.

Example:

from prefect import flow, task

@task
def say_hello(name):
    print f"Hello, {name}!"

@flow
def example_flow():
    say_hello.submit(name="Marvin)
    say_hello.wait()

example_flow()

15:38:32.360 | ERROR   | prefect.task_worker - Task worker stopped with 1 exception:
The enclosing flow must be running to begin task execution.
(demo-flows311) ➜  task_scheduling git:(main)
```


### Versions

```Text
Version:             3.0.0rc2+33.g8c7fc247f8
API version:         0.8.4
Python version:      3.11.9
Git commit:          8c7fc247
Built:               Mon, Jun 10, 2024 1:14 PM
OS/Arch:             darwin/arm64
Profile:             prod-test
Server type:         cloud
Pydantic version:    2.7.3
```


### Additional context

_No response_

Hints from issue discussion:

Commit message: fix comment

Patch to review:
diff --git a/src/prefect/task_worker.py b/src/prefect/task_worker.py
index 0b7585da3f6c..8f9e1e2caef4 100644
--- a/src/prefect/task_worker.py
+++ b/src/prefect/task_worker.py
@@ -156,12 +156,24 @@ async def _subscribe_to_task_scheduling(self):
             client_id=self._client_id,
             base_url=base_url,
         ):
+            logger.info(f"Received task run: {task_run.id} - {task_run.name}")
             if self._limiter:
                 await self._limiter.acquire_on_behalf_of(task_run.id)
-            logger.info(f"Received task run: {task_run.id} - {task_run.name}")
             self._runs_task_group.start_soon(self._submit_scheduled_task_run, task_run)
 
     async def _submit_scheduled_task_run(self, task_run: TaskRun):
+        try:
+            await self._submit_and_start_run(task_run)
+        except BaseException as exc:
+            logger.exception(
+                f"Failed to submit task run {task_run.id!r}",
+                exc_info=exc,
+            )
+        finally:
+            if self._limiter:
+                self._limiter.release_on_behalf_of(task_run.id)
+
+    async def _submit_and_start_run(self, task_run: TaskRun):
         logger.debug(
             f"Found task run: {task_run.name!r} in state: {task_run.state.name!r}"
         )
@@ -266,8 +278,6 @@ async def _submit_scheduled_task_run(self, task_run: TaskRun):
                 context=run_context,
             )
             await asyncio.wrap_future(future)
-        if self._limiter:
-            self._limiter.release_on_behalf_of(task_run.id)
 
     async def execute_task_run(self, task_run: TaskRun):
         """Execute a task run in the task worker."""
diff --git a/tests/test_task_worker.py b/tests/test_task_worker.py
index e4c3b5202d14..a4fcf5e790d5 100644
--- a/tests/test_task_worker.py
+++ b/tests/test_task_worker.py
@@ -1,6 +1,7 @@
 import asyncio
 import signal
 import uuid
+from contextlib import contextmanager
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import anyio
@@ -172,6 +173,29 @@ async def test_task_worker_handles_deleted_task_run_submission(
     )
 
 
+async def test_task_worker_stays_running_on_errors(monkeypatch):
+    # regression test for https://github.com/PrefectHQ/prefect/issues/13911
+    # previously any error with submitting the task run would be raised
+    # and uncaught, causing the task worker to stop and this test to fail
+
+    @task
+    def empty_task():
+        pass
+
+    @contextmanager
+    def always_error(*args, **kwargs):
+        raise ValueError("oops")
+
+    monkeypatch.setattr("prefect.task_engine.TaskRunEngine.start", always_error)
+
+    task_worker = TaskWorker(empty_task)
+
+    empty_task.apply_async()
+
+    with anyio.move_on_after(1):
+        await task_worker.start()
+
+
 @pytest.mark.usefixtures("mock_task_worker_start")
 class TestServe:
     async def test_serve_basic_sync_task(self, foo_task, mock_task_worker_start):


Code review comments:
diff_hunk: @@ -156,12 +156,24 @@ async def _subscribe_to_task_scheduling(self):
             client_id=self._client_id,
             base_url=base_url,
         ):
+            logger.info(f"Received task run: {task_run.id} - {task_run.name}")
             if self._limiter:
                 await self._limiter.acquire_on_behalf_of(task_run.id)
-            logger.info(f"Received task run: {task_run.id} - {task_run.name}")
             self._runs_task_group.start_soon(self._submit_scheduled_task_run, task_run)
 
     async def _submit_scheduled_task_run(self, task_run: TaskRun):
+        try:
+            await self._submit_and_start_run(task_run)
+        except BaseException as exc:
+            logger.exception(
+                f"Failed to submit task run {task_run.id!r}",
+                exc_info=exc,
+            )
+        finally:
+            if self._limiter:
+                self._limiter.release_on_behalf_of(task_run.id)
+
+    async def _submit_and_start_run(self, task_run: TaskRun):
path: src/prefect/task_worker.py
text: @author:
name suggestions appreciated for this or the parent function, I didn't want to inline everything in a giant try/except

@user1:
Maybe `_safe_submit_scheduled_task_run` and `_submit_scheduled_task_run`?