Skip to content

Commit 6c8a85f

Browse files
zdevitofacebook-github-bot
zdevito
authored and
facebook-github-bot
committed
remove _monarch files that are not needed
Summary: get rid of the unneeded redirects, fix more cases where the rust module name didn't match the python name Reviewed By: colin2328 Differential Revision: D75016551 fbshipit-source-id: 7eb37cd399764c91f597ad046f86c259119afddb
1 parent a350349 commit 6c8a85f

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

60 files changed

+682
-681
lines changed

‎controller/src/bootstrap.rs

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -64,7 +64,7 @@ static LABEL_NAME_WORKER: &str = concatcp!(WORLD_LABEL_PREFIX, "worker");
6464
static LABEL_NAME_CONTROLLER_ACTOR_ID: &str = concatcp!(WORLD_LABEL_PREFIX, "controllerActorId");
6565

6666
#[derive(Clone, Debug, Serialize, Deserialize, Args)]
67-
#[pyclass]
67+
#[pyclass(module = "monarch._rust_bindings.controller.bootstrap")]
6868
pub struct ControllerCommand {
6969
/// The worker world to create
7070
#[arg(long)]
@@ -180,7 +180,7 @@ impl ControllerCommand {
180180
/// The ones for System / Host should probably be moved to the hyperactor
181181
/// multiprocess crate.
182182
#[derive(Clone, Debug, Serialize, Deserialize, Subcommand)]
183-
#[pyclass]
183+
#[pyclass(module = "monarch._rust_bindings.controller.bootstrap")]
184184
pub enum RunCommand {
185185
System {
186186
/// The system address to bootstrap with.
@@ -219,7 +219,7 @@ pub enum RunCommand {
219219
Controller(ControllerCommand),
220220
}
221221

222-
#[pyclass(frozen, module = "monarch._monarch.worker")]
222+
#[pyclass(frozen, module = "monarch._rust_bindings.controller.bootstrap")]
223223
#[derive(Debug, Serialize, Deserialize)]
224224
pub enum ControllerServerRequest {
225225
Run(RunCommand),
@@ -237,7 +237,7 @@ impl ControllerServerRequest {
237237
}
238238
}
239239

240-
#[pyclass(frozen, module = "monarch._monarch.worker")]
240+
#[pyclass(frozen, module = "monarch._rust_bindings.controller.bootstrap")]
241241
#[derive(Debug, Serialize, Deserialize)]
242242
pub enum ControllerServerResponse {
243243
Finished { error: Option<String> },
@@ -641,3 +641,11 @@ pub fn parse_key_val(s: &str) -> anyhow::Result<(String, String)> {
641641
Some((a, b)) => Ok((a.to_owned(), b.to_owned())),
642642
}
643643
}
644+
645+
pub fn register_python_bindings(controller_mod: &Bound<'_, PyModule>) -> PyResult<()> {
646+
controller_mod.add_class::<ControllerServerRequest>()?;
647+
controller_mod.add_class::<ControllerServerResponse>()?;
648+
controller_mod.add_class::<RunCommand>()?;
649+
controller_mod.add_class::<ControllerCommand>()?;
650+
Ok(())
651+
}

‎hyperactor_extension/src/alloc.rs

Lines changed: 20 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,10 @@ use pyo3::types::PyDict;
1919
/// A python class that wraps a Rust Alloc trait object. It represents what
2020
/// is shown on the python side. Internals are not exposed.
2121
/// It ensures that the Alloc is only used once (i.e. moved) in rust.
22-
#[pyclass(name = "Alloc", module = "monarch._monarch.hyperactor")]
22+
#[pyclass(
23+
name = "Alloc",
24+
module = "monarch._rust_bindings.hyperactor_extension.alloc"
25+
)]
2326
pub struct PyAlloc {
2427
pub inner: Arc<Mutex<Option<PyAllocWrapper>>>,
2528
}
@@ -67,7 +70,10 @@ impl Alloc for PyAllocWrapper {
6770
}
6871
}
6972

70-
#[pyclass(name = "AllocConstraints", module = "monarch._monarch.hyperactor")]
73+
#[pyclass(
74+
name = "AllocConstraints",
75+
module = "monarch._rust_bindings.hyperactor_extension.alloc"
76+
)]
7177
pub struct PyAllocConstraints {
7278
inner: AllocConstraints,
7379
}
@@ -86,7 +92,10 @@ impl PyAllocConstraints {
8692
}
8793
}
8894

89-
#[pyclass(name = "AllocSpec", module = "monarch._monarch.hyperactor")]
95+
#[pyclass(
96+
name = "AllocSpec",
97+
module = "monarch._rust_bindings.hyperactor_extension.alloc"
98+
)]
9099
pub struct PyAllocSpec {
91100
pub inner: AllocSpec,
92101
}
@@ -131,3 +140,11 @@ impl PyAllocSpec {
131140
})
132141
}
133142
}
143+
144+
pub fn register_python_bindings(module: &Bound<'_, PyModule>) -> PyResult<()> {
145+
module.add_class::<PyAlloc>()?;
146+
module.add_class::<PyAllocConstraints>()?;
147+
module.add_class::<PyAllocSpec>()?;
148+
149+
Ok(())
150+
}

‎hyperactor_extension/src/lib.rs

Lines changed: 0 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,3 @@
11
#![allow(unsafe_op_in_unsafe_fn)]
22

33
pub mod alloc;
4-
5-
use pyo3::Bound;
6-
use pyo3::PyResult;
7-
use pyo3::prelude::*;
8-
use pyo3::types::PyModule;
9-
10-
use crate::alloc::PyAlloc;
11-
use crate::alloc::PyAllocConstraints;
12-
use crate::alloc::PyAllocSpec;
13-
14-
pub fn register_python_bindings(module: &Bound<'_, PyModule>) -> PyResult<()> {
15-
module.add_class::<PyAlloc>()?;
16-
module.add_class::<PyAllocConstraints>()?;
17-
module.add_class::<PyAllocSpec>()?;
18-
19-
Ok(())
20-
}

‎monarch_extension/src/client.rs

Lines changed: 23 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ use pyo3::types::PyNone;
3636
use tokio::sync::Mutex;
3737
use torch_sys::RValue;
3838

39-
#[pyclass(frozen, module = "monarch._monarch.client")]
39+
#[pyclass(frozen, module = "monarch._rust_bindings.monarch_extension.client")]
4040
struct WorkerResponse {
4141
seq: Seq,
4242
result: Option<Result<Serialized, Exception>>,
@@ -111,7 +111,11 @@ impl WorkerResponse {
111111
}
112112
}
113113

114-
#[pyclass(frozen, name = "WorldState", module = "monarch._monarch.client")]
114+
#[pyclass(
115+
frozen,
116+
name = "WorldState",
117+
module = "monarch._rust_bindings.monarch_extension.client"
118+
)]
115119
pub struct PyWorldState {
116120
inner: WorldSnapshot,
117121
}
@@ -137,7 +141,7 @@ impl PyWorldState {
137141
#[pyclass(
138142
frozen,
139143
name = "SystemSnapshotFilter",
140-
module = "monarch._monarch.client"
144+
module = "monarch._rust_bindings.monarch_extension.client"
141145
)]
142146
pub struct PySystemSnapshotFilter {
143147
inner: SystemSnapshotFilter,
@@ -217,7 +221,11 @@ impl From<PySystemSnapshotFilter> for SystemSnapshotFilter {
217221
}
218222
}
219223

220-
#[pyclass(frozen, name = "ProcInfo", module = "monarch._monarch.client")]
224+
#[pyclass(
225+
frozen,
226+
name = "ProcInfo",
227+
module = "monarch._rust_bindings.monarch_extension.client"
228+
)]
221229
pub struct PyProcInfo {
222230
inner: WorldSnapshotProcInfo,
223231
}
@@ -240,7 +248,7 @@ impl PyProcInfo {
240248
frozen,
241249
subclass,
242250
name = "Exception",
243-
module = "monarch._monarch.client"
251+
module = "monarch._rust_bindings.monarch_extension.client"
244252
)]
245253
pub struct PyException {
246254
inner: Exception,
@@ -260,7 +268,7 @@ impl PyException {
260268
}
261269
}
262270

263-
#[pyclass(frozen, extends = PyException, subclass, name = "Error", module = "monarch._monarch.client")]
271+
#[pyclass(frozen, extends = PyException, subclass, name = "Error", module = "monarch._rust_bindings.monarch_extension.client")]
264272
pub struct PyError;
265273

266274
#[pymethods]
@@ -345,7 +353,7 @@ impl PyError {
345353
eq,
346354
eq_int,
347355
name = "LogLevel",
348-
module = "monarch._monarch.client"
356+
module = "monarch._rust_bindings.monarch_extension.client"
349357
)]
350358
enum PyLogLevel {
351359
Info,
@@ -375,7 +383,11 @@ impl PyLogLevel {
375383
const INFO: PyLogLevel = PyLogLevel::Info;
376384
}
377385

378-
#[pyclass(frozen, name = "LogMessage", module = "monarch._monarch.client")]
386+
#[pyclass(
387+
frozen,
388+
name = "LogMessage",
389+
module = "monarch._rust_bindings.monarch_extension.client"
390+
)]
379391
pub struct LogMessage {
380392
level: PyLogLevel,
381393
message: String,
@@ -405,7 +417,7 @@ impl LogMessage {
405417
}
406418
}
407419

408-
#[pyclass(frozen, extends = PyException, subclass, name = "Failure", module = "monarch._monarch.client")]
420+
#[pyclass(frozen, extends = PyException, subclass, name = "Failure", module = "monarch._rust_bindings.monarch_extension.client")]
409421
pub struct PyFailure;
410422

411423
#[pymethods]
@@ -473,7 +485,7 @@ impl PyFailure {
473485
frozen,
474486
get_all,
475487
name = "DebuggerMessage",
476-
module = "monarch._monarch.client"
488+
module = "monarch._rust_bindings.monarch_extension.client"
477489
)]
478490
pub struct DebuggerMessage {
479491
debugger_actor_id: PyActorId,
@@ -492,7 +504,7 @@ impl DebuggerMessage {
492504
}
493505
}
494506

495-
#[pyclass(module = "monarch._monarch.client")]
507+
#[pyclass(module = "monarch._rust_bindings.monarch_extension.client")]
496508
pub struct ClientActor {
497509
instance: Arc<Mutex<InstanceWrapper<ClientMessage>>>,
498510
}

‎monarch_extension/src/controller.rs

Lines changed: 6 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,3 @@
1-
use controller::bootstrap::ControllerCommand;
2-
use controller::bootstrap::ControllerServerRequest;
3-
use controller::bootstrap::ControllerServerResponse;
4-
use controller::bootstrap::RunCommand;
51
/// These the controller messages that are exposed to python to allow the client to construct and
62
/// send messages to the controller. For more details of the definitions take a look at
73
/// [`monarch_messages::controller::ControllerMessage`].
@@ -17,7 +13,11 @@ use pyo3::prelude::*;
1713

1814
use crate::worker::PyWorkerMessage;
1915

20-
#[pyclass(frozen, get_all, module = "monarch._monarch.controller")]
16+
#[pyclass(
17+
frozen,
18+
get_all,
19+
module = "monarch._rust_bindings.monarch_extension.controller"
20+
)]
2121
struct Node {
2222
seq: Seq,
2323
defs: Vec<Ref>,
@@ -59,7 +59,7 @@ enum PyRanks {
5959
SliceList(Vec<PySlice>),
6060
}
6161

62-
#[pyclass(frozen, module = "monarch._monarch.controller")]
62+
#[pyclass(frozen, module = "monarch._rust_bindings.monarch_extension.controller")]
6363
struct Send {
6464
ranks: Ranks,
6565
message: Serialized,
@@ -124,9 +124,5 @@ impl Send {
124124
pub(crate) fn register_python_bindings(controller_mod: &Bound<'_, PyModule>) -> PyResult<()> {
125125
controller_mod.add_class::<Node>()?;
126126
controller_mod.add_class::<Send>()?;
127-
controller_mod.add_class::<ControllerServerRequest>()?;
128-
controller_mod.add_class::<ControllerServerResponse>()?;
129-
controller_mod.add_class::<RunCommand>()?;
130-
controller_mod.add_class::<ControllerCommand>()?;
131127
Ok(())
132128
}

‎monarch_extension/src/debugger.rs

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ use tokio::sync::Mutex;
2222
frozen,
2323
get_all,
2424
name = "DebuggerMessage",
25-
module = "monarch._monarch.debugger"
25+
module = "monarch._rust_bindings.monarch_extension.debugger"
2626
)]
2727
pub struct PyDebuggerMessage {
2828
action: DebuggerAction,
@@ -61,7 +61,7 @@ pub fn get_bytes_from_write_action(
6161
}
6262
}
6363

64-
#[pyclass(module = "monarch._monarch.debugger")]
64+
#[pyclass(module = "monarch._rust_bindings.monarch_extension.debugger")]
6565
pub struct PdbActor {
6666
instance: Arc<Mutex<InstanceWrapper<DebuggerMessage>>>,
6767
controller_actor_ref: ActorRef<ControllerActor>,
@@ -126,7 +126,12 @@ impl PdbActor {
126126
pub fn register_python_bindings(debugger: &Bound<'_, PyModule>) -> PyResult<()> {
127127
debugger.add_class::<PdbActor>()?;
128128
debugger.add_class::<PyDebuggerMessage>()?;
129-
debugger.add_function(wrap_pyfunction!(get_bytes_from_write_action, debugger)?)?;
129+
let f = wrap_pyfunction!(get_bytes_from_write_action, debugger)?;
130+
f.setattr(
131+
"__module__",
132+
"monarch._rust_bindings.monarch_extension.debugger",
133+
)?;
134+
debugger.add_function(f)?;
130135
Ok(())
131136
}
132137

‎monarch_extension/src/lib.rs

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -109,9 +109,18 @@ pub fn mod_init(module: &Bound<'_, PyModule>) -> PyResult<()> {
109109
module,
110110
"monarch_extension.simulator_client",
111111
)?)?;
112-
hyperactor_extension::register_python_bindings(&get_or_add_new_module(
112+
hyperactor_extension::alloc::register_python_bindings(&get_or_add_new_module(
113113
module,
114-
"hyperactor_extension",
114+
"hyperactor_extension.alloc",
115+
)?)?;
116+
::controller::bootstrap::register_python_bindings(&get_or_add_new_module(
117+
module,
118+
"controller.bootstrap",
119+
)?)?;
120+
121+
::monarch_worker::bootstrap::register_python_bindings(&get_or_add_new_module(
122+
module,
123+
"monarch_worker.bootstrap",
115124
)?)?;
116125

117126
Ok(())

‎monarch_extension/src/simulator_client.rs

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@ use pyo3::prelude::*;
2323
#[pyclass(
2424
name = "SimulatorClient",
2525
frozen,
26-
module = "monarch._monarch.simulator_client"
26+
module = "monarch._rust_bindings.monarch_extension.simulator_client"
2727
)]
2828
#[derive(Clone)]
2929
pub(crate) struct SimulatorClient {
@@ -91,9 +91,11 @@ impl SimulatorClient {
9191

9292
pub(crate) fn register_python_bindings(simulator_client_mod: &Bound<'_, PyModule>) -> PyResult<()> {
9393
simulator_client_mod.add_class::<SimulatorClient>()?;
94-
simulator_client_mod.add_function(wrap_pyfunction!(
95-
bootstrap_simulator_backend,
96-
simulator_client_mod
97-
)?)?;
94+
let f = wrap_pyfunction!(bootstrap_simulator_backend, simulator_client_mod)?;
95+
f.setattr(
96+
"__module__",
97+
"monarch._rust_bindings.monarch_extension.simulator_client",
98+
)?;
99+
simulator_client_mod.add_function(f)?;
98100
Ok(())
99101
}

0 commit comments

Comments
 (0)