pip installing jaxlib 0.1.63 fails on Mac OS

[gecrooks]

Upgrading from jaxlib 0.1.62 to 0.1.63 causes import of jax to fail with mysterious error messages. Downgrading back to 0.1.62 and everything is happy again. (This in on Mac OS)

Requirement already satisfied: jaxlib in /Users/work/miniconda3/envs/QPACK/lib/python3.8/site-packages (0.1.62)
Collecting jaxlib
  Using cached jaxlib-0.1.63-cp38-none-macosx_10_9_x86_64.whl (42.7 MB)
Requirement already satisfied: absl-py in /Users/work/miniconda3/envs/QPACK/lib/python3.8/site-packages (from jaxlib) (0.12.0)
Requirement already satisfied: numpy>=1.16 in /Users/work/miniconda3/envs/QPACK/lib/python3.8/site-packages (from jaxlib) (1.20.1)
Requirement already satisfied: flatbuffers in /Users/work/miniconda3/envs/QPACK/lib/python3.8/site-packages (from jaxlib) (20210203211748)
Requirement already satisfied: scipy in /Users/work/miniconda3/envs/QPACK/lib/python3.8/site-packages (from jaxlib) (1.6.1)
Requirement already satisfied: six in /Users/work/miniconda3/envs/QPACK/lib/python3.8/site-packages (from absl-py->jaxlib) (1.15.0)
Installing collected packages: jaxlib
  Attempting uninstall: jaxlib
    Found existing installation: jaxlib 0.1.62
    Uninstalling jaxlib-0.1.62:
      Successfully uninstalled jaxlib-0.1.62
Successfully installed jaxlib-0.1.63

$ python -c "import jax"
Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File "/Users/work/miniconda3/envs/QPACK/lib/python3.8/site-packages/jax/__init__.py", line 21, in <module>
    from .config import config
  File "/Users/work/miniconda3/envs/QPACK/lib/python3.8/site-packages/jax/config.py", line 19, in <module>
    from jax import lib
  File "/Users/work/miniconda3/envs/QPACK/lib/python3.8/site-packages/jax/lib/__init__.py", line 53, in <module>
    from jaxlib import xla_client
  File "/Users/work/miniconda3/envs/QPACK/lib/python3.8/site-packages/jaxlib/xla_client.py", line 31, in <module>
    from . import xla_extension as _xla
ImportError: dlopen(/Users/work/miniconda3/envs/QPACK/lib/python3.8/site-packages/jaxlib/xla_extension.so, 2): Symbol not found: __ZN10tensorflow8profiler11MonitorGrpcERKNSt3__112basic_stringIcNS1_11char_traitsIcEENS1_9allocatorIcEEEERKNS_14MonitorRequestEPNS_15MonitorResponseE
  Referenced from: /Users/work/miniconda3/envs/QPACK/lib/python3.8/site-packages/jaxlib/xla_extension.so
  Expected in: flat namespace
 in /Users/work/miniconda3/envs/QPACK/lib/python3.8/site-packages/jaxlib/xla_extension.so

[jakevdp]

Thanks for the report - I can reproduce this on OSX 10.15. Installation on ubuntu is fine, so it appears to be specific to the OSX wheels.

@skye can you take a look?

[skye]

Oof, after lovingly fixing this error and building new release wheels, I uploaded the old broken wheels! 🤦‍♀️ I don't think pypi lets you upload new wheels for a given release, but I missed something else in the 0.1.63 release anyway so am gonna cut a 0.1.64 soon (like tomorrow). So I recommend sticking to 0.1.62 for now, and upgrading to 0.1.64 once it's available (I can close this issue when that happens).

[matthewfeickert]

Thanks very much for the quick diagnosis, @skye! Super appreciated, and I think everyone's been there before with a PyPI upload mishap. :)

I don't think pypi lets you upload new wheels for a given release, but I missed something else in the 0.1.63 release anyway so am gonna cut a 0.1.64 soon (like tomorrow).

No, they don't, but can you please also yank jaxlib v0.1.63 so that it can't cause problems in the future (people can still force installation of the release if they want by using == or ===)? It will also fix everyone's problems with CI immediately.

I'm also don't know how Google operates when it comes to the building and releasing of wheels, but unless there is some bespoke in-house workflow and procedure I'd highly recommend cibuildwheel for fast and easy multi-platform builds that also makes validation and uploading to PyPI very simple. c.f. boost-histogram's wheel list for an excellent example of a project that really takes full advantage of this. 🚀 cc @henryiii

[henryiii]

I'd also list some happy users: https://cibuildwheel.readthedocs.io/en/stable/working-examples/ :)

You can do a post1 release, perhaps?

[skye]

I hadn't heard of yanking a release (or post1 releases, or cibuildwheel, for that matter), thanks for the suggestions! I just yanked 0.1.63 and am working on a 0.1.64 release now.

I think @hawkinsp set up most of our Python packaging and to be honest I don't know much about Python packaging in general, so it's hard for me to say if cibuildwheel makes sense for us. Our builds are really huge (especially the GPU wheels), so we'd ideally have a way to run on the beefy GCP VMs we have access to, but even fully-automated really slow builds might be better than our semi-automated fast-ish builds.

Stupid beginner Q: how does cibuildwheel know how to build the wheel? I don't see a build script or similar config. Maybe you can point me to an example in one of the repos that use it?

[matthewfeickert]

I hadn't heard of yanking a release (or post1 releases, or cibuildwheel, for that matter), thanks for the suggestions! I just yanked 0.1.63 and am working on a 0.1.64 release now.

Awesome! That means everyone's CI is working again. 🎉 Yeah, PEP 592 is still not super widely known about in general it seems, but very helpful!

I think @hawkinsp set up most of our Python packaging and to be honest I don't know much about Python packaging in general, so it's hard for me to say if cibuildwheel makes sense for us. Our builds are really huge (especially the GPU wheels), so we'd ideally have a way to run on the beefy GCP VMs we have access to, but even fully-automated really slow builds might be better than our semi-automated fast-ish builds.

Stupid beginner Q: how does cibuildwheel know how to build the wheel? I don't see a build script or similar config. Maybe you can point me to an example in one of the repos that use it?

I'll let the cibuildwheel experts (@joerick, @YannickJadoul, @henryiii, @Czaki) give you actual good advice, but wrt building the wheels that's all configured through option variables that you set. All the repos that @henryiii linked above have in them a YAML file (probably named wheels.yml) that contains a workflow. Here's two example workflows from that list:

• scikit-learn using GHA
• boost-histogram using GHA

[henryiii]

You can use cibuildwheel on custom hardware, too. The only reason it's not recommended for local use is that it installs Python into global locations on macOS and Windows, but any virtualized or runner hardware should be absolutely fine. If it's not a known CI, you might want to set --platform <osname> or export CI=1, otherwise it will think you are running locally and will require one of the two above methods to ensure you know what you are doing. The one requirement is that your package builds with python -m pip wheel ., but you can use CIBW_BEFORE_BUILD and CIBW_BEFORE_ALL to control any sort of setup, other building, etc. that you need.

Personally I recommend getting as close to being able to install from SDist as possible and using pyproject.toml from PEP 517/518. Then it should work "out of the box" except for whatever you were not able to do as part of that process, which you then can put in the above mentioned variables (say as a script or install commands, etc).

[skye]

The one requirement is that your package builds with python -m pip wheel .

I think this is the key piece I didn't know about :) We're not set up to build that way. Instead, we have a build script that configures bazel (e.g. whether to link in CUDA support, where CUDA is installed, etc.), then uses bazel to build the required C++ dependencies, and finally invokes another script as part of the same BUILD rule that uses setup.py bdist_wheel to bundle everything up into a wheel.

I still don't really know how things needs to be structured such that python -m pip wheel . works, but my vague understanding is that we need to refactor setup.py to drive everything (setting up bazel etc.) instead of the other way around. Or something? Is there a guide somewhere to setting all this up? I've only been able to find various reference documentation that seems to assume context I don't have :\ If I spent more time digging I could presumably figure it out eventually though.

BTW, y'all have been so kind and helpful! So even if this thread isn't worth pulling now, thanks for all the tips and advice so far.

[hawkinsp]

One other comment is that we need a very unusual toolchain to build C++14 code that is manylinux2010 compliant. And we do that in an unusual Docker container. So I don't hold great hopes of make the setup.py route work, but I'm willing to be convinced otherwise :-)

[Czaki]

You could see this project https://github.com/czaki/imagecodecs_build
I use in this project prebuild docker images with many external libraries.

And cibuildwheel use manylinux compatible containers during build linux wheel.

[skye]

Oops, forgot to close this issue, but jaxlib 0.1.64 has been available for about a week and shouldn't have this issue!