From b106c46d78252ffd9205f8c8ddb52df26d6d2fec Mon Sep 17 00:00:00 2001
From: quotentiroler <max.nussbaumer@maxhealth.tech>
Date: Thu, 15 Oct 2026 17:01:45 +0200
Subject: [PATCH] [CI] Cache torch.compile artifacts for the pytorch unit tests

The pytorch test job compiles 23 models with torch.compile, and every run
starts from an empty TorchInductor cache, so each graph goes through the
full Inductor lowering, C++ codegen and gcc build again. On CPU this is
the largest block of the job: the tablecenternet compile test alone
takes over 10 minutes cold.

setup-env gets an opt-in cache-compile input that points
TORCHINDUCTOR_CACHE_DIR at ~/.cache/torchinductor and persists it with
actions/cache. Only pytest-torch enables it. The key carries the runner
image, Python version, torch version and the vector ISA Inductor picks
for code generation, plus a hash of doctr/**/*.py. A prefix restore-key
lets a run reuse the entries of a previous commit: Inductor looks entries
up by the traced graph, torch build and its config, so graphs a change
does not touch are reused and changed ones are compiled again.

Measured in a Linux container (python:3.11, 4 CPUs, torch 2.14.1+cpu),
with the CI command (pytest tests/pytorch/ -rs --cov):
- the 23 compile tests: 2315s with an empty cache, 535s with a cache
  restored from a previous run (-77%); the cache is 96 MB
- tablecenternet compile test: 665s -> 52s
- full suite with the cache restored into a fresh container: same 452
  test ids, same outcomes (447 passed, 5 skipped) and same coverage
  totals (8711/10741 lines) as the cold baseline run
---
 .github/actions/setup-env/action.yml | 33 +++++++++++++++++++++++++++-
 .github/workflows/main.yml           |  1 +
 2 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/.github/actions/setup-env/action.yml b/.github/actions/setup-env/action.yml
index d283ca939c..ca711d49d7 100644
--- a/.github/actions/setup-env/action.yml
+++ b/.github/actions/setup-env/action.yml
@@ -1,7 +1,7 @@
 # Shared environment setup for all docTR CI jobs.
 
 name: "Setup docTR environment"
-description: "Python setup with pip caching, CPU-only PyTorch, package install and optional pretrained model cache"
+description: "Python setup with pip caching, CPU-only PyTorch, package install and optional pretrained model and torch.compile caches"
 
 inputs:
   python-version:
@@ -20,11 +20,16 @@ inputs:
     description: "Cache pretrained models (~/.cache/doctr) across runs"
     required: false
     default: "true"
+  cache-compile:
+    description: "Cache TorchInductor (torch.compile) artifacts across runs"
+    required: false
+    default: "false"
 
 runs:
   using: composite
   steps:
     - name: Set up Python
+      id: setup-python
       uses: actions/setup-python@v6
       with:
         python-version: ${{ inputs.python-version }}
@@ -63,3 +68,29 @@ runs:
         if [ -n "${{ inputs.extra-requirements }}" ]; then
           pip install -r "${{ inputs.extra-requirements }}"
         fi
+
+    - name: Resolve torch.compile cache key
+      if: inputs.cache-compile == 'true'
+      id: compile-cache
+      shell: bash
+      run: |
+        echo "TORCHINDUCTOR_CACHE_DIR=$HOME/.cache/torchinductor" >> "$GITHUB_ENV"
+        TORCH_ISA=$(python - <<'EOF'
+        import torch
+        try:
+            from torch._inductor.cpu_vec_isa import pick_vec_isa
+            isa = type(pick_vec_isa()).__name__
+        except ImportError:
+            isa = torch.backends.cpu.get_cpu_capability()
+        print(f"{torch.__version__}-{isa}")
+        EOF
+        )
+        echo "prefix=torchinductor-${ImageOS}-py${{ steps.setup-python.outputs.python-version }}-${TORCH_ISA}" >> "$GITHUB_OUTPUT"
+
+    - name: Cache torch.compile artifacts
+      if: inputs.cache-compile == 'true'
+      uses: actions/cache@v6
+      with:
+        path: ~/.cache/torchinductor
+        key: ${{ steps.compile-cache.outputs.prefix }}-${{ hashFiles('doctr/**/*.py') }}
+        restore-keys: ${{ steps.compile-cache.outputs.prefix }}-
diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index b5a00436e7..947715a259 100644
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -51,6 +51,7 @@ jobs:
         with:
           python-version: ${{ matrix.python }}
           extras: viz,html,testing
+          cache-compile: "true"
       - name: Run unittests
         run: pytest tests/pytorch/ -rs --cov --cov-report=xml:coverage-pt.xml
       - uses: actions/upload-artifact@v7