Fix python docs not being visible (pytorch#799)

Summary:
Pull Request resolved: facebookresearch/FBGEMM#799

- Fix python docs not being visible.  Bisection reveals that this bug first appeared in 887e5bf (pytorch#3477)

X-link: pytorch#3717

Reviewed By: sryap

Differential Revision: D69897513

Pulled By: q10

fbshipit-source-id: f04c1bb1986003161a2c81186b2981e39f95a4b7

Author: q10

.github/scripts/fbgemm_gpu_build.bash

@@ -145,6 +145,14 @@ __configure_fbgemm_gpu_build_cpu () {
   )
 }
 
+__configure_fbgemm_gpu_build_docs () {
+  # Update the package name and build args depending on if CUDA is specified
+  echo "[BUILD] Setting CPU-only (docs) build args ..."
+  build_args=(
+    --package_variant=docs
+  )
+}
+
 __configure_fbgemm_gpu_build_rocm () {
   local fbgemm_variant_targets="$1"
 
@@ -299,6 +307,10 @@ __configure_fbgemm_gpu_build () {
     echo "[BUILD] Configuring build as CPU variant ..."
     __configure_fbgemm_gpu_build_cpu
 
+  elif [ "$fbgemm_variant" == "docs" ]; then
+    echo "[BUILD] Configuring build as CPU (docs) variant ..."
+    __configure_fbgemm_gpu_build_docs
+
   elif [ "$fbgemm_variant" == "rocm" ]; then
     echo "[BUILD] Configuring build as ROCm variant ..."
     __configure_fbgemm_gpu_build_rocm "${fbgemm_variant_targets}"
@@ -401,9 +413,10 @@ __build_fbgemm_gpu_common_pre_steps () {
   (test_binpath "${env_name}" g++) || return 1
 
   # Set the default the FBGEMM_GPU variant to be CUDA
-  if [ "$fbgemm_variant" != "cpu" ] &&
-     [ "$fbgemm_variant" != "rocm" ] &&
-     [ "$fbgemm_variant" != "genai" ]; then
+  if  [ "$fbgemm_variant" != "cpu" ] &&
+      [ "$fbgemm_variant" != "docs" ] &&
+      [ "$fbgemm_variant" != "rocm" ] &&
+      [ "$fbgemm_variant" != "genai" ]; then
     echo "################################################################################"
     echo "[BUILD] Unknown FBGEMM_GPU variant: $fbgemm_variant"
     echo "[BUILD] Defaulting to CUDA"
@@ -501,7 +514,8 @@ __verify_library_symbols () {
 
   # Prepare a sample set of symbols whose existence in the built library should be checked
   # This is by no means an exhaustive set, and should be updated accordingly
-  if [ "${fbgemm_variant}" == "cpu" ]; then
+  if  [ "${fbgemm_variant}" == "cpu" ] ||
+      [ "${fbgemm_variant}" == "docs" ]; then
     local lib_symbols_to_check=(
       fbgemm_gpu::asynchronous_inclusive_cumsum_cpu
       fbgemm_gpu::jagged_2d_to_dense
@@ -539,6 +553,7 @@ run_fbgemm_gpu_postbuild_checks () {
     echo "Usage: ${FUNCNAME[0]} FBGEMM_VARIANT"
     echo "Example(s):"
     echo "    ${FUNCNAME[0]} cpu"
+    echo "    ${FUNCNAME[0]} docs"
     echo "    ${FUNCNAME[0]} cuda"
     echo "    ${FUNCNAME[0]} rocm"
     echo "    ${FUNCNAME[0]} genai"
@@ -590,6 +605,7 @@ build_fbgemm_gpu_package () {
     echo "Usage: ${FUNCNAME[0]} ENV_NAME RELEASE_CHANNEL VARIANT [VARIANT_TARGETS]"
     echo "Example(s):"
     echo "    ${FUNCNAME[0]} build_env release cpu                      # CPU-only variant"
+    echo "    ${FUNCNAME[0]} build_env release docs                     # CPU-only (docs) variant"
     echo "    ${FUNCNAME[0]} build_env nightly cuda                     # CUDA variant for default target(s)"
     echo "    ${FUNCNAME[0]} build_env test cuda '7.0;8.0'              # CUDA variant for custom target(s)"
     echo "    ${FUNCNAME[0]} build_env test rocm                        # ROCm variant for default target(s)"
@@ -659,6 +675,7 @@ build_fbgemm_gpu_install () {
     echo "Usage: ${FUNCNAME[0]} ENV_NAME VARIANT [TARGETS]"
     echo "Example(s):"
     echo "    ${FUNCNAME[0]} build_env cpu                          # CPU-only variant"
+    echo "    ${FUNCNAME[0]} build_env docs                         # CPU-only (docs) variant"
     echo "    ${FUNCNAME[0]} build_env cuda                         # CUDA variant for default target(s)"
     echo "    ${FUNCNAME[0]} build_env cuda '7.0;8.0'               # CUDA variant for custom target(s)"
     echo "    ${FUNCNAME[0]} build_env rocm                         # ROCm variant for default target(s)"

.github/workflows/fbgemm_gpu_docs.yml

@@ -83,7 +83,7 @@ jobs:
       run: . $PRELUDE; cd fbgemm_gpu; prepare_fbgemm_gpu_build $BUILD_ENV
 
     - name: Build + Install FBGEMM_GPU (CPU version)
-      run: . $PRELUDE; cd fbgemm_gpu; build_fbgemm_gpu_install $BUILD_ENV $BUILD_VARIANT
+      run: . $PRELUDE; cd fbgemm_gpu; build_fbgemm_gpu_install $BUILD_ENV docs
 
     - name: Build FBGEMM_GPU Documentation
       run: . $PRELUDE; cd fbgemm_gpu/docs; build_fbgemm_gpu_docs $BUILD_ENV

fbgemm_gpu/docs/src/fbgemm_gpu-development/BuildInstructions.rst

@@ -543,6 +543,9 @@ For CPU-only builds, the ``--cpu_only`` flag needs to be specified.
   python setup.py install \
       --package_variant=cpu
 
+  # NOTE: To build the package as part of generating the documentation, use
+  # `--package_variant=docs` flag instead!
+
 To build using Clang + ``libstdc++`` instead of GCC, simply append the
 ``--cxxprefix`` flag:
 

fbgemm_gpu/fbgemm_gpu/__init__.py

@@ -11,14 +11,15 @@
 import torch
 
 
-def _load_library(filename: str) -> None:
+def _load_library(filename: str, no_throw: bool = False) -> None:
     """Load a shared library from the given filename."""
     try:
         torch.ops.load_library(os.path.join(os.path.dirname(__file__), filename))
         logging.info(f"Successfully loaded: '{filename}'")
     except Exception as error:
         logging.error(f"Could not load the library '{filename}': {error}")
-        raise error
+        if not no_throw:
+            raise error
 
 
 # Since __init__.py is only used in OSS context, we define `open_source` here
@@ -67,13 +68,24 @@ def _load_library(filename: str) -> None:
 
 libraries_to_load = {
     "cpu": fbgemm_gpu_libraries,
+    "docs": fbgemm_gpu_libraries,
     "cuda": fbgemm_gpu_libraries + fbgemm_gpu_genai_libraries,
     "genai": fbgemm_gpu_genai_libraries,
     "rocm": fbgemm_gpu_libraries,
 }
 
 for library in libraries_to_load.get(__variant__, []):
-    _load_library(f"{library}.so")
+    # NOTE: In all cases, we want to throw an error if we cannot load the
+    # library.  However, this appears to break the OSS documentation build,
+    # where the Python documentation doesn't show up in the generated docs.
+    #
+    # To work around this problem, we introduce a fake build variant called
+    # `docs` and we only throw a library load error when the variant is not
+    # `docs`.  For more information, see:
+    #
+    #   https://github.com/pytorch/FBGEMM/pull/3477
+    #   https://github.com/pytorch/FBGEMM/pull/3717
+    _load_library(f"{library}.so", __variant__ == "docs")
 
 try:
     # Trigger meta operator registrations

fbgemm_gpu/setup.py

@@ -52,7 +52,7 @@ def from_args(cls, argv: List[str]):
         parser.add_argument(
             "--package_variant",
             type=str,
-            choices=["cpu", "cuda", "rocm", "genai"],
+            choices=["docs", "cpu", "cuda", "rocm", "genai"],
             default="cuda",
             help="The FBGEMM_GPU variant to build.",
         )
@@ -270,7 +270,16 @@ def _get_cxx11_abi():
             # https://stackoverflow.com/questions/44284275/passing-compiler-options-in-cmake-command-line
             cxx_flags.extend(["-DTORCH_USE_CUDA_DSA", "-DTORCH_USE_HIP_DSA"])
 
-        if self.args.package_variant == "cpu":
+        if self.args.package_variant in ["docs", "cpu"]:
+            # NOTE: The docs variant is a fake variant that is effectively the
+            # cpu variant, but marks __VARIANT__ as "docs" instead of "cpu".
+            #
+            # This minor change lets the library loader know not throw
+            # exceptions on failed load, which is the workaround for a bug in
+            # the Sphinx documentation generation process, see:
+            #
+            #   https://github.com/pytorch/FBGEMM/pull/3477
+            #   https://github.com/pytorch/FBGEMM/pull/3717
             print("[SETUP.PY] Building the CPU-ONLY variant of FBGEMM_GPU ...")
             cmake_args.append("-DFBGEMM_CPU_ONLY=ON")
 

netlify.toml

@@ -29,7 +29,7 @@
     # Build the code
     cd ..
     prepare_fbgemm_gpu_build    $BUILD_ENV
-    build_fbgemm_gpu_install    $BUILD_ENV cpu
+    build_fbgemm_gpu_install    $BUILD_ENV docs
 
     # Build the docs
     cd docs