Enable Monarch extra builds

Summary:
This is the first step to supporting individual installs for Monarch capabilities, i.e. core/rdma/tensor_engine.

Changes:
- Modifies monarch_extension to use core, rdma, tensor_engine, full feature sets.
- Switches setup.py to use MONARCH_FEATURES accepting strings rather than the 0/1 boolean only USE_TENSOR_ENGINE (also makes all the downstream changes needed to reflect it)
- Renames the default wheel name from `monarch` to `torchmonarch`, to better align with what we're officially publishing to PyPI
- Slight clean up to setup.py workflow

Differential Revision: D87884695

Author: allenwang28

.github/workflows/build-cpu.yml

@@ -34,4 +34,4 @@ jobs:
         pip install -r build-requirements.txt
 
         # Build monarch (No tensor engine, CPU version)
-        USE_TENSOR_ENGINE=0 python setup.py bdist_wheel
+        MONARCH_FEATURES=core python setup.py bdist_wheel -v

.github/workflows/build-cuda.yml

@@ -43,5 +43,5 @@ jobs:
 
         export CUDA_LIB_DIR=/usr/lib64
 
-        # Build monarch (CUDA version)
-        python setup.py bdist_wheel
+        # Build monarch (CUDA version with all features)
+        MONARCH_FEATURES=full python setup.py bdist_wheel -v

.github/workflows/doc_build.yml

@@ -43,7 +43,6 @@ jobs:
 
         # Set environment variables for CUDA build
         export USE_CUDA=1
-        export USE_TENSOR_ENGINE=1
         export RUSTFLAGS="-Zthreads=16 ${RUSTFLAGS:-}"
         export _GLIBCXX_USE_CXX11_ABI=1
         export CUDA_LIB_DIR=/usr/lib64

.github/workflows/publish_release.yml

@@ -50,7 +50,7 @@ jobs:
         export MONARCH_PACKAGE_NAME="torchmonarch"
         export CUDA_LIB_DIR=/usr/lib64
         export MONARCH_VERSION="${{ github.event.inputs.version }}"
-        python setup.py bdist_wheel
+        python setup.py bdist_wheel -v
 
         # hacky until the right distribution wheel can be made...
         find dist -name "*linux_x86_64.whl" -type f -exec bash -c 'mv "$1" "${1/linux_x86_64.whl/manylinux2014_x86_64.whl}"' _ {} \;

.github/workflows/test-cpu-python.yml

@@ -28,9 +28,6 @@ jobs:
         # Setup test environment
         setup_conda_environment
 
-        # Disable tensor engine
-        export USE_TENSOR_ENGINE=0
-
         # Install PyTorch nightly
         pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cpu
 

.github/workflows/wheels.yml

@@ -51,7 +51,7 @@ jobs:
         export MONARCH_VERSION=$(date +'%Y.%m.%d')
         export CUDA_LIB_DIR=/usr/lib64
 
-        python setup.py bdist_wheel
+        MONARCH_FEATURES=full python setup.py bdist_wheel -v
 
         # hacky until the right distribution wheel can be made...
         find dist -name "*linux_x86_64.whl" -type f -exec bash -c 'mv "$1" "${1/linux_x86_64.whl/manylinux2014_x86_64.whl}"' _ {} \;

.gitignore

@@ -9,7 +9,7 @@ python/monarch.egg-info/*
 *.egg
 build/*
 dist/*
-monarch.egg-info/*
+torchmonarch.egg-info/*
 
 .ipynb_checkpoints
 .monarch

README.md

@@ -98,7 +98,7 @@ sudo dnf install clang-devel libnccl-devel
 conda install -c conda-forge clangdev nccl
 conda update -n monarchenv --all -c conda-forge -y
 
-# If you are building with RDMA support, build monarch with `USE_TENSOR_ENGINE=1 pip install --no-build-isolation .` and dnf install the following packages
+# If you are building with RDMA support, build monarch with `MONARCH_FEATURES=rdma pip install --no-build-isolation .` and dnf install the following packages
 sudo dnf install -y libibverbs rdma-core libmlx5 libibverbs-devel rdma-core-devel
 
 # Install build dependencies
@@ -147,12 +147,12 @@ pip install -r build-requirements.txt
 # Install test dependencies
 pip install -r python/tests/requirements.txt
 
-# Build and install Monarch (with tensor engine support)
+# Build and install Monarch (with all features)
 pip install --no-build-isolation .
 
 # or
 # Build and install Monarch (without tensor engine support)
-USE_TENSOR_ENGINE=0 pip install --no-build-isolation .
+MONARCH_FEATURES=core pip install --no-build-isolation .
 
 # or setup for development
 pip install --no-build-isolation -e .
@@ -185,10 +185,10 @@ pip install -r build-requirements.txt
 # Install test dependencies
 pip install -r python/tests/requirements.txt
 
-# Build and install Monarch
-USE_TENSOR_ENGINE=0 pip install --no-build-isolation .
+# Build and install Monarch (core only, no RDMA or tensor engine)
+MONARCH_FEATURES=core pip install --no-build-isolation .
 # or setup for development
-USE_TENSOR_ENGINE=0 pip install --no-build-isolation -e .
+MONARCH_FEATURES=core pip install --no-build-isolation -e .
 
 # Verify installation
 pip list | grep monarch

monarch_extension/Cargo.toml

@@ -40,5 +40,8 @@ torch-sys-cuda = { version = "0.0.0", path = "../torch-sys-cuda", optional = tru
 tracing = { version = "0.1.41", features = ["attributes", "valuable"] }
 
 [features]
-default = ["tensor_engine"]
-tensor_engine = ["dep:monarch_messages", "dep:monarch_rdma_extension", "dep:monarch_tensor_worker", "dep:nccl-sys", "dep:rdmaxcel-sys", "dep:torch-sys", "dep:torch-sys-cuda"]
+core = []
+default = ["full"]
+full = ["rdma", "tensor_engine"]
+rdma = ["dep:monarch_rdma_extension", "dep:rdmaxcel-sys"]
+tensor_engine = ["dep:monarch_messages", "dep:monarch_tensor_worker", "dep:nccl-sys", "dep:torch-sys", "dep:torch-sys-cuda"]

monarch_extension/src/lib.rs

@@ -122,6 +122,10 @@ pub fn mod_init(module: &Bound<'_, PyModule>) -> PyResult<()> {
             module,
             "monarch_extension.mesh_controller",
         )?)?;
+    }
+
+    #[cfg(feature = "rdma")]
+    {
         monarch_rdma_extension::register_python_bindings(&get_or_add_new_module(module, "rdma")?)?;
     }
     simulation_tools::register_python_bindings(&get_or_add_new_module(