Compute resource aware simulation (meta-pytorch#857)

Summary:

The sim allocator will now register the location (region, dc, zone, rack, host, gpu) of every ProcId upon creation with the simnet.

Reviewed By: pablorfb-meta

Differential Revision: D80137963

Author: thomasywang

hyperactor/src/simnet.rs

@@ -26,6 +26,7 @@ use async_trait::async_trait;
 use dashmap::DashMap;
 use dashmap::DashSet;
 use enum_as_inner::EnumAsInner;
+use ndslice::view::Point;
 use serde::Deserialize;
 use serde::Deserializer;
 use serde::Serialize;
@@ -43,6 +44,7 @@ use tokio::time::interval;
 use crate::ActorId;
 use crate::Mailbox;
 use crate::OncePortRef;
+use crate::ProcId;
 use crate::channel::ChannelAddr;
 use crate::clock::Clock;
 use crate::clock::RealClock;
@@ -295,6 +297,7 @@ pub struct SimNetHandle {
     training_script_state_tx: tokio::sync::watch::Sender<TrainingScriptState>,
     /// Signal to stop the simnet loop
     stop_signal: Arc<AtomicBool>,
+    resources: DashMap<ProcId, Point>,
 }
 
 impl SimNetHandle {
@@ -405,6 +408,11 @@ impl SimNetHandle {
             "timeout waiting for received events to be scheduled".to_string(),
         ))
     }
+
+    /// Register the location in resource space for a Proc
+    pub fn register_proc(&self, proc_id: ProcId, point: Point) {
+        self.resources.insert(proc_id, point);
+    }
 }
 
 pub(crate) type Topology = DashMap<SimNetEdge, SimNetEdgeInfo>;
@@ -478,6 +486,7 @@ pub fn start() {
         pending_event_count,
         training_script_state_tx,
         stop_signal,
+        resources: DashMap::new(),
     });
 }
 

hyperactor_mesh/src/alloc/sim.rs

@@ -11,6 +11,7 @@
 #![allow(dead_code)] // until it is used outside of testing
 
 use async_trait::async_trait;
+use hyperactor::ProcId;
 use hyperactor::WorldId;
 use hyperactor::channel::ChannelAddr;
 use hyperactor::channel::ChannelTransport;
@@ -61,12 +62,23 @@ pub struct SimAlloc {
 
 impl SimAlloc {
     fn new(spec: AllocSpec) -> Self {
-        Self {
-            inner: LocalAlloc::new_with_transport(
-                spec,
-                ChannelTransport::Sim(Box::new(ChannelTransport::Unix)),
-            ),
-        }
+        let inner = LocalAlloc::new_with_transport(
+            spec,
+            ChannelTransport::Sim(Box::new(ChannelTransport::Unix)),
+        );
+        let client_proc_id = ProcId::Ranked(WorldId(format!("{}_manager", inner.name())), 0);
+
+        let ext = inner.extent();
+
+        hyperactor::simnet::simnet_handle()
+            .expect("simnet event loop not running")
+            .register_proc(
+                client_proc_id.clone(),
+                ext.point(ext.sizes().iter().map(|_| 0).collect())
+                    .expect("should be valid point"),
+            );
+
+        Self { inner }
     }
     /// A chaos monkey that can be used to stop procs at random.
     pub(crate) fn chaos_monkey(&self) -> impl Fn(usize, ProcStopReason) + 'static {
@@ -90,7 +102,13 @@ impl SimAlloc {
 #[async_trait]
 impl Alloc for SimAlloc {
     async fn next(&mut self) -> Option<ProcState> {
-        self.inner.next().await
+        let proc_state = self.inner.next().await;
+        if let Some(ProcState::Created { proc_id, point, .. }) = &proc_state {
+            hyperactor::simnet::simnet_handle()
+                .expect("simnet event loop not running")
+                .register_proc(proc_id.clone(), point.clone());
+        }
+        proc_state
     }
 
     fn extent(&self) -> &Extent {

python/monarch/_src/actor/proc_mesh.py

@@ -503,8 +503,26 @@ def local_proc_mesh(*, gpus: Optional[int] = None, hosts: int = 1) -> ProcMesh:
     return _proc_mesh_from_allocator(allocator=LocalAllocator(), gpus=gpus, hosts=hosts)
 
 
-def sim_proc_mesh(*, gpus: Optional[int] = None, hosts: int = 1) -> ProcMesh:
-    return _proc_mesh_from_allocator(allocator=SimAllocator(), gpus=gpus, hosts=hosts)
+def sim_proc_mesh(
+    *,
+    gpus: int = 1,
+    hosts: int = 1,
+    racks: int = 1,
+    zones: int = 1,
+    dcs: int = 1,
+    regions: int = 1,
+) -> ProcMesh:
+    spec: AllocSpec = AllocSpec(
+        AllocConstraints(),
+        hosts=hosts,
+        gpus=gpus,
+        racks=racks,
+        zones=zones,
+        dcs=dcs,
+        regions=regions,
+    )
+    alloc = SimAllocator().allocate(spec)
+    return ProcMesh.from_alloc(alloc, None, True)
 
 
 _BOOTSTRAP_MAIN = "monarch._src.actor.bootstrap_main"