[CI/CD] New PyPI Infra [skip ci]

[zhisbug]

Changes on the release_jaxlib workflow:

1. Usage 1: tagging. When you tag, what will happen is below:

• it will compile your tagged alpa branch into three alpa wheels, upload them to the official pypi, and generate bump version PR.
• it will compile 9 jaxlib wheels using the tensorflow-alpa branch you tagged as third-party in your tagged alpa branch, upload the 9 wheels as assets of this alpa tag release, generate a wheels.html using the 9 GitHub links, and commit&push the html to the master branch of the alpa-projects.github.io repo, which serves as a PyPI index page at the address https://alpa-projects.github.io/wheels.html.

2. Usage 2 (not recommended): manual workflow_dispatch in actions page. It will require you to input a valid branch name of the tensorflow-alpa repo. It will then compile 9 jaxlib wheels using your input tensorflow-alpa branch (instead of using the tagged tensorflow-alpa branch in the chosen alpa branch), upload the 9 wheels as assets to the latest tagged release, and do the rest mentioned in Usage 1 for wheels page. No alpa wheel release is involved here.

Changes on the docs.yml workflow:

Now it will overwrite the original master branch instead of deleting original files and putting new files. In this way, the wheels.html updates are kept independent from the doc updates

@zhuohan123 @ZYHowell Take a look at the above and let me know if it does not make sense.

[zhuohan123]

Correct me if I'm wrong. So after this change, when we modify the tensorflow-alpa repo, we need to push a commit into alpa's main branch and also tag that commit to trigger the CD process? What exactly do we need to do to update pypi? Why isn't this action being triggered for every commit to the main branch?

[zhisbug]

@zhuohan123

The action only happens when you tag a release. It does not trigger for every commit because that's too exhausted to build 9 jaxlib and 3 alpa wheels for every new commit . In the future, when we have more CI/CD resources we can increase the frequency to be nightly.

When you plan to tag a version and release, you need to make sure in alpa repo, that the correct tensorflow-alpa commit is tagged under thirdparty/tensorflow-alpa. That's all.

If you want to manually release jaxlib (instead of releasing by tagging alpa), you need to input a tensorflow-alpa commit hash at the action input.

[zhuohan123]

LGTM! New workflow sounds good to me.