From 49523d62691039a9a8c29891acc5d48641048cc4 Mon Sep 17 00:00:00 2001
From: Gabriele Oliaro <goliaro@cs.cmu.edu>
Date: Mon, 2 Sep 2024 03:05:25 -0700
Subject: [PATCH] Fix python install issue caused by new Legion version (#1482)

* fix

* .

* .

* fix

* cleanup

* fix

* cleanup
---
 CMakeLists.txt                   | 20 ++++++++++++++------
 cmake/pip_install/CMakeLists.txt | 20 ++++++++++++++++++--
 pyproject.toml                   |  3 ++-
 requirements.txt                 |  1 +
 4 files changed, 35 insertions(+), 9 deletions(-)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index d7a6391e06..c82a53644e 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -37,13 +37,24 @@ if(NOT CMAKE_BUILD_TYPE AND NOT CMAKE_CONFIGURATION_TYPES)
       STRING "Choose the type of build." FORCE)
 endif()
 
+# option for using Python
+option(FF_USE_PYTHON "Enable Python" ON)
+if (FF_USE_PYTHON)
+  find_package(Python3 COMPONENTS Interpreter Development)
+endif()
+
 if(INSTALL_DIR)
   message(STATUS "INSTALL_DIR: ${INSTALL_DIR}")
   set(CMAKE_INSTALL_PREFIX ${INSTALL_DIR} CACHE PATH "Installation directory" FORCE)
 else()
-  # Install DIR not set. Use default, unless a conda environment is active
-  if (DEFINED ENV{CONDA_PREFIX} AND NOT FF_BUILD_FROM_PYPI)
-    set(CONDA_PREFIX $ENV{CONDA_PREFIX})
+  # Install DIR not set. Use default, unless a conda environment is in use
+  if ((DEFINED ENV{CONDA_PREFIX} OR (Python3_EXECUTABLE AND Python3_EXECUTABLE MATCHES "conda")) AND NOT FF_BUILD_FROM_PYPI)
+    if (DEFINED ENV{CONDA_PREFIX})
+        set(CONDA_PREFIX $ENV{CONDA_PREFIX})
+    else()
+        get_filename_component(CONDA_PREFIX "${Python3_EXECUTABLE}" DIRECTORY)
+        get_filename_component(CONDA_PREFIX "${CONDA_PREFIX}" DIRECTORY)
+    endif()
     # Set CMAKE_INSTALL_PREFIX to the Conda environment's installation path
     set(CMAKE_INSTALL_PREFIX ${CONDA_PREFIX} CACHE PATH "Installation directory" FORCE)
     message(STATUS "Active conda environment detected. Setting CMAKE_INSTALL_PREFIX: ${CMAKE_INSTALL_PREFIX}")
@@ -64,9 +75,6 @@ option(FF_BUILD_FROM_PYPI "Build from pypi" OFF)
 # build shared or static flexflow lib
 option(BUILD_SHARED_LIBS "Build shared libraries instead of static ones" ON)
 
-# option for using Python
-option(FF_USE_PYTHON "Enable Python" ON)
-
 # option for building legion only 
 option(BUILD_LEGION_ONLY "Build Legion only" OFF)
 
diff --git a/cmake/pip_install/CMakeLists.txt b/cmake/pip_install/CMakeLists.txt
index 105133a310..217d7e14f0 100644
--- a/cmake/pip_install/CMakeLists.txt
+++ b/cmake/pip_install/CMakeLists.txt
@@ -2,9 +2,25 @@
 if (FF_USE_PYTHON)
     execute_process(COMMAND ${Python_EXECUTABLE} -c "import site, os; print([pkg for func in (site.getsitepackages(), site.getusersitepackages()) for pkg in ([func] if isinstance(func, str) else func) if os.access(pkg, os.W_OK)][0])" OUTPUT_VARIABLE PY_DEST OUTPUT_STRIP_TRAILING_WHITESPACE)  
     if(FF_BUILD_FROM_PYPI)
-        install(CODE "execute_process(COMMAND ${CMAKE_COMMAND} -E echo \"Editing path to Legion library using path: ${PY_DEST}/flexflow/lib \")")
+        cmake_path(SET CMAKE_SOURCE_DIR_ NORMALIZE ${CMAKE_CURRENT_SOURCE_DIR}/../../deps/legion)
+        cmake_path(SET CMAKE_BUILD_DIR_ NORMALIZE ${Legion_BINARY_DIR}/runtime)
+        cmake_path(SET CMAKE_INSTALL_PREFIX_ NORMALIZE ${PY_DEST}/../../..)
+        cmake_path(SET WORKING_DIRECTORY_ NORMALIZE ${CMAKE_CURRENT_SOURCE_DIR}/../../deps/legion/bindings/python/)
         # CMAKE_CURRENT_SOURCE_DIR=/usr/FlexFlow/cmake/pip_install 
         # Legion_BINARY_DIR=/usr/FlexFlow/build/<something>/deps/legion
-        install(CODE "execute_process(COMMAND ${Python_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/../../deps/legion/bindings/python/setup.py install --cmake-build-dir ${Legion_BINARY_DIR}/runtime --prefix ${PY_DEST}/flexflow ${Legion_PYTHON_EXTRA_INSTALL_ARGS} WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/../../deps/legion/bindings/python)")
+        # CMAKE_SOURCE_DIR_=/usr/FlexFlow/deps/legion
+        # CMAKE_BUILD_DIR_: /usr/FlexFlow/build/<something>/deps/legion/runtime
+        # CMAKE_INSTALL_PREFIX_: /opt/conda/ or /usr/local
+        # WORKING_DIRECTORY_: /usr/FlexFlow/deps/legion/bindings/python/
+        # PY_DEST: <path to python lib>/python3.11/site-packages
+        message(STATUS "CMAKE_CURRENT_SOURCE_DIR: ${CMAKE_CURRENT_SOURCE_DIR}")
+        message(STATUS "Legion_BINARY_DIR: ${Legion_BINARY_DIR}")
+        message(STATUS "CMAKE_SOURCE_DIR_: ${CMAKE_SOURCE_DIR_}")
+        message(STATUS "CMAKE_BUILD_DIR_: ${CMAKE_BUILD_DIR_}")
+        message(STATUS "CMAKE_INSTALL_PREFIX_: ${CMAKE_INSTALL_PREFIX_}")
+        message(STATUS "WORKING_DIRECTORY_: ${WORKING_DIRECTORY_}")
+        message(STATUS "PY_DEST: ${PY_DEST}")
+        install(CODE "execute_process(COMMAND ${CMAKE_COMMAND} -E echo \"Editing path to Legion library using path: ${CMAKE_INSTALL_PREFIX_} \")")
+        install(CODE "execute_process(COMMAND ${CMAKE_COMMAND} -E env CMAKE_SOURCE_DIR=${CMAKE_SOURCE_DIR_} CMAKE_BUILD_DIR=${CMAKE_BUILD_DIR_} CMAKE_INSTALL_PREFIX=${PY_DEST}/flexflow ${Python3_EXECUTABLE} setup.py install --prefix ${CMAKE_INSTALL_PREFIX_} ${Legion_PYTHON_EXTRA_INSTALL_ARGS} WORKING_DIRECTORY ${WORKING_DIRECTORY_} COMMAND_ECHO STDOUT COMMAND_ERROR_IS_FATAL ANY)")
     endif()
 endif()
diff --git a/pyproject.toml b/pyproject.toml
index 4b8214f3fe..373c53beb8 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -5,6 +5,7 @@ requires = [
     "setuptools_scm[toml]>=6.0",
     "cmake-build-extension",
     "ninja",
-    "requests"
+    "requests",
+    "pip",
 ]
 build-backend = "setuptools.build_meta"
diff --git a/requirements.txt b/requirements.txt
index 1037661337..ad65622367 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -15,3 +15,4 @@ onnx
 transformers>=4.31.0
 sentencepiece
 einops
+pip