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

hyperactor_mesh/resource/
mesh.rs
