Clean up Python files: pyink formatting and pylint fixes

[igorts-git]

Description

Overview

This pull request addresses code quality and formatting consistency across all Python files in the MaxText repository, ensuring full compliance with pre-commit linting standards (pyink, pylint, and codespell).
All changes are purely cosmetic and structural formatting fixes. No underlying execution logic or existing documentation comments were altered.

Key Changes

• Automated Code Formatting: Ran pyink repository-wide to ensure consistent line lengths, indentation, and spacing.
• Robust File Handling: Specified explicit encoding="utf-8" in open() calls (json_reporter.py, eval_runner.py) and refactored temporary file handling to use with context managers (test_eval_runner.py).
• Docstring & Import Corrections:
  • Corrected a broken import path in generate_sft_golden_data.py (tests.post_training.integration.sft_trainer_correctness_test).
  • Added missing docstrings to helper functions (_derive_from_maxtext_config, _build_arg_parser, _wait_until_healthy) and test classes to satisfy pylint.
  • Replaced keyword calls to dict() with standard dictionary literals ({}) for better runtime performance as recommended by pylint: use-dict-literal.
  • Added targeted # pylint: disable annotations where appropriate (e.g., for test fixtures requiring import-outside-toplevel or internal protected-access).

Tests

GitHub CI.

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.

[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 successfully improves the code quality of the MaxText evaluation and benchmark modules. The changes primarily focus on automated formatting (pyink), robust file handling by specifying explicit UTF-8 encoding, and correcting minor import and docstring issues.

🔍 General Feedback

• Robustness: Specifying encoding="utf-8" in open() calls is a great practice that prevents platform-dependent encoding issues.
• Maintainability: The transition to pyink formatting and the addition of missing docstrings enhance the readability and modularity of the codebase.
• Correctness: The correction of the import path in generate_sft_golden_data.py ensures that the SFT golden data generation remains functional after the recent reorganization of tests.

[github-actions]

🤖 I'm sorry @igorts-git, but I was unable to process your request. Please see the logs for more details.

[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 successfully improves the code quality and consistency of the MaxText repository by applying pyink formatting, adding missing docstrings, and addressing various pylint warnings. The changes are primarily cosmetic and structural, enhancing maintainability without altering execution logic.

🔍 General Feedback

• Consistency: The repository-wide application of pyink ensures a uniform coding style, which is excellent for long-term maintenance.
• Robustness: Adding explicit encoding="utf-8" to file operations and using context managers for temporary files are good practices that prevent common bugs.
• Documentation: The addition of docstrings to key helper functions and test classes significantly improves the readability and discoverability of the codebase.
• Performance: Replacing dict() with dictionary literals {} is a minor but welcomed optimization for Python runtimes.

[github-actions]

🤖 I'm sorry @igorts-git, but I was unable to process your request. Please see the logs for more details.

[codecov]

Codecov Report

❌ Patch coverage is 0% with 17 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...c/maxtext/kernels/attention/jax_flash_attention.py	0.00%	17 Missing ⚠️

📢 Thoughts on this report? Let us know!

[richjames0]

lgtm