Set TF_CPP_MIN_LOG_LEVEL=0 in __init__.py

[igorts-git]

Description

Set TF_CPP_MIN_LOG_LEVEL=0 in __init__.py instead of train.py.

Setting this env variable after import jax is called has no effect, because the C++ logging system is already initialized.

The jax/__init__.py sets the default value of TF_CPP_MIN_LOG_LEVEL to "1" (code pointer). Then it calls _cloud_tpu_init(), which also initializes the C++ logging system.

There are many useful C++ LOG(INFO) messages that are suppressed without TF_CPP_MIN_LOG_LEVEL=0. For example, the path of the created profile trace (when running with profiler=xplane flag).

FIXES: b/474109045

Tests

Local run of train.py with profile=xplane to confirm that the trace path is printed.

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 Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[github-actions]

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

[github-actions]

📋 Review Summary

This Pull Request effectively addresses an issue where TF_CPP_MIN_LOG_LEVEL was not being set early enough to influence JAX's C++ logging, by moving the environment variable setting to __init__.py. The change is well-explained and improves logging visibility.

🔍 General Feedback

• The core change to set TF_CPP_MIN_LOG_LEVEL=0 earlier is a good improvement for debugging and visibility.
• Consider adding an automated unit test to ensure the logging behavior is consistent across environments and to prevent future regressions.

[github-actions]

🟡 It would be beneficial to add a unit test to programmatically verify that setting TF_CPP_MIN_LOG_LEVEL to '0' here correctly enables the desired C++ logging output, especially the trace path for profiler=xplane. This would help prevent regressions and ensure the intended behavior persists.

[igorts-git]

I could not figure out how to add a reliable unit test for this. The tests are located outside of the "src/MaxText" directory and invoke train.main() as a function. This means that the order of imports in the test itself would affect the C++ logging behavior. Specifically, if the test imports JAX before "MaxText.train" the logging won't work.
Additionally, it is not easy to figure out how to capture C++ LOG(INFO) outputs in a unit tests. The standard Python unit testing libs only capture logs coming from Python.

[gobbleturk]

LGTM, why del os though?

[igorts-git]

LGTM, why del os though?

I am not a Python expert, but I understand that this is the common convention to avoid pollution of the namespace. Without this del statement one can refer to the module os as MaxText.os, which would be confusing. FWIW, Gemini seems to agree with me that this is a nice thing to have.