refactor(trainer): simplify pod filtering per reviewer feedback

- Remove _select_best_pod_for_role() helper with status-priority logic
- Sort pods by creation_timestamp and use seen set to skip duplicates
- Add test case to parametrized test_get_job() per AGENTS.md pattern
- Remove standalone test_get_job_with_pod_restarts() function

Author: HKanoje

kubeflow/trainer/backends/kubernetes/backend.py

@@ -86,58 +86,6 @@ def verify_backend(self) -> None:
             )
             return
 
-    def _select_best_pod_for_role(
-        self, pods: list[models.IoK8sApiCoreV1Pod]
-    ) -> Optional[models.IoK8sApiCoreV1Pod]:
-        """
-        Select the best Pod for a role based on status priority and creation timestamp.
-
-        Priority order:
-        1. Running or Succeeded Pods (prefer most recent)
-        2. Failed Pods (prefer most recent)
-        3. Pending Pods (prefer most recent)
-        4. Unknown Pods (prefer most recent)
-        """
-        if not pods:
-            return None
-
-        # Pod status priority (higher number = higher priority)
-        status_priority = {
-            constants.POD_RUNNING: 4,  # Highest priority
-            constants.POD_SUCCEEDED: 3,  # Second highest
-            constants.POD_FAILED: 2,  # Third priority
-            constants.POD_PENDING: 1,  # Low priority
-            constants.POD_UNKNOWN: 0,  # Lowest priority
-        }
-
-        # Group Pods by status priority
-        pods_by_status = {}
-        for pod in pods:
-            status = pod.status.phase if pod.status else constants.POD_UNKNOWN
-            priority = status_priority.get(status, 0)
-
-            if priority not in pods_by_status:
-                pods_by_status[priority] = []
-            pods_by_status[priority].append(pod)
-
-        # Find the highest priority status that has Pods
-        highest_priority = max(pods_by_status.keys()) if pods_by_status else 0
-        candidate_pods = pods_by_status[highest_priority]
-
-        # Among Pods with the same status, select the most recent one
-        if len(candidate_pods) == 1:
-            return candidate_pods[0]
-
-        # Sort by creation timestamp (most recent first)
-        candidate_pods.sort(
-            key=lambda p: (
-                p.metadata.creation_timestamp or datetime.datetime.min.replace(tzinfo=timezone.utc)
-            ),
-            reverse=True,
-        )
-
-        return candidate_pods[0]
-
     def list_runtimes(self) -> list[types.Runtime]:
         """List available runtimes, preferring namespaced over cluster-scoped for duplicates.
 
@@ -656,62 +604,6 @@ def __get_runtime_from_cr(
             ),
         )
 
-    def _select_best_pod_for_role(
-        self, pods: list[models.IoK8sApiCoreV1Pod]
-    ) -> Optional[models.IoK8sApiCoreV1Pod]:
-        """
-        Select the best Pod for a role based on status priority and creation timestamp.
-
-        Priority order (higher priority = preferred):
-        1. Running or Succeeded Pods (equal priority, prefer most recent)
-        2. Failed Pods (prefer most recent)
-        3. Pending Pods (prefer most recent)
-        4. Unknown Pods (prefer most recent)
-
-        Both Running and Succeeded are considered healthy states with equal priority.
-        When multiple pods share the same priority, the most recently created pod is selected.
-        """
-        if not pods:
-            return None
-
-        # Pod status priority (higher number = higher priority)
-        # Running and Succeeded have equal priority as both are healthy states
-        status_priority = {
-            constants.POD_RUNNING: 4,  # Highest priority (healthy)
-            constants.POD_SUCCEEDED: 4,  # Highest priority (healthy)
-            constants.POD_FAILED: 2,  # Lower priority
-            constants.POD_PENDING: 1,  # Low priority
-            constants.POD_UNKNOWN: 0,  # Lowest priority
-        }
-
-        # Group Pods by status priority
-        pods_by_status = {}
-        for pod in pods:
-            status = pod.status.phase if pod.status else constants.POD_UNKNOWN
-            priority = status_priority.get(status, 0)
-
-            if priority not in pods_by_status:
-                pods_by_status[priority] = []
-            pods_by_status[priority].append(pod)
-
-        # Find the highest priority status that has Pods
-        highest_priority = max(pods_by_status.keys()) if pods_by_status else 0
-        candidate_pods = pods_by_status[highest_priority]
-
-        # Among Pods with the same priority, select the most recent one
-        if len(candidate_pods) == 1:
-            return candidate_pods[0]
-
-        # Sort by creation timestamp (most recent first)
-        candidate_pods.sort(
-            key=lambda p: (
-                p.metadata.creation_timestamp or datetime.datetime.min.replace(tzinfo=timezone.utc)
-            ),
-            reverse=True,
-        )
-
-        return candidate_pods[0]
-
     def _read_pod_logs(self, pod_name: str, container_name: str, follow: bool) -> Iterator[str]:
         """Read logs from a pod container."""
         try:
@@ -787,9 +679,21 @@ def __get_trainjob_from_cr(
             if not pod_list:
                 return trainjob
 
-            # Collect all Pods for this role
-            pods_by_role: dict[str, list[models.IoK8sApiCoreV1Pod]] = {}
-            for pod in pod_list.items:
+            # Sort Pods by creation timestamp (newest first) to ensure we process
+            # the most recent Pod for each role when duplicates exist (e.g., restarts)
+            pods = sorted(
+                pod_list.items,
+                key=lambda p: (
+                    p.metadata.creation_timestamp
+                    or datetime.datetime.min.replace(tzinfo=timezone.utc)
+                ),
+                reverse=True,
+            )
+
+            # Track seen role keys to skip duplicate Pods (older ones)
+            seen_roles: set[str] = set()
+
+            for pod in 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):
@@ -805,24 +709,13 @@ def __get_trainjob_from_cr(
                 else:
                     key = role
 
-                if key not in pods_by_role:
-                    pods_by_role[key] = []
-                pods_by_role[key].append(pod)
-
-            # Select the best Pod for each role using status-priority logic
-            selected_pods: dict[str, models.IoK8sApiCoreV1Pod] = {}
-            for role_key, pods in pods_by_role.items():
-                best_pod = self._select_best_pod_for_role(pods)
-                if best_pod:
-                    selected_pods[role_key] = best_pod
+                # Skip if we've already processed a Pod for this role (newer one)
+                if key in seen_roles:
+                    continue
+                seen_roles.add(key)
 
-            # Process only the selected Pod for each role
-            for _role_key, pod in selected_pods.items():
                 # Get the Initializer step.
-                if pod.metadata.labels[constants.JOBSET_RJOB_NAME_LABEL] in {
-                    constants.DATASET_INITIALIZER,
-                    constants.MODEL_INITIALIZER,
-                }:
+                if role in {constants.DATASET_INITIALIZER, constants.MODEL_INITIALIZER}:
                     trainjob.steps.append(
                         utils.get_trainjob_initializer_step(
                             pod.metadata.name,
@@ -831,17 +724,14 @@ def __get_trainjob_from_cr(
                         )
                     )
                 # Get the Node step.
-                elif pod.metadata.labels[constants.JOBSET_RJOB_NAME_LABEL] in {
-                    constants.LAUNCHER,
-                    constants.NODE,
-                }:
+                elif role in {constants.LAUNCHER, constants.NODE}:
                     trainjob.steps.append(
                         utils.get_trainjob_node_step(
                             pod.metadata.name,
                             pod.spec,
                             pod.status,
                             trainjob.runtime,
-                            pod.metadata.labels[constants.JOBSET_RJOB_NAME_LABEL],
+                            role,
                             int(pod.metadata.labels[constants.JOB_INDEX_LABEL]),
                         )
                     )