[CI] Migrate from sphinx.ext.doctest to pytest-sphinx (ray-projec…

…t#35286)

This PR switches our doctest tooling from sphinx.ext.doctest to pytest-sphinx.

Context
We want to test all code snippets in CI.

testcode and doctest examples are tested in CI with sphinx.ext.doctest. But, we often write code-block:: python examples, and they aren't tested in CI.

Problem
sphinx.ext.doctest doesn't let you test specific files. As a result, it's hard to update examples or isolate GPU examples.

Solution
Migrate to pytest-sphinx. Two reasons:

Lets us run GPU doctests
Enables developers to quickly test their examples

---------

Signed-off-by: Balaji Veeramani <[email protected]>

.buildkite/pipeline.build.yml

@@ -491,6 +491,17 @@
       --test_env=CONDA_DEFAULT_ENV
       python/ray/tests/...
 
+- label: ":book: Doctest (CPU)"
+  instance_size: large
+  commands:
+    - cleanup() { if [ "${BUILDKITE_PULL_REQUEST}" = "false" ]; then ./ci/build/upload_build_info.sh; fi }; trap cleanup EXIT
+    - ./ci/env/env_info.sh
+    - DOC_TESTING=1 ./ci/env/install-dependencies.sh
+    - ./ci/env/install-horovod.sh
+    - bazel test --config=ci $(./scripts/bazel_export_options)
+      --test_tag_filters=doctest,-gpu
+      python/ray/... doc/...
+
 - label: ":python: Ray on Spark Test"
   conditions: ["RAY_CI_PYTHON_AFFECTED"]
   instance_size: medium

.buildkite/pipeline.gpu_large.yml

@@ -15,7 +15,7 @@
     - TRAIN_TESTING=1 DATA_PROCESSING_TESTING=1 ./ci/env/install-dependencies.sh
     - pip install -Ur ./python/requirements/ml/requirements_ml_docker.txt
     - ./ci/env/env_info.sh
-    - bazel test --config=ci $(./ci/run/bazel_export_options) --build_tests_only --test_tag_filters=datasets_train doc/...
+    - bazel test --config=ci $(./ci/run/bazel_export_options) --build_tests_only --test_tag_filters=datasets_train,-doctest doc/...
 
 - label: ":tv: :brain: RLlib: Multi-GPU Tests"
   conditions: ["NO_WHEELS_REQUIRED", "RAY_CI_RLLIB_AFFECTED"]
@@ -53,7 +53,17 @@
     - DOC_TESTING=1 TRAIN_TESTING=1 TUNE_TESTING=1 ./ci/env/install-dependencies.sh
     - pip install -Ur ./python/requirements/ml/requirements_ml_docker.txt
     - ./ci/env/env_info.sh
-    - bazel test --config=ci $(./ci/run/bazel_export_options) --build_tests_only --test_tag_filters=gpu,-timeseries_libs,-py37,-post_wheel_build doc/...
+    - bazel test --config=ci $(./ci/run/bazel_export_options) --build_tests_only --test_tag_filters=gpu,-timeseries_libs,-py37,-post_wheel_build,-doctest doc/...
+
+- label: ":book: Doctest (GPU)"
+  commands:
+    - cleanup() { if [ "${BUILDKITE_PULL_REQUEST}" = "false" ]; then ./ci/build/upload_build_info.sh; fi }; trap cleanup EXIT
+    - DOC_TESTING=1 TRAIN_TESTING=1 TUNE_TESTING=1 ./ci/env/install-dependencies.sh
+    - pip install -Ur ./python/requirements/ml/requirements_ml_docker.txt
+    - ./ci/env/install-horovod.sh
+    - ./ci/env/env_info.sh
+    - bazel test --config=ci $(./scripts/bazel_export_options)
+      --test_tag_filters=doctest,-cpu python/ray/... doc/...
 
 - label: ":zap: :python: Lightning 2.0 Train GPU tests"
   conditions:
@@ -66,4 +76,4 @@
     - pip uninstall -y pytorch-lightning
     - pip install lightning==2.0.0
     - ./ci/env/env_info.sh
-    - bazel test --config=ci $(./scripts/bazel_export_options) --test_tag_filters=ptl_v2 python/ray/train/...
+    - bazel test --config=ci $(./scripts/bazel_export_options) --test_tag_filters=ptl_v2 python/ray/train/...

.buildkite/pipeline.ml.yml

@@ -37,7 +37,7 @@
     - ./ci/env/env_info.sh
     - ./ci/run/run_bazel_test_with_sharding.sh
       --config=ci $(./ci/run/bazel_export_options)
-      --test_tag_filters=-gpu_only,-gpu,-minimal,-tune,-needs_credentials
+      --test_tag_filters=-gpu_only,-gpu,-minimal,-tune,-needs_credentials,-doctest
       python/ray/train/...
 
 - label: ":steam_locomotive: :octopus: Train + Tune tests and examples"
@@ -47,7 +47,7 @@
     - cleanup() { if [ "${BUILDKITE_PULL_REQUEST}" = "false" ]; then ./ci/build/upload_build_info.sh; fi }; trap cleanup EXIT
     - TRAIN_TESTING=1 TUNE_TESTING=1 ./ci/env/install-dependencies.sh
     - ./ci/env/env_info.sh
-    - bazel test --config=ci $(./ci/run/bazel_export_options) --build_tests_only --test_tag_filters=tune,-gpu_only,-ray_air,-gpu python/ray/train/...
+    - bazel test --config=ci $(./ci/run/bazel_export_options) --build_tests_only --test_tag_filters=tune,-gpu_only,-ray_air,-gpu,-doctest python/ray/train/...
 
 - label: ":brain: RLlib: Learning tests TF2-static-graph"
   conditions: ["NO_WHEELS_REQUIRED", "RAY_CI_RLLIB_AFFECTED"]
@@ -210,7 +210,7 @@
     - ./ci/run/run_bazel_test_with_sharding.sh --config=ci $(./ci/run/bazel_export_options)
       --build_tests_only
       --test_tag_filters=rlm
-      --test_env=RLLIB_ENABLE_RL_MODULE=1 
+      --test_env=RLLIB_ENABLE_RL_MODULE=1
       --test_env=RAY_USE_MULTIPROCESSING_CPU_COUNT=1 rllib/...
 
 - label: ":brain: RLlib: Examples"
@@ -373,7 +373,7 @@
     - cleanup() { if [ "${BUILDKITE_PULL_REQUEST}" = "false" ]; then ./ci/build/upload_build_info.sh; fi }; trap cleanup EXIT
     - TUNE_TESTING=1 INSTALL_HOROVOD=1 ./ci/env/install-dependencies.sh
     - ./ci/env/env_info.sh
-    - bazel test --config=ci $(./ci/run/bazel_export_options) --build_tests_only 
+    - bazel test --config=ci $(./ci/run/bazel_export_options) --build_tests_only
       --test_env=TUNE_NEW_EXECUTION=0
       --test_tag_filters=py37,-client python/ray/tune/...
 
@@ -495,7 +495,7 @@
     - cleanup() { if [ "${BUILDKITE_PULL_REQUEST}" = "false" ]; then ./ci/build/upload_build_info.sh; fi }; trap cleanup EXIT
     - DOC_TESTING=1 INSTALL_HOROVOD=1 ./ci/env/install-dependencies.sh
     - ./ci/env/env_info.sh
-    - bazel test --config=ci $(./ci/run/bazel_export_options) --build_tests_only --test_tag_filters=-timeseries_libs,-external,-ray_air,-gpu,-py37,-post_wheel_build doc/...
+    - bazel test --config=ci $(./ci/run/bazel_export_options) --build_tests_only --test_tag_filters=-timeseries_libs,-external,-ray_air,-gpu,-py37,-post_wheel_build,-doctest doc/...
 
 - label: ":book: Doc tests and examples with time series libraries"
   conditions:
@@ -505,7 +505,7 @@
     - cleanup() { if [ "${BUILDKITE_PULL_REQUEST}" = "false" ]; then ./ci/build/upload_build_info.sh; fi }; trap cleanup EXIT
     - DOC_TESTING=1 INSTALL_TIMESERIES_LIBS=1 ./ci/env/install-dependencies.sh
     - ./ci/env/env_info.sh
-    - bazel test --config=ci $(./ci/run/bazel_export_options) --build_tests_only --test_tag_filters=timeseries_libs,-external,-gpu,-py37,-post_wheel_build doc/...
+    - bazel test --config=ci $(./ci/run/bazel_export_options) --build_tests_only --test_tag_filters=timeseries_libs,-external,-gpu,-py37,-post_wheel_build,-doctest doc/...
 
 - label: ":book: :airplane: Ray AIR examples"
   # Todo: check if this could be a medium test. Last time it failed because of dependency issues.
@@ -516,7 +516,7 @@
     - cleanup() { if [ "${BUILDKITE_PULL_REQUEST}" = "false" ]; then ./ci/build/upload_build_info.sh; fi }; trap cleanup EXIT
     - DOC_TESTING=1 ./ci/env/install-dependencies.sh
     - ./ci/env/env_info.sh
-    - bazel test --config=ci $(./ci/run/bazel_export_options) --build_tests_only --test_tag_filters=ray_air,-external,-timeseries_libs,-needs_credentials,-gpu,-py37,-post_wheel_build doc/...
+    - bazel test --config=ci $(./ci/run/bazel_export_options) --build_tests_only --test_tag_filters=ray_air,-external,-timeseries_libs,-needs_credentials,-gpu,-py37,-post_wheel_build,-doctest doc/...
 
 - label: ":book: Doc examples with authentication "
   conditions: ["NO_WHEELS_REQUIRED", "RAY_CI_BRANCH_BUILD"]
@@ -527,7 +527,7 @@
     - DOC_TESTING=1 ./ci/env/install-dependencies.sh
     - ./ci/env/env_info.sh
     - python ./ci/env/setup_credentials.py wandb comet_ml
-    - bazel test --config=ci $(./ci/run/bazel_export_options) --build_tests_only --test_tag_filters=needs_credentials,-external,-timeseries_libs,-gpu,-py37,-post_wheel_build doc/...
+    - bazel test --config=ci $(./ci/run/bazel_export_options) --build_tests_only --test_tag_filters=needs_credentials,-external,-timeseries_libs,-gpu,-py37,-post_wheel_build,-doctest doc/...
 
 - label: ":book: Doc examples for external code "
   conditions: ["RAY_CI_PYTHON_AFFECTED", "RAY_CI_TUNE_AFFECTED", "RAY_CI_DOC_AFFECTED", "RAY_CI_SERVE_AFFECTED", "RAY_CI_ML_AFFECTED"]
@@ -536,7 +536,7 @@
     - cleanup() { if [ "${BUILDKITE_PULL_REQUEST}" = "false" ]; then ./ci/build/upload_build_info.sh; fi }; trap cleanup EXIT
     - DOC_TESTING=1 ./ci/env/install-dependencies.sh
     - ./ci/env/env_info.sh
-    - bazel test --config=ci $(./ci/run/bazel_export_options) --build_tests_only --test_tag_filters=external,-timeseries_libs,-gpu,-py37,-post_wheel_build doc/...
+    - bazel test --config=ci $(./ci/run/bazel_export_options) --build_tests_only --test_tag_filters=external,-timeseries_libs,-gpu,-py37,-post_wheel_build,-doctest doc/...
 
 
 - label: ":exploding_death_star: RLlib Contrib: A3C Tests"

bazel/BUILD

@@ -0,0 +1,4 @@
+exports_files([
+    "pytest_wrapper.py",
+    "conftest.py"
+])

bazel/conftest.py

@@ -0,0 +1,9 @@
+import pytest
+
+import ray
+
+
+@pytest.fixture(autouse=True, scope="module")
+def shutdown_ray():
+    ray.shutdown()
+    yield

bazel/pytest_wrapper.py

@@ -0,0 +1,10 @@
+import sys
+
+import pytest
+
+if __name__ == "__main__":
+    exit_code = pytest.main(sys.argv[1:])
+    if exit_code is pytest.ExitCode.NO_TESTS_COLLECTED:
+        exit_code = pytest.ExitCode.OK
+
+    sys.exit(exit_code)

bazel/python.bzl

@@ -3,6 +3,37 @@ load("@bazel_skylib//lib:paths.bzl", "paths")
 # py_test_module_list creates a py_test target for each
 # Python file in `files`
 
+
+def doctest(files, gpu = False, name="doctest", deps=[], srcs=[], data=[], args=[], size="medium", tags=[], **kwargs):
+    # NOTE: If you run `pytest` on `__init__.py`, it tries to test all files in that
+    # package. We don't want that, so we exclude it from the list of input files.
+    files = native.glob(include=files, exclude=["__init__.py"])
+    if gpu:
+        name += "[gpu]"
+        tags = tags + ["gpu"]
+    else:
+        tags = tags + ["cpu"]
+
+    native.py_test(
+        name = name,
+        srcs = ["//bazel:pytest_wrapper.py"] + srcs,
+        main = "//bazel:pytest_wrapper.py",
+        size = size,
+        args = [
+            "--doctest-modules",
+            "--capture=no",
+            "-c=$(location //bazel:conftest.py)",
+            "-v"
+        ] + args + ["$(location :%s)" % file for file in files],
+        data = ["//bazel:conftest.py"] + files + data,
+        python_version = "PY3",
+        srcs_version = "PY3",
+        tags = ["doctest"] + tags,
+        deps = ["//:ray_lib"] + deps,
+        **kwargs
+    )
+
+
 def py_test_module_list(files, size, deps, extra_srcs=[], name_suffix="", **kwargs):
     for file in files:
         # remove .py

ci/ci.sh

@@ -317,7 +317,6 @@ build_sphinx_docs() {
       # TODO: revert to "make html" once "sphinx_panels" plugin is fully removed.
       FAST=True make develop
       pip install datasets==2.0.0
-      RAY_MOCK_MODULES=0 RAY_DEDUP_LOGS=0 make doctest
     fi
   )
 }

doc/BUILD

@@ -1,5 +1,6 @@
 load("//bazel:python.bzl", "py_test_run_all_subdirectory")
 load("//bazel:python.bzl", "py_test_run_all_notebooks")
+load("//bazel:python.bzl", "doctest")
 
 exports_files(["test_myst_doc.py"])
 
@@ -299,3 +300,24 @@ py_test_run_all_subdirectory(
     extra_srcs = [],
     tags = ["exclusive", "team:ml", "external"],
 )
+
+
+# --------------------------------------------------------------------
+# Tests code snippets in user guides.
+# --------------------------------------------------------------------
+
+doctest(
+    files = glob(
+        include=["source/**/*.rst"],
+        exclude=[
+            "source/ray-contribute/getting-involved.rst",
+            "source/ray-core/handling-dependencies.rst",
+            "source/ray-core/tasks/nested-tasks.rst",
+            "source/ray-observability/user-guides/ray-tracing.rst",
+            "source/ray-observability/user-guides/cli-sdk.rst",
+            "source/rllib/rllib-env.rst",
+            "source/rllib/rllib-sample-collection.rst",
+        ]
+    ),
+    size="large"
+)

python/ray/_private/BUILD

@@ -0,0 +1,6 @@
+load("//bazel:python.bzl", "doctest")
+
+doctest(
+    files = glob(["**/*.py"]),
+    tags = ["team:core"]
+)

python/ray/air/BUILD

@@ -1,3 +1,10 @@
+load("//bazel:python.bzl", "doctest")
+
+doctest(
+    files = glob(["**/*.py"], exclude=glob(["examples/**/*", "tests/**/*", "callbacks/*.py", "integrations/wandb.py"])),  # TODO: Add note for callbacks
+    tags = ["team:ml"]
+)
+
 # --------------------------------------------------------------------
 # Tests from the python/ray/air/examples directory.
 # Please keep these sorted alphabetically.

python/ray/autoscaler/BUILD

@@ -0,0 +1,6 @@
+load("//bazel:python.bzl", "doctest")
+
+doctest(
+    files = glob(["**/*.py"], exclude=["_private/aliyun/**"]),
+    tags = ["team:core"]
+)

python/ray/dag/BUILD

@@ -1,3 +1,12 @@
+load("//bazel:python.bzl", "doctest")
+
+doctest(
+    files = glob(["**/*.py"]),
+    tags = ["team:core"],
+    deps = [":dag_lib"]
+)
+
+
 # This is a dummy test dependency that causes the above tests to be
 # re-run if any of these files changes.
 py_library(

python/ray/data/BUILD

@@ -1,3 +1,10 @@
+load("//bazel:python.bzl", "doctest")
+
+doctest(
+    files = glob(["**/*.py"], exclude=glob(["examples/**/*", "tests/**/*"])),
+    tags = ["team:data"],
+    size = "large"
+)
 # --------------------------------------------------------------------
 # Tests from the python/ray/data/tests directory.
 # Covers all tests starting with `test_`.

python/ray/experimental/BUILD

@@ -0,0 +1,5 @@
+load("//bazel:python.bzl", "doctest")
+
+doctest(
+    files = glob(["**/*.py"]),
+)

python/ray/runtime_env/BUILD

@@ -0,0 +1,6 @@
+load("//bazel:python.bzl", "doctest")
+
+doctest(
+    files = glob(["**/*.py"]),
+    tags = ["team:core"]
+)

python/ray/scripts/BUILD

@@ -0,0 +1,5 @@
+load("//bazel:python.bzl", "doctest")
+
+doctest(
+    files = glob(["**/*.py"]),
+)

python/ray/serve/BUILD

@@ -1,3 +1,10 @@
+load("//bazel:python.bzl", "doctest")
+
+doctest(
+    files = glob(["**/*.py"], exclude=["examples/**/*", "tests/**/*", "benchmarks/**/*"]),
+    tags = ["team:serve"]
+)
+
 # This is a dummy test dependency that causes the above tests to be
 # re-run if any of these files changes.
 py_library(

python/ray/train/BUILD

@@ -1,3 +1,27 @@
+load("//bazel:python.bzl", "doctest")
+
+doctest(
+    files = glob(
+        ["**/*.py"],
+        exclude=[
+            "examples/**",
+            "tests/**",
+            # GPU tests
+            "tensorflow/tensorflow_trainer.py",
+            # FIXME: Failing doctests
+            "xgboost/xgboost_predictor.py",
+            "xgboost/xgboost_checkpoint.py"
+        ]
+    ),
+    tags = ["team:ml"],
+)
+
+doctest(
+    files = ["tensorflow/tensorflow_trainer.py"],
+    tags = ["team:ml"],
+    gpu = True
+)
+
 # --------------------------------------------------------------------
 # Tests from the python/ray/train/examples directory.
 # Please keep these sorted alphabetically.

python/ray/train/tensorflow/tensorflow_trainer.py

@@ -85,7 +85,7 @@ def train_loop_per_worker():
 
     Example:
 
-    .. code-block:: python
+    .. testcode::
 
         import tensorflow as tf
 
@@ -94,9 +94,6 @@ def train_loop_per_worker():
         from ray.air.config import ScalingConfig
         from ray.train.tensorflow import TensorflowTrainer
 
-        # If using GPUs, set this to True.
-        use_gpu = False
-
         def build_model():
             # toy neural network : 1-layer
             return tf.keras.Sequential(
@@ -131,12 +128,18 @@ def train_loop_per_worker(config):
         train_dataset = ray.data.from_items([{"x": x, "y": x + 1} for x in range(32)])
         trainer = TensorflowTrainer(
             train_loop_per_worker=train_loop_per_worker,
-            scaling_config=ScalingConfig(num_workers=3, use_gpu=use_gpu),
+            scaling_config=ScalingConfig(num_workers=3, use_gpu=True),
             datasets={"train": train_dataset},
             train_loop_config={"num_epochs": 2},
         )
         result = trainer.fit()
 
+    .. testoutput::
+        :options:+ELLIPSIS
+        :hide:
+
+        ...
+
     Args:
         train_loop_per_worker: The training function to execute.
             This can either take in no arguments or a ``config`` dict.