fix(trainer): deduplicate restarted TrainJob pods

Author: reckless-sherixx

kubeflow/trainer/backends/kubernetes/backend.py

@@ -698,14 +698,32 @@ def __get_trainjob_from_cr(
             if not pod_list:
                 return trainjob
 
-            for pod in pod_list.items:
+            sorted_pods = sorted(
+                pod_list.items,
+                key=lambda pod: (
+                    pod.metadata is not None and pod.metadata.creation_timestamp is not None,
+                    pod.metadata.creation_timestamp if pod.metadata else None,
+                ),
+                reverse=True,
+            )
+            seen_step_keys: set[str] = set()
+            for pod in sorted_pods:
                 # Pod must have labels to detect the TrainJob step.
                 # Every Pod always has a single TrainJob step.
                 if not (pod.metadata and pod.metadata.name and pod.metadata.labels and pod.spec):
                     raise Exception(f"TrainJob Pod is invalid: {pod}")
 
+                role = pod.metadata.labels[constants.JOBSET_RJOB_NAME_LABEL]
+                step_key = role
+                if role in {constants.LAUNCHER, constants.NODE}:
+                    step_key = f"{role}-{pod.metadata.labels[constants.JOB_INDEX_LABEL]}"
+
+                if step_key in seen_step_keys:
+                    continue
+                seen_step_keys.add(step_key)
+
                 # Get the Initializer step.
-                if pod.metadata.labels[constants.JOBSET_RJOB_NAME_LABEL] in {
+                if role in {
                     constants.DATASET_INITIALIZER,
                     constants.MODEL_INITIALIZER,
                 }:
@@ -717,7 +735,7 @@ def __get_trainjob_from_cr(
                         )
                     )
                 # Get the Node step.
-                elif pod.metadata.labels[constants.JOBSET_RJOB_NAME_LABEL] in {
+                elif role in {
                     constants.LAUNCHER,
                     constants.NODE,
                 }:
@@ -727,7 +745,7 @@ def __get_trainjob_from_cr(
                             pod.spec,
                             pod.status,
                             trainjob.runtime,
-                            pod.metadata.labels[constants.JOBSET_RJOB_NAME_LABEL],
+                            role,
                             int(pod.metadata.labels[constants.JOB_INDEX_LABEL]),
                         )
                     )

kubeflow/trainer/backends/kubernetes/backend_test.py

@@ -18,6 +18,7 @@
 It tests KubernetesBackend's behavior across job listing, resource creation etc.
 """
 
+import copy
 from dataclasses import asdict
 import datetime
 import logging
@@ -71,6 +72,7 @@
 FAIL_LOGS = "fail_logs"
 LIST_RUNTIMES = "list_runtimes"
 BASIC_TRAIN_JOB_NAME = "basic-job"
+JOB_WITH_POD_RESTARTS = "job-with-pod-restarts"
 TRAIN_JOBS = "trainjobs"
 TRAIN_JOB_WITH_BUILT_IN_TRAINER = "train-job-with-built-in-trainer"
 TRAIN_JOB_WITH_CUSTOM_TRAINER = "train-job-with-custom-trainer"
@@ -128,8 +130,13 @@ def conditional_error_handler(*args, **kwargs):
 
 
 def list_namespaced_pod_response(*args, **kwargs):
-    """Return mock pod list response."""
-    pod_list = get_mock_pod_list()
+    """Return a mock pod list response for the requested TrainJob."""
+    label_selector = kwargs.get("label_selector", "")
+    pod_list = (
+        get_mock_pod_list_with_restarts()
+        if JOB_WITH_POD_RESTARTS in label_selector
+        else get_mock_pod_list()
+    )
     mock_thread = Mock()
     mock_thread.get.return_value = pod_list
     return mock_thread
@@ -210,6 +217,30 @@ def get_mock_pod_list():
     )
 
 
+def get_mock_pod_list_with_restarts() -> models.IoK8sApiCoreV1PodList:
+    """Create Pods where newer replacements share the same TrainJob component roles."""
+    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])
+    node_1_pod.metadata.name = "node-1-pod"
+    node_1_pod.metadata.labels[constants.JOB_INDEX_LABEL] = "1"
+    old_pods.append(node_1_pod)
+    restarted_pods = []
+
+    for old_pod in old_pods:
+        old_pod.metadata.creation_timestamp = old_timestamp
+        old_pod.metadata.labels[constants.JOBSET_NAME_LABEL] = JOB_WITH_POD_RESTARTS
+
+        restarted_pod = copy.deepcopy(old_pod)
+        restarted_pod.metadata.name = f"{old_pod.metadata.name}-restarted"
+        restarted_pod.metadata.creation_timestamp = new_timestamp
+        restarted_pod.status.phase = constants.POD_PENDING
+        restarted_pods.append(restarted_pod)
+
+    return models.IoK8sApiCoreV1PodList(items=[*old_pods, *restarted_pods])
+
+
 def get_resource_requirements() -> models.IoK8sApiCoreV1ResourceRequirements:
     """Create a mock ResourceRequirements object for testing."""
     return models.IoK8sApiCoreV1ResourceRequirements(
@@ -726,6 +757,7 @@ def get_train_job_data_type(
         num_nodes=2,
         image="example.com/test-runtime",
     )
+
     trainer.set_command(constants.TORCH_COMMAND)
     return types.TrainJob(
         name=train_job_name,
@@ -763,6 +795,25 @@ def get_train_job_data_type(
     )
 
 
+def get_train_job_with_restarted_pods_data_type(
+    runtime_name: str,
+    train_job_name: str,
+) -> types.TrainJob:
+    """Create the expected TrainJob after newer replacement Pods are selected."""
+    train_job = get_train_job_data_type(runtime_name, train_job_name)
+
+    for step in train_job.steps:
+        step.pod_name = f"{step.pod_name}-restarted"
+        step.status = constants.POD_PENDING
+
+    node_1_step = copy.deepcopy(train_job.steps[-1])
+    node_1_step.name = "node-1"
+    node_1_step.pod_name = "node-1-pod-restarted"
+    train_job.steps.append(node_1_step)
+
+    return train_job
+
+
 def _run_verify_backend_with_core_api(core_api: Mock) -> tuple[list[str], int]:
     """Helper to run verify_backend and capture warning logs."""
 
@@ -1447,6 +1498,15 @@ def test_train(kubernetes_backend, test_case):
                 train_job_name=BASIC_TRAIN_JOB_NAME,
             ),
         ),
+        TestCase(
+            name="returns only the newest Pod for each TrainJob component",
+            expected_status=SUCCESS,
+            config={"name": JOB_WITH_POD_RESTARTS},
+            expected_output=get_train_job_with_restarted_pods_data_type(
+                runtime_name=TORCH_RUNTIME,
+                train_job_name=JOB_WITH_POD_RESTARTS,
+            ),
+        ),
         TestCase(
             name="timeout error when getting job",
             expected_status=FAILED,