Refactor the parameters and usage instructions of the setup script. (#3162)

As this
issuecomment(https://github.com/llvm/torch-mlir/pull/3021#issuecomment-2031248199)
suggests, `setup.py` should only be used for building Python packages,
so:
* disabled the develop command
* refactor the environment variable parameters
* add more doc for the usage and env var of setup.py

docs/development.md

@@ -349,6 +349,15 @@ python -m pip install -r requirements.txt
 CMAKE_GENERATOR=Ninja python setup.py bdist_wheel
 ```
 
+To package a completed CMake build directory,
+you can use the `TORCH_MLIR_CMAKE_BUILD_DIR` and `TORCH_MLIR_CMAKE_ALREADY_BUILT` environment variables:
+```shell
+TORCH_MLIR_CMAKE_BUILD_DIR=build/ TORCH_MLIR_CMAKE_ALREADY_BUILT=1 python setup.py bdist_wheel
+```
+
+Note: The setup.py script is only used for building the Python packages,
+not support commands like `setup.py develop` to build the development environment.
+
 # Testing
 
 Torch-MLIR has two types of tests:

setup.py

@@ -7,6 +7,31 @@
 # ```
 # $ python setup.py bdist_wheel
 # ```
+# Environment variables you are probably interested in:
+#
+#   TORCH_MLIR_PYTHON_PACKAGE_VERSION:
+#       specify the version of torch-mlir, for example, this can be "20220330.357"
+#       for a snapshot release on 2022-03-30 with build number 357.
+#
+#   TORCH_MLIR_ENABLE_LTC:
+#       enables the Lazy Tensor Core Backend
+#
+#   LLVM_INSTALL_DIR:
+#       build the project *out-of-tree* using the built llvm-project
+#
+#   CMAKE_BUILD_TYPE:
+#       specify the build type: DEBUG/RelWithDebInfo/Release
+#
+#   TORCH_MLIR_CMAKE_BUILD_DIR:
+#       specify the cmake build directory
+#
+#   TORCH_MLIR_CMAKE_ALREADY_BUILT:
+#       the `TORCH_MLIR_CMAKE_BUILD_DIR` directory has already been compiled,
+#       and the CMake compilation process will not be executed again.
+#       On CIs, it is often advantageous to re-use/control the CMake build directory.
+#
+#   MAX_JOBS:
+#       maximum number of compile jobs we should use to compile your code
 #
 # It is recommended to build with Ninja and ccache. To do so, set environment
 # variables by prefixing to above invocations:
@@ -14,16 +39,6 @@
 # CMAKE_GENERATOR=Ninja CMAKE_C_COMPILER_LAUNCHER=ccache CMAKE_CXX_COMPILER_LAUNCHER=ccache
 # ```
 #
-# On CIs, it is often advantageous to re-use/control the CMake build directory.
-# This can be set with the TORCH_MLIR_CMAKE_BUILD_DIR env var.
-# Additionally, the TORCH_MLIR_CMAKE_BUILD_DIR_ALREADY_BUILT env var will
-# prevent this script from attempting to build the directory, and will simply
-# use the (presumed already built) directory as-is.
-#
-# The package version can be set with the TORCH_MLIR_PYTHON_PACKAGE_VERSION
-# environment variable. For example, this can be "20220330.357" for a snapshot
-# release on 2022-03-30 with build number 357.
-#
 # Implementation notes:
 # The contents of the wheel is just the contents of the `python_packages`
 # directory that our CMake build produces. We go through quite a bit of effort
@@ -42,20 +57,33 @@ from setuptools.command.build_ext import build_ext
 from setuptools.command.build_py import build_py
 
 
-def check_env_flag(name: str, default=None) -> bool:
+if "develop" in sys.argv:
+    print("Warning: The setup.py script is only used for building the wheel package.")
+    print("For initializing the development environment,"
+          "please use the cmake commands introduced in the docs/development.md.")
+    sys.exit(1)
+
+def _check_env_flag(name: str, default=None) -> bool:
     return str(os.getenv(name, default)).upper() in ["ON", "1", "YES", "TRUE", "Y"]
 
 
-PACKAGE_VERSION = os.environ.get("TORCH_MLIR_PYTHON_PACKAGE_VERSION") or "0.0.1"
+PACKAGE_VERSION = os.getenv("TORCH_MLIR_PYTHON_PACKAGE_VERSION", "0.0.1")
 
 # If true, enable LTC build by default
-TORCH_MLIR_ENABLE_LTC_DEFAULT = True
+TORCH_MLIR_ENABLE_LTC = _check_env_flag("TORCH_MLIR_ENABLE_LTC", True)
-TORCH_MLIR_ENABLE_ONLY_MLIR_PYTHON_BINDINGS = check_env_flag(
+TORCH_MLIR_ENABLE_ONLY_MLIR_PYTHON_BINDINGS = _check_env_flag(
-    'TORCH_MLIR_ENABLE_ONLY_MLIR_PYTHON_BINDINGS', False)
+    "TORCH_MLIR_ENABLE_ONLY_MLIR_PYTHON_BINDINGS", False
-LLVM_INSTALL_DIR = os.getenv('LLVM_INSTALL_DIR', None)
+)
+LLVM_INSTALL_DIR = os.getenv("LLVM_INSTALL_DIR", None)
 SRC_DIR = pathlib.Path(__file__).parent.absolute()
 CMAKE_BUILD_TYPE = os.getenv("CMAKE_BUILD_TYPE", "Release")
 
+TORCH_MLIR_CMAKE_ALREADY_BUILT = _check_env_flag(
+    "TORCH_MLIR_CMAKE_ALREADY_BUILT", False
+)
+TORCH_MLIR_CMAKE_BUILD_DIR = os.getenv("TORCH_MLIR_CMAKE_BUILD_DIR")
+MAX_JOBS = os.getenv("MAX_JOBS", str(multiprocessing.cpu_count()))
+
 
 # Build phase discovery is unreliable. Just tell it what phases to run.
 class CustomBuild(_build):
@@ -77,8 +105,6 @@ class CMakeBuild(build_py):
 
     def cmake_build(self, cmake_build_dir):
         llvm_dir = str(SRC_DIR / "externals" / "llvm-project" / "llvm")
-        enable_ltc = check_env_flag('TORCH_MLIR_ENABLE_LTC', TORCH_MLIR_ENABLE_LTC_DEFAULT)
-        max_jobs = os.getenv("MAX_JOBS") or str(multiprocessing.cpu_count())
 
         cmake_config_args = [
             f"cmake",
@@ -92,7 +118,7 @@ class CMakeBuild(build_py):
             f"-DCMAKE_VISIBILITY_INLINES_HIDDEN=ON",
             f"-DCMAKE_C_VISIBILITY_PRESET=hidden",
             f"-DCMAKE_CXX_VISIBILITY_PRESET=hidden",
-            f"-DTORCH_MLIR_ENABLE_LTC={'ON' if enable_ltc else 'OFF'}",
+            f"-DTORCH_MLIR_ENABLE_LTC={'ON' if TORCH_MLIR_ENABLE_LTC else 'OFF'}",
             f"-DTORCH_MLIR_ENABLE_PYTORCH_EXTENSIONS={'OFF' if TORCH_MLIR_ENABLE_ONLY_MLIR_PYTHON_BINDINGS else 'ON'}",
         ]
         if LLVM_INSTALL_DIR:
@@ -117,7 +143,7 @@ class CMakeBuild(build_py):
             f"--target",
             f"TorchMLIRPythonModules",
             f"--",
-            f"-j{max_jobs}"
+            f"-j{MAX_JOBS}",
         ]
         try:
             subprocess.check_call(cmake_config_args, cwd=cmake_build_dir)
@@ -131,22 +157,22 @@ class CMakeBuild(build_py):
             print(f"cmake build: {' '.join(cmake_build_args)}")
             print(f"cmake workspace: {cmake_build_dir}")
 
-
     def run(self):
         target_dir = self.build_lib
-        cmake_build_dir = os.getenv("TORCH_MLIR_CMAKE_BUILD_DIR")
+        cmake_build_dir = TORCH_MLIR_CMAKE_BUILD_DIR
         if not cmake_build_dir:
             cmake_build_dir = os.path.abspath(
-                os.path.join(target_dir, "..", "cmake_build"))
+                os.path.join(target_dir, "..", "cmake_build")
+            )
         if LLVM_INSTALL_DIR:
-            python_package_dir = os.path.join(cmake_build_dir,
+            python_package_dir = os.path.join(
-                                              "python_packages",
+                cmake_build_dir, "python_packages", "torch_mlir"
-                                              "torch_mlir")
+            )
         else:
-            python_package_dir = os.path.join(cmake_build_dir,
+            python_package_dir = os.path.join(
-                                              "tools", "torch-mlir", "python_packages",
+                cmake_build_dir, "tools", "torch-mlir", "python_packages", "torch_mlir"
-                                              "torch_mlir")
+            )
-        if not os.getenv("TORCH_MLIR_CMAKE_BUILD_DIR_ALREADY_BUILT"):
+        if not TORCH_MLIR_CMAKE_ALREADY_BUILT:
             os.makedirs(cmake_build_dir, exist_ok=True)
             cmake_cache_file = os.path.join(cmake_build_dir, "CMakeCache.txt")
             if os.path.exists(cmake_cache_file):
@@ -169,16 +195,14 @@ class CMakeBuild(build_py):
         if os.path.exists(target_dir):
             shutil.rmtree(target_dir, ignore_errors=False, onerror=None)
 
-        shutil.copytree(python_package_dir,
+        shutil.copytree(python_package_dir, target_dir, symlinks=False)
-                        target_dir,
-                        symlinks=False)
 
 
 class CMakeExtension(Extension):
 
-  def __init__(self, name, sourcedir=""):
+    def __init__(self, name, sourcedir=""):
-    Extension.__init__(self, name, sources=[])
+        Extension.__init__(self, name, sources=[])
-    self.sourcedir = os.path.abspath(sourcedir)
+        self.sourcedir = os.path.abspath(sourcedir)
 
 
 class NoopBuildExtension(build_ext):