GitHub actions for JAX and TensorFlow (#104)

* GitHub actions for linting JAX code

Pylint is disabled since JAX container is not publicly available.
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* GitHub action for building JAX code

Disabled since JAX container is not publicly available.
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add GitHub actions for TensorFlow and license checking
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Debug GitHub action for license checking
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Use TensorFlow container for GitHub JAX tests
Signed-off-by: Tim Moon <tmoon@nvidia.com>

---------
Signed-off-by: Tim Moon <tmoon@nvidia.com>

.github/workflows/build.yml

@@ -3,16 +3,16 @@
 # See LICENSE for license information.
 
 # A workflow to trigger TE build on GitHub
-name: 'Test build'
+name: 'Build'
 on:
   pull_request:
   workflow_dispatch:
 jobs:
-  pyTorch:
-    name: 'pyTorch build'
+  pytorch:
+    name: 'PyTorch'
     runs-on: ubuntu-latest
     container:
-      image: nvcr.io/nvidia/pytorch:22.09-py3
+      image: nvcr.io/nvidia/pytorch:23.03-py3
       options: --user root
     steps:
       - name: 'Checkout'
@@ -20,29 +20,57 @@ jobs:
       - name: 'Build'
         run: |
           mkdir -p wheelhouse && \
-          pip wheel -w wheelhouse . -v
+          NVTE_FRAMEWORK=pytorch pip wheel -w wheelhouse . -v
       - name: 'Upload wheel'
         uses: actions/upload-artifact@v3
         with:
-          name: te_wheel
+          name: te_wheel_pyt
           path: wheelhouse/transformer_engine*.whl
           retention-days: 7
       - name: 'Install'
         run: pip install --no-cache-dir wheelhouse/transformer_engine*.whl
       - name: 'Sanity check'
         run: python tests/pytorch/test_sanity_import.py
-  TensorFlow:
-    name: 'TensorFlow build'
+  jax:
+    name: 'JAX'
     runs-on: ubuntu-latest
     container:
-      image: nvcr.io/nvidia/tensorflow:23.02-tf2-py3
+      #image: nvcr.io/nvidia/jax:XX.XX-py3  # Not yet available
+      image: nvcr.io/nvidia/tensorflow:23.03-tf2-py3
       options: --user root
     steps:
       - name: 'Checkout'
         uses: actions/checkout@v3
       - name: 'Build'
         run: |
-          apt-get update && apt-get install -y ninja-build pybind11-dev
+          apt-get -y update && \
+          apt -y install ninja-build pybind11-dev && \
+          pip install --upgrade "jax[cuda12_local]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html && \
+          mkdir -p wheelhouse && \
+          NVTE_FRAMEWORK=jax pip wheel -w wheelhouse . -v
+      - name: 'Upload wheel'
+        uses: actions/upload-artifact@v3
+        with:
+          name: te_wheel_jax
+          path: wheelhouse/transformer_engine*.whl
+          retention-days: 7
+      - name: 'Install'
+        run: pip install --no-cache-dir wheelhouse/transformer_engine*.whl
+      - name: 'Sanity check'
+        run: python tests/jax/test_sanity_import.py
+  tensorflow:
+    name: 'TensorFlow'
+    runs-on: ubuntu-latest
+    container:
+      image: nvcr.io/nvidia/tensorflow:23.03-tf2-py3
+      options: --user root
+    steps:
+      - name: 'Checkout'
+        uses: actions/checkout@v3
+      - name: 'Build'
+        run: |
+          apt-get -y update && \
+          apt -y install ninja-build pybind11-dev && \
           mkdir -p wheelhouse && \
           NVTE_FRAMEWORK=tensorflow pip wheel -w wheelhouse . -v
       - name: 'Upload wheel'

.github/workflows/build_docs.yml → .github/workflows/docs.yml

@@ -3,13 +3,14 @@
 # See LICENSE for license information.
 
 # A workflow to trigger the build of TE documentation on GitHub
-name: 'Build documentation'
+name: 'Documentation'
 on:
   pull_request:
   workflow_dispatch:
   workflow_call:
 jobs:
   build_docs:
+    name: 'Build'
     runs-on: ubuntu-latest
     steps:
       - name: 'Checkout'

.github/workflows/license.yml

+# Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+# A workflow to trigger the TE license check on GitHub
+name: 'License'
+on:
+  pull_request:
+  workflow_dispatch:
+jobs:
+  check:
+    name: 'Check'
+    runs-on: ubuntu-latest
+    steps:
+      - name: 'Checkout'
+        uses: actions/checkout@v3
+      - name: 'Check License'
+        run: |
+          export TE_PATH=.
+          bash ./qa/L0_license/test.sh

.github/workflows/lint.yml

@@ -2,37 +2,87 @@
 #
 # See LICENSE for license information.
 
-# A workflow to trigger C++ lint test on GitHub
-name: 'Lint test'
+# A workflow to trigger lint tests on GitHub
+name: 'Lint'
 on:
   pull_request:
   workflow_dispatch:
 jobs:
-  cpplint:
-    name: C++
+  pytorch_cpplint:
+    name: 'PyTorch C++'
     runs-on: ubuntu-latest
     steps:
       - name: Checkout
         uses: actions/checkout@v3
-      - name: Lint
+      - name: 'Lint'
         run: |
-          cp qa/L0_lint/CPPLINT.cfg .
           export CPP_ONLY=1
           export TE_PATH=.
           bash ./qa/L0_lint/test.sh
-  pylint:
-    name: Python
+  pytorch_pylint:
+    name: 'PyTorch Python'
     runs-on: ubuntu-latest
     container:
-      image: nvcr.io/nvidia/pytorch:22.09-py3
+      image: nvcr.io/nvidia/pytorch:23.03-py3
       options: --user root
     steps:
-      - name: Checkout
+      - name: 'Checkout'
         uses: actions/checkout@v3
-      - name: Lint
+      - name: 'Lint'
         run: |
-          cp qa/L0_lint/pylintrc .
           export PYTHON_ONLY=1
           export TE_PATH=.
           bash ./qa/L0_lint/test.sh
-
+  jax_cpplint:
+    name: 'JAX C++'
+    runs-on: ubuntu-latest
+    steps:
+      - name: 'Checkout'
+        uses: actions/checkout@v3
+      - name: 'Lint'
+        run: |
+          export CPP_ONLY=1
+          export TE_PATH=.
+          bash ./qa/L0_jax_lint/test.sh
+  jax_pylint:
+    name: 'JAX Python'
+    runs-on: ubuntu-latest
+    container:
+      #image: nvcr.io/nvidia/jax:XX.XX-py3  # Not yet available
+      image: nvcr.io/nvidia/tensorflow:23.03-tf2-py3
+      options: --user root
+    steps:
+      - name: 'Checkout'
+        uses: actions/checkout@v3
+      - name: 'Lint'
+        run: |
+          pip install --upgrade "jax[cuda12_local]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+          pip install flax
+          export PYTHON_ONLY=1
+          export TE_PATH=.
+          bash ./qa/L0_jax_lint/test.sh
+  tensorflow_cpplint:
+    name: 'TensorFlow C++'
+    runs-on: ubuntu-latest
+    steps:
+      - name: 'Checkout'
+        uses: actions/checkout@v3
+      - name: 'Lint'
+        run: |
+          export CPP_ONLY=1
+          export TE_PATH=.
+          bash ./qa/L0_tensorflow_lint/test.sh
+  tensorflow_pylint:
+    name: 'TensorFlow Python'
+    runs-on: ubuntu-latest
+    container:
+      image: nvcr.io/nvidia/tensorflow:23.03-tf2-py3
+      options: --user root
+    steps:
+      - name: 'Checkout'
+        uses: actions/checkout@v3
+      - name: 'Lint'
+        run: |
+          export PYTHON_ONLY=1
+          export TE_PATH=.
+          bash ./qa/L0_tensorflow_lint/test.sh

tests/jax/test_sanity_import.py

+# Copyright (c) 2022-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+try:
+    import transformer_engine.jax
+    te_imported = True
+except:
+    te_imported = False
+
+assert te_imported, 'transformer_engine import failed'
+print("OK")