Skip to content

Allow identical config overrides and handle configuration exceptions …#3933

Merged
copybara-service[bot] merged 1 commit into
mainfrom
shralex_config_1
May 18, 2026
Merged

Allow identical config overrides and handle configuration exceptions …#3933
copybara-service[bot] merged 1 commit into
mainfrom
shralex_config_1

Conversation

@shralex

@shralex shralex commented May 18, 2026

Copy link
Copy Markdown
Collaborator

…cleanly with code 2 and Kubernetes termination logging

This PR addresses two issues in MaxText's configuration initialization and validation:

  1. Previously, calling a model config with a command-line override would hard-reject the execution if the key was defined by both, even if the provided override value was completely identical (such as tokenizer_type="huggingface" when running qwen3 models). This PR modifies validate_no_keys_overridden_twice to compare the resolved values. Overrides are now allowed if the override value matches the model configuration value exactly. A ValueError is only raised if the values differ and override_model_config=False.

  2. When a deterministic config validation error occurs during startup, JobSet/Kubernetes would previously crash-loop and exhaust maxRestarts because config validation errors looked like infrastructure/general flakes (exiting with code 1).
    This PR catches configuration validation errors, prints a highly visible error banner to sys.stderr, writes the specific failure message to /dev/termination-log for Kubernetes status propagation, and exits with standard exit code 2 (reserved for CLI syntax and config errors). Standard behavior is bypassed when running in test environments (pytest/unittest) to allow original validation exceptions to propagate for assertions.

Tests

Tested on TPU VM

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

codecov Bot commented May 18, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

@shralex
shralex force-pushed the shralex_config_1 branch from e02c3e7 to 7fdc477 Compare May 18, 2026 15:59
Comment thread src/maxtext/configs/pyconfig.py Outdated
…cleanly with code 2 and Kubernetes termination logging
@shralex
shralex force-pushed the shralex_config_1 branch from 7fdc477 to 338fcb6 Compare May 18, 2026 16:13
@copybara-service
copybara-service Bot merged commit 23b4787 into main May 18, 2026
50 checks passed
@copybara-service
copybara-service Bot deleted the shralex_config_1 branch May 18, 2026 18:19
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.

3 participants