fix(trainer): deduplicate restarted TrainJob pods - #632
Conversation
|
🎉 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:
Join the community:
Feel free to ask questions in the comments if you need any help or clarification! |
Signed-off-by: Vidyansh <ezboi2312@gmail.com>
ff9e826 to
fc80755
Compare
|
@andreyvelich Review this |
andreyvelich
left a comment
andreyvelich
left a comment
There was a problem hiding this comment.
Thanks for tackling this, @reckless-sherixx! This review was performed with the help of AI tools.
| old_timestamp = datetime.datetime(2025, 6, 1, 10, 0, 0) | ||
| new_timestamp = datetime.datetime(2025, 6, 1, 11, 0, 0) | ||
| old_pods = get_mock_pod_list().items | ||
| node_1_pod = copy.deepcopy(old_pods[-1]) |
There was a problem hiding this comment.
Select the node pod by role instead of position — old_pods[-1] silently breaks if get_mock_pod_list() is ever reordered.
| node_1_pod = copy.deepcopy(old_pods[-1]) | |
| node_0_pod = next( | |
| p for p in old_pods if p.metadata.labels[constants.JOBSET_RJOB_NAME_LABEL] == constants.NODE | |
| ) | |
| node_1_pod = copy.deepcopy(node_0_pod) |
| step.pod_name = f"{step.pod_name}-restarted" | ||
| step.status = constants.POD_PENDING | ||
|
|
||
| node_1_step = copy.deepcopy(train_job.steps[-1]) |
There was a problem hiding this comment.
Same positional fragility — clone the node-0 step by name rather than by index.
| node_1_step = copy.deepcopy(train_job.steps[-1]) | |
| node_0_step = next(s for s in train_job.steps if s.name == "node-0") | |
| node_1_step = copy.deepcopy(node_0_step) |
| sorted_pods = sorted( | ||
| pod_list.items, | ||
| key=lambda pod: ( | ||
| pod.metadata is not None and pod.metadata.creation_timestamp is not None, |
There was a problem hiding this comment.
This is not None guard exists to survive Pods with no creation_timestamp, but no test exercises it. A future simplification to key=lambda pod: pod.metadata.creation_timestamp would pass every test yet raise TypeError: '<' not supported between 'datetime' and 'NoneType' in production. Worth adding a case where one restarted Pod has creation_timestamp=None and the timestamped Pod still wins.
|
/ok-to-test |
Signed-off-by: Vidyansh <ezboi2312@gmail.com>
andreyvelich
left a comment
andreyvelich
left a comment
There was a problem hiding this comment.
Thanks for this fix @reckless-sherixx!
/lgtm
/approve
|
[APPROVALNOTIFIER] This PR is APPROVED This pull-request has been approved by: andreyvelich The full list of commands accepted by this bot can be found here. The pull request process is described here DetailsNeeds approval from an approver in each of these files:
Approvers can indicate their approval by writing |
Summary
Fixes #25 by returning only the newest Pod for each TrainJob component in
get_job().creation_timestampdescending and keeps the first occurrence.role-indexfor launcher/node keys.node-0andnode-1, with newer Pending Pods taking precedence over older Running Pods.Why
Kubernetes can recreate Pods under restart policies. Without de-duplication,
get_job()returns stale and current Pods for the same component, producing conflicting statuses.Validation
python -m uv lock --checkpython -m uv run ruff check --show-fixes --output-format=github .python -m uv run ruff format --check kubeflowpython -m uv run ty check kubeflow/hubpython -m uv run coverage run --source=kubeflow -m pytest ./kubeflow/(629 passed)Supersedes the focused deduplication work in #160 without its unrelated changes.