Adds live log streaming for remote pods

[JyotinderSingh]

Summary

• Adds live container log streaming to the terminal during remote job execution, giving users real-time visibility into what their code is doing on the cluster
• Introduces a LogStreamer context manager backed by a daemon thread that tails Kubernetes pod logs via the follow API
• In interactive terminals, logs are rendered in a Rich Live panel (last 25 lines, auto-refreshing); in non-interactive contexts (CI, piped output), logs are streamed as raw lines with rule delimiters
• Integrates log streaming into both the GKE and Pathways job-polling loops. Streaming starts automatically once the pod reaches Running state
• Moves rich from an optional CLI dependency to a core dependency

Example outputs

Regular In-Pod Logs

Training Run with Live Progress Bars

[gemini-code-assist]

Summary of Changes

Hello @JyotinderSingh, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly enhances the user experience by introducing real-time log streaming for remote Kubernetes pods. Users can now observe the output of their GKE and Pathways jobs directly in their terminal as they execute, eliminating the need for manual log retrieval. The new LogStreamer intelligently adapts its display based on the terminal environment, ensuring an optimal viewing experience whether interactive or not, and efficiently manages log retrieval in a background thread.

Highlights

• Live Log Streaming: Implemented live container log streaming to the terminal for remote job execution, providing real-time visibility into cluster operations.
• LogStreamer Context Manager: Introduced a LogStreamer context manager, backed by a daemon thread, to tail Kubernetes pod logs using the follow API.
• Adaptive Log Rendering: Logs are rendered in a Rich Live panel (last 25 lines, auto-refreshing) for interactive terminals, and as raw lines with rule delimiters for non-interactive contexts (CI, piped output).
• Backend Integration: Integrated log streaming into both the GKE and Pathways job-polling loops, with streaming automatically starting once a pod reaches the 'Running' state.
• Dependency Update: Moved the rich library from an optional CLI dependency to a core project dependency.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

• keras_remote/backend/gke_client.py
  • Imported LogStreamer for log management.
  • Integrated LogStreamer into the wait_for_job function to initiate log streaming when a pod transitions to a 'Running' state.
• keras_remote/backend/gke_client_test.py
  • Updated setUp method to mock the LogStreamer for testing purposes.
  • Added new test cases to verify that log streaming starts correctly when a pod is running and does not start prematurely.
• keras_remote/backend/log_streaming.py
  • Added a new module to provide live log streaming capabilities from Kubernetes pods.
  • Implemented _stream_pod_logs to fetch and process logs using the Kubernetes follow API.
  • Developed _render_live_panel for displaying logs in an interactive Rich Live panel.
  • Created _render_plain for streaming raw log lines with delimiters in non-interactive environments.
  • Introduced the LogStreamer context manager to manage the lifecycle of the log-streaming daemon thread.
• keras_remote/backend/log_streaming_test.py
  • Added comprehensive unit tests for the _stream_pod_logs function, covering API calls, routing based on terminal type, and error handling.
  • Included tests for _render_plain to ensure correct streaming and delimiter usage.
  • Provided tests for _render_live_panel to validate handling of partial log lines.
  • Implemented tests for the LogStreamer context manager, verifying thread creation, idempotency, and proper thread joining on exit.
• keras_remote/backend/pathways_client.py
  • Imported LogStreamer for log management.
  • Integrated LogStreamer into the wait_for_job function to enable log streaming for Pathways jobs once the leader pod is in a 'Running' state.
• pyproject.toml
  • Moved the rich dependency from the optional cli group to the core dependencies list.

Activity

• The pull request was created by JyotinderSingh.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces live log streaming for remote Kubernetes pods, significantly enhancing real-time visibility into job execution for both GKE and Pathways backends. The implementation uses a LogStreamer context manager with a daemon thread to tail pod logs, rendering them in a Rich Live panel for interactive terminals or as raw lines with delimiters for non-interactive contexts. The changes also correctly update pyproject.toml to make rich a core dependency, reflecting its integral role in the new logging functionality. The new test cases for log streaming are well-designed and cover important scenarios. Overall, this is a valuable feature that improves the user experience by providing immediate feedback on remote job status. I've identified a few areas for minor improvements related to constant definition, exception handling, and potential log redundancy.

[divyashreepathihalli]

the file doesnt really fit in backend. Maybe move it to utils?

[divyashreepathihalli]

Merging this for now, so the team can dogfood. We can move the file around in a followup PR.