Conversation
There was a problem hiding this comment.
Code Review
This pull request migrates the MaxDiffusion documentation and configuration from the deprecated XPK tool to Cluster Toolkit (using the gcluster CLI) for running large-scale jobs on GKE. It introduces a new guide for Cluster Toolkit, updates the main README and other docs to deprecate XPK, and updates pyconfig.py to note that JOBSET_NAME is not set by Cluster Toolkit. The review feedback points out inconsistencies in variable names (such as $PROJECT and $ZONE versus $PROJECT_ID and $LOCATION) between the README examples and the new setup guide, suggesting updates to ensure copy-paste compatibility for users.
Perseus14
marked this pull request as draft
October 3, 2026 15:46
Perseus14
force-pushed
the
repo-hygiene-minor-fixes
branch
from
October 3, 2026 16:00
dfbb356 to
53d36af
Compare
Perseus14
force-pushed
the
docs/cluster-toolkit-migration
branch
from
October 3, 2026 16:01
d6d0746 to
94ebd06
Compare
XPK is deprecated (maintenance mode through Q3 2026, then archived; new TPU/GPU
generations are only supported via Cluster Toolkit). Mirror MaxText's additive
migration: make Cluster Toolkit the recommended GKE path while keeping the XPK
material for users with existing XPK clusters.
* docs/getting_started/run_maxdiffusion_via_cluster_toolkit.md (new): full
guide covering prerequisites, gcluster job config, building the dependency +
runner image and pushing it to Artifact Registry (there is no public base
image), a first `gcluster job submit`, --compute-type/--topology selection,
env vars, storage mounts, monitoring, and a repo-specific XPK -> gcluster
flag table.
* README.md: rename "Deploying with XPK" -> "Deploying with Cluster Toolkit"
and "Multi-Host Training with XPK" -> "Multi-Host Training with Cluster
Toolkit"; add translated `gcluster job submit` commands (--workload -> --name,
--zone -> --location, --device-type -> --compute-type + --topology,
--base-docker-image -> --image, --max-restarts -> --restarts,
--enable-debug-logs -> --verbose). The previously undefined ${IMAGE_DIR} is
replaced by an explicit IMAGE definition. Legacy XPK commands are preserved
in collapsed <details> blocks. Link the guide from both Getting Started
sections and update two prose references to xpk.
* README.md: the Wan 2.1 LIBTPU_INIT_ARGS blocks were single-quoted, so the
backslash-newlines were kept literally in the value; use double quotes. The
deployment command now forwards the flags with --env (xpk never did).
* docs/getting_started/run_maxdiffusion_via_xpk.md: deprecation banner linking
to the new guide and the official migration guide; XPK links now point at
the AI-Hypercomputer org and the moved docker-images page.
* docs/README.md, docs/getting_started/first_run.md: list the Cluster Toolkit
guide as recommended and mark XPK as deprecated.
* src/maxdiffusion/pyconfig.py: comment-only update. XPK injected JOBSET_NAME
into pods via a fieldRef; gcluster's JobSet template does not, so run_name
must be passed explicitly (all documented commands do).
Why `--image` with a prebuilt runner image instead of gcluster's
`--base-image` + `--build-context .` (as MaxText documents): gcluster's crane
builder appends the build context at the image root and keeps the base image's
working directory, while MaxDiffusion images set WORKDIR /deps and already
contain a source copy there, so local changes would silently be ignored.
The gcluster commands were derived from the official migration guide's flag
mapping, the gcluster job guide and the cluster-toolkit source (JobSet template,
crane builder) but have not yet been validated end-to-end on a cluster; both
the guide and the README say so explicitly.
This branch is stacked on repo-hygiene-minor-fixes because that commit touched
docs/README.md and the README lines adjacent to these sections.
Perseus14
force-pushed
the
repo-hygiene-minor-fixes
branch
from
October 3, 2026 16:36
53d36af to
1fa74a3
Compare
Perseus14
force-pushed
the
docs/cluster-toolkit-migration
branch
from
October 3, 2026 16:36
94ebd06 to
3089ff5
Compare
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
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.
Warning
Stacked on #500 (base branch is
repo-hygiene-minor-fixes, notmain), because that PR rewritesdocs/README.mdand touches the README lines next to these sections. Please merge #500 first; once itsbranch is deleted GitHub retargets this PR to
mainautomatically. Do not merge this PR into the hygiene branch.Summary
XPK is officially deprecated
(maintenance through Q3 2026, then archived; new TPU/GPU generations only in Cluster Toolkit). This PR mirrors
MaxText's additive migration: Cluster Toolkit's
gclusterbecomes the documented GKE path, and the XPK material iskept behind a deprecation banner for users with existing XPK clusters.
docs/getting_started/run_maxdiffusion_via_cluster_toolkit.md: prerequisites,gcluster job config,building the dependency + runner image and pushing it to Artifact Registry (MaxDiffusion has no public base image),
a first
gcluster job submit,--compute-type/--topologyselection, env vars,--mount, monitoring, and arepo-specific XPK →
gclusterflag table.Deploying with XPK→Deploying with Cluster Toolkit,Multi-Host Training with XPK→Multi-Host Training with Cluster Toolkit, with translatedgcluster job submitcommands(
--workload→--name,--zone→--location,--device-type→--compute-type+--topology,--base-docker-image→--image,--max-restarts→--restarts,--enable-debug-logs→--verbose). The legacyXPK commands are preserved in collapsed
<details>blocks. The guide is linked from both Getting Startedsections. The previously undefined
${IMAGE_DIR}is replaced by an explicitIMAGEdefinition.LIBTPU_INIT_ARGSblocks were single-quoted, so the backslash-newlines were keptliterally in the value; they are now double-quoted, and the deployment command forwards the flags with
--env(the XPK command never did).
run_maxdiffusion_via_xpk.md: deprecation banner pointing at the new guide and the official migration guide;XPK links updated to the
AI-Hypercomputerorg / moved docs page.docs/README.md,first_run.md: Cluster Toolkit listed as recommended, XPK marked deprecated.pyconfig.py: comment-only. XPK injectedJOBSET_NAMEinto pods via a fieldRef;gcluster's JobSet template doesnot, so
run_namemust be passed explicitly (all documented commands do).Why
--imagewith a prebuilt runner image (not--base-image+--build-context .as MaxText uses)gcluster's crane builder appends the build context at the image root and keeps the base image's workingdirectory, while MaxDiffusion images set
WORKDIR /depsand already contain a source copy there. A command likepython src/maxdiffusion/train.pywould silently run the copy baked into the base image instead of local changes.The guide explains this in a note and uses
maxdiffusion_runner.Dockerfile+--imageinstead.Verification
migration guide,
the gcluster job guide
and the cluster-toolkit source (
pkg/orchestrator/gke/templates/jobset.tmpl,pkg/imagebuilder/crane_builder.go);--enable-debug-logs↔--verboseequivalence verified against both code bases (same four TPU debug env vars).gclusterblocks dry-run underbash -euwith a stubbedgcluster: well-formed argv,LIBTPU_INIT_ARGSexpands to a single line with no stray backslashes.<details>balanced; all same-file and cross-file anchors and relative links resolve (scripted);ruff+pyinkclean onpyconfig.py.Important
The
gclustercommands have not been run end-to-end on a cluster (no cluster available while writing this).Both the guide and the README say so explicitly. The cheapest smoke test is the first-run
v6e-8command insection 4 of the guide — a reviewer with a Cluster Toolkit cluster is very welcome to try it before merging.