🐛[BUG]: Download API issues

[adamjstewart]

Version

0.1.0

On which installation method(s) does this occur?

Source

Describe the issue

The download API is a bit finicky. First, when I use earth2mip.networks.get_model(...) with a remote model name, it takes a long time to instantiate (presumably because it is downloading a model checkpoint). However, there is no progress bar (tqdm would help here). Because I didn't know what was happening, I Ctrl+C'ed the process and tried again. Every subsequent time, I get the following error:

>>> m = get_model("e2mip://pangu", device="cuda:8")
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/adam/spack/opt/spack/linux-ubuntu22.04-zen2/gcc-12.3.0/py-earth2mip-0.1.0-eg37aficdacia2bylnpyisnznioeuoo2/lib/python3.10/site-packages/earth2mip/networks/__init__.py", line 345, in get_model
    return _load_package_builtin(package, device, name=url.netloc)
  File "/home/adam/spack/opt/spack/linux-ubuntu22.04-zen2/gcc-12.3.0/py-earth2mip-0.1.0-eg37aficdacia2bylnpyisnznioeuoo2/lib/python3.10/site-packages/earth2mip/networks/__init__.py", line 291, in _load_package_builtin
    return inference_loader(package, device=device)
  File "/home/adam/spack/opt/spack/linux-ubuntu22.04-zen2/gcc-12.3.0/py-earth2mip-0.1.0-eg37aficdacia2bylnpyisnznioeuoo2/lib/python3.10/site-packages/earth2mip/networks/pangu.py", line 266, in load
    model_6 = PanguStacked(PanguWeather(p6))
  File "/home/adam/spack/opt/spack/linux-ubuntu22.04-zen2/gcc-12.3.0/py-earth2mip-0.1.0-eg37aficdacia2bylnpyisnznioeuoo2/lib/python3.10/site-packages/earth2mip/networks/pangu.py", line 74, in __init__
    os.stat(self.path)
FileNotFoundError: [Errno 2] No such file or directory: '/home/adam/.cache/earth2mip/models/pangu/pangu_weather_6.onnx'

Ideally, this should just redownload the model instead of crashing just because the model hasn't been downloaded. Current workaround is to delete this model directory and try again.

Environment details

Linux Ubuntu 22.04

[nbren12]

Thanks for the report. I've done this myself many times.

Indeed the caching logic is not very robust right now. When writing to the cache we could write e.g. to .cache/earth2mip/models/pangu/pangu_weather_6.onnx.partial or similar, and then shutil.move to the final location. This will make the cache operations atomic, and avoid leaving it in a broken state.

And yes, a progress bar would be helpful too. Pangu is very slow to download.