ENH: migrate to runs-on GPU instance

.github/runs-on.yml

@@ -2,4 +2,4 @@ images:
   quantecon_ubuntu2204_ami:
     platform: "linux"
     arch: "x64"
-    ami: "ami-0d4c683ae7bde17e9"
+    ami: "ami-0d4c683ae7bde17e9" 

.github/workflows/cache.yml

@@ -6,17 +6,26 @@ on:
   workflow_dispatch:
 jobs:
   cache:
-    runs-on: quantecon-gpu
-    container:
-      image: ghcr.io/quantecon/lecture-python-container:cuda-12.8.1-anaconda-2024-10-py312
-      options: --gpus all
+    runs-on: "runs-on=${{ github.run_id }}/family=g4dn.2xlarge/image=ubuntu24-gpu-x64/disk=large"
     steps:
       - uses: actions/checkout@v4
         with:
           ref: ${{ github.event.pull_request.head.sha }}
-      - name: Check nvidia drivers
+      - name: Setup Anaconda
+        uses: conda-incubator/setup-miniconda@v3
+        with:
+          auto-update-conda: true
+          auto-activate-base: true
+          miniconda-version: 'latest'
+          python-version: "3.12"
+          environment-file: environment.yml
+          activate-environment: quantecon
+      - name: Install jax (and install checks for GPU)
         shell: bash -l {0}
         run: |
+          pip install -U "jax[cuda12]"
+          python --version
+          python scripts/test-jax-install.py
           nvidia-smi
       - name: Build HTML
         shell: bash -l {0}

.github/workflows/ci.yml

@@ -2,60 +2,64 @@ name: Build Project [using jupyter-book]
 on: [pull_request]
 jobs:
   preview:
-    runs-on: quantecon-gpu
-    container:
-      image: ghcr.io/quantecon/lecture-python-container:cuda-12.8.1-anaconda-2024-10-py312
-      options: --gpus all
+    runs-on: "runs-on=${{ github.run_id }}/family=g4dn.2xlarge/image=quantecon_ubuntu2204_ami/disk=large"
     steps:
       - uses: actions/checkout@v4
         with:
           ref: ${{ github.event.pull_request.head.sha }}
-      # Check nvidia drivers
-      - name: nvidia Drivers
-        shell: bash -l {0}
+      - name: Setup QuantEcon Environment (/runner/)
+        run: |
+          export PATH=$PATH:/opt/anaconda3/bin/
+          conda env create -f environment.yml -y
+          conda init bash && source ~/.bashrc 
+          conda activate quantecon
+      - name: Display nvidia
         run: nvidia-smi
       - name: Display Conda Environment Versions
-        shell: bash -l {0}
-        run: conda list
-      - name: Display Pip Versions
-        shell: bash -l {0}
-        run: pip list
-      - name: Download "build" folder (cache)
-        uses: dawidd6/action-download-artifact@v9
-        with:
-          workflow: cache.yml
-          branch: main
-          name: build-cache
-          path: _build
-      # Build Assets (Download Notebooks and PDF via LaTeX)
-      - name: Build Download Notebooks (sphinx-tojupyter)
-        shell: bash -l {0}
-        run: |
-          jb build lectures -n -W --keep-going --path-output ./ --builder=custom --custom-builder=jupyter
-          mkdir -p _build/html/_notebooks
-          cp -u _build/jupyter/*.ipynb _build/html/_notebooks
-      - name: Upload Execution Reports (Download Notebooks)
-        uses: actions/upload-artifact@v4
-        if: failure()
-        with:
-          name: execution-reports
-          path: _build/jupyter/reports
-      - name: Build PDF from LaTeX
-        shell: bash -l {0}
         run: |
-          jb build lectures --builder pdflatex --path-output ./ -W --keep-going
-          mkdir -p _build/html/_pdf
-          cp -u _build/latex/*.pdf _build/html/_pdf
-      - name: Upload Execution Reports (LaTeX)
-        uses: actions/upload-artifact@v4
-        if: failure()
-        with:
-          name: execution-reports
-          path: _build/latex/reports
+          export PATH=$PATH:/opt/anaconda3/bin/
+          conda list
+      # - name: Display Pip Versions
+      #   run: pip list
+      # - name: Download "build" folder (cache)
+      #   uses: dawidd6/action-download-artifact@v9
+      #   with:
+      #     workflow: cache.yml
+      #     branch: main
+      #     name: build-cache
+      #     path: _build
+      # Build Assets (Download Notebooks and PDF via LaTeX)
+      # - name: Build Download Notebooks (sphinx-tojupyter)
+      #   shell: bash -l {0}
+      #   run: |
+      #     jb build lectures -n -W --keep-going --path-output ./ --builder=custom --custom-builder=jupyter
+      #     mkdir -p _build/html/_notebooks
+      #     cp -u _build/jupyter/*.ipynb _build/html/_notebooks
+      # - name: Upload Execution Reports (Download Notebooks)
+      #   uses: actions/upload-artifact@v4
+      #   if: failure()
+      #   with:
+      #     name: execution-reports
+      #     path: _build/jupyter/reports
+      # - name: Build PDF from LaTeX
+      #   shell: bash -l {0}
+      #   run: |
+      #     jb build lectures --builder pdflatex --path-output ./ -W --keep-going
+      #     mkdir -p _build/html/_pdf
+      #     cp -u _build/latex/*.pdf _build/html/_pdf
+      # - name: Upload Execution Reports (LaTeX)
+      #   uses: actions/upload-artifact@v4
+      #   if: failure()
+      #   with:
+      #     name: execution-reports
+      #     path: _build/latex/reports
       # Final Build of HTML
       - name: Build HTML
         shell: bash -l {0}
         run: |
+          export PATH=$PATH:/opt/anaconda3/bin/
+          conda init bash && source ~/.bashrc 
+          conda activate quantecon
           jb build lectures --path-output ./ -n -W --keep-going
       - name: Upload Execution Reports (HTML)
         uses: actions/upload-artifact@v4

.github/workflows/collab.yml

@@ -2,7 +2,7 @@ name: Build Project on Google Collab (Execution)
 on: [pull_request]
 jobs:
   execution-checks:
-    runs-on: quantecon-gpu
+    runs-on: "runs-on=${{ github.run_id }}/family=g4dn.2xlarge/image=ubuntu24-gpu-x64/disk=large"
     container:
       image: docker://us-docker.pkg.dev/colab-images/public/runtime
       options: --gpus all

.github/workflows/publish.yml

@@ -6,20 +6,40 @@ on:
 jobs:
   publish:
     if: github.event_name == 'push' && startsWith(github.event.ref, 'refs/tags')
-    runs-on: quantecon-gpu
-    container:
-      image: ghcr.io/quantecon/lecture-python-container:cuda-12.8.1-anaconda-2024-10-py312
-      options: --gpus all
+    runs-on: "runs-on=${{ github.run_id }}/family=g4dn.2xlarge/image=ubuntu24-gpu-x64/disk=large"
     steps:
       - name: Checkout
         uses: actions/checkout@v4
-      - name: Install Git (required to commit notebooks)
-        shell: bash -l {0}
-        run: apt-get install -y git
-      - name: Check nvidia drivers
+      - name: Setup Anaconda
+        uses: conda-incubator/setup-miniconda@v3
+        with:
+          auto-update-conda: true
+          auto-activate-base: true
+          miniconda-version: 'latest'
+          python-version: "3.12"
+          environment-file: environment.yml
+          activate-environment: quantecon
+      - name: Install jax (and install checks for GPU)
         shell: bash -l {0}
         run: |
+          pip install -U "jax[cuda12]"
+          python --version
+          python scripts/test-jax-install.py
           nvidia-smi
+      - name: Install latex dependencies
+        shell: bash -l {0}
+        run: |
+          sudo apt-get -qq update
+          sudo apt-get install -y     \
+            texlive-latex-recommended \
+            texlive-latex-extra       \
+            texlive-fonts-recommended \
+            texlive-fonts-extra       \
+            texlive-xetex             \
+            latexmk                   \
+            xindy                     \
+            dvipng                    \
+            cm-super  
       - name: Display Conda Environment Versions
         shell: bash -l {0}
         run: conda list

scripts/test-jax-install.py

@@ -0,0 +1,21 @@
+import jax
+import jax.numpy as jnp
+
+devices = jax.devices()
+print(f"The available devices are: {devices}")
+
+@jax.jit
+def matrix_multiply(a, b):
+    return jnp.dot(a, b)
+
+# Example usage:
+key = jax.random.PRNGKey(0)
+x = jax.random.normal(key, (1000, 1000))
+y = jax.random.normal(key, (1000, 1000))
+z = matrix_multiply(x, y)
+
+# Now the function is JIT compiled and will likely run on GPU (if available)
+print(z)
+
+devices = jax.devices()
+print(f"The available devices are: {devices}")