ci: set LD_LIBRARY_PATH in Docker images for correct cuBLAS detection#2468
ci: set LD_LIBRARY_PATH in Docker images for correct cuBLAS detection#2468yzh119 merged 2 commits intoflashinfer-ai:mainfrom
Conversation
Summary of ChangesHello @bkryu, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request resolves critical runtime issues within the Docker CI images by ensuring that the correct Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Changelog
Activity
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
📝 WalkthroughWalkthroughFour CUDA Dockerfiles (12.6, 12.8, 12.9, 13.0) add ENV entries that prepend conda-installed NVIDIA cuBLAS library paths to LD_LIBRARY_PATH to prefer conda-provided cuBLAS over system libraries at runtime. Changes
Estimated code review effort🎯 1 (Trivial) | ⏱️ ~3 minutes Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request correctly sets the LD_LIBRARY_PATH in the Docker images to ensure the pip-installed nvidia-cublas library is prioritized, which resolves runtime errors with incorrect cuBLAS versions. The changes are applied consistently across Dockerfiles for different CUDA versions. My review includes a suggestion to improve the maintainability of these Dockerfiles by using ARG variables to avoid hardcoding Python versions and conda environment names. This will make future updates easier and less error-prone.
yzh119
left a comment
yzh119
left a comment
There was a problem hiding this comment.
LGTM, we need to keep in mind about updating these environments variables in the future though (e.g. if we update to py313 or new cuda docker containers).
…flashinfer-ai#2468) <!-- .github/pull_request_template.md --> ## 📌 Description Summary * Add `LD_LIBRARY_PATH` to Docker images to ensure pip-installed `nvidia-cublas` takes precedence over system libraries * Fixes issues where incorrect cuBLAS versions could be loaded at runtime Example of what happens without prepending the path to `LD_LIBRARY_PATH` in our cu130 containers: ``` $ docker run --gpus all -it flashinfer/flashinfer-ci-cu130:20260131-a52eff1 Unable to find image 'flashinfer/flashinfer-ci-cu130:20260131-a52eff1' locally 20260131-a52eff1: Pulling from flashinfer/flashinfer-ci-cu130 Digest: sha256:582aeb35289cf804735a31727abe8ff37ae722fe6c7bd7fb8ddf50654429ff7a Status: Downloaded newer image for flashinfer/flashinfer-ci-cu130:20260131-a52eff1 ========== == CUDA == ========== CUDA Version 13.0.1 Container image Copyright (c) 2016-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved. This container image and its contents are governed by the NVIDIA Deep Learning Container License. By pulling and using the container, you accept the terms and conditions of this license: https://developer.nvidia.com/ngc/nvidia-deep-learning-container-license A copy of this license is made available in this container at /NGC-DL-CONTAINER-LICENSE for your convenience. (py312) root@fdac9b9cd61e:/workspace# python -c "import torch; print(torch.matmul(torch.randn(128,128,device='cuda'), torch.randn(128,128,device='cuda')))" Traceback (most recent call last): File "<string>", line 1, in <module> RuntimeError: CUDA error: CUBLAS_STATUS_INVALID_VALUE when calling `cublasSgemm( handle, opa, opb, m, n, k, &alpha, a, lda, b, ldb, &beta, c, ldc)` (py312) root@fdac9b9cd61e:/workspace# export LD_LIBRARY_PATH=/opt/conda/envs/py312/lib/python3.12/site-packages/nvidia/cu13/lib/:$LD_LIBRARY_PATH (py312) root@fdac9b9cd61e:/workspace# python -c "import torch; print(torch.matmul(torch.randn(128,128,device='cuda'), torch.randn(128,128,device='cuda')))" tensor([[ 14.9044, 14.3420, 26.0861, ..., -10.4334, -4.5352, 4.2331], [ 1.9701, 13.6111, 1.0954, ..., 3.0715, -2.9266, 7.8847], [ 6.5089, -7.4811, -12.6226, ..., -5.3695, -4.4557, -22.4567], ..., [-12.0462, -2.0045, 15.7295, ..., -4.5688, 22.5680, -11.9852], [ -0.4228, 10.2761, 0.1951, ..., 16.5192, 12.7168, 0.9931], [ -0.2800, -5.7174, -2.9644, ..., 1.8484, -10.0042, -7.7290]], device='cuda:0') ``` <!-- What does this PR do? Briefly describe the changes and why they’re needed. --> ## 🔍 Related Issues <!-- Link any related issues here --> ## 🚀 Pull Request Checklist Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete. ### ✅ Pre-commit Checks - [ ] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [ ] I have installed the hooks with `pre-commit install`. - [ ] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. > If you are unsure about how to set up `pre-commit`, see [the pre-commit documentation](https://pre-commit.com/). ## 🧪 Tests - [ ] Tests have been added or updated as needed. - [ ] All tests are passing (`unittest`, etc.). ## Reviewer Notes <!-- Optional: anything you'd like reviewers to focus on, concerns, etc. --> <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Chores** * Updated Docker build configurations for CUDA 12.6, 12.8, 12.9, and 13.0 to set runtime library precedence so conda-installed NVIDIA cuBLAS libraries are favored over system libraries. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
2 participants
📌 Description
Summary
LD_LIBRARY_PATHto Docker images to ensure pip-installednvidia-cublastakes precedence over system librariesExample of what happens without prepending the path to
LD_LIBRARY_PATHin our cu130 containers:🔍 Related Issues
🚀 Pull Request Checklist
Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.
✅ Pre-commit Checks
pre-commitby runningpip install pre-commit(or used your preferred method).pre-commit install.pre-commit run --all-filesand fixed any reported issues.🧪 Tests
unittest, etc.).Reviewer Notes
Summary by CodeRabbit