hyperactor_mesh/resource/
mesh.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
9#![allow(dead_code)]
10
11use hyperactor::Named;
12/// This module defines common types for mesh resources. Meshes are managed as
13/// resources, usually by a controller actor implementing the [`crate::resource`]
14/// behavior.
15///
16/// The mesh controller manages all aspects of the mesh lifecycle, and the owning
17/// actor uses the resource behavior directly to query the state of the mesh.
18use ndslice::Extent;
19use serde::Deserialize;
20use serde::Serialize;
21
22use crate::resource::Resource;
23use crate::resource::Status;
24use crate::v1::ValueMesh;
25
26/// Mesh specs
27#[derive(Debug, Named, Serialize, Deserialize)]
28pub struct Spec<S> {
29    /// All meshes have an extent
30    extent: Extent,
31    // supervisor: PortHandle<SupervisionEvent(?)>
32    /// The mesh-specific spec.
33    spec: S,
34}
35
36/// Mesh states
37#[derive(Debug, Named, Serialize, Deserialize)]
38pub struct State<S> {
39    /// The current status for each rank in the mesh.
40    statuses: ValueMesh<Status>,
41    /// Mesh-specific state.
42    state: S,
43}
44
45/// A mesh trait bundles a set of types that together define a mesh resource.
46pub trait Mesh {
47    /// The mesh-specific specification for this resource.
48    type Spec: Named + Serialize + for<'de> Deserialize<'de> + Send + Sync + std::fmt::Debug;
49
50    /// The mesh-specific state for this resource.
51    type State: Named + Serialize + for<'de> Deserialize<'de> + Send + Sync + std::fmt::Debug;
52}
53
54impl<M: Mesh> Resource for M {
55    type Spec = Spec<M::Spec>;
56    type State = State<M::State>;
57}
58
59#[cfg(test)]
60mod test {
61    use hyperactor::Actor;
62    use hyperactor::Context;
63    use hyperactor::Handler;
64
65    use super::*;
66    use crate::resource::Controller;
67    use crate::resource::CreateOrUpdate;
68    use crate::resource::GetState;
69    use crate::resource::Stop;
70
71    // Consider upstreaming this into `hyperactor` -- lightweight handler definitions
72    // can be quite useful.
73    macro_rules! handler {
74        (
75            $actor:path,
76            $(
77                $name:ident: $msg:ty => $body:expr
78            ),* $(,)?
79        ) => {
80            $(
81                #[async_trait::async_trait]
82                impl Handler<$msg> for $actor {
83                    async fn handle(
84                        &mut self,
85                        #[allow(unused_variables)]
86                        cx: & Context<Self>,
87                        $name: $msg
88                    ) -> anyhow::Result<()> {
89                        $body
90                    }
91                }
92            )*
93        };
94    }
95
96    #[derive(Debug, Named, Serialize, Deserialize)]
97    struct TestMesh;
98
99    impl Mesh for TestMesh {
100        type Spec = ();
101        type State = ();
102    }
103
104    #[derive(Actor, Debug, Default, Named, Serialize, Deserialize)]
105    struct TestMeshController;
106
107    // Ensure that TestMeshController conforms to the Controller behavior for TestMesh.
108    handler! {
109        TestMeshController,
110        _message: CreateOrUpdate<Spec<()>> => unimplemented!(),
111        _message: GetState<State<()>> => unimplemented!(),
112        _message: Stop => unimplemented!(),
113    }
114
115    hyperactor::assert_behaves!(TestMeshController as Controller<TestMesh>);
116}
hyperactor_mesh/resource/mesh.rs

hyperactor_mesh/resource/
mesh.rs
