Change the python setup script to shell out to the config script inst…

…ead of parsing it adhoc
flexflow · Nov 8, 2022 · ef97c21 · ef97c21
1 parent 19f6cb3
commit ef97c21
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 61 deletions.
diff --git a/config/config.inc b/config/config.inc
@@ -1,6 +1,5 @@
 #!/bin/bash
 
-function run_cmake() {
 # set CC and CXX
 if [ -n "$CC" ]; then
   SET_CC="-DCMAKE_C_COMPILER=${CC}"
@@ -41,7 +40,8 @@ fi
 # set CUDA dir
 if [ -n "$CUDA_DIR" ]; then
   SET_CUDA="-DCUDA_PATH=${CUDA_DIR}"
-  SET_CUDA_LIB_PATH="CUDA_PATH=${CUDA_DIR}/lib64/stubs"
+  CUDA_PATH="${CUDA_DIR}/lib64/stubs"
+  SET_CUDA_LIB_PATH="CUDA_PATH=${CUDA_PATH}"
 fi
 
 # set cudnn dir
@@ -132,8 +132,8 @@ if [ -n "$FF_GPU_BACKEND" ]; then
   # to set these values but this is a sufficient compromise.
   if [ "$FF_GPU_BACKEND" = "hip_cuda" ] || [ "$FF_GPU_BACKEND" = "hip_rocm" ]; then
     if [ -n "$SET_CXX" ]; then
-      echo "FF_GPU_BACKEND is set to ${FF_GPU_BACKEND}. Normally we would set the compiler and linker"
-      echo "to hipcc, but the compiler is already set to ${SET_CXX}".
+      echo "FF_GPU_BACKEND is set to ${FF_GPU_BACKEND}. Normally we would set the compiler and linker" 1>&2
+      echo "to hipcc, but the compiler is already set to ${SET_CXX}". 1>&2
     else
       if [ "$FF_GPU_BACKEND" = "hip_cuda" ]; then
         # Configuring hipcc for nvidia:
@@ -158,8 +158,11 @@ if [ -n "$FF_GPU_BACKEND" ]; then
   fi
 fi
 
+CMAKE_FLAGS="-DCUDA_USE_STATIC_CUDA_RUNTIME=OFF ${SET_CC} ${SET_CXX} ${SET_INSTALL_DIR} ${SET_BUILD} ${SET_CUDA_ARCH} ${SET_CUDA} ${SET_CUDNN} ${SET_PYTHON} ${SET_NCCL} ${SET_GASNET} ${SET_EXAMPLES} ${SET_BUILD_UNIT_TESTS} ${SET_AVX2} ${SET_MAX_DIM} ${SET_ROCM_PATH} ${SET_FF_GPU_BACKEND}"
+
+function run_cmake() {
 SRC_LOCATION=${SRC_LOCATION:=`dirname $0`/../}
-CMAKE_COMMAND="${SET_CC_FLAGS} ${SET_NVCC_FLAGS} ${SET_LD_FLAGS} ${SET_CUDA_LIB_PATH} cmake -DCUDA_USE_STATIC_CUDA_RUNTIME=OFF ${SET_CC} ${SET_CXX} ${SET_INSTALL_DIR} ${SET_BUILD} ${SET_CUDA_ARCH} ${SET_CUDA} ${SET_CUDNN} ${SET_PYTHON} ${SET_NCCL} ${SET_GASNET} ${SET_EXAMPLES} ${SET_BUILD_UNIT_TESTS} ${SET_AVX2} ${SET_MAX_DIM} ${SET_ROCM_PATH} ${SET_FF_GPU_BACKEND} $* ${SRC_LOCATION}"
+CMAKE_COMMAND="${SET_CC_FLAGS} ${SET_NVCC_FLAGS} ${SET_LD_FLAGS} ${SET_CUDA_LIB_PATH} cmake ${CMAKE_FLAGS} $* ${SRC_LOCATION}"
 echo $CMAKE_COMMAND
 eval $CMAKE_COMMAND
 }
diff --git a/config/config.linux b/config/config.linux
@@ -1,7 +1,5 @@
 #!/bin/bash
 
-echo " Defaults for Linux machine"
-
 # set the CC and CXX, usually it is not needed as cmake can detect it
 # set CC and CXX to mpicc and mpic++ when enable gasnet
 # CC=mpicc
@@ -62,4 +60,11 @@ else
 fi
 
 . $(dirname $0)/config.inc
-run_cmake $*
+
+if [ -n "$1" ]; then
+    # You can pass the name of the variable you want to print out as $1. This
+    # is used in the python setup script to get the cmake config
+    echo "${!1}"
+else
+    run_cmake $*
+fi
diff --git a/setup.py b/setup.py
@@ -2,6 +2,7 @@
 from pathlib import Path
 from cmake_build_extension import BuildExtension, CMakeExtension
 import os
+import subprocess
 
 datadir = Path(__file__).parent / "python/flexflow"
 files = [str(p.relative_to(datadir)) for p in datadir.rglob("*.py")]
@@ -11,60 +12,10 @@
     os.path.dirname(os.path.abspath(__file__)), "config", "config.linux"
 )
 
+cmake_configure_options = subprocess.check_output([configs_path, 'CMAKE_FLAGS']).decode("utf-8").strip().split()
+cuda_path = subprocess.check_output([configs_path, 'CUDA_PATH']).decode("utf-8").strip()
 
-def parse_configs_from_file(configs_path):
-    supported_cmake_vars = {
-        "CC": "-DCMAKE_C_COMPILER",
-        "CXX": "-DCMAKE_CXX_COMPILER",
-        "INSTALL_DIR": "-DCMAKE_INSTALL_PREFIX",
-        "BUILD_TYPE": "-DCMAKE_BUILD_TYPE",
-        "FF_CUDA_ARCH": "-DFF_CUDA_ARCH",
-        "CUDA_DIR": "-DCUDA_PATH",
-        "CUDNN_DIR": "-DCUDNN_PATH",
-        "FF_USE_NCCL": "-DFF_USE_NCCL",
-        "FF_USE_GASNET": "-DFF_USE_GASNET",
-        "FF_BUILD_ALL_EXAMPLES": "-DFF_BUILD_ALL_EXAMPLES",
-        "FF_USE_AVX2": "-DFF_USE_AVX2",
-        "FF_MAX_DIM": "-DFF_MAX_DIM",
-    }
-    envs_to_set = {}
-    cmake_vars = {
-        "-DFF_BUILD_FROM_PYPI": "ON",
-        "-DCUDA_USE_STATIC_CUDA_RUNTIME": "OFF",
-        "-DFF_USE_PYTHON": "ON",
-        "-DFF_USE_NCCL": "ON",
-        "-DFF_USE_GASNET": "ON",
-        "-DFF_BUILD_ALL_EXAMPLES": "ON",
-        "-DFF_USE_AVX2": "OFF",
-    }
-    gasnet_conduit = None
-    with open(configs_path, "r") as f:
-        for line in f.readlines():
-            l = line.strip().split("=", 1)
-            if len(l) == 2 and "#" not in l[0]:
-                if len(l[1]) > 0:
-                    if l[0] in supported_cmake_vars:
-                        cmake_vars[supported_cmake_vars[l[0]]] = l[1]
-                        # Special case
-                        if l[0] == "CUDA_DIR":
-                            envs_to_set["CUDA_PATH"] = os.path.join(
-                                l[1], "lib64", "stubs"
-                            )
-                    elif l[0] == "FF_GASNET_CONDUIT":
-                        gasnet_conduit = l[1]
-    # Handle special cases
-    if gasnet_conduit and cmake_vars["-DFF_USE_GASNET"] == "ON":
-        cmake_vars["-DFF_GASNET_CONDUIT"] = gasnet_conduit
-    cmake_vars = ["=".join((k, cmake_vars[k])) for k in cmake_vars]
-
-    return envs_to_set, cmake_vars
-
-
-envs_to_set, cmake_configure_options = parse_configs_from_file(configs_path)
-
-# Export any relevant environment variables
-for k in envs_to_set:
-    os.environ[k] = envs_to_set[k]
+os.environ['CUDA_PATH'] = cuda_path
 
 setup(
     name="flexflow",