build monarch_extension as extension-module

Differential Revision: D84956118

Author: colin2328

controller/Cargo.toml

@@ -27,7 +27,7 @@ hyperactor_multiprocess = { version = "0.0.0", path = "../hyperactor_multiproces
 monarch_messages = { version = "0.0.0", path = "../monarch_messages" }
 nccl-sys = { path = "../nccl-sys" }
 ndslice = { version = "0.0.0", path = "../ndslice" }
-pyo3 = { version = "0.24", features = ["anyhow", "multiple-pymethods", "py-clone"] }
+pyo3 = { version = "0.24", features = ["anyhow", "extension-module", "multiple-pymethods", "py-clone"] }
 serde = { version = "1.0.219", features = ["derive", "rc"] }
 serde_json = { version = "1.0.140", features = ["alloc", "float_roundtrip", "raw_value", "unbounded_depth"] }
 tokio = { version = "1.47.1", features = ["full", "test-util", "tracing"] }

monarch_extension/Cargo.toml

@@ -34,7 +34,7 @@ monarch_tensor_worker = { version = "0.0.0", path = "../monarch_tensor_worker",
 monarch_types = { version = "0.0.0", path = "../monarch_types" }
 nccl-sys = { path = "../nccl-sys", optional = true }
 ndslice = { version = "0.0.0", path = "../ndslice" }
-pyo3 = { version = "0.24", features = ["anyhow", "multiple-pymethods", "py-clone"] }
+pyo3 = { version = "0.24", features = ["anyhow", "extension-module", "multiple-pymethods", "py-clone"] }
 rdmaxcel-sys = { path = "../rdmaxcel-sys", optional = true }
 serde = { version = "1.0.219", features = ["derive", "rc"] }
 tokio = { version = "1.47.1", features = ["full", "test-util", "tracing"] }

monarch_hyperactor/Cargo.toml

@@ -40,7 +40,7 @@ ndslice = { version = "0.0.0", path = "../ndslice" }
 nix = { version = "0.30.1", features = ["dir", "event", "hostname", "inotify", "ioctl", "mman", "mount", "net", "poll", "ptrace", "reboot", "resource", "sched", "signal", "term", "time", "user", "zerocopy"] }
 once_cell = "1.21"
 opentelemetry = "0.29"
-pyo3 = { version = "0.24", features = ["anyhow", "multiple-pymethods", "py-clone"] }
+pyo3 = { version = "0.24", features = ["anyhow", "extension-module", "multiple-pymethods", "py-clone"] }
 pyo3-async-runtimes = { version = "0.24", features = ["attributes", "tokio-runtime"] }
 serde = { version = "1.0.219", features = ["derive", "rc"] }
 serde_bytes = "0.11"

monarch_messages/Cargo.toml

@@ -14,7 +14,7 @@ enum-as-inner = "0.6.0"
 hyperactor = { version = "0.0.0", path = "../hyperactor" }
 monarch_types = { version = "0.0.0", path = "../monarch_types" }
 ndslice = { version = "0.0.0", path = "../ndslice" }
-pyo3 = { version = "0.24", features = ["anyhow", "multiple-pymethods", "py-clone"] }
+pyo3 = { version = "0.24", features = ["anyhow", "extension-module", "multiple-pymethods", "py-clone"] }
 serde = { version = "1.0.219", features = ["derive", "rc"] }
 serde_bytes = "0.11"
 thiserror = "2.0.12"

monarch_rdma/extension/Cargo.toml

@@ -18,7 +18,7 @@ hyperactor = { version = "0.0.0", path = "../../hyperactor" }
 hyperactor_mesh = { version = "0.0.0", path = "../../hyperactor_mesh" }
 monarch_hyperactor = { version = "0.0.0", path = "../../monarch_hyperactor" }
 monarch_rdma = { version = "0.0.0", path = ".." }
-pyo3 = { version = "0.24", features = ["anyhow", "multiple-pymethods", "py-clone"] }
+pyo3 = { version = "0.24", features = ["anyhow", "extension-module", "multiple-pymethods", "py-clone"] }
 serde = { version = "1.0.219", features = ["derive", "rc"] }
 serde_json = { version = "1.0.140", features = ["alloc", "float_roundtrip", "raw_value", "unbounded_depth"] }
 tracing = { version = "0.1.41", features = ["attributes", "valuable"] }

monarch_tensor_worker/Cargo.toml

@@ -26,7 +26,7 @@ monarch_types = { version = "0.0.0", path = "../monarch_types" }
 ndslice = { version = "0.0.0", path = "../ndslice" }
 nix = { version = "0.30.1", features = ["dir", "event", "hostname", "inotify", "ioctl", "mman", "mount", "net", "poll", "ptrace", "reboot", "resource", "sched", "signal", "term", "time", "user", "zerocopy"] }
 parking_lot = { version = "0.12.1", features = ["send_guard"] }
-pyo3 = { version = "0.24", features = ["anyhow", "multiple-pymethods", "py-clone"] }
+pyo3 = { version = "0.24", features = ["anyhow", "extension-module", "multiple-pymethods", "py-clone"] }
 serde = { version = "1.0.219", features = ["derive", "rc"] }
 serde_json = { version = "1.0.140", features = ["alloc", "float_roundtrip", "raw_value", "unbounded_depth"] }
 sorted-vec = "0.8.3"

monarch_types/Cargo.toml

@@ -10,7 +10,7 @@ license = "BSD-3-Clause"
 [dependencies]
 derive_more = { version = "1.0.0", features = ["full"] }
 hyperactor = { version = "0.0.0", path = "../hyperactor" }
-pyo3 = { version = "0.24", features = ["anyhow", "multiple-pymethods", "py-clone"] }
+pyo3 = { version = "0.24", features = ["anyhow", "extension-module", "multiple-pymethods", "py-clone"] }
 serde = { version = "1.0.219", features = ["derive", "rc"] }
 serde_bytes = "0.11"
 

setup.py

@@ -119,6 +119,13 @@ def run(self):
 with open("README.md", encoding="utf8") as f:
     readme = f.read()
 
+python_lib_dir = sysconfig.get_config_var("LIBDIR")  # path to libpython3.10.so.1.0
+
+rust_binding_args = ["--"]
+if not USE_TENSOR_ENGINE:
+    rust_binding_args.append("--no-default-features")
+rust_binding_args += [f"-C", f"link-args=-Wl,-rpath,{python_lib_dir}"]
+
 rust_extensions = [
     RustBin(
         target="process_allocator",
@@ -131,7 +138,7 @@ def run(self):
         path="monarch_extension/Cargo.toml",
         debug=False,
         features=["tensor_engine"] if USE_TENSOR_ENGINE else [],
-        args=[] if USE_TENSOR_ENGINE else ["--no-default-features"],
+        args=rust_binding_args,
     ),
 ]
 

torch-sys-cuda/Cargo.toml

@@ -15,7 +15,7 @@ cxx = "1.0.119"
 derive_more = { version = "1.0.0", features = ["full"] }
 fxhash = "0.2.1"
 nccl-sys = { path = "../nccl-sys" }
-pyo3 = { version = "0.24", features = ["anyhow", "multiple-pymethods", "py-clone"] }
+pyo3 = { version = "0.24", features = ["anyhow", "extension-module", "multiple-pymethods", "py-clone"] }
 serde = { version = "1.0.219", features = ["derive", "rc"] }
 thiserror = "2.0.12"
 torch-sys = { version = "0.0.0", path = "../torch-sys" }

torch-sys/Cargo.toml

@@ -19,7 +19,7 @@ hyperactor = { version = "0.0.0", path = "../hyperactor" }
 monarch_types = { version = "0.0.0", path = "../monarch_types" }
 nccl-sys = { path = "../nccl-sys", optional = true }
 paste = "1.0.14"
-pyo3 = { version = "0.24", features = ["anyhow", "multiple-pymethods", "py-clone"] }
+pyo3 = { version = "0.24", features = ["anyhow", "extension-module", "multiple-pymethods", "py-clone"] }
 regex = "1.11.1"
 serde = { version = "1.0.219", features = ["derive", "rc"] }
 thiserror = "2.0.12"