tools: Add script to do split build in one command (#141359)

Usage:
```bash
python3 tools/packaging/split_wheel.py bdist_wheel
python3 tools/packaging/split_wheel.py install
python3 tools/packaging/split_wheel.py develop
```
Ideally this should make it easier to do the split build locally while
we're doing development.

Signed-off-by: Eli Uriegas <eliuriegas@meta.com>
Pull Request resolved: https://github.com/pytorch/pytorch/pull/141359
Approved by: https://github.com/kit1980

.ci/pytorch/build.sh

@@ -285,8 +285,7 @@ else
       WERROR=1 python setup.py clean
 
       if [[ "$USE_SPLIT_BUILD" == "true" ]]; then
-        BUILD_LIBTORCH_WHL=1 BUILD_PYTHON_ONLY=0 python setup.py bdist_wheel
-        BUILD_LIBTORCH_WHL=0 BUILD_PYTHON_ONLY=1 python setup.py bdist_wheel --cmake
+        python3 tools/packaging/split_wheel.py bdist_wheel
       else
         WERROR=1 python setup.py bdist_wheel
       fi

tools/packaging/split_wheel.py

@@ -0,0 +1,105 @@
+"""Script to build split pytorch wheels
+
+What is split build / why is it important?
+    > Split build is splitting the PyTorch build into a libtorch &
+    > PyTorch python frontend package. This allows us to to publish
+    > both as separate packages and opens up our ability to have users
+    > install different libtorch backends per their PyTorch frontend
+    >
+    > Example: opening up the door to things like:
+    >     pip install torch[cuda]
+    >     pip install torch[rocm]
+    >     pip install torch[cpu]
+    >     etc.
+
+Why does this exist?
+    > Currently our split build requires you to invoke setup.py twice
+    > Which ends up complicating the build process and adds some level
+    > of complexity to our setup.py / build invocation for split builds.
+    > Ideally this script will eventually not be needed but for
+    > development purposes we should have an easy way to invoke this script
+"""
+
+import argparse
+import logging
+import os
+import subprocess
+import sys
+from pathlib import Path
+from typing import Dict, List, Optional
+
+
+logging.basicConfig(
+    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
+)
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.DEBUG)
+
+
+# NOTE: This will need to be updated if this script is ever moved
+ROOT_PATH = Path(__file__).absolute().parents[2]
+SETUP_PY_PATH = ROOT_PATH / "setup.py"
+
+
+def requirements_installed() -> bool:
+    try:
+        import setuptools  # type: ignore[import-untyped]  # noqa: F401
+
+        return True
+    except ImportError:
+        logger.error(
+            "Requirements not installed, run the following command to install:"
+        )
+        logger.error(
+            "    > %s -m pip install -r %s/requirements.txt", sys.executable, ROOT_PATH
+        )
+        return False
+
+
+def setup_py(cmd_args: List[str], extra_env: Optional[Dict[str, str]] = None) -> None:
+    if extra_env is None:
+        extra_env = {}
+    cmd = [sys.executable, str(SETUP_PY_PATH), *cmd_args]
+    logger.debug("+ %s", " ".join(cmd))
+    subprocess.run(
+        cmd,
+        # Give the parent environment to the subprocess
+        env={**os.environ, **extra_env},
+        check=True,
+    )
+
+
+def split_build(cmd: str) -> None:
+    logger.info("Running %s for libtorch wheel", cmd)
+    setup_py(
+        [cmd],
+        extra_env={"BUILD_LIBTORCH_WHL": "1", "BUILD_PYTHON_ONLY": "0"},
+    )
+    logger.info("Running %s for torch wheel", cmd)
+    # NOTE: Passing --cmake is necessary here since the torch frontend has it's
+    # own cmake files that it needs to generate
+    setup_py(
+        [cmd, "--cmake"],
+        extra_env={"BUILD_LIBTORCH_WHL": "0", "BUILD_PYTHON_ONLY": "1"},
+    )
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    command_subparser = parser.add_subparsers(dest="command")
+    # Ideally these should mirror setuptools commands if we need support here for that
+    command_subparser.add_parser("install")
+    command_subparser.add_parser("bdist_wheel")
+    command_subparser.add_parser("develop")
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    if not requirements_installed():
+        sys.exit(1)
+    split_build(args.command)
+
+
+if __name__ == "__main__":
+    main()