[fbsync] port special tests from CircleCI to GHA (#7396)

Reviewed By: vmoens

Differential Revision: D44416639

fbshipit-source-id: a1a088a1a8e04a38889652c1316ab00bb3f8f2ea

Author: NicolasHug

.github/unittest.sh renamed to .github/scripts/setup-env.sh

@@ -95,11 +95,3 @@ echo '::endgroup::'
 echo '::group::Collect PyTorch environment information'
 python -m torch.utils.collect_env
 echo '::endgroup::'
-
-echo '::group::Install testing utilities'
-pip install --progress-bar=off pytest pytest-mock pytest-cov
-echo '::endgroup::'
-
-echo '::group::Run tests'
-pytest --durations=25
-echo '::endgroup::'

.github/scripts/unittest.sh

@@ -0,0 +1,18 @@
+#!/usr/bin/env bash
+
+set -euo pipefail
+
+./.github/scripts/setup-env.sh
+
+# Prepare conda
+CONDA_PATH=$(which conda)
+eval "$(${CONDA_PATH} shell.bash hook)"
+conda activate ci
+
+echo '::group::Install testing utilities'
+pip install --progress-bar=off pytest pytest-mock pytest-cov
+echo '::endgroup::'
+
+echo '::group::Run unittests'
+pytest --durations=25
+echo '::endgroup::'

.github/workflows/test-linux.yml

@@ -1,4 +1,4 @@
-name: Unit-tests on Linux
+name: Tests on Linux
 
 on:
   pull_request:
@@ -10,7 +10,7 @@ on:
   workflow_dispatch:
 
 jobs:
-  tests:
+  unittests:
     strategy:
       matrix:
         python-version:
@@ -34,8 +34,70 @@ jobs:
       gpu-arch-version: ${{ matrix.gpu-arch-version }}
       timeout: 120
       script: |
+        set -euo pipefail
+
         export PYTHON_VERSION=${{ matrix.python-version }}
         export GPU_ARCH_TYPE=${{ matrix.gpu-arch-type }}
         export GPU_ARCH_VERSION=${{ matrix.gpu-arch-version }}
 
-        ./.github/unittest.sh
+        ./.github/scripts/unittest.sh
+
+  onnx:
+    uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
+    with:
+      repository: pytorch/vision
+      script: |
+        set -euo pipefail
+
+        export PYTHON_VERSION=3.8
+        export GPU_ARCH_TYPE=cpu
+
+        ./.github/scripts/setup-env.sh
+        
+        # Prepare conda
+        CONDA_PATH=$(which conda)
+        eval "$(${CONDA_PATH} shell.bash hook)"
+        conda activate ci
+        
+        echo '::group::Install ONNX'
+        pip install --progress-bar=off onnx onnxruntime
+        echo '::endgroup::'
+        
+        echo '::group::Install testing utilities'
+        pip install --progress-bar=off pytest
+        echo '::endgroup::'
+        
+        echo '::group::Run ONNX tests'
+        pytest --durations=25 -v test/test_onnx.py
+        echo '::endgroup::'
+
+  unittests-extended:
+    uses: pytorch/test-infra/.github/workflows/linux_job.yml@main
+    with:
+      repository: pytorch/vision
+      script: |
+        set -euo pipefail
+
+        export PYTHON_VERSION=3.8
+        export GPU_ARCH_TYPE=cpu
+
+        ./.github/scripts/setup-env.sh
+        
+        # Prepare conda
+        CONDA_PATH=$(which conda)
+        eval "$(${CONDA_PATH} shell.bash hook)"
+        conda activate ci
+        
+        echo '::group::Pre-download model weights'
+        pip install --progress-bar=off aiohttp aiofiles tqdm
+        python scripts/download_model_urls.py
+        echo '::endgroup::'
+        
+        echo '::group::Install testing utilities'
+        pip install --progress-bar=off pytest
+        echo '::endgroup::'
+        
+        echo '::group::Run extended unittests'
+        export PYTORCH_TEST_WITH_EXTENDED=1
+        pytest --durations=25 -v test/test_extended_*.py
+        echo '::endgroup::'

.github/workflows/test-macos.yml

@@ -1,4 +1,4 @@
-name: Unit-tests on macOS
+name: Tests on macOS
 
 on:
   pull_request:
@@ -10,7 +10,7 @@ on:
   workflow_dispatch:
 
 jobs:
-  tests:
+  unittests:
     strategy:
       matrix:
         python-version:
@@ -31,7 +31,9 @@ jobs:
       timeout: 240
       runner: ${{ matrix.runner }}
       script: |
+        set -euo pipefail
+
         export PYTHON_VERSION=${{ matrix.python-version }}
         export GPU_ARCH_TYPE=cpu
 
-        ./.github/unittest.sh
+        ./.github/scripts/unittest.sh

scripts/download_model_urls.py

@@ -0,0 +1,41 @@
+import asyncio
+import sys
+from pathlib import Path
+from time import perf_counter
+from urllib.parse import urlsplit
+
+import aiofiles
+import aiohttp
+from torchvision import models
+from tqdm.asyncio import tqdm
+
+
+async def main(download_root):
+    download_root.mkdir(parents=True, exist_ok=True)
+    urls = {weight.url for name in models.list_models() for weight in iter(models.get_model_weights(name))}
+
+    async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=None)) as session:
+        await tqdm.gather(*[download(download_root, session, url) for url in urls])
+
+
+async def download(download_root, session, url):
+    response = await session.get(url, params=dict(source="ci"))
+
+    assert response.ok
+
+    file_name = Path(urlsplit(url).path).name
+    async with aiofiles.open(download_root / file_name, "wb") as f:
+        async for data in response.content.iter_any():
+            await f.write(data)
+
+
+if __name__ == "__main__":
+    download_root = (
+        (Path(sys.argv[1]) if len(sys.argv) > 1 else Path("~/.cache/torch/hub/checkpoints")).expanduser().resolve()
+    )
+    print(f"Downloading model weights to {download_root}")
+    start = perf_counter()
+    asyncio.get_event_loop().run_until_complete(main(download_root))
+    stop = perf_counter()
+    minutes, seconds = divmod(stop - start, 60)
+    print(f"Download took {minutes:2.0f}m {seconds:2.0f}s")