Auto-generate the version file (pytorch#2011)

Summary:
- Auto-generate the version file in OSS, so that the `__version__` symbol is available in the package

Pull Request resolved: pytorch#2011

Reviewed By: shintaro-iwasaki

Differential Revision: D49174509

Pulled By: q10

fbshipit-source-id: d0dfd1ab0a2912016ad6e003bf88feeab696fa4a

.github/scripts/fbgemm_gpu_build.bash

@@ -41,8 +41,8 @@ prepare_fbgemm_gpu_build () {
   echo "[BUILD] Installing other build dependencies ..."
   (exec_with_retries conda run -n "${env_name}" python -m pip install -r requirements.txt) || return 1
 
-  (test_python_import "${env_name}" numpy) || return 1
-  (test_python_import "${env_name}" skbuild) || return 1
+  (test_python_import_package "${env_name}" numpy) || return 1
+  (test_python_import_package "${env_name}" skbuild) || return 1
 
   echo "[BUILD] Successfully ran git submodules update"
 }
@@ -360,7 +360,7 @@ build_fbgemm_gpu_install () {
   # Exit this directory to prevent import clashing, since there is an
   # fbgemm_gpu/ subdirectory present
   cd - || return 1
-  (test_python_import "${env_name}" fbgemm_gpu) || return 1
+  (test_python_import_package "${env_name}" fbgemm_gpu) || return 1
 
   echo "[BUILD] FBGEMM-GPU build + install completed"
 }

.github/scripts/fbgemm_gpu_install.bash

@@ -13,6 +13,17 @@
 # FBGEMM_GPU Install Functions
 ################################################################################
 
+__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
+
+  echo "[CHECK] Printing out the FBGEMM-GPU version ..."
+  installed_fbgemm_gpu_version=$(conda run -n "${env_name}" python -c "import fbgemm_gpu; print(fbgemm_gpu.__version__)")
+  echo "[CHECK] The installed version is: ${installed_fbgemm_gpu_version}"
+}
+
 install_fbgemm_gpu_wheel () {
   local env_name="$1"
   local wheel_path="$2"
@@ -38,9 +49,7 @@ install_fbgemm_gpu_wheel () {
   echo "[INSTALL] Installing FBGEMM-GPU wheel: ${wheel_path} ..."
   (exec_with_retries conda run -n "${env_name}" python -m pip install "${wheel_path}") || return 1
 
-  echo "[INSTALL] Checking imports ..."
-  (test_python_import "${env_name}" fbgemm_gpu) || return 1
-  (test_python_import "${env_name}" fbgemm_gpu.split_embedding_codegen_lookup_invokers) || return 1
+  __fbgemm_gpu_post_install_checks || return 1
 
   echo "[INSTALL] FBGEMM-GPU installation through wheel completed ..."
 }
@@ -126,9 +135,7 @@ install_fbgemm_gpu_pip () {
   # shellcheck disable=SC2086
   (exec_with_retries conda run -n "${env_name}" pip install ${fbgemm_gpu_package}) || return 1
 
-  echo "[INSTALL] Checking imports ..."
-  (test_python_import "${env_name}" fbgemm_gpu) || return 1
-  (test_python_import "${env_name}" fbgemm_gpu.split_embedding_codegen_lookup_invokers) || return 1
+  __fbgemm_gpu_post_install_checks || return 1
 
   echo "[INSTALL] FBGEMM-GPU installation through PIP completed ..."
 }

.github/scripts/fbgemm_gpu_lint.bash

@@ -42,7 +42,7 @@ install_lint_tools () {
   # Check Python packages are importable
   local import_tests=( click )
   for p in "${import_tests[@]}"; do
-    (test_python_import "${env_name}" "${p}") || return 1
+    (test_python_import_package "${env_name}" "${p}") || return 1
   done
 
   echo "[INSTALL] Successfully installed all the lint tools"

.github/scripts/fbgemm_gpu_test.bash

@@ -93,8 +93,8 @@ run_fbgemm_gpu_tests () {
   print_exec conda install -n "${env_name}" -y pytest
 
   echo "[TEST] Checking imports ..."
-  (test_python_import "${env_name}" fbgemm_gpu) || return 1
-  (test_python_import "${env_name}" fbgemm_gpu.split_embedding_codegen_lookup_invokers) || return 1
+  (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
 
   echo "[TEST] Enumerating test files ..."
   print_exec ls -lth ./*.py

.github/scripts/setup_env.bash

@@ -179,7 +179,7 @@ install_build_tools () {
   # Check Python packages are importable
   local import_tests=( click hypothesis jinja2 numpy skbuild wheel )
   for p in "${import_tests[@]}"; do
-    (test_python_import "${env_name}" "${p}") || return 1
+    (test_python_import_package "${env_name}" "${p}") || return 1
   done
 
   echo "[INSTALL] Successfully installed all the build tools"
@@ -212,8 +212,8 @@ publish_to_pypi () {
 
   echo "[INSTALL] Installing twine ..."
   print_exec conda install -n "${env_name}" -y twine
-  (test_python_import "${env_name}" twine) || return 1
-  (test_python_import "${env_name}" OpenSSL) || return 1
+  (test_python_import_package "${env_name}" twine) || return 1
+  (test_python_import_package "${env_name}" OpenSSL) || return 1
 
   echo "[PUBLISH] Uploading package(s) to PyPI: ${package_name} ..."
   conda run -n "${env_name}" \

.github/scripts/utils_base.bash

@@ -70,7 +70,26 @@ exec_with_retries () {
 # Assert Functions
 ################################################################################
 
-test_python_import () {
+test_python_import_symbol () {
+  local env_name="$1"
+  local package_name="$2"
+  local target_symbol="$3"
+  if [ "$target_symbol" == "" ]; then
+    echo "Usage: ${FUNCNAME[0]} ENV_NAME PACKAGE_NAME SYMBOL"
+    echo "Example(s):"
+    echo "    ${FUNCNAME[0]} build_env numpy __version__"
+    return 1
+  fi
+
+  if conda run -n "${env_name}" python -c "from ${package_name} import ${target_symbol}"; then
+    echo "[CHECK] Found symbol '${target_symbol}' in Python package '${package_name}'."
+  else
+    echo "[CHECK] Could not find symbol '${target_symbol}' in Python package '${package_name}'; the package might be missing or broken."
+    return 1
+  fi
+}
+
+test_python_import_package () {
   local env_name="$1"
   local python_import="$2"
   if [ "$python_import" == "" ]; then
@@ -81,9 +100,9 @@ test_python_import () {
   fi
 
   if conda run -n "${env_name}" python -c "import ${python_import}"; then
-    echo "[CHECK] Python package ${python_import} found."
+    echo "[CHECK] Python package '${python_import}' found."
   else
-    echo "[CHECK] Python package ${python_import} was not found or is broken!"
+    echo "[CHECK] Python package '${python_import}' was not found, or the package is broken!"
     return 1
   fi
 }

.github/scripts/utils_conda.bash

@@ -116,7 +116,7 @@ create_conda_environment () {
 
   # This test fails with load errors if the pyOpenSSL and cryptography package versions don't align
   echo "[SETUP] Testing pyOpenSSL import ..."
-  (test_python_import "${env_name}" OpenSSL) || return 1
+  (test_python_import_package "${env_name}" OpenSSL) || return 1
 
   echo "[SETUP] Installed Python version: $(conda run -n "${env_name}" python --version)"
   echo "[SETUP] Successfully created Conda environment: ${env_name}"

.github/scripts/utils_pytorch.bash

@@ -65,7 +65,7 @@ install_pytorch_conda () {
   (exec_with_retries conda install --force-reinstall -n "${env_name}" -y ${pytorch_package} -c "${pytorch_channel}") || return 1
 
   # Check that PyTorch is importable
-  (test_python_import "${env_name}" torch.distributed) || return 1
+  (test_python_import_package "${env_name}" torch.distributed) || return 1
 
   # Print out the actual installed PyTorch version
   installed_pytorch_version=$(conda run -n "${env_name}" python -c "import torch; print(torch.__version__)")
@@ -154,7 +154,7 @@ install_pytorch_pip () {
   (exec_with_retries conda run -n "${env_name}" pip install ${pytorch_package} --extra-index-url ${pytorch_channel}) || return 1
 
   # Check that PyTorch is importable
-  (test_python_import "${env_name}" torch.distributed) || return 1
+  (test_python_import_package "${env_name}" torch.distributed) || return 1
 
   # Print out the actual installed PyTorch version
   installed_pytorch_version=$(conda run -n "${env_name}" python -c "import torch; print(torch.__version__)")

fbgemm_gpu/fbgemm_gpu/__init__.py

@@ -18,4 +18,8 @@
 # Use existence to check if fbgemm_gpu_py.so has already been loaded
 open_source: bool = True
 
+# Re-export docs
 from . import _fbgemm_gpu_docs  # noqa: F401, E402
+
+# Re-export the version string from the auto-generated version file
+from ._fbgemm_gpu_version import __version__  # noqa: F401, E402

fbgemm_gpu/setup.py

@@ -4,13 +4,15 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# @licenselint-loose-mode
 
 import argparse
 import os
 import random
 import re
 import subprocess
 import sys
+import textwrap
 
 from datetime import date
 from typing import List, Optional
@@ -179,6 +181,26 @@ def _get_cxx11_abi():
 class FbgemmGpuInstaller(PipInstall):
     """FBGEMM_GPU PIP Installer"""
 
+    @classmethod
+    def generate_version_file(cls, package_version: str) -> None:
+        with open("fbgemm_gpu/_fbgemm_gpu_version.py", "w") as file:
+            print(
+                f"[SETUP.PY] Generating version file at: {os.path.realpath(file.name)}"
+            )
+            text = textwrap.dedent(
+                f"""
+                #!/usr/bin/env python3
+                # 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.
+
+                __version__: str = "{package_version}"
+                """
+            )
+            file.write(text)
+
     @classmethod
     def description(cls) -> str:
         # Get the long description from the relevant file
@@ -250,9 +272,15 @@ def main(argv: List[str]) -> None:
     # Repair command line args for setup.
     sys.argv = [sys.argv[0]] + unknown
 
+    # Determine the package version
+    package_version = generate_package_version(args.package_name)
+
+    # Generate the version file
+    FbgemmGpuInstaller.generate_version_file(package_version)
+
     setup(
         name=args.package_name,
-        version=generate_package_version(args.package_name),
+        version=package_version,
         author="FBGEMM Team",
         author_email="[email protected]",
         long_description=FbgemmGpuInstaller.description(),