Fix cupy install version detection (deepspeedai#3276)

mrwyattii · web-flow · commit bcccee4d8555 · 2023-04-18T17:13:35.000Z
* updated cupy install

* do non-isolated pip install

* Update action.yml
diff --git a/.github/workflows/nv-torch19-p40.yml b/.github/workflows/nv-torch19-p40.yml
@@ -45,7 +45,7 @@ jobs:
 
       - name: Install deepspeed
         run: |
-          pip install .[dev,1bit,autotuning]
+          pip install .[dev,1bit,autotuning] --no-build-isolation
           ds_report
 
       - name: Python environment
diff --git a/.github/workflows/nv-torch19-v100.yml b/.github/workflows/nv-torch19-v100.yml
@@ -45,7 +45,7 @@ jobs:
 
       - name: Install deepspeed
         run: |
-          pip install .[dev,1bit,autotuning]
+          pip install .[dev,1bit,autotuning] --no-build-isolation
           ds_report
 
       - name: Python environment
diff --git a/.github/workflows/setup-venv/action.yml b/.github/workflows/setup-venv/action.yml
@@ -15,6 +15,7 @@ runs:
         python -m venv unit-test-venv
         source ./unit-test-venv/bin/activate
         python -m pip install --upgrade pip
+        pip install wheel # required after pip>=23.1
         echo PATH=$PATH >> $GITHUB_ENV # Make it so venv is inherited for other steps
       shell: bash
     - id: print-env
diff --git a/setup.py b/setup.py
@@ -78,7 +78,12 @@ def fetch_requirements(path):
         if rocm_major <= 4:
             cupy = f"cupy-rocm-{rocm_major}-{rocm_minor}"
     else:
-        cupy = f"cupy-cuda{''.join(map(str,installed_cuda_version()))}"
+        cuda_major_ver, cuda_minor_ver = installed_cuda_version()
+        if (cuda_major_ver < 11) or ((cuda_major_ver == 11) and (cuda_minor_ver < 3)):
+            cupy = f"cupy-cuda{cuda_major_ver}{cuda_minor_ver}"
+        else:
+            cupy = f"cupy-cuda{cuda_major_ver}x"
+
     if cupy:
         extras_require['1bit'].append(cupy)
         extras_require['1bit_mpi'].append(cupy)
