What happened?
While reviewing the behavior of wait_for_job_status() across the available trainer backends, I noticed that LocalProcessBackend handles failed jobs differently from the Container and Kubernetes backends.
Both the Container and Kubernetes implementations immediately raise a RuntimeError when a TrainJob reaches the Failed state (unless TRAINJOB_FAILED is one of the expected target statuses).
The LocalProcess backend, however, continues polling until the timeout expires and finally raises a TimeoutError.
This makes the behavior inconsistent depending on which backend is being used.
Current behavior
LocalProcessBackend.wait_for_job_status() currently performs the following check:
if trainjob.status in status:
return trainjob
time.sleep(polling_interval)
...
raise TimeoutError(...)
If the job enters the Failed state while waiting for TRAINJOB_COMPLETE, the method keeps polling until the timeout is reached instead of failing immediately.
Behavior in other backends
Both of the other backend implementations explicitly handle this case.
Container backend
if constants.TRAINJOB_FAILED not in status and tj.status == constants.TRAINJOB_FAILED:
raise RuntimeError(...)
Kubernetes backend
if (
constants.TRAINJOB_FAILED not in status
and trainjob.status == constants.TRAINJOB_FAILED
):
raise RuntimeError(...)
Both implementations stop immediately once the job has failed.
Expected behavior
LocalProcessBackend.wait_for_job_status() should behave consistently with the other backends.
If the job reaches TRAINJOB_FAILED and the caller is not waiting for the failed state, the method should immediately raise a RuntimeError instead of waiting until the timeout expires.
Local investigation
While reviewing the implementation, I compared the three backend implementations:
✅ Container backend raises RuntimeError immediately.
✅ Kubernetes backend raises RuntimeError immediately.
❌ LocalProcess backend continues polling until timeout.
I also reviewed the existing LocalProcess backend tests.
Current wait_for_job_status tests only cover:
nonexistent job
invalid polling interval values
There does not appear to be a test covering the failed-job path, which is why this behavioral difference is currently not exercised.
While attempting to add such a regression test locally, I initially ran into an unrelated test import issue (TRAINJOB_FAILED import), but the implementation comparison itself already demonstrates the behavioral inconsistency.
Why this matters
This makes the same API behave differently depending on which backend is selected.
Code using wait_for_job_status() may correctly receive an immediate failure when using the Kubernetes or Container backend, but experience an unnecessary timeout when using the LocalProcess backend.
Aligning the LocalProcess implementation with the other backends would provide consistent behavior across all trainer backends.
What did you expect to happen?
I expected LocalProcessBackend.wait_for_job_status() to behave consistently with the Container and Kubernetes backends. If a TrainJob enters the TRAINJOB_FAILED state while waiting for another status (for example, TRAINJOB_COMPLETE), and TRAINJOB_FAILED is not included in the expected status set, the method should immediately raise a RuntimeError instead of continuing to poll until the timeout expires and raising a TimeoutError. This would provide consistent failure handling across all trainer backends.
Environment
Kubernetes version:
Kubeflow Trainer version:
$ kubectl get pods -n kubeflow -l app.kubernetes.io/name=trainer -o jsonpath="{.items[*].spec.containers[*].image}"
Kubeflow Python SDK version:
Impacted by this bug?
Give it a 👍 We prioritize the issues with most 👍
What happened?
While reviewing the behavior of wait_for_job_status() across the available trainer backends, I noticed that LocalProcessBackend handles failed jobs differently from the Container and Kubernetes backends.
Both the Container and Kubernetes implementations immediately raise a RuntimeError when a TrainJob reaches the Failed state (unless TRAINJOB_FAILED is one of the expected target statuses).
The LocalProcess backend, however, continues polling until the timeout expires and finally raises a TimeoutError.
This makes the behavior inconsistent depending on which backend is being used.
Current behavior
LocalProcessBackend.wait_for_job_status() currently performs the following check:
if trainjob.status in status:
return trainjob
time.sleep(polling_interval)
...
raise TimeoutError(...)
If the job enters the Failed state while waiting for TRAINJOB_COMPLETE, the method keeps polling until the timeout is reached instead of failing immediately.
Behavior in other backends
Both of the other backend implementations explicitly handle this case.
Container backend
if constants.TRAINJOB_FAILED not in status and tj.status == constants.TRAINJOB_FAILED:
raise RuntimeError(...)
Kubernetes backend
if (
constants.TRAINJOB_FAILED not in status
and trainjob.status == constants.TRAINJOB_FAILED
):
raise RuntimeError(...)
Both implementations stop immediately once the job has failed.
Expected behavior
LocalProcessBackend.wait_for_job_status() should behave consistently with the other backends.
If the job reaches TRAINJOB_FAILED and the caller is not waiting for the failed state, the method should immediately raise a RuntimeError instead of waiting until the timeout expires.
Local investigation
While reviewing the implementation, I compared the three backend implementations:
✅ Container backend raises RuntimeError immediately.
✅ Kubernetes backend raises RuntimeError immediately.
❌ LocalProcess backend continues polling until timeout.
I also reviewed the existing LocalProcess backend tests.
Current wait_for_job_status tests only cover:
nonexistent job
invalid polling interval values
There does not appear to be a test covering the failed-job path, which is why this behavioral difference is currently not exercised.
While attempting to add such a regression test locally, I initially ran into an unrelated test import issue (TRAINJOB_FAILED import), but the implementation comparison itself already demonstrates the behavioral inconsistency.
Why this matters
This makes the same API behave differently depending on which backend is selected.
Code using wait_for_job_status() may correctly receive an immediate failure when using the Kubernetes or Container backend, but experience an unnecessary timeout when using the LocalProcess backend.
Aligning the LocalProcess implementation with the other backends would provide consistent behavior across all trainer backends.
What did you expect to happen?
I expected LocalProcessBackend.wait_for_job_status() to behave consistently with the Container and Kubernetes backends. If a TrainJob enters the TRAINJOB_FAILED state while waiting for another status (for example, TRAINJOB_COMPLETE), and TRAINJOB_FAILED is not included in the expected status set, the method should immediately raise a RuntimeError instead of continuing to poll until the timeout expires and raising a TimeoutError. This would provide consistent failure handling across all trainer backends.
Environment
Kubernetes version:
Kubeflow Trainer version:
$ kubectl get pods -n kubeflow -l app.kubernetes.io/name=trainer -o jsonpath="{.items[*].spec.containers[*].image}"Kubeflow Python SDK version:
Impacted by this bug?
Give it a 👍 We prioritize the issues with most 👍