XLA smaller dockerfile

Dockerfile

FROM python:3.6-buster as builder

# - Install gsutil to copy wheels from gcr
# - Install openblas
# - Download torch & xla
# - Setup Python 3.6 env for Torch XLA wheels
# - Install torch & xla
RUN echo "deb [signed-by=/usr/share/keyrings/cloud.google.gpg] https://packages.cloud.google.com/apt cloud-sdk main" | tee -a /etc/apt/sources.list.d/google-cloud-sdk.list && \
  apt-get install -y apt-transport-https ca-certificates gnupg curl && \
  curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key --keyring /usr/share/keyrings/cloud.google.gpg add - && \
  apt-get update && apt-get install -y google-cloud-sdk && \
  gsutil cp gs://tpu-pytorch/wheels/torch-1.5-cp36-cp36m-linux_x86_64.whl . && \
  gsutil cp gs://tpu-pytorch/wheels/torch_xla-1.5-cp36-cp36m-linux_x86_64.whl .

FROM python:3.6-buster

COPY --from=builder *.whl /

RUN pip install mkl && \
  pip install torch-1.5-cp36-cp36m-linux_x86_64.whl && \
  pip install torch_xla-1.5-cp36-cp36m-linux_x86_64.whl && \
  rm *.whl && \
  apt-get update && \
  apt-get install -y libopenblas-dev libomp5 


ENV LD_LIBRARY_PATH $LD_LIBRARY_PATH:/usr/local/lib/
ENV XRT_DEVICE_MAP "CPU:0;/job:localservice/replica:0/task:0/device:XLA_CPU:0"
ENV XRT_WORKERS "localservice:0;grpc://localhost:40934"