Remove remote / local actor routing (risc0#3452)

Instead, just use loopback for local actor communication (if the actor
can be remote)

Author: bobbobbio

risc0/r0vm/src/actors/allocator.rs

@@ -35,18 +35,17 @@
 //!
 
 use std::collections::{BTreeMap, HashMap, HashSet, VecDeque, btree_map, hash_map};
-use std::error::Error as StdError;
 use std::fmt;
 use std::net::SocketAddr;
 use std::net::{IpAddr, Ipv4Addr};
 use std::sync::Arc;
 
 use super::{
     RemoteActor, RemoteAllocatorRequest, RpcDisconnect,
-    actor::{self, Actor, ActorRef, Context, Message, ReplySender},
+    actor::{Actor, ActorRef, Context, Message, ReplySender},
     error::{Error, Result},
     protocol::{GlobalId, WorkerId, WorkerIdFmt},
-    remote_actor_ask, routing_actor_impl,
+    remote_actor_ask,
 };
 use derive_more::{Add, AddAssign, From, Sub, SubAssign};
 use http_body_util::BodyExt as _;
@@ -1042,6 +1041,7 @@ impl AllocatorActor {
 #[cfg(test)]
 mod allocation_tests {
     use super::*;
+    use crate::actors::actor;
 
     fn test_gpu_id(n: u32) -> GpuUuid {
         assert!(n < 10);
@@ -1744,50 +1744,6 @@ mod allocation_tests {
 // |_|  | .__/ \___|
 //      |_|
 
-pub(crate) enum AllocatorRouterActor {
-    Local(ActorRef<AllocatorActor>),
-    Remote(ActorRef<RemoteAllocatorActor>),
-}
-
-impl Actor for AllocatorRouterActor {}
-
-impl AllocatorRouterActor {
-    pub async fn new(
-        addr: &Option<SocketAddr>,
-        local: &Option<ActorRef<AllocatorActor>>,
-    ) -> std::result::Result<ActorRef<Self>, Box<dyn StdError>> {
-        if let Some(addr) = addr {
-            let remote =
-                actor::spawn(RemoteAllocatorActor::new(*addr, "RemoteAllocatorActor").await?);
-            Ok(actor::spawn(Self::Remote(remote)))
-        } else {
-            Ok(actor::spawn(Self::Local(
-                local.as_ref().ok_or("no allocator configured")?.clone(),
-            )))
-        }
-    }
-}
-
-routing_actor_impl!(
-    AllocatorRouterActor,
-    RegisterWorker,
-    Result<RegisterWorkerReply>
-);
-
-routing_actor_impl!(
-    AllocatorRouterActor,
-    ScheduleTask,
-    Result<ScheduleTaskReply>
-);
-
-routing_actor_impl!(AllocatorRouterActor, AllocateHardware, Result<()>);
-
-routing_actor_impl!(AllocatorRouterActor, DeallocateHardware, Result<()>);
-
-routing_actor_impl!(AllocatorRouterActor, EndTask, Result<()>);
-
-routing_actor_impl!(AllocatorRouterActor, RegisterManager, Result<()>);
-
 pub type RemoteAllocatorActor = RemoteActor<AllocatorActor>;
 
 remote_actor_ask!(
@@ -2048,6 +2004,7 @@ pub async fn run_proxy(
 #[cfg(test)]
 mod proxy_tests {
     use super::*;
+    use crate::actors::actor;
     use axum_test::TestServer;
     use httpmock::MockServer;
     use rstest::rstest;

risc0/r0vm/src/actors/factory.rs

@@ -14,23 +14,22 @@
 // SPDX-License-Identifier: Apache-2.0 OR MIT
 
 use std::collections::HashMap;
-use std::error::Error as StdError;
-use std::net::SocketAddr;
 
 use multi_index_map::MultiIndexMap;
 
 use super::{
-    RemoteActor, RemoteFactoryRequest, RpcDisconnect, RpcMessageId, WorkerRouterActor,
-    actor::{self, Actor, ActorRef, Context, Message},
-    allocator::{AllocatorRouterActor, CpuCores, GpuTokens, ScheduleTask},
+    RemoteActor, RemoteFactoryRequest, RpcDisconnect,
+    actor::{Actor, ActorRef, Context, Message},
+    allocator::{CpuCores, GpuTokens, RemoteAllocatorActor, ScheduleTask},
     error::{Error, Result as ActorResult},
     job::JobActor,
     protocol::{
         GlobalId, JobId, TaskKind, WorkerId,
         factory::{DropJob, GetTasks, SubmitTaskMsg, TaskDoneMsg, TaskUpdateMsg},
         worker::TaskMsg,
     },
-    remote_actor_tell, routing_actor_impl,
+    remote_actor_tell,
+    worker::RemoteWorkerActor,
 };
 
 #[derive(MultiIndexMap)]
@@ -46,13 +45,13 @@ pub(crate) struct FactoryActor {
     workers: MultiIndexWorkerRowMap,
     pending_tasks: Vec<SubmitTaskMsg>,
     active_tasks: HashMap<GlobalId, SubmitTaskMsg>,
-    worker_actors: HashMap<WorkerId, ActorRef<WorkerRouterActor>>,
-    allocator: ActorRef<AllocatorRouterActor>,
+    worker_actors: HashMap<WorkerId, ActorRef<RemoteWorkerActor>>,
+    allocator: ActorRef<RemoteAllocatorActor>,
     require_gpu: bool,
 }
 
 impl FactoryActor {
-    pub fn new(allocator: ActorRef<AllocatorRouterActor>, require_gpu: bool) -> Self {
+    pub fn new(allocator: ActorRef<RemoteAllocatorActor>, require_gpu: bool) -> Self {
         Self {
             jobs: HashMap::default(),
             workers: Default::default(),
@@ -248,46 +247,6 @@ impl Message<RpcDisconnect> for FactoryActor {
 // |_|  | .__/ \___|
 //      |_|
 
-pub(crate) enum FactoryRouterActor {
-    Local(ActorRef<FactoryActor>),
-    Remote(ActorRef<RemoteFactoryActor>),
-}
-
-impl Actor for FactoryRouterActor {}
-
-impl FactoryRouterActor {
-    pub async fn new<RemoteMsgT, FutT>(
-        addr: &Option<SocketAddr>,
-        local: &Option<ActorRef<FactoryActor>>,
-        remote_msg_callback: impl FnMut(Option<RemoteMsgT>, Option<RpcMessageId>, SocketAddr) -> FutT
-        + Send
-        + 'static,
-    ) -> Result<ActorRef<Self>, Box<dyn StdError>>
-    where
-        RemoteMsgT: serde::de::DeserializeOwned,
-        FutT: Future<Output = ()> + Send,
-    {
-        if let Some(addr) = addr {
-            let remote = actor::spawn(
-                RemoteFactoryActor::new_with_remote_msg_callback(*addr, remote_msg_callback)
-                    .await?,
-            );
-            Ok(actor::spawn(Self::Remote(remote)))
-        } else {
-            Ok(actor::spawn(Self::Local(
-                local
-                    .as_ref()
-                    .ok_or("no manager found from allocator or locally")?
-                    .clone(),
-            )))
-        }
-    }
-}
-
-routing_actor_impl!(FactoryRouterActor, GetTasks, ());
-routing_actor_impl!(FactoryRouterActor, TaskUpdateMsg, ());
-routing_actor_impl!(FactoryRouterActor, TaskDoneMsg, ());
-
 pub type RemoteFactoryActor = RemoteActor<FactoryActor>;
 
 remote_actor_tell!(RemoteActor<FactoryActor>, GetTasks, RemoteFactoryRequest);