feat(trainer): add tail_lines support to get_job_logs

kubeflow/trainer/api/trainer_client.py

@@ -186,7 +186,8 @@ def get_job_logs(
         self,
         name: str,
         step: str = constants.NODE + "-0",
-        follow: bool | None = False,
+        follow: bool = False,
+        tail_lines: int | None = None,
     ) -> Iterator[str]:
         """Get logs from a specific step of a TrainJob.
 
@@ -202,6 +203,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.
+            tail_lines: Number of lines from the end of the logs to return. If None, returns all available logs.
 
         Returns:
             Iterator of log lines.
@@ -211,7 +213,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, tail_lines=tail_lines, step=step)
 
     def get_job_events(self, name: str) -> list[types.Event]:
         """Get events for a TrainJob.

kubeflow/trainer/backends/base.py

@@ -63,6 +63,7 @@ def get_job_logs(
         self,
         name: str,
         follow: bool = False,
+        tail_lines: int | None = None,
         step: str = constants.NODE + "-0",
     ) -> Iterator[str]:
         raise NotImplementedError()

kubeflow/trainer/backends/container/adapters/base.py

@@ -98,7 +98,12 @@ def get_container(self, container_id: str):
         raise NotImplementedError()
 
     @abc.abstractmethod
-    def container_logs(self, container_id: str, follow: bool) -> Iterator[str]:
+    def container_logs(
+        self,
+        container_id: str,
+        follow: bool,
+        tail_lines: int | None = None,
+    ) -> Iterator[str]:
         """Stream logs from a container."""
         raise NotImplementedError()
 

kubeflow/trainer/backends/container/adapters/docker.py

@@ -106,10 +106,19 @@ def get_container(self, container_id: str):
         """Get Docker container by ID."""
         return self.client.containers.get(container_id)
 
-    def container_logs(self, container_id: str, follow: bool) -> Iterator[str]:
+    def container_logs(
+        self,
+        container_id: str,
+        follow: bool,
+        tail_lines: int | None = None,
+    ) -> Iterator[str]:
         """Stream logs from Docker container."""
         container = self.get_container(container_id)
-        logs = container.logs(stream=bool(follow), follow=bool(follow))
+        logs = container.logs(
+            stream=bool(follow),
+            follow=bool(follow),
+            tail="all" if tail_lines is None else tail_lines,
+        )
         if follow:
             for chunk in logs:
                 if isinstance(chunk, bytes):

kubeflow/trainer/backends/container/adapters/podman.py

@@ -112,10 +112,19 @@ def get_container(self, container_id: str):
         """Get Podman container by ID."""
         return self.client.containers.get(container_id)
 
-    def container_logs(self, container_id: str, follow: bool) -> Iterator[str]:
+    def container_logs(
+        self,
+        container_id: str,
+        follow: bool,
+        tail_lines: int | None = None,
+    ) -> Iterator[str]:
         """Stream logs from Podman container."""
         container = self.get_container(container_id)
-        logs = container.logs(stream=bool(follow), follow=bool(follow))
+        logs = container.logs(
+            stream=bool(follow),
+            follow=bool(follow),
+            tail="all" if tail_lines is None else tail_lines,
+        )
         if follow:
             for chunk in logs:
                 if isinstance(chunk, bytes):

kubeflow/trainer/backends/container/backend.py

@@ -778,6 +778,7 @@ def get_job_logs(
         self,
         name: str,
         follow: bool = False,
+        tail_lines: int | None = None,
         step: str = constants.NODE + "-0",
     ) -> Iterator[str]:
         """Get logs for a training job by querying container runtime."""
@@ -798,7 +799,11 @@ def get_job_logs(
                 continue
 
             try:
-                yield from self._adapter.container_logs(container["id"], follow)
+                yield from self._adapter.container_logs(
+                    container["id"],
+                    follow,
+                    tail_lines,
+                )
             except Exception as e:
                 logger.warning(f"Failed to get logs for {container['name']}: {e}")
                 yield f"Error getting logs: {e}\n"

kubeflow/trainer/backends/container/backend_test.py

@@ -98,7 +98,12 @@ def get_container(self, container_id: str):
                 return Mock(id=container_id, status=container["status"])
         return None
 
-    def container_logs(self, container_id: str, follow: bool) -> Iterator[str]:
+    def container_logs(
+        self,
+        container_id: str,
+        follow: bool,
+        tail_lines: int | None = None,
+    ) -> Iterator[str]:
         if follow:
             yield f"Log line 1 from {container_id}\n"
             yield f"Log line 2 from {container_id}\n"

kubeflow/trainer/backends/kubernetes/backend.py

@@ -446,6 +446,7 @@ def get_job_logs(
         self,
         name: str,
         follow: bool = False,
+        tail_lines: int | None = None,
         step: str = constants.NODE + "-0",
     ) -> Iterator[str]:
         """Get the TrainJob logs"""
@@ -461,7 +462,10 @@ def get_job_logs(
         # Remove the number for the node step.
         container_name = re.sub(r"-\d+$", "", step)
         yield from self._read_pod_logs(
-            pod_name=pod_name, container_name=container_name, follow=follow
+            pod_name=pod_name,
+            container_name=container_name,
+            follow=follow,
+            tail_lines=tail_lines,
         )
 
     def wait_for_job_status(
@@ -623,7 +627,13 @@ def __get_runtime_from_cr(
             ),
         )
 
-    def _read_pod_logs(self, pod_name: str, container_name: str, follow: bool) -> Iterator[str]:
+    def _read_pod_logs(
+        self,
+        pod_name: str,
+        container_name: str,
+        follow: bool,
+        tail_lines: int | None = None,
+    ) -> Iterator[str]:
         """Read logs from a pod container."""
         try:
             if follow:
@@ -633,6 +643,7 @@ def _read_pod_logs(self, pod_name: str, container_name: str, follow: bool) -> It
                     namespace=self.namespace,
                     container=container_name,
                     follow=True,
+                    tail_lines=tail_lines,
                 )
 
                 # Stream logs incrementally.
@@ -642,6 +653,7 @@ def _read_pod_logs(self, pod_name: str, container_name: str, follow: bool) -> It
                     name=pod_name,
                     namespace=self.namespace,
                     container=container_name,
+                    tail_lines=tail_lines,
                 )
 
                 yield from logs.splitlines()

kubeflow/trainer/backends/localprocess/backend.py

@@ -194,6 +194,7 @@ def get_job_logs(
         self,
         name: str,
         follow: bool = False,
+        tail_lines: int | None = None,
         step: str = constants.NODE + "-0",
     ) -> Iterator[str]:
         _job = [j for j in self.__local_jobs if j.name == name]
@@ -207,7 +208,10 @@ def get_job_logs(
                 continue
             # Flatten the generator and pass through flags so it behaves as expected
             # (adjust args if stream_logs has different signature)
-            yield from _step.job.logs(follow=follow)
+            yield from _step.job.logs(
+                follow=follow,
+                tail_lines=tail_lines,
+            )
 
     def get_job_events(self, name: str) -> list[types.Event]:
         raise NotImplementedError()

kubeflow/trainer/backends/localprocess/job.py

@@ -149,17 +149,23 @@ def cancel(self):
     def returncode(self):
         return self._returncode
 
-    def logs(self, follow=False) -> list[str]:
+    def logs(
+        self,
+        follow=False,
+        tail_lines: int | None = None,
+    ) -> list[str]:
         if not follow:
-            return self._stdout.splitlines()
+            lines = self._stdout.splitlines()
+            return lines if tail_lines is None else lines[-tail_lines:]
 
         try:
             for chunk in self.stream_logs():
                 print(chunk, end="", flush=True)  # stream to console live
         except StopIteration:
             pass
 
-        return self._stdout.splitlines()
+        lines = self._stdout.splitlines()
+        return lines if tail_lines is None else lines[-tail_lines:]
 
     def stream_logs(self):
         """Generator that yields new output lines as they come in."""