jaxlib 0.1.61 without CUDA support not available in storage.googleapis.com

[sergiuiacob1]

A while ago some of us were depending on jaxlib 0.1.59, but that got removed from https://pypi.org/project/jaxlib/#history so the workaround was to upgrade jaxlib to 0.1.61.

Since upgrading is not always desirable/possible (i.e. dependencies, backwards compatibility etc), we switched to installing jaxlib 0.1.61 from https://storage.googleapis.com/jax-releases/jax_releases.html instead of https://pypi.org/project/jaxlib/#history in case 0.1.61 gets removed from pypi as well.

However, we're getting lots of warnings when loading jax:

2022-02-09 03:52:23.782602: W external/org_tensorflow/tensorflow/stream_executor/platform/default/dso_loader.cc:60] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /backend/erts-11.2.2.6/lib:

Is there a way to add 0.1.61 for (linux x86_64, linux arm, macos x86_64, macos arm) to https://storage.googleapis.com/ as well?

[hawkinsp]

What pip command did you run to install jaxlib?

If you look in the index, there are nocuda wheels there, e.g.,
https://storage.googleapis.com/jax-releases/nocuda/jaxlib-0.1.61-cp39-none-manylinux2010_x86_64.whl
and in general we are archiving releases in that GCS bucket going forward.

I should note we have never released Linux ARM wheels, and while we now release Mac ARM I don't think we did that far back.

[sergiuiacob1]

What pip command did you run to install jaxlib?

If you look in the index, there are nocuda wheels there, e.g., https://storage.googleapis.com/jax-releases/nocuda/jaxlib-0.1.61-cp39-none-manylinux2010_x86_64.whl and in general we are archiving releases in that GCS bucket going forward.

I should note we have never released Linux ARM wheels, and while we now release Mac ARM I don't think we did that far back.

I'm using pip install -f https://storage.googleapis.com/jax-releases/jax_releases.html jaxlib==0.1.61

[hawkinsp]

Ah, I think I see the problem. We use local version tags, e.g. 0.1.61+cuda112 to distinguish the CUDA wheel from the nocuda wheel, e.g., 0.1.61. Unfortunately pip accepts both versions as matching the constraint ==0.1.61, so you'll randomly get either of them if you install that way.

For now, I'd suggest installing a specific wheel file, e.g.,:

pip install https://storage.googleapis.com/jax-releases/nocuda/jaxlib-0.1.61-cp37-none-manylinux2010_x86_64.whl

For the future, we're planning to restructure how our jaxlib wheels are structured so the CUDA wheel is a plugin for the non-cuda wheel. At that point we won't need this local versioning scheme any more. Admittedly that doesn't help for older jaxlib versions; the only other thing we could do there is generate an index .html file that contains only the nocuda packages that you could point pip at.

Does the workaround of installing a specific wheel file work for you? (Ultimately we'd prefer you upgraded, of course...)

[sergiuiacob1]

I see. Correct me if I'm wrong, but in this case the only issue with the above solution is that the command would only work with linux systems, right?
My constraint is that I have the same requirements.txt file for multiple platforms (linux intel, linux arm, macos intel, macos arm).

Also, do you have any idea if and when 0.1.61 would be removed from PyPI? Is that still a concern?

[hawkinsp]

Correct, if you use a specific file URL, it will work only for a particular platform and python version.

The reason we remove wheels from pypi is because of space limits, so the answer is "when we run out of space". I would advocate you upgrade sooner than that.

As a temporary measure, I suggest we add another package index file, with the intention of removing it again at some point when we have separated the CUDA parts of jaxlib into their own distinct plugin package.

[hawkinsp]

@yashk2810 Could we, as part of our release process, generate a separate nocuda package index that contains only the nocuda wheels?

[yashk2810]

Done.

Cuda index file: https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

All others: https://storage.googleapis.com/jax-releases/jax_releases.html