Merged
Conversation
This commit switches the XPK workload generation template for pathways workloads from using the PathwaysJob CRD directly over to a standard JobSet object. - Updates PW_WORKLOAD_CREATE_YAML to use JobSet API - Modifies component YAML generation in pathways.py to output Pod containers rather than custom PathwaysJob components - Adds unit tests to workload_test.py to verify correct layout of pathways jobsets
* Convert proxy and RM sidecars to initContainers with restartPolicy: Always. * Ensure all container ports specify protocol: TCP. * Enforce restartPolicy: OnFailure for the worker template. * Inject required environment variables (JAX_PLATFORMS, JAX_BACKEND_TARGET, XCLOUD_ENVIRONMENT) natively into the primary user workload container. * Add completionMode: Indexed to both head and worker replicated jobs. * Set successPolicy, startupPolicy, and suspend fields to match legacy generated JobSet.
* Added get_pathways_instance_type to core/pathways.py * Format the gke_accelerator natively for the RM sidecar (e.g. tpuv6e:2x2)
…aling * Set worker backoffLimit to args.max_slice_restarts * 4 to replicate the legacy controller's logic.
* Since XPK now deploys Pathways workloads using native JobSet API, we no longer need to install the PathwaysJob CRD when creating or adapting a cluster.
* Changed prev_indentation to base_indentation to specify exactly how many spaces to prepend. * Removed the convoluted len(indentation) - 2 calls in favor of passing base_indentation=16 directly.
* Instead of manually checking for multiple variants of trailing whitespace after the 'env:' key, use the re module to robustly identify and inject the required Pathways environment variables.
* Since docker_container.py generates the container YAML template using `env: {env}\n`, the key will always be followed by an optional space and a newline, whether {env} is empty or not. Thus, the defensive else block was unreachable and unnecessary.
* Moved the inline 'import re' from get_user_workload_for_pathways to the module level.
* Removed regex pattern matching in favor of a direct string replacement on 'env:' since 'env: {env}' always appears in the template even if {env} evaluates to an empty string. This ensures our injected block always gets prepended correctly before any user variables.
* Since we always append the variables directly after the 'env:' key from the main container template, we can just define the injected block without the 'env:' header entirely and avoid having to strip it out.
* Solves the 'Returning Any from function declared to return str' type checker error in pathways.py by explicitly defining the tuple return type of get_user_workload_container.
* Replaced hardcoded strings and numbers (like 'test-pw-workload', 'test-docker', and 2) with formatted variables linked directly to the mocked args (e.g. {args.workload}, {args.num_slices}) to make the unit test more robust and maintainable.
* Added explicit assertions for the newly migrated JobSet fields including the coordinator block, network dns configurations, restart strategies, completion modes, and strict verification of the dynamically scaled backoffLimit calculation.
…pk into pathwaysjob-migration
…pk into pathwaysjob-migration
This commit switches the XPK workload generation template for pathways workloads from using the PathwaysJob CRD directly over to a standard JobSet object. - Updates PW_WORKLOAD_CREATE_YAML to use JobSet API - Modifies component YAML generation in pathways.py to output Pod containers rather than custom PathwaysJob components - Adds unit tests to workload_test.py to verify correct layout of pathways jobsets
* Convert proxy and RM sidecars to initContainers with restartPolicy: Always. * Ensure all container ports specify protocol: TCP. * Enforce restartPolicy: OnFailure for the worker template. * Inject required environment variables (JAX_PLATFORMS, JAX_BACKEND_TARGET, XCLOUD_ENVIRONMENT) natively into the primary user workload container. * Add completionMode: Indexed to both head and worker replicated jobs. * Set successPolicy, startupPolicy, and suspend fields to match legacy generated JobSet.
* Added get_pathways_instance_type to core/pathways.py * Format the gke_accelerator natively for the RM sidecar (e.g. tpuv6e:2x2)
…aling * Set worker backoffLimit to args.max_slice_restarts * 4 to replicate the legacy controller's logic.
* Since XPK now deploys Pathways workloads using native JobSet API, we no longer need to install the PathwaysJob CRD when creating or adapting a cluster.
…irectly * Added an optional field directly to . * Populated the for all supported TPU generations (v4, v5e, v5p, v6e, tpu7, tpu7x) directly in their system definitions in . * Removed the hardcoded fallback mapping and its associated exception from entirely, fetching the version dynamically from the system configuration instead.
|
🤖 Hi @FIoannides, I've received your request, and I'm working on it now! You can track my progress in the logs for more details. |
|
🤖 I'm sorry @FIoannides, but I was unable to process your request. Please see the logs for more details. |
FIoannides
requested review from
SikaGrr,
jamOne-,
scaliby and
stony-tark
as code owners
jamOne-
approved these changes
Mar 11, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Fixing the identation issue that was introduced on the (now rolled back) PR #1099
Issue
Testing
Run the generated JobSet manually on the cluster, no YAML issues caught