Repository: PrefectHQ/prefect
Pull request: Add `KeepAliveClientRequest` class for k8s async client
<!-- 
Thanks for opening a pull request to Prefect! 
If this is your first contribution, please make sure to review our contribution guidelines: https://docs.prefect.io/latest/contributing/overview/
-->

<!-- Include an overview of the proposed changes here -->
Closes #14954 

Swaps the request class on the async k8s base client for one that has TCP keepalive enabled after the client is built. Refer to https://github.com/aio-libs/aiohttp/issues/3904#issuecomment-759205696

### Checklist
<!-- These boxes may be checked after opening the pull request. -->

- [x] This pull request references any related issue by including "closes `<link to issue>`"
  - If no issue exists and your change is not a small fix, please [create an issue](https://github.com/PrefectHQ/prefect/issues/new/choose) first.
- [x] If this pull request adds new functionality, it includes unit tests that cover the changes
- [ ] If this pull request removes docs files, it includes redirect settings in `mint.json`.
- [x] If this pull request adds functions or classes, it includes helpful docstrings.


Problem statement:
K8s worker fails monitoring flow and sets it to crashed.
### Bug summary

We are seeing an issue with the latest k8s worker monitoring flow runs. After some amount of time (5-10 minutes?) we see the worker fail to monitor the flow run with the message provided in additional context. The worker then marks the flow as failed. The flow continues to run until it tries to update its task status at which point it fails because its task status is already in a terminal state.

### Version info (`prefect version` output)

```Text
prefect-client: 2.20.0
prefect-kubernetes: 0.4.1
chart version: 2024.8.3144127
```


### Additional context

```
Error during task execution: Response payload is not completed: <TransferEncodingError: 400, message='Not enough data for satisfy transfer length header.'>
NoneType: None
```

Hints from issue discussion:
Thanks for the issue @meggers!

I'm trying to reproduce this issue, but I haven't been able to do so yet. Do you have any additional logs or tracebacks that you can share?
Thank you @desertaxle ! The worker logs don't have anything more than what I added for additional context. If there is a setting that would expose more information such as a stack trace from the logs definitely happy to try and get more info!
Hmmm, it looks like we're light on logging in the worker. I'll need to add some logging to get more info on the failure. While in there, I'll see if I can add some error handling and retries to the watch logic. I'll respond here once there's a new version to try!
Thank you! We are eager to resolve this, let us know! 
Other users are reporting the same error. We might have to handle issues ourselves. [See](https://github.com/tomplus/kubernetes_asyncio/issues/156).
I wonder if it is something like this: https://github.com/tomplus/kubernetes_asyncio/issues/235#issuecomment-1717331429

I am following up internally to see what we might have in the middle, but could be something like istio.

Wonder if it is possible to configure the k8s aiohttp client with keepalives? Example in https://github.com/aio-libs/aiohttp/issues/3904#issuecomment-759205696
@meggers That's a good callout, the sync version of the worker did have this, but it doesn't seem like the async version does.
Makes sense. I was able to confirm we have a timeout of ~5 minutes in our mesh which aligns with the timing we see with this issue. I think a keepalive would be helpful.
I'm not sure about all of the teams experiencing this but at least my team sees the following error immediately after the `400 TransferEncodingError`

```
Could not determine exit code for 'messy-python-hgtwd'.Exit code will be reported as -1.First container status info did not report an exit code.First container info: {'allocated_resources': None,
 'container_id': 'containerd://336b5ee0dc994e4fa2a635724c9e354f9f1b859b213f2c64af16959a7c822517',
 'image': 'r1k8sacrdev.azurecr.io/istio/proxyv2:1.21.5-distroless',
 'image_id': 'sha256:3b779807c2c39177a38a0713f6b692cee5241afe749a9ecd78171549f2eb7ecf',
 'last_state': {'running': None, 'terminated': None, 'waiting': None},
 'name': 'istio-proxy',
 'ready': True,
 'resources': None,
 'restart_count': 0,
 'started': True,
 'state': {'running': {'started_at': datetime.datetime(2024, 8, 15, 14, 27, 54, tzinfo=tzlocal())},
           'terminated': None,
           'waiting': None},'volume_mounts': None}

```
Which looks like it might be related to this bit here: https://github.com/PrefectHQ/prefect/blob/56d6fb0e6ab006ba12620a07bd8840e5a75e3770/src/integrations/prefect-kubernetes/prefect_kubernetes/worker.py#L1021

When we run our job pods there are 3 containers inside a given pod, and it appears that on occasion this is checking on the **Istio-Proxy** instead of the job container... 

For my jobs, when we encounter `400 TransferEncodingError` followed by `Could not determine exit code` like above, the UI shows the flow as `Crashed` until the job container completes, and at that exact moment the flow is marked `Successful` - we can see the color change from Blue -> Orange -> Green in the graph in the UI.

Previous, when my team was running versions `2024.5.23194919`, `2.16.5`, and `0.3.7` , we periodically would see the following error message:
```
An error occurred while monitoring flow run '31893f0f-ddc0-4b60-87cf-b58138515726'. The flow run will not be marked as failed, but an issue may have occurred.
Traceback (most recent call last):
  File "/usr/local/lib/python3.10/dist-packages/prefect/workers/base.py", line 908, in _submit_run_and_capture_errors
    result = await self.run(
  File "/usr/local/lib/python3.10/dist-packages/prefect_kubernetes/worker.py", line 599, in run
    status_code = await run_sync_in_worker_thread(
  File "/usr/local/lib/python3.10/dist-packages/prefect/utilities/asyncutils.py", line 136, in run_sync_in_worker_thread
    return await anyio.to_thread.run_sync(
  File "/usr/local/lib/python3.10/dist-packages/anyio/to_thread.py", line 33, in run_sync
    return await get_asynclib().run_sync_in_worker_thread(
  File "/usr/local/lib/python3.10/dist-packages/anyio/_backends/_asyncio.py", line 877, in run_sync_in_worker_thread
    return await future
  File "/usr/local/lib/python3.10/dist-packages/anyio/_backends/_asyncio.py", line 807, in run
    result = context.run(func, *args)
  File "/usr/local/lib/python3.10/dist-packages/prefect_kubernetes/worker.py", line 1028, in _watch_job
    for event in self._job_events(
  File "/usr/local/lib/python3.10/dist-packages/kubernetes/watch/watch.py", line 178, in stream
    for line in iter_resp_lines(resp):
  File "/usr/local/lib/python3.10/dist-packages/kubernetes/watch/watch.py", line 56, in iter_resp_lines
    for segment in resp.stream(amt=None, decode_content=False):
  File "/usr/local/lib/python3.10/dist-packages/urllib3/response.py", line 1057, in stream
    yield from self.read_chunked(amt, decode_content=decode_content)
  File "/usr/local/lib/python3.10/dist-packages/urllib3/response.py", line 1206, in read_chunked
    self._update_chunk_length()
  File "/usr/local/lib/python3.10/dist-packages/urllib3/response.py", line 1136, in _update_chunk_length
    raise ProtocolError("Response ended prematurely") from None
urllib3.exceptions.ProtocolError: Response ended prematurely
```

While this error showed up in the worker logs and the UI, the flow was still marked as Running - though the UI would stop updating about Subflow/Task status. Then, when the job container completed, the flow would be marked as Successful/Failed as appropriate and the UI would fill in.

I'm not sure if the Istio-Proxy container being referenced in the new error message is actually relevant to the issue, or if it's something else and then it's just further tripping up trying to write out an error message.

Commit message: update values, add docstring to class

Patch to review:
diff --git a/src/integrations/prefect-kubernetes/prefect_kubernetes/utilities.py b/src/integrations/prefect-kubernetes/prefect_kubernetes/utilities.py
index 879b8cb3e030..cb5a9d7549e8 100644
--- a/src/integrations/prefect-kubernetes/prefect_kubernetes/utilities.py
+++ b/src/integrations/prefect-kubernetes/prefect_kubernetes/utilities.py
@@ -4,7 +4,9 @@
 import sys
 from typing import Optional, TypeVar
 
-from kubernetes_asyncio.client import ApiClient
+from aiohttp import ClientResponse
+from aiohttp.client_reqrep import ClientRequest
+from aiohttp.connector import Connection
 from slugify import slugify
 
 # Note: `dict(str, str)` is the Kubernetes API convention for
@@ -14,34 +16,35 @@
 V1KubernetesModel = TypeVar("V1KubernetesModel")
 
 
-def enable_socket_keep_alive(client: ApiClient) -> None:
+class KeepAliveClientRequest(ClientRequest):
     """
-    Setting the keep-alive flags on the kubernetes client object.
-    Unfortunately neither the kubernetes library nor the urllib3 library which
-    kubernetes is using internally offer the functionality to enable keep-alive
-    messages. Thus the flags are added to be used on the underlying sockets.
+    aiohttp only directly implements socket keepalive for incoming connections
+    in its RequestHandler. For client connections, we need to set the keepalive
+    ourselves.
+
+    Refer to https://github.com/aio-libs/aiohttp/issues/3904#issuecomment-759205696
     """
 
-    socket_options = [(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)]
+    async def send(self, conn: "Connection") -> "ClientResponse":
+        sock = conn.protocol.transport.get_extra_info("socket")
+        sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
 
-    if hasattr(socket, "TCP_KEEPINTVL"):
-        socket_options.append((socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, 30))
+        if hasattr(socket, "TCP_KEEPIDLE"):
+            sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPIDLE, 30)
 
-    if hasattr(socket, "TCP_KEEPCNT"):
-        socket_options.append((socket.IPPROTO_TCP, socket.TCP_KEEPCNT, 6))
+        if hasattr(socket, "TCP_KEEPINTVL"):
+            sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, 6)
 
-    if hasattr(socket, "TCP_KEEPIDLE"):
-        socket_options.append((socket.IPPROTO_TCP, socket.TCP_KEEPIDLE, 6))
+        if hasattr(socket, "TCP_KEEPCNT"):
+            sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPCNT, 6)
 
-    if sys.platform == "darwin":
-        # TCP_KEEP_ALIVE not available on socket module in macOS, but defined in
-        # https://github.com/apple/darwin-xnu/blob/2ff845c2e033bd0ff64b5b6aa6063a1f8f65aa32/bsd/netinet/tcp.h#L215
-        TCP_KEEP_ALIVE = 0x10
-        socket_options.append((socket.IPPROTO_TCP, TCP_KEEP_ALIVE, 30))
+        if sys.platform == "darwin":
+            # TCP_KEEP_ALIVE not available on socket module in macOS, but defined in
+            # https://github.com/apple/darwin-xnu/blob/2ff845c2e033bd0ff64b5b6aa6063a1f8f65aa32/bsd/netinet/tcp.h#L215
+            TCP_KEEP_ALIVE = 0x10
+            sock.setsockopt(socket.IPPROTO_TCP, TCP_KEEP_ALIVE, 30)
 
-    client.rest_client.pool_manager.connection_pool_kw[
-        "socket_options"
-    ] = socket_options
+        return await super().send(conn)
 
 
 def _slugify_name(name: str, max_length: int = 45) -> Optional[str]:
diff --git a/src/integrations/prefect-kubernetes/prefect_kubernetes/worker.py b/src/integrations/prefect-kubernetes/prefect_kubernetes/worker.py
index d8190b684c75..a4d51d0ac502 100644
--- a/src/integrations/prefect-kubernetes/prefect_kubernetes/worker.py
+++ b/src/integrations/prefect-kubernetes/prefect_kubernetes/worker.py
@@ -159,10 +159,10 @@
 from prefect_kubernetes.credentials import KubernetesClusterConfig
 from prefect_kubernetes.events import KubernetesEventsReplicator
 from prefect_kubernetes.utilities import (
+    KeepAliveClientRequest,
     _slugify_label_key,
     _slugify_label_value,
     _slugify_name,
-    enable_socket_keep_alive,
 )
 
 MAX_ATTEMPTS = 3
@@ -637,10 +637,6 @@ async def _get_configured_kubernetes_client(
         Returns a configured Kubernetes client.
         """
         client = None
-        if os.environ.get(
-            "PREFECT_KUBERNETES_WORKER_ADD_TCP_KEEPALIVE", "TRUE"
-        ).strip().lower() in ("true", "1"):
-            enable_socket_keep_alive(client)
 
         if configuration.cluster_config:
             config_dict = configuration.cluster_config.config
@@ -657,6 +653,12 @@ async def _get_configured_kubernetes_client(
             except config.ConfigException:
                 # If in-cluster config fails, load the local kubeconfig
                 client = await config.new_client_from_config()
+
+        if os.environ.get(
+            "PREFECT_KUBERNETES_WORKER_ADD_TCP_KEEPALIVE", "TRUE"
+        ).strip().lower() in ("true", "1"):
+            client.rest_client.pool_manager._request_class = KeepAliveClientRequest
+
         try:
             yield client
         finally:


Code review comments:
diff_hunk: @@ -14,34 +16,35 @@
 V1KubernetesModel = TypeVar("V1KubernetesModel")
 
 
-def enable_socket_keep_alive(client: ApiClient) -> None:
+class KeepAliveClientRequest(ClientRequest):
     """
-    Setting the keep-alive flags on the kubernetes client object.
-    Unfortunately neither the kubernetes library nor the urllib3 library which
-    kubernetes is using internally offer the functionality to enable keep-alive
-    messages. Thus the flags are added to be used on the underlying sockets.
+    aiohttp only directly implements socket keepalive for incoming connections
+    in its RequestHandler. For client connections, we need to set the keepalive
+    ourselves.
+
+    Refer to https://github.com/aio-libs/aiohttp/issues/3904#issuecomment-759205696
     """
 
-    socket_options = [(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)]
+    async def send(self, conn: "Connection") -> "ClientResponse":
path: src/integrations/prefect-kubernetes/prefect_kubernetes/utilities.py
text: @user1:
if we're actually importing these, we should be able to use them
```suggestion
    async def send(self, conn: Connection) -> ClientResponse:
```

@author:
true