From 9273636026902877c3573e5788144983b89ddbab Mon Sep 17 00:00:00 2001
From: Quentin Pradet <quentin.pradet@elastic.co>
Date: Thu, 5 Oct 2023 18:43:52 +0400
Subject: [PATCH] Reduce Docker image size and support arm64 (#615)

Co-authored-by: David Olaru <dolaru@elastic.co>

* Reduce Docker image size from 4.8GB to 2.2GB

* Use torch+cpu variant if target platform is linux/amd64

Avoids downloading large & unnecessary NVIDIA deps defined in the package on PyPI

* Build linux/arm64 image using buildx and QEMU
---
 .buildkite/release-docker/run.sh | 13 ++++++++-----
 Dockerfile                       | 27 +++++++++++++++++++++------
 2 files changed, 29 insertions(+), 11 deletions(-)

diff --git a/.buildkite/release-docker/run.sh b/.buildkite/release-docker/run.sh
index 7bcd34b..48fb17c 100644
--- a/.buildkite/release-docker/run.sh
+++ b/.buildkite/release-docker/run.sh
@@ -22,11 +22,14 @@ pushd eland
 git checkout "v${RELEASE_VERSION}"
 git --no-pager show
 
-docker build -t "$docker_registry/eland/eland:$RELEASE_VERSION" "$PWD"
-docker push "$docker_registry/eland/eland:$RELEASE_VERSION"
-
-docker tag "$docker_registry/eland/eland:$RELEASE_VERSION" "$docker_registry/eland/eland:latest"
-docker push "$docker_registry/eland/eland:latest"
+# Create builder that supports QEMU emulation (needed for linux/arm64)
+docker buildx rm --force eland-multiarch-builder || true
+docker buildx create --name eland-multiarch-builder --bootstrap --use
+docker buildx build --push \
+  --tag "$docker_registry/eland/eland:$RELEASE_VERSION" \
+  --tag "$docker_registry/eland/eland:latest" \
+  --platform linux/amd64,linux/arm64 \
+  "$PWD"
 
 popd
 popd
diff --git a/Dockerfile b/Dockerfile
index 3d5da56..aa706bb 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,13 +1,28 @@
-FROM python:3.10-slim-bookworm@sha256:d364435d339ad318ac4c533b9fbe709739f9ba006b0721fd25e8592d0bb857cb
+# syntax=docker/dockerfile:1
+FROM python:3.10-slim
 
-RUN apt-get update && \
-    apt-get install -y build-essential pkg-config cmake \
-                       libzip-dev libjpeg-dev && \
-    apt-get clean
+RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
+    --mount=type=cache,target=/var/lib/apt,sharing=locked \
+    apt-get update && apt-get install -y \
+      build-essential \
+      pkg-config \
+      cmake \
+      libzip-dev \
+      libjpeg-dev
 
 ADD . /eland
 WORKDIR /eland
 
-RUN python3 -m pip install --no-cache-dir --disable-pip-version-check .[all]
+ARG TARGETPLATFORM
+RUN --mount=type=cache,target=/root/.cache/pip \
+    if [ "$TARGETPLATFORM" = "linux/amd64" ]; then \
+      python3 -m pip install \
+        --no-cache-dir --disable-pip-version-check --extra-index-url https://download.pytorch.org/whl/cpu  \
+        torch==1.13.1+cpu .[all]; \
+    else \
+      python3 -m pip install \
+        --no-cache-dir --disable-pip-version-check \
+        .[all]; \
+    fi
 
 CMD ["/bin/sh"]