Dockerfile_CUDA

FROM docker.io/nvidia/cuda:10.2-cudnn7-devel-ubuntu18.04

WORKDIR /usr/src/app
COPY pyproject.toml poetry.lock /usr/src/app/
ENV PIP_DISABLE_PIP_VERSION_CHECK=on

# Get the software-properties-common packages for updating repos
# Please add any debian apt packages dependencies on here (non Python packages)
# TODO (davidnet): Parametrize for python version
RUN apt-get update && apt-get install -y --no-install-recommends \
    software-properties-common && \
    add-apt-repository ppa:deadsnakes -y && \
    apt-get update && apt-get install -y --no-install-recommends \
    python3.8-dev \
    python3.8-distutils \
    curl \

    && curl -Lk "https://bootstrap.pypa.io/get-pip.py" | python3.8 && \
    rm -rf /var/lib/apt/lists/*

RUN python3.8 -m pip install poetry && \
    poetry config virtualenvs.create false && \
    poetry install --no-interaction

# TODO (davidnet): Parametrize for nvidia version
ENV BASE_URL="https://storage.googleapis.com/jax-releases"
RUN python3.8 -m pip install --upgrade $BASE_URL/cuda102/jaxlib-0.1.51-cp38-none-manylinux2010_x86_64.whl
RUN python3.8 -m pip install --upgrade jax
ENV XLA_PYTHON_CLIENT_ALLOCATOR="platform"