hyperactor_mesh/
casting.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//! Casting utilities for actor meshes.
10
11use hyperactor::Actor;
12use hyperactor::ActorAddr;
13use hyperactor::Context;
14use hyperactor::mailbox::MailboxSenderError;
15use hyperactor::mailbox::MessageEnvelope;
16use hyperactor::mailbox::Undeliverable;
17pub use hyperactor_cast::cast_actor::CAST_ORIGINATING_SENDER;
18pub use hyperactor_cast::cast_actor::CAST_POINT;
19use hyperactor_config::Flattrs;
20use hyperactor_config::attrs::declare_attrs;
21use ndslice::Extent;
22use ndslice::Point;
23use ndslice::Selection;
24use ndslice::ShapeError;
25use ndslice::SliceError;
26use ndslice::reshape::ReshapeError;
27
28use crate::mesh_id::ActorMeshId;
29
30declare_attrs! {
31    /// Which mesh this message was cast to. Used for undeliverable message
32    /// handling when the serialized cast payload cannot be inspected.
33    pub attr CAST_ACTOR_MESH_ID: ActorMeshId;
34}
35
36/// Update a cast undeliverable with the original sender recorded in its headers.
37pub fn update_undeliverable_envelope_for_casting(
38    mut envelope: Undeliverable<MessageEnvelope>,
39) -> Undeliverable<MessageEnvelope> {
40    let Some(message) = envelope.as_message_mut() else {
41        return envelope;
42    };
43    let old_actor = message.sender().clone();
44    if let Some(actor_id) = message.headers().get(CAST_ORIGINATING_SENDER) {
45        tracing::debug!(
46            actor_id = %old_actor,
47            "remapped cast sender to id from CAST_ORIGINATING_SENDER {}", actor_id
48        );
49        message.update_sender(actor_id);
50    }
51    // Else do nothing: the message was not sent through cast routing.
52    envelope
53}
54
55pub fn set_cast_info_on_headers(headers: &mut Flattrs, cast_point: Point, sender: ActorAddr) {
56    headers.set(
57        hyperactor::mailbox::headers::SENDER_ACTOR_ID_HASH,
58        hyperactor_telemetry::hash_to_u64(sender.id()),
59    );
60    headers.set(CAST_POINT, cast_point);
61    headers.set(CAST_ORIGINATING_SENDER, sender);
62}
63
64pub trait CastInfo {
65    /// Get the cast rank and cast shape.
66    /// If something wasn't explicitly sent via a cast, represent it as the only
67    /// member of a zero-dimensional cast shape, which is the same as a singleton.
68    fn cast_point(&self) -> Point;
69    fn sender(&self) -> ActorAddr;
70}
71
72impl<A: Actor> CastInfo for Context<'_, A> {
73    fn cast_point(&self) -> Point {
74        match self.headers().get(CAST_POINT) {
75            Some(point) => point,
76            None => Extent::unity()
77                .point_of_rank(0)
78                .expect("the unity extent contains rank 0"),
79        }
80    }
81
82    fn sender(&self) -> ActorAddr {
83        self.headers()
84            .get(CAST_ORIGINATING_SENDER)
85            .expect("cast headers contain the originating sender")
86    }
87}
88
89/// The type of error of casting operations.
90#[derive(Debug, thiserror::Error)]
91pub enum CastError {
92    #[error("invalid selection {0}: {1}")]
93    InvalidSelection(Selection, ShapeError),
94
95    #[error("send on rank {0}: {1}")]
96    MailboxSenderError(usize, MailboxSenderError),
97
98    #[error("unsupported selection: {0}")]
99    SelectionNotSupported(String),
100
101    #[error(transparent)]
102    RootMailboxSenderError(#[from] MailboxSenderError),
103
104    #[error(transparent)]
105    ShapeError(#[from] ShapeError),
106
107    #[error(transparent)]
108    SliceError(#[from] SliceError),
109
110    #[error(transparent)]
111    SerializationEncodeError(#[from] bincode::error::EncodeError),
112
113    #[error(transparent)]
114    SerializationDecodeError(#[from] bincode::error::DecodeError),
115
116    #[error(transparent)]
117    Other(#[from] anyhow::Error),
118
119    #[error(transparent)]
120    ReshapeError(#[from] ReshapeError),
121}
hyperactor_mesh/casting.rs

hyperactor_mesh/
casting.rs
