Create process_allocator binary for OSS

Summary: Adds the OSS equivalent of `hyperactor_meta mesh-worker` called `process_allocator`. Also created a meta version of it (but haven't switched over to it in the job launcher).

Reviewed By: colin2328

Differential Revision: D75299505

fbshipit-source-id: 1b4a89ddd2da424ce8f8ded9343d134979ebabb5

Author: kiukchung

monarch_hyperactor/Cargo.toml

@@ -1,4 +1,4 @@
-# @generated by autocargo from //monarch/monarch_hyperactor:monarch_hyperactor
+# @generated by autocargo from //monarch/monarch_hyperactor:[monarch_hyperactor,process_allocator-oss]
 
 [package]
 name = "monarch_hyperactor"
@@ -11,6 +11,7 @@ license = "BSD-3"
 anyhow = "1.0.95"
 async-trait = "0.1.86"
 bincode = "1.3.3"
+clap = { version = "4.5.38", features = ["derive", "env", "string", "unicode", "wrap_help"] }
 hyperactor = { version = "0.0.0", path = "../hyperactor" }
 hyperactor_extension = { version = "0.0.0", path = "../hyperactor_extension" }
 hyperactor_mesh = { version = "0.0.0", path = "../hyperactor_mesh" }

monarch_hyperactor/src/bin/process_allocator/common.rs

@@ -0,0 +1,161 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+use std::result::Result;
+
+use clap::Parser;
+use clap::command;
+use hyperactor::channel::ChannelAddr;
+use hyperactor_mesh::alloc::remoteprocess::RemoteProcessAllocator;
+use tokio::process::Command;
+#[derive(Parser, Debug)]
+#[command(about = "Runs hyperactor's process allocator")]
+pub struct Args {
+    #[arg(
+        long,
+        default_value = "[::]",
+        help = "The address bind to. The process allocator runs on `bind_addr:port`"
+    )]
+    pub addr: String,
+
+    #[arg(
+        long,
+        default_value_t = 26600,
+        help = "Port to bind to. The process allocator runs on `bind_addr:port`"
+    )]
+    pub port: u16,
+
+    #[arg(
+        long,
+        default_value = "monarch_bootstrap",
+        help = "The path to the binary that this process allocator spawns on an `allocate` request"
+    )]
+    pub program: String,
+}
+
+pub fn main_impl(
+    serve_address: ChannelAddr,
+    program: String,
+) -> tokio::task::JoinHandle<Result<(), anyhow::Error>> {
+    tracing::info!("bind address is: {}", serve_address);
+    tracing::info!("program to spawn on allocation request: [{}]", &program);
+
+    tokio::spawn(async {
+        RemoteProcessAllocator::new()
+            .start(Command::new(program), serve_address)
+            .await
+    })
+}
+
+#[cfg(test)]
+mod tests {
+    use std::collections::HashSet;
+
+    use clap::Parser;
+    use hyperactor::WorldId;
+    use hyperactor::channel::ChannelTransport;
+    use hyperactor_mesh::alloc;
+    use hyperactor_mesh::alloc::Alloc;
+    use hyperactor_mesh::alloc::remoteprocess;
+    use ndslice::shape;
+
+    use super::*;
+
+    #[tokio::test]
+    async fn test_args_defaults() -> Result<(), anyhow::Error> {
+        let args = vec!["process_allocator"];
+
+        let parsed_args = Args::parse_from(args);
+
+        assert_eq!(parsed_args.addr, "[::]");
+        assert_eq!(parsed_args.port, 26600);
+        assert_eq!(parsed_args.program, "monarch_bootstrap");
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn test_args() -> Result<(), anyhow::Error> {
+        let args = vec![
+            "process_allocator",
+            "--addr=127.0.0.1",
+            "--port=29500",
+            "--program=/bin/echo",
+        ];
+
+        let parsed_args = Args::parse_from(args);
+
+        assert_eq!(parsed_args.addr, "127.0.0.1");
+        assert_eq!(parsed_args.port, 29500);
+        assert_eq!(parsed_args.program, "/bin/echo");
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn test_main_impl() -> Result<(), anyhow::Error> {
+        hyperactor::initialize();
+
+        let serve_address = ChannelAddr::any(ChannelTransport::Unix);
+        let program = String::from("/bin/date"); // date is usually a unix built-in command
+        let server_handle = main_impl(serve_address.clone(), program);
+
+        let spec = alloc::AllocSpec {
+            // NOTE: x cannot be more than 1 since we created a single process-allocator server instance!
+            shape: shape! { x=1, y=4 },
+            constraints: alloc::AllocConstraints::none(),
+        };
+
+        let mut initializer = remoteprocess::MockRemoteProcessAllocInitializer::new();
+        initializer.expect_initialize_alloc().return_once(move || {
+            Ok(vec![remoteprocess::RemoteProcessAllocHost {
+                hostname: serve_address.to_string(),
+                id: serve_address.to_string(),
+            }])
+        });
+
+        let heartbeat = std::time::Duration::from_millis(100);
+        let world_id = WorldId("__unused__".to_string());
+
+        let mut alloc = remoteprocess::RemoteProcessAlloc::new(
+            spec.clone(),
+            world_id,
+            ChannelTransport::Unix,
+            0,
+            heartbeat,
+            initializer,
+        )
+        .await
+        .unwrap();
+
+        // make sure we accounted for `world_size` number of Created and Stopped proc states
+        let world_size = spec.shape.slice().iter().count();
+        let mut created_ranks: HashSet<usize> = HashSet::new();
+        let mut stopped_ranks: HashSet<usize> = HashSet::new();
+
+        while created_ranks.len() < world_size || stopped_ranks.len() < world_size {
+            let proc_state = alloc.next().await.unwrap();
+            match proc_state {
+                alloc::ProcState::Created { proc_id, coords: _ } => {
+                    // alloc.next() will keep creating procs and incrementing rank id
+                    // so we mod the rank by world_size to map it to its logical rank
+                    created_ranks.insert(proc_id.rank() % world_size);
+                }
+                alloc::ProcState::Stopped { proc_id, reason: _ } => {
+                    stopped_ranks.insert(proc_id.rank() % world_size);
+                }
+                _ => {}
+            }
+        }
+
+        let expected_ranks: HashSet<usize> = (0..world_size).collect();
+        assert_eq!(created_ranks, expected_ranks);
+        assert_eq!(stopped_ranks, expected_ranks);
+
+        server_handle.abort();
+        Ok(())
+    }
+}

monarch_hyperactor/src/bin/process_allocator/main.rs

@@ -0,0 +1,26 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+mod common;
+
+use clap::Parser;
+use common::Args;
+use common::main_impl;
+use hyperactor::channel::ChannelAddr;
+
+#[tokio::main]
+async fn main() {
+    let args = Args::parse();
+    hyperactor::initialize();
+
+    let bind = format!("{}:{}", args.addr, args.port);
+    let socket_addr: std::net::SocketAddr = bind.parse().unwrap();
+    let serve_address = ChannelAddr::Tcp(socket_addr);
+
+    let _ = main_impl(serve_address, args.program).await.unwrap();
+}

monarch_hyperactor/src/bin/process_allocator/main_meta.rs

@@ -0,0 +1,65 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+// std::time::SystemTime::now is used by #[cli::main] suppress the warning since
+// it is used around the CLI not in hyperactor core code
+#![allow(clippy::disallowed_methods)]
+
+mod common;
+
+use anyhow::Result;
+use cli::ExitCode;
+use common::Args;
+use common::main_impl;
+use fbinit::FacebookInit;
+use hyperactor::channel::ChannelAddr;
+use hyperactor_meta_lib::system_resolution::canonicalize_hostname;
+use hyperactor_telemetry::env::Env;
+use hyperactor_telemetry::env::HYPERACTOR_EXECUTION_ID_ENV;
+use hyperactor_telemetry::env::MAST_HPC_JOB_NAME_ENV;
+
+fn hostname() -> String {
+    canonicalize_hostname(
+        hostname::get()
+            .ok()
+            .and_then(|hostname| hostname.into_string().ok())
+            .expect("failed to retrieve hostname")
+            .as_str(),
+    )
+}
+
+#[cli::main("process_allocator")]
+async fn main(_fb: FacebookInit, args: Args) -> Result<ExitCode> {
+    match Env::current() {
+        Env::Mast => {
+            let job_name =
+                std::env::var(MAST_HPC_JOB_NAME_ENV).expect("MAST_HPC_JOB_NAME not set in MAST");
+            std::env::set_var(HYPERACTOR_EXECUTION_ID_ENV, job_name);
+        }
+        _ => {}
+    }
+    hyperactor::initialize();
+
+    let current_host = hostname();
+    tracing::info!(
+        "NOTE: argument `--addr` is ignored when running internally at Meta! \
+        Process allocator runs on current host: `{}:{}` using Meta-TLS over TCP",
+        &current_host,
+        args.port
+    );
+    let serve_address = ChannelAddr::MetaTls(current_host, args.port);
+    let result = main_impl(serve_address, args.program).await;
+
+    match result {
+        Ok(_) => Ok(ExitCode::SUCCESS),
+        Err(e) => {
+            tracing::error!("Error running process allocator: {:?}", e);
+            Ok(ExitCode::FAILURE)
+        }
+    }
+}

tools/components/hyperactor.py

@@ -41,7 +41,7 @@ def proc_mesh(
         mesh_role = specs.Role(
             name=mesh.name,
             image=image,
-            entrypoint="hyperactor",  # TODO install hyperactor cargo and add ~/.cargo/bin to PATH in Dockerfile
+            entrypoint="process_allocator",  # 'cargo install monarch_hyperactor' to get this binary
             args=[
                 "mesh-worker",
                 f"--port={port}",