From bb2a70b339af9096b9dbd18a3be16a21bc069c6a Mon Sep 17 00:00:00 2001
From: John Agapiou <jagapiou@google.com>
Date: Thu, 12 Sep 2024 06:07:16 -0700
Subject: [PATCH] Cache examples venv in GitHub actions for faster test runs

PiperOrigin-RevId: 673813732
Change-Id: Ideb3963d62b203a75112d93bd85c5390c129e1cd
---
 .github/actions/install-examples/action.yml   | 48 +++++++++++++++++--
 .github/actions/install-meltingpot/action.yml | 32 ++++++++-----
 2 files changed, 64 insertions(+), 16 deletions(-)

diff --git a/.github/actions/install-examples/action.yml b/.github/actions/install-examples/action.yml
index 009f2041..608fc4fc 100644
--- a/.github/actions/install-examples/action.yml
+++ b/.github/actions/install-examples/action.yml
@@ -1,15 +1,57 @@
 name: install-examples
 
+inputs:
+  python-version:
+    description: Python version
+    required: false
+    default: '3.11'
+    type: string
+
 runs:
   using: composite
   steps:
     - name: Install Melting Pot
+      id: install-meltingpot
       uses: ./.github/actions/install-meltingpot
+      with:
+        python-version: ${{ inputs.python-version }}
+
+    - name: Restore Examples installation
+      id: restore
+      uses: actions/cache/restore@0c45773b623bea8c8e75f6c82b208c3cf94ea4f9
+      with:
+        path: venv
+        key: install-examples-${{ steps.install-meltingpot.outputs.key }}-${{ hashFiles('examples/requirements.txt') }}
+
+    - name: Activate venv
+      shell: bash
+      run: |
+        echo "${PWD}/venv/bin" >> $GITHUB_PATH
 
     - name: Install requirements for examples
+      if: steps.restore.outputs.cache-hit != 'true'
       shell: bash
-      run: pip install -r examples/requirements.txt
+      run: |
+        pip install -r examples/requirements.txt
 
-    - name: Show installed dependencies
+    - name: Show installation
       shell: bash
-      run: pip list
+      run: |
+        which python
+        python --version
+        which pip
+        pip --version
+        which pylint
+        pylint --version
+        which pytest
+        pytest --version
+        which pytype
+        pytype --version
+        pip list
+
+    - name: Save Examples installation
+      if: steps.restore.outputs.cache-hit != 'true'
+      uses: actions/cache/save@0c45773b623bea8c8e75f6c82b208c3cf94ea4f9
+      with:
+        path: venv
+        key: ${{ steps.restore.outputs.cache-primary-key }}
diff --git a/.github/actions/install-meltingpot/action.yml b/.github/actions/install-meltingpot/action.yml
index 62f894f5..262466f5 100644
--- a/.github/actions/install-meltingpot/action.yml
+++ b/.github/actions/install-meltingpot/action.yml
@@ -7,28 +7,34 @@ inputs:
     default: '3.11'
     type: string
 
+outputs:
+  key:
+    description: Cache key that identifies the installation
+    value: ${{ steps.cache-key.outputs.key }}
+
 runs:
   using: composite
   steps:
-    - name: Get current runner
-      id: os-info
+    - name: Set up Python ${{ inputs.python-version }}
+      id: setup-python
+      uses: actions/setup-python@f677139bbe7f9c59b41e40162b753c062f5d49a3
+      with:
+        python-version: ${{ inputs.python-version }}
+
+    - name: Get cache key
+      id: cache-key
       shell: bash
       run: |
         if [ "${RUNNER_OS}" = 'macOS' ]; then
-          echo "name=$(sw_vers -productName)" >> $GITHUB_OUTPUT
-          echo "version=$(sw_vers -productVersion)" >> $GITHUB_OUTPUT
+          readonly name="$(sw_vers -productName)"
+          readonly version="$(sw_vers -productVersion)"
         elif [ "${RUNNER_OS}" = 'Linux' ]; then
-          echo "name=$(lsb_release -i -s)" >> $GITHUB_OUTPUT
-          echo "version=$(lsb_release -r -s)" >> $GITHUB_OUTPUT
+          readonly name="$(lsb_release -i -s)"
+          readonly version="$(lsb_release -r -s)"
         else
           exit 1
         fi
-
-    - name: Set up Python ${{ inputs.python-version }}
-      id: setup-python
-      uses: actions/setup-python@f677139bbe7f9c59b41e40162b753c062f5d49a3
-      with:
-        python-version: ${{ inputs.python-version }}
+        echo "key=${name}-${version}-py${{ steps.setup-python.outputs.python-version }}-${{ hashFiles('setup.py') }}" >> $GITHUB_OUTPUT
 
     - name: Restore Melting Pot installation
       id: restore
@@ -37,7 +43,7 @@ runs:
         path: |
           meltingpot/assets
           venv
-        key: install-meltingpot-${{ steps.os-info.outputs.name }}-${{ steps.os-info.outputs.version }}-py${{ steps.setup-python.outputs.python-version }}-${{ hashFiles('setup.py') }}
+        key: install-meltingpot-${{ steps.cache-key.outputs.key }}
 
     - name: Create venv
       if: steps.restore.outputs.cache-hit != 'true'