From 0a5881df1206dcad540cf62d7c05bb5fe5db049e Mon Sep 17 00:00:00 2001
From: Aleksandr Malyshev <164964928+maleksan85@users.noreply.github.com>
Date: Fri, 27 Sep 2024 14:35:40 -0700
Subject: [PATCH] Cythonize vllm build (#214)

* adding cython into docker file with flag

* correcting if
---
 Dockerfile.rocm | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/Dockerfile.rocm b/Dockerfile.rocm
index 506f2212f10b1..90134dd543261 100644
--- a/Dockerfile.rocm
+++ b/Dockerfile.rocm
@@ -12,6 +12,7 @@ ARG BUILD_PYTORCH="1"
 # This ARG should also be "0" or "1". If "1", the vLLM development directory is obtained via git clone.
 # If "0", it is copied in from the local working directory.
 ARG REMOTE_VLLM="0"
+ARG USE_CYTHON="0"
 
 # -----------------------
 # vLLM base image
@@ -153,6 +154,7 @@ FROM fetch_vllm_${REMOTE_VLLM} AS fetch_vllm
 # vLLM (and gradlib) build stages
 FROM fetch_vllm AS build_vllm
 ARG COMMON_WORKDIR
+ARG USE_CYTHON
 # Install hipblaslt
 RUN --mount=type=bind,from=export_hipblaslt,src=/,target=/install \
 if ls /install/*.deb; then \
@@ -168,7 +170,9 @@ if ls /install/*.whl; then \
 fi
 # Build vLLM
 RUN cd vllm \
-    && python3 setup.py clean --all && python3 setup.py bdist_wheel --dist-dir=dist
+    && python3 setup.py clean --all  \
+    && if [ ${USE_CYTHON} -eq "1" ]; then python3 setup_cython.py build_ext --inplace; fi \
+    && python3 setup.py bdist_wheel --dist-dir=dist
 # Build gradlib
 RUN cd vllm/gradlib \
     && python3 setup.py clean --all && python3 setup.py bdist_wheel --dist-dir=dist