Adjust install and test steps based on installed variant (pytorch#2702)

Summary: - Adjust install and test steps based on installed variant Pull Request resolved: pytorch#2702 Reviewed By: jianyuh Differential Revision: D58318235 Pulled By: q10 fbshipit-source-id: b6631cd7fdccea0a41bcce13ac8626476da95b17
q10 · Jun 10, 2024 · d17fc49 · d17fc49
1 parent 98b7780
commit d17fc49
Show file tree

Hide file tree

Showing 4 changed files with 92 additions and 47 deletions.
diff --git a/.github/scripts/fbgemm_gpu_install.bash b/.github/scripts/fbgemm_gpu_install.bash
@@ -45,8 +45,22 @@ __fbgemm_gpu_post_install_checks () {
 
   echo "[INSTALL] Checking imports and symbols ..."
   (test_python_import_package "${env_name}" fbgemm_gpu) || return 1
-  (test_python_import_package "${env_name}" fbgemm_gpu.split_embedding_codegen_lookup_invokers) || return 1
   (test_python_import_symbol "${env_name}" fbgemm_gpu __version__) || return 1
+  (test_python_import_symbol "${env_name}" fbgemm_gpu __variant__) || return 1
+
+  echo "[CHECK] Printing out the FBGEMM-GPU version ..."
+  # shellcheck disable=SC2086,SC2155
+  local installed_fbgemm_gpu_version=$(conda run ${env_prefix} python -c "import fbgemm_gpu; print(fbgemm_gpu.__version__)")
+  # shellcheck disable=SC2086,SC2155
+  local installed_fbgemm_gpu_variant=$(conda run ${env_prefix} python -c "import fbgemm_gpu; print(fbgemm_gpu.__variant__)")
+  echo "################################################################################"
+  echo "[CHECK] The installed VERSION of FBGEMM_GPU is: ${installed_fbgemm_gpu_version}"
+  echo "[CHECK] The installed VARIANT of FBGEMM_GPU is: ${installed_fbgemm_gpu_variant}"
+  echo "################################################################################"
+
+  if [ "$installed_fbgemm_gpu_variant" != "genai" ]; then
+    (test_python_import_package "${env_name}" fbgemm_gpu.split_embedding_codegen_lookup_invokers) || return 1
+  fi
 
   echo "[INSTALL] Checking operator registrations ..."
   local test_operator="torch.ops.fbgemm.asynchronous_inclusive_cumsum"
@@ -63,13 +77,6 @@ __fbgemm_gpu_post_install_checks () {
     echo "################################################################################"
     return 1
   fi
-
-  echo "[CHECK] Printing out the FBGEMM-GPU version ..."
-  # shellcheck disable=SC2086,SC2155
-  local installed_fbgemm_gpu_version=$(conda run ${env_prefix} python -c "import fbgemm_gpu; print(fbgemm_gpu.__version__)")
-  echo "################################################################################"
-  echo "[CHECK] The installed version of FBGEMM_GPU is: ${installed_fbgemm_gpu_version}"
-  echo "################################################################################"
 }
 
 install_fbgemm_gpu_wheel () {

diff --git a/.github/scripts/fbgemm_gpu_test.bash b/.github/scripts/fbgemm_gpu_test.bash
@@ -110,16 +110,16 @@ __setup_fbgemm_gpu_test () {
 
   # Configure the environment for ignored test suites for each FBGEMM_GPU
   # variant
-  if [ "$fbgemm_variant" == "cpu" ]; then
+  if [ "$fbgemm_gpu_variant" == "cpu" ]; then
     echo "[TEST] Configuring for CPU-based testing ..."
     __configure_fbgemm_gpu_test_cpu
 
-  elif [ "$fbgemm_variant" == "rocm" ]; then
+  elif [ "$fbgemm_gpu_variant" == "rocm" ]; then
     echo "[TEST] Configuring for ROCm-based testing ..."
     __configure_fbgemm_gpu_test_rocm
 
   else
-    echo "[TEST] Configuring for CUDA-based testing ..."
+    echo "[TEST] FBGEMM_GPU variant is ${fbgemm_gpu_variant}; configuring for CUDA-based testing ..."
     __configure_fbgemm_gpu_test_cuda
   fi
 
@@ -139,7 +139,9 @@ __setup_fbgemm_gpu_test () {
 
   echo "[TEST] Checking imports ..."
   (test_python_import_package "${env_name}" fbgemm_gpu) || return 1
-  (test_python_import_package "${env_name}" fbgemm_gpu.split_embedding_codegen_lookup_invokers) || return 1
+  if [ "$fbgemm_gpu_variant" != "genai" ]; then
+    (test_python_import_package "${env_name}" fbgemm_gpu.split_embedding_codegen_lookup_invokers) || return 1
+  fi
 
   # Configure the PyTest args
   pytest_args=(
@@ -153,33 +155,21 @@ __setup_fbgemm_gpu_test () {
   echo "[TEST] PyTest args:  ${pytest_args[@]}"
 }
 
-
 ################################################################################
 # FBGEMM_GPU Test Functions
 ################################################################################
 
-run_fbgemm_gpu_tests () {
-  env_name="$1"
-  fbgemm_variant="$2"
-  if [ "$fbgemm_variant" == "" ]; then
-    echo "Usage: ${FUNCNAME[0]} ENV_NAME [FBGEMM_VARIANT]"
-    echo "Example(s):"
-    echo "    ${FUNCNAME[0]} build_env cpu    # Run all tests applicable to CPU"
-    echo "    ${FUNCNAME[0]} build_env cuda   # Run all tests applicable to CUDA"
-    echo "    ${FUNCNAME[0]} build_env rocm   # Run all tests applicable to ROCm"
-    return 1
-  else
-    echo "################################################################################"
-    echo "# Run FBGEMM-GPU Tests"
-    echo "#"
-    echo "# [$(date --utc +%FT%T.%3NZ)] + ${FUNCNAME[0]} ${*}"
-    echo "################################################################################"
-    echo ""
-  fi
+__run_fbgemm_gpu_tests_in_directory () {
+  echo "################################################################################"
+  # shellcheck disable=SC2154
+  echo "# Run FBGEMM-GPU Tests: ${pwd}"
+  echo "#"
+  echo "# [$(date --utc +%FT%T.%3NZ)] + ${FUNCNAME[0]} ${*}"
+  echo "################################################################################"
+  echo ""
 
   # shellcheck disable=SC2155
   local env_prefix=$(env_name_or_prefix "${env_name}")
-  __setup_fbgemm_gpu_test
 
   echo "[TEST] Enumerating ALL test files ..."
   # shellcheck disable=SC2155
@@ -205,26 +195,70 @@ run_fbgemm_gpu_tests () {
   done
 }
 
-test_all_fbgemm_gpu_modules () {
-  local env_name="$1"
-  local fbgemm_variant="$2"
+__determine_test_directories () {
+  target_directories=()
 
-  local target_directories=(
-    fbgemm_gpu/test
-  )
+  if [ "$fbgemm_gpu_variant" != "genai" ]; then
+    target_directories+=(
+      fbgemm_gpu/test
+    )
+  fi
 
-  if [ "$fbgemm_variant" == "cuda" ]; then
+  if [ "$fbgemm_gpu_variant" == "cuda" ] || [ "$fbgemm_gpu_variant" == "genai" ]; then
     target_directories+=(
       fbgemm_gpu/experimental/example/test
       fbgemm_gpu/experimental/gemm/test
       fbgemm_gpu/experimental/gen_ai/test
     )
   fi
 
+  echo "[TEST] Determined the testing directories:"
+  for test_dir in "${target_directories[@]}"; do
+    echo "$test_dir"
+  done
+}
+
+test_all_fbgemm_gpu_modules () {
+  env_name="$1"
+  fbgemm_gpu_variant="$2"
+  if [ "$env_name" == "" ]; then
+    echo "Usage: ${FUNCNAME[0]} ENV_NAME [FBGEMM_GPU_VARIANT]"
+    echo "Example(s):"
+    echo "    ${FUNCNAME[0]} build_env        # Test all FBGEMM_GPU modules applicable to to the installed variant"
+    echo "    ${FUNCNAME[0]} build_env cpu    # Test all FBGEMM_GPU modules applicable to CPU"
+    echo "    ${FUNCNAME[0]} build_env cuda   # Test all FBGEMM_GPU modules applicable to CUDA"
+    echo "    ${FUNCNAME[0]} build_env rocm   # Test all FBGEMM_GPU modules applicable to ROCm"
+    return 1
+  else
+    echo "################################################################################"
+    echo "# Test All FBGEMM-GPU Modules"
+    echo "#"
+    echo "# [$(date --utc +%FT%T.%3NZ)] + ${FUNCNAME[0]} ${*}"
+    echo "################################################################################"
+    echo ""
+  fi
+
+  # shellcheck disable=SC2155
+  local env_prefix=$(env_name_or_prefix "${env_name}")
+
+  # Determine the FBGEMM_GPU varaiant if needed
+  if [ "$fbgemm_gpu_variant" == "" ]; then
+    # shellcheck disable=SC2086
+    fbgemm_gpu_variant=$(conda run ${env_prefix} python -c "import fbgemm_gpu; print(fbgemm_gpu.__variant__)")
+    echo "[TEST] Determined FBGEMM_GPU variant from installation: ${fbgemm_gpu_variant}"
+  fi
+
+  # Determine the test directories to include for testing
+  __determine_test_directories
+
+  # Set the ignored tests and PyTest args
+  __setup_fbgemm_gpu_test
+
+  # Iterate through the test directories and run bulk tests
   for test_dir in "${target_directories[@]}"; do
-    cd "${test_dir}"                                        || return 1
-    run_fbgemm_gpu_tests "${env_name}" "${fbgemm_variant}"  || return 1
-    cd -                                                    || return 1
+    cd "${test_dir}"                                                          || return 1
+    __run_fbgemm_gpu_tests_in_directory "${env_name}" "${fbgemm_gpu_variant}" || return 1
+    cd -                                                                      || return 1
   done
 }
 

diff --git a/fbgemm_gpu/fbgemm_gpu/__init__.py b/fbgemm_gpu/fbgemm_gpu/__init__.py
@@ -32,7 +32,7 @@
 import fbgemm_gpu.docs  # noqa: F401, E402
 
 # Export the version string from the version file auto-generated by setup.py
-from fbgemm_gpu.docs.version import __version__  # noqa: F401, E402
+from fbgemm_gpu.docs.version import __variant__, __version__  # noqa: F401, E402
 
 try:
     # Trigger meta operator registrations

diff --git a/fbgemm_gpu/setup.py b/fbgemm_gpu/setup.py
@@ -33,7 +33,7 @@ class FbgemmGpuBuild:
 
     @classmethod
     def from_args(cls, argv: List[str]):
-        parser = argparse.ArgumentParser(description="fbgemm_gpu setup")
+        parser = argparse.ArgumentParser(description="FBGEMM_GPU Build Setup")
         parser.add_argument(
             "--verbose",
             action="store_true",
@@ -389,8 +389,10 @@ class FbgemmGpuInstall(PipInstall):
     """FBGEMM_GPU PIP Install Routines"""
 
     @classmethod
-    def generate_version_file(cls, package_version: str) -> None:
+    def generate_version_file(cls, build: FbgemmGpuBuild) -> None:
         with open("fbgemm_gpu/docs/version.py", "w") as file:
+            package_version = build.package_version()
+
             print(
                 f"[SETUP.PY] Generating version file at: {os.path.realpath(file.name)}"
             )
@@ -404,6 +406,7 @@ def generate_version_file(cls, package_version: str) -> None:
                 # LICENSE file in the root directory of this source tree.
 
                 __version__: str = "{package_version}"
+                __variant__: str = "{build.args.package_variant}"
                 """
             )
             file.write(text)
@@ -474,7 +477,8 @@ def main(argv: List[str]) -> None:
 
     # Extract the package name
     package_name = build.package_name()
-    # Generate the full package version string
+
+    # Extract the package version
     package_version = build.package_version()
 
     if build.args.dryrun:
@@ -485,7 +489,7 @@ def main(argv: List[str]) -> None:
         sys.exit(0)
 
     # Generate the version file
-    FbgemmGpuInstall.generate_version_file(package_version)
+    FbgemmGpuInstall.generate_version_file(build)
 
     setup(
         name=package_name,