jaxlib v0.4.32 #281

regro-cf-autotick-bot · 2024-09-12T04:09:38Z

It is very likely that the current package version for this feedstock is out of date.

Checklist before merging this PR:

Dependencies have been updated if changed: see upstream
Tests have passed
Updated license if changed and license_file is packaged

Information about this PR:

Feel free to push to the bot's branch to update this PR if needed.
The bot will almost always only open one PR per version.
The bot will stop issuing PRs if more than 3 version bump PRs generated by the bot are open. If you don't want to package a particular version please close the PR.
If you want these PRs to be merged automatically, make an issue with @conda-forge-admin,please add bot automerge in the title and merge the resulting PR. This command will add our bot automerge feature to your feedstock.
If this PR was opened in error or needs to be updated please add the bot-rerun label to this PR. The bot will close this PR and schedule another one. If you do not have permissions to add this label, you can use the phrase @conda-forge-admin, please rerun bot in a PR comment to have the conda-forge-admin add it for you.

Pending Dependency Version Updates

Here is a list of all the pending dependency version updates for this repo. Please double check all dependencies before merging.

Name	Upstream Version	Current Version
bazel	7.3.1
cudnn	9.3.0.75

_{This PR was created by the regro-cf-autotick-bot. The regro-cf-autotick-bot is a service to automatically track the dependency graph, migrate packages, and propose package version updates for conda-forge. Feel free to drop us a line if there are any issues! This PR was generated by - please use this URL for debugging.}

…nda-forge-pinning 2024.09.11.15.30.13

conda-forge-webservices · 2024-09-12T04:09:45Z

Hi! This is the friendly automated conda-forge-linting service.

I just wanted to let you know that I linted all conda-recipes in your PR (recipe/meta.yaml) and found it was in an excellent condition.

I do have some suggestions for making it better though...

For recipe/meta.yaml:

No valid build backend found for Python recipe for package jaxlib using pip. Python recipes using pip need to explicitly specify a build backend in the host section. If your recipe has built with only pip in the host section in the past, you likely should add setuptools to the host section of your recipe.

xhochy · 2024-09-25T20:31:42Z

It doesn't build locally for because of conda-forge/binutils-feedstock#76

github-actions · 2024-09-27T08:45:57Z

Hi! This is the friendly automated conda-forge-linting service.

I just wanted to let you know that I linted all conda-recipes in your PR (recipe/meta.yaml) and found it was in an excellent condition.

…nda-forge-pinning 2024.09.29.16.11.30

xhochy · 2024-10-01T07:35:31Z

This now uses a new hermetic CUDA approach outlined here: https://github.com/openxla/xla/blob/main/docs/hermetic_cuda.md we need to adjust our CUDA builds to it.

jakirkham · 2024-10-10T09:45:16Z

recipe/build.sh

        export CUDA_HOME="${BUILD_PREFIX}/targets/x86_64-linux"
        export TF_CUDA_PATHS="${BUILD_PREFIX}/targets/x86_64-linux,${PREFIX}/targets/x86_64-linux"
-	# Needed for some nvcc binaries
-	export PATH=$PATH:${BUILD_PREFIX}/nvvm/bin
-	# XLA can only cope with a single cuda header include directory, merge both
-	rsync -a ${PREFIX}/targets/x86_64-linux/include/ ${BUILD_PREFIX}/targets/x86_64-linux/include/
+        # Needed for some nvcc binaries
+        export PATH=$PATH:${BUILD_PREFIX}/nvvm/bin
+        # XLA can only cope with a single cuda header include directory, merge both
+        rsync -a ${PREFIX}/targets/x86_64-linux/include/ ${BUILD_PREFIX}/targets/x86_64-linux/include/
+
+        # Although XLA supports a non-hermetic build, it still tries to find headers in the hermetic locations.
+        # We do this in the BUILD_PREFIX to not have any impact on the resulting jaxlib package.
+        # Otherwise, these copied files would be included in the package.
+        rm -rf ${BUILD_PREFIX}/targets/x86_64-linux/include/third_party
+        mkdir -p ${BUILD_PREFIX}/targets/x86_64-linux/include/third_party/gpus/cuda/extras/CUPTI
+        cp -r ${PREFIX}/targets/x86_64-linux/include ${BUILD_PREFIX}/targets/x86_64-linux/include/third_party/gpus/cuda/
+        cp -r ${PREFIX}/targets/x86_64-linux/include ${BUILD_PREFIX}/targets/x86_64-linux/include/third_party/gpus/cuda/extras/CUPTI/
+        mkdir -p ${BUILD_PREFIX}/targets/x86_64-linux/include/third_party/gpus/cudnn
+	cp ${PREFIX}/include/cudnn.h ${BUILD_PREFIX}/targets/x86_64-linux/include/third_party/gpus/cudnn/
+        export LOCAL_CUDA_PATH="${BUILD_PREFIX}/targets/x86_64-linux"
+        export LOCAL_CUDNN_PATH="${PREFIX}/targets/x86_64-linux"
+        export LOCAL_NCCL_PATH="${PREFIX}/targets/x86_64-linux"
+        mkdir -p ${BUILD_PREFIX}/targets/x86_64-linux/bin
+        ln -s $(which ptxas) ${BUILD_PREFIX}/targets/x86_64-linux/bin/ptxas
+        ln -s $(which nvlink) ${BUILD_PREFIX}/targets/x86_64-linux/bin/nvlink
+        ln -s $(which fatbinary) ${BUILD_PREFIX}/targets/x86_64-linux/bin/fatbinary


Am sorry to see this is needed 😢

Is it able to parse CUDA_HOME with multiple paths. IOW would this be possible?

export CUDA_HOME="${PREFIX}:${PREFIX}/targets/x86_64-linux:${BUILD_PREFIX}/targets/x86_64-linux"

No, in the end the bazel scripts do some globbing over $LOCAL_CUDA_PATH/bin* to determine the available CUDA binaries. You can do a bit of trickery and use one of LOCAL_CUDNN_PATH / LOCAL_NCCL_PATH to add a third directory but this then sadly led to the wrong protobuf headers being picked up.

jakirkham · 2024-10-10T09:46:05Z

recipe/meta.yaml

-    - cuda-nvcc  # [(cuda_compiler_version or "").startswith("12")]
-    # Workaround for https://github.com/conda-forge/cuda-cupti-feedstock/issues/14
-    - cuda-cupti >=12.0.90,<13.0a0  # [(cuda_compiler_version or "").startswith("12")]
+    - cuda-nvcc-tools  # [(cuda_compiler_version or "").startswith("12")]


Thanks Uwe! 🙏

Looks good (as does the host addition above)

jakirkham · 2024-10-10T09:47:08Z

recipe/meta.yaml

@@ -57,6 +60,7 @@ requirements:
    - cuda-cudart-dev   # [(cuda_compiler_version or "").startswith("12")]
    - cuda-nvml-dev     # [(cuda_compiler_version or "").startswith("12")]
    - cuda-nvtx-dev     # [(cuda_compiler_version or "").startswith("12")]
+    - cuda-nvcc-tools  # [(cuda_compiler_version or "").startswith("12")]


Is there a way to confirm that JAX is using nvvm from here? Does it have some config paths we can check in the test phase?

I have no idea.

@mwittgen would you know how to do this/where to look? We sadly don't have a GPU device in the CI to test it.

Ah was hoping this kind of check would not actually use the GPU. Sort of like python -m sysconfig for JAX. Poking around it looks like they have jax.print_environment_info(). Maybe we could try that?

If we are unable to figure out an appropriate config check, we could configure CI to use a GPU. Looks like some work has already been done here. So this would just be adding another runner type

Another option we have made heavy use of elsewhere is caching the built packages on CI and then downloading them for local testing ( conda-forge/conda-forge.github.io#1424 ). This may be more useful in the end as it is easier to iterate

Poking around it looks like they have jax.print_environment_info(). Maybe we could try that?

I looked at its implementation and it doesn't provide anything useful for the case we are testing on.

If we are unable to figure out an appropriate config check, we could configure CI to use a GPU. Looks like some work has already been done here. So this would just be adding another runner type

This seems like rather expensive to add a GPU here to do a simple test.

Overall, it seems that this will need input from someone where it doesn't work. The current PR is an improvement on the status quo and I don't want to hold up the release anymore, thus I'm going to merge.

regro-cf-autotick-bot added 2 commits September 12, 2024 04:08

updated v0.4.32

810da94

MNT: Re-rendered with conda-build 24.7.1, conda-smithy 3.39.1, and co…

d1f821e

…nda-forge-pinning 2024.09.11.15.30.13

regro-cf-autotick-bot requested review from ericmjl, ngam and xhochy as code owners September 12, 2024 04:09

xhochy and others added 2 commits September 12, 2024 12:49

Update patches

fb30d19

Don't set cuda_path

462d60b

Add setuptools as a build backend

4e67265

xhochy added 3 commits September 29, 2024 18:44

MNT: Re-rendered with conda-build 24.7.1, conda-smithy 3.40.1, and co…

9f9cc1e

…nda-forge-pinning 2024.09.29.16.11.30

Fix linux CPU build

937e296

Remove cudnn_path

94d311d

xhochy and others added 4 commits October 2, 2024 13:37

Merge remote-tracking branch 'upstream/main' into 0.4.32_h06a6e9

1930b4e

Use local hermetic Python

4ad204f

Fix CUDA 12 build

9c4c61c

Apply fixes from conda-forge#283

8b8ab4a

xhochy mentioned this pull request Oct 10, 2024

libdevice not found when using jax with nvidia drivers installed #283

Open

1 task

jakirkham reviewed Oct 10, 2024

View reviewed changes

xhochy merged commit 9ead0a1 into conda-forge:main Oct 14, 2024
18 checks passed

regro-cf-autotick-bot deleted the 0.4.32_h06a6e9 branch October 14, 2024 13:20

flferretti mentioned this pull request Oct 15, 2024

Update version to 0.4.32 conda-forge/jax-feedstock#158

Merged

5 tasks

hmaarrfk mentioned this pull request Oct 18, 2024

Cannot find libdevice in TF 2.11 + compilation fails without ptxas conda-forge/tensorflow-feedstock#296

Open

1 task

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

jaxlib v0.4.32 #281

jaxlib v0.4.32 #281

regro-cf-autotick-bot commented Sep 12, 2024

conda-forge-webservices bot commented Sep 12, 2024

xhochy commented Sep 25, 2024

github-actions bot commented Sep 27, 2024 •

edited by conda-forge-admin

Loading

xhochy commented Oct 1, 2024

jakirkham Oct 10, 2024

xhochy Oct 10, 2024

jakirkham Oct 10, 2024

jakirkham Oct 10, 2024

xhochy Oct 10, 2024

jakirkham Oct 10, 2024

xhochy Oct 14, 2024

jaxlib v0.4.32 #281

jaxlib v0.4.32 #281

Conversation

regro-cf-autotick-bot commented Sep 12, 2024

Pending Dependency Version Updates

conda-forge-webservices bot commented Sep 12, 2024

xhochy commented Sep 25, 2024

github-actions bot commented Sep 27, 2024 • edited by conda-forge-admin Loading

xhochy commented Oct 1, 2024

jakirkham Oct 10, 2024

Choose a reason for hiding this comment

xhochy Oct 10, 2024

Choose a reason for hiding this comment

jakirkham Oct 10, 2024

Choose a reason for hiding this comment

jakirkham Oct 10, 2024

Choose a reason for hiding this comment

xhochy Oct 10, 2024

Choose a reason for hiding this comment

jakirkham Oct 10, 2024

Choose a reason for hiding this comment

xhochy Oct 14, 2024

Choose a reason for hiding this comment

github-actions bot commented Sep 27, 2024 •

edited by conda-forge-admin

Loading