FROM quay.io/pypa/manylinux_2_28_s390x as base

# Language variables
ENV LC_ALL=C.UTF-8
ENV LANG=C.UTF-8
ENV LANGUAGE=C.UTF-8

# there is a bugfix in gcc >= 14 for precompiled headers and s390x vectorization interaction.
# with earlier gcc versions test/inductor/test_cpu_cpp_wrapper.py will fail.
ARG DEVTOOLSET_VERSION=14
# Installed needed OS packages. This is to support all
# the binary builds (torch, vision, audio, text, data)
RUN yum -y install epel-release
RUN yum -y update
RUN yum install -y \
  sudo \
  autoconf \
  automake \
  bison \
  bzip2 \
  curl \
  diffutils \
  file \
  git \
  make \
  patch \
  perl \
  unzip \
  util-linux \
  wget \
  which \
  xz \
  yasm \
  less \
  zstd \
  libgomp \
  gcc-toolset-${DEVTOOLSET_VERSION}-gcc \
  gcc-toolset-${DEVTOOLSET_VERSION}-gcc-c++ \
  gcc-toolset-${DEVTOOLSET_VERSION}-binutils \
  gcc-toolset-${DEVTOOLSET_VERSION}-gcc-gfortran \
  cmake \
  rust \
  cargo \
  llvm-devel \
  libzstd-devel \
  python3.12-devel \
  python3.12-test \
  python3.12-setuptools \
  python3.12-pip \
  python3-virtualenv \
  python3.12-pyyaml \
  python3.12-numpy \
  python3.12-wheel \
  python3.12-cryptography \
  blas-devel \
  openblas-devel \
  lapack-devel \
  atlas-devel \
  libjpeg-devel \
  libxslt-devel \
  libxml2-devel \
  openssl-devel \
  valgrind \
  ninja-build

ENV PATH=/opt/rh/gcc-toolset-${DEVTOOLSET_VERSION}/root/usr/bin:$PATH
ENV LD_LIBRARY_PATH=/opt/rh/gcc-toolset-${DEVTOOLSET_VERSION}/root/usr/lib64:/opt/rh/gcc-toolset-${DEVTOOLSET_VERSION}/root/usr/lib:$LD_LIBRARY_PATH

# git236+ would refuse to run git commands in repos owned by other users
# Which causes version check to fail, as pytorch repo is bind-mounted into the image
# Override this behaviour by treating every folder as safe
# For more details see https://github.com/pytorch/pytorch/issues/78659#issuecomment-1144107327
RUN git config --global --add safe.directory "*"

# installed python doesn't have development parts. Rebuild it from scratch
RUN /bin/rm -rf /opt/_internal /opt/python /usr/local/*/*

# EPEL for cmake
FROM base as patchelf
# Install patchelf
ADD ./common/install_patchelf.sh install_patchelf.sh
RUN bash ./install_patchelf.sh && rm install_patchelf.sh
RUN cp $(which patchelf) /patchelf

FROM patchelf as python
# build python
COPY manywheel/build_scripts /build_scripts
ADD ./common/install_cpython.sh /build_scripts/install_cpython.sh
ENV SSL_CERT_FILE=
RUN bash build_scripts/build.sh && rm -r build_scripts

FROM base as final
COPY --from=python             /opt/python                           /opt/python
COPY --from=python             /opt/_internal                        /opt/_internal
COPY --from=python             /opt/python/cp39-cp39/bin/auditwheel  /usr/local/bin/auditwheel
COPY --from=patchelf           /usr/local/bin/patchelf               /usr/local/bin/patchelf

RUN alternatives --set python /usr/bin/python3.12
RUN alternatives --set python3 /usr/bin/python3.12

RUN pip-3.12 install typing_extensions

ENTRYPOINT []
CMD ["/bin/bash"]

# install test dependencies:
# - grpcio requires system openssl, bundled crypto fails to build
RUN dnf install -y \
  hdf5-devel \
  python3-h5py \
  git

RUN env GRPC_PYTHON_BUILD_SYSTEM_OPENSSL=True pip3 install grpcio

# cmake-3.28.0 from pip for onnxruntime
RUN python3 -mpip install cmake==3.28.0

# build onnxruntime 1.21.0 from sources.
# it is not possible to build it from sources using pip,
# so just build it from upstream repository.
# h5py is dependency of onnxruntime_training.
# h5py==3.11.0 builds with hdf5-devel 1.10.5 from repository.
# h5py 3.11.0 doesn't build with numpy >= 2.3.0.
# install newest flatbuffers version first:
# for some reason old version is getting pulled in otherwise.
# packaging package is required for onnxruntime wheel build.
RUN pip3 install flatbuffers && \
  pip3 install cython 'pkgconfig>=1.5.5' 'setuptools>=77' 'numpy<2.3.0' && \
  pip3 install --no-build-isolation h5py==3.11.0 && \
  pip3 install packaging && \
  git clone https://github.com/microsoft/onnxruntime && \
  cd onnxruntime && git checkout v1.21.0 && \
  git submodule update --init --recursive && \
  ./build.sh --config Release --parallel 0 --enable_pybind \
  --build_wheel --enable_training --enable_training_apis \
  --enable_training_ops --skip_tests --allow_running_as_root \
  --compile_no_warning_as_error && \
  pip3 install ./build/Linux/Release/dist/onnxruntime_training-*.whl && \
  cd .. && /bin/rm -rf ./onnxruntime