Upgrading base image and moving mlio build to a seperate stage

Author: Venkatesh Bingi

docker/1.4-2/base/Dockerfile.cpu

@@ -1,13 +1,10 @@
-ARG UBUNTU_VERSION=20.04
-ARG UBUNTU_IMAGE_DIGEST=874aca52f79ae5f8258faff03e10ce99ae836f6e7d2df6ecd3da5c1cad3a912b
+ARG UBUNTU_VERSION=24.04
+ARG UBUNTU_IMAGE_DIGEST=b359f1067efa76f37863778f7b6d0e8d911e3ee8efa807ad01fbf5dc1ef9006b
 
 # Build stage for SQLite compilation
 FROM ubuntu:${UBUNTU_VERSION}@sha256:${UBUNTU_IMAGE_DIGEST} as sqlite-builder
 RUN apt-get update && apt-get install -y --no-install-recommends \
-        build-essential \
-        wget \
-        ca-certificates \
-        && \
+        build-essential wget ca-certificates && \
     cd /tmp && \
     wget https://www.sqlite.org/2025/sqlite-autoconf-3500200.tar.gz && \
     tar xzf sqlite-autoconf-3500200.tar.gz && \
@@ -21,39 +18,19 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
     apt-get clean && \
     rm -rf /var/lib/apt/lists/*
 
-# Main image
-FROM ubuntu:${UBUNTU_VERSION}@sha256:${UBUNTU_IMAGE_DIGEST}
+# MLIO builder stage with Ubuntu 20.04
+FROM ubuntu:20.04@sha256:874aca52f79ae5f8258faff03e10ce99ae836f6e7d2df6ecd3da5c1cad3a912b as mlio-builder
 
 ARG PYTHON_VERSION=3.10
-ARG PYARROW_VERSION=17.0.0
 ARG MLIO_VERSION=v0.9.0
+ARG PYARROW_VERSION=17.0.0
 
 ENV DEBIAN_FRONTEND=noninteractive
 
 # Install python and other scikit-learn runtime dependencies
 RUN apt-get update && \
-    apt-get -y upgrade && \
     apt-get -y install --no-install-recommends \
-        build-essential \
-        curl \
-        git \
-        jq \
-        libatlas-base-dev \
-        nginx \
-        openjdk-8-jdk-headless \
-        unzip \
-        wget \
-        expat \
-        tzdata \
-        apparmor\
-        libgstreamer1.0-0 \
-        linux-libc-dev \
-        libxml2 \
-        libsqlite3-0 \
-        software-properties-common \
-        ca-certificates \
-        lsb-release \
-        && \
+        build-essential curl git wget ca-certificates lsb-release software-properties-common && \
     # Add Apache Arrow repository
     wget https://packages.apache.org/artifactory/arrow/$(lsb_release --id --short | tr 'A-Z' 'a-z')/apache-arrow-apt-source-latest-$(lsb_release --codename --short).deb && \
     apt install -y -V ./apache-arrow-apt-source-latest-$(lsb_release --codename --short).deb && \
@@ -63,37 +40,20 @@ RUN apt-get update && \
     add-apt-repository ppa:deadsnakes/ppa && \
     apt-get update && \
     apt-get -y install --no-install-recommends \
-        python3.10 \
-        python3.10-dev \
-        python3.10-distutils \
-        && \
+        python3.10 python3.10-dev python3.10-distutils && \
     # MLIO build dependencies
     wget http://es.archive.ubuntu.com/ubuntu/pool/main/libf/libffi/libffi7_3.3-4_amd64.deb && \
     dpkg -i libffi7_3.3-4_amd64.deb && \
     apt-get -y install --no-install-recommends \
-        apt-transport-https \
-        ca-certificates \
-        gnupg \
-        && \
+        apt-transport-https gnupg && \
     wget -O - https://apt.kitware.com/keys/kitware-archive-latest.asc 2>/dev/null | \
         gpg --dearmor - | \
         tee /usr/share/keyrings/kitware-archive-keyring.gpg >/dev/null && \
     echo 'deb [signed-by=/usr/share/keyrings/kitware-archive-keyring.gpg] https://apt.kitware.com/ubuntu/ bionic main' | tee /etc/apt/sources.list.d/kitware.list >/dev/null && \
     apt-get update && \
     rm /usr/share/keyrings/kitware-archive-keyring.gpg && \
     apt-get install -y --no-install-recommends \
-        autoconf \
-        automake \
-        cmake \
-        cmake-data \
-        doxygen \
-        kitware-archive-keyring \
-        libcurl4-openssl-dev \
-        libssl-dev \
-        libtool \
-        ninja-build \
-        zlib1g-dev \
-        && \
+        autoconf automake cmake cmake-data doxygen kitware-archive-keyring libcurl4-openssl-dev libssl-dev libtool ninja-build zlib1g-dev && \
     update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.10 1 && \
     curl -sS https://bootstrap.pypa.io/get-pip.py | python3 && \
     apt-get clean && \
@@ -102,17 +62,9 @@ RUN apt-get update && \
 RUN ln -fs /usr/share/zoneinfo/UTC /etc/localtime && \
     dpkg-reconfigure --frontend noninteractive tzdata
 
-# Install uv for fast Python package management
-RUN curl -LsSf https://astral.sh/uv/install.sh | sh && \
-    mv /root/.local/bin/uv /usr/local/bin/uv
-
-ENV PATH=/usr/local/bin:${PATH}
 ENV PIP_ROOT_USER_ACTION=ignore
 
-# Install MLIO with Apache Arrow integration
-# First install Arrow C++ libraries (needed for MLIO compilation)
-RUN uv pip install --system pyarrow==${PYARROW_VERSION}
-
+# Build MLIO from scratch
 # Clone MLIO repository
 RUN cd /tmp && \
     git clone --branch ${MLIO_VERSION} https://github.com/awslabs/ml-io.git mlio
@@ -147,15 +99,64 @@ RUN cd /tmp/mlio/build/release && \
     cmake --build . --target mlio-py && \
     cmake --build . --target mlio-arrow
 
-# Build and install MLIO Python wheel
+# Build MLIO Python wheel
 RUN cd /tmp/mlio/src/mlio-py && \
-    python3 setup.py bdist_wheel && \
-    uv pip install --system dist/*.whl
+    python3 setup.py bdist_wheel
 
-# Copy TBB libraries and cleanup
-RUN cp -r /tmp/mlio/build/third-party/lib/libtbb* /usr/local/lib/ && \
-    ldconfig && \
-    rm -rf /tmp/mlio
+# Copy TBB libraries and MLIO shared libraries to a location we can copy from
+RUN mkdir -p /mlio-artifacts && \
+    cp -r /tmp/mlio/build/third-party/lib/libtbb* /mlio-artifacts/ && \
+    cp /usr/local/lib/libmlio* /mlio-artifacts/ 2>/dev/null || true && \
+    cp /tmp/mlio/src/mlio-py/dist/*.whl /mlio-artifacts/
+
+# Main image
+FROM ubuntu:${UBUNTU_VERSION}@sha256:${UBUNTU_IMAGE_DIGEST}
+
+ARG PYTHON_VERSION=3.10
+ARG PYARROW_VERSION=17.0.0
+
+ENV DEBIAN_FRONTEND=noninteractive
+
+# Install runtime dependencies only
+RUN apt-get update && \
+    apt-get -y upgrade && \
+    apt-get -y install --no-install-recommends \
+        curl git jq libatlas-base-dev nginx openjdk-8-jdk-headless unzip wget expat tzdata apparmor \
+        libgstreamer1.0-0 libxml2 libsqlite3-0 software-properties-common ca-certificates lsb-release \
+        build-essential linux-libc-dev && \
+    # Add Apache Arrow repository for runtime libraries only
+    wget https://packages.apache.org/artifactory/arrow/$(lsb_release --id --short | tr 'A-Z' 'a-z')/apache-arrow-apt-source-latest-$(lsb_release --codename --short).deb && \
+    apt install -y -V ./apache-arrow-apt-source-latest-$(lsb_release --codename --short).deb && \
+    apt-get update && \
+    apt-get install -y -V libarrow-dev=17.0.0-1 libarrow-dataset-dev=17.0.0-1 libparquet-dev=17.0.0-1 libarrow-acero-dev=17.0.0-1 && \
+    # Add deadsnakes PPA for Python 3.10
+    add-apt-repository ppa:deadsnakes/ppa && \
+    apt-get update && \
+    apt-get -y install --no-install-recommends \
+        python3.10 python3.10-distutils python3.10-dev && \
+    update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.10 1 && \
+    curl -sS https://bootstrap.pypa.io/get-pip.py | python3 && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/*
+
+RUN ln -fs /usr/share/zoneinfo/UTC /etc/localtime && \
+    dpkg-reconfigure --frontend noninteractive tzdata
+
+# Install uv for fast Python package management
+RUN curl -LsSf https://astral.sh/uv/install.sh | sh && \
+    mv /root/.local/bin/uv /usr/local/bin/uv
+
+ENV PATH=/usr/local/bin:${PATH}
+ENV PIP_ROOT_USER_ACTION=ignore
+
+# Copy MLIO wheel, TBB libraries, and MLIO shared libraries from builder stage
+COPY --from=mlio-builder /mlio-artifacts/*.whl /tmp/
+COPY --from=mlio-builder /mlio-artifacts/libtbb* /usr/local/lib/
+COPY --from=mlio-builder /mlio-artifacts/libmlio* /usr/local/lib/
+
+# Install MLIO wheel
+RUN uv pip install --system /tmp/*.whl && \
+    rm /tmp/*.whl
 
 # Copy compiled SQLite from builder stage
 COPY --from=sqlite-builder /usr/local/bin/sqlite3 /usr/local/bin/sqlite3

test/unit/test_modules.py

@@ -9,4 +9,4 @@ def test_pandas_version():
 
 def test_pyarrow_to_parquet_conversion_regression_issue_106():
     df = pd.DataFrame({'x': [1, 2]})
-    df.to_parquet('test.parquet', engine='pyarrow')
+    df.to_parquet('test.parquet', engine='pyarrow')