Fix API Docs generation

This is a minimal set of changes to fix API docs generation
that's currently failing all our Read The Docs syncing.
Adapted from #3029

Author: shralex

.github/workflows/check_docs_build.yml

@@ -18,7 +18,7 @@ jobs:
           persist-credentials: false
 
       - name: Install uv and set the Python version
-        uses: astral-sh/setup-uv@v6
+        uses: astral-sh/setup-uv@eb1897b8dc4b5d5bfe39a428a8f2304605e0983c  # v7.0.0
         with:
           python-version: '3.12'
           enable-cache: true

dependencies/requirements/requirements_docs.txt

@@ -1,5 +1,5 @@
 # Sphinx-related requirements.
-sphinx
+sphinx<9
 myst-nb
 myst-parser[linkify]
 sphinx-book-theme
@@ -8,10 +8,4 @@ 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
+jax

docs/conf.py

@@ -28,17 +28,16 @@
 import os
 import os.path
 import sys
-import warnings
+import logging
+from sphinx.util import logging as sphinx_logging
 
 # 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 = "2023–2026, Google LLC"
@@ -62,15 +61,7 @@
 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",
-]
+suppress_warnings = ["autodoc.import_object"]
 
 # -- Options for HTML output -------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output
@@ -94,18 +85,14 @@
 autodoc_member_order = "bysource"
 autodoc_typehints = "description"
 autodoc_mock_imports = [
-    "cloud_tpu_diagnostics",
-    "google_cloud_mldiagnostics",
-    "jetstream",
-    "librosa",
-    "ml_goodput_measurement",
-    "pathwaysutils",
     "safetensors",
     "tensorflow_datasets",
     "torch",
     "tpu_inference",
-    "transformer_engine",
     "vllm",
+    "grain",
+    "librosa",
+    "sentencepiece",
 ]
 autosummary_generate = True
 
@@ -129,10 +116,12 @@
     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"),
+    os.path.join("reference", "api_generated", "modules.rst"),
+    os.path.join("reference", "api_generated", "install_maxtext_extra_deps.rst"),
+    os.path.join("reference", "api_generated", "install_maxtext_extra_deps.install_github_deps.rst"),
 ]
 
 
@@ -191,8 +180,31 @@ def run_apidoc(_):
     sys.exit(1)
 
 
-# Connect the apidoc generation to the Sphinx build process
+class FilterSphinxWarnings(logging.Filter):
+  """Filter autosummary 'duplicate object description' warnings.
+
+  These warnings are unnecessary as they do not cause missing documentation
+  or rendering issues, so it is safe to filter them out.
+  """
+
+  def __init__(self, app):
+    self.app = app
+    super().__init__()
+
+  def filter(self, record: logging.LogRecord) -> bool:
+    msg = record.getMessage()
+    filter_out = ("duplicate object description",)
+    return not msg.strip().startswith(filter_out)
+
+
 def setup(app):
+  """Set up the Sphinx application with custom behavior."""
+
+  # Connect the apidoc generation to the Sphinx build process
   run_apidoc(None)
   print("running:", app)
-  # app.connect("builder-inited", run_apidoc)
+
+  # Set up custom logging filters
+  logger = logging.getLogger("sphinx")
+  warning_handler, *_ = [h for h in logger.handlers if isinstance(h, sphinx_logging.WarningStreamHandler)]
+  warning_handler.filters.insert(0, FilterSphinxWarnings(app))

docs/reference/api.rst

@@ -23,4 +23,4 @@ This section contains the complete API documentation for the ``MaxText`` library
    :caption: Package Modules
    :glob:
 
-   api_generated/modules
+   api_generated/MaxText