hyperactor_mesh/
test_utils.rs

1/*
2 * Copyright (c) Meta Platforms, Inc. and affiliates.
3 * All rights reserved.
4 *
5 * This source code is licensed under the BSD-style license found in the
6 * LICENSE file in the root directory of this source tree.
7 */
8
9use async_trait::async_trait;
10use hyperactor::Actor;
11use hyperactor::Bind;
12use hyperactor::Context;
13use hyperactor::Handler;
14use hyperactor::Unbind;
15use hyperactor::channel::ChannelTransport;
16use serde::Deserialize;
17use serde::Serialize;
18use typeuri::Named;
19
20use crate::host_mesh::HostMesh;
21
22/// Message that can be sent to an EmptyActor.
23#[derive(Serialize, Deserialize, Debug, Named, Clone, Bind, Unbind)]
24pub struct EmptyMessage();
25
26#[derive(Debug, PartialEq, Default)]
27#[hyperactor::export(EmptyMessage { cast = true })]
28#[hyperactor::spawnable]
29pub struct EmptyActor();
30
31impl Actor for EmptyActor {}
32
33#[async_trait]
34impl Handler<EmptyMessage> for EmptyActor {
35    async fn handle(&mut self, _: &Context<Self>, _: EmptyMessage) -> Result<(), anyhow::Error> {
36        Ok(())
37    }
38}
39
40/// Create a local in-process host mesh with `n` hosts, all running in
41/// the current process using `Local` channel transport.
42///
43/// This is similar to [`HostMesh::local_in_process`] but supports
44/// multiple hosts. All hosts use [`LocalProcManager`] with
45/// [`ChannelTransport::Local`], so there is no IPC overhead.
46///
47/// # Examples
48///
49/// ```ignore
50/// let mut host_mesh = test_utils::local_host_mesh(4).await;
51/// let proc_mesh = host_mesh
52///     .spawn(instance, "test", ndslice::extent!(gpu = 8))
53///     .await
54///     .unwrap();
55/// // ... do something with the proc mesh ...
56/// // shutdown the host mesh.
57/// let _ = host_mesh.shutdown(&instance).await;
58/// ```
59pub async fn local_host_mesh(n: usize) -> HostMesh {
60    let addrs = (0..n).map(|_| ChannelTransport::Local.any()).collect();
61    let host_mesh = HostMesh::local_n_in_process(addrs).await.unwrap();
62    HostMesh::take(host_mesh)
63}
hyperactor_mesh/test_utils.rs

hyperactor_mesh/
test_utils.rs
