Skip to content

refactor(omni): multi-stage Docker build with decoupled builder/runtime #549

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
KristianZeng wants to merge 1 commit into
base: main
Choose a base branch
from
+152 −93
Draft

refactor(omni): multi-stage Docker build with decoupled builder/runtime #549

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
202 changes: 114 additions & 88 deletions omni/docker/Dockerfile
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,67 +2,54 @@
# Copyright (C) 2025 Intel Corporation
# SPDX-License-Identifier: Apache-2.0

ARG BASE_IMAGE=intel/llm-scaler-platform:26.13.7.1
ARG BUILDER_IMAGE=intel/omix:0.1.0-devel-ubuntu24.04
ARG RUNTIME_IMAGE=intel/pytorch:xpu-2.11.0-ubuntu24.04
ARG PYTHON_VERSION=3.12
ARG TORCH_VERSION=2.10.0
ARG TORCHVISION_VERSION=0.25.0
ARG TORCHAUDIO_VERSION=2.10.0
ARG ONEAPI_VERSION=2025.2.0-517
ARG TORCH_VERSION=2.11.0+xpu
ARG IMAGE_TAG

# ======== Builder Stage ========
FROM ${BASE_IMAGE} AS builder
FROM ${BUILDER_IMAGE} AS builder

ARG https_proxy
ARG http_proxy
ARG no_proxy
ARG PYTHON_VERSION
ARG TORCH_VERSION
ARG TORCHVISION_VERSION
ARG TORCHAUDIO_VERSION
ARG ONEAPI_VERSION
ARG IMAGE_TAG

RUN test -n "${IMAGE_TAG}"

ENV LD_LIBRARY_PATH="/usr/local/lib:/usr/local/lib/python${PYTHON_VERSION}/dist-packages/torch/lib:${LD_LIBRARY_PATH}"
ENV VIRTUAL_ENV=/opt/buildenv \
PATH="/opt/buildenv/bin:${PATH}" \
LD_LIBRARY_PATH="/usr/local/lib:/opt/buildenv/lib/python${PYTHON_VERSION}/site-packages/torch/lib:${LD_LIBRARY_PATH}"

# System packages
RUN wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB \
| gpg --dearmor \
| tee /usr/share/keyrings/oneapi-archive-keyring.gpg > /dev/null && \
echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" \
| tee /etc/apt/sources.list.d/oneAPI.list && \
add-apt-repository -y ppa:kobuk-team/intel-graphics && \
SHELL ["bash", "-c"]

RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
--mount=type=cache,target=/var/lib/apt/lists,sharing=locked \
apt-get update -y && \
apt-get install -y python${PYTHON_VERSION} python${PYTHON_VERSION}-dev python3-pip && \
update-alternatives --install /usr/bin/python3 python3 /usr/bin/python${PYTHON_VERSION} 1 && \
update-alternatives --install /usr/bin/python python /usr/bin/python${PYTHON_VERSION} 1 && \
apt-get install -y --no-install-recommends --fix-missing \
curl ffmpeg git libsndfile1 libsm6 libxext6 libgl1 \
lsb-release numactl wget vim linux-libc-dev \
curl ffmpeg git patch ca-certificates \
python3-dev python3-venv python3-pip \
build-essential ninja-build cmake pkg-config \
libsndfile1 libsm6 libxext6 libgl1 numactl wget vim linux-libc-dev \
libopengl0 libglx0 libxrender1 libxfixes3 libx11-dev \
libxi6 libxxf86vm1 libxcursor1 libxrandr2 libxinerama1 \
libxkbcommon0 && \
apt remove python3-blinker -y && \
apt-get install -y intel-oneapi-dpcpp-ct=${ONEAPI_VERSION} && \
apt-get clean && rm -rf /var/lib/apt/lists/* && \
rm -f /usr/lib/python3*/EXTERNALLY-MANAGED
libxkbcommon0

RUN python3 -m venv ${VIRTUAL_ENV} && \
pip install --upgrade pip setuptools wheel

# PyTorch
RUN --mount=type=cache,target=/root/.cache/pip \
pip install \
torch==${TORCH_VERSION} \
torchvision==${TORCHVISION_VERSION} \
torchaudio==${TORCHAUDIO_VERSION} \
torchvision \
torchaudio \
--index-url https://download.pytorch.org/whl/xpu && \
pip install pytest

# omni_xpu_kernel is installed AFTER sgl-kernel-xpu (below): its cute FMHA
# extension needs the cutlass-sycl headers that sgl-kernel-xpu fetches into
# build/_deps/repo-cutlass-sycl-src.
COPY ./omni_xpu_kernel /tmp/omni_xpu_kernel

# xDiT / long-context-attention
COPY ./patches/yunchang_for_multi_arc.patch /tmp/
COPY ./patches/xdit_for_multi_arc.patch /tmp/

Expand All @@ -84,7 +71,6 @@ RUN --mount=type=cache,target=/root/.cache/pip \
pip install -e . && \
rm -f /tmp/yunchang_for_multi_arc.patch /tmp/xdit_for_multi_arc.patch

# ComfyUI core (upstream, no patches)

@xiangyuT xiangyuT Jul 17, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Restore comments.

RUN --mount=type=cache,target=/root/.cache/pip \
cd /llm && \
git clone --depth 1 https://github.com/comfyanonymous/ComfyUI.git && \
Expand All @@ -93,14 +79,12 @@ RUN --mount=type=cache,target=/root/.cache/pip \
git checkout bb131be9e83d2f773c90f1d6f1e4b248a498c8c5 && \
pip install -r requirements.txt

# ComfyUI-OmniXPU: Intel XPU acceleration via custom node (replaces core patch)
COPY ./ComfyUI-OmniXPU /llm/ComfyUI/custom_nodes/ComfyUI-OmniXPU
RUN cd /llm/ComfyUI/custom_nodes/ComfyUI-OmniXPU && \
git init && git add -A && \
git -c user.name="llm-scaler" -c user.email="noreply@intel.com" \
commit --allow-empty -m "ComfyUI-OmniXPU ${IMAGE_TAG}"

# ComfyUI custom nodes (active)
COPY ./patches/raylight_for_multi_arc.patch /tmp/
COPY ./patches/comfyui_controlnet_aux_depth_anything_v2_xpu.patch /tmp/

Expand Down Expand Up @@ -138,7 +122,6 @@ RUN --mount=type=cache,target=/root/.cache/pip \
pip install -r requirements.txt && \
rm -f /tmp/raylight_for_multi_arc.patch /tmp/comfyui_controlnet_aux_depth_anything_v2_xpu.patch

# nunchaku-torch runtime + ComfyUI-nunchaku-XPU custom node (SVDQuant W4A4/W4A16 diffusion inference)
RUN --mount=type=cache,target=/root/.cache/pip \
cd /llm && \
git clone --depth 1 https://github.com/xiangyuT/nunchaku-torch.git && \
Expand All @@ -154,7 +137,6 @@ RUN --mount=type=cache,target=/root/.cache/pip \
pip install -r requirements.txt && \
pip install --no-deps kernels==0.14.0

# ComfyUI disabled custom nodes (optional: --build-arg INSTALL_DISABLED_NODES=true)
ARG INSTALL_DISABLED_NODES=true
COPY ./patches/comfyui_voxcpm_for_xpu.patch /tmp/
COPY ./patches/comfyui_indextts.patch /tmp/
Expand Down Expand Up @@ -209,20 +191,25 @@ RUN --mount=type=cache,target=/root/.cache/pip \
fi && \
rm -f /tmp/*.patch

# Xinference
COPY ./patches/xinference_device_utils.patch /tmp/

RUN --mount=type=cache,target=/root/.cache/pip \
pip install "xinference[transformers]==1.15.0" \
kokoro Jinja2==3.1.6 jieba ordered-set pypinyin cn2an pypinyin-dict && \
patch /usr/local/lib/python${PYTHON_VERSION}/dist-packages/xinference/device_utils.py \
patch ${VIRTUAL_ENV}/lib/python${PYTHON_VERSION}/site-packages/xinference/device_utils.py \
< /tmp/xinference_device_utils.patch && \
rm -f /tmp/xinference_device_utils.patch

# SGLang Diffusion + sgl-kernel-xpu
COPY ./patches/sglang_diffusion_for_multi_arc.patch /tmp/

ARG DPCPP_SYCL_TARGET=bmg
ARG SGL_KERNEL_XPU_COMMIT=9aef6481524906d11956bf4b4ab8e243c5128b4e

RUN --mount=type=cache,target=/root/.cache/pip \
source /opt/intel/oneapi/setvars.sh --force && \
export TORCH_XPU_ARCH_LIST=${DPCPP_SYCL_TARGET} && \
export OMNI_XPU_DEVICE=${DPCPP_SYCL_TARGET} && \
mkdir -p /wheels && \
cd /llm && \
git clone --depth 1 https://github.com/sgl-project/sglang.git && \
cd sglang && \
Expand All @@ -241,84 +228,123 @@ RUN --mount=type=cache,target=/root/.cache/pip \
'peft==0.19.1' && \
pip install \
torch==${TORCH_VERSION} \
torchvision==${TORCHVISION_VERSION} \
torchaudio==${TORCHAUDIO_VERSION} \
torchvision \
torchaudio \
--index-url https://download.pytorch.org/whl/xpu && \
pip install triton==3.5.0 && \
pip install pytorch-triton-xpu==3.5.0 \
--index-url https://download.pytorch.org/whl/xpu --force-reinstall && \
cd /llm && \
git clone --depth 1 https://github.com/sgl-project/sgl-kernel-xpu.git && \
cd sgl-kernel-xpu && \
pip install -v . && \
git fetch --depth 1 origin ${SGL_KERNEL_XPU_COMMIT} && \
git checkout ${SGL_KERNEL_XPU_COMMIT} && \
CMAKE_BUILD_PARALLEL_LEVEL=8 MAX_JOBS=8 pip wheel -v . \
--config-settings=cmake.define.DPCPP_SYCL_TARGET=${DPCPP_SYCL_TARGET} \
--wheel-dir /wheels --no-deps && \
pip install --no-deps /wheels/sgl_kernel-*.whl && \
cp -r /llm/sglang/python/sglang/multimodal_gen/apps/ComfyUI_SGLDiffusion \
/llm/ComfyUI/custom_nodes/ && \
rm -f /tmp/sglang_diffusion_for_multi_arc.patch

# omni_xpu_kernel (installed here so its cute FMHA ext can use sgl-kernel-xpu's
# cutlass-sycl headers). CUTLASS_SYCL_ROOT points at the fetched tree; the cute
# ext is AOT-compiled for OMNI_XPU_DEVICE (bmg = Arc B-series).
RUN --mount=type=cache,target=/root/.cache/pip \
source /opt/intel/oneapi/setvars.sh --force && \
export TORCH_XPU_ARCH_LIST=${DPCPP_SYCL_TARGET} && \
export CUTLASS_SYCL_ROOT=/llm/sgl-kernel-xpu/build/_deps/repo-cutlass-sycl-src && \
export OMNI_XPU_REQUIRE_CUTE=1 && \
export OMNI_XPU_DEVICE=bmg && \
export OMNI_XPU_DEVICE=${DPCPP_SYCL_TARGET} && \
cd /tmp/omni_xpu_kernel && \
grep -Fqx "__version__ = \"${IMAGE_TAG}\"" omni_xpu_kernel/_version.py && \
pip install . --no-build-isolation && \
CMAKE_BUILD_PARALLEL_LEVEL=8 MAX_JOBS=8 \
pip wheel . --wheel-dir /wheels --no-deps --no-build-isolation && \
pip install --no-deps /wheels/omni_xpu_kernel-*.whl && \
rm -rf /tmp/omni_xpu_kernel

# Application files
RUN ls -lh /wheels/

COPY ./workflows/* /llm/ComfyUI/user/default/workflows/
COPY ./example_inputs/* /llm/ComfyUI/input/
COPY ./tools/* /llm/tools/
COPY ./entrypoints/* /llm/entrypoints/

# ======== Runtime Stage ========
FROM ${BASE_IMAGE} AS runtime
RUN pip freeze --exclude-editable --all | \
sed -E \
-e '/^(omni_xpu_kernel|sgl-kernel) @ file:/d' \
-e '/^(torch|torchaudio|torchvision|triton|triton-xpu|pytorch-triton-xpu)==/d' \
-e '/^(dpcpp-cpp-rt|impi-rt|intel-cmplr-lib-rt|intel-cmplr-lib-ur|intel-cmplr-lic-rt|intel-opencl-rt|intel-openmp|intel-pti|intel-sycl-rt|mkl|oneccl|oneccl-devel|onemkl-license|onemkl-sycl-blas|onemkl-sycl-dft|onemkl-sycl-lapack|onemkl-sycl-rng|onemkl-sycl-sparse|tbb|tcmlib|umf)==/d' \
-e '/^(pip|setuptools|wheel)==/d' \
> /runtime-requirements.txt && \
rm -rf /llm/sgl-kernel-xpu

FROM ${RUNTIME_IMAGE} AS runtime

ARG https_proxy
ARG http_proxy
ARG PYTHON_VERSION=3.12
ARG ONEAPI_VERSION=2025.2.0-517
ARG no_proxy
ARG PYTHON_VERSION
ARG IMAGE_TAG
ARG DPCPP_SYCL_TARGET=bmg

LABEL org.opencontainers.image.version="${IMAGE_TAG}"

# System packages
RUN wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB \
| gpg --dearmor \
| tee /usr/share/keyrings/oneapi-archive-keyring.gpg > /dev/null && \
echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" \
| tee /etc/apt/sources.list.d/oneAPI.list && \
add-apt-repository -y ppa:kobuk-team/intel-graphics && \
COPY --from=builder /llm/ /llm/
COPY --from=builder /runtime-requirements.txt /tmp/runtime-requirements.txt

ENV http_proxy=${http_proxy} \
https_proxy=${https_proxy} \
no_proxy=${no_proxy} \
PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/xpu" \
PIP_DISABLE_PIP_VERSION_CHECK=1 \
TORCH_XPU_ARCH_LIST=${DPCPP_SYCL_TARGET} \
OMNI_XPU_DEVICE=${DPCPP_SYCL_TARGET} \
LD_LIBRARY_PATH="/usr/local/lib:/opt/venv/lib/python${PYTHON_VERSION}/site-packages/torch/lib:${LD_LIBRARY_PATH}" \
CCL_SYCL_ALLTOALL_ARC_LL=1 \
CCL_SYCL_CCL_BARRIER=1 \
OMNIXPU_INTERPOLATE_FIX=0 \
OMNI_IMAGE_VERSION="${IMAGE_TAG}"

SHELL ["bash", "-c"]

RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
--mount=type=cache,target=/var/lib/apt/lists,sharing=locked \
--mount=type=cache,target=/root/.cache/pip \
--mount=type=bind,from=builder,source=/wheels,target=/tmp/wheels \
apt-get update -y && \
apt-get install -y python${PYTHON_VERSION} python${PYTHON_VERSION}-dev python3-pip && \
update-alternatives --install /usr/bin/python3 python3 /usr/bin/python${PYTHON_VERSION} 1 && \
update-alternatives --install /usr/bin/python python /usr/bin/python${PYTHON_VERSION} 1 && \
apt-get install -y --no-install-recommends --fix-missing \
curl ffmpeg git libsndfile1 libsm6 libxext6 libgl1 \
lsb-release numactl wget vim linux-libc-dev \
libopengl0 libglx0 libxrender1 libxfixes3 libx11-dev \
libxi6 libxxf86vm1 libxcursor1 libxrandr2 libxinerama1 \
libxkbcommon0 && \
apt remove python3-blinker -y && \
apt-get install -y intel-oneapi-dpcpp-ct=${ONEAPI_VERSION} && \
apt-get clean && rm -rf /var/lib/apt/lists/* && \
rm -f /usr/lib/python3*/EXTERNALLY-MANAGED

# Copy built artifacts from builder
COPY --from=builder /usr/local/lib/python${PYTHON_VERSION}/ /usr/local/lib/python${PYTHON_VERSION}/
COPY --from=builder /usr/local/bin/ /usr/local/bin/
COPY --from=builder /usr/local/lib/*.so* /usr/local/lib/
COPY --from=builder /llm/ /llm/
curl ffmpeg git libsndfile1 libsm6 libxext6 libgl1 numactl wget \
libopengl0 libglx0 libxrender1 libxfixes3 libxi6 libxxf86vm1 \
libxcursor1 libxrandr2 libxinerama1 libxkbcommon0 && \
pip install --no-deps -r /tmp/runtime-requirements.txt && \
pip install --no-deps \
/tmp/wheels/sgl_kernel-*.whl \
/tmp/wheels/omni_xpu_kernel-*.whl && \
pip install --no-deps -e /llm/long-context-attention && \
pip install --no-deps -e /llm/xDiT && \
pip install --no-deps -e /llm/nunchaku-torch && \
pip install --no-deps -e /llm/sglang/python && \
rm -f /tmp/runtime-requirements.txt && \
rm -rf /var/lib/apt/lists/* /root/.cmake /var/tmp/* && \
find /opt/venv -name '*.pyc' -delete && \
{ find /opt/venv -name '__pycache__' -type d -exec rm -rf {} + 2>/dev/null || true; }

# omni_xpu_kernel links against the oneDNN 3.9 ABI used during build.
COPY --from=builder /opt/intel/oneapi/dnnl/2025.3/lib/libdnnl.so.3.9 /usr/local/lib/libdnnl.so.3.9

RUN ln -sf libdnnl.so.3.9 /usr/local/lib/libdnnl.so.3 && \
ln -sf libdnnl.so.3.9 /usr/local/lib/libdnnl.so && \
ldconfig

ENV LD_LIBRARY_PATH="/usr/local/lib:/opt/venv/lib:/opt/venv/lib/python3.12/site-packages/torch/lib"

# Copy only the oneDNN runtime library required by omni_xpu_kernel.
COPY --from=builder \
/opt/intel/oneapi/dnnl/2025.3/lib/libdnnl.so.3.9 \
/usr/local/lib/libdnnl.so.3.9

# Fix for multi-XPU support: copy SYCL native binary files to a common library path
RUN cp /opt/intel/oneapi/compiler/2025.3/lib/libsycl-native-*.spv /usr/local/lib/
RUN ln -sf libdnnl.so.3.9 /usr/local/lib/libdnnl.so.3 && \
ln -sf libdnnl.so.3.9 /usr/local/lib/libdnnl.so && \
ldconfig

ENV LD_LIBRARY_PATH="/usr/local/lib:/usr/local/lib/python${PYTHON_VERSION}/dist-packages/torch/lib:${LD_LIBRARY_PATH}"
ENV CCL_SYCL_ALLTOALL_ARC_LL=1
ENV CCL_SYCL_CCL_BARRIER=1
ENV OMNIXPU_INTERPOLATE_FIX=0
ENV OMNI_IMAGE_VERSION="${IMAGE_TAG}"
ENV LD_LIBRARY_PATH="/usr/local/lib:/opt/venv/lib:/opt/venv/lib/python3.12/site-packages/torch/lib:${LD_LIBRARY_PATH}"

WORKDIR /llm/entrypoints
43 changes: 38 additions & 5 deletions omni/omni_xpu_kernel/setup.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,30 @@ def get_icpx_path():
return None


def compiler_env_with_explicit_onednn(onednn_include):
"""Remove duplicate oneDNN include paths injected by setvars.sh."""
env = os.environ.copy()
explicit_path = os.path.realpath(onednn_include)

for name in ("CPATH", "C_INCLUDE_PATH", "CPLUS_INCLUDE_PATH"):
value = env.get(name)
if not value:
continue

entries = value.split(os.pathsep)
entries = [
entry for entry in entries
if not entry or os.path.realpath(entry) != explicit_path
]

if entries:
env[name] = os.pathsep.join(entries)
else:
env.pop(name, None)

return env


class ICPXBuildExt(build_ext):
"""Build extension using Intel icpx compiler directly."""

Expand Down Expand Up @@ -268,13 +292,13 @@ def build_extension(self, ext):
"-fPIC", "-shared",
"-std=c++17",
f"-I{python_include}",
f"-I{torch_include}",
f"-I{torch_include}/torch/csrc/api/include",
f"-I{src_dir}",
]
if has_onednn:
cmd.append(f"-I{onednn_include}")
cmd += [
f"-I{torch_include}",
f"-I{torch_include}/torch/csrc/api/include",
f"-I{src_dir}",
f"-L{torch_lib}",
"-ltorch", "-ltorch_python", "-ltorch_cpu", "-ltorch_xpu", "-lc10", "-lc10_xpu",
]
Expand All @@ -288,8 +312,17 @@ def build_extension(self, ext):

print(f"Compile command: {' '.join(cmd)}")

# Run compiler
result = subprocess.run(cmd, capture_output=True, text=True)
# Run compiler with the explicit oneDNN include taking precedence.
compiler_env = os.environ.copy()
if has_onednn and not is_cute:
compiler_env = compiler_env_with_explicit_onednn(onednn_include)

result = subprocess.run(
cmd,
capture_output=True,
text=True,
env=compiler_env,
)

if result.returncode != 0:
print("STDOUT:", result.stdout)
Expand Down