Add Pathways Integration by divyashreepathihalli · Pull Request #31 · keras-team/remote

divyashreepathihalli · 2026-02-19T21:17:12Z

Pathways Backend Integration: Introduced a new PathwaysBackend to support ML Pathways (LeaderWorkerSet) for remote job execution.
Automatic Backend Selection: The @keras_remote.run decorator now automatically selects the 'pathways' backend for multi-host TPU configurations, while retaining 'gke' as the default for other scenarios.
New Pathways Client: A dedicated 'pathways_client.py' module was added to handle the submission, monitoring, and cleanup of LeaderWorkerSet jobs on Kubernetes.
LeaderWorkerSet (LWS) Installation: Added utilities and constants to facilitate the installation of the LeaderWorkerSet Custom Resource Definition (CRD) on GKE clusters.
Pathways Example: Included a new example (pathways_example.py) demonstrating how to train a simple Keras model remotely using the Pathways integration with a JAX backend.

Example execution.

Fixes #19

gemini-code-assist · 2026-02-19T21:17:33Z

Summary of Changes

Hello @divyashreepathihalli, 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 introduces comprehensive support for ML Pathways (LeaderWorkerSet) as an alternative backend for keras_remote job execution. It enables users to leverage Pathways for distributed training, particularly for multi-host TPU setups, by providing a new backend implementation, client-side logic for managing LeaderWorkerSet resources, and an updated 'run' decorator with intelligent backend auto-detection. This enhancement expands the capabilities of keras_remote for scalable machine learning workloads on Kubernetes.

Highlights

Pathways Backend Integration: Introduced a new PathwaysBackend to support ML Pathways (LeaderWorkerSet) for remote job execution.
Automatic Backend Selection: The @keras_remote.run decorator now automatically selects the 'pathways' backend for multi-host TPU configurations, while retaining 'gke' as the default for other scenarios.
New Pathways Client: A dedicated 'pathways_client.py' module was added to handle the submission, monitoring, and cleanup of LeaderWorkerSet jobs on Kubernetes.
LeaderWorkerSet (LWS) Installation: Added utilities and constants to facilitate the installation of the LeaderWorkerSet Custom Resource Definition (CRD) on GKE clusters.
Pathways Example: Included a new example (pathways_example.py) demonstrating how to train a simple Keras model remotely using the Pathways integration with a JAX backend.

🧠 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

examples/pathways_example.py
- Added a new example script demonstrating remote Keras model training using the Pathways backend with a JAX environment.
keras_remote/backend/execution.py
- Imported 'pathways_client'.
- Refactored 'GKEBackend' to inherit from a new 'BaseK8sBackend' class.
- Introduced 'PathwaysBackend' class with methods for submitting, waiting for, and cleaning up LeaderWorkerSet jobs.
keras_remote/backend/pathways_client.py
- Added a new module containing functions for interacting with the Kubernetes LeaderWorkerSet API, including '_get_lws_version', 'submit_pathways_job', 'wait_for_job', 'cleanup_job', and '_create_lws_spec'.
keras_remote/cli/constants.py
- Added 'LWS_INSTALL_URL' constant pointing to the LeaderWorkerSet CRD manifest.
keras_remote/cli/infra/post_deploy.py
- Imported 'LWS_INSTALL_URL'.
- Added 'install_lws' function to apply the LeaderWorkerSet CRD to a Kubernetes cluster.
keras_remote/core/core.py
- Imported 'PathwaysBackend' and 'accelerators'.
- Modified the 'run' decorator to accept an optional 'backend' parameter.
- Implemented logic within 'run' to automatically select 'pathways' backend for multi-host TPUs or fall back to 'gke'.
- Added '_execute_on_pathways' function to delegate execution to the 'PathwaysBackend'.

Activity

No human activity has been recorded on this pull request yet.

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.

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 Pathways integration as a new backend for remote execution, which is a great addition. The changes are well-organized, with a new PathwaysBackend and pathways_client to handle LeaderWorkerSet (LWS) jobs on GKE. The auto-detection of the backend based on the accelerator type in core.py is a nice touch that aligns with the goal of minimizing user configuration. I have a couple of suggestions to improve maintainability and consistency.

keras_remote/backend/pathways_client.py

keras_remote/core/core.py

JyotinderSingh

Thanks for these changes! Left a few comments.

Also, I'm a little unfamiliar with Pathways, but from my understanding we might need to add a couple of environment variables for collective communication (MEGASCALE_COORDINATOR_ADDRESS, TPU_WORKER_ID, etc.). The leader and worker templates are identical for now, they might need to integrate this information.

examples/pathways_example.py

keras_remote/backend/pathways_client.py

JyotinderSingh · 2026-02-20T16:05:39Z

keras_remote/cli/infra/post_deploy.py

+def install_lws():
+  """Install the LeaderWorkerSet custom resource controller.
+
+  This enables Pathways scheduling on the GKE cluster.
+  """
+  subprocess.run(
+    ["kubectl", "apply", "--server-side", "-f", LWS_INSTALL_URL],
+    check=True,
+  )


We're not calling this function anywhere.

ahh, missed adding this to up.py - done

divyashreepathihalli · 2026-02-20T18:47:54Z

Thanks for the review Jyotinder. Updated the code and re-ran the example to verify execution.

jeffcarp

Thanks! Lgtm after comments

examples/pathways_example.py

jeffcarp · 2026-02-20T18:55:37Z

examples/pathways_example.py

+  history = model.fit(x, y, epochs=5, batch_size=32, validation_split=0.2)
+
+  print("\nTraining completed successfully on Pathways!")
+  return history.history


Just to confirm, there are no user code changes to run on Pathways within their remote function? All it needs is backend="pathways"? That's pretty cool

yeah if it is None. It auto detects if the user requested for a multi node TPU and picks the pathways backend

keras_remote/backend/pathways_client.py

jeffcarp · 2026-02-20T18:59:49Z

keras_remote/backend/pathways_client.py

+    )
+    logger.info(f"Deleted LeaderWorkerSet: {job_name}")
+  except ApiException as e:
+    if e.status != 404:


What if it is 404? Should it be re-raised?

Updated this. A 404 Not Found during deletion means the LeaderWorkerSet has already been deleted. we should not re raise this as it would crash the clean up routine.

keras_remote/core/core.py

jeffcarp · 2026-02-20T21:51:17Z

Thanks!

divyashreepathihalli added 5 commits February 19, 2026 20:13

initial pathways integration

7453b29

add pathways example

20a54ab

fix core

01ac7b7

code reformat

3a69ed8

auto detect lws version

08881d9

divyashreepathihalli marked this pull request as draft February 19, 2026 21:17

code reformat

8085ae3

gemini-code-assist bot reviewed Feb 19, 2026

View reviewed changes

keras_remote/backend/pathways_client.py Outdated Show resolved Hide resolved

keras_remote/core/core.py Show resolved Hide resolved

divyashreepathihalli marked this pull request as ready for review February 19, 2026 22:41

address gemini comment

2769b1f

divyashreepathihalli requested review from JyotinderSingh and jeffcarp February 19, 2026 22:56

address Gemini comment

d4e97b6

divyashreepathihalli mentioned this pull request Feb 19, 2026

Integrate Pathways on top of GKE standard #19

Closed

JyotinderSingh requested changes Feb 20, 2026

View reviewed changes

divyashreepathihalli added 3 commits February 20, 2026 18:36

address comments

83eca73

code reformat

50cf0c2

map internal APIs to standard Multi-Host Eenv specs

b010eb6

divyashreepathihalli requested a review from JyotinderSingh February 20, 2026 18:49

jeffcarp approved these changes Feb 20, 2026

View reviewed changes

JyotinderSingh approved these changes Feb 20, 2026

View reviewed changes

address jeff's comment

364d81d

jeffcarp merged commit c6b4dad into main Feb 20, 2026
2 checks passed

JyotinderSingh deleted the pathways-1 branch February 25, 2026 12:33

Conversation

divyashreepathihalli commented Feb 19, 2026 • edited by jeffcarp Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

gemini-code-assist bot commented Feb 19, 2026

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

JyotinderSingh left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

JyotinderSingh Feb 20, 2026

Choose a reason for hiding this comment

Uh oh!

divyashreepathihalli Feb 20, 2026

Choose a reason for hiding this comment

Uh oh!

divyashreepathihalli commented Feb 20, 2026

Uh oh!

jeffcarp left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

jeffcarp Feb 20, 2026

Choose a reason for hiding this comment

Uh oh!

divyashreepathihalli Feb 20, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

jeffcarp Feb 20, 2026

Choose a reason for hiding this comment

Uh oh!

divyashreepathihalli Feb 20, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

jeffcarp commented Feb 20, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

divyashreepathihalli commented Feb 19, 2026 •

edited by jeffcarp

Loading