From 1eb56c899b3517d807a97a9b657ccfcdcf2b4d77 Mon Sep 17 00:00:00 2001
From: Benson Ma <bensonma415@meta.com>
Date: Mon, 26 Feb 2024 11:21:49 -0800
Subject: [PATCH] [fbgemm_gpu] Fix mpmath dependency

- Fix mpmath dependency to allow tests to pass again; see

https://github.com/jianyicheng/mase-docker/pull/9
https://github.com/nod-ai/SHARK/issues/2095
---
 .github/scripts/utils_build.bash              |  8 ++++---
 .github/workflows/fbgemm_gpu_ci_rocm.yml      |  3 ---
 docs/requirements.txt                         |  6 ++++++
 fbgemm_gpu/docs/requirements.txt              |  6 ++++++
 .../TestInstructions.rst                      | 21 +++++++++++++++----
 fbgemm_gpu/requirements.txt                   | 13 ++++++++++++
 6 files changed, 47 insertions(+), 10 deletions(-)

diff --git a/.github/scripts/utils_build.bash b/.github/scripts/utils_build.bash
index 6d825a0fe0..3ea84ce03a 100644
--- a/.github/scripts/utils_build.bash
+++ b/.github/scripts/utils_build.bash
@@ -247,9 +247,11 @@ install_build_tools () {
 
   echo "[INSTALL] Installing build tools ..."
   # NOTES:
-  # - Only the openblas package will install cblas.h directly into
-  #   $CONDA_PREFIX/include directory
-  # - ncurses is needed to silence bad libtinfo6.so errors for ROCm+Clang builds
+  #
+  # - Only the openblas package will install <cblas.h> directly into
+  #   $CONDA_PREFIX/include directory, which is required for FBGEMM tests
+  #
+  # - ncurses is needed to silence libtinfo6.so errors for ROCm+Clang builds
   #
   # shellcheck disable=SC2086
   (exec_with_retries 3 conda install ${env_prefix} -c conda-forge -y \
diff --git a/.github/workflows/fbgemm_gpu_ci_rocm.yml b/.github/workflows/fbgemm_gpu_ci_rocm.yml
index 6f99127bbb..db70f22fdb 100644
--- a/.github/workflows/fbgemm_gpu_ci_rocm.yml
+++ b/.github/workflows/fbgemm_gpu_ci_rocm.yml
@@ -178,9 +178,6 @@ jobs:
     - name: Create Conda Environment
       run: . $PRELUDE; create_conda_environment $BUILD_ENV ${{ matrix.python-version }}
 
-    - name: Install Build Tools
-      run: . $PRELUDE; install_build_tools $BUILD_ENV
-
     - name: Install PyTorch-ROCm Nightly
       run:  . $PRELUDE; install_pytorch_pip $BUILD_ENV nightly rocm/${{ matrix.rocm-version }}
 
diff --git a/docs/requirements.txt b/docs/requirements.txt
index c6d34064aa..c0a3144411 100644
--- a/docs/requirements.txt
+++ b/docs/requirements.txt
@@ -1,2 +1,8 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
 breathe
 sphinx_rtd_theme
diff --git a/fbgemm_gpu/docs/requirements.txt b/fbgemm_gpu/docs/requirements.txt
index f62b11dae9..9f3bca4399 100644
--- a/fbgemm_gpu/docs/requirements.txt
+++ b/fbgemm_gpu/docs/requirements.txt
@@ -1,3 +1,9 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
 # Note: Sphinx 7+ currently runs into an `Undefinederror("'style' is undefined")`
 # We should eventually move to Sphinx 7+ to resolve
 # https://github.com/sphinx-doc/sphinx/issues/1514
diff --git a/fbgemm_gpu/docs/src/fbgemm_gpu-development/TestInstructions.rst b/fbgemm_gpu/docs/src/fbgemm_gpu-development/TestInstructions.rst
index 04abf5a2c8..1aabfab69c 100644
--- a/fbgemm_gpu/docs/src/fbgemm_gpu-development/TestInstructions.rst
+++ b/fbgemm_gpu/docs/src/fbgemm_gpu-development/TestInstructions.rst
@@ -4,8 +4,21 @@ Testing FBGEMM_GPU
 The tests (in the ``fbgemm_gpu/test/`` directory) and benchmarks (in the
 ``fbgemm_gpu/bench/`` directory) provide good examples on how to use FBGEMM_GPU.
 
-FBGEMM_GPU Tests
-~~~~~~~~~~~~~~~~
+Setup the FBGEMM_GPU Test Environment
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+After an environment is available from building / installing the FBGEMM_GPU
+package, additional packages need to be installed for tests to run correctly:
+
+.. code:: sh
+
+  # !! Run inside the Conda environment !!
+
+  # From the /fbgemm_gpu/ directory
+  python -m pip install -r requirements.txt
+
+Running FBGEMM_GPU Tests
+~~~~~~~~~~~~~~~~~~~~~~~~
 
 To run the tests after building / installing the FBGEMM_GPU package:
 
@@ -60,8 +73,8 @@ For ROCm machines, testing against a ROCm GPU needs to be enabled with
 
   python -m pytest -v -rsx -s -W ignore::pytest.PytestCollectionWarning split_table_batched_embeddings_test.py
 
-FBGEMM_GPU Benchmarks
-~~~~~~~~~~~~~~~~~~~~~
+Running FBGEMM_GPU Benchmarks
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
 To run the benchmarks:
 
diff --git a/fbgemm_gpu/requirements.txt b/fbgemm_gpu/requirements.txt
index 0e06a77a43..f85b60da58 100644
--- a/fbgemm_gpu/requirements.txt
+++ b/fbgemm_gpu/requirements.txt
@@ -1,6 +1,19 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# NOTES:
+#
+# - A fixed version of mpmath is needed to work around an AttributeError; see:
+#   * https://github.com/nod-ai/SHARK/issues/2095
+#   * https://github.com/jianyicheng/mase-docker/pull/9
+
 cmake
 hypothesis
 jinja2
+mpmath==1.3.0
 ninja
 numpy
 scikit-build