kubeflow / kubeflow/sdk

feat(trainer): Add retry with exponential backoff to get_job_logs() for transient Kubernetes API errors

Open
#551 3 comments 0 reactions 0 assignees View on GitHub
kind/feature needs-triage
Dominant language
Python
Stars
148
Forks
262
Avg merge
1d 2h
Merged PRs (30d)
1

Description

### What you would like to be added?

`TrainerClient.get_job_logs()` currently fails immediately on any exception raised while reading pod logs — including **transient** errors like connection resets, temporary API server unavailability, or `429 Too Many Requests` from the Kubernetes API server.

**Root cause location:** `_read_pod_logs()` in `kubeflow/trainer/backends/kubernetes/backend.py` (line 604):

\```python
def _read_pod_logs(self, pod_name: str, container_name: str, follow: bool) -> Iterator[str]:
"""Read logs from a pod container."""
try:
if follow:
log_stream = watch.Watch().stream(...)
yield from log_stream
else:
logs = self.core_api.read_namespaced_pod_log(...)
yield from logs.splitlines()
except Exception as e:
raise RuntimeError(
f"Failed to read logs for the pod {self.namespace}/{pod_name}"
) from e
\```

Any exception — transient or permanent — is caught once and immediately re-raised as `RuntimeError`. There's no distinction between a genuinely failed pod and a momentary network blip, and no retry attempt at all.

### Why is this needed?

`get_job_logs()` should retry on transient errors with exponential backoff before giving up, and should expose a configurable `max_retries` parameter. After retries are exhausted, it should raise a clear `RuntimeError` stating how many attempts were made.

## Proposed Implementation

**Dependencies:** Verified in `pyproject.toml` — current deps are `kubernetes`, `pydantic`, `kubeflow-trainer-api`, `kubeflow-katib-api`. `tenacity` is NOT currently a dependency. Since `time` is already imported in `backend.py`, this can be implemented with a manual retry loop and `time.sleep()` backoff, avoiding a new dependency.

**1. `kubeflow/trainer/backends/kubernetes/backend.py` (line 604) — add retry loop to `_read_pod_logs()`:**
\```python
def _read_pod_logs(
self, pod_name: str, container_name: str, follow: bool, max_retries: int = 3
) -> Iterator[str]:
"""Read logs from a pod container, retrying on transient errors."""
attempt = 0
while True:
attempt += 1
try:
if follow:
log_stream = watch.Watch().stream(
self.core_api.read_namespaced_pod_log,
name=pod_name,
namespace=self.namespace,
container=container_name,
follow=True,
)
yield from log_stream # type: ignore
else:
logs = self.core_api.read_namespaced_pod_log(
name=pod_name,
namespace=self.namespace,
container=container_name,
)
yield from logs.splitlines()
return

except Exception as e:
if attempt >= max_retries:
raise RuntimeError(
f"Failed to read logs for the pod {self.namespace}/{pod_name} "
f"after {max_retries} attempts"
) from e
backoff = 2 ** (attempt - 1) # 1s, 2s, 4s...
logger.debug(
f"Retrying log read for pod {pod_name} (attempt {attempt}/{max_retries}) "
f"after {backoff}s due to: {e}"
)
time.sleep(backoff)
\```

**2. Same file (line 424) — thread `max_retries` through `get_job_logs()`:**
\```python
def get_job_logs(
self,
name: str,
follow: bool = False,
step: str = constants.NODE + "-0",
max_retries: int = 3,
) -> Iterator[str]:
"""Get the TrainJob logs"""
...
yield from self._read_pod_logs(
pod_name=pod_name, container_name=container_name, follow=follow, max_retries=max_retries
)
\```

**3. `kubeflow/trainer/api/trainer_client.py` (line 184) — add `max_retries` to public API:**
\```python
def get_job_logs(
self,
name: str,
step: str = constants.NODE + "-0",
follow: bool | None = False,
max_retries: int = 3,
) -> Iterator[str]:
...
return self.backend.get_job_logs(name=name, follow=follow, step=step, max_retries=max_retries)
\```

**4. `kubeflow/trainer/backends/kubernetes/backend_test.py` — extend the existing `mock_read_namespaced_pod_log` (line 456) to simulate a transient failure that succeeds after N attempts, and add two new tests near `test_get_job_logs` (line 1492): one confirming retry-then-success, one confirming `RuntimeError` after exhausting `max_retries`.**

### Love this feature?

I'd like to work on this. I've traced the exact code path and drafted the implementation above (manual retry with exponential backoff, no new dependency). Could a maintainer please assign this to me? I'll follow up with a PR referencing this issue.

Contributor guide

Open the contributing guide

Research direction

Start with _read_pod_logs() and get_job_logs() in kubeflow/trainer/backends/kubernetes/backend.py, then trace the public method in kubeflow/trainer/api/trainer_client.py. Read the existing log tests in backend_test.py, especially around test_get_job_logs, and run them before making changes. Done means transient failures are retried, retry exhaustion reports the attempt count, max_retries is exposed through the API, and both success and exhaustion cases are tested.

Written by the indexing model from the issue text.

Assessment

Tech stack
kubernetes, python
Domain
api, backend, infrastructure, testing-qa
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Clearly specified
Newbie friendliness
68/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.