[WIP] ci example how to cache weights

.github/download-models-weights.py

@@ -0,0 +1,21 @@
+import argparse
+
+import torch
+
+fonts = {
+    "sold2_wireframe": "http://cmp.felk.cvut.cz/~mishkdmy/models/sold2_wireframe.pth",
+}
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser("WeightsDownloader")
+    parser.add_argument("--target_directory", "-t", required=False, default="target_directory")
+
+    args = parser.parse_args()
+
+    torch.hub.set_dir(args.target_directory)
+
+    for name, url in fonts.items():
+        print(f"Downloading weights of `{name}` from `url`. Caching to dir `{args.target_directory}`")
+        torch.hub.load_state_dict_from_url(url, model_dir=args.target_directory, map_location=torch.device("cpu"))
+
+    raise SystemExit(0)

.github/workflows/pr_test_cpu.yml

@@ -13,41 +13,88 @@ concurrency:
   cancel-in-progress: true
 
 jobs:
+
+  pre-tests:
+    runs-on: ubuntu-latest
+    outputs:
+      hash: ${{ steps.hashid.outputs.weights-hash }}
+    steps:
+    - uses: kornia/workflows/.github/actions/env@v1.12.1
+    - uses: actions/cache@v4
+      id: cache-weights
+      with:
+        path: weights/
+        key: model-weights-${{ hashFiles('.github/download-models-weights.py') }}
+        enableCrossOsArchive: true
+
+    - name: Download models weights...
+      if: steps.cache-weights.outputs.cache-hit != 'true'
+      run: python .github/download-models-weights.py -t weights/
+
+    - name: write hashid
+      id: hashid
+      run: echo "weights-hash=${{ hashFiles('.github/download-models-weights.py') }}" >> "$GITHUB_OUTPUT"
+
   tests-cpu:
+    needs: [pre-tests]
     strategy:
       fail-fast: true
       matrix:
         os: ['Ubuntu-latest', 'Windows-latest']
         pytorch-dtype: ['float32', 'float64']
 
-    uses: kornia/workflows/.github/workflows/tests.yml@v1.11.0
+    uses: kornia/workflows/.github/workflows/tests.yml@v1.12.1
     with:
       os: ${{ matrix.os }}
       python-version: '["3.8", "3.12"]'
       pytorch-version: '["1.9.1", "2.4.0"]'
       pytorch-dtype: ${{ matrix.pytorch-dtype }}
+      cache-path: weights/
+      cache-key: model-weights-${{ needs.pre-tests.outputs.hash }}
+      cache-restore-keys: |
+        model-weights-${{ needs.pre-tests.outputs.hash }}
+        model-weights-
 
   tests-cpu-macos:
-    uses: kornia/workflows/.github/workflows/tests.yml@v1.11.0
+    needs: [pre-tests]
+    uses: kornia/workflows/.github/workflows/tests.yml@v1.12.1
     with:
       os: 'MacOS-latest'
       python-version: '["3.8", "3.12"]'
       pytorch-dtype: 'float32'
+      cache-path: weights/
+      cache-key: model-weights-${{ needs.pre-tests.outputs.hash }}
+      cache-restore-keys: |
+        model-weights-${{ needs.pre-tests.outputs.hash }}
+        model-weights-
 
 
   coverage:
-    uses: kornia/workflows/.github/workflows/coverage.yml@v1.11.0
+    needs: [pre-tests]
+    uses: kornia/workflows/.github/workflows/coverage.yml@v1.12.1
+    with:
+      cache-path: weights/
+      cache-key: model-weights-${{ needs.pre-tests.outputs.hash }}
+      cache-restore-keys: |
+        model-weights-${{ needs.pre-tests.outputs.hash }}
+        model-weights-
 
   typing:
-    uses: kornia/workflows/.github/workflows/mypy.yml@v1.11.0
+    uses: kornia/workflows/.github/workflows/mypy.yml@v1.12.1
 
   tutorials:
     uses: kornia/workflows/.github/workflows/tutorials.yml@v1.12.1
 
   docs:
-    uses: kornia/workflows/.github/workflows/docs.yml@v1.11.0
+    needs: [pre-tests]
+    uses: kornia/workflows/.github/workflows/docs.yml@v1.12.1
     with:
       python-version: '["3.11"]'
+      cache-path: weights/
+      cache-key: model-weights-${{ needs.pre-tests.outputs.hash }}
+      cache-restore-keys: |
+        model-weights-${{ needs.pre-tests.outputs.hash }}
+        model-weights-
 
   collector:
     needs: [coverage, tests-cpu, tutorials, typing, docs]
@@ -60,6 +107,7 @@ jobs:
 
   tests-nightly:
     if:  contains(github.event.pull_request.labels.*.name, 'nightly')
+    needs: [pre-tests]
     name: ${{ matrix.os }} - torch-nightly, ${{ matrix.pytorch-dtype }}
 
     strategy:
@@ -68,8 +116,13 @@ jobs:
         os: ['Ubuntu-latest', 'Windows-latest'] #, 'MacOS-latest'] add it when https://github.com/pytorch/pytorch/pull/89262 be merged
         pytorch-dtype: ['float32', 'float64']
 
-    uses: kornia/workflows/.github/workflows/tests.yml@v1.11.0
+    uses: kornia/workflows/.github/workflows/tests.yml@v1.12.1
     with:
       os: ${{ matrix.os }}
       pytorch-version: '["nightly"]'
       pytorch-dtype: ${{ matrix.pytorch-dtype }}
+      cache-path: weights/
+      cache-key: model-weights-${{ needs.pre-tests.outputs.hash }}
+      cache-restore-keys: |
+        model-weights-${{ needs.pre-tests.outputs.hash }}
+        model-weights-

.github/workflows/scheduled_test_cpu.yml

@@ -11,57 +11,107 @@ concurrency:
   cancel-in-progress: true
 
 jobs:
+  pre-tests:
+    runs-on: ubuntu-latest
+    outputs:
+      hash: ${{ steps.hashid.outputs.weights-hash }}
+    steps:
+    - uses: kornia/workflows/.github/actions/env@v1.12.1
+    - uses: actions/cache@v4
+      id: cache-weights
+      with:
+        path: weights/
+        key: model-weights-${{ hashFiles('.github/download-models-weights.py') }}
+        enableCrossOsArchive: true
+
+    - name: Download models weights...
+      if: steps.cache-weights.outputs.cache-hit != 'true'
+      run: python .github/download-models-weights.py -t weights/
+
+    - name: write hashid
+      id: hashid
+      run: echo "weights-hash=${{ hashFiles('.github/download-models-weights.py') }}" >> "$GITHUB_OUTPUT"
+
   tests-cpu-ubuntu:
+    needs: [pre-tests]
     strategy:
       fail-fast: false
       matrix:
         # os: ['Ubuntu-latest', 'Windows-latest', 'MacOS-latest']
         pytorch-dtype: ['float32', 'float64']
 
-    uses: kornia/workflows/.github/workflows/tests.yml@v1.11.0
+    uses: kornia/workflows/.github/workflows/tests.yml@v1.12.1
     with:
       os: 'Ubuntu-latest'
       python-version: '["3.8", "3.9", "3.10", "3.11", "3.12"]'
       pytorch-version: '["1.9.1", "1.10.2", "1.11.0", "1.12.1", "1.13.1", "2.0.1", "2.1.2", "2.2.2", "2.3.1", "2.4.0"]'
       pytorch-dtype: ${{ matrix.pytorch-dtype }}
       pytest-extra: '--runslow'
-
-
+      cache-path: weights/
+      cache-key: model-weights-${{ needs.pre-tests.outputs.hash }}
+      cache-restore-keys: |
+        model-weights-${{ needs.pre-tests.outputs.hash }}
+        model-weights-
 
   tests-cpu-windows:
+    needs: [pre-tests]
     strategy:
       fail-fast: true
       matrix:
         pytorch-dtype: ['float32', 'float64']
 
-    uses: kornia/workflows/.github/workflows/tests.yml@v1.11.0
+    uses: kornia/workflows/.github/workflows/tests.yml@v1.12.1
     with:
       os: 'Windows-latest'
       python-version: '["3.12"]'
       pytorch-version: '["1.9.1", "2.4.0"]'
       pytorch-dtype: ${{ matrix.pytorch-dtype }}
+      cache-path: weights/
+      cache-key: model-weights-${{ needs.pre-tests.outputs.hash }}
+      cache-restore-keys: |
+        model-weights-${{ needs.pre-tests.outputs.hash }}
+        model-weights-
 
   tests-cpu-mac:
+    needs: [pre-tests]
     strategy:
       fail-fast: true
       matrix:
         pytorch-dtype: ['float32', 'float64']
 
-    uses: kornia/workflows/.github/workflows/tests.yml@v1.11.0
+    uses: kornia/workflows/.github/workflows/tests.yml@v1.12.1
     with:
       os: 'MacOS-latest'
       pytorch-dtype: ${{ matrix.pytorch-dtype }}
+      cache-path: weights/
+      cache-key: model-weights-${{ needs.pre-tests.outputs.hash }}
+      cache-restore-keys: |
+        model-weights-${{ needs.pre-tests.outputs.hash }}
+        model-weights-
 
   coverage:
-    uses: kornia/workflows/.github/workflows/coverage.yml@v1.11.0
+    needs: [pre-tests]
+    uses: kornia/workflows/.github/workflows/coverage.yml@v1.12.1
+    with:
+      cache-path: weights/
+      cache-key: model-weights-${{ needs.pre-tests.outputs.hash }}
+      cache-restore-keys: |
+        model-weights-${{ needs.pre-tests.outputs.hash }}
+        model-weights-
 
   typing:
-    uses: kornia/workflows/.github/workflows/mypy.yml@v1.11.0
+    uses: kornia/workflows/.github/workflows/mypy.yml@v1.12.1
 
   tutorials:
     uses: kornia/workflows/.github/workflows/tutorials.yml@v1.12.1
 
   docs:
-    uses: kornia/workflows/.github/workflows/docs.yml@v1.11.0
+    needs: [pre-tests]
+    uses: kornia/workflows/.github/workflows/docs.yml@v1.12.1
     with:
       python-version: '["3.11"]'
+      cache-path: weights/
+      cache-key: model-weights-${{ needs.pre-tests.outputs.hash }}
+      cache-restore-keys: |
+        model-weights-${{ needs.pre-tests.outputs.hash }}
+        model-weights-

conftest.py

@@ -1,3 +1,4 @@
+import os
 import sys
 from functools import partial
 from itertools import product
@@ -18,6 +19,9 @@
     _backends_non_experimental = []
 
 
+WEIGHTS_CACHE_DIR = "weights/"
+
+
 def get_test_devices() -> Dict[str, torch.device]:
     """Create a dictionary with the devices to test the source code. CUDA devices will be test only in case the
     current hardware supports it.
@@ -189,7 +193,9 @@ def pytest_sessionstart(session):
             ex
         ) and "Dynamo is not supported on Python 3.12+" not in str(ex):
             raise ex
-    # TODO: cache all torch.load weights/states here to not impact on test suite
+
+    os.makedirs(WEIGHTS_CACHE_DIR, exist_ok=True)
+    torch.hub.set_dir(WEIGHTS_CACHE_DIR)
 
 
 def _get_env_info() -> Dict[str, Dict[str, str]]:
@@ -249,7 +255,7 @@ def pytest_report_header(config):
     import onnx
 
     env_info = _get_env_info()
-
+    CACHED_WEIGTHS = os.listdir(WEIGHTS_CACHE_DIR)
     if "cpu" in env_info:
         desired_cpu_info = ["Model name", "Architecture", "CPU(s)", "Thread(s) per core", "CPU max MHz", "CPU min MHz"]
         cpu_info = "cpu info:\n" + "\n".join(
@@ -276,6 +282,7 @@ def pytest_report_header(config):
     - onnx-{onnx.__version__}
 {gcc_info}
 available optimizers: {TEST_OPTIMIZER_BACKEND}
+model weights cached: {CACHED_WEIGTHS}
 """