Skip to content

docs: add Cluster Toolkit (gcluster) guide and mark XPK as deprecated - #501

Draft
Perseus14 wants to merge 1 commit into
repo-hygiene-minor-fixesfrom
docs/cluster-toolkit-migration
Draft

Perseus14 wants to merge 1 commit into
repo-hygiene-minor-fixesfrom
docs/cluster-toolkit-migration

Conversation

@Perseus14

Copy link
Copy Markdown
Collaborator

Warning

Stacked on #500 (base branch is repo-hygiene-minor-fixes, not main), because that PR rewrites
docs/README.md and touches the README lines next to these sections. Please merge #500 first; once its
branch is deleted GitHub retargets this PR to main automatically. 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 gcluster becomes the documented GKE path, and the XPK material is
kept behind a deprecation banner for users with existing XPK clusters.

  • New 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/--topology selection, env vars, --mount, monitoring, and a
    repo-specific XPK → gcluster flag table.
  • README: Deploying with XPK → Deploying with Cluster Toolkit, Multi-Host Training with XPK →
    Multi-Host Training with Cluster Toolkit, with 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 legacy
    XPK commands are preserved in collapsed <details> blocks. The guide is linked from both Getting Started
    sections. The previously undefined ${IMAGE_DIR} is replaced by an explicit IMAGE definition.
  • README: the two Wan 2.1 LIBTPU_INIT_ARGS blocks were single-quoted, so the backslash-newlines were kept
    literally 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-Hypercomputer org / moved docs page.
  • docs/README.md, first_run.md: Cluster Toolkit listed as recommended, XPK marked deprecated.
  • pyconfig.py: comment-only. 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 (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 working
directory, while MaxDiffusion images set WORKDIR /deps and already contain a source copy there. A command like
python src/maxdiffusion/train.py would 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 + --image instead.

Verification

  • Flag mapping cross-checked against the official
    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 ↔ --verbose equivalence verified against both code bases (same four TPU debug env vars).
  • Both README gcluster blocks dry-run under bash -eu with a stubbed gcluster: well-formed argv,
    LIBTPU_INIT_ARGS expands to a single line with no stray backslashes.
  • Code fences / <details> balanced; all same-file and cross-file anchors and relative links resolve (scripted);
    ruff + pyink clean on pyconfig.py.

Important

The gcluster commands 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-8 command in
section 4 of the guide — a reviewer with a Cluster Toolkit cluster is very welcome to try it before merging.

@Perseus14
Perseus14 requested a review from entrpn as a code owner October 3, 2026 15:32
@github-actions

github-actions Bot commented Oct 3, 2026

Copy link
Copy Markdown

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

Copy link
Copy Markdown

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

Comment thread README.md Outdated
Comment thread README.md Outdated
@Perseus14
Perseus14 marked this pull request as draft October 3, 2026 15:46
@Perseus14
Perseus14 force-pushed the repo-hygiene-minor-fixes branch from dfbb356 to 53d36af Compare October 3, 2026 16:00
@Perseus14
Perseus14 force-pushed the docs/cluster-toolkit-migration branch from d6d0746 to 94ebd06 Compare October 3, 2026 16:01
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
Perseus14 force-pushed the repo-hygiene-minor-fixes branch from 53d36af to 1fa74a3 Compare October 3, 2026 16:36
@Perseus14
Perseus14 force-pushed the docs/cluster-toolkit-migration branch from 94ebd06 to 3089ff5 Compare October 3, 2026 16:36
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.

1 participant