RHOAIENG-31715: add s390x Support for TrustyAI Notebook (opendatahub-io#2568)

* Enabled TrustyAI Notebook for s390x

Signed-off-by: Nishan Acharya <Nishan.Acharya@ibm.com>

* Address Comments

Signed-off-by: Nishan Acharya <Nishan.Acharya@ibm.com>

* Removed EPEL from output stage

Signed-off-by: Nishan Acharya <Nishan.Acharya@ibm.com>

* Add s390x label for konflux

Signed-off-by: Nishan Acharya <Nishan.Acharya@ibm.com>

---------

Signed-off-by: Nishan Acharya <Nishan.Acharya@ibm.com>

Author: Nash-123

.tekton/odh-workbench-jupyter-trustyai-cpu-py312-ubi9-pull-request.yaml

@@ -38,6 +38,7 @@ spec:
     - linux/x86_64
     - linux-m2xlarge/arm64
     - linux/ppc64le
+    - linux/s390x
   - name: dockerfile
     value: jupyter/trustyai/ubi9-python-3.12/Dockerfile.cpu
   - name: path-context

jupyter/trustyai/ubi9-python-3.12/Dockerfile.cpu

@@ -18,6 +18,7 @@ RUN curl -Lo mongodb-cli-mongocli-v${MONGOCLI_VERSION}.zip https://github.com/mo
 RUN unzip ./mongodb-cli-mongocli-v${MONGOCLI_VERSION}.zip
 RUN cd ./mongodb-cli-mongocli-v${MONGOCLI_VERSION}/ && \
     CGO_ENABLED=1 GOOS=linux go build -a -tags strictfipsruntime -o /tmp/mongocli ./cmd/mongocli/
+
 ####################
 # wheel-cache-base #
 ####################
@@ -35,7 +36,7 @@ COPY ${TRUSTYAI_SOURCE_CODE}/devel_env_setup.sh .
 
 RUN --mount=type=cache,target=/root/.cache/uv \
     pip install --no-cache-dir uv && \
-    # the devel script is ppc64le specific - sets up build-time dependencies
+    # the devel script is ppc64le and s390x specific - sets up build-time dependencies
     source ./devel_env_setup.sh && \
     # This may have to download and compile some dependencies, and as we don't lock requirements from `build-system.requires`,
     #  we often don't know the correct hashes and `--require-hashes` would therefore fail on non amd64, where building is common.
@@ -121,7 +122,7 @@ WORKDIR /opt/app-root/bin
 USER root
 
 # Install useful OS packages
-RUN dnf install -y jq unixODBC postgresql git-lfs libsndfile libxcrypt-compat && dnf clean all && rm -rf /var/cache/yum
+RUN dnf install -y jq unixODBC unixODBC-devel postgresql git-lfs libsndfile libxcrypt-compat && dnf clean all && rm -rf /var/cache/yum
 
 # Copy dynamically-linked mongocli built in earlier build stage
 COPY --from=mongocli-builder /tmp/mongocli /opt/app-root/bin/
@@ -142,7 +143,7 @@ FROM jupyter-datascience AS jupyter-trustyai
 ARG DATASCIENCE_SOURCE_CODE=jupyter/datascience/ubi9-python-3.12
 ARG TRUSTYAI_SOURCE_CODE=jupyter/trustyai/ubi9-python-3.12
 
-ENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/lib
+ENV LD_LIBRARY_PATH=/usr/lib:/usr/lib64${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}
 
 LABEL name="odh-notebook-jupyter-trustyai-ubi9-python-3.12" \
     summary="Jupyter trustyai notebook image for ODH notebooks" \
@@ -157,28 +158,40 @@ LABEL name="odh-notebook-jupyter-trustyai-ubi9-python-3.12" \
 USER 0
 
 # Install jre that is needed to run the trustyai library
+# Also install runtime libraries for s390x/ppc64le
 RUN INSTALL_PKGS="java-17-openjdk" && \
+    ARCH=$(uname -m) && \
+    if [ "$ARCH" = "s390x" ] || [ "$ARCH" = "ppc64le" ]; then \
+        # Add runtime libraries needed for s390x/ppc64le (OpenBLAS for PyTorch/NumPy)
+        INSTALL_PKGS="$INSTALL_PKGS openblas openblas-threads"; \
+    fi && \
     dnf install -y --setopt=tsflags=nodocs $INSTALL_PKGS && \
-    dnf -y clean all --enablerepo='*'
+    dnf -y clean all --enablerepo='*' && \
+    # Create symlink for compatibility (openblas package provides libopenblasp.so.0 but PyTorch looks for libopenblas.so.0)
+    if [ "$ARCH" = "s390x" ] || [ "$ARCH" = "ppc64le" ]; then \
+        ln -sf /usr/lib64/libopenblasp.so.0 /usr/lib64/libopenblas.so.0 && \
+        ldconfig; \
+    fi
 
 # Install Python packages and Jupyterlab extensions from requirements.txt
 COPY ${TRUSTYAI_SOURCE_CODE}/pylock.toml ./
 
 # install openblas for ppc64le
 RUN --mount=type=cache,from=whl-cache,source=/root/OpenBLAS/,target=/OpenBlas/,rw \
-    bash -c ' \
-        if [[ $(uname -m) == "ppc64le" ]]; then \
+    bash -c 'ARCH=$(uname -m); \
+        if [ "$ARCH" = "ppc64le" ]; then \
             PREFIX=/usr/ make install -C /OpenBlas; \
-        fi '
+        fi'
 
 # Install packages and cleanup
 # install packages as USER 0 (this will allow us to consume uv cache)
 RUN --mount=type=cache,from=whl-cache,source=/wheelsdir/,target=/wheelsdir/,rw \
     --mount=type=cache,target=/root/.cache/uv \
-    bash -c ' \
-        if [[ $(uname -m) == "ppc64le" ]]; then \
+    bash -c 'ARCH=$(uname -m); \
+        if [ "$ARCH" = "ppc64le" ] || [ "$ARCH" = "s390x" ]; then \
             UV_LINK_MODE=copy uv pip install /wheelsdir/*.whl accelerate --cache-dir /root/.cache/uv; \
-        fi '
+        fi'
+
 RUN --mount=type=cache,target=/root/.cache/uv \
     echo "Installing softwares and packages" && \
     # we can ensure wheels are consumed from the cache only by restricting internet access for uv install with '--offline' flag

jupyter/trustyai/ubi9-python-3.12/devel_env_setup.sh

@@ -2,13 +2,25 @@
 set -eoux pipefail
 
 #####################################################################################################
-# This script is expected to be run on ppc64le hosts as `root`                                      #
+# This script is expected to be run on ppc64le and s390x hosts as `root`                           #
 # It installs the required build-time dependencies for python wheels                                #
 # OpenBlas is built from source (instead of distro provided) with recommended flags for performance #
 #####################################################################################################
+
+# Initialize environment variables with default values
+if [[ $(uname -m) == "s390x" ]]; then
+    export GRPC_PYTHON_BUILD_SYSTEM_OPENSSL=1
+    export CFLAGS="-O3"
+    export CXXFLAGS="-O3"
+else
+    # For other architectures, set custom library paths
+    export LD_LIBRARY_PATH=${LD_LIBRARY_PATH:-/usr/local/lib64:/usr/local/lib}
+    export PKG_CONFIG_PATH=${PKG_CONFIG_PATH:-/usr/local/lib64/pkgconfig:/usr/local/lib/pkgconfig}
+fi
+
 WHEELS_DIR=/wheelsdir
 mkdir -p ${WHEELS_DIR}
-if [[ $(uname -m) == "ppc64le" ]]; then
+if [[ $(uname -m) == "ppc64le" ]] || [[ $(uname -m) == "s390x" ]]; then
     CURDIR=$(pwd)
 
     # install development packages
@@ -17,31 +29,64 @@ if [[ $(uname -m) == "ppc64le" ]]; then
     dnf install -y fribidi-devel gcc-toolset-13 lcms2-devel libimagequant-devel patchelf \
         libraqm-devel openjpeg2-devel tcl-devel tk-devel unixODBC-devel
 
-    curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y
+     # Install build tools and libraries needed for compiling PyTorch/PyArrow
+     if [[ $(uname -m) == "s390x" ]]; then
+         dnf install -y gcc gcc-gfortran gcc-c++ make cmake ninja-build \
+             autoconf automake libtool pkg-config \
+             python3.12-devel python3-devel pybind11-devel \
+             openssl-devel openblas-devel \
+             libjpeg-devel zlib-devel libtiff-devel freetype-devel \
+             lcms2-devel libwebp-devel \
+             fribidi-devel openjpeg2-devel libraqm-devel libimagequant-devel \
+             tcl-devel tk-devel unixODBC-devel \
+             git tar wget unzip
+     else
+         # ppc64le packages
+         dnf install -y fribidi-devel lcms2-devel libimagequant-devel \
+             libraqm-devel openjpeg2-devel tcl-devel tk-devel unixODBC-devel
+     fi
+
+     # Install Rust for both ppc64le and s390x
+     curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y
+     source $HOME/.cargo/env
 
-    source /opt/rh/gcc-toolset-13/enable
-    source $HOME/.cargo/env
-    
-    uv pip install cmake
+     # Install cmake via pip (already available via dnf for s390x, but this ensures it's in PATH)
+     uv pip install cmake
+
+     # Set python alternatives for s390x
+     if [[ $(uname -m) == "s390x" ]]; then
+         alternatives --install /usr/bin/python python /usr/bin/python3.12 1
+         alternatives --install /usr/bin/python3 python3 /usr/bin/python3.12 1
+         alternatives --install /usr/bin/python3-config python3-config /usr/bin/python3.12-config 1
+         alternatives --install /usr/bin/python3-devel python3-devel /usr/bin/python3.12-devel 1
+         python --version && python3 --version
+     fi
 
     export MAX_JOBS=${MAX_JOBS:-$(nproc)}
-    export OPENBLAS_VERSION=${OPENBLAS_VERSION:-0.3.30}
-
-    # Install OpenBlas
-    # IMPORTANT: Ensure Openblas is installed in the final image
-    cd /root
-    curl -L https://github.com/OpenMathLib/OpenBLAS/releases/download/v${OPENBLAS_VERSION}/OpenBLAS-${OPENBLAS_VERSION}.tar.gz | tar xz
-    # rename directory for mounting (without knowing version numbers) in multistage builds
-    mv OpenBLAS-${OPENBLAS_VERSION}/ OpenBLAS/
-    cd OpenBLAS/
-    make -j${MAX_JOBS} TARGET=POWER9 BINARY=64 USE_OPENMP=1 USE_THREAD=1 NUM_THREADS=120 DYNAMIC_ARCH=1 INTERFACE64=0
-    make install
-    cd ..
-
-    # set path for openblas
-    export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/opt/OpenBLAS/lib/:/usr/local/lib64:/usr/local/lib
-    export PKG_CONFIG_PATH=$(find / -type d -name "pkgconfig" 2>/dev/null | tr '\n' ':')
-    export CMAKE_ARGS="-DPython3_EXECUTABLE=python"
+
+    # For s390x, we use the system openblas-devel package
+    # Only build OpenBLAS from source for ppc64le
+    if [[ $(uname -m) == "ppc64le" ]]; then
+        export OPENBLAS_VERSION=${OPENBLAS_VERSION:-0.3.30}
+        cd /root
+        curl -L https://github.com/OpenMathLib/OpenBLAS/releases/download/v${OPENBLAS_VERSION}/OpenBLAS-${OPENBLAS_VERSION}.tar.gz | tar xz
+        mv OpenBLAS-${OPENBLAS_VERSION}/ OpenBLAS/
+        cd OpenBLAS/
+        make -j${MAX_JOBS} TARGET=POWER9 BINARY=64 USE_OPENMP=1 USE_THREAD=1 NUM_THREADS=120 DYNAMIC_ARCH=1 INTERFACE64=0
+        make PREFIX=/usr/local install NO_STATIC=1
+        cd ..
+    else
+        # Create empty OpenBLAS directory for s390x (for Docker mount compatibility)
+        mkdir -p /root/OpenBLAS/
+    fi
+
+    # Verify OpenBLAS is found by pkg-config for s390x
+    if [[ $(uname -m) == "s390x" ]]; then
+        echo "Checking OpenBLAS pkg-config..."
+        pkg-config --exists openblas || echo "Warning: openblas.pc not found"
+    fi
+
+    export CMAKE_ARGS="-DPython3_EXECUTABLE=python -DCMAKE_PREFIX_PATH=/usr/local"
     export CMAKE_POLICY_VERSION_MINIMUM=3.5
 
     TMP=$(mktemp -d)
@@ -50,36 +95,82 @@ if [[ $(uname -m) == "ppc64le" ]]; then
     cd ${CURDIR}
     TORCH_VERSION=$(grep -A1 '"torch"' pylock.toml | grep -Eo '\b[0-9\.]+\b')
     cd ${TMP}
-    git clone --recursive https://github.com/pytorch/pytorch.git -b v${TORCH_VERSION}
-    cd pytorch
-    uv pip install -r requirements.txt
-    python setup.py develop
-    rm -f dist/torch*+git*whl
-    MAX_JOBS=${MAX_JOBS:-$(nproc)} \
-        PYTORCH_BUILD_VERSION=${TORCH_VERSION} PYTORCH_BUILD_NUMBER=1 uv build --wheel --out-dir ${WHEELS_DIR}
+    if [[ $(uname -m) == "s390x" ]]; then
+        echo "Building PyTorch for s390x"
+        export CMAKE_C_FLAGS="-fPIC -O2"
+        export CMAKE_CXX_FLAGS="-fPIC -O2"
+        export CFLAGS="-O2 -pipe"
+        export CXXFLAGS="-O2 -pipe"
+        git clone --depth 1 --branch "v${TORCH_VERSION}" --recurse-submodules --shallow-submodules https://github.com/pytorch/pytorch.git
+        cd pytorch
+        pip install --no-cache-dir -r requirements.txt
+        python setup.py develop
+        rm -f dist/torch*+git*whl
+        MAX_JOBS=${MAX_JOBS} PYTORCH_BUILD_VERSION=${TORCH_VERSION} PYTORCH_BUILD_NUMBER=1 uv build --wheel --out-dir ${WHEELS_DIR}
+        echo "PyTorch build completed successfully"
+    else
+	git clone --depth 1 --branch "v${TORCH_VERSION}" --recurse-submodules --shallow-submodules https://github.com/pytorch/pytorch.git
+        cd pytorch
+        uv pip install -r requirements.txt
+        python setup.py develop
+        rm -f dist/torch*+git*whl
+        MAX_JOBS=${MAX_JOBS:-$(nproc)} \
+            PYTORCH_BUILD_VERSION=${TORCH_VERSION} PYTORCH_BUILD_NUMBER=1 uv build --wheel --out-dir ${WHEELS_DIR}
+    fi
 
     cd ${CURDIR}
     # Pyarrow
     PYARROW_VERSION=$(grep -A1 '"pyarrow"' pylock.toml | grep -Eo '\b[0-9\.]+\b')
     cd ${TMP}
-    git clone --recursive https://github.com/apache/arrow.git -b apache-arrow-${PYARROW_VERSION}
+    git clone --depth 1 --branch "apache-arrow-${PYARROW_VERSION}" --recurse-submodules --shallow-submodules https://github.com/apache/arrow.git 
     cd arrow/cpp
     mkdir build && cd build && \
-    cmake -DCMAKE_BUILD_TYPE=release \
-        -DCMAKE_INSTALL_PREFIX=/usr/local \
-        -DARROW_PYTHON=ON \
-        -DARROW_BUILD_TESTS=OFF \
-        -DARROW_JEMALLOC=ON \
-        -DARROW_BUILD_STATIC="OFF" \
-        -DARROW_PARQUET=ON \
-        .. && \
-    make install -j ${MAX_JOBS:-$(nproc)} && \
+    # Set architecture-specific CMake flags
+    if [[ $(uname -m) == "s390x" ]]; then
+        ARROW_CMAKE_FLAGS="-DCMAKE_BUILD_TYPE=Release \
+            -DCMAKE_INSTALL_PREFIX=/usr/local \
+            -DARROW_PYTHON=ON \
+            -DARROW_PARQUET=ON \
+            -DARROW_ORC=ON \
+            -DARROW_FILESYSTEM=ON \
+            -DARROW_JSON=ON \
+            -DARROW_CSV=ON \
+            -DARROW_DATASET=ON \
+            -DARROW_DEPENDENCY_SOURCE=BUNDLED \
+            -DARROW_WITH_LZ4=OFF \
+            -DARROW_WITH_ZSTD=OFF \
+            -DARROW_WITH_SNAPPY=OFF \
+            -DARROW_BUILD_TESTS=OFF \
+            -DARROW_BUILD_BENCHMARKS=OFF"
+    else
+        ARROW_CMAKE_FLAGS="-DCMAKE_BUILD_TYPE=release \
+            -DCMAKE_INSTALL_PREFIX=/usr/local \
+            -DARROW_PYTHON=ON \
+            -DARROW_BUILD_TESTS=OFF \
+            -DARROW_JEMALLOC=ON \
+            -DARROW_BUILD_STATIC=OFF \
+            -DARROW_PARQUET=ON"
+    fi && \
+    cmake ${ARROW_CMAKE_FLAGS} .. && \
+    make -j${MAX_JOBS} VERBOSE=1 && \
+    make install -j${MAX_JOBS} && \
     cd ../../python/ && \
-    uv pip install -v -r requirements-wheel-build.txt && \
-    PYARROW_PARALLEL=${PYARROW_PARALLEL:-$(nproc)} \
-    python setup.py build_ext \
-    --build-type=release --bundle-arrow-cpp \
-    bdist_wheel --dist-dir ${WHEELS_DIR}
+    uv pip install -v -r requirements-build.txt && \
+    if [[ $(uname -m) == "s390x" ]]; then
+        PYARROW_WITH_PARQUET=1 \
+        PYARROW_WITH_DATASET=1 \
+        PYARROW_WITH_FILESYSTEM=1 \
+        PYARROW_WITH_JSON=1 \
+        PYARROW_WITH_CSV=1 \
+        PYARROW_PARALLEL=${MAX_JOBS} \
+        python setup.py build_ext --build-type=release --bundle-arrow-cpp bdist_wheel
+    else
+        PYARROW_PARALLEL=${PYARROW_PARALLEL:-$(nproc)} \
+        python setup.py build_ext \
+        --build-type=release --bundle-arrow-cpp \
+        bdist_wheel
+    fi && \
+    mkdir -p /wheelsdir && cp dist/pyarrow-*.whl /wheelsdir/ && cp dist/pyarrow-*.whl ${WHEELS_DIR}/
 
     # Pillow (use auditwheel repaired wheel to avoid pulling runtime libs from EPEL)
     cd ${CURDIR}
@@ -97,11 +188,15 @@ if [[ $(uname -m) == "ppc64le" ]]; then
     ls -ltr ${WHEELS_DIR}
 
     cd ${CURDIR}
-    uv pip install --refresh ${WHEELS_DIR}/*.whl accelerate==$(grep -A1 '"accelerate"' pylock.toml | grep -Eo '\b[0-9\.]+\b')
+    # Install wheels for s390x and ppc64le
+    if [[ $(uname -m) == "ppc64le" ]] || [[ $(uname -m) == "s390x" ]]; then
+        pip install --no-cache-dir "${WHEELS_DIR}"/*.whl
+        uv pip install --refresh ${WHEELS_DIR}/*.whl accelerate==$(grep -A1 '"accelerate"' pylock.toml | grep -Eo '\b[0-9\.]+\b')
+    fi
 
     uv pip list
     cd ${CURDIR}
 else
-    # only for mounting on non-ppc64le
+    # only for mounting on non-ppc64le and non-s390x
     mkdir -p /root/OpenBLAS/
 fi