Skip to content

[rhoai-2.25] fix(tests): bound exposing_contextmanager's portforward retry loop #2691

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

Open
jiridanek wants to merge 4 commits into
base: rhoai-2.25
Choose a base branch
from
Open

[rhoai-2.25] fix(tests): bound exposing_contextmanager's portforward retry loop #2691

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
51a0006
fix(tests): bound exposing_contextmanager's portforward retry loop
jiridanek Aug 2, 2026
fb85a53
fixup: don't assume TimeoutError means service-not-ready in proxy log
jiridanek Aug 2, 2026
824c53a
fixup: bound the individual portforward() call, not just the retry loop
jiridanek Aug 2, 2026
e691e4a
fixup: apply ruff format (previous commit missed socket_proxy.py)
jiridanek Aug 2, 2026
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
80 changes: 70 additions & 10 deletions tests/containers/kubernetes_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -505,23 +505,83 @@ def check_not_none(value: Any, message: str) -> None:
raise ValueError(message)


def _portforward_with_timeout(
core_v1_api: kubernetes.client.CoreV1Api, pod: kubernetes.client.models.V1Pod, timeout: float
) -> kubernetes.stream.ws_client.PortForward:
"""Runs kubernetes.stream.portforward() with a wall-clock bound.

The kubernetes client doesn't expose a per-call connect timeout for portforward() -- it calls
websocket.connect() with no timeout, so a wedged connection to the API server can hang this call
forever. The only alternative the library offers is a process-wide `websocket.setdefaulttimeout()`,
which isn't safe to use here since multiple SocketProxy instances (one per ImageDeployment) can be
making concurrent portforward() calls from different threads. So we run the call in a helper thread
and bound it with join(timeout) instead.

If the call doesn't return in time, it is abandoned rather than joined: the daemon thread and any
partially established connection are leaked for the remainder of the test process. That's an
accepted tradeoff for test infrastructure that's about to raise TimeoutError and fail the test
(and the process will exit soon after) rather than hang forever with no recovery.
"""
Comment on lines +508 to +524

@coderabbitai coderabbitai Bot Aug 2, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
set -euo pipefail
echo '--- callers of exposing_contextmanager and any outer retry (Wait.until) ---'
rg -n -C 5 'exposing_contextmanager|class Wait\b|def until\b' tests

echo '--- pytest fail-fast configuration (addopts, -x, maxfail) ---'
rg -n 'addopts|maxfail|-x\b' pyproject.toml pytest.ini tox.ini .github 2>/dev/null || true

Repository: red-hat-data-services/notebooks

Length of output: 5171

🏁 Script executed:

#!/bin/bash
set -euo pipefail
printf '%s\n' '--- Wait.until implementation ---'
sed -n '380,455p' tests/containers/kubernetes_utils.py
printf '%s\n' '--- exposing_contextmanager and timeout helper ---'
sed -n '500,595p' tests/containers/kubernetes_utils.py
printf '%s\n' '--- SocketProxy retry loop and cleanup ---'
sed -n '1,135p' tests/containers/socket_proxy.py
printf '%s\n' '--- pytest configuration and workflow invocations ---'
cat pytest.ini
rg -n -C 3 'pytest(\s|$)|pytest\.|make test|make[[:space:]]+test' Makefile .github tests 2>/dev/null | head -200

Repository: red-hat-data-services/notebooks

Length of output: 28040

🏁 Script executed:

#!/bin/bash
set -euo pipefail
printf '%s\n' '--- remainder of Wait.until ---'
sed -n '430,510p' tests/containers/kubernetes_utils.py
printf '%s\n' '--- all Wait.until call sites ---'
rg -n -C 8 'Wait\.until\(' tests
printf '%s\n' '--- ImageDeployment setup and lifecycle around SocketProxy ---'
sed -n '210,285p' tests/containers/kubernetes_utils.py
printf '%s\n' '--- SocketProxy callers and client retry references ---'
rg -n -C 8 'get_actual_port|listen_and_serve_until_canceled|cancellation_token|remote_socket_factory|Wait\.until' tests/containers

Repository: red-hat-data-services/notebooks

Length of output: 28122

🏁 Script executed:

#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from __future__ import annotations

import ast
from pathlib import Path

utils = Path("tests/containers/kubernetes_utils.py").read_text()
proxy = Path("tests/containers/socket_proxy.py").read_text()
pytest_ini = Path("pytest.ini").read_text()

u = ast.parse(utils)
p = ast.parse(proxy)

def calls(tree: ast.AST, name: str) -> list[int]:
    return [
        node.lineno
        for node in ast.walk(tree)
        if isinstance(node, ast.Call)
        and isinstance(node.func, ast.Attribute)
        and node.func.attr == name
    ]

wait_until_calls = calls(u, "until")
portforward_calls = [
    node.lineno
    for node in ast.walk(u)
    if isinstance(node, ast.Call)
    and isinstance(node.func, ast.Attribute)
    and node.func.attr == "portforward"
]
print(f"Wait.until call lines: {wait_until_calls}")
print(f"kubernetes.stream.portforward call lines: {portforward_calls}")
print(f"pytest fail-fast options present: {any(x in pytest_ini for x in ('-x', '--exitfirst', 'maxfail'))}")

for node in ast.walk(p):
    if isinstance(node, ast.Try):
        caught = [
            handler.type.names[0]
            for handler in node.handlers
            if isinstance(handler.type, ast.Tuple)
            for _ in [0]
        ]
        if any(isinstance(handler.type, ast.Name) and handler.type.id == "TimeoutError" for handler in node.handlers):
            print(f"SocketProxy catches TimeoutError at line {node.lineno}")
            print(f"  handler types: {caught or ['TimeoutError']}")
            print(f"  handler contains continue: {any(isinstance(x, ast.Continue) for h in node.handlers for x in ast.walk(h))}")

for node in ast.walk(u):
    if isinstance(node, ast.FunctionDef) and node.name == "_portforward_with_timeout":
        thread_starts = [
            n.lineno for n in ast.walk(node)
            if isinstance(n, ast.Call)
            and isinstance(n.func, ast.Attribute)
            and n.func.attr == "start"
        ]
        joins = [
            n.lineno for n in ast.walk(node)
            if isinstance(n, ast.Call)
            and isinstance(n.func, ast.Attribute)
            and n.func.attr == "join"
        ]
        raises = [
            n.lineno for n in ast.walk(node)
            if isinstance(n, ast.Raise)
        ]
        print(f"_portforward_with_timeout starts threads at: {thread_starts}")
        print(f"_portforward_with_timeout joins at: {joins}")
        print(f"_portforward_with_timeout raises at: {raises}")
PY

Repository: red-hat-data-services/notebooks

Length of output: 450

🏁 Script executed:

#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
import re

utils = Path("tests/containers/kubernetes_utils.py").read_text()
proxy = Path("tests/containers/socket_proxy.py").read_text()
pytest_ini = Path("pytest.ini").read_text()

print("Wait.until call count:", len(re.findall(r"\bWait\.until\s*\(", utils)))
print("portforward call count:", len(re.findall(r"\bkubernetes\.stream\.portforward\s*\(", utils)))
print("SocketProxy TimeoutError handling:", bool(re.search(
    r"except\s*\([^)]*\bTimeoutError\b[^)]*\).*?:.*?continue",
    proxy,
    re.S,
)))
print("pytest fail-fast option configured:", bool(re.search(
    r"(?m)^\s*(?:addopts\s*=.*(?:^|\s)(?:-x|--exitfirst)|(?:maxfail\s*=))",
    pytest_ini,
)))
print("helper starts daemon thread:", bool(re.search(
    r"thread\s*=\s*threading\.Thread\(.*daemon=True\).*?thread\.start\(\)",
    utils,
    re.S,
)))
print("helper joins without a later join:", bool(re.search(
    r"thread\.join\(timeout\).*?if not pf_result:.*?raise TimeoutError",
    utils,
    re.S,
)))
PY

Repository: red-hat-data-services/notebooks

Length of output: 376

Close abandoned port-forward attempts.

SocketProxy catches TimeoutError and continues serving clients, so the exception does not fail the test or end the pytest process. The readiness Wait.until retries for up to TIMEOUT_2MIN; each retry can leave a daemon thread or a late-returning, unclosed PortForward. Close abandoned results when they complete, or otherwise bound their lifetime.

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/containers/kubernetes_utils.py` around lines 508 - 524, Update
_portforward_with_timeout so timed-out port-forward attempts are eventually
closed even when callers catch TimeoutError and retry. Track the helper thread’s
late result and close any returned PortForward, and ensure the abandoned worker
or connection has a bounded lifetime rather than leaking across readiness
retries.

pf_result: list[kubernetes.stream.ws_client.PortForward] = []
error_result: list[Exception] = []

def _target() -> None:
try:
pf = kubernetes.stream.portforward(
api_method=core_v1_api.connect_get_namespaced_pod_portforward,
name=pod.metadata.name,
namespace=pod.metadata.namespace,
ports=",".join(str(p) for p in [8888]),
)
pf_result.append(pf)
except Exception as e: # propagated to the caller below, not swallowed
error_result.append(e)

thread = threading.Thread(target=_target, daemon=True)
thread.start()
thread.join(timeout)
if error_result:
raise error_result[0]
if not pf_result:
raise TimeoutError(
f"kubernetes.stream.portforward() to {pod.metadata.name} did not return within {timeout:.1f}s"
)
return pf_result[0]


@contextlib.contextmanager
def exposing_contextmanager(
core_v1_api: kubernetes.client.CoreV1Api, pod: kubernetes.client.models.V1Pod
core_v1_api: kubernetes.client.CoreV1Api,
pod: kubernetes.client.models.V1Pod,
timeout: float = 30,
) -> Generator[socket]:
# If we e.g., specify the wrong port, the pf = portforward() call succeeds,
# but pf.connected will later flip to False
# we need to check that _everything_ works before moving on
pf = None
s = None
#
# https://github.com/red-hat-data-services/notebooks/issues/2684: bound this retry loop (and each
# individual portforward() attempt via _portforward_with_timeout, since a single hung call would
# otherwise never let this loop's own deadline check run again) so a pod that never becomes
# reachable can't block the single-threaded SocketProxy forever and starve every later
# Wait.until retry.
deadline = time.monotonic() + timeout
pf: kubernetes.stream.ws_client.PortForward | None = None
s: kubernetes.stream.ws_client.PortForward._Port._Socket | socket.socket | None = None
while not pf or not pf.connected or not s:
pf: kubernetes.stream.ws_client.PortForward = kubernetes.stream.portforward(
api_method=core_v1_api.connect_get_namespaced_pod_portforward,
name=pod.metadata.name,
namespace=pod.metadata.namespace,
ports=",".join(str(p) for p in [8888]),
)
s: kubernetes.stream.ws_client.PortForward._Port._Socket | socket.socket | None = pf.socket(8888)
remaining = deadline - time.monotonic()
if remaining <= 0:
if s is not None:
s.close()
if pf is not None:
pf.close()
raise TimeoutError(f"Failed to establish a working portforward to {pod.metadata.name} within {timeout}s")
if s is not None:
s.close()
s = None
if pf is not None:
pf.close()
pf = _portforward_with_timeout(core_v1_api, pod, remaining)
s = pf.socket(8888)
Comment thread
coderabbitai[bot] marked this conversation as resolved.
assert s, "Failed to establish connection"

try:
Expand Down
8 changes: 6 additions & 2 deletions tests/containers/socket_proxy.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -94,10 +94,14 @@ def listen_and_serve_until_canceled(self):
try:
# handle client synchronously, which means that there can be at most one at a time
self._handle_client(client_socket)
except (BrokenPipeError, ConnectionResetError) as e:
except (BrokenPipeError, ConnectionResetError, TimeoutError) as e:
# BrokenPipeError happens when the proxy connects to the pod, but the service inside is not yet listening.
# TimeoutError is raised by remote_socket_factory() (e.g. exposing_contextmanager) when it can't
# establish a working connection within its own bound -- see https://github.com/red-hat-data-services/notebooks/issues/2684.
# The client (Wait.until) will retry.
logging.info(f"Proxy connection to remote failed, likely due to service not being ready: {e}")
logging.info(
f"Proxy connection to remote failed, will retry on the next connection attempt: {e}"
)
# The client_socket is closed by the `with` statement in `_handle_client`.
# We continue the loop to accept the next connection attempt.
continue
Expand Down
Loading