[CI/CD] change to PYPI token and exclude cupy from dependency (#475)

- lift cupy as a standalone dependency
- update pypi token

.github/workflows/release_alpa.yml

@@ -7,7 +7,7 @@ on:
 
 env:
   TWINE_USERNAME: "__token__"
-  TWINE_PASSWORD: ${{ secrets.TEST_PYPI_TOKEN }}
+  TWINE_PASSWORD: ${{ secrets.PYPI_TOKEN }}
 
 jobs:
 
@@ -60,7 +60,7 @@ jobs:
         run: |
           echo "Publish to PyPI"
           ls -ltr dist/
-          twine upload --repository testpypi dist/*
+          python -m twine upload --verbose dist/*
 
   update-version:
     runs-on: [ubuntu-latest]

docker/scripts/test_alpa_docker_entrypoint.sh

@@ -31,6 +31,9 @@ git checkout -qf FETCH_HEAD
 pip install /alpa-dist/jaxlib-alpa-ci/jaxlib-0.3.5+cuda111.cudnn805-cp38-none-manylinux2010_x86_64.whl
 pip install jax==0.3.5
 
+# install cupy
+pip install cupy-cuda111
+python -m cupyx.tools.install_library --library nccl --cuda 11.1
 pip install -e .[dev]
 ray start --head
 cd tests

docs/install.rst

@@ -4,19 +4,54 @@ Install Alpa
 Requirements
 ------------
 - CUDA >= 11.1
-- CuDNN >= 8.0.5
+- cuDNN >= 8.0.5
 - python >= 3.7
 
+Prerequisites
+-------------
+
+Regardless of installing from wheels or from source, we need to install a few prerequisite packages:
+
+1. CUDA toolkit:
+
+  Follow the official guides to install `CUDA <https://developer.nvidia.com/cuda-toolkit>`_ and `cuDNN <https://developer.nvidia.com/cudnn>`_.
+
+2. Install the ILP solver used by Alpa:
+
+    If you have sudo permission, use
+
+      .. code:: bash
+
+        sudo apt install coinor-cbc
+
+    Otherwise, please try to install via `binary <https://projects.coin-or.org/Cbc#DownloadandInstall>`_ or `conda <https://anaconda.org/conda-forge/coincbc>`_.
+
+3. Install cupy:
+
+  .. code:: bash
+
+    # Use your own CUDA version. Here cuda-cuda114 means cuda 11.4
+    pip3 install cupy-cuda114
+
+  Then, check whether your system already has NCCL installed.
+
+  .. code:: bash
+
+    python3 -c "from cupy.cuda import nccl"
+
+  If it prints nothing, then NCCL has already been installed.
+  Otherwise, follow the printed instructions to install NCCL.
+
 .. _install-from-wheels:
 
 Install from Python Wheels
 --------------------------
-Alpa provides wheels for the following CUDA (CuDNN) and Python versions:
+Alpa provides wheels for the following CUDA (cuDNN) and Python versions:
 
-- CUDA (CuDNN): 11.1 (8.0.5), 11.2 (8.1.0), 11.3 (8.2.0)
+- CUDA (cuDNN): 11.1 (8.0.5), 11.2 (8.1.0), 11.3 (8.2.0)
 - Python: 3.7, 3.8, 3.9
 
-If you need to use other CUDA, CuDNN, or Python versions, please follow the next section to :ref:`install from source<install-from-source>`.
+If you need to use other CUDA, cuDNN, or Python versions, please follow the next section to :ref:`install from source<install-from-source>`.
 
 1. To install from wheels, first install Alpa:
 
@@ -26,36 +61,26 @@ If you need to use other CUDA, CuDNN, or Python versions, please follow the next
     pip3 install alpa
 
 2. Then install the Alpa-modified Jaxlib from our `self-hosted PyPI server <http://169.229.48.123:8080/simple/>`_,
-   and make sure that the jaxlib version corresponds to the version of the existing CUDA and CuDNN installation you want to use.
-   You can specify a particular CUDA and CuDNN version for jaxlib explicitly:
+   and make sure that the jaxlib version corresponds to the version of the existing CUDA and cuDNN installation you want to use.
+   You can specify a particular CUDA and cuDNN version for jaxlib explicitly via:
 
   .. code:: bash
 
     pip3 install --trusted-host 169.229.48.123 -i http://169.229.48.123:8080/simple/ jaxlib==0.3.5+cuda{cuda_version}.cudnn{cudnn_version}
 
-  For example, to install the wheel built compatible with CUDA >= 11.1 and CuDNN >= 8.0.5, use the following command:
+  For example, to install the wheel compatible with CUDA >= 11.1 and cuDNN >= 8.0.5, use the following command:
 
   .. code:: bash
 
     pip3 install --trusted-host 169.229.48.123 -i http://169.229.48.123:8080/simple/ jaxlib==0.3.5+cuda111.cudnn805
 
-  You can see all available wheel versions at our `PyPI index <http://169.229.48.123:8080/simple/jaxlib/>`_.
+  You can see all available wheel versions we provided at our `PyPI index <http://169.229.48.123:8080/simple/jaxlib/>`_.
 
 .. note::
 
   As of now, Alpa modified the original jaxlib at the version ``jaxlib==0.3.5``. Alpa regularly rebases the official jaxlib repository to catch up with the upstream.
   If you need features from newer versions of jaxlib, please open an issue at the `Alpa GitHub Issue Page <https://github.com/alpa-projects/alpa/issues>`_.
 
-3. Finally, install the ILP solver:
-
-    If you have sudo permission, use
-
-      .. code:: bash
-
-        sudo apt install coinor-cbc
-
-    Otherwise, please try to install via `binary <https://projects.coin-or.org/Cbc#DownloadandInstall>`_ or `conda <https://anaconda.org/conda-forge/coincbc>`_.
-
 
 .. _install-from-source:
 
@@ -67,37 +92,7 @@ Install from Source
   
     git clone --recursive git@github.com:alpa-projects/alpa.git
 
-2. Install dependencies
-
-  - CUDA Toolkit:
-      Follow the official guides to install `cuda <https://developer.nvidia.com/cuda-toolkit>`_ and `cudnn <https://developer.nvidia.com/cudnn>`_.
-  - Python packages:
-
-      .. code:: bash
-    
-        pip3 install cmake tqdm pybind11 numba numpy pulp ray tensorstore flax==0.4.1 jax==0.3.5
-        pip3 install cupy-cuda114  # use your own CUDA version. Here cuda-cuda114 means cuda 11.4.
-
-  - NCCL:
-      First, check whether your system already has NCCL installed.
-
-      .. code:: bash
-
-        python3 -c "from cupy.cuda import nccl"
-
-      If it prints nothing, then nccl is already installed.
-      Otherwise, follow the printed instructions to install nccl.
-
-  - ILP Solver:
-      If you have sudo permission, use
-
-      .. code:: bash
-    
-        sudo apt install coinor-cbc
-
-      Otherwise, please try to install via `binary <https://projects.coin-or.org/Cbc#DownloadandInstall>`_ or `conda <https://anaconda.org/conda-forge/coincbc>`_.
-
-3. Build and install jaxlib
+2. Build and install jaxlib
 
   .. code:: bash
   
@@ -107,7 +102,7 @@ Install from Source
 
     pip3 install -e .
 
-4. Install Alpa
+3. Install Alpa
 
   .. code:: bash
   
@@ -152,7 +147,7 @@ Common issues of running Alpa on Slurm include:
 
 You might also find the discussion under `Issue #452 <https://github.com/alpa-projects/alpa/issues/452>`__ helpful.
 
-Jaxlib, Jax, FLAX Version Problems
+Jaxlib, Jax, Flax Version Problems
 ##################################
 Alpa is compatible with the following Jaxlib, Jax, and Flax versions:
 - Jax==0.3.5

setup.py

@@ -79,7 +79,6 @@ def get_cuda_version_str(no_dot=False):
     "ray[default]",
     "jax==0.3.5",
     "flax==0.4.1",
-    f"cupy-cuda{get_cuda_version_str(no_dot=True)}",
     "pulp",
     "tensorstore",
     "numpy<1.22",
@@ -89,7 +88,6 @@ def get_cuda_version_str(no_dot=False):
 dev_require_list = [
     "prospector",
     "yapf",
-    "coverage",
     "cmake",
     "pybind11"
 ]