[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/[email protected]
+ - 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/[email protected]
+ needs: [pre-tests]
+ uses: kornia/workflows/.github/workflows/[email protected]
  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/[email protected]
+ needs: [pre-tests]
+ uses: kornia/workflows/.github/workflows/[email protected]
+ 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/[email protected]
 
  docs:
- uses: kornia/workflows/.github/workflows/[email protected]
+ needs: [pre-tests]
+ uses: kornia/workflows/.github/workflows/[email protected]
  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/[email protected]
+ - 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/[email protected]
+ needs: [pre-tests]
+ uses: kornia/workflows/.github/workflows/[email protected]
+ 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/[email protected]
 
  docs:
- uses: kornia/workflows/.github/workflows/[email protected]
+ needs: [pre-tests]
+ uses: kornia/workflows/.github/workflows/[email protected]
  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}
 """