Skip to content

feat(client): Add context manager support (__enter__ / __exit__) to TrainerClient and SparkClient #669

Description

@Sumitha007

What you would like to be added?

Add context manager support (__enter__ and __exit__ magic methods) to TrainerClient, SparkClient, and core client classes in kubeflow/sdk.

This allows developers to manage client lifecycles using standard Python with statements:

from kubeflow.storage import TrainerClient

with TrainerClient() as client:
    job = client.train(
        name="mnist-job",
        num_workers=2,
    )
    client.wait_for_job(job.name)

### Why is this needed?

```markdown
1. **Resource Teardown & Cleanliness:** Currently, `TrainerClient` and `SparkClient` rely on manual cleanup or non-guaranteed garbage collection. If a script fails midway, open network connections or background polling sessions may remain dangling.
2. **Pythonic Usability:** Context managers are a standard Python pattern for client SDKs (such as `requests.Session`, `httpx.Client`, and `docker.from_env()`). Supporting `with` blocks reduces boilerplate `try...finally` code and makes error recovery more reliable for end-users.

### Love this feature?

Give it a 👍 We prioritize the features with most 👍

Metadata

Metadata

Assignees

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions