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

Summary:

Update the install methods now that fbgemm has onboarded project nova.

Also this publishes cuda 12.1 as well.

Also updating gpu ci, even though it is currently not used.

Differential Revision: D49852074

.github/scripts/filter.py

@@ -22,8 +22,8 @@ def main():
     new_matrix_entries = []
 
     for entry in full_matrix["include"]:
-        if entry["gpu_arch_version"] != "12.1":
-            new_matrix_entries.append(entry)
+        # if entry["gpu_arch_version"] != "12.1":
+        new_matrix_entries.append(entry)
 
     new_matrix = {"include": new_matrix_entries}
     print(json.dumps(new_matrix))

.github/scripts/install_fbgemm.sh

@@ -0,0 +1,18 @@
+#!/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.
+
+echo "CU_VERSION"
+echo "$CU_VERSION"
+
+echo "CHANNEL"
+echo "$CHANNEL"
+
+if [ "$CHANNEL" = "nightly" ]; then
+    ${CONDA_RUN} pip install fbgemm-gpu --index-url https://download.pytorch.org/whl/nightly/"$CU_VERSION"
+elif [ "$CHANNEL" = "test" ]; then
+    ${CONDA_RUN} pip install fbgemm-gpu --index-url https://download.pytorch.org/whl/test/"$CU_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.yml

@@ -166,6 +166,11 @@ jobs:
       shell: bash
       run: |
         conda install -n build_binary -y pytorch pytorch-cuda=11.8 -c pytorch-nightly -c nvidia
+    - name: Install fbgemm
+      shell: bash
+      run: |
+        conda run -n build_binary \
+          pip install fbgemm-gpu --index-url https://download.pytorch.org/whl/nightly/cu118
     # download wheel from GHA
     - name: Download wheel
       uses: actions/download-artifact@v2

.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=(