[Bug]: Inconsistent jax/jaxlib versioning leads to crashes

[sssangha]

Checked for duplicates

Yes - I've already checked

Describe the bug

When running dolphin (e.g. dolphin run dolphin_config.yaml), the following error is raised through jaxlib:

Traceback (most recent call last):
  File "/u/leffe-data/ssangha/conda_installation/stable_july30_2020/envs/dolphin-env/bin/dolphin", line 8, in <module>
    sys.exit(main())
             ^^^^^^
  File "/u/leffe-data/ssangha/conda_installation/stable_july30_2020/envs/dolphin-env/lib/python3.11/site-packages/dolphin/cli.py", line 28, in main
    run_func(**arg_dict)
  File "/u/leffe-data/ssangha/conda_installation/stable_july30_2020/envs/dolphin-env/lib/python3.11/site-packages/dolphin/workflows/_cli_run.py", line 30, in run
    displacement.run(cfg, debug=debug)
  File "/u/leffe-data/ssangha/conda_installation/stable_july30_2020/envs/dolphin-env/lib/python3.11/site-packages/dolphin/_log.py", line 118, in wrapper
    result = f(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^
  File "/u/leffe-data/ssangha/conda_installation/stable_july30_2020/envs/dolphin-env/lib/python3.11/site-packages/dolphin/workflows/displacement.py", line 64, in run
    utils.disable_gpu()
  File "/u/leffe-data/ssangha/conda_installation/stable_july30_2020/envs/dolphin-env/lib/python3.11/site-packages/dolphin/utils.py", line 137, in disable_gpu
    jax.config.update("jax_platform_name", "cpu")

Upon closer inspection, it looks like my jax/jaxlib modules are out of sync as so:

(dolphin-env) [ssangha@leffe standard_dolphin]$ pip list | grep jax
jax                0.4.25
jaxlib             0.4.23.dev20240125

It was advised through the jax documentation to ensure compatibility between the two packages, but didn't advise how to automate this (i.e. an alternate to simply hardcoding the versioning of the two manually).

However, a separate discussion thread advised to install the most recent version through pip, which worked to clear this issue on my end.

Please refer to the PR #252 I've posted.

*Note while I have demonstrated this isn't an issue associated with the config/run parameters themselves, I've gone ahead and attached the config I used here anyways.
dolphin_config_yaml.txt

What did you expect?

I expected jax/jaxlib to work without leading to a crash.

Reproducible steps

1. dolphin config step: `dolphin config --slc-files CSLCs/*.h5 -sds /data/VV --mask-file /u/trappist-r0/bato/work/OPERA_Applications/CSLC/Discover/BigIsland/WBD/sub_waterMask_utm.tif --threads-per-worker 8 -s 6 3 --unwrap-method phass --work-directory dolphin_config_phass_stride63 --n-parallel-unwrap 2 -o dolphin_config.yaml`
2. dolphin run step `dolphin run dolphin_config.yaml`

Environment

dolphin/isce info:
     dolphin: 0.16.0
 opera_utils: 0.3.0
       isce3: 0.19.1
       tophu: None

System:
      python: 3.11.7 | packaged by conda-forge | (main, Dec 23 2023, 14:43:09) [GCC 12.3.0]
  executable: /u/leffe-data/ssangha/conda_installation/stable_july30_2020/envs/dolphin-env/bin/python
     machine: Linux-3.10.0-1160.59.1.el7.x86_64-x86_64-with-glibc2.17

Python deps:
       numpy: 1.26.4
       numba: 0.58.1
         jax: 0.4.23
  osgeo.gdal: 3.8.3
        h5py: 3.10.0
 ruamel_yaml: None
    pydantic: 2.6.0
  setuptools: 69.0.3
optional GPU info:
gpu_is_available() = True
         jax: 0.4.23
gpu_is_available: True
None

[scottstanie]

huh, that's interesting, thanks for the report!

@mirzaees is this the same problem you encountered when you first tried it on AWS?

Does conda really install a different jax/jaxlib version when you start from a fresh install?? 😞

[mirzaees]

@scottstanie yes exactly, I could not install compatible versions with conda but pip worked
@sssangha thanks for reporting

[scottstanie]

if you can provide the exact environment-creating command at some point, that would be quite helpful. I think if installing jax from conda in a fresh env leads to this, that's worth reporting on https://github.com/conda-forge/jax-feedstock

But for now we can use the pip approach if that worked for both of you 👍

[scottstanie]

Just putting some more speculations here for now- I wouldn't be surprised if this was somehow related to the number of dependencies GDAL from conda uses...

jaxlib looks like it has about 6-7 dependencies on conda. Checking on one of them-

$ conda-tree whoneeds libgrpc
libgoogle-cloud
libarrow
grpcio
jaxlib

the libgoogle-cloud and libarrow are only needed because conda's gdal installs the whole kitchen sink. not 100% sure this is why, but this is another reason why it'll be good to watch out for this getting implemented conda-forge/gdal-feedstock#722 (which is the more realistic scenario than completely removing gdal as a required dep in the short term)

[sssangha]

if you can provide the exact environment-creating command at some point, that would be quite helpful. I think if installing jax from conda in a fresh env leads to this, that's worth reporting on https://github.com/conda-forge/jax-feedstock

But for now we can use the pip approach if that worked for both of you 👍

Yeah, in my experience it was just installing jax from conda in a fresh env that led to this. And ok, I'll see to reporting this there as you advise.

Thanks!

[scottstanie]

Just clarifying- it may not actually be a jax feedstock bug. it could be a bad mix of gdal+jax

[scottstanie]

@sssangha What is the actual error in your traceback? I see that you pasted up until

  File "/u/leffe-data/ssangha/conda_installation/stable_july30_2020/envs/dolphin-env/lib/python3.11/site-packages/dolphin/utils.py", line 137, in disable_gpu
    jax.config.update("jax_platform_name", "cpu")

Was the error saying "mismatch between jaxlib and jax"?

Testing my guess about gdal is easy:

$ mamba create -n test-jax-311 python=3.11 jax gdal
...

  + gdal                              3.8.4  py311h43f0207_0       conda-forge        2MB
  + jaxlib                           0.4.23  cpu_py311h6cc994a_0   conda-forge       45MB

  + jax                              0.4.25  pyhd8ed1ab_0          conda-forge        1MB

So it does install different versions. BUT, after I created this, I can import jax alright..

>>> import jax

I do get a problem with the "config" call

>>> jax.config.update("jax_platform_name", "cpu")
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
AttributeError: module 'jax.config' has no attribute 'update'

but this is definitely different than what @mirzaees saw and is a much quicker fix.

I'm also now reading that the jax and jaxlib packages dont need to be exactly the same. just "compatible"

So it's possible the change to use pip was a red herring and not the actual issue

[scottstanie]

@sssangha any word on what the final traceback was? I'm thinking of reverting the change to conda-env.yml from #252 if I can't reproduce the problem... I just installed twice in a fresh env and I'm getting no errors from jax

[sssangha]

@scottstanie apologies just saw this, yeah I'm no longer able to replicate and access the original error message, but going off of memory, I recall something like:
jaxlib does not contain config

But looking through the forums I linked above, I just followed their suggestion to enforce jax and jaxlib.

So all that said, perhaps it was transient. Perhaps reverting back like you suggest would be for the best, but in case it crops up again we have the issue ticket as a point of reference.