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

hyperactor_mesh/
test_utils.rs
