Skip to content

Commit a350349

Browse files
zdevitofacebook-github-bot
zdevito
authored and
facebook-github-bot
committed
monarch_hyperactor
Summary: Moves the monarch_hyperactor python bindings into modules that match where they are defined in python. This completes the move of binding code. Reviewed By: colin2328 Differential Revision: D74975746 fbshipit-source-id: 5ec7d62c54ee3e7ad3f689463d410c8626b30df8
1 parent 5015e28 commit a350349

Some content is hidden

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

44 files changed

+792
-180
lines changed

‎monarch_extension/src/lib.rs

Lines changed: 34 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -59,7 +59,40 @@ pub fn mod_init(module: &Bound<'_, PyModule>) -> PyResult<()> {
5959
module,
6060
"monarch_extension.controller",
6161
)?)?;
62-
monarch_hyperactor::register_python_bindings(&get_or_add_new_module(module, "hyperactor")?)?;
62+
63+
monarch_hyperactor::bootstrap::register_python_bindings(&get_or_add_new_module(
64+
module,
65+
"monarch_hyperactor.bootstrap",
66+
)?)?;
67+
68+
monarch_hyperactor::proc::register_python_bindings(&get_or_add_new_module(
69+
module,
70+
"monarch_hyperactor.proc",
71+
)?)?;
72+
73+
monarch_hyperactor::actor::register_python_bindings(&get_or_add_new_module(
74+
module,
75+
"monarch_hyperactor.actor",
76+
)?)?;
77+
78+
monarch_hyperactor::mailbox::register_python_bindings(&get_or_add_new_module(
79+
module,
80+
"monarch_hyperactor.mailbox",
81+
)?)?;
82+
83+
monarch_hyperactor::alloc::register_python_bindings(&get_or_add_new_module(
84+
module,
85+
"monarch_hyperactor.alloc",
86+
)?)?;
87+
monarch_hyperactor::actor_mesh::register_python_bindings(&get_or_add_new_module(
88+
module,
89+
"monarch_hyperactor.actor_mesh",
90+
)?)?;
91+
monarch_hyperactor::proc_mesh::register_python_bindings(&get_or_add_new_module(
92+
module,
93+
"monarch_hyperactor.proc_mesh",
94+
)?)?;
95+
6396
monarch_hyperactor::runtime::register_python_bindings(&get_or_add_new_module(
6497
module,
6598
"monarch_hyperactor.runtime",

‎monarch_hyperactor/src/actor.rs

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -331,3 +331,12 @@ impl Handler<Cast<PythonMessage>> for PythonActor {
331331
Ok(())
332332
}
333333
}
334+
335+
pub fn register_python_bindings(hyperactor_mod: &Bound<'_, PyModule>) -> PyResult<()> {
336+
hyperactor_mod.add_class::<PickledMessage>()?;
337+
hyperactor_mod.add_class::<PickledMessageClientActor>()?;
338+
hyperactor_mod.add_class::<PythonActorHandle>()?;
339+
hyperactor_mod.add_class::<PythonMessage>()?;
340+
hyperactor_mod.add_class::<PythonActorHandle>()?;
341+
Ok(())
342+
}

‎monarch_hyperactor/src/actor_mesh.rs

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,3 +47,7 @@ impl PythonActorMesh {
4747
PyShape::from(self.inner.shape().clone())
4848
}
4949
}
50+
pub fn register_python_bindings(hyperactor_mod: &Bound<'_, PyModule>) -> PyResult<()> {
51+
hyperactor_mod.add_class::<PythonActorMesh>()?;
52+
Ok(())
53+
}

‎monarch_hyperactor/src/alloc.rs

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -123,3 +123,10 @@ impl PyProcessAllocator {
123123
})?
124124
}
125125
}
126+
127+
pub fn register_python_bindings(hyperactor_mod: &Bound<'_, PyModule>) -> PyResult<()> {
128+
hyperactor_mod.add_class::<PyProcessAllocator>()?;
129+
hyperactor_mod.add_class::<PyLocalAllocator>()?;
130+
131+
Ok(())
132+
}

‎monarch_hyperactor/src/bootstrap.rs

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,9 @@ use pyo3::PyAny;
44
use pyo3::PyResult;
55
use pyo3::Python;
66
use pyo3::pyfunction;
7+
use pyo3::types::PyModule;
8+
use pyo3::types::PyModuleMethods;
9+
use pyo3::wrap_pyfunction;
710

811
#[pyfunction]
912
#[pyo3(signature = ())]
@@ -12,3 +15,9 @@ pub fn bootstrap_main(py: Python) -> PyResult<Bound<PyAny>> {
1215
bootstrap_or_die().await;
1316
})
1417
}
18+
19+
pub fn register_python_bindings(hyperactor_mod: &Bound<'_, PyModule>) -> PyResult<()> {
20+
hyperactor_mod.add_function(wrap_pyfunction!(bootstrap_main, hyperactor_mod)?)?;
21+
22+
Ok(())
23+
}

‎monarch_hyperactor/src/lib.rs

Lines changed: 1 addition & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -3,53 +3,11 @@
33
pub mod actor;
44
pub mod actor_mesh;
55
pub mod alloc;
6-
mod bootstrap;
6+
pub mod bootstrap;
77
pub mod mailbox;
88
pub mod ndslice;
99
pub mod proc;
1010
pub mod proc_mesh;
1111
pub mod runtime;
1212
pub mod selection;
1313
pub mod shape;
14-
15-
use pyo3::Bound;
16-
use pyo3::PyResult;
17-
use pyo3::prelude::*;
18-
use pyo3::types::PyModule;
19-
use pyo3::wrap_pyfunction;
20-
21-
pub fn register_python_bindings(hyperactor_mod: &Bound<'_, PyModule>) -> PyResult<()> {
22-
hyperactor_mod.add_function(wrap_pyfunction!(bootstrap::bootstrap_main, hyperactor_mod)?)?;
23-
24-
hyperactor_mod.add_function(wrap_pyfunction!(proc::init_proc, hyperactor_mod)?)?;
25-
26-
hyperactor_mod.add_class::<proc::PyProc>()?;
27-
hyperactor_mod.add_class::<proc::PyActorId>()?;
28-
hyperactor_mod.add_class::<proc::PySerialized>()?;
29-
30-
hyperactor_mod.add_class::<actor::PickledMessage>()?;
31-
hyperactor_mod.add_class::<actor::PickledMessageClientActor>()?;
32-
hyperactor_mod.add_class::<actor::PythonActorHandle>()?;
33-
hyperactor_mod.add_class::<actor::PythonMessage>()?;
34-
hyperactor_mod.add_class::<actor::PythonActorHandle>()?;
35-
36-
hyperactor_mod.add_class::<mailbox::PyMailbox>()?;
37-
hyperactor_mod.add_class::<mailbox::PyPortId>()?;
38-
hyperactor_mod.add_class::<mailbox::PythonPortHandle>()?;
39-
hyperactor_mod.add_class::<mailbox::PythonPortReceiver>()?;
40-
hyperactor_mod.add_class::<mailbox::PythonOncePortHandle>()?;
41-
hyperactor_mod.add_class::<mailbox::PythonOncePortReceiver>()?;
42-
43-
hyperactor_mod.add_class::<alloc::PyProcessAllocator>()?;
44-
hyperactor_mod.add_class::<alloc::PyLocalAllocator>()?;
45-
46-
hyperactor_mod.add_class::<proc_mesh::PyProcMesh>()?;
47-
48-
hyperactor_mod.add_class::<actor_mesh::PythonActorMesh>()?;
49-
50-
hyperactor_mod.add_class::<shape::PyShape>()?;
51-
52-
hyperactor_mod.add_class::<selection::PySelection>()?;
53-
54-
Ok(())
55-
}

‎monarch_hyperactor/src/mailbox.rs

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -132,7 +132,11 @@ impl PyMailbox {
132132
}
133133
}
134134

135-
#[pyclass(frozen, name = "PortId", module = "monarch._monarch.hyperactor")]
135+
#[pyclass(
136+
frozen,
137+
name = "PortId",
138+
module = "monarch._rust_bindings.monarch_hyperactor.mailbox"
139+
)]
136140
#[derive(Clone)]
137141
pub struct PyPortId {
138142
inner: PortId,
@@ -309,3 +313,14 @@ impl PythonOncePortReceiver {
309313
.map_err(|err| PyErr::new::<PyEOFError, _>(format!("Port closed: {}", err)))
310314
}
311315
}
316+
317+
pub fn register_python_bindings(hyperactor_mod: &Bound<'_, PyModule>) -> PyResult<()> {
318+
hyperactor_mod.add_class::<PyMailbox>()?;
319+
hyperactor_mod.add_class::<PyPortId>()?;
320+
hyperactor_mod.add_class::<PythonPortHandle>()?;
321+
hyperactor_mod.add_class::<PythonPortReceiver>()?;
322+
hyperactor_mod.add_class::<PythonOncePortHandle>()?;
323+
hyperactor_mod.add_class::<PythonOncePortReceiver>()?;
324+
325+
Ok(())
326+
}

‎monarch_hyperactor/src/proc.rs

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -638,3 +638,12 @@ async fn check_actor_supervision_state(
638638
}
639639
Ok(())
640640
}
641+
642+
pub fn register_python_bindings(hyperactor_mod: &Bound<'_, PyModule>) -> PyResult<()> {
643+
hyperactor_mod.add_function(wrap_pyfunction!(init_proc, hyperactor_mod)?)?;
644+
645+
hyperactor_mod.add_class::<PyProc>()?;
646+
hyperactor_mod.add_class::<PyActorId>()?;
647+
hyperactor_mod.add_class::<PySerialized>()?;
648+
Ok(())
649+
}

‎monarch_hyperactor/src/proc_mesh.rs

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -130,3 +130,8 @@ impl PyProcMesh {
130130
Ok(format!("<ProcMesh {}>", self.inner))
131131
}
132132
}
133+
134+
pub fn register_python_bindings(hyperactor_mod: &Bound<'_, PyModule>) -> PyResult<()> {
135+
hyperactor_mod.add_class::<PyProcMesh>()?;
136+
Ok(())
137+
}

‎pyproject.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,3 +18,4 @@ requires-python = ">= 3.10"
1818
markers = [
1919
"oss_skip: marks tests to skip in OSS CI",
2020
]
21+
asyncio_mode = "auto"

‎python/monarch/__init__.py

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,6 @@
3939
from monarch.common.opaque_ref import OpaqueRef
4040
from monarch.common.pipe import create_pipe, Pipe, remote_generator
4141
from monarch.common.remote import remote
42-
from monarch.common.selection import Selection
4342
from monarch.common.shape import NDSlice, Shape
4443
from monarch.common.stream import get_active_stream, Stream
4544
from monarch.common.tensor import reduce, reduce_, Tensor

‎python/monarch/_monarch/hyperactor/__init__.py

Lines changed: 13 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -1,35 +1,25 @@
11
# pyre-strict
22
import abc
33

4-
from monarch._rust_bindings.hyperactor import ( # @manual=//monarch/monarch_extension:monarch_extension
5-
ActorId,
6-
init_proc,
7-
LocalAllocatorBase,
8-
Mailbox,
9-
OncePortHandle,
10-
OncePortReceiver,
11-
PickledMessage,
12-
PickledMessageClientActor,
13-
PortHandle,
14-
PortId,
15-
PortReceiver,
16-
Proc,
17-
ProcessAllocatorBase,
18-
ProcMesh,
19-
PythonActorHandle as ActorHandle,
20-
PythonActorMesh,
21-
PythonMessage,
22-
Serialized,
23-
)
24-
254
from monarch._rust_bindings.hyperactor_extension import ( # @manual=//monarch/monarch_extension:monarch_extension
265
Alloc,
276
AllocConstraints,
287
AllocSpec,
298
)
309

31-
from monarch._rust_bindings.monarch_hyperactor.selection import ( # @manual=//monarch/monarch_extension:monarch_extension
32-
Selection,
10+
from monarch._rust_bindings.monarch_hyperactor.actor import PythonMessage
11+
12+
from monarch._rust_bindings.monarch_hyperactor.alloc import ( # @manual=//monarch/monarch_extension:monarch_extension
13+
LocalAllocatorBase,
14+
)
15+
16+
from monarch._rust_bindings.monarch_hyperactor.mailbox import Mailbox, PortId
17+
18+
from monarch._rust_bindings.monarch_hyperactor.proc import ( # @manual=//monarch/monarch_extension:monarch_extension
19+
ActorId,
20+
init_proc,
21+
Proc,
22+
Serialized,
3323
)
3424

3525
from monarch._rust_bindings.monarch_hyperactor.shape import ( # @manual=//monarch/monarch_extension:monarch_extension

‎python/monarch/_monarch/hyperactor/bootstrap_main.py

Lines changed: 3 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -8,17 +8,13 @@
88

99

1010
async def main():
11-
await hyperactor.bootstrap_main()
11+
await bootstrap_main()
1212

1313

1414
def invoke_main() -> None:
15-
global hyperactor
15+
global bootstrap_main
1616
# TODO: figure out what from worker_main.py we should reproduce here.
17-
18-
# pyre-ignore[21]
19-
from ..._rust_bindings import ( # @manual=//monarch/monarch_extension:monarch_extension
20-
hyperactor,
21-
)
17+
from monarch._rust_bindings.monarch_hyperactor.bootstrap import bootstrap_main
2218

2319
with (
2420
importlib.resources.path("monarch", "py-spy") as pyspy,

‎python/monarch/_rust_bindings/monarch_extension/client.pyi

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
from typing import Any, ClassVar, Dict, final, List
22

3-
from monarch._rust_bindings.hyperactor import ActorId, Proc, Serialized
43
from monarch._rust_bindings.monarch_extension.worker import Ref
4+
from monarch._rust_bindings.monarch_hyperactor.proc import ActorId, Proc, Serialized
55

66
from monarch._rust_bindings.monarch_messages.debugger import DebuggerActionType
77

‎python/monarch/_rust_bindings/monarch_extension/controller.pyi

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,10 @@
11
from enum import Enum
22
from typing import Any, final, List, Optional, Tuple, Union
33

4-
from monarch._rust_bindings.hyperactor import ActorId, Serialized
5-
64
from monarch._rust_bindings.monarch_extension.worker import Ref, WorkerMessage
75

6+
from monarch._rust_bindings.monarch_hyperactor.proc import Serialized
7+
88
from monarch._rust_bindings.monarch_hyperactor.shape import Slice
99

1010
@final

‎python/monarch/_rust_bindings/monarch_extension/debugger.pyi

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
from typing import final, Optional, Union
22

3-
from monarch._rust_bindings.hyperactor import Serialized
3+
from monarch._rust_bindings.monarch_hyperactor.proc import Serialized
44
from monarch._rust_bindings.monarch_messages.debugger import (
55
DebuggerAction,
66
DebuggerActionType,

‎python/monarch/_rust_bindings/monarch_extension/simulator_client.pyi

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
from typing import final, Iterator, overload
1+
from typing import final
22

33
@final
44
class SimulatorClient:

‎python/monarch/_rust_bindings/monarch_extension/worker.pyi

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
from typing import Callable, final, Optional, Sequence, Tuple
22

33
import torch
4-
from monarch._rust_bindings.hyperactor import ActorId
4+
from monarch._rust_bindings.monarch_hyperactor.proc import ActorId
55
from monarch._rust_bindings.monarch_hyperactor.shape import Slice
66

77
@final

0 commit comments

Comments
 (0)