Install jaxlib with poetry, cleanup

[warmlogic]

Background

• Apparently jax is not PEP 503 compatible (issue). Boo google.
• I was hoping we could use the method here but I don't know how to do this with multiple sources/platforms
• Here's a different option that I didn't do
• Instead I just specified the wheel URLs per platform

Updates

• Use newer version of Poetry (1.2 preview)
  • Installed this so we could do this with the jax url, but now I'm not actually using that method... I think it's still ok!
  • See the readme for how to uninstall/reinstall quickly
  • I deleted the old env first
  • I ran poetry env list and then poetry env remove <envname>
• Install specific versions of jaxlib using poetry depending on the platform
• Autoformatting/linting cleanup
• I renamed two functions with the same name in b2f8b1e

[xeb]

Hmmm, running into this:

$ python worker.py
Traceback (most recent call last):
  File "/media/xeb/.../projects/dalle-mini-tools/worker.py", line 9, in <module>
    from generate import Generator
  File "/media/xeb/.../projects/dalle-mini-tools/generate.py", line 9, in <module>
    import jax
  File "/home/xeb/.local/lib/python3.10/site-packages/jax/__init__.py", line 35, in <module>
    from jax import config as _config_module
  File "/home/xeb/.local/lib/python3.10/site-packages/jax/config.py", line 17, in <module>
    from jax._src.config import config
  File "/home/xeb/.local/lib/python3.10/site-packages/jax/_src/config.py", line 27, in <module>
    from jax._src import lib
  File "/home/xeb/.local/lib/python3.10/site-packages/jax/_src/lib/__init__.py", line 117, in <module>
    import jaxlib.pocketfft as pocketfft
  File "/home/xeb/.local/lib/python3.10/site-packages/jaxlib/pocketfft.py", line 25, in <module>
    from . import pocketfft_flatbuffers_py_generated as pd
  File "/home/xeb/.local/lib/python3.10/site-packages/jaxlib/pocketfft_flatbuffers_py_generated.py", line 1, in <module>
    import flatbuffers
ModuleNotFoundError: No module named 'flatbuffers'

Will try to look more tonight at why the jaxlib reference (which installed fine with poetry install) has not resolved flatbuffers

[xeb]

Hmmm, my old virtualenv is working. Here's the output:

(venv) xeb@eunice:~/projects/dalle-mini-tools$ pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
Defaulting to user installation because normal site-packages is not writeable
Looking in links: https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
Requirement already satisfied: jax[cuda] in /home/xeb/.local/lib/python3.10/site-packages (0.3.13)
Requirement already satisfied: opt-einsum in /home/xeb/.local/lib/python3.10/site-packages (from jax[cuda]) (3.3.0)
Requirement already satisfied: numpy>=1.19 in /home/xeb/.local/lib/python3.10/site-packages (from jax[cuda]) (1.23.0)
Requirement already satisfied: typing-extensions in /home/xeb/.local/lib/python3.10/site-packages (from jax[cuda]) (4.2.0)
Requirement already satisfied: absl-py in /home/xeb/.local/lib/python3.10/site-packages (from jax[cuda]) (1.1.0)
Requirement already satisfied: scipy>=1.2.1 in /home/xeb/.local/lib/python3.10/site-packages (from jax[cuda]) (1.8.1)
Requirement already satisfied: jaxlib==0.3.10+cuda11.cudnn82 in /home/xeb/.local/lib/python3.10/site-packages (from jax[cuda]) (0.3.10+cuda11.cudnn82)
Requirement already satisfied: flatbuffers<3.0,>=1.12 in /home/xeb/.local/lib/python3.10/site-packages (from jaxlib==0.3.10+cuda11.cudnn82->jax[cuda]) (2.0)
(venv) xeb@eunice:~/projects/dalle-mini-tools$ python worker.py 
Running worker with queue_name='dalle-mini-tools'

e.g. it works! The worker is up and running now.

[xeb]

Ok, making progress. Here's the latest issue. wandb forgot to include setuptools in requirements.txt and therefore poetry cannot resolve the dependency. I tried adding setuptools to poetry, but no luck. The current error is:

$ python worker.py
Traceback (most recent call last):
  File "/media/xeb/GreyArea/projects/dalle-mini-tools/worker.py", line 9, in <module>
    from generate import Generator
  File "/media/xeb/GreyArea/projects/dalle-mini-tools/generate.py", line 12, in <module>
    from dalle_mini import DalleBart, DalleBartProcessor
  File "/home/xeb/.cache/pypoetry/virtualenvs/dalle-mini-tools-4Bbrb4Fl-py3.10/lib/python3.10/site-packages/dalle_mini/__init__.py", line 3, in <module>
    from .model import DalleBart, DalleBartProcessor
  File "/home/xeb/.cache/pypoetry/virtualenvs/dalle-mini-tools-4Bbrb4Fl-py3.10/lib/python3.10/site-packages/dalle_mini/model/__init__.py", line 1, in <module>
    from .configuration import DalleBartConfig
  File "/home/xeb/.cache/pypoetry/virtualenvs/dalle-mini-tools-4Bbrb4Fl-py3.10/lib/python3.10/site-packages/dalle_mini/model/configuration.py", line 21, in <module>
    from .utils import PretrainedFromWandbMixin
  File "/home/xeb/.cache/pypoetry/virtualenvs/dalle-mini-tools-4Bbrb4Fl-py3.10/lib/python3.10/site-packages/dalle_mini/model/utils.py", line 5, in <module>
    import wandb
  File "/home/xeb/.cache/pypoetry/virtualenvs/dalle-mini-tools-4Bbrb4Fl-py3.10/lib/python3.10/site-packages/wandb/__init__.py", line 26, in <module>
    from wandb import sdk as wandb_sdk
  File "/home/xeb/.cache/pypoetry/virtualenvs/dalle-mini-tools-4Bbrb4Fl-py3.10/lib/python3.10/site-packages/wandb/sdk/__init__.py", line 7, in <module>
    from .wandb_artifacts import Artifact  # noqa: F401
  File "/home/xeb/.cache/pypoetry/virtualenvs/dalle-mini-tools-4Bbrb4Fl-py3.10/lib/python3.10/site-packages/wandb/sdk/wandb_artifacts.py", line 30, in <module>
    from wandb.apis import InternalApi, PublicApi
  File "/home/xeb/.cache/pypoetry/virtualenvs/dalle-mini-tools-4Bbrb4Fl-py3.10/lib/python3.10/site-packages/wandb/apis/__init__.py", line 31, in <module>
    from .internal import Api as InternalApi  # noqa
  File "/home/xeb/.cache/pypoetry/virtualenvs/dalle-mini-tools-4Bbrb4Fl-py3.10/lib/python3.10/site-packages/wandb/apis/internal.py", line 1, in <module>
    from wandb.sdk.internal.internal_api import Api as InternalApi
  File "/home/xeb/.cache/pypoetry/virtualenvs/dalle-mini-tools-4Bbrb4Fl-py3.10/lib/python3.10/site-packages/wandb/sdk/internal/internal_api.py", line 12, in <module>
    from pkg_resources import parse_version
ModuleNotFoundError: No module named 'pkg_resources'

[xeb]

Also -- I just realized I was still on poetry==1.1.13 after upgrading to poetry==1.2.0b2 I get the following issue with installing:

$ poetry install
Creating virtualenv dalle-mini-tools-4Bbrb4Fl-py3.10 in /home/xeb/.cache/pypoetry/virtualenvs

Invalid PEP 440 version: '0.1.43ubuntu1'

Note: I am on Pop!_OS

$ uname -a
Linux eunice 5.17.5-76051705-generic #202204271406~1653440576~22.04~6277a18 SMP PREEMPT Wed May 25 01 x86_64 x86_64 x86_64 GNU/Linux

$ cat /etc/issue
Pop!_OS 22.04 LTS \n \l

[xeb]

... is it so bad for us to resurrect requirements.txt and then let each person decide if they want to use poetry? As long as the codebase is the same, I'd rather edit pyproject.toml and requirements.txt in sync than spend time fighting these dependency chains... especially given the recent wandb one.

[uduse]

@warmlogic Hi, appreciate the update but can you also provide a minimum example of how to use this? I added the source but it seems my JAX is still installing the vanilla version install of the GPU version.

[warmlogic]

@warmlogic Hi, appreciate the update but can you also provide a minimum example of how to use this? I added the source but it seems my JAX is still installing the vanilla version install of the GPU version.

@uduse Sorry, I haven't used this in a bunch of months. I'd recommend checking out the main README for installation instructions. Good luck!