search

matthewfeickert / distributed-inference-with-pyhf-and-funcX

Example code for vCHEP 2021 paper "Distributed statistical inference with pyhf enabled through funcX"
MIT License
0 stars 2 forks source link

feat: Update JAX versions to support a wider range of Pythons #23

Closed matthewfeickert closed 2 years ago

matthewfeickert commented 2 years ago

Update the version of JAX used to support a larger range of CPythons. Remove the specifier on the cuda vs. nocuda jaxlib wheel as pip will be able to determine at install time what version can be installed. This also will allow for nocuda versions to be installed on macOS machines that want to replicate the environment locally.

e.g. from https://storage.googleapis.com/jax-releases/jax_releases.html for jaxlib v0.3.7 there is

cuda11/jaxlib-0.3.7+cuda11.cudnn805-cp310-none-manylinux2014_x86_64.whl
cuda11/jaxlib-0.3.7+cuda11.cudnn805-cp37-none-manylinux2014_x86_64.whl
cuda11/jaxlib-0.3.7+cuda11.cudnn805-cp38-none-manylinux2014_x86_64.whl
cuda11/jaxlib-0.3.7+cuda11.cudnn805-cp39-none-manylinux2014_x86_64.whl
cuda11/jaxlib-0.3.7+cuda11.cudnn82-cp310-none-manylinux2014_x86_64.whl
cuda11/jaxlib-0.3.7+cuda11.cudnn82-cp37-none-manylinux2014_x86_64.whl
cuda11/jaxlib-0.3.7+cuda11.cudnn82-cp38-none-manylinux2014_x86_64.whl
cuda11/jaxlib-0.3.7+cuda11.cudnn82-cp39-none-manylinux2014_x86_64.whl
nocuda/jaxlib-0.3.7-cp310-none-manylinux2014_x86_64.whl
nocuda/jaxlib-0.3.7-cp37-none-manylinux2014_x86_64.whl
nocuda/jaxlib-0.3.7-cp38-none-manylinux2014_x86_64.whl
nocuda/jaxlib-0.3.7-cp39-none-manylinux2014_x86_64.whl
mac/jaxlib-0.3.7-cp310-none-macosx_10_9_x86_64.whl
mac/jaxlib-0.3.7-cp310-none-macosx_11_0_arm64.whl
mac/jaxlib-0.3.7-cp37-none-macosx_10_9_x86_64.whl
mac/jaxlib-0.3.7-cp38-none-macosx_10_9_x86_64.whl
mac/jaxlib-0.3.7-cp38-none-macosx_11_0_arm64.whl
mac/jaxlib-0.3.7-cp39-none-macosx_10_9_x86_64.whl
mac/jaxlib-0.3.7-cp39-none-macosx_11_0_arm64.whl

while for jaxlib v0.1.66 there was only Linux wheels

cuda101/jaxlib-0.1.66+cuda101-cp36-none-manylinux2010_x86_64.whl
cuda101/jaxlib-0.1.66+cuda101-cp37-none-manylinux2010_x86_64.whl
cuda101/jaxlib-0.1.66+cuda101-cp38-none-manylinux2010_x86_64.whl
cuda101/jaxlib-0.1.66+cuda101-cp39-none-manylinux2010_x86_64.whl
cuda102/jaxlib-0.1.66+cuda102-cp36-none-manylinux2010_x86_64.whl
cuda102/jaxlib-0.1.66+cuda102-cp37-none-manylinux2010_x86_64.whl
cuda102/jaxlib-0.1.66+cuda102-cp38-none-manylinux2010_x86_64.whl
cuda102/jaxlib-0.1.66+cuda102-cp39-none-manylinux2010_x86_64.whl
cuda110/jaxlib-0.1.66+cuda110-cp36-none-manylinux2010_x86_64.whl
cuda110/jaxlib-0.1.66+cuda110-cp37-none-manylinux2010_x86_64.whl
cuda110/jaxlib-0.1.66+cuda110-cp38-none-manylinux2010_x86_64.whl
cuda110/jaxlib-0.1.66+cuda110-cp39-none-manylinux2010_x86_64.whl
cuda111/jaxlib-0.1.66+cuda111-cp36-none-manylinux2010_x86_64.whl
cuda111/jaxlib-0.1.66+cuda111-cp37-none-manylinux2010_x86_64.whl
cuda111/jaxlib-0.1.66+cuda111-cp38-none-manylinux2010_x86_64.whl
cuda111/jaxlib-0.1.66+cuda111-cp39-none-manylinux2010_x86_64.whl
nocuda/jaxlib-0.1.66-cp36-none-manylinux2010_x86_64.whl
nocuda/jaxlib-0.1.66-cp37-none-manylinux2010_x86_64.whl
nocuda/jaxlib-0.1.66-cp38-none-manylinux2010_x86_64.whl
nocuda/jaxlib-0.1.66-cp39-none-manylinux2010_x86_64.whl

Squash and merge commit message

* Update the version of JAX used to support a larger range of CPythons.
* Remove the specifier on the cuda vs. nocuda jaxlib wheel as pip will be able
to determine at install time what version can be installed. This also will allow for
nocuda versions to be installed on macOS machines that want to replicate the environment
locally.
matthewfeickert commented 2 years ago

@BenGalewsky can you remind me how the RIVER environment is actually setup with Conda? We have this documented on EXPANSE

https://github.com/matthewfeickert/distributed-inference-with-pyhf-and-funcX/blob/91d231e2b5e1a8f4fe1142472b82e1885a19cd4e/README.md?plain=1#L17-L30

but not RIVER and it would be good to make sure that the claim of

Remove the specifier on the cuda vs. nocuda jaxlib wheel as pip will be able to determine at install time what version can be installed.

actually holds up (though I think RIVER just has CPUs?)

(cc @alexander-held in case this is of interest)

matthewfeickert commented 2 years ago

Okay so as the Docker images deployed at RIVER are built on demand, c.f. https://github.com/matthewfeickert/distributed-inference-with-pyhf-and-funcX/issues/25#issuecomment-1106752153, then I think this is fine to merge. In the event that this becomes problematic for anyone I'm happy to revert quickly. :+1: