Re-organize documentation around docs (#2258)

Summary: - Clean up the scripts for handling platform name when building packages - Clarify and re-organize the docs around documenting the codebase Pull Request resolved: #2258 Test Plan: https://deploy-preview-2258--pytorch-fbgemm-docs.netlify.app/ Reviewed By: jianyuh Differential Revision: D52614741 Pulled By: q10 fbshipit-source-id: 2e3bb4d2d54b2233d639971421cfa637150d5805
pytorch · Jan 9, 2024 · 127474c · 127474c
1 parent 0a23158
commit 127474c
Show file tree

Hide file tree

Showing 11 changed files with 602 additions and 458 deletions.
diff --git a/.github/scripts/fbgemm_gpu_build.bash b/.github/scripts/fbgemm_gpu_build.bash
@@ -194,6 +194,76 @@ __configure_fbgemm_gpu_build () {
   echo "[BUILD] FBGEMM_GPU build arguments have been set:  ${build_args[@]}"
 }
 
+__build_fbgemm_gpu_set_package_name () {
+  # Determine the package name based on release type and variant
+  export package_name="fbgemm_gpu"
+
+  # Append qualifiers for the non-release version
+  if [ "$fbgemm_release_type" != "release" ]; then
+    export package_name="${package_name}_${fbgemm_release_type}"
+  fi
+
+  # Append cpu or rocm for the non-CUDA case
+  if [ "$fbgemm_variant" == "cpu" ]; then
+    export package_name="${package_name}-cpu"
+  elif [ "$fbgemm_variant" == "rocm" ]; then
+    export package_name="${package_name}-rocm"
+  fi
+
+  echo "[BUILD] Determined and set Python package name to use: ${package_name}"
+}
+
+__build_fbgemm_gpu_set_python_tag () {
+  # shellcheck disable=SC2207,SC2086
+  local python_version=($(conda run --no-capture-output ${env_prefix} python --version))
+
+  # shellcheck disable=SC2206
+  local python_version_arr=(${python_version[1]//./ })
+
+  # Set the python tag (e.g. Python 3.12 -> py312)
+  export python_tag="py${python_version_arr[0]}${python_version_arr[1]}"
+  echo "[BUILD] Extracted and set Python tag: ${python_tag}"
+}
+
+__build_fbgemm_gpu_set_python_plat_name () {
+  if [[ $KERN_NAME == 'Darwin' ]]; then
+    # This follows PyTorch package naming conventions
+    # See https://pypi.org/project/torch/#files
+    if [[ $MACHINE_NAME == 'arm64' ]]; then
+      export python_plat_name="macosx_11_0_${MACHINE_NAME}"
+    else
+      export python_plat_name="macosx_10_9_${MACHINE_NAME}"
+    fi
+
+  elif [[ $KERN_NAME == 'Linux' ]]; then
+    # manylinux2014 is specified, bc manylinux1 does not support aarch64
+    # See https://github.com/pypa/manylinux
+    export python_plat_name="manylinux2014_${MACHINE_NAME}"
+
+  else
+    echo "[BUILD] Unsupported OS platform: ${KERN_NAME}"
+    return 1
+  fi
+
+  echo "[BUILD] Extracted and set Python platform name: ${python_plat_name}"
+}
+
+__build_fbgemm_gpu_set_run_multicore () {
+  # shellcheck disable=SC2155
+  local core=$(lscpu | grep "Core(s)" | awk '{print $NF}') && echo "core = ${core}" || echo "core not found"
+  # shellcheck disable=SC2155
+  local sockets=$(lscpu | grep "Socket(s)" | awk '{print $NF}') && echo "sockets = ${sockets}" || echo "sockets not found"
+  local re='^[0-9]+$'
+
+  export run_multicore=""
+  if [[ $core =~ $re && $sockets =~ $re ]] ; then
+    local n_core=$((core * sockets))
+    export run_multicore=" -j ${n_core}"
+  fi
+
+  echo "[BUILD] Set multicore run option for setup.py: ${run_multicore}"
+}
+
 __build_fbgemm_gpu_common_pre_steps () {
   # Private function that uses variables instantiated by its caller
 
@@ -203,28 +273,23 @@ __build_fbgemm_gpu_common_pre_steps () {
   (test_binpath "${env_name}" c++) || return 1
   (test_binpath "${env_name}" g++) || return 1
 
-  # Determine the package name based on release type and variant
-  package_name="fbgemm_gpu"
-  if [ "$fbgemm_release_type" != "release" ]; then
-    package_name="${package_name}_${fbgemm_release_type}"
+  # Set the default the FBGEMM_GPU variant to be CUDA
+  if [ "$fbgemm_variant" != "cpu" ] && [ "$fbgemm_variant" != "rocm" ]; then
+    export fbgemm_variant="cuda"
   fi
-  if [ "$fbgemm_variant" == "cpu" ]; then
-    package_name="${package_name}-cpu"
-  elif [ "$fbgemm_variant" == "rocm" ]; then
-    package_name="${package_name}-rocm"
-  else
-    # Set to the default variant
-    fbgemm_variant="cuda"
-  fi
-  echo "[BUILD] Determined Python package name to use: ${package_name}"
 
-  # Extract the Python tag
-  # shellcheck disable=SC2207,SC2086
-  python_version=($(conda run --no-capture-output ${env_prefix} python --version))
-  # shellcheck disable=SC2206
-  python_version_arr=(${python_version[1]//./ })
-  python_tag="py${python_version_arr[0]}${python_version_arr[1]}"
-  echo "[BUILD] Extracted Python tag: ${python_tag}"
+  # Extract and set the package name given the FBGEMM_GPU variant
+  __build_fbgemm_gpu_set_package_name
+
+  # Extract and set the Python tag
+  __build_fbgemm_gpu_set_python_tag
+
+  # Extract and set the platform name
+  __build_fbgemm_gpu_set_python_plat_name
+
+  # Set multicore run option for setup.py if the number of cores on the machine
+  # permit for this
+  __build_fbgemm_gpu_set_run_multicore
 
   echo "[BUILD] Running pre-build cleanups ..."
   print_exec rm -rf dist
@@ -332,33 +397,14 @@ build_fbgemm_gpu_package () {
   echo "################################################################################"
   echo ""
 
-  # manylinux2014 is specified, bc manylinux1 does not support aarch64
-  # See https://github.com/pypa/manylinux
-  local plat_name="manylinux2014_${MACHINE_NAME}"
-
-  echo "[BUILD] Checking arch_list = ${arch_list}"
-  echo "[BUILD] Checking build_args:"
-  echo "${build_args[@]}"
-
-  # shellcheck disable=SC2155
-  local core=$(lscpu | grep "Core(s)" | awk '{print $NF}') && echo "core = ${core}" || echo "core not found"
-  # shellcheck disable=SC2155
-  local sockets=$(lscpu | grep "Socket(s)" | awk '{print $NF}') && echo "sockets = ${sockets}" || echo "sockets not found"
-  local re='^[0-9]+$'
-  local run_multicore=""
-  if [[ $core =~ $re && $sockets =~ $re ]] ; then
-    local n_core=$((core * sockets))
-    local run_multicore=" -j ${n_core}"
-  fi
-
   # Distribute Python extensions as wheels on Linux
   echo "[BUILD] Building FBGEMM-GPU wheel (VARIANT=${fbgemm_variant}) ..."
   # shellcheck disable=SC2086
   print_exec conda run --no-capture-output ${env_prefix} \
     python setup.py "${run_multicore}" bdist_wheel \
       --package_name="${package_name}" \
       --python-tag="${python_tag}" \
-      --plat-name="${plat_name}" \
+      --plat-name="${python_plat_name}" \
       --verbose \
       "${build_args[@]}"
 
@@ -410,7 +456,9 @@ build_fbgemm_gpu_install () {
   echo "[BUILD] Building + installing FBGEMM-GPU (VARIANT=${fbgemm_variant}) ..."
   # shellcheck disable=SC2086
   print_exec conda run --no-capture-output ${env_prefix} \
-    python setup.py install "${build_args[@]}"
+    python setup.py "${run_multicore}" install \
+      --verbose \
+      "${build_args[@]}"
 
   # Run checks on the built libraries
   (run_fbgemm_gpu_postbuild_checks "${fbgemm_variant}") || return 1
@@ -460,7 +508,9 @@ build_fbgemm_gpu_develop () {
   echo "[BUILD] Building (develop) FBGEMM-GPU (VARIANT=${fbgemm_variant}) ..."
   # shellcheck disable=SC2086
   print_exec conda run --no-capture-output ${env_prefix} \
-    python setup.py build develop "${build_args[@]}"
+    python setup.py "${run_multicore}" build develop \
+      --verbose \
+      "${build_args[@]}"
 
   # Run checks on the built libraries
   (run_fbgemm_gpu_postbuild_checks "${fbgemm_variant}") || return 1

diff --git a/fbgemm_gpu/docs/src/fbgemm_gpu-development/BuildInstructions.rst b/fbgemm_gpu/docs/src/fbgemm_gpu-development/BuildInstructions.rst
@@ -380,6 +380,33 @@ build cache:
 
   python setup.py clean
 
+Set Wheel Build Variables
+~~~~~~~~~~~~~~~~~~~~~~~~~
+
+When building out the Python wheel, the package name, Python version tag, and
+Python platform name must first be properly set:
+
+.. code:: sh
+
+  # Set the package name depending on the build variant
+  export package_name=fbgemm_gpu_{cpu, cuda, rocm}
+
+  # Set the Python version tag.  It should follow the convention `py<major><minor>`,
+  # e.g. Python 3.12 -> py312
+  export python_tag=py312
+
+  # Determine the processor architecture
+  export ARCH=$(uname -m)
+
+  # Set the Python platform name for the Linux case
+  export python_plat_name="manylinux2014_${ARCH}"
+  # For the macOS (x86_64) case
+  export python_plat_name="macosx_10_9_${ARCH}"
+  # For the macOS (arm64) case
+  export python_plat_name="macosx_11_0_${ARCH}"
+  # For the Windows case
+  export python_plat_name="win_${ARCH}"
+
 .. _fbgemm-gpu.build.process.cpu:
 
 CPU-Only Build
@@ -391,20 +418,16 @@ For CPU-only builds, the ``--cpu_only`` flag needs to be specified.
 
   # !! Run in fbgemm_gpu/ directory inside the Conda environment !!
 
-  export ARCH=$(uname -m)
-
-  python_tag=py310
-  package_name=fbgemm_gpu_cpu
-
   # Build the wheel artifact only
   python setup.py bdist_wheel \
-      --package_name="${package_name}" \
       --package_variant=cpu \
+      --package_name="${package_name}" \
       --python-tag="${python_tag}" \
-      --plat-name="manylinux1_${ARCH}"
+      --plat-name="${python_plat_name}"
 
   # Build and install the library into the Conda environment
-  python setup.py install --package_variant=cpu
+  python setup.py install \
+      --package_variant=cpu
 
 .. _fbgemm-gpu.build.process.cuda:
 
@@ -419,9 +442,6 @@ CUDA device, however, is not required for building the package.
 
   # !! Run in fbgemm_gpu/ directory inside the Conda environment !!
 
-  # Determine the processor architecture
-  export ARCH=$(uname -m)
-
   # [OPTIONAL] Specify the CUDA installation paths
   # This may be required if CMake is unable to find nvcc
   export CUDACXX=/path/to/nvcc
@@ -437,10 +457,6 @@ CUDA device, however, is not required for building the package.
   # Specify NVML path
   export NVML_LIB_PATH=/path/to/libnvidia-ml.so
 
-  # Update to reflect the version of Python in the Conda environment
-  python_tag=py310
-  package_name=fbgemm_gpu
-
   # Build for SM70/80 (V100/A100 GPU); update as needed
   # If not specified, only the CUDA architecture supported by current system will be targeted
   # If not specified and no CUDA device is present either, all CUDA architectures will be targeted
@@ -452,10 +468,10 @@ CUDA device, however, is not required for building the package.
 
   # Build the wheel artifact only
   python setup.py bdist_wheel \
-      --package_name="${package_name}" \
       --package_variant=cuda \
+      --package_name="${package_name}" \
       --python-tag="${python_tag}" \
-      --plat-name="manylinux1_${ARCH}" \
+      --plat-name="${python_plat_name}" \
       --nvml_lib_path=${NVML_LIB_PATH} \
       -DTORCH_CUDA_ARCH_LIST="${cuda_arch_list}"
 
@@ -478,22 +494,18 @@ the package.
 
   # !! Run in fbgemm_gpu/ directory inside the Conda environment !!
 
-  export ARCH=$(uname -m)
   export ROCM_PATH=/path/to/rocm
 
   # Build for the target architecture of the ROCm device installed on the machine (e.g. 'gfx906;gfx908;gfx90a')
   # See https://wiki.gentoo.org/wiki/ROCm for list
   export PYTORCH_ROCM_ARCH=$(${ROCM_PATH}/bin/rocminfo | grep -o -m 1 'gfx.*')
 
-  python_tag=py310
-  package_name=fbgemm_gpu_rocm
-
   # Build the wheel artifact only
   python setup.py bdist_wheel \
-      --package_name="${package_name}" \
       --package_variant=rocm \
+      --package_name="${package_name}" \
       --python-tag="${python_tag}" \
-      --plat-name="manylinux1_${ARCH}" \
+      --plat-name="${python_plat_name}" \
       -DHIP_ROOT_DIR="${ROCM_PATH}" \
       -DCMAKE_C_FLAGS="-DTORCH_USE_HIP_DSA" \
       -DCMAKE_CXX_FLAGS="-DTORCH_USE_HIP_DSA"