From 13f2e602c19bcdf3adad8595b455d9106026241c Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Tue, 25 Aug 2026 20:08:10 -0700 Subject: [PATCH] fix(tests): download the shared torchvision checkpoints before pytest forks ts-api is red on Windows on most main runs: FAILED api/test_operator_fallback.py::TestFallbackModels::test_fallback_mobilenet_v2 _pickle.UnpicklingError: pickle data was truncated pytest runs one test file per xdist worker (addopts carries --dist=loadfile), and api/test_operator_fallback.py and api/test_module_fallback.py both open with models.mobilenet_v2(pretrained=True). Two workers therefore download the same checkpoint at the same time. torch.hub finishes a download with shutil.move, which calls os.rename and falls back to a plain copy when that raises. On Linux os.rename over an existing file succeeds and is atomic; on Windows it raises FileExistsError, so the second worker copies over the file the first one is reading and torch.load sees a truncated pickle. That is why this only ever fails on Windows, and why the config that fails moves around between runs. resnet18 is in the same position: three files in api, two in models and two in integrations. The ts suites already run tests/modules/hub.py as a setup step to put what they need on disk before pytest starts. It just does not cover the torchvision checkpoints, so fetch those there too. One process, no concurrent writers, and the workers only ever read. Checkpoints only one file uses cannot race, so they are left alone. --- tests/modules/hub.py | 43 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 43 insertions(+) diff --git a/tests/modules/hub.py b/tests/modules/hub.py index e810df0404..562f5e04f7 100644 --- a/tests/modules/hub.py +++ b/tests/modules/hub.py @@ -39,6 +39,47 @@ # "bert_base_uncased": {"model": cm.BertModule(), "path": "trace"}, } +# torchvision checkpoints that more than one test file in the same suite asks for +# with pretrained=True: resnet18 in tests/py/ts/{api,models,integrations} and +# mobilenet_v2 in tests/py/ts/api. Checkpoints that only one file uses cannot be +# downloaded twice at once, so they are left to that file. IMAGENET1K_V1 is the +# set of weights pretrained=True selects. +PRETRAINED_WEIGHTS = ( + "MobileNet_V2_Weights", + "ResNet18_Weights", +) + + +def download_pretrained_weights(): + """Fetch the shared torchvision checkpoints before pytest forks its workers. + + pytest runs one test file per xdist worker (--dist=loadfile), and several + files in the same suite ask torchvision for the same checkpoint, so two + workers can download one URL at the same time. torch.hub finishes a download + with shutil.move, and on Windows that is a plain copy whenever the + destination already exists, so the second worker rewrites the file while the + first one is reading it and torch.load fails on truncated pickle data. + Downloading here, in one process before the workers start, leaves them with + nothing to do but read. + """ + if importlib.util.find_spec("torchvision") is None: + print("torchvision is not installed, skipping pretrained weight download") + return + + import torchvision.models as tv_models + + for name in PRETRAINED_WEIGHTS: + weights = getattr(tv_models, name).IMAGENET1K_V1 + print("Downloading {}".format(weights)) + try: + # check_hash is what torchvision itself passes, and it is the only + # thing that stops a bad download from being cached and reused. + weights.get_state_dict(progress=False, check_hash=True) + except Exception as e: + # Warming the cache is an optimization, so a failure here must not + # take the whole suite down. The tests still fetch what they need. + print("Could not pre-download {}: {}".format(name, e)) + def get(n, m, manifest): print("Downloading {}".format(n)) @@ -93,6 +134,8 @@ def download_models(version_matches, manifest): def main(): + download_pretrained_weights() + manifest = None version_matches = False manifest_exists = False