CODEWITHSUNDEEP
dockerdockerfileenvironment-variablestorch
DeltaIV
DeltaIV

Reputation: 5646

How to save the PyTorch version and the CUDA version to environment variables

I'm writing a Dockerfile which needs to install different pip wheels, depending on the PyTorch version and the CUDA version installed in the base image:

RUN pip install torch-scatter -f https://data.pyg.org/whl/torch-${TORCH}+${CUDA}.html
RUN pip install torch-sparse -f https://data.pyg.org/whl/torch-${TORCH}+${CUDA}.html

For this reason, I need to capture these versions at build time, into the environment variables TORCH and CUDA.

I know I could use the ENV command in a Dockerfile to assign to environmental variables:

ENV TORCH=1.12.0
ENV CUDA=113

(note that CUDA must not contain any dot, unlike TORCH), then build the container. Now, if I log into the running Docker container, I can get these versions from the command line:

python -c "import torch; print(torch.__version__)"
>>> 1.12.0
python -c "import torch; print(torch.version.cuda)"
>>> 11.3

However, I don't want to hardcode the versions in the Dockerfile, because if I change the base image, the hardcoded values will be wrong, I will try installing the wrong wheels, and installation will fail. I want to find them at build time, and assign them to TORCH & CUDA. How can I do it?

Upvotes: 1

Views: 991

Answers (1)

David Maze
David Maze

Reputation: 159242

You can use ordinary shell syntax to set environment variables within a RUN command, with the limitation that those settings will be lost at the end of that command. So within a single RUN command you can use shell command substitution to set the environment variable and use it, but its value will not be available any more after that command.

# all within a single RUN line
RUN TORCH=$(python -c "import torch; print(torch.__version__)"); \
    CUDA=$(python -c "import torch; print(torch.version.cuda)" | sed 's/\.//g'); \
    pip install torch-scatter -f https://data.pyg.org/whl/torch-${TORCH}+${CUDA}.html; \
    pip install torch-sparse -f https://data.pyg.org/whl/torch-${TORCH}+${CUDA}.html

Upvotes: 1

Related Questions