[Bugfix] Pass PYTHONPATH from setup.py to CMake (vllm-project#7730)

omrishiv · Aug 26, 2024 · 410066a · 410066a
1 parent e03135e
commit 410066a
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -233,7 +233,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
   # Generate sources:
   execute_process(
     COMMAND ${CMAKE_COMMAND} -E env 
-    PYTHONPATH=${CMAKE_CURRENT_SOURCE_DIR}/csrc/cutlass_extensions/:${CUTLASS_DIR}/python/:$PYTHONPATH 
+    PYTHONPATH=${CMAKE_CURRENT_SOURCE_DIR}/csrc/cutlass_extensions/:${CUTLASS_DIR}/python/:${VLLM_PYTHON_PATH}:$PYTHONPATH 
       ${Python_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/csrc/quantization/machete/generate.py
     RESULT_VARIABLE machete_generation_result
     OUTPUT_VARIABLE machete_generation_output

diff --git a/setup.py b/setup.py
@@ -184,6 +184,10 @@ def configure(self, ext: CMakeExtension) -> None:
         # match.
         cmake_args += ['-DVLLM_PYTHON_EXECUTABLE={}'.format(sys.executable)]
 
+        # Pass the python path to cmake so it can reuse the build dependencies
+        # on subsequent calls to python.
+        cmake_args += ['-DVLLM_PYTHON_PATH={}'.format(":".join(sys.path))]
+
         #
         # Setup parallelism and build tool
         #