[Tunix] Moves the smoke_tests folder to top level.

FUTURE_COPYBARA_INTEGRATE_REVIEW=#888 from ev-br:doc_ci_check 03f8ca7
PiperOrigin-RevId: 846423327

Author: wang2yn84

.github/workflows/build_docs.yml

@@ -0,0 +1,32 @@
+name: Check the Documentation Build
+
+on:
+  pull_request:
+    types: [opened, synchronize]
+  workflow_dispatch:
+
+permissions:
+  contents: read
+
+jobs:
+  build-docs:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8  # v5.0.0
+        with:
+          persist-credentials: false
+
+      - name: Set up Python
+        uses: actions/setup-python@e797f83bcb11b83ae66e0230d6156d7c80228e7c  # v6.0.0
+        with:
+          python-version: '3.12'
+          cache: 'pip' # caching pip dependencies
+
+      - name: Install dependencies
+        run: pip install .[docs]
+
+      - name: Build documentation
+        run: |
+          pushd docs   # XXX _collection/examples not found if run from root
+          sphinx-build -W -b html . ./_build/html

.github/workflows/cpu-tests.yml

@@ -45,7 +45,8 @@ jobs:
       - name: Install the tunix wheel
         run: |
           python -m pip install --upgrade pip
-          python -m pip install google_tunix-*-py3-none-any.whl
+          WHEEL=$(ls -1 google_tunix-*-py3-none-any.whl | head -n1)
+          python -m pip install "${WHEEL}[prod]"
           python -m pip install pytest
 
       - name: Verify Tunix imports from installed package

.github/workflows/tpu-tests.yml

@@ -215,6 +215,8 @@ jobs:
           pip list | egrep 'jax|flax|libtpu'
           cd ..
           git clone https://github.com/sgl-project/sglang-jax.git && cd sglang-jax/python && pip install -e . --force-reinstall && cd ../..
+          # TODO(b/470113586): Remove this test once the bug is fixed.
+          pip install jax[tpu]==0.8.1 --force-reinstall
           pip list | egrep 'jax|flax|libtpu'
 
           # Install bookworm, vllm container defaults to bullseye causes segfault for sglang-jax.

.gitignore

@@ -21,6 +21,7 @@
 # Documentation
 /docs/_build/
 /docs/_collections/
+/docs/jupyter_execute
 /docs/sg_execution_times.rst
 /docs/generated
 build/*

docs/conf.py

@@ -6,6 +6,8 @@
 
 import logging
 
+from sphinx.util import logging as sphinx_logging
+
 # -- Project information -----------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information
 
@@ -32,8 +34,10 @@
     "_build",
     "Thumbs.db",
     ".DS_Store",
-    "_collections/examples/model_load/from_safetensor_load/*"
-    "_collections/examples/rl/README.md",
+    (
+        "_collections/examples/model_load/from_safetensor_load/*"
+        "_collections/examples/rl/README.md"
+    ),
     "_collections/examples/sft/**",
     "_collections/examples/deepscaler/**",
 ]
@@ -47,7 +51,7 @@
 html_static_path = ["_static"]
 html_logo = "_static/img/tunix.png"
 html_css_files = [
-    'custom.css',
+    "custom.css",
 ]
 
 html_theme_options = {
@@ -119,23 +123,29 @@
 }
 
 
-
 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",)
-
-        if msg.strip().startswith(filter_out):
-            return False
-        return True
+  """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 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))