π[BUG]: Download API issues
adamjstewart opened this issue Β· comments
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
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.