Update install methods after fbgemm onboarded onto nova (pytorch#1420)

Summary:

wip

Differential Revision: D49852074

.github/scripts/install_fbgemm.sh

@@ -0,0 +1,20 @@
+#!/bin/bash
+# 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.
+
+if [ "$MATRIX_GPU_ARCH_TYPE" = "cpu" ]; then
+    CUDA_VERSION="cpu"
+elif [ "$MATRIX_GPU_ARCH_VERSION" = "11.8" ]; then
+    CUDA_VERSION="cu118"
+elif [ "$MATRIX_GPU_ARCH_VERSION" = "12.1" ]; then
+    CUDA_VERSION="cu121"
+fi
+
+echo "$CUDA_VERSION"
+
+if [ "$MATRIX_CHANNEL" = "nightly" ]; then
+    ${CONDA_RUN} pip install fbgemm-gpu --index-url https://download.pytorch.org/whl/nightly/"$CUDA_VERSION"
+fi

.github/workflows/build-wheels-linux.yml

@@ -54,7 +54,7 @@ jobs:
       test-infra-ref: main
       build-matrix: ${{ needs.filter-matrix.outputs.matrix }}
       pre-script: ""
-      post-script: ""
+      post-script: ./.github/scripts/install_fbgemm.sh
       package-name: torchrec
       smoke-test-script: ""
       trigger-event: ${{ github.event_name }}

.github/workflows/docs.yml

@@ -52,12 +52,13 @@ jobs:
     - name: Install PyTorch
       shell: bash
       run: |
-        conda install -n build_binary -y -c pytorch-nightly "pytorch-nightly"::pytorch[build="*cpu*"]
+        conda install -n build_binary --yes pytorch cpuonly -c pytorch-nightly
+    - name: Install fbgemm
+      run: |
+        conda run -n build_binary pip install fbgemm-gpu --index-url https://download.pytorch.org/whl/nightly/cpu
     - name: Install TorchRec
       run: |
-        conda run -n build_binary python -m pip install torchrec_nightly
-        conda run -n build_binary python -m pip uninstall fbgemm-gpu-nightly -y
-        conda run -n build_binary python -m pip install fbgemm-gpu-nightly-cpu
+        pip install torchrec --index-url https://download.pytorch.org/whl/nightly/cpu
     - name: Test fbgemm_gpu and torchrec installation
       shell: bash
       run: |

.github/workflows/pyre.yml

@@ -19,8 +19,8 @@ jobs:
         uses: actions/checkout@v2
       - name: Install dependencies
         run: >
-          conda install --yes -c pytorch-nightly "pytorch-nightly"::pytorch[build="*cpu*"] &&
-          sed -i 's/fbgemm-gpu-nightly/fbgemm-gpu-nightly-cpu/g' requirements.txt &&
+          conda install --yes pytorch cpuonly -c pytorch-nightly &&
+          pip install fbgemm-gpu --index-url https://download.pytorch.org/whl/nightly/cpu &&
           pip install -r requirements.txt &&
           pip install pyre-check-nightly==$(cat .pyre_configuration | grep version | awk '{print $2}' | sed 's/\"//g')
       - name: Pyre check

.github/workflows/unittest_ci_cpu.yml

@@ -47,14 +47,16 @@ jobs:
         conda run -n build_binary python --version
         conda install -n build_binary \
           --yes \
-          -c pytorch-nightly \
-          "pytorch-nightly"::pytorch[build="*cpu*"]
+          pytorch cpuonly -c pytorch-nightly
         conda run -n build_binary \
           python -c "import torch.distributed"
-        sed -i 's/fbgemm-gpu-nightly/fbgemm-gpu-nightly-cpu/g' requirements.txt
+        conda run -n build_binary \
+          pip install fbgemm-gpu --index-url https://download.pytorch.org/whl/nightly/cpu
+        conda run -n build_binary \
+          python -c "import fbgemm_gpu"
+        echo "fbgemm_gpu succeeded"
         conda run -n build_binary \
           pip install -r requirements.txt
-        export CU_VERSION="cpu"
         conda run -n build_binary \
           python setup.py bdist_wheel \
           --python-tag=${{ matrix.python-tag }}
@@ -64,9 +66,6 @@ jobs:
         conda run -n build_binary \
           python -c "import numpy"
         echo "numpy succeeded"
-        conda run -n build_binary \
-          python -c "import fbgemm_gpu"
-        echo "fbgemm_gpu succeeded"
         conda install -n build_binary -y pytest
         conda run -n build_binary \
           python -m pytest torchrec -v -s -W ignore::pytest.PytestCollectionWarning --continue-on-collection-errors -k 'not test_sharding_gloo_cw and not test_load_state_dict_prefix and not test_load_state_dict_cw_multiple_shards' --ignore-glob=**/test_utils/

install-requirements.txt

@@ -1,3 +1,3 @@
-fbgemm-gpu-nightly
+fbgemm-gpu
 torchmetrics==1.0.3
 tqdm

requirements.txt

@@ -1,6 +1,6 @@
 black
 cmake
-fbgemm-gpu-nightly
+fbgemm-gpu
 hypothesis==6.70.1
 iopath
 numpy

setup.py

@@ -9,7 +9,6 @@
 import os
 import subprocess
 import sys
-from datetime import date
 from pathlib import Path
 from typing import List
 
@@ -45,17 +44,6 @@ def _export_version(version, sha):
         fileobj.write("git_version = {}\n".format(repr(sha)))
 
 
-def get_channel():
-    # Channel typically takes on the following values:
-    # - NIGHTLY: for nightly published binaries
-    # - TEST: for binaries build from release candidate branches
-    return os.getenv("CHANNEL")
-
-
-def get_cu_version():
-    return os.getenv("CU_VERSION", "cpu")
-
-
 def parse_args(argv: List[str]) -> argparse.Namespace:
     parser = argparse.ArgumentParser(description="torchrec setup")
     return parser.parse_known_args(argv)
@@ -64,9 +52,6 @@ def parse_args(argv: List[str]) -> argparse.Namespace:
 def main(argv: List[str]) -> None:
     args, unknown = parse_args(argv)
 
-    # Set up package version
-    channel = get_channel()
-
     with open(
         os.path.join(os.path.dirname(__file__), "README.MD"), encoding="utf8"
     ) as f:
@@ -81,21 +66,7 @@ def main(argv: List[str]) -> None:
     version, sha = _get_version()
     _export_version(version, sha)
 
-    if channel != "nightly":
-        if "fbgemm-gpu-nightly" in install_requires:
-            install_requires.remove("fbgemm-gpu-nightly")
-        install_requires.append("fbgemm-gpu")
-
-    cu_version = get_cu_version()
-    if cu_version == "cpu":
-        if "fbgemm-gpu-nightly" in install_requires:
-            install_requires.remove("fbgemm-gpu-nightly")
-            install_requires.append("fbgemm-gpu-nightly-cpu")
-        if "fbgemm-gpu" in install_requires:
-            install_requires.remove("fbgemm-gpu")
-            install_requires.append("fbgemm-gpu-cpu")
-
-    print(f"-- torchrec building version: {version} CU Version: {cu_version}")
+    print(f"-- torchrec building version: {version}")
 
     packages = find_packages(
         exclude=(