hyperactor_mesh/comm/
multicast.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//! The comm actor that provides message casting and result accumulation.
10
11use hyperactor::Actor;
12use hyperactor::ActorAddr;
13use hyperactor::Context;
14use hyperactor::RemoteHandles;
15use hyperactor::RemoteMessage;
16use hyperactor::actor::Referable;
17use hyperactor::id::Uid;
18use hyperactor_config::Flattrs;
19use ndslice::Extent;
20use ndslice::Point;
21use ndslice::Region;
22use ndslice::Shape;
23use ndslice::Slice;
24use ndslice::selection::Selection;
25use ndslice::selection::routing::RoutingFrame;
26use serde::Deserialize;
27use serde::Serialize;
28use typeuri::Named;
29use uuid::Uuid;
30
31use crate::ValueMesh;
32use crate::comm::CommMeshConfig;
33use crate::mesh_id::ActorMeshId;
34
35// A temporary trait used to share code in v0/v1 migration. Can be deleted after
36// v0 casting is deleted.
37pub(crate) trait CastEnvelope {
38    fn dest_port(&self) -> &DestinationPort;
39    fn headers(&self) -> &Flattrs;
40    fn sender(&self) -> &ActorAddr;
41    fn cast_point(&self, config: &CommMeshConfig) -> anyhow::Result<Point>;
42    fn data(&self) -> &wirevalue::Any<wirevalue::encoding::Multipart>;
43    fn data_mut(&mut self) -> &mut wirevalue::Any<wirevalue::encoding::Multipart>;
44}
45
46/// A union of slices that can be used to represent arbitrary subset of
47/// ranks in a gang. It is represented by a Slice together with a Selection.
48/// This is used to define the destination of a cast message or the source of
49/// accumulation request.
50#[derive(Serialize, Deserialize, Debug, Clone)]
51pub struct Uslice {
52    /// A slice representing a whole gang.
53    pub slice: Slice,
54    /// A selection used to represent any subset of the gang.
55    pub selection: Selection,
56}
57
58/// An envelope that carries a message destined to a group of actors.
59#[derive(Debug, Serialize, Deserialize, Clone, Named)]
60pub struct CastMessageEnvelope {
61    /// The destination actor mesh id.
62    actor_mesh_id: ActorMeshId,
63    /// The end-to-end message headers.
64    headers: Flattrs,
65    /// The sender of this message.
66    sender: ActorAddr,
67    /// The destination port of the message. It could match multiple actors with
68    /// rank wildcard.
69    dest_port: DestinationPort,
70    /// The serialized message.
71    data: wirevalue::Any<wirevalue::encoding::Multipart>,
72    /// The shape of the cast.
73    shape: Shape,
74}
75wirevalue::register_type!(CastMessageEnvelope);
76
77impl CastEnvelope for CastMessageEnvelope {
78    fn sender(&self) -> &ActorAddr {
79        &self.sender
80    }
81
82    fn headers(&self) -> &Flattrs {
83        &self.headers
84    }
85
86    fn dest_port(&self) -> &DestinationPort {
87        &self.dest_port
88    }
89
90    fn data(&self) -> &wirevalue::Any<wirevalue::encoding::Multipart> {
91        &self.data
92    }
93
94    fn data_mut(&mut self) -> &mut wirevalue::Any<wirevalue::encoding::Multipart> {
95        &mut self.data
96    }
97
98    fn cast_point(&self, config: &CommMeshConfig) -> anyhow::Result<Point> {
99        let rank_on_root_mesh = config.self_rank();
100        let cast_rank = self.relative_rank(rank_on_root_mesh)?;
101        let cast_shape = self.shape();
102        let cast_point = cast_shape
103            .extent()
104            .point_of_rank(cast_rank)
105            .expect("rank out of bounds");
106        Ok(cast_point)
107    }
108}
109
110impl CastMessageEnvelope {
111    /// Create a new CastMessageEnvelope.
112    pub fn new<A, M>(
113        actor_mesh_id: ActorMeshId,
114        sender: ActorAddr,
115        shape: Shape,
116        headers: Flattrs,
117        message: M,
118    ) -> Result<Self, anyhow::Error>
119    where
120        A: Referable + RemoteHandles<M>,
121        M: RemoteMessage,
122    {
123        let actor_uid = actor_mesh_id.uid().clone();
124        let data = wirevalue::Any::<wirevalue::encoding::Multipart>::serialize(&message)?;
125        Ok(Self {
126            actor_mesh_id,
127            headers,
128            sender,
129            dest_port: DestinationPort::new::<A, M>(actor_uid),
130            data,
131            shape,
132        })
133    }
134
135    /// Create a new CastMessageEnvelope from serialized data. Only use this
136    /// when the message do not contain reply ports. Or it does but you are okay
137    /// with the destination actors reply to the client actor directly.
138    pub fn from_serialized(
139        actor_mesh_id: ActorMeshId,
140        sender: ActorAddr,
141        dest_port: DestinationPort,
142        shape: Shape,
143        headers: Flattrs,
144        data: wirevalue::Any<wirevalue::encoding::Multipart>,
145    ) -> Self {
146        Self {
147            actor_mesh_id,
148            sender,
149            headers,
150            dest_port,
151            data,
152            shape,
153        }
154    }
155
156    pub(crate) fn shape(&self) -> &Shape {
157        &self.shape
158    }
159
160    /// Given a rank in the root shape, return the corresponding point in the
161    /// provided shape, which is a view of the root shape.
162    pub(crate) fn relative_rank(&self, rank_on_root_mesh: usize) -> anyhow::Result<usize> {
163        let shape = self.shape();
164        let coords = shape.slice().coordinates(rank_on_root_mesh).map_err(|e| {
165            anyhow::anyhow!(
166                "fail to calculate coords for root rank {} due to error: {}; shape is {:?}",
167                rank_on_root_mesh,
168                e,
169                shape,
170            )
171        })?;
172        let extent =
173            Extent::new(shape.labels().to_vec(), shape.slice().sizes().to_vec()).map_err(|e| {
174                anyhow::anyhow!(
175                    "fail to calculate extent for root rank {} due to error: {}; shape is {}",
176                    rank_on_root_mesh,
177                    e,
178                    shape,
179                )
180            })?;
181        let point = extent.point(coords).map_err(|e| {
182            anyhow::anyhow!(
183                "fail to calculate point for root rank {} due to error: {}; extent is {}, shape is {}",
184                rank_on_root_mesh,
185                e,
186                extent,
187                shape,
188            )
189        })?;
190        Ok(point.rank())
191    }
192
193    /// The unique key used to indicate the stream to which to deliver this message.
194    /// Concretely, the comm actors along the path should use this key to manage
195    /// sequence numbers and reorder buffers.
196    pub(crate) fn stream_key(&self) -> (ActorMeshId, ActorAddr) {
197        (self.actor_mesh_id.clone(), self.sender.clone())
198    }
199}
200
201/// Destination port id of a message. It is a `PortId` with the rank masked out,
202/// and the messege is always sent to the root actor because only root actor
203/// can be accessed externally. The rank is resolved by the destination Selection
204/// of the message. We can use `DestinationPort::port_id(rank)` to get the actual
205/// `PortId` of the message.
206#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Named)]
207pub struct DestinationPort {
208    /// The actor uid to which the message should be delivered.
209    actor_uid: Uid,
210    /// The port index of the destination actors, it is derived from the
211    /// message type and cached here.
212    port: u64,
213}
214wirevalue::register_type!(DestinationPort);
215
216impl DestinationPort {
217    /// Create a new DestinationPort for an actor uid and message type.
218    pub fn new<A, M>(actor_uid: Uid) -> Self
219    where
220        A: Referable + RemoteHandles<M>,
221        M: RemoteMessage,
222    {
223        Self {
224            actor_uid,
225            port: M::port(),
226        }
227    }
228
229    /// The port id of the destination.
230    pub fn port(&self) -> u64 {
231        self.port
232    }
233
234    /// Get the actor uid of the destination.
235    pub fn actor_uid(&self) -> &Uid {
236        &self.actor_uid
237    }
238}
239
240/// The is used to start casting a message to a group of actors.
241#[derive(Serialize, Deserialize, Debug, Clone, Named)]
242pub struct CastMessage {
243    /// The cast destination.
244    pub dest: Uslice,
245    /// The message to cast.
246    pub message: CastMessageEnvelope,
247}
248wirevalue::register_type!(CastMessage);
249
250/// Forward a message to procs of next hops. This is used by comm actor to
251/// forward a message to other comm actors following the selection topology.
252/// This message is not visible to the clients.
253#[derive(Serialize, Deserialize, Debug, Clone, Named)]
254pub(crate) struct ForwardMessage {
255    /// The comm actor who originally casted the message.
256    pub(crate) sender: ActorAddr,
257    /// The destination of the message.
258    pub(crate) dests: Vec<RoutingFrame>,
259    /// The sequence number of this message.
260    pub(crate) seq: usize,
261    /// The sequence number of the previous message receieved.
262    pub(crate) last_seq: usize,
263    /// The message to distribute.
264    pub(crate) message: CastMessageEnvelope,
265}
266wirevalue::register_type!(ForwardMessage);
267
268/// The is used to start casting a message to a group of actors.
269#[derive(Serialize, Deserialize, Debug, Clone, Named)]
270pub(crate) struct CastMessageV1 {
271    /// The additional end-to-end message headers.
272    pub(super) headers: Flattrs,
273    /// The client who sent this message.
274    pub(super) sender: ActorAddr,
275    /// The client-assigned session id of this message.
276    pub(super) session_id: Uuid,
277    /// The client-assigned sequence numbers of this message.
278    pub(super) seqs: ValueMesh<u64>,
279    /// The destination mesh's region.
280    pub(super) dest_region: Region,
281    /// The destination port of the message. It could match multiple actors with
282    /// rank wildcard.
283    pub(super) dest_port: DestinationPort,
284    /// The serialized message.
285    pub(super) data: wirevalue::Any<wirevalue::encoding::Multipart>,
286}
287
288impl CastEnvelope for CastMessageV1 {
289    fn sender(&self) -> &ActorAddr {
290        &self.sender
291    }
292
293    fn headers(&self) -> &Flattrs {
294        &self.headers
295    }
296
297    fn dest_port(&self) -> &DestinationPort {
298        &self.dest_port
299    }
300
301    fn data(&self) -> &wirevalue::Any<wirevalue::encoding::Multipart> {
302        &self.data
303    }
304
305    fn data_mut(&mut self) -> &mut wirevalue::Any<wirevalue::encoding::Multipart> {
306        &mut self.data
307    }
308
309    fn cast_point(&self, config: &CommMeshConfig) -> anyhow::Result<Point> {
310        let rank_on_root_mesh = config.self_rank();
311        let cast_point = self.dest_region.point_of_base_rank(rank_on_root_mesh)?;
312        Ok(cast_point)
313    }
314}
315
316#[cfg(test)]
317impl CastMessageV1 {
318    /// Create a new CastMessageEnvelope.
319    pub(crate) fn new<A, M>(
320        sender: ActorAddr,
321        dest_mesh: &ActorMeshId,
322        dest_region: Region,
323        headers: Flattrs,
324        message: M,
325        session_id: Uuid,
326        seqs: ValueMesh<u64>,
327    ) -> Result<Self, anyhow::Error>
328    where
329        A: Referable + RemoteHandles<M>,
330        M: RemoteMessage,
331    {
332        let data = wirevalue::Any::<wirevalue::encoding::Multipart>::serialize(&message)?;
333        Ok(Self {
334            headers,
335            sender,
336            session_id,
337            seqs,
338            dest_region,
339            dest_port: DestinationPort::new::<A, M>(dest_mesh.uid().clone()),
340            data,
341        })
342    }
343}
344
345/// Forward a message to procs of next hops. This is used by comm actor to
346/// forward a message to other comm actors following the selection topology.
347/// This message is not visible to the clients.
348#[derive(Serialize, Deserialize, Debug, Clone, Named)]
349pub(super) struct ForwardMessageV1 {
350    /// The destination of the message.
351    pub(super) dests: Vec<RoutingFrame>,
352    /// The message to distribute.
353    pub(super) message: CastMessageV1,
354}
355
356pub use hyperactor_cast::cast_actor::CAST_ORIGINATING_SENDER;
357pub use hyperactor_cast::cast_actor::CAST_POINT;
358
359pub fn set_cast_info_on_headers(headers: &mut Flattrs, cast_point: Point, sender: ActorAddr) {
360    // Pre-set the telemetry sender hash to the originating actor,
361    // so post_unchecked() does not overwrite it with the comm actor.
362    // TODO: consider merging SENDER_ACTOR_ID_HASH and
363    // CAST_ORIGINATING_SENDER -- they carry overlapping sender identity.
364    headers.set(
365        hyperactor::mailbox::headers::SENDER_ACTOR_ID_HASH,
366        hyperactor_telemetry::hash_to_u64(sender.id()),
367    );
368    headers.set(CAST_POINT, cast_point);
369    headers.set(CAST_ORIGINATING_SENDER, sender);
370}
371
372pub trait CastInfo {
373    /// Get the cast rank and cast shape.
374    /// If something wasn't explicitly sent via a cast, then
375    /// we represent it as the only member of a 0-dimensonal cast shape,
376    /// which is the same as a singleton.
377    fn cast_point(&self) -> Point;
378    fn sender(&self) -> ActorAddr;
379}
380
381impl<A: Actor> CastInfo for Context<'_, A> {
382    fn cast_point(&self) -> Point {
383        match self.headers().get(CAST_POINT) {
384            Some(point) => point,
385            None => Extent::unity().point_of_rank(0).unwrap(),
386        }
387    }
388
389    fn sender(&self) -> ActorAddr {
390        self.headers()
391            .get(CAST_ORIGINATING_SENDER)
392            .expect("has sender header")
393    }
394}
hyperactor_mesh/comm/multicast.rs

hyperactor_mesh/comm/
multicast.rs
