Merge pull request #2049 from SamuelMarks:docgen

PiperOrigin-RevId: 859256825

Author: Google-ML-Automation

.github/workflows/check_docs_build.yml

@@ -13,19 +13,28 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Checkout repository
-        uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8  # v5.0.0
+        uses: actions/checkout@v5
         with:
           persist-credentials: false
 
-      - name: Set up Python
-        uses: actions/setup-python@e797f83bcb11b83ae66e0230d6156d7c80228e7c  # v6.0.0
+      - name: Install uv and set the Python version
+        uses: astral-sh/setup-uv@v6
         with:
           python-version: '3.12'
-          cache: 'pip' # caching pip dependencies
+          enable-cache: true
+
+      - name: Set venv
+        run: uv venv --python 3.12 $GITHUB_WORKSPACE/venv
 
       - name: Install dependencies
-        run: pip install -r dependencies/requirements/requirements_docs.txt
+        run: . $GITHUB_WORKSPACE/venv/bin/activate && uv pip install -r dependencies/requirements/requirements_docs.txt
 
       - name: Build documentation
         run: |
+          . $GITHUB_WORKSPACE/venv/bin/activate
+          uv pip install -e . --no-deps
+          uv pip install torch
           sphinx-build -W -b html docs docs/_build/html
+        env:
+          JAX_PLATFORMS: cpu
+          CUDA_VISIBLE_DEVICES: ""

.gitignore

@@ -1,7 +1,8 @@
 *__pycache__*
 tmp/
 logs/
-
+.venvs
+venv*
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]

DOCS.md

@@ -0,0 +1,45 @@
+<!--
+# Copyright 2023–2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+-->
+Documentation… documentation!
+=============================
+
+## Dependencies
+First install the dependencies:
+```sh
+$ python3 -m pip install -r requirements_docs.txt
+```
+(or `uv pip install` …)
+
+## Build
+```sh
+$ sphinx-build -M html docs out
+```
+
+## Serve
+You can use any static file HTTP server, e.g.:
+```sh
+$ python3 -m http.server -d 'out/html'
+```
+
+## Build & server (watch for changes)
+```sh
+$ python3 -m pip install sphinx-autobuild
+$ sphinx-autobuild docs out
+```
+
+## Release to readthedocs
+
+See GitHub Action

README.md

@@ -126,3 +126,7 @@ MaxText aims to provide you with the best OSS models, whether as a reference imp
 ## Get involved
 
 Please join our [Discord Channel](https://discord.com/invite/2H9PhvTcDU) and if you have feedback, you can file a feature request, documentation request, or bug report [here](https://github.com/AI-Hypercomputer/maxtext/issues/new/choose).
+
+## License
+
+[Apache License 2.0](LICENSE)

dependencies/requirements/requirements_docs.txt

@@ -5,3 +5,13 @@ myst-parser[linkify]
 sphinx-book-theme
 sphinx-design
 sphinx-copybutton
+
+# for import docs
+-r base_requirements/requirements.txt
+google-tunix
+mlcommons-loadgen
+mlperf-logging @ https://github.com/mlcommons/logging/archive/38ab22670527888c8eb7825a4ece176fcc36a95d.zip
+tf-keras
+torch==2.7.1
+trl==0.19.0
+vllm

docs/conf.py

@@ -1,4 +1,4 @@
-# Copyright 2023–2025 Google LLC
+# Copyright 2023–2026 Google LLC
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -25,9 +25,23 @@
 # -- Project information -----------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information
 
+import os
+import os.path
+import sys
+import warnings
+
+# Prevent JAX/Torch/TF from trying to access TPU/GPU during doc build
+os.environ["JAX_PLATFORMS"] = "cpu"
+os.environ["CUDA_VISIBLE_DEVICES"] = ""
+MAXTEXT_REPO_ROOT = os.environ.get("MAXTEXT_REPO_ROOT", os.path.join(os.path.dirname(os.path.dirname(__file__))))
+sys.path.insert(0, os.path.abspath(os.path.join(MAXTEXT_REPO_ROOT, "test")))
+sys.path.insert(0, os.path.abspath(os.path.join(MAXTEXT_REPO_ROOT, "src")))
+
+warnings.filterwarnings("ignore", category=FutureWarning, module="keras.src.export.tf2onnx_lib")
+
 project = "MaxText"
 # pylint: disable=redefined-builtin
-copyright = "2025, Google LLC"
+copyright = "2023–2026, Google LLC"
 author = "MaxText developers"
 
 # -- General configuration ---------------------------------------------------
@@ -37,11 +51,27 @@
     "myst_nb",
     "sphinx_design",
     "sphinx_copybutton",
+    "sphinx.ext.napoleon",
+    # This needs to be before autodoc^
+    "sphinx.ext.autodoc",
+    "sphinx.ext.autosummary",
+    "sphinx.ext.viewcode",
 ]
 
 templates_path = ["_templates"]
 source_suffix = [".rst", ".ipynb", ".md"]
 
+# Suppress specific warnings
+suppress_warnings = [
+    "app.add_node",
+    "ref.python",
+    "myst.xref_ambiguous",
+    "docutils",
+    "autodoc",
+    "autodoc.duplicate_object_description",
+    "toc.not_included",
+]
+
 # -- Options for HTML output -------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output
 
@@ -60,6 +90,24 @@
 ]
 myst_linkify_fuzzy_links = False
 
+# -- Options for autodoc ----------------------------------------------------
+autodoc_member_order = "bysource"
+autodoc_typehints = "description"
+autodoc_mock_imports = [
+    "cloud_tpu_diagnostics",
+    "google_cloud_mldiagnostics",
+    "jetstream",
+    "ml_goodput_measurement",
+    "pathwaysutils",
+    "safetensors",
+    "tensorflow_datasets",
+    "torch",
+    "tpu_inference",
+    "transformer_engine",
+    "vllm",
+]
+autosummary_generate = True
+
 # Theme-specific options
 # https://sphinx-book-theme.readthedocs.io/en/stable/reference.html
 html_theme_options = {
@@ -77,7 +125,72 @@
 
 # Remove specific documents from ToC
 exclude_patterns = [
-    "run_maxtext/run_maxtext_via_multihost_job.md",
-    "run_maxtext/run_maxtext_via_multihost_runner.md",
-    "reference/core_concepts/llm_calculator.ipynb",
+    os.path.join("guides", "run_maxtext", "run_maxtext_via_multihost_job.md"),
+    os.path.join("guides", "run_maxtext", "run_maxtext_via_multihost_runner.md"),
+    os.path.join("explanations", "llm_calculator.ipynb"),
+    os.path.join("reference", "api.rst"),
+    os.path.join("run_maxtext", "run_maxtext_via_multihost_job.md"),
+    os.path.join("run_maxtext", "run_maxtext_via_multihost_runner.md"),
+    os.path.join("reference", "core_concepts", "llm_calculator.ipynb"),
 ]
+
+
+# -- Autogenerate API documentation ------------------------------------------
+def run_apidoc(_):
+  """Runs sphinx-apidoc to generate API documentation.
+
+  This function is connected to the Sphinx build process and is triggered to
+  automatically generate the reStructuredText (RST) files for the API
+  documentation from the docstrings in the MaxText source code.
+
+  Args:
+    _: The Sphinx application object. Not used.
+  """
+  # directly within the Sphinx process, especially on macOS, as it avoids
+  # potential multiprocessing/forking issues like the "mutex lock failed" error.
+  # pylint: disable=import-outside-toplevel
+  import subprocess
+
+  os.environ["OBJC_DISABLE_INITIALIZE_FORK_SAFETY"] = "1"
+
+  assert os.path.isfile(os.path.join(MAXTEXT_REPO_ROOT, "pyproject.toml"))
+
+  # The path where the generated RST files will be stored
+  output_path = os.path.join(MAXTEXT_REPO_ROOT, "docs", "reference", "api_generated")
+
+  # Command to run sphinx-apidoc
+  # Note: We use `sys.executable -m sphinx.ext.apidoc` to ensure we're using
+  # the apidoc from the same Python environment as Sphinx.
+  command = [
+      sys.executable,
+      "-m",
+      "sphinx.ext.apidoc",
+      "--module-first",
+      "--force",
+      "--separate",
+      "--output-dir",
+      output_path,
+      os.path.join(MAXTEXT_REPO_ROOT, "src"),
+      # Paths to exclude
+      os.path.join(MAXTEXT_REPO_ROOT, "tests"),
+      os.path.join(MAXTEXT_REPO_ROOT, "src", "MaxText", "experimental"),
+      os.path.join(MAXTEXT_REPO_ROOT, "src", "MaxText", "inference_mlperf"),
+      os.path.join(MAXTEXT_REPO_ROOT, "src", "MaxText", "scratch_code"),
+      os.path.join(MAXTEXT_REPO_ROOT, "src", "MaxText", "utils", "ckpt_conversion"),
+      os.path.join(MAXTEXT_REPO_ROOT, "src", "MaxText", "rl"),
+  ]
+
+  # Run the command and check for errors
+  try:
+    print("Running sphinx-apidoc...")
+    subprocess.check_call(command, env={**os.environ, **{"OBJC_DISABLE_INITIALIZE_FORK_SAFETY": "1"}})
+  except subprocess.CalledProcessError as e:
+    print(f"sphinx-apidoc failed with error: {e}", file=sys.stderr)
+    sys.exit(1)
+
+
+# Connect the apidoc generation to the Sphinx build process
+def setup(app):
+  run_apidoc(None)
+  print("running:", app)
+  # app.connect("builder-inited", run_apidoc)

docs/index.md

@@ -13,12 +13,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
  -->
+
 # MaxText
 
 ```{raw} html
 :file: index.html
 ```
 
+:link: reference/api
+
 <div class="doc-body">
 <section class="latest-news">
 
@@ -34,11 +37,11 @@
 :maxdepth: 2
 :hidden:
 
-install_maxtext.md
-tutorials.md
-run_maxtext.md
-guides.md
-reference.md
-development.md
-release_notes.md
+install_maxtext
+tutorials
+run_maxtext
+guides
+reference
+development
+release_notes
 ```

docs/reference.md

@@ -1,5 +1,5 @@
 <!--
- Copyright 2024 Google LLC
+ Copyright 2023–2025 Google LLC
 
  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.
@@ -54,8 +54,8 @@ Key concepts including checkpointing strategies, quantization, tiling, and Mixtu
 :hidden:
 :maxdepth: 1
 
-reference/performance_metrics.md
-reference/models.md
-reference/architecture.md
-reference/core_concepts.md
+reference/performance_metrics
+reference/models
+reference/architecture
+reference/core_concepts
 ```

docs/reference/api.rst

@@ -0,0 +1,26 @@
+..
+   Copyright 2023–2026 Google LLC
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+        https://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
+
+API Reference
+=============
+
+This section contains the complete API documentation for the ``MaxText`` library, automatically generated from the source code docstrings.
+
+.. toctree::
+   :maxdepth: 4
+   :caption: Package Modules
+   :glob:
+
+   api_generated/modules

pyproject.toml

@@ -28,6 +28,7 @@ dependencies = []
 [tool.hatch.metadata.hooks.requirements_txt.optional-dependencies]
 tpu = ["dependencies/requirements/generated_requirements/tpu-requirements.txt"]
 cuda12 = ["dependencies/requirements/generated_requirements/cuda12-requirements.txt"]
+docs = ["dependencies/requirements/requirements_docs.txt"]
 
 [project.urls]
 Repository = "https://github.com/AI-Hypercomputer/maxtext.git"