Skip to content

Add PathwaysJobSet.with_custom_libtpu to override libtpu at runtime on GKE. - #331

Open
copybara-service[bot] wants to merge 1 commit into
mainfrom
test_984946941
Open

copybara-service[bot] wants to merge 1 commit into
mainfrom
test_984946941

Conversation

@copybara-service

Copy link
Copy Markdown

Add PathwaysJobSet.with_custom_libtpu to override libtpu at runtime on GKE.

Testing a custom libtpu.so with Pathways on GKE currently requires
rebuilding the server and proxy images, which bake libtpu at /lib/libtpu.so.

with_custom_libtpu(uri) adds an init container that downloads a libtpu.so
(or a .whl/.zip containing one) from gs:// or http(s):// into an
emptyDir volume and subPath-mounts it over /lib/libtpu.so in the RM,
proxy and worker containers (also setting TPU_LIBRARY_PATH), so the servers
dlopen the custom library without any image rebuild. Calling it again
replaces the override.

@google-cla

google-cla Bot commented Sep 26, 2026

Copy link
Copy Markdown

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

… on GKE.

Testing a custom `libtpu.so` with Pathways on GKE currently requires
rebuilding the server and proxy images, which bake libtpu at `/lib/libtpu.so`.

`with_custom_libtpu(uri)` adds an init container that downloads a `libtpu.so`
(or a `.whl`/`.zip` containing one) from `gs://` or `http(s)://` into an
`emptyDir` volume and `subPath`-mounts it over `/lib/libtpu.so` in the RM,
proxy and worker containers (also setting `TPU_LIBRARY_PATH`), so the servers
`dlopen` the custom library without any image rebuild. Calling it again
replaces the override.

PiperOrigin-RevId: 984946941

This branch has not been deployed

No deployments
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.

0 participants