Enum StreamMessage

pub enum StreamMessage {
    CallFunction(CallFunctionParams, HashMap<Ref, DeviceMesh>, HashMap<Ref, (DeviceMesh, Vec<String>, Arc<ActorHandle<NcclCommActor>>)>),
    BorrowCreate {
        borrow: u64,
        tensor: Ref,
        first_use_sender: PortHandle<(Option<Event>, TensorCellResult)>,
    },
    BorrowFirstUse {
        borrow: u64,
        result: Ref,
        first_use_receiver: Arc<Mutex<PortReceiver<(Option<Event>, TensorCellResult)>>>,
    },
    BorrowLastUse {
        borrow: u64,
        result: Ref,
        last_use_sender: PortHandle<(Option<Event>, TensorCellResult)>,
    },
    BorrowDrop {
        borrow: u64,
        last_use_receiver: Arc<Mutex<PortReceiver<(Option<Event>, TensorCellResult)>>>,
    },
    DeleteRefs(Vec<Ref>),
    RequestStatus(OncePortHandle<()>),
    InitComm(ActorHandle<NcclCommActor>),
    Reduce {
        comm: Arc<ActorHandle<NcclCommActor>>,
        dim_size: i64,
        result: Ref,
        local_tensor: Ref,
        factory: Factory,
        reduction: Reduction,
        scatter: bool,
        in_place: bool,
        out: Option<Ref>,
    },
    SendTensor {
        result: Ref,
        from_rank: Option<usize>,
        to_rank: Option<usize>,
        tensor: Ref,
        factory: Factory,
        comm: Arc<ActorHandle<NcclCommActor>>,
    },
    SendValue {
        seq: Seq,
        worker_actor_id: ActorId,
        mutates: Vec<Ref>,
        function: Option<ResolvableFunction>,
        args: Vec<WireValue>,
        kwargs: HashMap<String, WireValue>,
        device_meshes: HashMap<Ref, DeviceMesh>,
        pipe: Option<PortHandle<PipeMessage>>,
    },
    SetValue {
        seq: Seq,
        results: Vec<Option<Ref>>,
        pipe: PortHandle<PipeMessage>,
    },
    DefineRecording {
        recording: Ref,
    },
    FinalizeRecording {
        recording: Ref,
    },
    CallRecording {
        seq: Seq,
        recording: Ref,
        results: Vec<Ref>,
        actuals: Vec<Ref>,
    },
    RecordingFormal {
        result: Ref,
        argument_index: usize,
    },
    RecordingResult {
        result: Ref,
        output_index: usize,
    },
    SetRefUnitTestsOnly(Ref, WireValue),
    SetTensorRefUnitTestsOnly(Ref, TensorCellResult),
    GetRefUnitTestsOnly(Ref, OncePortHandle<Option<Result<WireValue, String>>>),
    GetTensorRefUnitTestsOnly(Ref, OncePortHandle<Option<TensorCellResult>>),
    SendResultOfActorCall(ActorId, ActorCallParams),
    CallActorMethod(ActorMethodParams),
}

Messages handled by the stream. Generally these are stream-local versions of crate::WorkerMessage.

Variants

CallFunction(CallFunctionParams, HashMap<Ref, DeviceMesh>, HashMap<Ref, (DeviceMesh, Vec<String>, Arc<ActorHandle<NcclCommActor>>)>)

BorrowCreate

Fields

borrow: u64

Id for the borrow.

tensor: Ref

Tensor to borrow.

first_use_sender: PortHandle<(Option<Event>, TensorCellResult)>

Port for sending the first use CUDA event + borrowed tensor to the borrower.

BorrowFirstUse

Fields

borrow: u64

Id for the borrow.

result: Ref

Ref for storing the borrowed tensor.

first_use_receiver: Arc<Mutex<PortReceiver<(Option<Event>, TensorCellResult)>>>

Port for receiving the first use CUDA event + borrowed tensor from the provider stream.

BorrowLastUse

Fields

borrow: u64

Id for the borrow.

result: Ref

Ref for the borrowed tensor.

last_use_sender: PortHandle<(Option<Event>, TensorCellResult)>

Port for sending the last use CUDA event and borrowed tensor.

BorrowDrop

Fields

borrow: u64

last_use_receiver: Arc<Mutex<PortReceiver<(Option<Event>, TensorCellResult)>>>

Port for receiving the last use CUDA event and borrowed tensor.

DeleteRefs(Vec<Ref>)

RequestStatus(OncePortHandle<()>)

InitComm(ActorHandle<NcclCommActor>)

Reduce

Fields

comm: Arc<ActorHandle<NcclCommActor>>

dim_size: i64

result: Ref

local_tensor: Ref

factory: Factory

reduction: Reduction

scatter: bool

in_place: bool

out: Option<Ref>

SendTensor

Fields

result: Ref

from_rank: Option<usize>

to_rank: Option<usize>

tensor: Ref

factory: Factory

comm: Arc<ActorHandle<NcclCommActor>>

SendValue

Fields

seq: Seq

worker_actor_id: ActorId

mutates: Vec<Ref>

function: Option<ResolvableFunction>

args: Vec<WireValue>

kwargs: HashMap<String, WireValue>

device_meshes: HashMap<Ref, DeviceMesh>

pipe: Option<PortHandle<PipeMessage>>

SetValue

Fields

seq: Seq

results: Vec<Option<Ref>>

pipe: PortHandle<PipeMessage>

DefineRecording

Fields

recording: Ref

FinalizeRecording

Fields

recording: Ref

CallRecording

Fields

seq: Seq

recording: Ref

results: Vec<Ref>

actuals: Vec<Ref>

RecordingFormal

Fields

result: Ref

argument_index: usize

RecordingResult

Fields

result: Ref

output_index: usize

SetRefUnitTestsOnly(Ref, WireValue)

SetTensorRefUnitTestsOnly(Ref, TensorCellResult)

GetRefUnitTestsOnly(Ref, OncePortHandle<Option<Result<WireValue, String>>>)

GetTensorRefUnitTestsOnly(Ref, OncePortHandle<Option<TensorCellResult>>)

SendResultOfActorCall(ActorId, ActorCallParams)

CallActorMethod(ActorMethodParams)

Trait Implementations

impl Debug for StreamMessage

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl Handler<StreamMessage> for StreamActor

fn handle<'life0, 'life1, 'async_trait>( &'life0 mut self, cx: &'life1 Context<'_, Self>, message: StreamMessage, ) -> Pin<Box<dyn Future<Output = Result<()>> + Send + 'async_trait>>

where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait,

Handle the next M-typed message.

impl Named for StreamMessage

fn typename() -> &'static str

The globally unique type name for the type. This should typically be the fully qualified Rust name of the type.

fn typehash() -> u64

A globally unique hash for this type. TODO: actually enforce perfect hashing

fn arm(&self) -> Option<&'static str>

If the named type is an enum, this returns the name of the arm of the value self.

fn typeid() -> TypeId

The TypeId for this type. TypeIds are unique only within a binary, and should not be used for global identification.

fn port() -> u64

The globally unique port for this type. Typed ports are in the range of 1<<63..1<<64-1.

unsafe fn arm_unchecked(self_: *const ()) -> Option<&'static str>

An unsafe version of ‘arm’, accepting a pointer to the value, for use in type-erased settings.