TST: runs-on service for GPU

.github/runs-on.yml

@@ -0,0 +1,3 @@
+runners:
+ larger-disk:
+ disk: large

.github/workflows/ci.yml

@@ -2,35 +2,56 @@ name: Build Preview [using jupyter-book]
 on: [pull_request]
 jobs:
  preview:
- runs-on: quantecon-gpu
- container:
- image: ghcr.io/quantecon/lecture-python-container:cuda-12.6.0-anaconda-2024-10-py312-b
- 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 jaxlib
  shell: bash -l {0}
  run: |
- nvidia-smi
- - name: Check python version
+ pip install -U "jax[cuda12]"
+ python --version
+ python test-jax-install.py
+ - name: Install latex dependencies
+ 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: Check nvidia drivers
  shell: bash -l {0}
  run: |
- python --version
+ 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@v3
- with:
- workflow: cache.yml
- branch: main
- name: build-cache
- path: _build
+ # - name: Download "build" folder (cache)
+ #  uses: dawidd6/action-download-artifact@v3
+ #  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}
@@ -48,6 +69,9 @@ jobs:
  - name: Build HTML
  shell: bash -l {0}
  run: |
+ echo $PATH
+ python --version
+ python test-jax-install.py
  jb build lectures --path-output ./ -n -W --keep-going
  - name: Upload build folder
  uses: actions/upload-artifact@v4

lectures/status.md

@@ -25,8 +25,7 @@ You can check the backend used by JAX using:
 
 ```{code-cell} ipython3
 import jax
-# Check if JAX is using GPU
-print(f"JAX backend: {jax.devices()[0].platform}")
+print(f"JAX backend: {jax.devices()[0].platform}") # Check if JAX is using GPU
 ```
 
 and the hardware we are running on:

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}")