pre-download model weights in CI docs build

.circleci/config.yml.in

@@ -174,6 +174,26 @@ commands:
       - store_test_results:
           path: test-results
 
+  download_model_weights:
+    parameters:
+      extract_roots:
+        type: string
+        default: "torchvision/models"
+      background:
+        type: boolean
+        default: true
+    steps:
+      - apt_install:
+          args: parallel wget
+          descr: Install download utilitites
+      - run:
+          name: Download model weights
+          background: << parameters.background >>
+          command: |
+            mkdir -p ~/.cache/torch/hub/checkpoints
+            python scripts/collect_model_urls.py << parameters.extract_roots >> \
+                | parallel -j0 'wget --no-verbose -O ~/.cache/torch/hub/checkpoints/`basename {}` {}\?source=ci'
+
 binary_common: &binary_common
   parameters:
     # Edit these defaults to do a release
@@ -340,14 +360,8 @@ jobs:
     resource_class: xlarge
     steps:
       - checkout
-      - run:
-          name: Download model weights
-          background: true
-          command: |
-            sudo apt update -qy && sudo apt install -qy parallel wget
-            mkdir -p ~/.cache/torch/hub/checkpoints
-            python scripts/collect_model_urls.py torchvision/prototype/models \
-                | parallel -j0 'wget --no-verbose -O ~/.cache/torch/hub/checkpoints/`basename {}` {}\?source=ci'
+      - download_model_weights:
+          extract_roots: torchvision/prototype/models
       - install_torchvision
       - install_prototype_dependencies
       - pip_install:
@@ -1011,12 +1025,13 @@ jobs:
   build_docs:
     <<: *binary_common
     docker:
-      - image: "pytorch/manylinux-cuda100"
+      - image: circleci/python:3.7
     resource_class: 2xlarge+
     steps:
       - attach_workspace:
           at: ~/workspace
       - checkout
+      - download_model_weights
       - run:
           name: Setup
           command: .circleci/unittest/linux/scripts/setup_env.sh

scripts/collect_model_urls.py

@@ -2,21 +2,19 @@
 import re
 import sys
 
-MODEL_URL_PATTERN = re.compile(r"https://download[.]pytorch[.]org/models/.*?[.]pth")
+MODEL_URL_PATTERN = re.compile(r"https://download[.]pytorch[.]org/models/.+?[.]pth")
 
 
-def main(root):
+def main(*roots):
     model_urls = set()
-    for path in pathlib.Path(root).glob("**/*"):
-        if path.name.startswith("_") or not path.suffix == ".py":
-            continue
-
-        with open(path, "r") as file:
-            for line in file:
-                model_urls.update(MODEL_URL_PATTERN.findall(line))
+    for root in roots:
+        for path in pathlib.Path(root).rglob("*.py"):
+            with open(path, "r") as file:
+                for line in file:
+                    model_urls.update(MODEL_URL_PATTERN.findall(line))
 
     print("\n".join(sorted(model_urls)))
 
 
 if __name__ == "__main__":
-    main(sys.argv[1])
+    main(*sys.argv[1:])