cache a venv to avoid pip usage (#10689)

* try built in pip caching

* try venv

* export venv

* set VIRTUAL_ENV

* revert that

* venv key

* fix

* ci cache hit?

* fix windows

.github/actions/setup-tinygrad/action.yml

@@ -44,31 +44,14 @@ runs:
       uses: actions/setup-python@v5
       with:
         python-version: ${{ inputs.python-version }}
-    - name: Upgrade pip
-      shell: bash
-      run: python -m pip install --upgrade pip
 
     # **** Caching packages ****
-    # TODO: key should include input.deps, but it can't since it can't contain commas
 
-    - name: Cache Python packages (Linux)
-      if: inputs.key != '' && runner.os == 'Linux'
+    - name: Cache Python packages
       uses: actions/cache@v4
       with:
-        path: ${{ env.Python3_ROOT_DIR }}/lib/python${{ inputs.python-version }}/site-packages
-        key: python-package-${{ inputs.key }}-${{ hashFiles('**/setup.py') }}-${{ env.PYTHON_CACHE_VERSION }}
-    - name: Cache Python packages (macOS)
-      if: inputs.key != '' && runner.os == 'macOS'
-      uses: actions/cache@v4
-      with:
-        path: /Users/runner/Library/Python/${{ inputs.python-version }}/lib/python/site-packages
-        key: osx-python-package-${{ inputs.key }}-${{ hashFiles('**/setup.py') }}-${{ env.PYTHON_CACHE_VERSION }}
-    - name: Cache Python packages (Windows)
-      if: inputs.key != '' && runner.os == 'Windows'
-      uses: actions/cache@v4
-      with:
-        path: ${{ env.Python3_ROOT_DIR }}\Lib\site-packages
-        key: windows-python-package-${{ inputs.key }}-${{ hashFiles('**/setup.py') }}-${{ env.PYTHON_CACHE_VERSION }}
+        path: .venv
+        key: venv-${{ runner.os }}-python-${{ inputs.python-version }}-${{ inputs.deps }}-${{ inputs.pydeps }}-${{ hashFiles('**/setup.py') }}-${{ env.PYTHON_CACHE_VERSION }}
 
     # **** Caching downloads ****
 
@@ -87,14 +70,41 @@ runs:
 
     # **** Python deps ****
 
-    - name: Install dependencies (with extra)
+    - name: Install dependencies in venv (with extra)
       if: inputs.deps != ''
       shell: bash
-      run: python -m pip install ${{ (runner.os == 'macOS' && '--user') || (runner.os != 'macOS' && '') }} -e ".[${{ inputs.deps }}]" ${{ inputs.pydeps }} --extra-index-url https://download.pytorch.org/whl/cpu --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/Triton-Nightly/pypi/simple/
-    - name: Install dependencies (without extra)
+      run: |
+        if [[ ! -d .venv ]]; then
+          python -m venv .venv
+          if [[ "$RUNNER_OS" == "Windows" ]]; then
+            source .venv/Scripts/activate
+          else
+            . .venv/bin/activate
+          fi
+          python -m pip install -e ".[${{ inputs.deps }}]" ${{ inputs.pydeps }} --extra-index-url https://download.pytorch.org/whl/cpu --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/Triton-Nightly/pypi/simple/
+        fi
+    - name: Install dependencies in venv (without extra)
       if: inputs.deps == ''
       shell: bash
-      run: python -m pip install ${{ (runner.os == 'macOS' && '--user') || (runner.os != 'macOS' && '') }} -e . ${{ inputs.pydeps }}
+      run: |
+        if [[ ! -d .venv ]]; then
+          python -m venv .venv
+          if [[ "$RUNNER_OS" == "Windows" ]]; then
+            source .venv/Scripts/activate
+          else
+            . .venv/bin/activate
+          fi
+          python -m pip install -e . ${{ inputs.pydeps }}
+        fi
+    - name: Export venv
+      shell: bash
+      run: |
+        echo "VIRTUAL_ENV=${{ github.workspace }}/.venv" >> "$GITHUB_ENV"
+        if [[ "$RUNNER_OS" == "Windows" ]]; then
+          echo "${{ github.workspace }}/.venv/Scripts" >> "$GITHUB_PATH"
+        else
+          echo "${{ github.workspace }}/.venv/bin" >> "$GITHUB_PATH"
+        fi
 
     # **** OpenCL ****
 

.github/workflows/test.yml

@@ -458,7 +458,8 @@ jobs:
         uses: ./.github/actions/setup-tinygrad
         with:
           key: onnxoptl
-          deps: testing,testing_tf
+          deps: testing
+          pydeps: "tensorflow==2.15.1 tensorflow_addons"
           python-version: '3.11'
           opencl: 'true'
           llvm: 'true'
@@ -493,7 +494,7 @@ jobs:
   testmodels:
     name: Models (llvm+cpu+gpu)
     runs-on: ubuntu-22.04
-    timeout-minutes: 10
+    timeout-minutes: 15
     env:
       IGNORE_OOB: 0
     steps:

autogen_stubs.sh

@@ -9,7 +9,7 @@ if [[ ! $(clang2py -V) ]]; then
   pip install clang==14.0.6
   git clone https://github.com/nimlgen/ctypeslib.git
   cd ctypeslib
-  pip install --user .
+  pip install .
   clang2py -V
   popd
 fi

setup.py

@@ -85,9 +85,5 @@ setup(name='tinygrad',
             "black",
             "numpy",
         ],
-        'testing_tf': [
-            "tensorflow==2.15.1",
-            "tensorflow_addons",
-        ],
       },
       include_package_data=True)