[nothariharan]

What this PR does / why we need it:

This PR adds a connect() API to the Kubeflow Trainer SDK. TrainerClient.connect() creates a TrainJob that keeps running and returns a TrainJobSession that runs commands in the trainer Pod of that job.

Today the only way to run a command in a training job is to do it by hand. You create a TrainJob whose trainer command keeps the Pod alive, for example sleep infinity, you wait for the job, you find the launcher Pod with the JobSet labels, you kubectl exec into that Pod, and you delete the job afterwards. That flow is documented in the Kubeflow Trainer integration guide in the website repository (kubeflow/website#4492). It works, but it needs kubectl, it needs knowledge of the JobSet labels, and it cannot be done from a notebook.

The API follows the same shape as SparkClient.connect() that was added for Kubeflow Spark, so both SDKs feel the same for users. The design is described in the KEP that is included in this PR.

The API

from kubeflow.trainer import CustomTrainerContainer, TrainerClient

client = TrainerClient()
runtime = client.get_runtime("mpi-interactive")

with client.connect(
    runtime=runtime,
    trainer=CustomTrainerContainer(image=runtime.trainer.image, num_nodes=4),
) as session:
    result = session.execute(["mpirun", "hostname"])
    print(result.stdout)
    print(result.stderr)
    print(result.returncode)

TrainerClient.connect(runtime, initializer, trainer, options, timeout, polling_interval) takes the same runtime, initializer, trainer and options arguments as train(). It creates the TrainJob, waits for the trainer Pod to run and returns a TrainJobSession.

TrainJobSession has three public members:

Member	Description
name	The name of the TrainJob the session is connected to.
execute(command, *, container=None)	Runs a command in the trainer Pod and returns an ExecResult. A string command runs with sh -c, a list of arguments runs directly, the same way kubectl exec runs a command.
close(timeout=180)	Deletes the TrainJob and waits until its Pods are deleted. It is idempotent and it runs when the session is used as a context manager.

ExecResult holds the stdout, the stderr and the returncode of the command. The returncode is None when the Kubernetes API server did not report an exit code.

How it works

Creating the job. connect() calls train() with the arguments of the caller. When the caller did not pass a TrainerCommand option, the SDK injects one with the sleep infinity command, so the trainer Pod stays alive. The command lives in kubeflow.trainer.constants.constants as INTERACTIVE_COMMAND. When the caller passes their own TrainerCommand, that command is used and nothing is injected.

Selecting the trainer Pod. The Pod that runs the trainer command is the Pod with the trainer.kubeflow.org/trainjob-ancestor-step: trainer label. For the MPI runtimes it is the launcher Pod, for the other runtimes it is the node Pod. Every runtime that ships with Kubeflow Trainer sets that label. The label matters for the MPI runtimes, because their node Pods only run the SSH server, and the launcher Pod is created only after the node Pods are ready. Without the label, a session could be created on a node Pod before the launcher Pod exists, and the commands would run on the wrong Pod. When no Pod has the label, the launcher Pod is used as a fallback.

Waiting for the session. connect() polls the Pods of the TrainJob until the trainer Pod is in the Running phase. A Pod in the Succeeded or the Failed phase raises a RuntimeError, because the job cannot run any command in that state. When the Pod does not start within the timeout, connect() raises a TimeoutError and the message tells the user that the job was created and how to delete it.

Running the commands. The commands run through the Kubernetes pods/exec API. The client uses kubernetes.stream.stream with _preload_content=False, which returns the WebSocket client that streams the output of the command. The client reads the standard output and the standard error channels until the connection is closed, and then it reads the exit code from the error channel. Executing commands in a Pod requires the pods/exec permission in the namespace of the TrainJob.

Closing the session. close() deletes the TrainJob and polls the Pods until they are gone. Waiting for the Pods matters, because the SSH keys of the MPI runtimes are mounted from a secret that is named after the job, so the old Pods have to be gone before a new job with the same name can start. When the Pods are still there after the timeout, the client logs a warning and returns instead of raising.

Backends. connect() is part of the RuntimeBackend interface. The Kubernetes backend implements it. The LocalProcess and the Container backends raise a NotImplementedError, the same way they do for the other Kubernetes only APIs.

What was tested

Unit tests

The new unit tests are in kubeflow/trainer/backends/kubernetes/backend_test.py, kubeflow/trainer/types/types_test.py, kubeflow/trainer/api/trainer_client_test.py, kubeflow/trainer/backends/localprocess/backend_test.py and kubeflow/trainer/backends/container/backend_test.py. They cover:

1. connect() creates the TrainJob with the sleep infinity command.
2. connect() does not override a TrainerCommand that the caller passed.
3. connect() rejects a trainer that is not a CustomTrainerContainer.
4. connect() waits for the trainer Pod and never connects to the node Pods of the MPI runtimes.
5. connect() times out when the Pod does not start, and it fails when the Pod has already finished.
6. connect() falls back to the launcher Pod when the ancestor label is not set.
7. execute() runs the command in the trainer Pod with the pods/exec API, with the right container, and with the command as a list and as a string.
8. execute() returns the output, the error output and the exit code. It reports an empty command, a missing Pod, a Pod that is not running and a failure of the exec call.
9. close() deletes the TrainJob, waits for the Pods and is idempotent.
10. The context manager closes the session, and execute() raises a RuntimeError in a closed session.
11. The client passes all the arguments to the backend, and the LocalProcess and the Container backends raise a NotImplementedError.

The full suite of the repository passes, and the pre-commit hooks pass:

$ uv run pytest kubeflow -q
729 passed

$ uv run pre-commit run --all-files
check yaml...............................................................Passed
fix end of files.........................................................Passed
trim trailing whitespace.................................................Passed
ruff check...............................................................Passed
ruff format..............................................................Passed

The documentation builds without new warnings, and the new classes are in the API reference.

End to end test on a real cluster

The end to end test ran against a real cluster, not against mocks:

1. A kind cluster with one control plane node and two worker nodes, on Kubernetes 1.37.0.
2. Kubeflow Trainer v2.3.0, installed from the released manifests, with the JobSet controller and the trainer controller running.
3. The mpi-interactive ClusterTrainingRuntime of the Kubeflow website guide. It has one launcher Pod and three node Pods, with OpenMPI and the SSH server of the MPI Operator, and runLauncherAsNode: true so the launcher counts as one of the MPI nodes.

The test created a session with four nodes and checked the following:

[info] runtime 'mpi-interactive', image 'local/mpi4-runtime:v1'
[ok] connected to 'mpi-interactive' in 14.7s
[info] TrainJob steps: [('node-0', 'Running'), ('node-1', 'Running'), ('node-2', 'Running'), ('node-3', 'Running')]
[ok] mpirun exit=0, hosts=['mpi-interactive-launcher-0-0', 'mpi-interactive-node-0-0', 'mpi-interactive-node-0-1', 'mpi-interactive-node-0-2']
[ok] string command -> 'hello from mpi-interactive-launcher-0-0'
[ok] exit code -> 3
[ok] stderr -> 'oops'
[ok] repeat exec -> 'mpi-interactive-launcher-0-0'
[ok] TrainJob deleted, all Pods gone (0.0s after close)
[ok] exec after close rejected: TrainJob session mpi-interactive is closed
[ok] reconnected: 'second session on mpi-interactive-launcher-0-0'
[done] e2e passed

So the checks were:

1. connect() returned once the launcher Pod was running, and the steps of the TrainJob were the launcher and the three nodes.
2. mpirun hostname returned the hostnames of the launcher Pod and of the three node Pods, so the MPI setup works and the command ran in the right Pod.
3. A string command ran with the shell, and the $(hostname) substitution worked.
4. A command that exits with 3 reported the exit code 3.
5. A command that writes to the standard error and exits with 1 reported the error output and the exit code.
6. Several commands ran in the same session.
7. Closing the session deleted the TrainJob and all of its Pods.
8. Running a command in a closed session raised a RuntimeError.
9. A new session with the same name could be created right after the first one was closed.

The end to end test found two problems in the first version of the implementation, and both of them are fixed and covered by a unit test now:

1. The Kubernetes client returns the output of the command as a string when stream() is called without _preload_content=False. The client has to ask for the WebSocket client, otherwise there is nothing to read the output from.
2. The first version of the Pod selection used the job-index label of the node Pods, and it connected to a node Pod of the MPI runtime before the launcher Pod existed. The commands then ran on the SSH server Pod instead of the launcher Pod. The selection now uses the trainjob-ancestor-step label of the trainer Pod, and connect() waits for that Pod.

One note about the test environment. The image of the guide, ghcr.io/kubeflow/trainer/deepspeed-runtime:v2.3.0, is 27.7 GB because it is based on CUDA, and it did not fit on the machine that ran the test. The test used a smaller image with the same OpenMPI version, the same SSH setup and the same Pod layout, and only the base image is different. The runtime manifest of the guide is otherwise unchanged.

Compatibility and limitations

1. The trainer must be a CustomTrainerContainer. The CustomTrainer and the BuiltinTrainer generate their own command, so their training process exits when the training finishes.
2. Only the Kubernetes backend supports connect(). The LocalProcess and the Container backends raise a NotImplementedError.
3. The session always creates a new TrainJob. Attaching to a session of another process is not part of this PR.
4. The commands run without a TTY, so the output is captured in the ExecResult instead of being streamed while the command runs.
5. The user of the SDK needs the pods/exec permission in the namespace of the TrainJob.

Documentation

1. The new classes are in the API reference, kubeflow.trainer.TrainJobSession and kubeflow.trainer.ExecResult.
2. There is a new guide page in docs/source/train/interactive-session.rst with the usage, the keep alive command and the requirements.
3. The KEP of this change is in proposals/868-interactive-trainjob-sessions.

Which issue(s) this PR fixes (optional, in Fixes #<issue number>, #<issue number>, ... format, will close the issue(s) when PR gets merged):

Fixes #868

Checklist:

• Docs included if any changes are user facing
• KEP added under proposals/868-interactive-trainjob-sessions

[github-actions]

🎉 Welcome to the Kubeflow SDK! 🎉

Thanks for opening your first PR! We're happy to have you as part of our community 🚀

Here's what happens next:

• If you haven't already, please check out our Contributing Guide for repo-specific guidelines and the Kubeflow Contributor Guide for general community standards
• Our team will review your PR soon! cc @kubeflow/kubeflow-sdk-team

Join the community:

• Slack: Join our #kubeflow-ml-experience and #kubeflow-trainer Slack channels
• Meetings: Attend the Kubeflow SDK and ML Experience bi-weekly meetings

Feel free to ask questions in the comments if you need any help or clarification!
Thanks again for contributing to Kubeflow! 🙏

[google-oss-prow]

[APPROVALNOTIFIER] This PR is NOT APPROVED

This pull-request has been approved by:
Once this PR has been reviewed and has the lgtm label, please assign kramaranya for approval. For more information see the Kubernetes Code Review Process.

The full list of commands accepted by this bot can be found here.

Details

Needs approval from an approver in each of these files:

• OWNERS

Approvers can indicate their approval by writing /approve in a comment
Approvers can cancel approval by writing /approve cancel in a comment

[christian-heusel]

/hold

This will first be discussed in a KEP, see #868 (comment)

[andreyvelich]

Thanks for this POC @nothariharan, but I will close it for now, since we need to gather more user stories and use-cases in the Trainer KEP first.
/close

[google-oss-prow]

@andreyvelich: Closed this PR.

Details

In response to this:

Thanks for this POC @nothariharan, but I will close it for now, since we need to gather more user stories and use-cases in the Trainer KEP first.
/close

Instructions for interacting with me using PR comments are available here. If you have questions or suggestions related to my behavior, please file an issue against the kubernetes/test-infra repository.

[nothariharan]

Thanks @andreyvelich and @christian-heusel, makes sense. We were not aware of the KEP flow, so apologies for jumping ahead with the code. The branch stays up in the fork, and we are happy to help with the KEP and the user stories whenever you need them.