From 83ca97e1aa52a781c5b7b9a54b724c9b3eff8a7e Mon Sep 17 00:00:00 2001
From: Dmitry Meyer <me@undef.im>
Date: Tue, 19 May 2026 08:54:31 +0000
Subject: [PATCH] Fix Kubernetes backend `utils.py` typing

See: https://github.com/lexdene/kubernetes-stubs/pull/11
---
 pyproject.toml                                |  2 +-
 .../core/backends/kubernetes/utils.py         | 25 ++++++++++++++-----
 2 files changed, 20 insertions(+), 7 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 4d19b23d7..c4db75319 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -152,7 +152,7 @@ dev = [
     "testcontainers>=4.9.2",
     "pytest-xdist>=3.6.1",
     "pyinstrument>=5.0.0",
-    "kubernetes-stubs-elephant-fork",
+    "kubernetes-stubs-elephant-fork>=35.0.0.post3",
     # docs
     "dstack[server]; python_version >= '3.11'",
     "dstack-plugin-server; python_version >= '3.11'",
diff --git a/src/dstack/_internal/core/backends/kubernetes/utils.py b/src/dstack/_internal/core/backends/kubernetes/utils.py
index b1ea5c7cb..054e48eba 100644
--- a/src/dstack/_internal/core/backends/kubernetes/utils.py
+++ b/src/dstack/_internal/core/backends/kubernetes/utils.py
@@ -1,6 +1,17 @@
 from collections.abc import Generator
 from contextlib import contextmanager
-from typing import Annotated, Any, Callable, Generic, Literal, Optional, Protocol, TypeVar, Union
+from typing import (
+    Annotated,
+    Any,
+    Callable,
+    Generic,
+    Literal,
+    Optional,
+    Protocol,
+    TypeVar,
+    Union,
+    cast,
+)
 
 import yaml
 from kubernetes.client import CoreV1Api, V1Status
@@ -9,9 +20,7 @@
     # XXX: This function is missing in the stubs package
     new_client_from_config_dict,  # pyright: ignore[reportAttributeAccessIssue]
 )
-
-# XXX: The watch module is missing in the stubs package
-from kubernetes.watch import Watch  # pyright: ignore[reportMissingImports]
+from kubernetes.watch import Watch
 from pydantic import Field
 from typing_extensions import ParamSpec, TypedDict
 from urllib3.exceptions import HTTPError
@@ -157,7 +166,8 @@ def watch_events(
     method: Callable[P, ObjectList[T]], *args: P.args, **kwargs: P.kwargs
 ) -> Generator[Generator[tuple[str, T], None, None], None, None]:
     watch = Watch()
-    gen = _watch_events_gen(watch.stream(method, *args, **kwargs))
+    inner_gen = cast(Generator[_EventDict[T], None, None], watch.stream(method, *args, **kwargs))
+    gen = _watch_events_gen(inner_gen)
     try:
         yield gen
     finally:
@@ -182,8 +192,11 @@ class _ErrorEventDict(TypedDict):
     object: V1Status
 
 
+_EventDict = Union[_StateEventDict[T], _BookmarkEventDict[T], _ErrorEventDict]
+
+
 def _watch_events_gen(
-    gen: Generator[Union[_StateEventDict[T], _BookmarkEventDict[T], _ErrorEventDict], None, None],
+    gen: Generator[_EventDict[T], None, None],
 ) -> Generator[tuple[str, T], None, None]:
     try:
         for event in gen: