trainer: add strict mode for get_job_logs when pod is missing

Signed-off-by: Shanhuizi Jiang <sjiang83@fordham.edu>

Author: sjiang83

kubeflow/optimizer/backends/kubernetes/backend.py

@@ -223,6 +223,7 @@ def get_job_logs(
         name: str,
         trial_name: str | None = None,
         follow: bool = False,
+        strict: bool = False,
     ) -> Iterator[str]:
         """Get the OptimizationJob logs from a Trial"""
         # Determine what trial to get logs from.
@@ -247,6 +248,8 @@ def get_job_logs(
                 pod_name = c.pod_name
                 break
         if pod_name is None:
+            if strict:
+                raise RuntimeError(f"No pod found for Trial {trial_name} step={step}")
             return
 
         container_name = constants.METRICS_COLLECTOR_CONTAINER

kubeflow/trainer/api/trainer_client.py

@@ -185,6 +185,7 @@ def get_job_logs(
         name: str,
         step: str = constants.NODE + "-0",
         follow: bool | None = False,
+        strict: bool = False,
     ) -> Iterator[str]:
         """Get logs from a specific step of a TrainJob.
 
@@ -200,6 +201,7 @@ def get_job_logs(
             name: Name of the TrainJob.
             step: Step of the TrainJob to collect logs from, like dataset-initializer or node-0.
             follow: Whether to stream logs in realtime as they are produced.
+            strict: If True, raise an error when no pod is found for the requested step.
 
         Returns:
             Iterator of log lines.
@@ -209,7 +211,7 @@ def get_job_logs(
             TimeoutError: Timeout to get a TrainJob.
             RuntimeError: Failed to get a TrainJob.
         """
-        return self.backend.get_job_logs(name=name, follow=follow, step=step)
+        return self.backend.get_job_logs(name=name, follow=follow, step=step, strict=strict)
 
     def get_job_events(self, name: str) -> list[types.Event]:
         """Get events for a TrainJob.

kubeflow/trainer/backends/base.py

@@ -64,6 +64,7 @@ def get_job_logs(
         name: str,
         follow: bool = False,
         step: str = constants.NODE + "-0",
+        strict: bool = False,
     ) -> Iterator[str]:
         raise NotImplementedError()
 

kubeflow/trainer/backends/container/backend.py

@@ -773,6 +773,7 @@ def get_job_logs(
         name: str,
         follow: bool = False,
         step: str = constants.NODE + "-0",
+        strict: bool = False,
     ) -> Iterator[str]:
         """Get logs for a training job by querying container runtime."""
         containers = self._get_job_containers(name)

kubeflow/trainer/backends/kubernetes/backend.py

@@ -37,6 +37,9 @@
 
 logger = logging.getLogger(__name__)
 
+class PodNotFoundError(RuntimeError):
+    """No pod resolved for the requested job/step."""
+    pass
 
 class KubernetesBackend(RuntimeBackend):
     def __init__(self, cfg: KubernetesBackendConfig):
@@ -432,6 +435,7 @@ def get_job_logs(
         name: str,
         follow: bool = False,
         step: str = constants.NODE + "-0",
+        strict: bool = False,
     ) -> Iterator[str]:
         """Get the TrainJob logs"""
         # Get the TrainJob Pod name.
@@ -440,7 +444,12 @@ def get_job_logs(
             if c.status != constants.POD_PENDING and c.name == step:
                 pod_name = c.pod_name
                 break
+
         if pod_name is None:
+            if strict:
+                raise PodNotFoundError(
+                    f"No pod found for TrainJob {self.namespace}/{name} step={step}"
+                )
             return
 
         # Remove the number for the node step.

kubeflow/trainer/backends/kubernetes/backend_test.py

@@ -32,7 +32,7 @@
 import pytest
 
 from kubeflow.common.types import KubernetesBackendConfig
-from kubeflow.trainer.backends.kubernetes.backend import KubernetesBackend
+from kubeflow.trainer.backends.kubernetes.backend import KubernetesBackend, PodNotFoundError
 import kubeflow.trainer.backends.kubernetes.utils as utils
 from kubeflow.trainer.constants import constants
 from kubeflow.trainer.options import (
@@ -1439,6 +1439,16 @@ def test_get_job_logs(kubernetes_backend, test_case):
     print("test execution complete")
 
 
+def test_get_job_logs_strict_raises_when_pod_missing(kubernetes_backend, monkeypatch):
+    tj = get_train_job_data_type(runtime_name=TORCH_RUNTIME, train_job_name=BASIC_TRAIN_JOB_NAME)
+    for s in tj.steps:
+        s.status = constants.POD_PENDING
+
+    monkeypatch.setattr(kubernetes_backend, "get_job", lambda name: tj)
+
+    with pytest.raises(PodNotFoundError):
+        list(kubernetes_backend.get_job_logs(BASIC_TRAIN_JOB_NAME, strict=True))
+
 @pytest.mark.parametrize(
     "test_case",
     [

kubeflow/trainer/backends/localprocess/backend.py

@@ -195,6 +195,7 @@ def get_job_logs(
         name: str,
         follow: bool = False,
         step: str = constants.NODE + "-0",
+        strict: bool = False,
     ) -> Iterator[str]:
         _job = [j for j in self.__local_jobs if j.name == name]
         if not _job: