Skip to content

Auto-set ZARR3 and OCDBT to False for Pathways in config#3358

Draft
igorts-git wants to merge 1 commit intomainfrom
igorts/b490520104-zarr3-aocdbt-flags
Draft

Auto-set ZARR3 and OCDBT to False for Pathways in config#3358
igorts-git wants to merge 1 commit intomainfrom
igorts/b490520104-zarr3-aocdbt-flags

Conversation

@igorts-git
Copy link
Collaborator

@igorts-git igorts-git commented Mar 9, 2026

Description

Address b/490520104: ZARR3 and OCDBT flags must be manually set to False for Pathways (should be default).

When enable_single_controller==True we automatically set ZARR3 and OCDBT flags to False.

One potentially controversial change as part of this PR is that the checkpoint conversion script to_maxtext now needs to be explicitly told enable_single_controller=True. I did not see a way around it.

FIXES: b/490520104

Pair-programmed with Gemini-CLI.

Tests

Added a unit test. Manually ran a tiny Pathways workload.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

  • I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
  • I have necessary comments in my code, particularly in hard-to-understand areas.
  • I have run end-to-end tests tests and provided workload links above if applicable.
  • I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

@codecov
Copy link

codecov bot commented Mar 10, 2026

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

@igorts-git igorts-git force-pushed the igorts/b490520104-zarr3-aocdbt-flags branch 4 times, most recently from e0d77a0 to 36ad42d Compare March 11, 2026 20:31
@igorts-git igorts-git force-pushed the igorts/b490520104-zarr3-aocdbt-flags branch from 36ad42d to fc5c0a9 Compare March 11, 2026 20:39
@github-actions
Copy link

🤖 Hi @igorts-git, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

Copy link

@github-actions github-actions bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

## 📋 Review Summary

This pull request correctly centralizes the logic for disabling ZARR3 and OCDBT when using Pathways (enable_single_controller=True) into the configuration system. This simplifies command-line invocations by removing the need to manually set these flags and ensures consistent default behavior across the codebase.

🔍 General Feedback

  • The architectural choice to move the configuration overrides to a single source of truth in types.py is sound and enhances maintainability.
  • The unit tests clearly and effectively verify that the overridden defaults behave as expected, covering both the default scenario and the edge case where colocated_python_checkpointing is set to True.
  • Good job updating the documentation and standalone scripts to reflect the new unified approach.

# To create a scanned checkpoint, you can use `src/maxtext/checkpoint_conversion/to_maxtext.py` and if
# using Pathways, please set `enable_single_controller=True`
# python src/maxtext/checkpoint_conversion/to_maxtext.py \
# maxtext/configs/base.yml \

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟢 The path to the config file might need to be prefixed with `src/` to match the Python script execution path, assuming the user is running this from the repository root.
Suggested change
# maxtext/configs/base.yml \
# src/maxtext/configs/base.yml \

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant