Fix race condition in maybe_download (#846)

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:

```python
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]
```

Author: kvablack

src/openpi/shared/download.py

@@ -81,18 +81,19 @@ def maybe_download(url: str, *, force_download: bool = False, **kwargs) -> pathl
                 else:
                     local_path.unlink()
 
-            # Download the data to a local cache.
-            logger.info(f"Downloading {url} to {local_path}")
-            scratch_path = local_path.with_suffix(".partial")
-            # Route openpi-assets through gsutil to avoid gcsfs auth issues with this bucket.
-            # All other gs:// URLs (e.g. big_vision) continue to use gcsfs as normal.
-            if parsed.scheme == "gs" and parsed.netloc == "openpi-assets":
-                _download_gsutil(url, scratch_path, **kwargs)
-            else:
-                _download_fsspec(url, scratch_path, **kwargs)
-
-            shutil.move(scratch_path, local_path)
-            _ensure_permissions(local_path)
+            if not local_path.exists():
+                # Download the data to a local cache.
+                logger.info(f"Downloading {url} to {local_path}")
+                scratch_path = local_path.with_suffix(".partial")
+                # Route openpi-assets through gsutil to avoid gcsfs auth issues with this bucket.
+                # All other gs:// URLs (e.g. big_vision) continue to use gcsfs as normal.
+                if parsed.scheme == "gs" and parsed.netloc == "openpi-assets":
+                    _download_gsutil(url, scratch_path, **kwargs)
+                else:
+                    _download_fsspec(url, scratch_path, **kwargs)
+
+                shutil.move(scratch_path, local_path)
+                _ensure_permissions(local_path)
 
     except PermissionError as e:
         msg = (