feat: enable cuequivariance kernels on the base image #44
Conversation
avelinoapheris
commented
Nov 18, 2025
avelinoapheris
commented
- Modify the base image to install the cuequivariance kernels. This also involve upgrading the base cuda version.
- Add checks to ensure that the cuequivariance kernels are not loaded when using a setup without GPUs.
…ge + fixes avoid errors when cuequivariance is installed in setups without gpus
|
Hi @avelinoapheris thank you for the PR. Thank you for the fix regarding checking if cuEquivariance is installed. Regarding updating the dockerfile - do you know if this new docker recipe would be compatible on GPUs that might have an older CUDA driver (<12.5) or are otherwise incompatible (e.g. NVIDIA/cuEquivariance#209)? One option could be to support multiple docker recipes. We had a previous PR contribution for a dockerfile specifically for Blackwell GPUs |
jandom
left a comment
jandom
left a comment
There was a problem hiding this comment.
Great contribution @avelinoapheris – couple of questions/comments below
Many thanks!
| # Install the package | ||
| WORKDIR /opt/openfold3 | ||
| RUN python3 setup.py install | ||
| RUN pip install '.[cuequivariance]' --no-build-isolation && python setup.py install |
There was a problem hiding this comment.
Yeah, I've also found that even with --no-build-isolation, you still need to run python setup.py install
| cueq_is_installed = ( | ||
| importlib.util.find_spec("cuequivariance_torch") is not None | ||
| and torch.cuda.is_available() | ||
| ) |
There was a problem hiding this comment.
Can we mutualize this into a helper function? it seems to be duplicated 2-3 times in this PR
|
|
||
| # Runtime stage - use devel image for full CUDA support | ||
| FROM nvidia/cuda:12.1.1-cudnn8-devel-ubuntu22.04 AS runtime | ||
| FROM nvidia/cuda:12.6.1-cudnn-devel-ubuntu22.04 AS runtime |
There was a problem hiding this comment.
This has me worried the most, we're adding multiple base build images support – is this strictly necessary for cuequivariance or just nice to have?
3 participants
