From 8903b50a2b999c3282070f99cb9e79482228722f Mon Sep 17 00:00:00 2001
From: Ardavan Oskooi <ardavano@google.com>
Date: Sun, 20 Jun 2021 17:35:39 -0700
Subject: [PATCH] cache Python dependencies for CI (#1623)

---
 .github/workflows/build-ci.yml | 19 +++++++++++++++++--
 python/requirements.txt        |  9 +++++++++
 2 files changed, 26 insertions(+), 2 deletions(-)
 create mode 100644 python/requirements.txt

diff --git a/.github/workflows/build-ci.yml b/.github/workflows/build-ci.yml
index 4cabc9e42..9beb4a1ea 100644
--- a/.github/workflows/build-ci.yml
+++ b/.github/workflows/build-ci.yml
@@ -108,11 +108,26 @@ jobs:
       with:
         python-version: ${{ matrix.python-version }}
 
+    - name: Get pip cache dir
+      id: pip-cache
+      run: |
+        python -m pip install --upgrade pip wheel
+        echo "::set-output name=dir::$(pip cache dir)"
+
+    - name: pip cache
+      uses: actions/cache@v2
+      id: cache
+      with:
+        path: ${{ steps.pip-cache.outputs.dir }}
+        key: ${{ runner.os }}-py${{ matrix.python-version }}-mpi-${{ matrix.enable-mpi }}-pip-${{ hashFiles('**/requirements.txt') }}
+        restore-keys: ${{ runner.os }}-py${{ matrix.python-version }}-mpi-${{ matrix.enable-mpi }}-pip-
+
     - name: Install Python dependencies
-      run: pip install autograd h5py jax jaxlib matplotlib numpy parameterized pytest scipy
+      if: steps.cache.outputs.cache-hit != 'true'
+      run: pip install -r python/requirements.txt
 
     - name: Install mpi4py
-      if: matrix.enable-mpi
+      if: ${{ matrix.enable-mpi && steps.cache.outputs.cache-hit != true }}
       run: pip install mpi4py
 
     - name: Run autoreconf
diff --git a/python/requirements.txt b/python/requirements.txt
new file mode 100644
index 000000000..772aac158
--- /dev/null
+++ b/python/requirements.txt
@@ -0,0 +1,9 @@
+autograd
+h5py
+jax
+jaxlib
+matplotlib
+numpy
+parameterized
+pytest
+scipy