@@ -80,16 +80,16 @@ ENV PROJ_LIB=/opt/conda/share/proj
80
80
# the remaining pip commands: https://www.anaconda.com/using-pip-in-a-conda-environment/
81
81
RUN conda config --add channels nvidia && \
82
82
conda config --add channels rapidsai && \
83
+ conda install -c conda-forge mamba && \
83
84
# Base image channel order: conda-forge (highest priority), defaults.
84
85
# 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 && \
86
87
/tmp/clean-layer.sh
87
88
88
89
{{ if eq .Accelerator "gpu" }}
89
90
90
91
# b/232247930: uninstall pyarrow to avoid double installation with the GPU specific version.
91
92
RUN pip uninstall -y pyarrow && \
92
- conda install -c conda-forge mamba && \
93
93
mamba install -y cudf cuml && \
94
94
# b/255757999: Force switch the cudatoolkit version back to the existing installation so it is compatible with pytorch
95
95
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 && \
98
98
99
99
# Install implicit
100
100
{{ if eq .Accelerator "gpu" }}
101
- RUN conda install implicit implicit-proc=*=gpu && \
101
+ RUN mamba install implicit implicit-proc=*=gpu && \
102
102
/tmp/clean-layer.sh
103
103
{{ else }}
104
- RUN conda install implicit && \
104
+ RUN mamba install implicit && \
105
105
/tmp/clean-layer.sh
106
106
{{ end}}
107
107
108
108
# Install PyTorch
109
109
{{ if eq .Accelerator "gpu" }}
110
110
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} && \
112
112
pip install /tmp/torch/*.whl && \
113
113
# b/255757999 openmp (libomp.so) is an dependency of libtorchtext and libtorchaudio but
114
114
# the built from source versions don't seem to properly link it in. This forces the dep
115
115
# which makes sure that libomp is loaded when these libraries are loaded.
116
- conda install -y openmp && \
116
+ mamba install -y openmp && \
117
117
pip install patchelf && \
118
118
patchelf --add-needed libomp.so /opt/conda/lib/python3.7/site-packages/torchtext/lib/libtorchtext.so && \
119
119
patchelf --add-needed libomp.so /opt/conda/lib/python3.7/site-packages/torchaudio/lib/libtorchaudio.so && \
0 commit comments