Skip to content

Commit 35d6478

Browse files
pzhan9facebook-github-bot
authored andcommitted
Add new_with_shape methods (#700)
Summary: Pull Request resolved: #700 Pull Request resolved: #557 This diff is to address [this comment](https://www.internalfb.com/diff/D78355743?dst_version_fbid=23914538798238882&transaction_fbid=1255127772923671) from D78355743. Basically the python side wants to use its own logic to slice shape, and build the mesh with the sliced shape it provides. In the long run, we probably do want to have both Rust and Python layers have their own slicing logic. But the consolidation is not in the scope of this stack. Reviewed By: moonli Differential Revision: D79287947 fbshipit-source-id: 56f8582390c8f4a3ef69c8e6b6023e0b003e0f26
1 parent 3105c56 commit 35d6478

File tree

3 files changed

+46
-0
lines changed

3 files changed

+46
-0
lines changed

hyperactor_mesh/src/reference.rs

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ use ndslice::Range;
2323
use ndslice::Selection;
2424
use ndslice::Shape;
2525
use ndslice::ShapeError;
26+
use ndslice::selection::ReifyView;
2627
use serde::Deserialize;
2728
use serde::Serialize;
2829

@@ -162,6 +163,25 @@ impl<A: RemoteActor> ActorMeshRef<A> {
162163
phantom: PhantomData,
163164
})
164165
}
166+
167+
pub fn new_with_shape(&self, new_shape: Shape) -> anyhow::Result<Self> {
168+
let base_slice = self.shape().slice();
169+
base_slice.reify_view(new_shape.slice()).map_err(|e| {
170+
anyhow::anyhow!(
171+
"failed to reify the new shape into the base shape; this \
172+
normally means the new shape is not a valid slice of the base \
173+
error is: {e:?}"
174+
)
175+
})?;
176+
177+
Ok(Self {
178+
mesh_id: self.mesh_id.clone(),
179+
root: self.root.clone(),
180+
sliced: Some(new_shape),
181+
comm_actor_ref: self.comm_actor_ref.clone(),
182+
phantom: PhantomData,
183+
})
184+
}
165185
}
166186

167187
impl<A: RemoteActor> Clone for ActorMeshRef<A> {

monarch_hyperactor/src/actor_mesh.rs

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -224,6 +224,10 @@ impl PythonActorMesh {
224224
self.bind()?.slice(kwargs)
225225
}
226226

227+
fn new_with_shape(&self, shape: PyShape) -> PyResult<PythonActorMeshRef> {
228+
self.bind()?.new_with_shape(shape)
229+
}
230+
227231
#[getter]
228232
pub fn client(&self) -> PyMailbox {
229233
self.client.clone()
@@ -356,6 +360,14 @@ impl PythonActorMeshRef {
356360
Ok(Self { inner: sliced })
357361
}
358362

363+
fn new_with_shape(&self, shape: PyShape) -> PyResult<PythonActorMeshRef> {
364+
let sliced = self
365+
.inner
366+
.new_with_shape(shape.get_inner().clone())
367+
.map_err(|e| PyErr::new::<PyValueError, _>(e.to_string()))?;
368+
Ok(Self { inner: sliced })
369+
}
370+
359371
#[getter]
360372
fn shape(&self) -> PyShape {
361373
PyShape::from(self.inner.shape().clone())

python/monarch/_rust_bindings/monarch_hyperactor/actor_mesh.pyi

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,13 @@ class PythonActorMeshRef:
3838
"""
3939
...
4040

41+
def new_with_shape(self, shape: Shape) -> PythonActorMeshRef:
42+
"""
43+
Return a new mesh ref with the given sliced shape. If the provided shape
44+
is not a valid slice of the current shape, an exception will be raised.
45+
"""
46+
...
47+
4148
@property
4249
def shape(self) -> Shape:
4350
"""
@@ -77,6 +84,13 @@ class PythonActorMesh:
7784
"""
7885
...
7986

87+
def new_with_shape(self, shape: Shape) -> PythonActorMeshRef:
88+
"""
89+
Return a new mesh ref with the given sliced shape. If the provided shape
90+
is not a valid slice of the current shape, an exception will be raised.
91+
"""
92+
...
93+
8094
def get_supervision_event(self) -> ActorSupervisionEvent | None:
8195
"""
8296
Returns supervision event if there is any.

0 commit comments

Comments
 (0)