ENH: GPU Support for Preview, Cache and Publish

.github/workflows/cache.yml

@@ -4,11 +4,35 @@ on:
     branches:
       - main
 jobs:
-  tests:
+  deploy-runner:
     runs-on: ubuntu-latest
     steps:
-      - name: Checkout
-        uses: actions/checkout@v2
+      - uses: iterative/setup-cml@v1
+      - uses: actions/checkout@v3
+        with:
+          ref: ${{ github.event.pull_request.head.sha }}
+      - name: Deploy runner on EC2
+        env:
+          REPO_TOKEN: ${{ secrets.QUANTECON_SERVICES_PAT }}
+          AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }}
+          AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
+        run: |
+          cml runner launch \
+              --cloud=aws \
+              --cloud-region=us-west-2 \
+              --cloud-type=p3.2xlarge \
+              --labels=cml-gpu \
+              --cloud-hdd-size=40
+  cache:
+    needs: deploy-runner
+    runs-on: [self-hosted, cml-gpu]
+    container:
+      image: docker://nvidia/cuda:11.2.1-devel-ubuntu20.04
+      options: --gpus all
+    steps:
+      - uses: actions/checkout@v3
+        with:
+          ref: ${{ github.event.pull_request.head.sha }}
       - name: Setup Anaconda
         uses: conda-incubator/setup-miniconda@v2
         with:
@@ -18,6 +42,11 @@ jobs:
           python-version: 3.9
           environment-file: environment.yml
           activate-environment: quantecon
+      - name: Install Jax and Upgrade CUDA
+        shell: bash -l {0}
+        run: |
+          pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+          nvidia-smi
       - name: Build HTML
         shell: bash -l {0}
         run: |

.github/workflows/ci-ec2-gpu.yml → .github/workflows/ci.yml

@@ -17,13 +17,14 @@ jobs:
           cml runner launch \
               --cloud=aws \
               --cloud-region=us-west-2 \
-              --cloud-type=p2.xlarge \
-              --labels=cml-gpu
+              --cloud-type=p3.2xlarge \
+              --labels=cml-gpu \
+              --cloud-hdd-size=40
   preview:
     needs: deploy-runner
     runs-on: [self-hosted, cml-gpu]
     container:
-      image: docker://iterativeai/cml:latest-gpu
+      image: docker://nvidia/cuda:11.2.1-devel-ubuntu20.04
       options: --gpus all
     steps:
       - uses: actions/checkout@v3
@@ -38,10 +39,18 @@ jobs:
           python-version: 3.9
           environment-file: environment.yml
           activate-environment: quantecon
+      - name: Install Jax and Upgrade CUDA
+        shell: bash -l {0}
+        run: |
+          pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+          nvidia-smi
       - name: Install latex dependencies
+        shell: bash -l {0}
         run: |
-          sudo apt-get -qq update
-          sudo apt-get install -y     \
+          apt-get -qq update
+          export DEBIAN_FRONTEND=noninteractive
+          apt-get install -y tzdata
+          apt-get install -y     \
             texlive-latex-recommended \
             texlive-latex-extra       \
             texlive-fonts-recommended \
@@ -80,6 +89,12 @@ jobs:
         shell: bash -l {0}
         run: |
           jb build lectures --path-output ./ -n -W --keep-going
+      - name: Upload Execution Reports
+        uses: actions/upload-artifact@v2
+        if: failure()
+        with:
+          name: execution-reports
+          path: _build/html/reports
       - name: Preview Deploy to Netlify
         uses: nwtgck/actions-netlify@v1.1
         with:

.github/workflows/publish.yml

@@ -4,9 +4,32 @@ on:
     tags:
       - 'publish*'
 jobs:
+  deploy-runner:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: iterative/setup-cml@v1
+      - uses: actions/checkout@v3
+        with:
+          ref: ${{ github.event.pull_request.head.sha }}
+      - name: Deploy runner on EC2
+        env:
+          REPO_TOKEN: ${{ secrets.QUANTECON_SERVICES_PAT }}
+          AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }}
+          AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
+        run: |
+          cml runner launch \
+              --cloud=aws \
+              --cloud-region=us-west-2 \
+              --cloud-type=p3.2xlarge \
+              --labels=cml-gpu \
+              --cloud-hdd-size=40
   publish:
     if: github.event_name == 'push' && startsWith(github.event.ref, 'refs/tags')
-    runs-on: ubuntu-latest
+    needs: deploy-runner
+    runs-on: [self-hosted, cml-gpu]
+    container:
+      image: docker://nvidia/cuda:11.2.1-devel-ubuntu20.04
+      options: --gpus all
     steps:
       - name: Checkout
         uses: actions/checkout@v2
@@ -19,10 +42,17 @@ jobs:
           python-version: 3.9
           environment-file: environment.yml
           activate-environment: quantecon
+      - name: Install Jax and Upgrade CUDA
+        shell: bash -l {0}
+        run: |
+          pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+          nvidia-smi
       - name: Install latex dependencies
         run: |
-          sudo apt-get -qq update
-          sudo apt-get install -y     \
+          apt-get -qq update
+          export DEBIAN_FRONTEND=noninteractive
+          apt-get install -y tzdata
+          apt-get install -y     \
             texlive-latex-recommended \
             texlive-latex-extra       \
             texlive-fonts-recommended \

environment.yml

@@ -14,11 +14,12 @@ dependencies:
     - ghp-import==1.1.0
     - sphinxcontrib-youtube==1.1.0
     - sphinx-togglebutton==0.3.1
-    # Sandpit Requirements
     - quantecon
     - array-to-latex
-    - PuLP
-    - cvxpy
-    - cvxopt
-    - cylp
     - prettytable
+    # Sandpit Requirements
+    # - PuLP
+    # - cvxpy
+    # - cvxopt
+    # - cylp
+