pyconfig → pydantic

[SamuelMarks]

Description

Background

MaxText is currently engineered around pyconfig. Pyconfig—https://pypi.org/project/pyconfig/—was last updated in 2017 and has 20k monthly downloads.

Pydantic—https://pypi.org/project/pydantic/—is constantly updated and has hundreds of millions of monthly downloads.

Pydantic is the most widely used data validation library for Python.

– https://docs.pydantic.dev

Summary

The TL;DR version is:

• pydantic has become basically the standard for configuration formats, specifying the inputs and outputs (e.g., for REST APIs in FastAPI framework)
• pydantic links in well with the Python type-checker so you can oft find errors before runtime
• pydantic errors are clean and clear
• with my python compiler, you can have clean CLIs with --help and GUIs and generate SQL models (SQLAlchemy) as desired
• SDK documentation becomes much cleaner (which is good in preparation for the pypi release + hosted doc pages)

Migration

Proposed changes to the MaxText codebase:

1. Completely remove the pyconfig dependency
   1. maybe not immediately so there’s a chance for people to easily migration of their existing setups, e.g., with new functions from_pyconfig_to_pydantic
2. Migrate all examples, and codebase occurrences with new pydantic types
3. Put pydantic types on a computer+human understandable hierarchy, e.g.:
   1. One global types.py, or
   2. One types.py per config occurrence (e.g., one per module if each module has a different config)
4. Create new CLI that uses common CLI syntax (e.g., this can be automatically created using my Python compiler https://github.com/offscale/cdd-python)
5. Migrate all shell scripts and docs to use this new CLI
   1. TBD: remove the shell scripts in favour of Python SDK usage.

Technical implementation notes

Consistency

As part of migrating the codebase across to the new configuration format various small changes were made to improve consistency. These include:

0. self.cfg → self.config:

maxtext $ rg -Fq 'self.config' -tpy --stats --json | jq '.data.stats.matches'
547
maxtext $ rg -Fq 'self.cfg' -tpy --stats --json | jq '.data.stats.matches' 
144

Tests

Tested the change on your CI.

Checklist

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

• I have performed a self-review of my code.
• 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.