Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add _request_timeout to Kubernetes watches #15744

Merged
merged 3 commits into from
Oct 17, 2024
Merged

Add _request_timeout to Kubernetes watches #15744

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
6d23e66
add _request_timeout to watches
kevingrismore Oct 17, 2024
155d122
add _request_timeout to events replicator watch
kevingrismore Oct 17, 2024
62b922b
update tests
kevingrismore Oct 17, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions src/integrations/prefect-kubernetes/prefect_kubernetes/events.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import asyncio
from typing import Dict, List, Optional

import aiohttp
import kubernetes_asyncio
import kubernetes_asyncio.watch
from kubernetes_asyncio.client import ApiClient, V1Pod
Expand Down Expand Up @@ -81,6 +82,7 @@ async def _replicate_pod_events(self):
namespace=self._namespace,
label_selector=f"job-name={self._job_name}",
timeout_seconds=self._timeout_seconds,
_request_timeout=aiohttp.ClientTimeout(),
):
phase = event["object"].status.phase

Expand Down
3 changes: 3 additions & 0 deletions src/integrations/prefect-kubernetes/prefect_kubernetes/worker.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -907,6 +907,7 @@ async def _job_events(
func=batch_client.list_namespaced_job,
namespace=namespace,
field_selector=f"metadata.name={job_name}",
_request_timeout=aiohttp.ClientTimeout(),
**watch_kwargs,
):
yield event
Expand All @@ -915,6 +916,7 @@ async def _job_events(
job_list = await batch_client.list_namespaced_job(
namespace=namespace,
field_selector=f"metadata.name={job_name}",
_request_timeout=aiohttp.ClientTimeout(),
)

resource_version = job_list.metadata.resource_version
Expand Down Expand Up @@ -1118,6 +1120,7 @@ async def _get_job_pod(
namespace=configuration.namespace,
label_selector=f"job-name={job_name}",
timeout_seconds=configuration.pod_watch_timeout_seconds,
_request_timeout=aiohttp.ClientTimeout(),
):
pod: V1Pod = event["object"]
last_pod_name = pod.metadata.name
Expand Down
10 changes: 10 additions & 0 deletions src/integrations/prefect-kubernetes/tests/test_worker.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2301,6 +2301,7 @@ async def mock_stream(*args, **kwargs):
func=mock_batch_client.return_value.list_namespaced_job,
namespace=mock.ANY,
field_selector=mock.ANY,
_request_timeout=mock.ANY,
)

if job_timeout is not None:
Expand All @@ -2322,6 +2323,7 @@ async def mock_stream(*args, **kwargs):
namespace=mock.ANY,
label_selector=mock.ANY,
timeout_seconds=42,
_request_timeout=mock.ANY,
),
mock.call(**expected_job_call_kwargs),
]
Expand Down Expand Up @@ -2365,11 +2367,13 @@ async def mock_stream(*args, **kwargs):
namespace=mock.ANY,
label_selector=mock.ANY,
timeout_seconds=mock.ANY,
_request_timeout=mock.ANY,
),
mock.call(
func=mock_batch_client.return_value.list_namespaced_job,
namespace=mock.ANY,
field_selector=mock.ANY,
_request_timeout=mock.ANY,
# Note: timeout_seconds is excluded here
),
]
Expand Down Expand Up @@ -2410,11 +2414,13 @@ async def mock_stream(*args, **kwargs):
namespace="my-awesome-flows",
label_selector=mock.ANY,
timeout_seconds=60,
_request_timeout=mock.ANY,
),
mock.call(
func=mock_batch_client.return_value.list_namespaced_job,
namespace="my-awesome-flows",
field_selector=mock.ANY,
_request_timeout=mock.ANY,
),
]
)
Expand Down Expand Up @@ -2551,13 +2557,15 @@ async def mock_log_stream(*args, **kwargs):
namespace=mock.ANY,
label_selector=mock.ANY,
timeout_seconds=mock.ANY,
_request_timeout=mock.ANY,
),
# Starts with the full timeout minus the amount we slept streaming logs
mock.call(
func=mock_batch_client.return_value.list_namespaced_job,
field_selector=mock.ANY,
namespace=mock.ANY,
timeout_seconds=pytest.approx(50, 1),
_request_timeout=mock.ANY,
),
]
)
Expand Down Expand Up @@ -2776,12 +2784,14 @@ async def mock_stream(*args, **kwargs):
func=mock_batch_client.return_value.list_namespaced_job,
namespace=mock.ANY,
field_selector="metadata.name=mock-job",
_request_timeout=mock.ANY,
),
mock.call(
func=mock_batch_client.return_value.list_namespaced_job,
namespace=mock.ANY,
field_selector="metadata.name=mock-job",
resource_version="1",
_request_timeout=mock.ANY,
),
]
)
Expand Down
Loading