New python venv fails to pip install the flax mnist example requirements

[sycamoreoak]

New python venv fails to pip install the mnist example requirements

System information

• OS Platform and Distribution: Linux Ubuntu 22.04.4 LTS
• Flax, jax, jaxlib versions (obtain with pip show flax jax jaxlib:

(mnenv) ~/src/python/flax/examples/mnist$ pip show flax jax jaxlib
WARNING: Package(s) not found: flax, jax, jaxlib

• Python version: Python 3.10.12
• GPU/TPU model and memory: RTX 4070
• CUDA version (if applicable):

 nvcc --version
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2023 NVIDIA Corporation
Built on Tue_Jun_13_19:16:58_PDT_2023
Cuda compilation tools, release 12.2, V12.2.91
Build cuda_12.2.r12.2/compiler.32965470_0

Problem you have encountered:

The pip install requirements fails for the mnist example

~/src/python/flax/examples/mnist$ python3 -m venv mnenv
~/src/python/flax/examples/mnist$ source ./mnenv/bin/activate
(mnenv) ~/src/python/flax/examples/mnist$ pip install -r ./requirements.txt Looking in links: https://storage.googleapis.com/jax-releases/jax_releases.html
DEPRECATION: The HTML index page being used (https://storage.googleapis.com/jax-releases/jax_releases.html) is not a proper HTML 5 document. This is in violation of PEP 503 which requires these pages to be well-formed HTML 5 documents. Please reach out to the owners of this index page, and ask them to update this index page to a valid HTML 5 document. pip 22.2 will enforce this behaviour change. Discussion can be found at https://github.com/pypa/pip/issues/10825
Collecting absl-py==1.0.0
  Using cached absl_py-1.0.0-py3-none-any.whl (126 kB)
Collecting clu==0.0.6
  Using cached clu-0.0.6-py3-none-any.whl (77 kB)
Collecting flax==0.4.1
  Using cached flax-0.4.1-py3-none-any.whl (184 kB)
Collecting jax==0.3.4
  Using cached jax-0.3.4.tar.gz (924 kB)
  Preparing metadata (setup.py) ... done
ERROR: Could not find a version that satisfies the requirement jaxlib==0.3.2+cuda11.cudnn82 (from versions: 0.1.75, 0.1.76, 0.3.0, 0.3.2, 0.3.5, 0.3.7, 0.3.8, 0.3.10, 0.3.14, 0.3.15, 0.3.18, 0.3.20, 0.3.22, 0.3.24, 0.3.25, 0.4.1, 0.4.2, 0.4.3, 0.4.4, 0.4.6, 0.4.7, 0.4.9, 0.4.10, 0.4.11, 0.4.12, 0.4.13, 0.4.14, 0.4.16, 0.4.17, 0.4.18, 0.4.19, 0.4.20, 0.4.21, 0.4.22, 0.4.23, 0.4.24, 0.4.25, 0.4.26, 0.4.27, 0.4.28, 0.4.29)
ERROR: No matching distribution found for jaxlib==0.3.2+cuda11.cudnn82

What you expected to happen:

mnist, make it so.

Logs, error messages, etc:

ERROR: Could not find a version that satisfies the requirement jaxlib==0.3.2+cuda11.cudnn82 (from versions: 0.1.75, 0.1.76, 0.3.0, 0.3.2, 0.3.5, 0.3.7, 0.3.8, 0.3.10, 0.3.14, 0.3.15, 0.3.18, 0.3.20, 0.3.22, 0.3.24, 0.3.25, 0.4.1, 0.4.2, 0.4.3, 0.4.4, 0.4.6, 0.4.7, 0.4.9, 0.4.10, 0.4.11, 0.4.12, 0.4.13, 0.4.14, 0.4.16, 0.4.17, 0.4.18, 0.4.19, 0.4.20, 0.4.21, 0.4.22, 0.4.23, 0.4.24, 0.4.25, 0.4.26, 0.4.27, 0.4.28, 0.4.29)

Steps to reproduce:

Seems like a fresh venv on an up-to-date stable ubuntu with an nvidia GPU should reproduce this...

[sycamoreoak]

Following the below steps:

python -m venv env_imagenet
source env_imagenet/bin/activate
unset LD_LIBRARY_PATH
pip install -U "jax[cuda12]"
cd flax
pip install -e .
cd examples/imagenet/
pip install -r requirements.txt

here is my pip list after all of this:

Package                      Version
---------------------------- -----------
absl-py                      1.0.0
astunparse                   1.6.3
attrs                        23.2.0
cachetools                   5.3.3
certifi                      2024.6.2
charset-normalizer           3.3.2
chex                         0.1.7
clu                          0.0.6
contextlib2                  21.6.0
contourpy                    1.2.1
cycler                       0.12.1
dill                         0.3.8
dm-tree                      0.1.8
etils                        1.7.0
flatbuffers                  24.3.25
flax                         0.6.5
fonttools                    4.53.0
fsspec                       2024.6.0
future                       1.0.0
gast                         0.4.0
google-auth                  2.30.0
google-auth-oauthlib         0.4.6
google-pasta                 0.2.0
googleapis-common-protos     1.63.1
grpcio                       1.64.1
h5py                         3.11.0
idna                         3.7
importlib_resources          6.4.0
jax                          0.4.29
jax-cuda12-pjrt              0.4.29
jax-cuda12-plugin            0.4.29
jaxlib                       0.4.29
keras                        2.11.0
kiwisolver                   1.4.5
libclang                     18.1.1
Markdown                     3.6
markdown-it-py               3.0.0
MarkupSafe                   2.1.5
matplotlib                   3.8.4
mdurl                        0.1.2
ml-collections               0.1.0
ml-dtypes                    0.4.0
msgpack                      1.0.8
nest-asyncio                 1.6.0
numpy                        1.22.0
nvidia-cublas-cu12           12.5.2.13
nvidia-cuda-cupti-cu12       12.5.39
nvidia-cuda-nvcc-cu12        12.5.40
nvidia-cuda-runtime-cu12     12.5.39
nvidia-cudnn-cu12            9.1.1.17
nvidia-cufft-cu12            11.2.3.18
nvidia-cusolver-cu12         11.6.2.40
nvidia-cusparse-cu12         12.4.1.24
nvidia-nccl-cu12             2.21.5
nvidia-nvjitlink-cu12        12.5.40
oauthlib                     3.2.2
opt-einsum                   3.3.0
optax                        0.1.3
orbax                        0.1.9
orbax-checkpoint             0.5.16
packaging                    24.1
pillow                       10.3.0
pip                          22.0.2
promise                      2.3
protobuf                     3.19.6
pyasn1                       0.6.0
pyasn1_modules               0.4.0
Pygments                     2.18.0
pyparsing                    3.1.2
python-dateutil              2.9.0.post0
PyYAML                       6.0.1
requests                     2.32.3
requests-oauthlib            2.0.0
rich                         13.7.1
rsa                          4.9
scipy                        1.11.4
setuptools                   59.6.0
six                          1.16.0
tensorboard                  2.11.2
tensorboard-data-server      0.6.1
tensorboard-plugin-wit       1.8.1
tensorflow                   2.11.1
tensorflow-datasets          4.4.0
tensorflow-estimator         2.11.0
tensorflow-io-gcs-filesystem 0.37.0
tensorflow-metadata          1.13.0
tensorstore                  0.1.61
termcolor                    2.4.0
toolz                        0.12.1
tqdm                         4.66.4
typing_extensions            4.12.2
urllib3                      2.2.1
Werkzeug                     3.0.3
wheel                        0.43.0
wrapt                        1.16.0
zipp                         3.19.2

and now i find:

python main.py --workdir=./imagenet --config=configs/default.py
2024-06-13 10:28:20.031363: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory
2024-06-13 10:28:20.612126: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer.so.7'; dlerror: libnvinfer.so.7: cannot open shared object file: No such file or directory
2024-06-13 10:28:20.612204: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer_plugin.so.7'; dlerror: libnvinfer_plugin.so.7: cannot open shared object file: No such file or directory
2024-06-13 10:28:20.612214: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Cannot dlopen some TensorRT libraries. If you would like to use Nvidia GPU with TensorRT, please make sure the missing libraries mentioned above are installed properly.
Traceback (most recent call last):
  File "~/src/python/flax/examples/imagenet/main.py", line 29, in <module>
    import train
  File "~/src/python/flax/examples/imagenet/train.py", line 28, in <module>
    from flax import jax_utils
  File "~/src/python/flax/examples/imagenet/flax_imagenet_env/lib/python3.10/site-packages/flax/__init__.py", line 18, in <module>
    from .configurations import (
  File "~/src/python/flax/examples/imagenet/flax_imagenet_env/lib/python3.10/site-packages/flax/configurations.py", line 74, in <module>
    flax_filter_frames = define_bool_state(
  File "~/src/python/flax/examples/imagenet/flax_imagenet_env/lib/python3.10/site-packages/flax/configurations.py", line 40, in define_bool_state
    return jax_config.define_bool_state('flax_' + name, default, help)
AttributeError: 'Config' object has no attribute 'define_bool_state'

[sycamoreoak]

Is the fact that this is occurring potentially a CI bug?

[sycamoreoak]

OK this works:

python -m venv mnist_env
source ./mnist_env/bin/activate
unset LD_LIBRARY_PATH
pip install "jax[cuda12]" flax tensorflow clu tensorflow-datasets
python main.py --workdir=/tmp/mnist --config=configs/default.py

might want to update the requirements.txt in the mnist example. this probably applies to other examples as well...

here is my pip freeze:

absl-py==2.1.0
array_record==0.5.1
astunparse==1.6.3
cachetools==5.3.3
certifi==2024.6.2
charset-normalizer==3.3.2
chex==0.1.86
click==8.1.7
clu==0.0.12
contextlib2==21.6.0
dm-tree==0.1.8
docstring_parser==0.16
etils==1.7.0
flatbuffers==24.3.25
flax==0.8.4
fsspec==2024.6.0
gast==0.4.0
google-auth==2.30.0
google-auth-oauthlib==1.0.0
google-pasta==0.2.0
grpcio==1.64.1
h5py==3.11.0
idna==3.7
immutabledict==4.2.0
importlib_resources==6.4.0
jax==0.4.29
jax-cuda12-pjrt==0.4.29
jax-cuda12-plugin==0.4.29
jaxlib==0.4.29
keras==2.13.1
libclang==18.1.1
Markdown==3.6
markdown-it-py==3.0.0
MarkupSafe==2.1.5
mdurl==0.1.2
ml-collections==0.1.1
ml-dtypes==0.4.0
msgpack==1.0.8
nest-asyncio==1.6.0
numpy==1.24.3
nvidia-cublas-cu12==12.5.2.13
nvidia-cuda-cupti-cu12==12.5.39
nvidia-cuda-nvcc-cu12==12.5.40
nvidia-cuda-runtime-cu12==12.5.39
nvidia-cudnn-cu12==9.1.1.17
nvidia-cufft-cu12==11.2.3.18
nvidia-cusolver-cu12==11.6.2.40
nvidia-cusparse-cu12==12.4.1.24
nvidia-nccl-cu12==2.21.5
nvidia-nvjitlink-cu12==12.5.40
oauthlib==3.2.2
opt-einsum==3.3.0
optax==0.2.2
orbax-checkpoint==0.5.16
packaging==24.1
promise==2.3
protobuf==3.20.3
psutil==5.9.8
pyarrow==16.1.0
pyasn1==0.6.0
pyasn1_modules==0.4.0
Pygments==2.18.0
PyYAML==6.0.1
requests==2.32.3
requests-oauthlib==2.0.0
rich==13.7.1
rsa==4.9
scipy==1.13.1
simple_parsing==0.1.5
six==1.16.0
tensorboard==2.13.0
tensorboard-data-server==0.7.2
tensorflow==2.13.1
tensorflow-datasets==4.9.6
tensorflow-estimator==2.13.0
tensorflow-io-gcs-filesystem==0.37.0
tensorflow-metadata==1.15.0
tensorstore==0.1.61
termcolor==2.4.0
toml==0.10.2
toolz==0.12.1
tqdm==4.66.4
typing_extensions==4.5.0
urllib3==2.2.1
Werkzeug==3.0.3
wrapt==1.16.0
zipp==3.19.2