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 ****