admin-requests
admin-requests copied to clipboard
Default installation of JAX + Tensorflow produces inconsistent environment
Solution to issue cannot be found in the documentation.
- [X] I checked the documentation.
Issue
Installing default versions of both JAX and Tensorflow from conda-forge currently results in an inconsistent environment in which neither of them can be imported. Using either of the two separately is fine.
I can recreate this in the miniforge3 docker image like this
FROM condaforge/miniforge3
RUN conda install tensorflow jax
RUN conda list
RUN conda info
RUN python -c 'import tensorflow'
# this also fails
# RUN python -c 'import jax'
which gives the following error:
#7 [5/5] RUN python -c 'import tensorflow'
#7 sha256:c7f1bae8ee24e9614e6e17d5be2614442a2df455b7b3733c3f1e4b8b7a40132a
#7 2.331 Traceback (most recent call last):
#7 2.331 File "<string>", line 1, in <module>
#7 2.331 File "/opt/conda/lib/python3.9/site-packages/tensorflow/__init__.py", line 55, in <module>
#7 2.331 from ._api.v2 import compat
#7 2.331 File "/opt/conda/lib/python3.9/site-packages/tensorflow/_api/v2/compat/__init__.py", line 39, in <module>
#7 2.331 from . import v1
#7 2.331 File "/opt/conda/lib/python3.9/site-packages/tensorflow/_api/v2/compat/v1/__init__.py", line 34, in <module>
#7 2.331 from . import compat
#7 2.331 File "/opt/conda/lib/python3.9/site-packages/tensorflow/_api/v2/compat/v1/compat/__init__.py", line 39, in <module>
#7 2.333 from . import v1
#7 2.333 File "/opt/conda/lib/python3.9/site-packages/tensorflow/_api/v2/compat/v1/compat/v1/__init__.py", line 51, in <module>
#7 2.334 from tensorflow._api.v2.compat.v1 import lite
#7 2.334 File "/opt/conda/lib/python3.9/site-packages/tensorflow/_api/v2/compat/v1/lite/__init__.py", line 11, in <module>
#7 2.334 from . import experimental
#7 2.334 File "/opt/conda/lib/python3.9/site-packages/tensorflow/_api/v2/compat/v1/lite/experimental/__init__.py", line 10, in <module>
#7 2.334 from . import authoring
#7 2.334 File "/opt/conda/lib/python3.9/site-packages/tensorflow/_api/v2/compat/v1/lite/experimental/authoring/__init__.py", line 10, in <module>
#7 2.334 from tensorflow.lite.python.authoring.authoring import compatible
#7 2.334 File "/opt/conda/lib/python3.9/site-packages/tensorflow/lite/python/authoring/authoring.py", line 43, in <module>
#7 2.335 from tensorflow.lite.python import convert
#7 2.335 File "/opt/conda/lib/python3.9/site-packages/tensorflow/lite/python/convert.py", line 33, in <module>
#7 2.335 from tensorflow.lite.python import util
#7 2.335 File "/opt/conda/lib/python3.9/site-packages/tensorflow/lite/python/util.py", line 55, in <module>
#7 2.335 from jax import xla_computation as _xla_computation
#7 2.335 File "/opt/conda/lib/python3.9/site-packages/jax/__init__.py", line 37, in <module>
#7 2.335 from jax import config as _config_module
#7 2.335 File "/opt/conda/lib/python3.9/site-packages/jax/config.py", line 18, in <module>
#7 2.335 from jax._src.config import config
#7 2.335 File "/opt/conda/lib/python3.9/site-packages/jax/_src/config.py", line 27, in <module>
#7 2.335 from jax._src import lib
#7 2.335 File "/opt/conda/lib/python3.9/site-packages/jax/_src/lib/__init__.py", line 101, in <module>
#7 2.335 version = check_jaxlib_version(
#7 2.335 File "/opt/conda/lib/python3.9/site-packages/jax/_src/lib/__init__.py", line 90, in check_jaxlib_version
#7 2.336 raise RuntimeError(msg)
#7 2.336 RuntimeError: jaxlib is version 0.1.75, but this version of jax requires version >= 0.3.2.
#7 ERROR: executor failed running [/bin/sh -c python -c 'import tensorflow']: exit code: 1
Installed packages
#5 [3/5] RUN conda list
#5 sha256:1b4a2bef8e285ea7908ee767af6ddccd2b54750f9f686f79b9cdf637ae8dbacc
#5 1.412 # packages in environment at /opt/conda:
#5 1.412 #
#5 1.412 # Name Version Build Channel
#5 1.412 _libgcc_mutex 0.1 conda_forge conda-forge
#5 1.412 _openmp_mutex 4.5 1_gnu conda-forge
#5 1.412 abseil-cpp 20210324.2 h9c3ff4c_0 conda-forge
#5 1.412 absl-py 1.0.0 pyhd8ed1ab_0 conda-forge
#5 1.412 aiohttp 3.8.1 py39hb9d737c_1 conda-forge
#5 1.412 aiosignal 1.2.0 pyhd8ed1ab_0 conda-forge
#5 1.412 astunparse 1.6.3 pyhd8ed1ab_0 conda-forge
#5 1.412 async-timeout 4.0.2 pyhd8ed1ab_0 conda-forge
#5 1.412 attrs 21.4.0 pyhd8ed1ab_0 conda-forge
#5 1.412 blinker 1.4 py_1 conda-forge
#5 1.412 brotlipy 0.7.0 py39h3811e60_1003 conda-forge
#5 1.412 bzip2 1.0.8 h7f98852_4 conda-forge
#5 1.412 c-ares 1.18.1 h7f98852_0 conda-forge
#5 1.412 ca-certificates 2021.10.8 ha878542_0 conda-forge
#5 1.412 cached-property 1.5.2 hd8ed1ab_1 conda-forge
#5 1.412 cached_property 1.5.2 pyha770c72_1 conda-forge
#5 1.412 cachetools 4.2.4 pyhd8ed1ab_0 conda-forge
#5 1.412 certifi 2021.10.8 py39hf3d152e_2 conda-forge
#5 1.412 cffi 1.15.0 py39h4bc2ebd_0 conda-forge
#5 1.412 charset-normalizer 2.0.12 pyhd8ed1ab_0 conda-forge
#5 1.412 click 8.1.3 py39hf3d152e_0 conda-forge
#5 1.412 colorama 0.4.4 pyh9f0ad1d_0 conda-forge
#5 1.412 conda 4.12.0 py39hf3d152e_0 conda-forge
#5 1.412 conda-package-handling 1.8.0 py39hb9d737c_0 conda-forge
#5 1.412 cryptography 36.0.2 py39hd97740a_0 conda-forge
#5 1.412 frozenlist 1.3.0 py39hb9d737c_1 conda-forge
#5 1.412 gast 0.4.0 pyh9f0ad1d_0 conda-forge
#5 1.412 giflib 5.2.1 h36c2ea0_2 conda-forge
#5 1.412 google-auth 1.35.0 pyh6c4a22f_0 conda-forge
#5 1.412 google-auth-oauthlib 0.4.6 pyhd8ed1ab_0 conda-forge
#5 1.412 google-pasta 0.2.0 pyh8c360ce_0 conda-forge
#5 1.412 grpc-cpp 1.42.0 ha1441d3_1 conda-forge
#5 1.412 grpcio 1.42.0 py39hff7568b_0 conda-forge
#5 1.412 h5py 3.6.0 nompi_py39h7e08c79_100 conda-forge
#5 1.412 hdf5 1.12.1 nompi_h2386368_104 conda-forge
#5 1.412 icu 69.1 h9c3ff4c_0 conda-forge
#5 1.412 idna 3.3 pyhd8ed1ab_0 conda-forge
#5 1.412 importlib-metadata 4.11.3 py39hf3d152e_1 conda-forge
#5 1.412 jax 0.3.7 pyhd8ed1ab_0 conda-forge
#5 1.412 jaxlib 0.1.75 py39hde0f152_0 conda-forge
#5 1.412 jpeg 9e h166bdaf_1 conda-forge
#5 1.412 keras 2.7.0 pyhd8ed1ab_0 conda-forge
#5 1.412 keras-preprocessing 1.1.2 pyhd8ed1ab_0 conda-forge
#5 1.412 keyutils 1.6.1 h166bdaf_0 conda-forge
#5 1.412 krb5 1.19.3 h3790be6_0 conda-forge
#5 1.412 ld_impl_linux-64 2.36.1 hea4e1c9_2 conda-forge
#5 1.412 libblas 3.9.0 14_linux64_openblas conda-forge
#5 1.412 libcblas 3.9.0 14_linux64_openblas conda-forge
#5 1.412 libcurl 7.83.0 h7bff187_0 conda-forge
#5 1.412 libedit 3.1.20191231 he28a2e2_2 conda-forge
#5 1.412 libev 4.33 h516909a_1 conda-forge
#5 1.412 libffi 3.4.2 h7f98852_5 conda-forge
#5 1.412 libgcc-ng 11.2.0 h1d223b6_14 conda-forge
#5 1.412 libgfortran-ng 11.2.0 h69a702a_16 conda-forge
#5 1.412 libgfortran5 11.2.0 h5c6108e_16 conda-forge
#5 1.412 libgomp 11.2.0 h1d223b6_14 conda-forge
#5 1.412 liblapack 3.9.0 14_linux64_openblas conda-forge
#5 1.412 libnghttp2 1.47.0 h727a467_0 conda-forge
#5 1.412 libnsl 2.0.0 h7f98852_0 conda-forge
#5 1.412 libopenblas 0.3.20 pthreads_h78a6416_0 conda-forge
#5 1.412 libpng 1.6.37 h21135ba_2 conda-forge
#5 1.412 libprotobuf 3.19.4 h780b84a_0 conda-forge
#5 1.412 libssh2 1.10.0 ha56f1ee_2 conda-forge
#5 1.412 libstdcxx-ng 11.2.0 he4da1e4_16 conda-forge
#5 1.412 libuuid 2.32.1 h7f98852_1000 conda-forge
#5 1.412 libzlib 1.2.11 h36c2ea0_1013 conda-forge
#5 1.412 markdown 3.3.7 pyhd8ed1ab_0 conda-forge
#5 1.412 multidict 6.0.2 py39hb9d737c_1 conda-forge
#5 1.412 ncurses 6.3 h9c3ff4c_0 conda-forge
#5 1.412 numpy 1.22.3 py39hc58783e_2 conda-forge
#5 1.412 oauthlib 3.2.0 pyhd8ed1ab_0 conda-forge
#5 1.412 openssl 1.1.1o h166bdaf_0 conda-forge
#5 1.412 opt_einsum 3.3.0 pyhd8ed1ab_1 conda-forge
#5 1.412 pip 22.0.4 pyhd8ed1ab_0 conda-forge
#5 1.412 protobuf 3.19.4 py39he80948d_0 conda-forge
#5 1.412 pyasn1 0.4.8 py_0 conda-forge
#5 1.412 pyasn1-modules 0.2.7 py_0 conda-forge
#5 1.412 pycosat 0.6.3 py39h3811e60_1009 conda-forge
#5 1.412 pycparser 2.21 pyhd8ed1ab_0 conda-forge
#5 1.412 pyjwt 2.3.0 pyhd8ed1ab_1 conda-forge
#5 1.412 pyopenssl 22.0.0 pyhd8ed1ab_0 conda-forge
#5 1.412 pysocks 1.7.1 py39hf3d152e_4 conda-forge
#5 1.412 python 3.9.10 h85951f9_2_cpython conda-forge
#5 1.412 python-flatbuffers 2.0 pyhd8ed1ab_0 conda-forge
#5 1.412 python_abi 3.9 2_cp39 conda-forge
#5 1.412 pyu2f 0.1.5 pyhd8ed1ab_0 conda-forge
#5 1.412 re2 2021.11.01 h9c3ff4c_0 conda-forge
#5 1.412 readline 8.1 h46c0cb4_0 conda-forge
#5 1.412 requests 2.27.1 pyhd8ed1ab_0 conda-forge
#5 1.412 requests-oauthlib 1.3.1 pyhd8ed1ab_0 conda-forge
#5 1.412 rsa 4.8 pyhd8ed1ab_0 conda-forge
#5 1.412 ruamel_yaml 0.15.80 py39h3811e60_1006 conda-forge
#5 1.412 scipy 1.8.0 py39hee8e79c_1 conda-forge
#5 1.412 setuptools 60.10.0 py39hf3d152e_0 conda-forge
#5 1.412 six 1.16.0 pyh6c4a22f_0 conda-forge
#5 1.412 snappy 1.1.8 he1b5a44_3 conda-forge
#5 1.412 sqlite 3.37.1 h4ff8645_0 conda-forge
#5 1.412 tensorboard 2.6.0 pyhd8ed1ab_1 conda-forge
#5 1.412 tensorboard-data-server 0.6.0 py39hd97740a_2 conda-forge
#5 1.412 tensorboard-plugin-wit 1.8.1 pyhd8ed1ab_0 conda-forge
#5 1.412 tensorflow 2.7.0 cpu_py39h4655687_0 conda-forge
#5 1.412 tensorflow-base 2.7.0 cpu_py39hf4995fd_0 conda-forge
#5 1.412 tensorflow-estimator 2.7.0 cpu_py39ha241409_0 conda-forge
#5 1.412 termcolor 1.1.0 py_2 conda-forge
#5 1.412 tk 8.6.12 h27826a3_0 conda-forge
#5 1.412 tqdm 4.63.0 pyhd8ed1ab_0 conda-forge
#5 1.412 typing-extensions 4.2.0 hd8ed1ab_1 conda-forge
#5 1.412 typing_extensions 4.2.0 pyha770c72_1 conda-forge
#5 1.412 tzdata 2022a h191b570_0 conda-forge
#5 1.412 urllib3 1.26.9 pyhd8ed1ab_0 conda-forge
#5 1.412 werkzeug 2.1.2 pyhd8ed1ab_1 conda-forge
#5 1.412 wheel 0.37.1 pyhd8ed1ab_0 conda-forge
#5 1.412 wrapt 1.14.1 py39hb9d737c_0 conda-forge
#5 1.412 xz 5.2.5 h516909a_1 conda-forge
#5 1.412 yaml 0.2.5 h7f98852_2 conda-forge
#5 1.412 yarl 1.7.2 py39hb9d737c_2 conda-forge
#5 1.412 zipp 3.8.0 pyhd8ed1ab_0 conda-forge
#5 1.412 zlib 1.2.11 h36c2ea0_1013 conda-forge
Environment info
#6 [4/5] RUN conda info
#6 sha256:e04dec7bc9e71b96d819465250e444e4080f1874b5342685a25c0744c32e910d
#6 0.730
#6 0.730 active environment : None
#6 0.730 user config file : /root/.condarc
#6 0.730 populated config files : /opt/conda/.condarc
#6 0.730 conda version : 4.12.0
#6 0.730 conda-build version : not installed
#6 0.730 python version : 3.9.10.final.0
#6 0.730 virtual packages : __linux=5.10.47=0
#6 0.730 __glibc=2.31=0
#6 0.730 __unix=0=0
#6 0.730 __archspec=1=x86_64
#6 0.730 base environment : /opt/conda (writable)
#6 0.730 conda av data dir : /opt/conda/etc/conda
#6 0.730 conda av metadata url : None
#6 0.730 channel URLs : https://conda.anaconda.org/conda-forge/linux-64
#6 0.730 https://conda.anaconda.org/conda-forge/noarch
#6 0.730 package cache : /opt/conda/pkgs
#6 0.730 /root/.conda/pkgs
#6 0.730 envs directories : /opt/conda/envs
#6 0.730 /root/.conda/envs
#6 0.730 platform : linux-64
#6 0.730 user-agent : conda/4.12.0 requests/2.27.1 CPython/3.9.10 Linux/5.10.47-linuxkit ubuntu/20.04.4 glibc/2.31
#6 0.730 UID:GID : 0:0
#6 0.730 netrc file : None
#6 0.730 offline mode : False
#6 0.730
I can recreate. Thank you for the detailed information.
I think that:
- The current
jax
feedstock should add constraints to thejaxlib
dependency and write a repopatch to patch existing builds. cc @conda-forge/jax @conda-forge/jaxlib - We are aware that the tensorflow feedstock has become incompatible with the latest packages from conda-forge. A few PRs are trying to alleviate that, but it might require more work due to
abseil
library mismatches.
If you could:
- Try to manually specify the versions you want. You will have to downgrade the version of
jax
you want is my guess.
mamba install "tensorflow>=2.7.0" "jax=x.x.x" "jaxlib=y.y.y."
that would help alleviate things
2. Try to submit repo-data patches to require the appropriate versions of jaxlib
for jax
.
3. Submit a pull request to jax
fixing the dependnecy for future packages.
I believe this should be resolved now, no? At least the current packages are compatible. A repodata patch may be worthy tho
If I add RUN conda update conda
to the Dockerfile I described above then this is indeed fixed.
If I don't do this then the default base image, condaforge/miniforge3:latest
is still inconsistent.
If I add
RUN conda update conda
to the Dockerfile I described above then this is indeed fixed.If I don't do this then the default base image,
condaforge/miniforge3:latest
is still inconsistent.
Okay, I see.
Btw, is this for a GPU or CPU deployment? I highly recommend you follow hmaarrfk's recommendation above. Our jaxlib>=0.3.10 (thus jax>=0.3.10) and tensorflow>=2.8.1 are somewhat different from previous versions (enabling XLA for tf and adding cuda for jax). Also, both are slightly more advanced in terms of migrations (they're incompatible with things like gdal, etc. as the latter haven't been migrated to the latest icu).
I've only tested on CPU.
I was indeed able to work around things with @hmaarrfk's suggestions above, so thanks for that!
Who is in charge of condaforge/miniforge3:latest? We can just bump that up?
Re CPU vs GPU, I am asking because our current recommendation is that users specify *cuda*
explicitly, e.g. tensorflow>=2.8.1=*cuda*
and jaxlib>=0.3.10=*cuda*
. In case you need that info in the future :)