[Docs] Mock compiled modules (ray-project#39666)

This PR does two things: - Compiled ray modules are mocked out using sphinx's autodoc_mock_modules. This means we no longer need to roll our own mocking system for the compiled ray modules, and removes some unsightly type annotations in certain places in the docs. - Intersphinx mappings were added for a lot of ray's dependencies. For types that are documented with sphinx and hosted elsewhere, our type annotations will now be clickable links to the appropriate documentation for that type. See for example the Python type annotations below, which are links to the pages of the official Python docs for those types: Signed-off-by: pdmurray <[email protected]> Co-authored-by: Max Pumperla <[email protected]> Co-authored-by: Balaji Veeramani <[email protected]> Signed-off-by: Victor <[email protected]>
vymao · Oct 11, 2023 · 6de01e5 · 6de01e5
1 parent 99546fa
commit 6de01e5
Show file tree

Hide file tree

Showing 4 changed files with 71 additions and 49 deletions.
diff --git a/dashboard/memory_utils.py b/dashboard/memory_utils.py
@@ -14,9 +14,6 @@
 TASKID_BYTES_SIZE = TaskID.size()
 ACTORID_BYTES_SIZE = ActorID.size()
 JOBID_BYTES_SIZE = JobID.size()
-# We need to multiply 2 because we need bits size instead of bytes size.
-TASKID_RANDOM_BITS_SIZE = (TASKID_BYTES_SIZE - ACTORID_BYTES_SIZE) * 2
-ACTORID_RANDOM_BITS_SIZE = (ACTORID_BYTES_SIZE - JOBID_BYTES_SIZE) * 2
 
 
 def decode_object_ref_if_needed(object_ref: str) -> bytes:
@@ -161,14 +158,18 @@ def _get_reference_type(self) -> str:
     def _is_object_ref_actor_handle(self) -> bool:
         object_ref_hex = self.object_ref.hex()
 
+        # We need to multiply 2 because we need bits size instead of bytes size.
+        taskid_random_bits_size = (TASKID_BYTES_SIZE - ACTORID_BYTES_SIZE) * 2
+        actorid_random_bits_size = (ACTORID_BYTES_SIZE - JOBID_BYTES_SIZE) * 2
+
         # random (8B) | ActorID(6B) | flag (2B) | index (6B)
         # ActorID(6B) == ActorRandomByte(4B) + JobID(2B)
         # If random bytes are all 'f', but ActorRandomBytes
         # are not all 'f', that means it is an actor creation
         # task, which is an actor handle.
-        random_bits = object_ref_hex[:TASKID_RANDOM_BITS_SIZE]
+        random_bits = object_ref_hex[:taskid_random_bits_size]
         actor_random_bits = object_ref_hex[
-            TASKID_RANDOM_BITS_SIZE : TASKID_RANDOM_BITS_SIZE + ACTORID_RANDOM_BITS_SIZE
+            taskid_random_bits_size : taskid_random_bits_size + actorid_random_bits_size
         ]
         if random_bits == "f" * 16 and not actor_random_bits == "f" * 24:
             return True

diff --git a/doc/source/conf.py b/doc/source/conf.py
@@ -3,7 +3,6 @@
 from importlib import import_module
 import os
 import sys
-from unittest.mock import MagicMock
 from jinja2.filters import FILTERS
 
 sys.path.insert(0, os.path.abspath("."))
@@ -14,24 +13,6 @@
     build_gallery,
 )
 
-# Compiled ray modules need to be mocked out; readthedocs doesn't have support for
-# compiling these. See https://readthedocs-lst.readthedocs.io/en/latest/faq.html
-# for more information. Other external dependencies should not be added here.
-# Instead add them to autodoc_mock_imports below.
-mock_modules = [
-    "ray._raylet",
-    "ray.core.generated",
-    "ray.core.generated.common_pb2",
-    "ray.core.generated.runtime_env_common_pb2",
-    "ray.core.generated.gcs_pb2",
-    "ray.core.generated.logging_pb2",
-    "ray.core.generated.ray.protocol.Task",
-    "ray.serve.generated",
-    "ray.serve.generated.serve_pb2",
-    "ray.serve.generated.serve_pb2_grpc",
-]
-sys.modules.update((mod_name, MagicMock()) for mod_name in mock_modules)
-
 # If extensions (or modules to document with autodoc) are in another directory,
 # add these directories to sys.path here. If the directory is relative to the
 # documentation root, use os.path.abspath to make it absolute, like shown here.
@@ -457,30 +438,43 @@ def setup(app):
 
 # Mock out external dependencies here.
 autodoc_mock_imports = [
-    "transformers",
-    "horovod",
+    "aiohttp",
+    "dask",
     "datasets",
-    "tensorflow",
-    "torch",
-    "torchvision",
+    "fastapi",
+    "fsspec",
+    "grpc",
+    "gymnasium",
+    "horovod",
+    "huggingface",
+    "joblib",
     "lightgbm",
     "lightgbm_ray",
+    "nevergrad",
+    "numpy",
+    "pandas",
+    "pyarrow",
     "pytorch_lightning",
-    "xgboost",
-    "xgboost_ray",
-    "wandb",
-    "huggingface",
-    "joblib",
-    "watchfiles",
+    "scipy",
     "setproctitle",
-    "gymnasium",
-    "fastapi",
+    "skimage",
+    "sklearn",
+    "skopt",
+    "starlette",
+    "tensorflow",
+    "torch",
+    "torchvision",
+    "transformers",
     "tree",
     "uvicorn",
-    "starlette",
-    "fsspec",
-    "skimage",
-    "aiohttp",
+    "wandb",
+    "watchfiles",
+    "xgboost",
+    "xgboost_ray",
+    # Internal compiled modules
+    "ray._raylet",
+    "ray.core.generated",
+    "ray.serve.generated",
 ]
 
 
@@ -492,15 +486,42 @@ def setup(app):
     )
 
 # Other sphinx docs can be linked to if the appropriate URL to the docs
-# is specified in the `intersphinx_mapping` - for example, types in function signatures
+# is specified in the `intersphinx_mapping` - for example, types annotations
 # that are defined in dependencies can link to their respective documentation.
 intersphinx_mapping = {
+    "aiohttp": ("https://docs.aiohttp.org/en/stable/", None),
+    "dask": ("https://docs.dask.org/en/stable/", None),
+    "datasets": ("https://huggingface.co/docs/datasets/main/en/", None),
+    "distributed": ("https://distributed.dask.org/en/stable/", None),
+    "grpc": ("https://grpc.github.io/grpc/python/", None),
+    "gymnasium": ("https://gymnasium.farama.org/", None),
+    "horovod": ("https://horovod.readthedocs.io/en/stable/", None),
+    "lightgbm": ("https://lightgbm.readthedocs.io/en/latest/", None),
+    "mars": ("https://mars-project.readthedocs.io/en/latest/", None),
+    "modin": ("https://modin.readthedocs.io/en/stable/", None),
+    "nevergrad": ("https://facebookresearch.github.io/nevergrad/", None),
+    "numpy": ("https://numpy.org/doc/stable/", None),
+    "pandas": ("https://pandas.pydata.org/pandas-docs/stable/", None),
+    "pyarrow": ("https://arrow.apache.org/docs", None),
+    "pymongoarrow": ("https://mongo-arrow.readthedocs.io/en/latest/", None),
+    "pyspark": ("https://spark.apache.org/docs/latest/api/python/", None),
+    "python": ("https://docs.python.org/3", None),
+    "pytorch_lightning": ("https://lightning.ai/docs/pytorch/stable/", None),
+    "scipy": ("https://docs.scipy.org/doc/scipy/", None),
     "sklearn": ("https://scikit-learn.org/stable/", None),
+    "skopt": ("https://scikit-optimize.github.io/stable/", None),
+    "tensorflow": (
+        "https://www.tensorflow.org/api_docs/python",
+        "https://raw.githubusercontent.com/GPflow/tensorflow-intersphinx/master/tf2_py_objects.inv",
+    ),
+    "torch": ("https://pytorch.org/docs/stable/", None),
+    "torchvision": ("https://pytorch.org/vision/stable/", None),
+    "transformers": ("https://huggingface.co/docs/transformers/main/en/", None),
 }
 
 # Ray must not be imported in conf.py because third party modules initialized by
 # `import ray` will no be mocked out correctly. Perform a check here to ensure
-# ray is not imported.
+# ray is not imported by future maintainers.
 assert (
     "ray" not in sys.modules
 ), "If ray is already imported, we will not render documentation correctly!"
diff --git a/python/ray/data/dataset.py b/python/ray/data/dataset.py
@@ -4025,7 +4025,7 @@ def to_dask(
             None,
         ] = None,
         verify_meta: bool = True,
-    ) -> "dask.DataFrame":
+    ) -> "dask.dataframe.DataFrame":
         """Convert this :class:`~ray.data.Dataset` into a
         `Dask DataFrame <https://docs.dask.org/en/stable/generated/dask.dataframe.DataFrame.html#dask.dataframe.DataFrame>`_.
 
@@ -4125,7 +4125,7 @@ def block_to_df(block: Block):
         return ddf
 
     @ConsumptionAPI(pattern="Time complexity:")
-    def to_mars(self) -> "mars.DataFrame":
+    def to_mars(self) -> "mars.dataframe.DataFrame":
         """Convert this :class:`~ray.data.Dataset` into a
         `Mars DataFrame <https://mars-project.readthedocs.io/en/latest/reference/dataframe/index.html>`_.
 
@@ -4158,7 +4158,7 @@ def to_mars(self) -> "mars.DataFrame":
         return op(index_value=index_value, columns_value=columns_value, dtypes=dtypes)
 
     @ConsumptionAPI(pattern="Time complexity:")
-    def to_modin(self) -> "modin.DataFrame":
+    def to_modin(self) -> "modin.pandas.dataframe.DataFrame":
         """Convert this :class:`~ray.data.Dataset` into a
         `Modin DataFrame <https://modin.readthedocs.io/en/stable/flow/modin/pandas/dataframe.html>`_.
 

diff --git a/python/ray/data/read_api.py b/python/ray/data/read_api.py
@@ -1756,7 +1756,7 @@ def create_connection():
 
 
 @PublicAPI
-def from_dask(df: "dask.DataFrame") -> MaterializedDataset:
+def from_dask(df: "dask.dataframe.DataFrame") -> MaterializedDataset:
     """Create a :class:`~ray.data.Dataset` from a
     `Dask DataFrame <https://docs.dask.org/en/stable/generated/dask.dataframe.DataFrame.html#dask.dataframe.DataFrame>`_.
 
@@ -1792,7 +1792,7 @@ def to_ref(df):
 
 
 @PublicAPI
-def from_mars(df: "mars.DataFrame") -> MaterializedDataset:
+def from_mars(df: "mars.dataframe.DataFrame") -> MaterializedDataset:
     """Create a :class:`~ray.data.Dataset` from a
     `Mars DataFrame <https://mars-project.readthedocs.io/en/latest/reference/dataframe/index.html>`_.
 
@@ -1809,7 +1809,7 @@ def from_mars(df: "mars.DataFrame") -> MaterializedDataset:
 
 
 @PublicAPI
-def from_modin(df: "modin.DataFrame") -> MaterializedDataset:
+def from_modin(df: "modin.pandas.dataframe.DataFrame") -> MaterializedDataset:
     """Create a :class:`~ray.data.Dataset` from a
     `Modin DataFrame <https://modin.readthedocs.io/en/stable/flow/modin/pandas/dataframe.html>`_.