Switch to installable Python module hierarchy ; replace imports and executable calls to match ; disable failing tests with new test-global flags

[SamuelMarks]

Description

FIXES (in part): #1107

For commentary see #1108

Created with:

$ git switch test-fixer
$ fd -Fepy -j1 '__init__' > /tmp/init.txt
$ git checkout -b PKG-FIX_setup
$ while read -r f; do git checkout test-fixer "${f}"; done</tmp/init.txt
$ git checkout test-fixer requirements.txt .github setup.py
$ git commit -S -am "[setup] Python package layout support and minor other fixes" && git push mine "PKG-FIX_setup" && gh pr create --title "[setup] Python package layout suppo
rt and minor other fixes" --body "For commentary see #1108"

Recommend merging this PR first; before the rest.

Tests

python -m pytest and pytest should both now work when run from repository root.

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.

[shralex]

Lets change the description of this PR to be descriptive and include the checklist please -- one of the automatic tests also requires it.

[gobbleturk]

Down to just 1 file missing license

Copybara failed running your registered migration:
migration:
//depot/google3/third_party/py/maxtext/copy.bara.sky:github_pr_to_piper
origin: [1482]
error: Error validating 'verify_match 'Licensed under the Apache
License, Version 2.0'': setup.py - Expected string was not present.

1 file(s) failed the validation of verify_match 'Licensed under the Apache
License, Version 2.0', located at
//depot/google3/third_party/py/maxtext/copy.bara.sky:307:22.

[gobbleturk]

I believe there are conflicts with main based on our internal integration flow, can you try rebasing to main? I've added rebase + squash instructions in an earlier comment (you might know of a better way though)

There might be non-trivial merge conflicts you have to resolve - our internal tool flagged three files
Conflicts:

• MaxText/inference/paged_attention.py
• MaxText/max_utils.py
• MaxText/maxengine.py

[gobbleturk]

okay we are making progress, it looks like we need a small internal integration change

Can you explain how setup.py is used https://github.com/AI-Hypercomputer/maxtext/pull/1482/files#diff-60f61ab7a8d1910d86d9fda2261620314edcae5894d5aaa236b821c7256badd7?

[gobbleturk]

We are working on internal copybara changes to support this

[SamuelMarks]

okay we are making progress, it looks like we need a small internal integration change

Can you explain how setup.py is used https://github.com/AI-Hypercomputer/maxtext/pull/1482/files#diff-60f61ab7a8d1910d86d9fda2261620314edcae5894d5aaa236b821c7256badd7?

setup.py is used to support the old-style python setup.py install syntax as well as the newer, python3 -m pip install; conda install; and uv pip install softwares.

More details: https://packaging.python.org/en/latest/guides/distributing-packages-using-setuptools/

I squashed the commits on this branch a minute before your comment! - Let me know if there's anything left for me to do.

[gobbleturk]

Removing and Adding back Pull Ready to pick up new copybara changes

[SurbhiJainUSC]

Similarly, these tests also pass with changes in main.

[SamuelMarks]

Wrong. This test does not pass. Sent through commit a66854f reverting enacting your position.

[SamuelMarks]

Awesome, thanks Samuel! Can you list the tests you had to disable (and regex to find them?) in the PR description? I assume you disabled the subprocess tests?

You're welcome =)

Yes I disable the subprocess tests as discussed on Thursday's call.

Tests disabled, toggleable by setting env var or changing vals in MaxText/tests/globals.py:

Function(s)	Location
test_autoselected_attention ; test_with_dot_product	MaxText/tests/integration_tests/checkpoint_compatibility_test.py
test_autoselected_attention; test_with_dot_product	MaxText/tests/integration_tests/checkpointing_test.py
test_autoselected_attention ; test_with_dot_product	MaxText/tests/integration_tests/generate_param_only_checkpoint_test.py
Inference_Microbenchmark::test	MaxText/tests/inference_microbenchmark_smoke_test.py
test_shmap_collective_matmul_example	MaxText/tests/integration_tests/shmap_collective_matmul_test.py
test_sft_format_with_prompt_completion	MaxText/tests/sft_data_processing_test.py