Skip to content

Add preflight and container caching UX improvements#57

Merged
jeffcarp merged 1 commit intomainfrom
preflight
Feb 25, 2026
Merged

Add preflight and container caching UX improvements#57
jeffcarp merged 1 commit intomainfrom
preflight

Conversation

@jeffcarp
Copy link
Member

@jeffcarp jeffcarp commented Feb 25, 2026

Adds three quality of life improvements:

  1. Checks whether a NodePool exists before building the container (otherwise users will waste 5-10 mins of build time before finding out there's no NodePool to run their job)
  2. Groups accelerator containers by category (e.g. gpu, tpu) instead of by the specific accelerator type - this means switching from TPU v3 to v4 will not trigger a rebuild of the container
  3. Adds the accelerator name to the error message when it doesn't exist

@gemini-code-assist
Copy link
Contributor

Summary of Changes

Hello @jeffcarp, I'm Gemini Code Assist1! 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 introduces several quality-of-life improvements aimed at enhancing the user experience and efficiency when working with remote Keras jobs on GKE. It adds a crucial preflight validation step to verify the availability of necessary hardware resources before committing to potentially lengthy build processes. Furthermore, it refines the container caching mechanism to be more flexible by grouping accelerators into broader categories, thereby minimizing unnecessary container rebuilds. Finally, it provides more informative error messages to guide users when accelerator resources are not correctly configured.

Highlights

  • Preflight NodePool Validation: Implemented a check to ensure a required NodePool exists for an accelerator before container building, preventing wasted time on builds that would ultimately fail due to missing resources.
  • Accelerator Category-Based Caching: Updated container image naming and hashing to use accelerator categories (e.g., 'gpu', 'tpu') instead of specific types, reducing rebuilds when switching between accelerators within the same category.
  • Enhanced Error Reporting: Improved error messages for failed job scheduling to explicitly state the missing accelerator selector, aiding in quicker debugging and resolution.
  • Documentation Clarification: Added a note to the README.md emphasizing the prerequisite of setting up NodePools for accelerators.

🧠 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
  • README.md
    • Added a note about NodePool prerequisites for accelerators.
  • keras_remote/backend/execution.py
    • Introduced a validate_preflight method in BaseK8sBackend and its GKE-specific subclasses.
    • Integrated the preflight validation call into the execute_remote function.
  • keras_remote/backend/gke_client.py
    • Implemented a new validate_preflight function to check for the existence of GKE NodePools matching accelerator selectors.
    • Enhanced the _check_pod_scheduling error message to include the specific accelerator selector when a mismatch occurs.
  • keras_remote/backend/gke_client_test.py
    • Updated test utilities to support node_selector in pending pod mocks.
    • Modified scheduling failure tests to verify the improved error messages.
  • keras_remote/infra/container_builder.py
    • Refactored container image generation to use accelerator categories for hashing and tagging, improving caching efficiency.
    • Updated internal functions (_hash_requirements, _build_and_push, _generate_dockerfile) to reflect the category-based logic.
Activity
  • Introduced preflight checks to validate NodePool existence before container builds.
  • Implemented accelerator category-based container caching to reduce rebuilds.
  • Improved error messages for missing accelerator NodePools.
  • Updated documentation regarding NodePool prerequisites.
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.

- Group accelerators by category (CPU, GPU, TPU) for container image sharing, reducing redundant builds.
- Implement preflight check to validate node pool existence before building containers.
- Improve error messages for node selector mismatches in pod scheduling.
- Update simple_demo.py to return training loss and use more idiomatic result retrieval.
- Add preflight check note to README with link to Quick Start.
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces several valuable user experience improvements. The preflight check for NodePool existence is a great addition to prevent users from wasting time on a container build that's destined to fail. The change to group container images by accelerator category is a smart optimization that will improve caching and reduce unnecessary rebuilds. The more detailed error messages are also a welcome improvement for debugging. I have one suggestion regarding code duplication in keras_remote/backend/execution.py to enhance maintainability. Overall, these are solid quality-of-life improvements.

Comment on lines +159 to +168
def validate_preflight(self, ctx: JobContext) -> None:
"""Preflight checks for Pathways (currently same as GKE)."""
# Pathways also runs on GKE nodes with specific labels
gke_client.validate_preflight(
accelerator=ctx.accelerator,
project=ctx.project,
cluster=self.cluster,
zone=ctx.zone,
namespace=self.namespace,
)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

This validate_preflight implementation is identical to the one in GKEBackend. To avoid code duplication and improve maintainability, consider abstracting this logic into a common mixin or a shared base class for GKE-based backends.

For example, you could create a GKEPreflightMixin:

class GKEPreflightMixin:
    def validate_preflight(self, ctx: JobContext) -> None:
        """Check if the required node pool exists for the accelerator."""
        gke_client.validate_preflight(
            accelerator=ctx.accelerator,
            project=ctx.project,
            cluster=self.cluster,
            zone=ctx.zone,
            namespace=self.namespace,
        )

class GKEBackend(GKEPreflightMixin, BaseK8sBackend):
    # ... (no need to implement validate_preflight)

class PathwaysBackend(GKEPreflightMixin, BaseK8sBackend):
    # ... (no need to implement validate_preflight)

This would make the code more DRY and easier to maintain if the preflight logic needs to be updated in the future.

@jeffcarp jeffcarp merged commit 686e314 into main Feb 25, 2026
4 checks passed
@JyotinderSingh JyotinderSingh deleted the preflight branch February 27, 2026 04:04
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants