Remove usage of environment variables in wheel building (#9583)

### Summary

We seem to be using a combination of CMAKE_ARGS and environment
variables when creating wheels. Ultimately, CMake only uses the cmake
args, however we redefine some of these flags as env vars to help
`setup.py` determine if a certain feature is turned on. Specifically, it
looks for pybinding vars to bundle pybindings.

Let's remove this redundancy and just use the CMAKE_ARGS as the single
source of truth. For more details and other considerations, see
#9494 (abandoned).

Note that even in the wheel building jobs, we use cmake args instead of
environment variables to control features:


https://github.com/pytorch/executorch/blob/644b7ddf14180d97e348faa627f576e13d367d69/.ci/scripts/wheel/envvar_base.sh#L20


https://github.com/pytorch/executorch/blob/644b7ddf14180d97e348faa627f576e13d367d69/.ci/scripts/wheel/envvar_macos.sh#L14-L15

### Test plan

build + check CMakeCache.txt to ensure flags are set

```bash
# Expected: EXECUTORCH_BUILD_PYBIND=OFF EXECUTORCH_BUILD_XNNPACK=OFF EXECUTORCH_BUILD_COREML=OFF
$ rm -rf pip-out dist && ./install_executorch.sh --pybind off

# Expected: EXECUTORCH_BUILD_PYBIND=ON EXECUTORCH_BUILD_XNNPACK=ON EXECUTORCH_BUILD_COREML=OFF
$ rm -rf pip-out dist && ./install_executorch.sh

# Expected: EXECUTORCH_BUILD_PYBIND=ON EXECUTORCH_BUILD_XNNPACK=OFF EXECUTORCH_BUILD_COREML=ON
$ rm -rf pip-out dist && ./install_executorch.sh --pybind coreml

# Expected: EXECUTORCH_BUILD_PYBIND=ON EXECUTORCH_BUILD_XNNPACK=ON EXECUTORCH_BUILD_COREML=ON
$ rm -rf pip-out dist && ./install_executorch.sh --pybind xnnpack coreml

# Throws an error
$ rm -rf pip-out dist && ./install_executorch.sh --pybind coreml off
```


cc @larryliu0820 @lucylq

Author: jathu

.ci/scripts/unittest-linux.sh

@@ -21,8 +21,7 @@ if [[ "$BUILD_TOOL" == "cmake" ]]; then
     source .ci/scripts/setup-vulkan-linux-deps.sh
 
     PYTHON_EXECUTABLE=python \
-    EXECUTORCH_BUILD_PYBIND=ON \
-    CMAKE_ARGS="-DEXECUTORCH_BUILD_XNNPACK=ON -DEXECUTORCH_BUILD_KERNELS_QUANTIZED=ON" \
+    CMAKE_ARGS="-DEXECUTORCH_BUILD_PYBIND=ON -DEXECUTORCH_BUILD_XNNPACK=ON -DEXECUTORCH_BUILD_KERNELS_QUANTIZED=ON" \
     .ci/scripts/setup-linux.sh "$@"
 
     # Install llama3_2_vision dependencies.

.ci/scripts/unittest-macos.sh

@@ -21,8 +21,7 @@ trap 'rm -rfv ${TMP_DIR}' EXIT
 
 # Setup MacOS dependencies as there is no Docker support on MacOS atm
 PYTHON_EXECUTABLE=python \
-EXECUTORCH_BUILD_PYBIND=ON \
-CMAKE_ARGS="-DEXECUTORCH_BUILD_COREML=ON -DEXECUTORCH_BUILD_MPS=ON -DEXECUTORCH_BUILD_XNNPACK=ON -DEXECUTORCH_BUILD_KERNELS_QUANTIZED=ON" \
+CMAKE_ARGS="-DEXECUTORCH_BUILD_PYBIND=ON -DEXECUTORCH_BUILD_COREML=ON -DEXECUTORCH_BUILD_MPS=ON -DEXECUTORCH_BUILD_XNNPACK=ON -DEXECUTORCH_BUILD_KERNELS_QUANTIZED=ON" \
 ${CONDA_RUN} --no-capture-output \
 .ci/scripts/setup-macos.sh "$@"
 

.ci/scripts/wheel/envvar_base.sh

@@ -8,13 +8,10 @@
 # should typically only contain shell variable assignments. Be sure to export
 # any variables so that subprocesses will see them.
 
-# Enable pybindings so that users can execute ExecuTorch programs from python.
-export EXECUTORCH_BUILD_PYBIND=1
-
 # Ensure that CMAKE_ARGS is defined before referencing it. Defaults to empty
 # if not defined.
 export CMAKE_ARGS="${CMAKE_ARGS:-}"
 
 # Link the XNNPACK backend into the pybindings runtime so that users can execute
 # ExecuTorch programs that delegate to it.
-CMAKE_ARGS="${CMAKE_ARGS} -DEXECUTORCH_BUILD_XNNPACK=ON"
+CMAKE_ARGS="${CMAKE_ARGS} -DEXECUTORCH_BUILD_PYBIND=ON -DEXECUTORCH_BUILD_XNNPACK=ON"

.github/workflows/pull.yml

@@ -365,8 +365,7 @@ jobs:
         # build module for executorch.extension.pybindings.portable_lib
         BUILD_TOOL="cmake"
         PYTHON_EXECUTABLE=python \
-        EXECUTORCH_BUILD_XNNPACK=ON \
-        EXECUTORCH_BUILD_PYBIND=ON \
+        CMAKE_ARGS="-DEXECUTORCH_BUILD_PYBIND=ON -DEXECUTORCH_BUILD_XNNPACK=ON" \
         bash .ci/scripts/setup-linux.sh --build-tool "${BUILD_TOOL}"
 
         # see if we can import the module successfully
@@ -504,7 +503,7 @@ jobs:
 
         # Setup MacOS dependencies as there is no Docker support on MacOS atm
         PYTHON_EXECUTABLE=python \
-        EXECUTORCH_BUILD_PYBIND=ON \
+        CMAKE_ARGS="-DEXECUTORCH_BUILD_PYBIND=ON" \
         EXECUTORCH_BUILD_ARM_BAREMETAL=ON \
         .ci/scripts/setup-linux.sh --build-tool "${BUILD_TOOL}"
 

.github/workflows/trunk.yml

@@ -261,7 +261,7 @@ jobs:
 
         # build module for executorch.extension.pybindings.portable_lib
         BUILD_TOOL=${{ matrix.build-tool }}
-        EXECUTORCH_BUILD_PYBIND=ON PYTHON_EXECUTABLE=python ${CONDA_RUN} bash .ci/scripts/setup-macos.sh --build-tool "${BUILD_TOOL}"
+        CMAKE_ARGS="-DEXECUTORCH_BUILD_PYBIND=ON" PYTHON_EXECUTABLE=python ${CONDA_RUN} bash .ci/scripts/setup-macos.sh --build-tool "${BUILD_TOOL}"
 
         # see if we can import the module successfully
         ${CONDA_RUN} python -c "from executorch.extension.pybindings import portable_lib; print('success!')"

install_executorch.py

@@ -15,6 +15,7 @@
 import subprocess
 import sys
 from contextlib import contextmanager
+from typing import List, Tuple
 
 from install_requirements import (
     install_requirements,
@@ -168,28 +169,30 @@ def build_args_parser() -> argparse.ArgumentParser:
     return parser
 
 
-def handle_pybind(args, cmake_args, executorch_build_pybind):
+# Returns (wants_off, wanted_pybindings)
+def _list_pybind_defines(args) -> Tuple[bool, List[str]]:
+    if args.pybind is None:
+        return False, []
+
     # Flatten list of lists.
     args.pybind = list(itertools.chain(*args.pybind))
     if "off" in args.pybind:
         if len(args.pybind) != 1:
             raise Exception(f"Cannot combine `off` with other pybinds: {args.pybind}")
-        executorch_build_pybind = "OFF"
-    else:
-        for pybind_arg in args.pybind:
-            if pybind_arg not in VALID_PYBINDS:
-                raise Exception(
-                    f"Unrecognized pybind argument {pybind_arg}; valid options are: {', '.join(VALID_PYBINDS)}"
-                )
-            if pybind_arg == "training":
-                cmake_args += " -DEXECUTORCH_BUILD_EXTENSION_TRAINING=ON"
-                os.environ["EXECUTORCH_BUILD_TRAINING"] = "ON"
-            elif pybind_arg == "mps":
-                cmake_args += " -DEXECUTORCH_BUILD_MPS=ON"
-            else:
-                cmake_args += f" -DEXECUTORCH_BUILD_{pybind_arg.upper()}=ON"
-            executorch_build_pybind = "ON"
-    return executorch_build_pybind, cmake_args
+        return True, []
+
+    cmake_args = []
+    for pybind_arg in args.pybind:
+        if pybind_arg not in VALID_PYBINDS:
+            raise Exception(
+                f"Unrecognized pybind argument {pybind_arg}; valid options are: {', '.join(VALID_PYBINDS)}"
+            )
+        if pybind_arg == "training":
+            cmake_args.append("-DEXECUTORCH_BUILD_EXTENSION_TRAINING=ON")
+        else:
+            cmake_args.append(f"-DEXECUTORCH_BUILD_{pybind_arg.upper()}=ON")
+
+    return False, cmake_args
 
 
 def main(args):
@@ -199,14 +202,23 @@ def main(args):
     parser = build_args_parser()
     args = parser.parse_args()
 
-    EXECUTORCH_BUILD_PYBIND = ""
-    CMAKE_ARGS = os.getenv("CMAKE_ARGS", "")
+    cmake_args = [os.getenv("CMAKE_ARGS", "")]
     use_pytorch_nightly = True
 
-    if args.pybind:
-        EXECUTORCH_BUILD_PYBIND, CMAKE_ARGS = handle_pybind(
-            args, CMAKE_ARGS, EXECUTORCH_BUILD_PYBIND
-        )
+    has_pybindings = False
+    wants_pybindings_off, pybind_defines = _list_pybind_defines(args)
+    if not wants_pybindings_off:
+        has_pybindings = True
+        if len(pybind_defines) > 0:
+            # If the user explicitly provides a list of bindings, just use them
+            cmake_args += pybind_defines
+        else:
+            # If the user has not set pybindings off but also has not provided
+            # a list, then turn on xnnpack by default
+            cmake_args.append("-DEXECUTORCH_BUILD_XNNPACK=ON")
+
+    if has_pybindings:
+        cmake_args.append("-DEXECUTORCH_BUILD_PYBIND=ON")
 
     if args.clean:
         clean()
@@ -218,18 +230,11 @@ def main(args):
         # latest PT commit otherwise
         use_pytorch_nightly = False
 
-    # If --pybind is not set explicitly for backends (e.g., --pybind xnnpack)
-    # or is not turned off explicitly (--pybind off)
-    # then install XNNPACK by default.
-    if EXECUTORCH_BUILD_PYBIND == "":
-        EXECUTORCH_BUILD_PYBIND = "ON"
-        CMAKE_ARGS += " -DEXECUTORCH_BUILD_XNNPACK=ON"
-
     # Use ClangCL on Windows.
     # ClangCL is an alias to Clang that configures it to work in an MSVC-compatible
     # mode. Using it on Windows to avoid compiler compatibility issues for MSVC.
     if os.name == "nt":
-        CMAKE_ARGS += " -T ClangCL"
+        cmake_args.append("-T ClangCL")
 
     #
     # Install executorch pip package. This also makes `flatc` available on the path.
@@ -238,8 +243,7 @@ def main(args):
     #
 
     # Set environment variables
-    os.environ["EXECUTORCH_BUILD_PYBIND"] = EXECUTORCH_BUILD_PYBIND
-    os.environ["CMAKE_ARGS"] = CMAKE_ARGS
+    os.environ["CMAKE_ARGS"] = " ".join(cmake_args)
 
     # Check if the required submodules are present and update them if not
     check_and_update_submodules()

setup.py

@@ -60,42 +60,70 @@
 
 from distutils import log
 from distutils.sysconfig import get_python_lib
+from functools import cache
 from pathlib import Path
-from typing import List, Optional
+from typing import Dict, List, Optional
 
 from setuptools import Extension, setup
 from setuptools.command.build import build
 from setuptools.command.build_ext import build_ext
 from setuptools.command.build_py import build_py
 
 
+@cache
+def _cmake_args_defines() -> Dict[str, str]:
+    result = {}
+
+    args = re.split(r"\s+", os.environ.get("CMAKE_ARGS", ""))
+    for arg in args:
+        if arg.startswith("-D") and "=" in arg:
+            arg_key, value = arg.split("=")
+            key = arg_key[2:]  # Remove the leading "-D"
+            result[key] = value
+
+    return result
+
+
 class ShouldBuild:
     """Indicates whether to build various components."""
 
     @staticmethod
-    def _is_env_enabled(env_var: str, default: bool = False) -> bool:
-        val = os.environ.get(env_var, None)
-        if val is None:
-            return default
-        if val in ("OFF", "0", ""):
+    def _is_truthy(value: Optional[str]) -> bool:
+        if (value is None) or (value.lower() in ("off", "0", "")):
             return False
         return True
 
+    @staticmethod
+    def _is_cmake_arg_enabled(var: str, default: bool) -> bool:
+        if os.environ.get(var) is not None:
+            raise RuntimeError(
+                f"Python wheel building does not support setting '{var}' using environment variables. Use CMAKE_ARGS='-D{var}=ON' instead."
+            )
+
+        value = _cmake_args_defines().get(var, None)
+        if value is None:
+            return default
+        return ShouldBuild._is_truthy(value)
+
     @classmethod
     def pybindings(cls) -> bool:
-        return cls._is_env_enabled("EXECUTORCH_BUILD_PYBIND", default=False)
+        return cls._is_cmake_arg_enabled("EXECUTORCH_BUILD_PYBIND", default=False)
 
     @classmethod
     def training(cls) -> bool:
-        return cls._is_env_enabled("EXECUTORCH_BUILD_TRAINING", default=False)
+        return cls._is_cmake_arg_enabled(
+            "EXECUTORCH_BUILD_EXTENSION_TRAINING", default=False
+        )
 
     @classmethod
     def llama_custom_ops(cls) -> bool:
-        return cls._is_env_enabled("EXECUTORCH_BUILD_KERNELS_CUSTOM_AOT", default=True)
+        return cls._is_cmake_arg_enabled(
+            "EXECUTORCH_BUILD_KERNELS_CUSTOM_AOT", default=True
+        )
 
     @classmethod
     def flatc(cls) -> bool:
-        return cls._is_env_enabled("EXECUTORCH_BUILD_FLATC", default=True)
+        return cls._is_cmake_arg_enabled("EXECUTORCH_BUILD_FLATC", default=True)
 
 
 class Version:
@@ -667,9 +695,6 @@ def run(self):
                 "-DEXECUTORCH_BUILD_KERNELS_QUANTIZED_AOT=ON",
             ]
             if ShouldBuild.training():
-                cmake_args += [
-                    "-DEXECUTORCH_BUILD_EXTENSION_TRAINING=ON",
-                ]
                 build_args += ["--target", "_training_lib"]
             build_args += ["--target", "portable_lib"]
             # To link backends into the portable_lib target, callers should