Skip to content

Commit 0b5d959

Browse files
mariusaefacebook-github-bot
mariusae
authored and
facebook-github-bot
committed
mesh: extent/view based allocation (#764)
Summary: Pull Request resolved: #764 This moves allocation away from being defined by their `Shape` to being defined by `Extent`s (and views of these). Most of these changes are mechanical, with some opportunities to simplify taken along the way. There is still some additional cleanup to do before removing shapes entirely. That will also need to be plumbed into the Python stack. Reviewed By: shayne-fletcher Differential Revision: D79618262 fbshipit-source-id: 0e07d71defa39a17404647dfa595ade70233f585
1 parent 6ef8a32 commit 0b5d959

File tree

24 files changed

+415
-380
lines changed

24 files changed

+415
-380
lines changed

‎hyperactor_mesh/benches/main.rs

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,9 +18,9 @@ use hyperactor_mesh::actor_mesh::RootActorMesh;
1818
use hyperactor_mesh::alloc::AllocSpec;
1919
use hyperactor_mesh::alloc::Allocator;
2020
use hyperactor_mesh::alloc::LocalAllocator;
21+
use hyperactor_mesh::extent;
2122
use hyperactor_mesh::selection::dsl::all;
2223
use hyperactor_mesh::selection::dsl::true_;
23-
use hyperactor_mesh::shape;
2424

2525
mod bench_actor;
2626
use bench_actor::BenchActor;
@@ -39,10 +39,9 @@ fn bench_actor_scaling(c: &mut Criterion) {
3939
group.bench_function(BenchmarkId::from_parameter(host_count), |b| {
4040
let mut b = b.to_async(Runtime::new().unwrap());
4141
b.iter_custom(|iters| async move {
42-
let shape = shape! { hosts=host_count, gpus=8 };
4342
let alloc = LocalAllocator
4443
.allocate(AllocSpec {
45-
shape: shape.clone(),
44+
extent: extent!(hosts = host_count, gpus = 8),
4645
constraints: Default::default(),
4746
})
4847
.await

‎hyperactor_mesh/examples/dining_philosophers.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -27,9 +27,9 @@ use hyperactor_mesh::alloc::AllocSpec;
2727
use hyperactor_mesh::alloc::Allocator;
2828
use hyperactor_mesh::alloc::LocalAllocator;
2929
use hyperactor_mesh::comm::multicast::CastInfo;
30+
use hyperactor_mesh::extent;
3031
use hyperactor_mesh::selection::dsl::all;
3132
use hyperactor_mesh::selection::dsl::true_;
32-
use hyperactor_mesh::shape;
3333
use ndslice::selection::selection_from;
3434
use serde::Deserialize;
3535
use serde::Serialize;
@@ -228,7 +228,7 @@ async fn main() -> Result<ExitCode> {
228228
let group_size = 5;
229229
let alloc = LocalAllocator
230230
.allocate(AllocSpec {
231-
shape: shape! {replica = group_size},
231+
extent: extent! {replica = group_size},
232232
constraints: Default::default(),
233233
})
234234
.await?;

‎hyperactor_mesh/examples/sieve.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ use hyperactor_mesh::ProcMesh;
2727
use hyperactor_mesh::alloc::AllocSpec;
2828
use hyperactor_mesh::alloc::Allocator;
2929
use hyperactor_mesh::alloc::LocalAllocator;
30-
use hyperactor_mesh::shape;
30+
use hyperactor_mesh::extent;
3131
use serde::Deserialize;
3232
use serde::Serialize;
3333

@@ -106,7 +106,7 @@ impl Actor for SieveActor {
106106
async fn main() -> Result<ExitCode> {
107107
let alloc = LocalAllocator
108108
.allocate(AllocSpec {
109-
shape: shape! { replica = 1 },
109+
extent: extent! { replica = 1 },
110110
constraints: Default::default(),
111111
})
112112
.await?;

‎hyperactor_mesh/src/actor_mesh.rs

Lines changed: 21 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -461,6 +461,7 @@ pub(crate) mod test_util {
461461
use hyperactor::Context;
462462
use hyperactor::Handler;
463463
use hyperactor::PortRef;
464+
use ndslice::extent;
464465

465466
use super::*;
466467
use crate::comm::multicast::CastInfo;
@@ -590,16 +591,14 @@ pub(crate) mod test_util {
590591
// The actor creates a mesh.
591592
use std::sync::Arc;
592593

593-
use ndslice::shape;
594-
595594
use crate::alloc::AllocSpec;
596595
use crate::alloc::Allocator;
597596
use crate::alloc::LocalAllocator;
598597

599598
let mut allocator = LocalAllocator;
600599
let alloc = allocator
601600
.allocate(AllocSpec {
602-
shape: shape! { replica = 1 },
601+
extent: extent! { replica = 1 },
603602
constraints: Default::default(),
604603
})
605604
.await
@@ -652,7 +651,7 @@ mod tests {
652651
($allocator:expr_2021) => {
653652
use std::assert_matches::assert_matches;
654653

655-
use ndslice::shape;
654+
use ndslice::extent;
656655
use $crate::alloc::AllocSpec;
657656
use $crate::alloc::Allocator;
658657
use $crate::assign::Ranks;
@@ -674,11 +673,11 @@ mod tests {
674673

675674
hyperactor_telemetry::initialize_logging(hyperactor::clock::ClockKind::default());
676675

677-
use ndslice::shape;
676+
use ndslice::extent;
678677

679678
let alloc = $allocator
680679
.allocate(AllocSpec {
681-
shape: shape! { replica = 1 },
680+
extent: extent! { replica = 1 },
682681
constraints: Default::default(),
683682
})
684683
.await
@@ -695,7 +694,7 @@ mod tests {
695694
async fn test_basic() {
696695
let alloc = $allocator
697696
.allocate(AllocSpec {
698-
shape: shape! { replica = 4 },
697+
extent: extent!(replica = 4),
699698
constraints: Default::default(),
700699
})
701700
.await
@@ -720,7 +719,7 @@ mod tests {
720719

721720
let alloc = $allocator
722721
.allocate(AllocSpec {
723-
shape: shape! { replica = 2 },
722+
extent: extent!(replica = 2),
724723
constraints: Default::default(),
725724
})
726725
.await
@@ -755,7 +754,7 @@ mod tests {
755754
const Z: usize = 3;
756755
let alloc = $allocator
757756
.allocate(AllocSpec {
758-
shape: shape! { x = X, y = Y, z = Z },
757+
extent: extent!(x = X, y = Y, z = Z),
759758
constraints: Default::default(),
760759
})
761760
.await
@@ -798,7 +797,7 @@ mod tests {
798797
async fn test_cast() {
799798
let alloc = $allocator
800799
.allocate(AllocSpec {
801-
shape: shape! { replica = 2, host = 2, gpu = 8 },
800+
extent: extent!(replica = 2, host = 2, gpu = 8),
802801
constraints: Default::default(),
803802
})
804803
.await
@@ -838,7 +837,7 @@ mod tests {
838837
// Sizes intentionally small to keep the time
839838
// required for this test in the process case
840839
// reasonable (< 60s).
841-
shape: shape! { replica = 2, host = 2, gpu = 8 },
840+
extent: extent!(replica = 2, host = 2, gpu = 8),
842841
constraints: Default::default(),
843842
})
844843
.await
@@ -868,7 +867,7 @@ mod tests {
868867
for _ in 0..2 {
869868
let alloc = $allocator
870869
.allocate(AllocSpec {
871-
shape: shape! { replica = 1 },
870+
extent: extent!(replica = 1),
872871
constraints: Default::default(),
873872
})
874873
.await
@@ -911,11 +910,11 @@ mod tests {
911910
use $crate::comm::test_utils::TestActorParams as CastTestActorParams;
912911
use $crate::comm::test_utils::TestMessage as CastTestMessage;
913912

914-
let shape = shape! {replica = 4, host = 4, gpu = 4 };
915-
let num_actors = shape.slice().len();
913+
let extent = extent!(replica = 4, host = 4, gpu = 4);
914+
let num_actors = extent.len();
916915
let alloc = $allocator
917916
.allocate(AllocSpec {
918-
shape,
917+
extent,
919918
constraints: Default::default(),
920919
})
921920
.await
@@ -938,7 +937,7 @@ mod tests {
938937
async fn test_delivery_failure() {
939938
let alloc = $allocator
940939
.allocate(AllocSpec {
941-
shape: shape! { replica = 1 },
940+
extent: extent!(replica = 1 ),
942941
constraints: Default::default(),
943942
})
944943
.await
@@ -966,7 +965,7 @@ mod tests {
966965
async fn test_send_with_headers() {
967966
let alloc = $allocator
968967
.allocate(AllocSpec {
969-
shape: shape! { replica = 1 },
968+
extent: extent!(replica = 1 ),
970969
constraints: Default::default(),
971970
})
972971
.await
@@ -1025,7 +1024,7 @@ mod tests {
10251024

10261025
let alloc = LocalAllocator
10271026
.allocate(AllocSpec {
1028-
shape: shape! { replica = 2 },
1027+
extent: extent!(replica = 2),
10291028
constraints: Default::default(),
10301029
})
10311030
.await
@@ -1092,7 +1091,7 @@ mod tests {
10921091

10931092
let alloc = LocalAllocator
10941093
.allocate(AllocSpec {
1095-
shape: shape! { replica = 1 },
1094+
extent: extent!(replica = 1),
10961095
constraints: Default::default(),
10971096
})
10981097
.await
@@ -1158,7 +1157,7 @@ mod tests {
11581157

11591158
let alloc = LocalAllocator
11601159
.allocate(AllocSpec {
1161-
shape: shape! { replica = 2 },
1160+
extent: extent!(replica = 2),
11621161
constraints: Default::default(),
11631162
})
11641163
.await
@@ -1241,10 +1240,10 @@ mod tests {
12411240
// One rank is enough for this test, because the timeout failure
12421241
// occurred in the `client->1st comm actor` channel, and thus will
12431242
// not be sent further to the rest of the network.
1244-
let shape = shape! {replica = 1 };
1243+
let extent = extent! {replica = 1 };
12451244
let alloc = process_allocator()
12461245
.allocate(AllocSpec {
1247-
shape,
1246+
extent,
12481247
constraints: Default::default(),
12491248
})
12501249
.await

0 commit comments

Comments
 (0)