Skip to content

Commit 962daca

Browse files
authored
Use mamba for all conda installs
1 parent abaa7f0 commit 962daca

File tree

1 file changed

+6
-6
lines changed

1 file changed

+6
-6
lines changed

Dockerfile.tmpl

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -80,16 +80,16 @@ ENV PROJ_LIB=/opt/conda/share/proj
8080
# the remaining pip commands: https://www.anaconda.com/using-pip-in-a-conda-environment/
8181
RUN conda config --add channels nvidia && \
8282
conda config --add channels rapidsai && \
83+
conda install -c conda-forge mamba && \
8384
# Base image channel order: conda-forge (highest priority), defaults.
8485
# End state: rapidsai (highest priority), nvidia, conda-forge, defaults.
85-
conda install mkl cartopy=0.19 imagemagick=7.1 pyproj==3.1.0 && \
86+
mamba install mkl cartopy=0.19 imagemagick=7.1 pyproj==3.1.0 && \
8687
/tmp/clean-layer.sh
8788

8889
{{ if eq .Accelerator "gpu" }}
8990

9091
# b/232247930: uninstall pyarrow to avoid double installation with the GPU specific version.
9192
RUN pip uninstall -y pyarrow && \
92-
conda install -c conda-forge mamba && \
9393
mamba install -y cudf cuml && \
9494
# b/255757999: Force switch the cudatoolkit version back to the existing installation so it is compatible with pytorch
9595
conda remove -y --force cudatoolkit=11.2.72 && mamba install --channel https://conda.anaconda.org/conda-forge/linux-64 cudatoolkit=11.2.2 && \
@@ -98,22 +98,22 @@ RUN pip uninstall -y pyarrow && \
9898

9999
# Install implicit
100100
{{ if eq .Accelerator "gpu" }}
101-
RUN conda install implicit implicit-proc=*=gpu && \
101+
RUN mamba install implicit implicit-proc=*=gpu && \
102102
/tmp/clean-layer.sh
103103
{{ else }}
104-
RUN conda install implicit && \
104+
RUN mamba install implicit && \
105105
/tmp/clean-layer.sh
106106
{{ end}}
107107

108108
# Install PyTorch
109109
{{ if eq .Accelerator "gpu" }}
110110
COPY --from=torch_whl /tmp/whl/*.whl /tmp/torch/
111-
RUN conda install -c pytorch magma-cuda${CUDA_MAJOR_VERSION}${CUDA_MINOR_VERSION} && \
111+
RUN mamba install -c pytorch magma-cuda${CUDA_MAJOR_VERSION}${CUDA_MINOR_VERSION} && \
112112
pip install /tmp/torch/*.whl && \
113113
# b/255757999 openmp (libomp.so) is an dependency of libtorchtext and libtorchaudio but
114114
# the built from source versions don't seem to properly link it in. This forces the dep
115115
# which makes sure that libomp is loaded when these libraries are loaded.
116-
conda install -y openmp && \
116+
mamba install -y openmp && \
117117
pip install patchelf && \
118118
patchelf --add-needed libomp.so /opt/conda/lib/python3.7/site-packages/torchtext/lib/libtorchtext.so && \
119119
patchelf --add-needed libomp.so /opt/conda/lib/python3.7/site-packages/torchaudio/lib/libtorchaudio.so && \

0 commit comments

Comments
 (0)