Edits to tower-based routing: src/router.rs room return type changed to Arc<Lighthouse>, Lighthouse::new now takes id prefix, test relocated to lighthouse_test.py and now uses coordination API, LighthouseServer now resolves host/port from the bound socket to give a routable http://host:port address

Author: Matt Kotzbauer

src/interceptor.rs~

@@ -17,6 +17,7 @@ pub use crate::router::Router;
 use anyhow::Result;
 use atty::Stream;
 use core::time::Duration;
+use gethostname::gethostname;
 use pyo3::exceptions::{PyRuntimeError, PyTimeoutError};
 use std::cmp;
 use std::env;
@@ -664,7 +665,6 @@ impl LighthouseServer {
 
             let listener = rt.block_on(tokio::net::TcpListener::bind(&bind))?;
             let bound_sock = listener.local_addr()?;
-            let bound = format!("http://{}", bound_sock);
             let incoming = TcpListenerStream::new(listener);
             let router = Router::new(opt.clone());
 
@@ -676,8 +676,15 @@ impl LighthouseServer {
                     .map_err(|e: tonic::transport::Error| anyhow::anyhow!(e))
             });
 
+            let host = if bind.starts_with("0.0.0.0") || bind.starts_with("[::]") {
+                gethostname().to_string_lossy().into_owned()
+            } else {
+                bind.rsplit_once(':').map(|(h, _)| h.to_string()).unwrap()
+            };
+            let public_addr = format!("http://{}:{}", host, bound_sock.port());
+
             Ok(Self {
-                bind: bound,
+                bind: public_addr,
                 handle,
                 _runtime: rt,
             })

src/lib.rs

@@ -58,6 +58,7 @@ struct State {
 }
 
 pub struct Lighthouse {
+    id: String,
     state: Mutex<State>,
     opt: LighthouseOpt,
     listener: Mutex<Option<tokio::net::TcpListener>>,
@@ -261,12 +262,13 @@ fn quorum_compute(
 }
 
 impl Lighthouse {
-    pub async fn new(opt: LighthouseOpt) -> Result<Arc<Self>> {
+    pub async fn new(id: String, opt: LighthouseOpt) -> Result<Arc<Self>> {
         let listener = tokio::net::TcpListener::bind(&opt.bind).await?;
 
         let (tx, _) = broadcast::channel(16);
 
         Ok(Arc::new(Self {
+            id: id,
             state: Mutex::new(State {
                 participants: HashMap::new(),
                 channel: tx,

src/lighthouse.rs

@@ -10,7 +10,7 @@ use dashmap::{mapref::entry::Entry, DashMap};
 use futures::FutureExt;
 use tonic::{
     body::BoxBody,
-    codegen::http::{HeaderMap, Request, Response}, // http-0.2 types
+    codegen::http::{HeaderMap, Request, Response},
     server::NamedService,
 };
 use tower::Service;
@@ -28,7 +28,7 @@ type GrpcSvc = LighthouseServiceServer<Arc<Lighthouse>>;
 
 #[derive(Clone)]
 pub struct Router {
-    rooms: Arc<DashMap<String, Arc<GrpcSvc>>>,
+    rooms: Arc<DashMap<String, Arc<Lighthouse>>>,
     tmpl_opt: LighthouseOpt,
 }
 
@@ -47,26 +47,23 @@ impl Router {
     }
 
     async fn room_service(
-        rooms: Arc<DashMap<String, Arc<GrpcSvc>>>,
+        rooms: Arc<DashMap<String, Arc<Lighthouse>>>,
         tmpl: LighthouseOpt,
         id: &str,
-    ) -> Arc<GrpcSvc> {
-        if let Some(svc) = rooms.get(id) {
-            return svc.clone();
+    ) -> Arc<Lighthouse> {
+        if let Some(lh) = rooms.get(id) {
+            return lh.clone();
         }
 
-        // Build room state once.
-        let lh = Lighthouse::new(tmpl.clone())
+        let lh = Lighthouse::new(id.to_owned(), tmpl.clone())
             .await
             .expect("failed to create Lighthouse");
 
-        let svc_new = Arc::new(LighthouseServiceServer::new(lh));
-
         match rooms.entry(id.to_owned()) {
             Entry::Occupied(e) => e.get().clone(),
             Entry::Vacant(v) => {
-                v.insert(svc_new.clone());
-                svc_new
+                v.insert(lh.clone());
+                lh
             }
         }
     }
@@ -89,10 +86,9 @@ impl Service<Request<BoxBody>> for Router {
         let room = Self::room_id(req.headers()).to_owned();
 
         async move {
-            let svc_arc = Self::room_service(rooms, tmpl, &room).await;
+            let lh = Self::room_service(rooms, tmpl, &room).await;
 
-            // `Arc<GrpcSvc>` itself isn’t a Service; clone the inner value.
-            let mut svc = (*svc_arc).clone();
+            let mut svc = LighthouseServiceServer::new(lh);
             let resp = svc
                 .call(req)
                 .await

src/router.rs

@@ -4,6 +4,7 @@
 
 import torch.distributed as dist
 
+import torchft.coordination as cd
 from torchft import Manager, ProcessGroupGloo
 from torchft._torchft import LighthouseClient, LighthouseServer, Quorum, QuorumMember
 
@@ -155,3 +156,32 @@ def test_heartbeat_round_trip(self) -> None:
 
         finally:
             lighthouse.shutdown()
+
+    def test_multi_room_quorums(self) -> None:
+        """One server, two logical rooms should yield two isolated quorums."""
+        server = cd.LighthouseServer(bind="[::]:0", min_replicas=1)
+        addr = server.address()
+
+        try:
+            # Two clients in two independent rooms
+            cli_a = cd.LighthouseClient(addr, timedelta(seconds=1), room_id="jobA")
+            cli_b = cd.LighthouseClient(addr, timedelta(seconds=1), room_id="jobB")
+
+            # Explicit heartbeat so each room has one participant
+            cli_a.heartbeat("a0")
+            cli_b.heartbeat("b0")
+
+            q_a = cli_a.quorum("a0", timedelta(seconds=1))
+            q_b = cli_b.quorum("b0", timedelta(seconds=1))
+
+            # Both rooms got a quorum-id of 1 but with disjoint members
+            self.assertEqual(q_a.quorum_id, 1)
+            self.assertEqual(q_b.quorum_id, 1)
+
+            self.assertEqual(len(q_a.participants), 1)
+            self.assertEqual(len(q_b.participants), 1)
+            self.assertEqual(q_a.participants[0].replica_id, "a0")
+            self.assertEqual(q_b.participants[0].replica_id, "b0")
+
+        finally:
+            server.shutdown()