Auto-generated API docs with sphinx #2049

SamuelMarks · 2025-07-30T19:21:27Z

This PR adds generated API docs.

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

Thanks Samuel! Added a few comments. Can we please close this space https://github.com/maxtext and replace references to it from the PR description to the reads-the-docs website.

docs/index.md

mkdocs.yml

bvandermoon

For my understanding, where do these generated API docs end up? Do we run this manually? Or will it automatically be uploaded to ReadTheDocs?

requirements-docs.txt

docs/index.md

SamuelMarks · 2025-09-25T21:16:27Z

Hi - I'm at Quansight Labs and new to the project, and was asked to take a look

I tried building this locally, and can confirm that
First install the dependencies:
$ python3 -m pip install -r requirements_docs.txt
Build
$ sphinx-build -M html docs out
Serve

You can use any static file HTTP server, e.g.:
$ python3 -m http.server -d 'out/html'
works fine for me. The generated API reference in the sidebar looks fine

However, if I click on any API reference item, all I see is something like
Is this intended, or is the toplevel docstring in that module

maxtext/src/MaxText/accelerator_to_spec_map.py

Lines 15 to 19 in 28e5097

""" Static map of TPU names such as v4-8 to properties such as chip layout."""

""" !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

IF YOU MODIFY THIS FILE YOU SHOULD ALSO ADD CORRESPONDING MODIFICATIONS TO

UserFacingNameToSystemCharacteristics in xpk/xpk.py !!!!! """

meant to show up?

Related to that, on the benchmark_chunked_prefill page, I don't see members listed, are we expecting

maxtext/src/MaxText/benchmark_chunked_prefill.py

Lines 101 to 116 in 28e5097

def benchmark_chunked_prefill(

engine: maxengine.MaxEngine,

params: Any,

chunked_tokens_list: list[chunked_prefill.ChunkedTokens],

):

"""Benchmarks chunked prefill without prefix caching.

Args:

engine: The MaxEngine instance.

params: The model parameters.

chunked_tokens_list: A list of ChunkedTokens objects representing the

input sequence split into chunks.

Returns:

The average time taken for chunked prefill.

"""

to show up?

@MarcoGorelli Thanks for following up. I plan on taking the codebase and going module by module "Gemini: add missing file, module, class, and function docstrings in Google docstring format". Alternatively others can do the same. It should be trivial to maintain 100% doc coverage for this codebase.

As for the specific issues you are coming across, I'll try and replicate. Sometimes sphinx rst's format needs to be explicitly overriden in favour of Google's docstring format; which may be the cause of the issue.

For my understanding, where do these generated API docs end up? Do we run this manually? Or will it automatically be uploaded to ReadTheDocs?

@bvandermoon Thanks for following up. The documentation is auto generated onto readthedocs.

MarcoGorelli · 2025-09-30T09:18:34Z

I plan on taking the codebase and going module by module

Sure, if it was expected that with this PR then the listed members wouldn't show up, then no objections

Other than that, there's some commented-out code that's being added here, is that intended?

MarcoGorelli · 2025-10-01T15:19:08Z

pyproject.toml

+#[tool.hatch.metadata.hooks.requirements_txt]
+#files = ["requirements.txt"]


is it intentional to have added this commented-out code?

MarcoGorelli · 2025-10-16T19:00:46Z

docs/conf.py

+def setup(app):
+  run_apidoc(None)
+  print("running:", app)
+  # app.connect("builder-inited", run_apidoc)


should this be un-commented or deleted?

MarcoGorelli · 2025-10-20T10:46:37Z

docs/conf.py

+def setup(app):
+  run_apidoc(None)
+  print("running:", app)
+  # app.connect("builder-inited", run_apidoc)


hi @SamuelMarks

sorry, it's still not clear to me why we're adding this commented-out line

# app.connect("builder-inited", run_apidoc)

is it meant to be uncommented under certain circumstances?

bvandermoon · 2025-10-21T21:56:46Z

.github/workflows/check_docs_build.yml

+          pip install -e . --no-dependencies
+          pip install torch


Why are these dependencies needed for building documentation?

Because torch is imported somewhere in the codebase and sphinx raises a warning about this

bvandermoon · 2025-10-21T21:57:32Z

dependencies/requirements/requirements_docs.txt

+# for import docs
+-r base_requirements/requirements.txt
+google-tunix>=0.1.1
+mlcommons-loadgen
+mlperf-logging @ https://github.com/mlcommons/logging/archive/38ab22670527888c8eb7825a4ece176fcc36a95d.zip
+tf-keras
+torch==2.7.1
+trl==0.19.0
+vllm


Why are these needed for docs?

Because these are each imported somewhere in the codebase and sphinx raises a warning about this

Shouldn't this just be dependencies that are required for the documents themselves? Seems that is how @melissawm set it up previously. If it's just a warning, can we leave it off/turn it off? Or is the recommended approach to actually require the full MaxText set of dependencies?

It gives nice highlighting and documentation hints that would be missed if these were not installed (also the build fails)

…p ; [docs/reference/api.rst] Placeholder for generated API docs ; [pyproject.toml] Uncomment docs deps ; [docs/guides.md] Remove redundant extensions ; [src/MaxText/inference_mlperf/offline_mode.py] Switch from `flags` to `argparse` to overcome duplicate flag error; [.github/workflows/check_docs_build.yml] Use `uv` to resolve max recursive dependency resolution issue ; [src/MaxText/inference_mlperf/matmul/matmul_dtypes.py] Hide under `if __name__ == "__main__"` to resolve hangups on TPUs and CPUs

SamuelMarks requested review from A9isha, Lumosis, RissyRan, SurbhiJainUSC, aireenmei, bvandermoon, gagika, gobbleturk, gpolovets1, hengtaoguo, jrplatin, khatwanimohit, mailvijayasingh, mitalisi, parambole, patemotter, richjames0, shralex, vipannalla and yangyuwei as code owners July 30, 2025 19:21

SamuelMarks mentioned this pull request Aug 7, 2025

Reorganize documentation #2040

Merged

4 tasks

shralex requested changes Aug 13, 2025

View reviewed changes

docs/index.md Outdated Show resolved Hide resolved

mkdocs.yml Outdated Show resolved Hide resolved

SamuelMarks requested a review from NuojCheng as a code owner August 14, 2025 18:14

SamuelMarks force-pushed the docgen branch from 7214f9d to 3d81488 Compare August 14, 2025 18:17

shralex approved these changes Aug 15, 2025

View reviewed changes

bvandermoon reviewed Aug 15, 2025

View reviewed changes

requirements-docs.txt Outdated Show resolved Hide resolved

requirements-docs.txt Outdated Show resolved Hide resolved

docs/index.md Outdated Show resolved Hide resolved

SamuelMarks requested review from SujeethJinesh, jshin1394, liudangyi and shauryagup as code owners September 2, 2025 19:00

SamuelMarks force-pushed the docgen branch from 2f950ec to 8c35f4a Compare September 24, 2025 21:05

SamuelMarks force-pushed the docgen branch from 8c35f4a to 9e17dbf Compare September 30, 2025 18:37

SamuelMarks requested a review from suexu1025 as a code owner September 30, 2025 18:37

MarcoGorelli reviewed Oct 1, 2025

View reviewed changes

SamuelMarks force-pushed the docgen branch from 9e17dbf to 9390e98 Compare October 3, 2025 20:07

SamuelMarks force-pushed the docgen branch from 9390e98 to 7f46449 Compare October 14, 2025 21:44

SamuelMarks requested a review from jacoguzo as a code owner October 14, 2025 21:44

SamuelMarks force-pushed the docgen branch 3 times, most recently from 838b576 to 13c45bd Compare October 16, 2025 18:13

MarcoGorelli reviewed Oct 16, 2025

View reviewed changes

SamuelMarks force-pushed the docgen branch from 13c45bd to 71ff616 Compare October 16, 2025 21:53

MarcoGorelli reviewed Oct 20, 2025

View reviewed changes

SamuelMarks force-pushed the docgen branch from 71ff616 to ef7029d Compare October 21, 2025 04:12

bvandermoon reviewed Oct 21, 2025

View reviewed changes

SamuelMarks force-pushed the docgen branch 10 times, most recently from 5be41e6 to 4358e11 Compare October 31, 2025 17:33

SamuelMarks force-pushed the docgen branch from 4358e11 to a5374c9 Compare November 4, 2025 19:04

SamuelMarks force-pushed the docgen branch from a5374c9 to 40a0470 Compare November 4, 2025 21:31

		#[tool.hatch.metadata.hooks.requirements_txt]
		#files = ["requirements.txt"]

Auto-generated API docs with sphinx #2049

Are you sure you want to change the base?

Auto-generated API docs with sphinx #2049

Uh oh!

Conversation

SamuelMarks commented Jul 30, 2025 • edited by shralex Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Checklist

Uh oh!

shralex left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

bvandermoon left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

SamuelMarks commented Sep 25, 2025

Build

Serve

Uh oh!

MarcoGorelli commented Sep 30, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

SamuelMarks commented Jul 30, 2025 •

edited by shralex

Loading