From a06cd4a02126c4e4f8075c531aa723ba00ffa96a Mon Sep 17 00:00:00 2001
From: Hao Zhang <zhisbug@gmail.com>
Date: Thu, 30 Jun 2022 01:27:05 -0700
Subject: [PATCH 01/36] copy several scripts from tlc release

---
 build_jaxlib/release/README.md              |   3 +
 build_jaxlib/release/generate_pypi_index.py | 171 ++++++++++++++++++++
 build_jaxlib/release/wheel_upload.py        |  46 ++++++
 3 files changed, 220 insertions(+)
 create mode 100644 build_jaxlib/release/README.md
 create mode 100644 build_jaxlib/release/generate_pypi_index.py
 create mode 100644 build_jaxlib/release/wheel_upload.py

diff --git a/build_jaxlib/release/README.md b/build_jaxlib/release/README.md
new file mode 100644
index 000000000..05ce95882
--- /dev/null
+++ b/build_jaxlib/release/README.md
@@ -0,0 +1,3 @@
+# How to Release JaxLib in Our Own Index
+
+Since Alpa 
diff --git a/build_jaxlib/release/generate_pypi_index.py b/build_jaxlib/release/generate_pypi_index.py
new file mode 100644
index 000000000..61e340d3b
--- /dev/null
+++ b/build_jaxlib/release/generate_pypi_index.py
@@ -0,0 +1,171 @@
+"""Update the wheels page, prune old nightly builds if necessary."""
+import github3
+import os
+import logging
+import argparse
+import subprocess
+from datetime import datetime
+
+import requests
+
+
+def py_str(cstr):
+    return cstr.decode("utf-8")
+
+
+def extract_group_key_order(name):
+    """Extract group key and order from name.
+
+    Parameters
+    ----------
+    name : str
+        name of the file.
+
+    Returns
+    -------
+    group_key : tuple
+        The group the build should belong to
+
+    order : tuple
+        The order used to sort the builds.
+        The higher the latest
+    """
+    assert name.endswith(".whl")
+    name = name[:-4]
+    arr = name.split("-")
+
+    pkg_name = arr[0]
+    group_key = [arr[0]] + arr[2:]
+
+    ver = arr[1]
+    plus_pos = ver.find("+")
+    if plus_pos != -1:
+        ver = ver[:plus_pos]
+
+    dev_pos = ver.find(".dev")
+    if dev_pos != -1:
+        # all nightly share the same group
+        group_key.append("nightly")
+        # dev number as the order.
+        pub_ver = [int(x) for x in ver[:dev_pos].split(".")]
+        order = pub_ver + [int(ver[dev_pos + 4 :])]
+    else:
+        # stable version has its own group
+        group_key.append(ver)
+        order = [0]
+
+    return tuple(group_key), tuple(order)
+
+
+def group_wheels(wheels):
+    group_map = {}
+    for asset in wheels:
+        gkey, order = extract_group_key_order(asset.name)
+        if gkey not in group_map:
+            group_map[gkey] = []
+        group_map[gkey].append((order, asset))
+    return group_map
+
+
+def url_is_valid(url):
+    """Check if a given URL is valid, i.e. it returns 200 OK when requested."""
+    r = requests.get(url)
+
+    if r.status_code != 200:
+        print("Warning: HTTP code %s for url %s" % (r.status_code, url))
+
+    return r.status_code == 200
+
+
+def run_prune(args, group_map):
+    keep_list = []
+    remove_list = []
+    for key, assets in group_map.items():
+        print(f"Group {key}:")
+        for idx, item in enumerate(reversed(sorted(assets, key=lambda x: x[0]))):
+            order, asset = item
+            if idx < args.keep_top:
+                print("keep  %s" % asset.browser_download_url)
+                keep_list.append(asset)
+            else:
+                print("remove  %s" % asset.browser_download_url)
+                remove_list.append(asset)
+        print()
+    return keep_list, remove_list
+
+
+def list_wheels(repo):
+    gh = github3.login(token=os.environ["GITHUB_TOKEN"])
+    repo = gh.repository(*repo.split("/"))
+    wheels = []
+
+    for release in repo.releases():
+        tag = release.tag_name
+        for asset in release.assets():
+            if asset.name.endswith(".whl") and url_is_valid(asset.browser_download_url):
+                wheels.append(asset)
+    return wheels
+
+
+def update_wheel_page(keep_list, site_repo, dry_run=False):
+    """Update the wheel page"""
+    new_html = ""
+    for asset in keep_list:
+        new_html += '<a href="%s">%s</a><br>\n' % (
+            asset.browser_download_url,
+            asset.name,
+        )
+
+    def run_cmd(cmd):
+        proc = subprocess.Popen(
+            cmd, cwd=site_repo, stdout=subprocess.PIPE, stderr=subprocess.STDOUT
+        )
+        (out, _) = proc.communicate()
+        if proc.returncode != 0:
+            msg = "git error: %s" % cmd
+            msg += py_str(out)
+            raise RuntimeError(msg)
+
+    run_cmd(["git", "fetch"])
+    run_cmd(["git", "checkout", "-B", "main", "origin/main"])
+    wheel_html_path = os.path.join(site_repo, "wheels.html")
+    if open(wheel_html_path, "r").read() != new_html:
+        print(f"Wheel page changed, update {wheel_html_path}..")
+        if not dry_run:
+            open(wheel_html_path, "w").write(new_html)
+            run_cmd(["git", "commit", "-am", "wheel update at %s" % datetime.now()])
+            run_cmd(["git", "push", "origin", "main"])
+
+
+def delete_assets(remove_list, dry_run):
+    for asset in remove_list:
+        if not dry_run:
+            asset.delete()
+    if remove_list:
+        print("Finish deleting %d removed assets" % len(remove_list))
+
+
+def main():
+    logging.basicConfig(level=logging.WARNING)
+    parser = argparse.ArgumentParser(
+        description="Prune nightly build and synchronize the wheel page."
+    )
+    parser.add_argument("--keep-top", type=int, default=1)
+    parser.add_argument("--dry-run", action="store_true")
+    parser.add_argument("--site-path", type=str, default="tlc-pack.github.io")
+    parser.add_argument("--repo", type=str, default="tlc-pack/tlcpack")
+
+    if "GITHUB_TOKEN" not in os.environ:
+        raise RuntimeError("need GITHUB_TOKEN")
+    args = parser.parse_args()
+    wheels = list_wheels(args.repo)
+    group_map = group_wheels(wheels)
+    keep_list, remove_list = run_prune(args, group_map)
+    # NOTE: important to update html first before deletion
+    # so that the wheel page always points to correct asset
+    update_wheel_page(keep_list, args.site_path, args.dry_run)
+    delete_assets(remove_list, args.dry_run)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/build_jaxlib/release/wheel_upload.py b/build_jaxlib/release/wheel_upload.py
new file mode 100644
index 000000000..d99760f56
--- /dev/null
+++ b/build_jaxlib/release/wheel_upload.py
@@ -0,0 +1,46 @@
+"""Update the wheels page, prune old nightly builds if necessary (source from tlcpack)."""
+import github3
+import os
+import logging
+import argparse
+import subprocess
+
+
+def upload(args, path):
+    gh = github3.login(token=os.environ["GITHUB_TOKEN"])
+    repo = gh.repository(*args.repo.split("/"))
+    release = repo.release_from_tag(args.tag)
+    name = os.path.basename(path)
+    content_bytes = open(path, "rb").read()
+
+    for asset in release.assets():
+        if asset.name == name:
+            if not args.dry_run:
+                asset.delete()
+                print(f"Remove duplicated file {name}")
+    print(f"Start to upload {path} to {args.repo}, this can take a while...")
+    if not args.dry_run:
+        release.upload_asset("application/octet-stream", name, content_bytes)
+    print(f"Finish uploading {path}")
+
+
+def main():
+    logging.basicConfig(level=logging.WARNING)
+    parser = argparse.ArgumentParser(description="Upload wheel as an asset of a tag.")
+    parser.add_argument("--tag", type=str)
+    parser.add_argument("--repo", type=str, default="alpa-projects/alpa")
+    parser.add_argument("--dry-run", action="store_true")
+    parser.add_argument("path", type=str)
+
+    if "GITHUB_TOKEN" not in os.environ:
+        raise RuntimeError("need GITHUB_TOKEN")
+    args = parser.parse_args()
+    if os.path.isdir(args.path):
+        for name in os.listdir(args.path):
+            if name.endswith(".whl"):
+                upload(args, os.path.join(args.path, name))
+    else:
+        upload(args, args.path)
+
+if __name__ == "__main__":
+    main()
\ No newline at end of file

From 5b4c528f7f317892645877b712fa7e3e668d9728 Mon Sep 17 00:00:00 2001
From: Hao Zhang <sjtu.haozhang@gmail.com>
Date: Thu, 30 Jun 2022 13:59:56 -0700
Subject: [PATCH 02/36] add ways to download singleton ckpts for smaller models

---
 examples/opt_serving/README.rst | 17 ++++++++++++++++-
 1 file changed, 16 insertions(+), 1 deletion(-)

diff --git a/examples/opt_serving/README.rst b/examples/opt_serving/README.rst
index 34a0e9e90..cbf43c81c 100644
--- a/examples/opt_serving/README.rst
+++ b/examples/opt_serving/README.rst
@@ -94,7 +94,8 @@ There are two ways to obtain Alpa-compatible OPT weights: converting the weights
 
 Convert weights into Alpa formats by yourself
 ---------------------------------------------
-We provide detailed instructions below on how to convert the original OPT-175B weights into Alpa-compatible formats. You can follow the same procedures to get Alpa-compatible weights for other model sizes.
+We provide detailed instructions below on how to convert the original OPT-175B weights into Alpa-compatible formats.
+For processing other sizes of OPT (125M - 66B), you can skip Step 1 and start from :ref:`the latter part of Step 2<download-singleton>`.
 
   .. note::
 
@@ -122,6 +123,20 @@ We provide detailed instructions below on how to convert the original OPT-175B w
 
     The above script will save the model weights as a single consolidated checkpoint at ``PATH_TO_SAVE_CHECKPOINT``, hence will require at least 350GB disk space available.
 
+.. _download-singleton:
+
+  .. note::
+    If you use Alpa to target smaller versions of OPT (125M, 350M, 1.3B, 2.7B, 6.7B, 13B, 30B), you can skip the above procedures
+    and download the consolidated singleton checkpoint using the links below, then proceed to the next step.
+
+      * `OPT-125M <https://huggingface.co/patrickvonplaten/opt_metaseq_125m/blob/main/model/restored.pt>`_
+      * `OPT-350M <https://dl.fbaipublicfiles.com/opt/v1_20220502/350m/reshard.pt>`_
+      * `OPT-1.3B <https://huggingface.co/patrickvonplaten/opt_metaseq_1300m/blob/main/model/restored.pt>`_
+      * `OPT-2.7B <https://huggingface.co/patrickvonplaten/opt_metaseq_2700m/blob/main/model/restored.pt>`_
+      * `OPT-6.7B <https://huggingface.co/patrickvonplaten/opt_metaseq_6700m/blob/main/model/restored.pt>`_
+      * `OPT-13B <https://huggingface.co/patrickvonplaten/opt_metaseq_13000m/blob/main/model/restored.pt>`_
+      * `OPT-30B <https://huggingface.co/patrickvonplaten/opt_metaseq_30000m/blob/main/model/restored.pt>`_
+
 
 3. Convert the single checkpoint into Alpa-compatible formats
     Alpa ingests weights simply from numpy formats. Use the script `step_3_convert_to_numpy_weights.py <https://github.com/alpa-projects/alpa/tree/main/examples/opt_serving/scripts/step_3_convert_to_numpy_weights.py>`_ to convert the

From 2f9b873144e1ab57740281a7ff76edef966dd435 Mon Sep 17 00:00:00 2001
From: Hao Zhang <sjtu.haozhang@gmail.com>
Date: Fri, 1 Jul 2022 14:05:42 -0700
Subject: [PATCH 03/36] defat

---
 .github/workflows/release_jaxlib.yml | 42 ++++++++++++++++++++++------
 1 file changed, 34 insertions(+), 8 deletions(-)

diff --git a/.github/workflows/release_jaxlib.yml b/.github/workflows/release_jaxlib.yml
index 26ff3c177..a147b1349 100644
--- a/.github/workflows/release_jaxlib.yml
+++ b/.github/workflows/release_jaxlib.yml
@@ -4,10 +4,12 @@ on:
   release:
     types: [created]
   workflow_dispatch:
+    inputs:
+      tensorflow:
+        description: 'TensorFlow-alpa branch to build'
+        required: true
+        default: 'master'
 
-env:
-  TWINE_USERNAME: "__token__"
-  TWINE_PASSWORD: ${{ secrets.PYPI_TOKEN }}
 
 jobs:
 
@@ -41,11 +43,18 @@ jobs:
       - name: Compile Jaxlib
         run: |
           mkdir -p /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.}
+          if [[ ${{ github.event.action }} == "release" ]]; then
+            docker run --gpus all --tmpfs /build:exec \
+              --rm -v /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.}:/dist \
+              build-jaxlib-image-cuda${CUDA_VERSION} ${PYTHON_VERSION} \
+              cuda ${CUDA_VERSION}
+          else
+            docker run --gpus all --tmpfs /build:exec \
+              --rm -v /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.}:/dist \
+              build-jaxlib-image-cuda${CUDA_VERSION} ${PYTHON_VERSION} \
+              cuda ${CUDA_VERSION} ${{ github.event.inputs.tensorflow }}
+          fi
 
-          docker run --gpus all --tmpfs /build:exec \
-          --rm -v /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.}:/dist \
-          build-jaxlib-image-cuda${CUDA_VERSION} ${PYTHON_VERSION} \
-          cuda ${CUDA_VERSION}
         env:
           CUDA_VERSION: ${{ matrix.cuda }}
           PYTHON_VERSION: ${{ matrix.python }}
@@ -57,10 +66,27 @@ jobs:
       matrix:
         cuda: ["11.1", "11.2", "11.3"]
     steps:
+
+    - uses: actions/checkout@v3
+
+    - name: Set up Python ${{ matrix.python }}
+      uses: actions/setup-python@v2
+      with:
+        python-version: ${{ matrix.python }}
+
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        python -m pip install github3.py
+        tag=$(head -n 1 VERSION)
+        
+
     - name: Publish CUDA${{ matrix.cuda }}
       run: |
-        echo "Move to self-hosted pypi"
+        echo "Publish jaxlib"
         ls /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.}
+        
         mv /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.}/*.whl /data/alpa-pypi/packages/
       env:
         CUDA_VERSION: ${{ matrix.cuda }}
+        GITHUB_TOKEN: ${{ secrets.PAT_TOKEN }}

From 8468474f1c18d945a03744846b2e097afb4c49cf Mon Sep 17 00:00:00 2001
From: Hao Zhang <sjtu.haozhang@gmail.com>
Date: Fri, 1 Jul 2022 23:09:18 -0700
Subject: [PATCH 04/36] update [skip ci]

---
 .github/workflows/build_jaxlib.yml               |  2 +-
 .github/workflows/release_jaxlib.yml             | 16 ++++++++++------
 docker/scripts/build_jaxlib_docker_entrypoint.sh | 15 +++++++++++++--
 3 files changed, 24 insertions(+), 9 deletions(-)

diff --git a/.github/workflows/build_jaxlib.yml b/.github/workflows/build_jaxlib.yml
index 33c19445f..1a63ed183 100644
--- a/.github/workflows/build_jaxlib.yml
+++ b/.github/workflows/build_jaxlib.yml
@@ -45,7 +45,7 @@ jobs:
           mkdir -p dist
           docker run --gpus all --tmpfs /build:exec \
           --rm -v $(pwd)/dist:/dist build-jaxlib-image \
-          3.8 cuda 11.1 ${TF_BRANCH##*/}
+          3.8 cuda 11.1 main ${TF_BRANCH##*/}
 
       # change this to publishing to pypi
       - name: Publish to local
diff --git a/.github/workflows/release_jaxlib.yml b/.github/workflows/release_jaxlib.yml
index a147b1349..506c7a4c1 100644
--- a/.github/workflows/release_jaxlib.yml
+++ b/.github/workflows/release_jaxlib.yml
@@ -47,17 +47,19 @@ jobs:
             docker run --gpus all --tmpfs /build:exec \
               --rm -v /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.}:/dist \
               build-jaxlib-image-cuda${CUDA_VERSION} ${PYTHON_VERSION} \
-              cuda ${CUDA_VERSION}
+              cuda ${CUDA_VERSION} ${ALPA_BRANCH}
           else
             docker run --gpus all --tmpfs /build:exec \
               --rm -v /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.}:/dist \
               build-jaxlib-image-cuda${CUDA_VERSION} ${PYTHON_VERSION} \
-              cuda ${CUDA_VERSION} ${{ github.event.inputs.tensorflow }}
+              cuda ${CUDA_VERSION} ${ALPA_BRANCH} ${TF_BRANCH}
           fi
 
         env:
           CUDA_VERSION: ${{ matrix.cuda }}
           PYTHON_VERSION: ${{ matrix.python }}
+          ALPA_BRANCH: ${{ github.ref }}
+          TF_BRANCH: ${{ github.event.inputs.tensorflow }}
 
   publish:
     runs-on: [self-hosted]
@@ -78,15 +80,17 @@ jobs:
       run: |
         python -m pip install --upgrade pip
         python -m pip install github3.py
-        tag=$(head -n 1 VERSION)
-        
+
+    - name: Get latest tag
+      id: latesttag
+      uses: "WyriHaximus/github-action-get-previous-tag@v1"
 
     - name: Publish CUDA${{ matrix.cuda }}
       run: |
         echo "Publish jaxlib"
         ls /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.}
-        
-        mv /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.}/*.whl /data/alpa-pypi/packages/
+        python build_jaxlib/release/wheel_upload.py --tag ${TAG} /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.}
       env:
         CUDA_VERSION: ${{ matrix.cuda }}
         GITHUB_TOKEN: ${{ secrets.PAT_TOKEN }}
+        TAG: ${{ steps.latesttag.outputs.tag }}
diff --git a/docker/scripts/build_jaxlib_docker_entrypoint.sh b/docker/scripts/build_jaxlib_docker_entrypoint.sh
index 317169851..bcf7a38b9 100644
--- a/docker/scripts/build_jaxlib_docker_entrypoint.sh
+++ b/docker/scripts/build_jaxlib_docker_entrypoint.sh
@@ -13,7 +13,7 @@ export CUDA_PATH=/usr/local/cuda
 export LD_LIBRARY_PATH=$CUDA_PATH/lib64:$LD_LIBRARY_PATH
 
 usage() {
-  echo "usage: ${0##*/} [3.7|3.8|3.9] [cuda|nocuda] [11.1|11.2|11.3] [tensorflow-alpa branch name]"
+  echo "usage: ${0##*/} [3.7|3.8|3.9] [cuda|nocuda] [11.1|11.2|11.3] [alpa branch name] [tensorflow-alpa branch name]"
   exit 1
 }
 
@@ -28,9 +28,20 @@ echo "Python version $PY_VERSION"
 # switch tensorflow-alpa branch if necessary
 git clone --recursive https://github.com/alpa-projects/alpa.git
 
+# switch alpa branch
 if [[ $# -eq 4 ]]
 then
-  TF_BRANCH="$4"
+  ALPA_BRANCH="$4"
+  echo "Switch to alpa branch ALPA_BRANCH"
+  cd /build/alpa
+  git fetch origin +${ALPA_BRANCH}
+  git checkout -qf FETCH_HEAD
+  git submodule update --recursive
+
+# switch tensorflow-alpa branch, this will overwrite the above
+if [[ $# -eq 5 ]]
+then
+  TF_BRANCH="$5"
   echo "Switch to tensorflow-alpa branch $TF_BRANCH"
   cd /build/alpa/third_party/tensorflow-alpa
   git fetch origin +${TF_BRANCH}

From 22f9168263c045153bf54de4414673abfeb66d83 Mon Sep 17 00:00:00 2001
From: Hao Zhang <sjtu.haozhang@gmail.com>
Date: Fri, 1 Jul 2022 23:57:48 -0700
Subject: [PATCH 05/36] new pypi infra [skip ci]

---
 .github/workflows/release_jaxlib_test.yml   | 97 +++++++++++++++++++++
 build_jaxlib/release/generate_pypi_index.py | 33 +++----
 2 files changed, 114 insertions(+), 16 deletions(-)
 create mode 100644 .github/workflows/release_jaxlib_test.yml

diff --git a/.github/workflows/release_jaxlib_test.yml b/.github/workflows/release_jaxlib_test.yml
new file mode 100644
index 000000000..211ddcd0b
--- /dev/null
+++ b/.github/workflows/release_jaxlib_test.yml
@@ -0,0 +1,97 @@
+name: Release Jaxlib exp
+
+on:
+  release:
+    types: [created]
+  workflow_dispatch:
+    inputs:
+      tensorflow:
+        description: 'TensorFlow-alpa branch to build'
+        required: true
+        default: 'master'
+
+
+jobs:
+
+  clean-up:
+    runs-on: [self-hosted]
+
+    steps:
+      - name: clean up images
+        run: |
+          docker image prune -f
+
+#  build-jaxlib:
+#    runs-on: [self-hosted]
+#    needs: [clean-up]
+#    strategy:
+#      matrix:
+#        cuda: ["11.1", "11.2", "11.3"]
+#        python: ["3.7", "3.8", "3.9"]
+#
+#    steps:
+#      - uses: actions/checkout@v3
+#
+#      - name: build image
+#        run: |
+#          docker build -t build-jaxlib-image-cuda${CUDA_VERSION} \
+#            -f docker/build_jaxlib.Dockerfile docker/ \
+#            --build-arg JAX_CUDA_VERSION=${CUDA_VERSION}
+#        env:
+#          CUDA_VERSION: ${{ matrix.cuda }}
+#
+#      - name: Compile Jaxlib
+#        run: |
+#          mkdir -p /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.}
+#          if [[ ${{ github.event.action }} == "release" ]]; then
+#            docker run --gpus all --tmpfs /build:exec \
+#              --rm -v /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.}:/dist \
+#              build-jaxlib-image-cuda${CUDA_VERSION} ${PYTHON_VERSION} \
+#              cuda ${CUDA_VERSION} ${ALPA_BRANCH}
+#          else
+#            docker run --gpus all --tmpfs /build:exec \
+#              --rm -v /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.}:/dist \
+#              build-jaxlib-image-cuda${CUDA_VERSION} ${PYTHON_VERSION} \
+#              cuda ${CUDA_VERSION} ${ALPA_BRANCH} ${TF_BRANCH}
+#          fi
+#
+#        env:
+#          CUDA_VERSION: ${{ matrix.cuda }}
+#          PYTHON_VERSION: ${{ matrix.python }}
+#          ALPA_BRANCH: ${{ github.ref }}
+#          TF_BRANCH: ${{ github.event.inputs.tensorflow }}
+
+  publish:
+    runs-on: [self-hosted]
+    needs: [build-jaxlib]
+    strategy:
+      matrix:
+        cuda: ["11.1", "11.2", "11.3"]
+    steps:
+
+    - uses: actions/checkout@v3
+
+    - name: Set up Python ${{ matrix.python }}
+      uses: actions/setup-python@v2
+      with:
+        python-version: ${{ matrix.python }}
+
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        python -m pip install github3.py
+
+    - name: Get latest tag
+      id: latesttag
+      uses: "WyriHaximus/github-action-get-previous-tag@v1"
+
+    - name: Publish CUDA${{ matrix.cuda }}
+      run: |
+        echo "Publish jaxlib"
+        ls /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.}
+        echo "Upload wheels to tag ${TAG}"
+        python build_jaxlib/release/wheel_upload.py --tag ${TAG} /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.}
+      env:
+        CUDA_VERSION: ${{ matrix.cuda }}
+        GITHUB_TOKEN: ${{ secrets.PAT_TOKEN }}
+        TAG: ${{ steps.latesttag.outputs.tag }}
diff --git a/build_jaxlib/release/generate_pypi_index.py b/build_jaxlib/release/generate_pypi_index.py
index 61e340d3b..458924c47 100644
--- a/build_jaxlib/release/generate_pypi_index.py
+++ b/build_jaxlib/release/generate_pypi_index.py
@@ -94,16 +94,17 @@ def run_prune(args, group_map):
     return keep_list, remove_list
 
 
-def list_wheels(repo):
+def list_wheels(repo, tag):
     gh = github3.login(token=os.environ["GITHUB_TOKEN"])
     repo = gh.repository(*repo.split("/"))
     wheels = []
-
-    for release in repo.releases():
-        tag = release.tag_name
-        for asset in release.assets():
-            if asset.name.endswith(".whl") and url_is_valid(asset.browser_download_url):
-                wheels.append(asset)
+    all_tags = [release.tag_name for release in repo.releases]
+    if tag not in all_tags:
+        raise RuntimeError("The tag provided does not exist.")
+    release = repo.release_from_tag(tag)
+    for asset in release.assets():
+        if asset.name.endswith(".whl") and url_is_valid(asset.browser_download_url):
+            wheels.append(asset)
     return wheels
 
 
@@ -148,23 +149,23 @@ def delete_assets(remove_list, dry_run):
 def main():
     logging.basicConfig(level=logging.WARNING)
     parser = argparse.ArgumentParser(
-        description="Prune nightly build and synchronize the wheel page."
+        description="Generate a wheel page given a release tag, assuming the wheels have been uploaded."
     )
-    parser.add_argument("--keep-top", type=int, default=1)
     parser.add_argument("--dry-run", action="store_true")
-    parser.add_argument("--site-path", type=str, default="tlc-pack.github.io")
-    parser.add_argument("--repo", type=str, default="tlc-pack/tlcpack")
+    parser.add_argument("--site-path", type=str, default="alpa-projects.github.io")
+    parser.add_argument("--repo", type=str, default="alpa-projects/alpa")
+    parser.add_argument("--tag", type=str)
 
     if "GITHUB_TOKEN" not in os.environ:
         raise RuntimeError("need GITHUB_TOKEN")
     args = parser.parse_args()
-    wheels = list_wheels(args.repo)
-    group_map = group_wheels(wheels)
-    keep_list, remove_list = run_prune(args, group_map)
+    wheels = list_wheels(args.repo, args.tag)
+    # group_map = group_wheels(wheels)
+    # keep_list, remove_list = run_prune(args, group_map)
     # NOTE: important to update html first before deletion
     # so that the wheel page always points to correct asset
-    update_wheel_page(keep_list, args.site_path, args.dry_run)
-    delete_assets(remove_list, args.dry_run)
+    update_wheel_page(wheels, args.site_path, args.dry_run)
+    # delete_assets(remove_list, args.dry_run)
 
 
 if __name__ == "__main__":

From ef445d605c99a0908179b3539f5c3922257eadaa Mon Sep 17 00:00:00 2001
From: Hao Zhang <sjtu.haozhang@gmail.com>
Date: Sat, 2 Jul 2022 00:59:57 -0700
Subject: [PATCH 06/36] add timoout for large file; add readme so core
 contributors can do it [skip ci]

---
 build_jaxlib/release/README.md       | 13 +++++++++++--
 build_jaxlib/release/wheel_upload.py | 10 ++++++----
 2 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/build_jaxlib/release/README.md b/build_jaxlib/release/README.md
index 05ce95882..48411a30c 100644
--- a/build_jaxlib/release/README.md
+++ b/build_jaxlib/release/README.md
@@ -1,3 +1,12 @@
-# How to Release JaxLib in Our Own Index
+# How to Release JaxLib and generate a PyPI Index
 
-Since Alpa 
+## Upload jaxlib wheels as assets of a release tag
+```shell
+GITHUB_TOKEN="admin_token" python wheel_upload.py --tag [TAG] --path [PATH_TO_WHEELS]
+```
+
+## Generate a html index page and commit it to Alpa doc page
+```shell
+python generate_pypi_index.py --tag [TAG]
+```
+All wheel assets under `[TAG]` will be included in a html index page appeared in the doc repo.
diff --git a/build_jaxlib/release/wheel_upload.py b/build_jaxlib/release/wheel_upload.py
index d99760f56..b254fafe2 100644
--- a/build_jaxlib/release/wheel_upload.py
+++ b/build_jaxlib/release/wheel_upload.py
@@ -1,13 +1,15 @@
 """Update the wheels page, prune old nightly builds if necessary (source from tlcpack)."""
 import github3
+import github3.session as session
 import os
 import logging
 import argparse
-import subprocess
 
 
 def upload(args, path):
-    gh = github3.login(token=os.environ["GITHUB_TOKEN"])
+    # gh = github3.login(token=os.environ["GITHUB_TOKEN"])
+    gh = github3.GitHub(token=os.environ["GITHUB_TOKEN"],
+                        session=session.GitHubSession(default_connect_timeout=100, default_read_timeout=100))
     repo = gh.repository(*args.repo.split("/"))
     release = repo.release_from_tag(args.tag)
     name = os.path.basename(path)
@@ -30,7 +32,7 @@ def main():
     parser.add_argument("--tag", type=str)
     parser.add_argument("--repo", type=str, default="alpa-projects/alpa")
     parser.add_argument("--dry-run", action="store_true")
-    parser.add_argument("path", type=str)
+    parser.add_argument("--path", type=str)
 
     if "GITHUB_TOKEN" not in os.environ:
         raise RuntimeError("need GITHUB_TOKEN")
@@ -43,4 +45,4 @@ def main():
         upload(args, args.path)
 
 if __name__ == "__main__":
-    main()
\ No newline at end of file
+    main()

From 98fdb3a2a4d39ebb3ccc266b993f2f31332f4439 Mon Sep 17 00:00:00 2001
From: Hao Zhang <sjtu.haozhang@gmail.com>
Date: Sat, 2 Jul 2022 01:58:24 -0700
Subject: [PATCH 07/36] manual workflow works [skip ci]

---
 build_jaxlib/release/README.md              |  4 ++--
 build_jaxlib/release/generate_pypi_index.py | 13 +++++++++----
 2 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/build_jaxlib/release/README.md b/build_jaxlib/release/README.md
index 48411a30c..c609c5896 100644
--- a/build_jaxlib/release/README.md
+++ b/build_jaxlib/release/README.md
@@ -2,11 +2,11 @@
 
 ## Upload jaxlib wheels as assets of a release tag
 ```shell
-GITHUB_TOKEN="admin_token" python wheel_upload.py --tag [TAG] --path [PATH_TO_WHEELS]
+GITHUB_TOKEN="[admin_token]" python wheel_upload.py --tag [TAG] --path [PATH_TO_WHEELS]
 ```
 
 ## Generate a html index page and commit it to Alpa doc page
 ```shell
-python generate_pypi_index.py --tag [TAG]
+GITHUB_TOKEN="[admin_token]" python generate_pypi_index.py --tag [TAG]
 ```
 All wheel assets under `[TAG]` will be included in a html index page appeared in the doc repo.
diff --git a/build_jaxlib/release/generate_pypi_index.py b/build_jaxlib/release/generate_pypi_index.py
index 458924c47..eb4dd7a8f 100644
--- a/build_jaxlib/release/generate_pypi_index.py
+++ b/build_jaxlib/release/generate_pypi_index.py
@@ -98,11 +98,12 @@ def list_wheels(repo, tag):
     gh = github3.login(token=os.environ["GITHUB_TOKEN"])
     repo = gh.repository(*repo.split("/"))
     wheels = []
-    all_tags = [release.tag_name for release in repo.releases]
+    all_tags = [release.tag_name for release in repo.releases()]
     if tag not in all_tags:
         raise RuntimeError("The tag provided does not exist.")
     release = repo.release_from_tag(tag)
     for asset in release.assets():
+        print(f"Validating {asset.name} with url: {asset.browser_download_url}")
         if asset.name.endswith(".whl") and url_is_valid(asset.browser_download_url):
             wheels.append(asset)
     return wheels
@@ -128,14 +129,18 @@ def run_cmd(cmd):
             raise RuntimeError(msg)
 
     run_cmd(["git", "fetch"])
-    run_cmd(["git", "checkout", "-B", "main", "origin/main"])
+    run_cmd(["git", "checkout", "-B", "master", "origin/master"])
     wheel_html_path = os.path.join(site_repo, "wheels.html")
-    if open(wheel_html_path, "r").read() != new_html:
+    if not os.path.exists(wheel_html_path) or open(wheel_html_path, "r").read() != new_html:
         print(f"Wheel page changed, update {wheel_html_path}..")
         if not dry_run:
             open(wheel_html_path, "w").write(new_html)
+            print(123)
+            run_cmd(["git", "add", "wheels.html"])
+            print(456)
             run_cmd(["git", "commit", "-am", "wheel update at %s" % datetime.now()])
-            run_cmd(["git", "push", "origin", "main"])
+            print(789)
+            run_cmd(["git", "push", "origin", "master"])
 
 
 def delete_assets(remove_list, dry_run):

From 36180274288e084ca2ea36a5a4fb00e2f50d8688 Mon Sep 17 00:00:00 2001
From: Hao Zhang <sjtu.haozhang@gmail.com>
Date: Sat, 2 Jul 2022 13:31:18 -0700
Subject: [PATCH 08/36] simplify pypi scripts

---
 .github/workflows/docs.yml                  |  1 +
 build_jaxlib/release/generate_pypi_index.py | 76 ---------------------
 2 files changed, 1 insertion(+), 76 deletions(-)

diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml
index a192f198a..bdb5ce3af 100644
--- a/.github/workflows/docs.yml
+++ b/.github/workflows/docs.yml
@@ -35,3 +35,4 @@ jobs:
           external_repository: alpa-projects/alpa-projects.github.io
           publish_branch: master
           publish_dir: /data/alpa-dist/docs
+          keep_files: true
diff --git a/build_jaxlib/release/generate_pypi_index.py b/build_jaxlib/release/generate_pypi_index.py
index eb4dd7a8f..8111da121 100644
--- a/build_jaxlib/release/generate_pypi_index.py
+++ b/build_jaxlib/release/generate_pypi_index.py
@@ -13,60 +13,6 @@ def py_str(cstr):
     return cstr.decode("utf-8")
 
 
-def extract_group_key_order(name):
-    """Extract group key and order from name.
-
-    Parameters
-    ----------
-    name : str
-        name of the file.
-
-    Returns
-    -------
-    group_key : tuple
-        The group the build should belong to
-
-    order : tuple
-        The order used to sort the builds.
-        The higher the latest
-    """
-    assert name.endswith(".whl")
-    name = name[:-4]
-    arr = name.split("-")
-
-    pkg_name = arr[0]
-    group_key = [arr[0]] + arr[2:]
-
-    ver = arr[1]
-    plus_pos = ver.find("+")
-    if plus_pos != -1:
-        ver = ver[:plus_pos]
-
-    dev_pos = ver.find(".dev")
-    if dev_pos != -1:
-        # all nightly share the same group
-        group_key.append("nightly")
-        # dev number as the order.
-        pub_ver = [int(x) for x in ver[:dev_pos].split(".")]
-        order = pub_ver + [int(ver[dev_pos + 4 :])]
-    else:
-        # stable version has its own group
-        group_key.append(ver)
-        order = [0]
-
-    return tuple(group_key), tuple(order)
-
-
-def group_wheels(wheels):
-    group_map = {}
-    for asset in wheels:
-        gkey, order = extract_group_key_order(asset.name)
-        if gkey not in group_map:
-            group_map[gkey] = []
-        group_map[gkey].append((order, asset))
-    return group_map
-
-
 def url_is_valid(url):
     """Check if a given URL is valid, i.e. it returns 200 OK when requested."""
     r = requests.get(url)
@@ -77,23 +23,6 @@ def url_is_valid(url):
     return r.status_code == 200
 
 
-def run_prune(args, group_map):
-    keep_list = []
-    remove_list = []
-    for key, assets in group_map.items():
-        print(f"Group {key}:")
-        for idx, item in enumerate(reversed(sorted(assets, key=lambda x: x[0]))):
-            order, asset = item
-            if idx < args.keep_top:
-                print("keep  %s" % asset.browser_download_url)
-                keep_list.append(asset)
-            else:
-                print("remove  %s" % asset.browser_download_url)
-                remove_list.append(asset)
-        print()
-    return keep_list, remove_list
-
-
 def list_wheels(repo, tag):
     gh = github3.login(token=os.environ["GITHUB_TOKEN"])
     repo = gh.repository(*repo.split("/"))
@@ -165,12 +94,7 @@ def main():
         raise RuntimeError("need GITHUB_TOKEN")
     args = parser.parse_args()
     wheels = list_wheels(args.repo, args.tag)
-    # group_map = group_wheels(wheels)
-    # keep_list, remove_list = run_prune(args, group_map)
-    # NOTE: important to update html first before deletion
-    # so that the wheel page always points to correct asset
     update_wheel_page(wheels, args.site_path, args.dry_run)
-    # delete_assets(remove_list, args.dry_run)
 
 
 if __name__ == "__main__":

From 0d0324cd59171cdbaa1831990cf7c25bda7a2cd5 Mon Sep 17 00:00:00 2001
From: Hao Zhang <sjtu.haozhang@gmail.com>
Date: Sat, 2 Jul 2022 14:18:27 -0700
Subject: [PATCH 09/36] experimental workflow [skip ci]

---
 .github/workflows/release_jaxlib_test.yml   | 19 +++++++++++++++++--
 build_jaxlib/release/generate_pypi_index.py | 16 ++++++----------
 2 files changed, 23 insertions(+), 12 deletions(-)

diff --git a/.github/workflows/release_jaxlib_test.yml b/.github/workflows/release_jaxlib_test.yml
index 211ddcd0b..59fdd330b 100644
--- a/.github/workflows/release_jaxlib_test.yml
+++ b/.github/workflows/release_jaxlib_test.yml
@@ -79,13 +79,13 @@ jobs:
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
-        python -m pip install github3.py
+        python -m pip install github3.py requests
 
     - name: Get latest tag
       id: latesttag
       uses: "WyriHaximus/github-action-get-previous-tag@v1"
 
-    - name: Publish CUDA${{ matrix.cuda }}
+    - name: Upload wheels for CUDA${{ matrix.cuda }}
       run: |
         echo "Publish jaxlib"
         ls /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.}
@@ -95,3 +95,18 @@ jobs:
         CUDA_VERSION: ${{ matrix.cuda }}
         GITHUB_TOKEN: ${{ secrets.PAT_TOKEN }}
         TAG: ${{ steps.latesttag.outputs.tag }}
+
+    - uses: actions/checkout@v3
+      name: "Checkout doc repo"
+      with:
+        repository: "alpa-projects/alpa-projects.github.io"
+
+    - uses
+      name: "Generate"
+      env:
+        GITHUB_TOKEN: ${{ secrets.PAT_TOKEN }}
+        TAG: ${{ steps.latesttag.outputs.tag }}
+      run:
+        ls -ltr
+        cd ..
+        python build_jaxlib/release/generate_pypi_index.py --tag ${TAG}
diff --git a/build_jaxlib/release/generate_pypi_index.py b/build_jaxlib/release/generate_pypi_index.py
index 8111da121..08a2b7d8d 100644
--- a/build_jaxlib/release/generate_pypi_index.py
+++ b/build_jaxlib/release/generate_pypi_index.py
@@ -60,16 +60,12 @@ def run_cmd(cmd):
     run_cmd(["git", "fetch"])
     run_cmd(["git", "checkout", "-B", "master", "origin/master"])
     wheel_html_path = os.path.join(site_repo, "wheels.html")
-    if not os.path.exists(wheel_html_path) or open(wheel_html_path, "r").read() != new_html:
-        print(f"Wheel page changed, update {wheel_html_path}..")
-        if not dry_run:
-            open(wheel_html_path, "w").write(new_html)
-            print(123)
-            run_cmd(["git", "add", "wheels.html"])
-            print(456)
-            run_cmd(["git", "commit", "-am", "wheel update at %s" % datetime.now()])
-            print(789)
-            run_cmd(["git", "push", "origin", "master"])
+    print(f"Wheel page changed, update {wheel_html_path}..")
+    if not dry_run:
+        open(wheel_html_path, "w").write(new_html)
+        run_cmd(["git", "add", "wheels.html"])
+        run_cmd(["git", "commit", "-am", "wheel update at %s" % datetime.now()])
+        run_cmd(["git", "push", "origin", "master"])
 
 
 def delete_assets(remove_list, dry_run):

From 93bdce8a78646a114f5c48c63e344d783ea1870e Mon Sep 17 00:00:00 2001
From: Hao Zhang <sjtu.haozhang@gmail.com>
Date: Sat, 2 Jul 2022 14:25:35 -0700
Subject: [PATCH 10/36] turn it on [skip ci]

---
 .github/workflows/release_jaxlib_test.yml | 55 +++++++++++++----------
 1 file changed, 31 insertions(+), 24 deletions(-)

diff --git a/.github/workflows/release_jaxlib_test.yml b/.github/workflows/release_jaxlib_test.yml
index 59fdd330b..a9980a100 100644
--- a/.github/workflows/release_jaxlib_test.yml
+++ b/.github/workflows/release_jaxlib_test.yml
@@ -3,6 +3,8 @@ name: Release Jaxlib exp
 on:
   release:
     types: [created]
+  push:
+    branches: [ new-pypi-infra ]
   workflow_dispatch:
     inputs:
       tensorflow:
@@ -21,17 +23,17 @@ jobs:
         run: |
           docker image prune -f
 
-#  build-jaxlib:
-#    runs-on: [self-hosted]
-#    needs: [clean-up]
-#    strategy:
-#      matrix:
-#        cuda: ["11.1", "11.2", "11.3"]
-#        python: ["3.7", "3.8", "3.9"]
-#
-#    steps:
-#      - uses: actions/checkout@v3
-#
+  build-jaxlib:
+    runs-on: [self-hosted]
+    needs: [clean-up]
+    strategy:
+      matrix:
+        cuda: ["11.1", "11.2", "11.3"]
+        python: ["3.7", "3.8", "3.9"]
+
+    steps:
+      - uses: actions/checkout@v3
+
 #      - name: build image
 #        run: |
 #          docker build -t build-jaxlib-image-cuda${CUDA_VERSION} \
@@ -39,27 +41,32 @@ jobs:
 #            --build-arg JAX_CUDA_VERSION=${CUDA_VERSION}
 #        env:
 #          CUDA_VERSION: ${{ matrix.cuda }}
-#
-#      - name: Compile Jaxlib
-#        run: |
-#          mkdir -p /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.}
-#          if [[ ${{ github.event.action }} == "release" ]]; then
+
+      - name: Compile Jaxlib
+        run: |
+          mkdir -p /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.}
+          if [[ ${{ github.event.action }} == "release" ]]; then
+            echo ${CUDA_VERSION}
+            echo ${ALPA_BRANCH}
 #            docker run --gpus all --tmpfs /build:exec \
 #              --rm -v /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.}:/dist \
 #              build-jaxlib-image-cuda${CUDA_VERSION} ${PYTHON_VERSION} \
 #              cuda ${CUDA_VERSION} ${ALPA_BRANCH}
-#          else
+          else
 #            docker run --gpus all --tmpfs /build:exec \
 #              --rm -v /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.}:/dist \
 #              build-jaxlib-image-cuda${CUDA_VERSION} ${PYTHON_VERSION} \
 #              cuda ${CUDA_VERSION} ${ALPA_BRANCH} ${TF_BRANCH}
-#          fi
-#
-#        env:
-#          CUDA_VERSION: ${{ matrix.cuda }}
-#          PYTHON_VERSION: ${{ matrix.python }}
-#          ALPA_BRANCH: ${{ github.ref }}
-#          TF_BRANCH: ${{ github.event.inputs.tensorflow }}
+            echo ${CUDA_VERSION}
+            echo ${ALPA_BRANCH}
+            echo ${TF_BRANCH}
+          fi
+
+        env:
+          CUDA_VERSION: ${{ matrix.cuda }}
+          PYTHON_VERSION: ${{ matrix.python }}
+          ALPA_BRANCH: ${{ github.ref }}
+          TF_BRANCH: ${{ github.event.inputs.tensorflow }}
 
   publish:
     runs-on: [self-hosted]

From 3601ce29c9aecd6e2ffe5ca6d6be261845dffb5e Mon Sep 17 00:00:00 2001
From: Hao Zhang <sjtu.haozhang@gmail.com>
Date: Sat, 2 Jul 2022 14:32:45 -0700
Subject: [PATCH 11/36] update

---
 .github/workflows/release_jaxlib_test.yml | 18 +++++++++---------
 1 file changed, 9 insertions(+), 9 deletions(-)

diff --git a/.github/workflows/release_jaxlib_test.yml b/.github/workflows/release_jaxlib_test.yml
index a9980a100..d4d81190a 100644
--- a/.github/workflows/release_jaxlib_test.yml
+++ b/.github/workflows/release_jaxlib_test.yml
@@ -3,8 +3,8 @@ name: Release Jaxlib exp
 on:
   release:
     types: [created]
-  push:
-    branches: [ new-pypi-infra ]
+  pull_request:
+    branches: [main]
   workflow_dispatch:
     inputs:
       tensorflow:
@@ -34,13 +34,13 @@ jobs:
     steps:
       - uses: actions/checkout@v3
 
-#      - name: build image
-#        run: |
-#          docker build -t build-jaxlib-image-cuda${CUDA_VERSION} \
-#            -f docker/build_jaxlib.Dockerfile docker/ \
-#            --build-arg JAX_CUDA_VERSION=${CUDA_VERSION}
-#        env:
-#          CUDA_VERSION: ${{ matrix.cuda }}
+      - name: build image
+        run: |
+          docker build -t build-jaxlib-image-cuda${CUDA_VERSION} \
+            -f docker/build_jaxlib.Dockerfile docker/ \
+            --build-arg JAX_CUDA_VERSION=${CUDA_VERSION}
+        env:
+          CUDA_VERSION: ${{ matrix.cuda }}
 
       - name: Compile Jaxlib
         run: |

From d9ef2c10a50265cbf8caa293f6cd3da792986b97 Mon Sep 17 00:00:00 2001
From: Hao Zhang <sjtu.haozhang@gmail.com>
Date: Sat, 2 Jul 2022 14:35:32 -0700
Subject: [PATCH 12/36] update

---
 .github/workflows/release_jaxlib_test.yml | 9 ---------
 1 file changed, 9 deletions(-)

diff --git a/.github/workflows/release_jaxlib_test.yml b/.github/workflows/release_jaxlib_test.yml
index d4d81190a..bb34acabb 100644
--- a/.github/workflows/release_jaxlib_test.yml
+++ b/.github/workflows/release_jaxlib_test.yml
@@ -48,20 +48,11 @@ jobs:
           if [[ ${{ github.event.action }} == "release" ]]; then
             echo ${CUDA_VERSION}
             echo ${ALPA_BRANCH}
-#            docker run --gpus all --tmpfs /build:exec \
-#              --rm -v /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.}:/dist \
-#              build-jaxlib-image-cuda${CUDA_VERSION} ${PYTHON_VERSION} \
-#              cuda ${CUDA_VERSION} ${ALPA_BRANCH}
           else
-#            docker run --gpus all --tmpfs /build:exec \
-#              --rm -v /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.}:/dist \
-#              build-jaxlib-image-cuda${CUDA_VERSION} ${PYTHON_VERSION} \
-#              cuda ${CUDA_VERSION} ${ALPA_BRANCH} ${TF_BRANCH}
             echo ${CUDA_VERSION}
             echo ${ALPA_BRANCH}
             echo ${TF_BRANCH}
           fi
-
         env:
           CUDA_VERSION: ${{ matrix.cuda }}
           PYTHON_VERSION: ${{ matrix.python }}

From 474debc5009fc84d008b7dc1dc481d85068d5281 Mon Sep 17 00:00:00 2001
From: Hao Zhang <sjtu.haozhang@gmail.com>
Date: Sat, 2 Jul 2022 14:55:10 -0700
Subject: [PATCH 13/36] update

---
 .github/workflows/release_jaxlib_test.yml | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/.github/workflows/release_jaxlib_test.yml b/.github/workflows/release_jaxlib_test.yml
index bb34acabb..f3da25067 100644
--- a/.github/workflows/release_jaxlib_test.yml
+++ b/.github/workflows/release_jaxlib_test.yml
@@ -99,8 +99,7 @@ jobs:
       with:
         repository: "alpa-projects/alpa-projects.github.io"
 
-    - uses
-      name: "Generate"
+    - name: "Generate"
       env:
         GITHUB_TOKEN: ${{ secrets.PAT_TOKEN }}
         TAG: ${{ steps.latesttag.outputs.tag }}

From d1d88581450376263b6e362e78e6d9022b6013c3 Mon Sep 17 00:00:00 2001
From: Hao Zhang <sjtu.haozhang@gmail.com>
Date: Sat, 2 Jul 2022 14:55:43 -0700
Subject: [PATCH 14/36] update

---
 .github/workflows/release_jaxlib_test.yml | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/release_jaxlib_test.yml b/.github/workflows/release_jaxlib_test.yml
index f3da25067..b5bc0618d 100644
--- a/.github/workflows/release_jaxlib_test.yml
+++ b/.github/workflows/release_jaxlib_test.yml
@@ -88,7 +88,7 @@ jobs:
         echo "Publish jaxlib"
         ls /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.}
         echo "Upload wheels to tag ${TAG}"
-        python build_jaxlib/release/wheel_upload.py --tag ${TAG} /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.}
+        python build_jaxlib/release/wheel_upload.py --tag ${TAG} /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.} --dry-run
       env:
         CUDA_VERSION: ${{ matrix.cuda }}
         GITHUB_TOKEN: ${{ secrets.PAT_TOKEN }}
@@ -106,4 +106,4 @@ jobs:
       run:
         ls -ltr
         cd ..
-        python build_jaxlib/release/generate_pypi_index.py --tag ${TAG}
+        python build_jaxlib/release/generate_pypi_index.py --tag ${TAG} --dry-run

From ac9a78b04b2022fc344e8bd34023c46f354cb4dd Mon Sep 17 00:00:00 2001
From: Hao Zhang <sjtu.haozhang@gmail.com>
Date: Sat, 2 Jul 2022 15:08:53 -0700
Subject: [PATCH 15/36] update

---
 .github/workflows/release_jaxlib_test.yml | 89 +++++++++++++----------
 1 file changed, 49 insertions(+), 40 deletions(-)

diff --git a/.github/workflows/release_jaxlib_test.yml b/.github/workflows/release_jaxlib_test.yml
index b5bc0618d..b39a1cf20 100644
--- a/.github/workflows/release_jaxlib_test.yml
+++ b/.github/workflows/release_jaxlib_test.yml
@@ -46,9 +46,19 @@ jobs:
         run: |
           mkdir -p /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.}
           if [[ ${{ github.event.action }} == "release" ]]; then
+            echo ${{ github.event.action }}
+            echo ${CUDA_VERSION}
+            echo ${ALPA_BRANCH}
+          elif [[ ${{ github.event.action }} == "pull_request"]]; then
+            echo ${{ github.event.action }}
+            echo ${CUDA_VERSION}
+            echo ${ALPA_BRANCH}
+          elif [[ ${{ github.event.action }} == "push"]]; then
+            echo ${{ github.event.action }}
             echo ${CUDA_VERSION}
             echo ${ALPA_BRANCH}
           else
+            echo ${{ github.event.action }}
             echo ${CUDA_VERSION}
             echo ${ALPA_BRANCH}
             echo ${TF_BRANCH}
@@ -66,44 +76,43 @@ jobs:
       matrix:
         cuda: ["11.1", "11.2", "11.3"]
     steps:
+      - name: Set up Python 3.8
+        uses: actions/setup-python@v2
+        with:
+          python-version: 3.8
+
+      - uses: actions/checkout@v3
+
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          python -m pip install github3.py requests
 
-    - uses: actions/checkout@v3
-
-    - name: Set up Python ${{ matrix.python }}
-      uses: actions/setup-python@v2
-      with:
-        python-version: ${{ matrix.python }}
-
-    - name: Install dependencies
-      run: |
-        python -m pip install --upgrade pip
-        python -m pip install github3.py requests
-
-    - name: Get latest tag
-      id: latesttag
-      uses: "WyriHaximus/github-action-get-previous-tag@v1"
-
-    - name: Upload wheels for CUDA${{ matrix.cuda }}
-      run: |
-        echo "Publish jaxlib"
-        ls /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.}
-        echo "Upload wheels to tag ${TAG}"
-        python build_jaxlib/release/wheel_upload.py --tag ${TAG} /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.} --dry-run
-      env:
-        CUDA_VERSION: ${{ matrix.cuda }}
-        GITHUB_TOKEN: ${{ secrets.PAT_TOKEN }}
-        TAG: ${{ steps.latesttag.outputs.tag }}
-
-    - uses: actions/checkout@v3
-      name: "Checkout doc repo"
-      with:
-        repository: "alpa-projects/alpa-projects.github.io"
-
-    - name: "Generate"
-      env:
-        GITHUB_TOKEN: ${{ secrets.PAT_TOKEN }}
-        TAG: ${{ steps.latesttag.outputs.tag }}
-      run:
-        ls -ltr
-        cd ..
-        python build_jaxlib/release/generate_pypi_index.py --tag ${TAG} --dry-run
+      - name: Get latest tag
+        id: latesttag
+        uses: "WyriHaximus/github-action-get-previous-tag@v1"
+
+      - name: Upload wheels for CUDA${{ matrix.cuda }}
+        run: |
+          echo "Publish jaxlib"
+          ls /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.}
+          echo "Upload wheels to tag ${TAG}"
+          python build_jaxlib/release/wheel_upload.py --tag ${TAG} /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.} --dry-run
+        env:
+          CUDA_VERSION: ${{ matrix.cuda }}
+          GITHUB_TOKEN: ${{ secrets.PAT_TOKEN }}
+          TAG: ${{ steps.latesttag.outputs.tag }}
+
+      - uses: actions/checkout@v3
+        name: "Checkout doc repo"
+        with:
+          repository: "alpa-projects/alpa-projects.github.io"
+
+      - name: "Generate"
+        env:
+          GITHUB_TOKEN: ${{ secrets.PAT_TOKEN }}
+          TAG: ${{ steps.latesttag.outputs.tag }}
+        run:
+          ls -ltr
+          cd ..
+          python build_jaxlib/release/generate_pypi_index.py --tag ${TAG} --dry-run

From b6d500d8e92c7b078ad4261bf5574cf008f8a185 Mon Sep 17 00:00:00 2001
From: Hao Zhang <sjtu.haozhang@gmail.com>
Date: Sat, 2 Jul 2022 15:11:40 -0700
Subject: [PATCH 16/36] update

---
 .github/workflows/release_jaxlib_test.yml | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/.github/workflows/release_jaxlib_test.yml b/.github/workflows/release_jaxlib_test.yml
index b39a1cf20..f0d4e6f31 100644
--- a/.github/workflows/release_jaxlib_test.yml
+++ b/.github/workflows/release_jaxlib_test.yml
@@ -53,10 +53,6 @@ jobs:
             echo ${{ github.event.action }}
             echo ${CUDA_VERSION}
             echo ${ALPA_BRANCH}
-          elif [[ ${{ github.event.action }} == "push"]]; then
-            echo ${{ github.event.action }}
-            echo ${CUDA_VERSION}
-            echo ${ALPA_BRANCH}
           else
             echo ${{ github.event.action }}
             echo ${CUDA_VERSION}

From 72ba4b09753e550ada3fe78dfae1684e7633c040 Mon Sep 17 00:00:00 2001
From: Hao Zhang <sjtu.haozhang@gmail.com>
Date: Sat, 2 Jul 2022 15:13:15 -0700
Subject: [PATCH 17/36] update

---
 .github/workflows/release_jaxlib_test.yml | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/.github/workflows/release_jaxlib_test.yml b/.github/workflows/release_jaxlib_test.yml
index f0d4e6f31..70b2e6b3a 100644
--- a/.github/workflows/release_jaxlib_test.yml
+++ b/.github/workflows/release_jaxlib_test.yml
@@ -45,11 +45,7 @@ jobs:
       - name: Compile Jaxlib
         run: |
           mkdir -p /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.}
-          if [[ ${{ github.event.action }} == "release" ]]; then
-            echo ${{ github.event.action }}
-            echo ${CUDA_VERSION}
-            echo ${ALPA_BRANCH}
-          elif [[ ${{ github.event.action }} == "pull_request"]]; then
+          if [[ ${{ github.event.action }} == "pull_request" ]]; then
             echo ${{ github.event.action }}
             echo ${CUDA_VERSION}
             echo ${ALPA_BRANCH}

From da30acd1e1b60774429cf95a6fa8b94038a528c7 Mon Sep 17 00:00:00 2001
From: Hao Zhang <sjtu.haozhang@gmail.com>
Date: Sat, 2 Jul 2022 15:33:41 -0700
Subject: [PATCH 18/36] update

---
 .github/workflows/release_jaxlib_test.yml   | 14 +++++++++-----
 build_jaxlib/release/generate_pypi_index.py |  6 ++++--
 2 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/.github/workflows/release_jaxlib_test.yml b/.github/workflows/release_jaxlib_test.yml
index 70b2e6b3a..1dac0dd03 100644
--- a/.github/workflows/release_jaxlib_test.yml
+++ b/.github/workflows/release_jaxlib_test.yml
@@ -45,12 +45,16 @@ jobs:
       - name: Compile Jaxlib
         run: |
           mkdir -p /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.}
-          if [[ ${{ github.event.action }} == "pull_request" ]]; then
-            echo ${{ github.event.action }}
+          if [[ ${{ github.event_name }} == "pull_request" ]]; then
+            echo ${{ github.event_name }}
+            echo ${{ github.event }}
+            echo ${{ github.action }}
             echo ${CUDA_VERSION}
             echo ${ALPA_BRANCH}
           else
-            echo ${{ github.event.action }}
+            echo ${{ github.event_name }}
+            echo ${{ github.event }}
+            echo ${{ github.action }}
             echo ${CUDA_VERSION}
             echo ${ALPA_BRANCH}
             echo ${TF_BRANCH}
@@ -86,10 +90,10 @@ jobs:
 
       - name: Upload wheels for CUDA${{ matrix.cuda }}
         run: |
-          echo "Publish jaxlib"
+          echo "Publish jaxlib to tag ${TAG}"
           ls /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.}
           echo "Upload wheels to tag ${TAG}"
-          python build_jaxlib/release/wheel_upload.py --tag ${TAG} /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.} --dry-run
+          python build_jaxlib/release/wheel_upload.py --tag ${TAG} --path /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.} --dry-run
         env:
           CUDA_VERSION: ${{ matrix.cuda }}
           GITHUB_TOKEN: ${{ secrets.PAT_TOKEN }}
diff --git a/build_jaxlib/release/generate_pypi_index.py b/build_jaxlib/release/generate_pypi_index.py
index 08a2b7d8d..b4ac587e9 100644
--- a/build_jaxlib/release/generate_pypi_index.py
+++ b/build_jaxlib/release/generate_pypi_index.py
@@ -1,11 +1,12 @@
 """Update the wheels page, prune old nightly builds if necessary."""
-import github3
 import os
 import logging
 import argparse
 import subprocess
 from datetime import datetime
 
+import github3
+import github3.session as session
 import requests
 
 
@@ -24,7 +25,8 @@ def url_is_valid(url):
 
 
 def list_wheels(repo, tag):
-    gh = github3.login(token=os.environ["GITHUB_TOKEN"])
+    gh = github3.GitHub(token=os.environ["GITHUB_TOKEN"],
+                        session=session.GitHubSession(default_connect_timeout=100, default_read_timeout=100))
     repo = gh.repository(*repo.split("/"))
     wheels = []
     all_tags = [release.tag_name for release in repo.releases()]

From fac950000cf7db04d215657d324b5d24dbac7964 Mon Sep 17 00:00:00 2001
From: Hao Zhang <sjtu.haozhang@gmail.com>
Date: Sat, 2 Jul 2022 16:07:12 -0700
Subject: [PATCH 19/36] update [skip ci]

---
 .github/workflows/release_jaxlib_test.yml | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/.github/workflows/release_jaxlib_test.yml b/.github/workflows/release_jaxlib_test.yml
index 1dac0dd03..9841ea23c 100644
--- a/.github/workflows/release_jaxlib_test.yml
+++ b/.github/workflows/release_jaxlib_test.yml
@@ -47,14 +47,10 @@ jobs:
           mkdir -p /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.}
           if [[ ${{ github.event_name }} == "pull_request" ]]; then
             echo ${{ github.event_name }}
-            echo ${{ github.event }}
-            echo ${{ github.action }}
             echo ${CUDA_VERSION}
             echo ${ALPA_BRANCH}
           else
             echo ${{ github.event_name }}
-            echo ${{ github.event }}
-            echo ${{ github.action }}
             echo ${CUDA_VERSION}
             echo ${ALPA_BRANCH}
             echo ${TF_BRANCH}
@@ -108,7 +104,7 @@ jobs:
         env:
           GITHUB_TOKEN: ${{ secrets.PAT_TOKEN }}
           TAG: ${{ steps.latesttag.outputs.tag }}
-        run:
+        run: |
           ls -ltr
           cd ..
           python build_jaxlib/release/generate_pypi_index.py --tag ${TAG} --dry-run

From 4142e9bbe888fe1965f6fb4c43d53032bba4ac44 Mon Sep 17 00:00:00 2001
From: Hao Zhang <sjtu.haozhang@gmail.com>
Date: Sat, 2 Jul 2022 16:09:05 -0700
Subject: [PATCH 20/36] update

---
 .github/workflows/release_jaxlib_test.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/release_jaxlib_test.yml b/.github/workflows/release_jaxlib_test.yml
index 9841ea23c..5ad72708a 100644
--- a/.github/workflows/release_jaxlib_test.yml
+++ b/.github/workflows/release_jaxlib_test.yml
@@ -96,7 +96,7 @@ jobs:
           TAG: ${{ steps.latesttag.outputs.tag }}
 
       - uses: actions/checkout@v3
-        name: "Checkout doc repo"
+        name: "Checkout docs"
         with:
           repository: "alpa-projects/alpa-projects.github.io"
 

From 45605e5aab09078d7a9833312850351fd3d96da0 Mon Sep 17 00:00:00 2001
From: Hao Zhang <sjtu.haozhang@gmail.com>
Date: Sat, 2 Jul 2022 16:18:38 -0700
Subject: [PATCH 21/36] update

---
 .github/workflows/release_jaxlib_test.yml | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/release_jaxlib_test.yml b/.github/workflows/release_jaxlib_test.yml
index 5ad72708a..51c069e4d 100644
--- a/.github/workflows/release_jaxlib_test.yml
+++ b/.github/workflows/release_jaxlib_test.yml
@@ -73,13 +73,13 @@ jobs:
         with:
           python-version: 3.8
 
-      - uses: actions/checkout@v3
-
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
           python -m pip install github3.py requests
 
+      - uses: actions/checkout@v3
+
       - name: Get latest tag
         id: latesttag
         uses: "WyriHaximus/github-action-get-previous-tag@v1"

From 5e4071c4dfb6fddd6fe9020928ab1b43debf7af4 Mon Sep 17 00:00:00 2001
From: Hao Zhang <sjtu.haozhang@gmail.com>
Date: Sat, 2 Jul 2022 16:22:25 -0700
Subject: [PATCH 22/36] update

---
 .github/workflows/release_jaxlib_test.yml | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/release_jaxlib_test.yml b/.github/workflows/release_jaxlib_test.yml
index 51c069e4d..b2d966ea2 100644
--- a/.github/workflows/release_jaxlib_test.yml
+++ b/.github/workflows/release_jaxlib_test.yml
@@ -12,7 +12,6 @@ on:
         required: true
         default: 'master'
 
-
 jobs:
 
   clean-up:
@@ -79,7 +78,10 @@ jobs:
           python -m pip install github3.py requests
 
       - uses: actions/checkout@v3
-
+      - name: Test
+        run: |
+          ls -ltr
+          
       - name: Get latest tag
         id: latesttag
         uses: "WyriHaximus/github-action-get-previous-tag@v1"

From 99aaeb62283651bd769bc2e0262bd36eef540bc1 Mon Sep 17 00:00:00 2001
From: Hao Zhang <sjtu.haozhang@gmail.com>
Date: Sat, 2 Jul 2022 16:26:43 -0700
Subject: [PATCH 23/36] update

---
 .github/workflows/release_jaxlib_test.yml | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/release_jaxlib_test.yml b/.github/workflows/release_jaxlib_test.yml
index b2d966ea2..6883d9eba 100644
--- a/.github/workflows/release_jaxlib_test.yml
+++ b/.github/workflows/release_jaxlib_test.yml
@@ -78,10 +78,14 @@ jobs:
           python -m pip install github3.py requests
 
       - uses: actions/checkout@v3
+        with:
+          fetch-depth: 0
+
       - name: Test
         run: |
           ls -ltr
-          
+          git branch
+
       - name: Get latest tag
         id: latesttag
         uses: "WyriHaximus/github-action-get-previous-tag@v1"

From cb08cec43b7ca3bebce06461af17f01a17aaeb73 Mon Sep 17 00:00:00 2001
From: Hao Zhang <sjtu.haozhang@gmail.com>
Date: Sat, 2 Jul 2022 16:38:03 -0700
Subject: [PATCH 24/36] update

---
 .github/workflows/release_jaxlib_test.yml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/.github/workflows/release_jaxlib_test.yml b/.github/workflows/release_jaxlib_test.yml
index 6883d9eba..a1f6a851d 100644
--- a/.github/workflows/release_jaxlib_test.yml
+++ b/.github/workflows/release_jaxlib_test.yml
@@ -113,4 +113,5 @@ jobs:
         run: |
           ls -ltr
           cd ..
+          ls -ltr
           python build_jaxlib/release/generate_pypi_index.py --tag ${TAG} --dry-run

From 63c2686b9cbf2ab37295dcb45f87c4b55400493d Mon Sep 17 00:00:00 2001
From: Hao Zhang <sjtu.haozhang@gmail.com>
Date: Sat, 2 Jul 2022 16:41:41 -0700
Subject: [PATCH 25/36] update

---
 .github/workflows/release_jaxlib_test.yml | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/release_jaxlib_test.yml b/.github/workflows/release_jaxlib_test.yml
index a1f6a851d..ae2f3eb55 100644
--- a/.github/workflows/release_jaxlib_test.yml
+++ b/.github/workflows/release_jaxlib_test.yml
@@ -113,5 +113,6 @@ jobs:
         run: |
           ls -ltr
           cd ..
-          ls -ltr
+          ls -l
+          ls -l alpa/
           python build_jaxlib/release/generate_pypi_index.py --tag ${TAG} --dry-run

From f827d60cdbc7899e30b77f2550f724701d2d29f6 Mon Sep 17 00:00:00 2001
From: Hao Zhang <sjtu.haozhang@gmail.com>
Date: Sat, 2 Jul 2022 16:46:07 -0700
Subject: [PATCH 26/36] update

---
 .github/workflows/release_jaxlib_test.yml | 9 ++-------
 1 file changed, 2 insertions(+), 7 deletions(-)

diff --git a/.github/workflows/release_jaxlib_test.yml b/.github/workflows/release_jaxlib_test.yml
index ae2f3eb55..429bc1de5 100644
--- a/.github/workflows/release_jaxlib_test.yml
+++ b/.github/workflows/release_jaxlib_test.yml
@@ -101,18 +101,13 @@ jobs:
           GITHUB_TOKEN: ${{ secrets.PAT_TOKEN }}
           TAG: ${{ steps.latesttag.outputs.tag }}
 
-      - uses: actions/checkout@v3
-        name: "Checkout docs"
-        with:
-          repository: "alpa-projects/alpa-projects.github.io"
-
       - name: "Generate"
         env:
           GITHUB_TOKEN: ${{ secrets.PAT_TOKEN }}
           TAG: ${{ steps.latesttag.outputs.tag }}
         run: |
+          
+          git clone https://$GITHUB_TOKEN@github.com/alpa-projects/alpa-projects.github.io
           ls -ltr
-          cd ..
-          ls -l
           ls -l alpa/
           python build_jaxlib/release/generate_pypi_index.py --tag ${TAG} --dry-run

From 3197fa6bcaeb70c0835faedf45185a4bd2876d1e Mon Sep 17 00:00:00 2001
From: Hao Zhang <sjtu.haozhang@gmail.com>
Date: Sat, 2 Jul 2022 17:32:15 -0700
Subject: [PATCH 27/36] update

---
 .github/workflows/release_jaxlib_test.yml   | 13 ++-----------
 build_jaxlib/release/generate_pypi_index.py | 18 ++++++++++--------
 2 files changed, 12 insertions(+), 19 deletions(-)

diff --git a/.github/workflows/release_jaxlib_test.yml b/.github/workflows/release_jaxlib_test.yml
index 429bc1de5..f9be281f5 100644
--- a/.github/workflows/release_jaxlib_test.yml
+++ b/.github/workflows/release_jaxlib_test.yml
@@ -3,8 +3,6 @@ name: Release Jaxlib exp
 on:
   release:
     types: [created]
-  pull_request:
-    branches: [main]
   workflow_dispatch:
     inputs:
       tensorflow:
@@ -44,7 +42,7 @@ jobs:
       - name: Compile Jaxlib
         run: |
           mkdir -p /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.}
-          if [[ ${{ github.event_name }} == "pull_request" ]]; then
+          if [[ ${{ github.event_name }} == "release" ]]; then
             echo ${{ github.event_name }}
             echo ${CUDA_VERSION}
             echo ${ALPA_BRANCH}
@@ -81,11 +79,6 @@ jobs:
         with:
           fetch-depth: 0
 
-      - name: Test
-        run: |
-          ls -ltr
-          git branch
-
       - name: Get latest tag
         id: latesttag
         uses: "WyriHaximus/github-action-get-previous-tag@v1"
@@ -101,13 +94,11 @@ jobs:
           GITHUB_TOKEN: ${{ secrets.PAT_TOKEN }}
           TAG: ${{ steps.latesttag.outputs.tag }}
 
-      - name: "Generate"
+      - name: "Generate and update PyPI index"
         env:
           GITHUB_TOKEN: ${{ secrets.PAT_TOKEN }}
           TAG: ${{ steps.latesttag.outputs.tag }}
         run: |
-          
           git clone https://$GITHUB_TOKEN@github.com/alpa-projects/alpa-projects.github.io
           ls -ltr
-          ls -l alpa/
           python build_jaxlib/release/generate_pypi_index.py --tag ${TAG} --dry-run
diff --git a/build_jaxlib/release/generate_pypi_index.py b/build_jaxlib/release/generate_pypi_index.py
index b4ac587e9..f00f96c04 100644
--- a/build_jaxlib/release/generate_pypi_index.py
+++ b/build_jaxlib/release/generate_pypi_index.py
@@ -40,7 +40,7 @@ def list_wheels(repo, tag):
     return wheels
 
 
-def update_wheel_page(keep_list, site_repo, dry_run=False):
+def update_wheel_page(keep_list, site_repo, tag, dry_run=False):
     """Update the wheel page"""
     new_html = ""
     for asset in keep_list:
@@ -62,12 +62,14 @@ def run_cmd(cmd):
     run_cmd(["git", "fetch"])
     run_cmd(["git", "checkout", "-B", "master", "origin/master"])
     wheel_html_path = os.path.join(site_repo, "wheels.html")
-    print(f"Wheel page changed, update {wheel_html_path}..")
-    if not dry_run:
-        open(wheel_html_path, "w").write(new_html)
-        run_cmd(["git", "add", "wheels.html"])
-        run_cmd(["git", "commit", "-am", "wheel update at %s" % datetime.now()])
-        run_cmd(["git", "push", "origin", "master"])
+    if not os.path.exists(wheel_html_path) or open(wheel_html_path, "r").read() != new_html:
+        print(f"Wheel page changed, update {wheel_html_path}..")
+        if not dry_run:
+            open(wheel_html_path, "w").write(new_html)
+            run_cmd(["git", "add", "wheels.html"])
+            run_cmd(["git", "commit", "-am",
+                     f"wheel update at {datetime.now()} from tag {tag}"])
+            run_cmd(["git", "push", "origin", "master"])
 
 
 def delete_assets(remove_list, dry_run):
@@ -92,7 +94,7 @@ def main():
         raise RuntimeError("need GITHUB_TOKEN")
     args = parser.parse_args()
     wheels = list_wheels(args.repo, args.tag)
-    update_wheel_page(wheels, args.site_path, args.dry_run)
+    update_wheel_page(wheels, args.site_path, args.tag, args.dry_run)
 
 
 if __name__ == "__main__":

From fa85bba6c076e4eb27159cc452908c1f6031e37d Mon Sep 17 00:00:00 2001
From: Hao Zhang <sjtu.haozhang@gmail.com>
Date: Sat, 2 Jul 2022 17:34:40 -0700
Subject: [PATCH 28/36] remove dry-run

---
 .github/workflows/release_jaxlib_test.yml | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/.github/workflows/release_jaxlib_test.yml b/.github/workflows/release_jaxlib_test.yml
index f9be281f5..8460b752f 100644
--- a/.github/workflows/release_jaxlib_test.yml
+++ b/.github/workflows/release_jaxlib_test.yml
@@ -85,10 +85,9 @@ jobs:
 
       - name: Upload wheels for CUDA${{ matrix.cuda }}
         run: |
-          echo "Publish jaxlib to tag ${TAG}"
           ls /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.}
           echo "Upload wheels to tag ${TAG}"
-          python build_jaxlib/release/wheel_upload.py --tag ${TAG} --path /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.} --dry-run
+          python build_jaxlib/release/wheel_upload.py --tag ${TAG} --path /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.}
         env:
           CUDA_VERSION: ${{ matrix.cuda }}
           GITHUB_TOKEN: ${{ secrets.PAT_TOKEN }}
@@ -101,4 +100,4 @@ jobs:
         run: |
           git clone https://$GITHUB_TOKEN@github.com/alpa-projects/alpa-projects.github.io
           ls -ltr
-          python build_jaxlib/release/generate_pypi_index.py --tag ${TAG} --dry-run
+          python build_jaxlib/release/generate_pypi_index.py --tag ${TAG}

From 76c76178bc31595ec1bc4b068a8ef178c2f14dd9 Mon Sep 17 00:00:00 2001
From: Hao Zhang <sjtu.haozhang@gmail.com>
Date: Sat, 2 Jul 2022 17:52:22 -0700
Subject: [PATCH 29/36] update

---
 .github/workflows/release_jaxlib_test.yml | 18 +++++++++++-------
 1 file changed, 11 insertions(+), 7 deletions(-)

diff --git a/.github/workflows/release_jaxlib_test.yml b/.github/workflows/release_jaxlib_test.yml
index 8460b752f..852fad30e 100644
--- a/.github/workflows/release_jaxlib_test.yml
+++ b/.github/workflows/release_jaxlib_test.yml
@@ -58,12 +58,17 @@ jobs:
           ALPA_BRANCH: ${{ github.ref }}
           TF_BRANCH: ${{ github.event.inputs.tensorflow }}
 
+      - name: Move CUDA${{ matrix.cuda }}
+        run: |
+          echo "Move to one single folder"
+          ls /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.}
+          mv /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.}/*.whl /data/alpa-pypi/packages/
+        env:
+          CUDA_VERSION: ${{ matrix.cuda }}
+
   publish:
     runs-on: [self-hosted]
     needs: [build-jaxlib]
-    strategy:
-      matrix:
-        cuda: ["11.1", "11.2", "11.3"]
     steps:
       - name: Set up Python 3.8
         uses: actions/setup-python@v2
@@ -83,13 +88,12 @@ jobs:
         id: latesttag
         uses: "WyriHaximus/github-action-get-previous-tag@v1"
 
-      - name: Upload wheels for CUDA${{ matrix.cuda }}
+      - name: Upload wheels
         run: |
-          ls /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.}
+          ls /data/alpa-pypi/packages/
           echo "Upload wheels to tag ${TAG}"
-          python build_jaxlib/release/wheel_upload.py --tag ${TAG} --path /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.}
+          python build_jaxlib/release/wheel_upload.py --tag ${TAG} --path /data/alpa-pypi/packages/
         env:
-          CUDA_VERSION: ${{ matrix.cuda }}
           GITHUB_TOKEN: ${{ secrets.PAT_TOKEN }}
           TAG: ${{ steps.latesttag.outputs.tag }}
 

From 717c9c158f3d1932a274781a6c894d2c82d7529c Mon Sep 17 00:00:00 2001
From: Hao Zhang <sjtu.haozhang@gmail.com>
Date: Sat, 2 Jul 2022 18:24:40 -0700
Subject: [PATCH 30/36] update temp

---
 .github/workflows/release_jaxlib_test.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/release_jaxlib_test.yml b/.github/workflows/release_jaxlib_test.yml
index 852fad30e..f64f10493 100644
--- a/.github/workflows/release_jaxlib_test.yml
+++ b/.github/workflows/release_jaxlib_test.yml
@@ -62,7 +62,7 @@ jobs:
         run: |
           echo "Move to one single folder"
           ls /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.}
-          mv /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.}/*.whl /data/alpa-pypi/packages/
+          mv /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.}/* /data/alpa-pypi/packages/
         env:
           CUDA_VERSION: ${{ matrix.cuda }}
 

From 48ef61e7f288f397994809eed154735e62874fb0 Mon Sep 17 00:00:00 2001
From: Hao Zhang <sjtu.haozhang@gmail.com>
Date: Sat, 2 Jul 2022 18:27:02 -0700
Subject: [PATCH 31/36] update

---
 .github/workflows/release_jaxlib_test.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/release_jaxlib_test.yml b/.github/workflows/release_jaxlib_test.yml
index f64f10493..c91dc7238 100644
--- a/.github/workflows/release_jaxlib_test.yml
+++ b/.github/workflows/release_jaxlib_test.yml
@@ -62,7 +62,7 @@ jobs:
         run: |
           echo "Move to one single folder"
           ls /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.}
-          mv /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.}/* /data/alpa-pypi/packages/
+#          mv /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.}/* /data/alpa-pypi/packages/
         env:
           CUDA_VERSION: ${{ matrix.cuda }}
 

From 8c29675c440243a3d781be3e7a18d40ac2c062a0 Mon Sep 17 00:00:00 2001
From: Hao Zhang <sjtu.haozhang@gmail.com>
Date: Sat, 2 Jul 2022 18:43:28 -0700
Subject: [PATCH 32/36] fix

---
 .github/workflows/release_jaxlib_test.yml | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/release_jaxlib_test.yml b/.github/workflows/release_jaxlib_test.yml
index c91dc7238..8d8d34473 100644
--- a/.github/workflows/release_jaxlib_test.yml
+++ b/.github/workflows/release_jaxlib_test.yml
@@ -103,5 +103,8 @@ jobs:
           TAG: ${{ steps.latesttag.outputs.tag }}
         run: |
           git clone https://$GITHUB_TOKEN@github.com/alpa-projects/alpa-projects.github.io
-          ls -ltr
+          cd alpa-projects.github.io
+          git config user.name github-actions
+          git config user.email github-actions@github.com
+          cd ..
           python build_jaxlib/release/generate_pypi_index.py --tag ${TAG}

From 9626baf96551ddc860f71ca75c47692bdfce70bc Mon Sep 17 00:00:00 2001
From: Hao Zhang <sjtu.haozhang@gmail.com>
Date: Sun, 3 Jul 2022 00:03:09 -0700
Subject: [PATCH 33/36] update

---
 .github/workflows/release_jaxlib.yml      | 96 -----------------------
 .github/workflows/release_jaxlib_test.yml | 22 +++---
 2 files changed, 12 insertions(+), 106 deletions(-)
 delete mode 100644 .github/workflows/release_jaxlib.yml

diff --git a/.github/workflows/release_jaxlib.yml b/.github/workflows/release_jaxlib.yml
deleted file mode 100644
index 506c7a4c1..000000000
--- a/.github/workflows/release_jaxlib.yml
+++ /dev/null
@@ -1,96 +0,0 @@
-name: Release Jaxlib
-
-on:
-  release:
-    types: [created]
-  workflow_dispatch:
-    inputs:
-      tensorflow:
-        description: 'TensorFlow-alpa branch to build'
-        required: true
-        default: 'master'
-
-
-jobs:
-
-  clean-up:
-    runs-on: [self-hosted]
-
-    steps:
-      - name: clean up images
-        run: |
-          docker image prune -f
-
-  build-jaxlib:
-    runs-on: [self-hosted]
-    needs: [clean-up]
-    strategy:
-      matrix:
-        cuda: ["11.1", "11.2", "11.3"]
-        python: ["3.7", "3.8", "3.9"]
-
-    steps:
-      - uses: actions/checkout@v3
-
-      - name: build image
-        run: |
-          docker build -t build-jaxlib-image-cuda${CUDA_VERSION} \
-            -f docker/build_jaxlib.Dockerfile docker/ \
-            --build-arg JAX_CUDA_VERSION=${CUDA_VERSION}
-        env:
-          CUDA_VERSION: ${{ matrix.cuda }}
-
-      - name: Compile Jaxlib
-        run: |
-          mkdir -p /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.}
-          if [[ ${{ github.event.action }} == "release" ]]; then
-            docker run --gpus all --tmpfs /build:exec \
-              --rm -v /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.}:/dist \
-              build-jaxlib-image-cuda${CUDA_VERSION} ${PYTHON_VERSION} \
-              cuda ${CUDA_VERSION} ${ALPA_BRANCH}
-          else
-            docker run --gpus all --tmpfs /build:exec \
-              --rm -v /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.}:/dist \
-              build-jaxlib-image-cuda${CUDA_VERSION} ${PYTHON_VERSION} \
-              cuda ${CUDA_VERSION} ${ALPA_BRANCH} ${TF_BRANCH}
-          fi
-
-        env:
-          CUDA_VERSION: ${{ matrix.cuda }}
-          PYTHON_VERSION: ${{ matrix.python }}
-          ALPA_BRANCH: ${{ github.ref }}
-          TF_BRANCH: ${{ github.event.inputs.tensorflow }}
-
-  publish:
-    runs-on: [self-hosted]
-    needs: [build-jaxlib]
-    strategy:
-      matrix:
-        cuda: ["11.1", "11.2", "11.3"]
-    steps:
-
-    - uses: actions/checkout@v3
-
-    - name: Set up Python ${{ matrix.python }}
-      uses: actions/setup-python@v2
-      with:
-        python-version: ${{ matrix.python }}
-
-    - name: Install dependencies
-      run: |
-        python -m pip install --upgrade pip
-        python -m pip install github3.py
-
-    - name: Get latest tag
-      id: latesttag
-      uses: "WyriHaximus/github-action-get-previous-tag@v1"
-
-    - name: Publish CUDA${{ matrix.cuda }}
-      run: |
-        echo "Publish jaxlib"
-        ls /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.}
-        python build_jaxlib/release/wheel_upload.py --tag ${TAG} /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.}
-      env:
-        CUDA_VERSION: ${{ matrix.cuda }}
-        GITHUB_TOKEN: ${{ secrets.PAT_TOKEN }}
-        TAG: ${{ steps.latesttag.outputs.tag }}
diff --git a/.github/workflows/release_jaxlib_test.yml b/.github/workflows/release_jaxlib_test.yml
index 8d8d34473..49e3a41ed 100644
--- a/.github/workflows/release_jaxlib_test.yml
+++ b/.github/workflows/release_jaxlib_test.yml
@@ -1,4 +1,4 @@
-name: Release Jaxlib exp
+name: Release Jaxlib
 
 on:
   release:
@@ -42,15 +42,17 @@ jobs:
       - name: Compile Jaxlib
         run: |
           mkdir -p /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.}
+          echo "Compile Python ${PYTHON_VERSION}, CUDA ${CUDA_VERSION}, ALPA BRANCH: ${ALPA_BRANCH}, TF_BRANCH: ${TF_BRANCH}"
           if [[ ${{ github.event_name }} == "release" ]]; then
-            echo ${{ github.event_name }}
-            echo ${CUDA_VERSION}
-            echo ${ALPA_BRANCH}
+            docker run --gpus all --tmpfs /build:exec \
+              --rm -v /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.}:/dist \
+              build-jaxlib-image-cuda${CUDA_VERSION} ${PYTHON_VERSION} \
+              cuda ${CUDA_VERSION} ${ALPA_BRANCH}
           else
-            echo ${{ github.event_name }}
-            echo ${CUDA_VERSION}
-            echo ${ALPA_BRANCH}
-            echo ${TF_BRANCH}
+            docker run --gpus all --tmpfs /build:exec \
+              --rm -v /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.}:/dist \
+              build-jaxlib-image-cuda${CUDA_VERSION} ${PYTHON_VERSION} \
+              cuda ${CUDA_VERSION} ${ALPA_BRANCH} ${TF_BRANCH}
           fi
         env:
           CUDA_VERSION: ${{ matrix.cuda }}
@@ -62,7 +64,7 @@ jobs:
         run: |
           echo "Move to one single folder"
           ls /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.}
-#          mv /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.}/* /data/alpa-pypi/packages/
+          mv /data/alpa-dist/jaxlib-alpa/cuda${CUDA_VERSION//.}/*.whl /data/alpa-pypi/packages/
         env:
           CUDA_VERSION: ${{ matrix.cuda }}
 
@@ -90,8 +92,8 @@ jobs:
 
       - name: Upload wheels
         run: |
-          ls /data/alpa-pypi/packages/
           echo "Upload wheels to tag ${TAG}"
+          ls /data/alpa-pypi/packages/
           python build_jaxlib/release/wheel_upload.py --tag ${TAG} --path /data/alpa-pypi/packages/
         env:
           GITHUB_TOKEN: ${{ secrets.PAT_TOKEN }}

From 63166d83a807b9336fefc9e98cc7141828b268a6 Mon Sep 17 00:00:00 2001
From: Hao Zhang <sjtu.haozhang@gmail.com>
Date: Sun, 3 Jul 2022 00:07:33 -0700
Subject: [PATCH 34/36] update [skp ci]

---
 .github/workflows/{release_jaxlib_test.yml => release_jaxlib.yml} | 0
 1 file changed, 0 insertions(+), 0 deletions(-)
 rename .github/workflows/{release_jaxlib_test.yml => release_jaxlib.yml} (100%)

diff --git a/.github/workflows/release_jaxlib_test.yml b/.github/workflows/release_jaxlib.yml
similarity index 100%
rename from .github/workflows/release_jaxlib_test.yml
rename to .github/workflows/release_jaxlib.yml

From bf13cecbd0c4176da318e763fb8204d44e71340f Mon Sep 17 00:00:00 2001
From: Hao Zhang <sjtu.haozhang@gmail.com>
Date: Sun, 3 Jul 2022 00:15:14 -0700
Subject: [PATCH 35/36] update [skip ci]

---
 build_jaxlib/release/README.md | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/build_jaxlib/release/README.md b/build_jaxlib/release/README.md
index c609c5896..12e4fc230 100644
--- a/build_jaxlib/release/README.md
+++ b/build_jaxlib/release/README.md
@@ -1,12 +1,14 @@
 # How to Release JaxLib and generate a PyPI Index
 
-## Upload jaxlib wheels as assets of a release tag
+1.  Upload jaxlib wheels as assets under a release tag.
 ```shell
-GITHUB_TOKEN="[admin_token]" python wheel_upload.py --tag [TAG] --path [PATH_TO_WHEELS]
+GITHUB_TOKEN=[ADMIN_TOKEN] python wheel_upload.py --tag [TAG] --path [PATH_TO_WHEELS]
 ```
 
-## Generate a html index page and commit it to Alpa doc page
+2. Generate a html index page and commit it to the master branch of Alpa doc repository.
 ```shell
-GITHUB_TOKEN="[admin_token]" python generate_pypi_index.py --tag [TAG]
+GITHUB_TOKEN=[ADMIN_TOKEN] python generate_pypi_index.py --tag [TAG]
 ```
 All wheel assets under `[TAG]` will be included in a html index page appeared in the doc repo.
+
+Please make sure the TAG is aligned in Step 1 and Step 2.

From 60607d37ed36bb2c38a9b3fa1b5e0596e3e22715 Mon Sep 17 00:00:00 2001
From: Hao Zhang <sjtu.haozhang@gmail.com>
Date: Sun, 3 Jul 2022 00:17:07 -0700
Subject: [PATCH 36/36] update

---
 build_jaxlib/release/generate_pypi_index.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/build_jaxlib/release/generate_pypi_index.py b/build_jaxlib/release/generate_pypi_index.py
index f00f96c04..c749e7891 100644
--- a/build_jaxlib/release/generate_pypi_index.py
+++ b/build_jaxlib/release/generate_pypi_index.py
@@ -1,4 +1,4 @@
-"""Update the wheels page, prune old nightly builds if necessary."""
+"""Generate and upload a PyPI index page given a tag."""
 import os
 import logging
 import argparse