Fix race condition in maybe_download

[abhaybd]

Fixes #845.

maybe_download only checks if the file exists before acquiring the file lock, so multiple threads could pass the existence check concurrently, and then each will sequentially download the same file, resulting in redundant downloads. This PR adds an existence check after acquiring the file lock, so if multiple threads are concurrently downloading the same file, only one of them will actually perform the download and the others will shortcut.

This is the code to reproduce:

from concurrent.futures import ThreadPoolExecutor
import logging

from openpi.shared.download import maybe_download


logging.basicConfig(level=logging.INFO)

def main():
    n_workers = 2
    with ThreadPoolExecutor(max_workers=n_workers) as executor:
        futures = []
        for _ in range(n_workers):
            futures.append(executor.submit(maybe_download, "gs://openpi-assets/checkpoints/pi05_droid/params/ocdbt.process_0/d/b862fdd01b4477d01ea1f43ddb38b88f"))
        for future in futures:
            future.result()


if __name__ == "__main__":
    main()

Before the fix, this was the output:

INFO:openpi.shared.download:Downloading gs://openpi-assets/checkpoints/pi05_droid/params/ocdbt.process_0/d/b862fdd01b4477d01ea1f43ddb38b88f to /root/.cache/openpi/openpi-assets/checkpoints/pi05_droid/params/ocdbt.process_0/d/b862fdd01b4477d01ea1f43ddb38b88f
100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2.61G/2.61G [00:28<00:00, 99.9MiB/s]
INFO:openpi.shared.download:Downloading gs://openpi-assets/checkpoints/pi05_droid/params/ocdbt.process_0/d/b862fdd01b4477d01ea1f43ddb38b88f to /root/.cache/openpi/openpi-assets/checkpoints/pi05_droid/params/ocdbt.process_0/d/b862fdd01b4477d01ea1f43ddb38b88f
100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2.61G/2.61G [00:26<00:00, 108MiB/s]

With the fix, this is the output:

INFO:openpi.shared.download:Downloading gs://openpi-assets/checkpoints/pi05_droid/params/ocdbt.process_0/d/b862fdd01b4477d01ea1f43ddb38b88f to /root/.cache/openpi/openpi-assets/checkpoints/pi05_droid/params/ocdbt.process_0/d/b862fdd01b4477d01ea1f43ddb38b88f
100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2.61G/2.61G [00:26<00:00, 108MiB/s]