burn-dispatch 0.23.0-pre.1

Backend dispatch for the Burn framework
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106

//! Serving Burn's own backends: the backend is the one the devices belong to.
//!
//! Lives in `burn-dispatch` because matching on [`DispatchDevice`] requires the local
//! `cube_backend` cfg set by this crate's `build.rs`, plus visibility of every in-tree `BackendIr`
//! type. The user surface is `burn::server::RemoteServer`.

#[cfg(not(target_family = "wasm"))]
use burn_remote::server::Transport;
use burn_remote::{
    Endpoint,
    server::{BackendServer, RemoteProtocol, ServeError, ServerSettings},
};

use crate::DispatchDevice;
use crate::backends::*;

/// Bind `$b` to the one backend every device of `$devices` belongs to, and `$hosted` to the
/// devices as that backend's own, then run `$body`. Autodiff is stripped: the autodiff graph is the
/// client's.
macro_rules! with_backend {
    ($devices:expr, |$b:ident, $hosted:ident| $body:expr) => {{
        let devices: Vec<DispatchDevice> =
            $devices.into_iter().map(DispatchDevice::inner).collect();
        match devices.first() {
            None => Err(ServeError::NoDevices),
            #[cfg(cube_backend)]
            Some(DispatchDevice::Cube(_)) => {
                type $b = Cube;
                let $hosted = hosted!(devices, Cube)?;
                $body
            }
            #[cfg(feature = "flex")]
            Some(DispatchDevice::Flex(_)) => {
                type $b = Flex;
                let $hosted = hosted!(devices, Flex)?;
                $body
            }
            #[cfg(feature = "remote")]
            Some(DispatchDevice::Remote(_)) => Err(ServeError::UnsupportedDevice {
                reason: "a remote device cannot host a remote server".into(),
            }),
            #[cfg(feature = "capture")]
            Some(DispatchDevice::Capture(_)) => Err(ServeError::UnsupportedDevice {
                reason: "a capture device cannot host a remote server".into(),
            }),
            #[cfg(feature = "group")]
            Some(DispatchDevice::Group(_)) => Err(ServeError::UnsupportedDevice {
                reason: "a device group cannot host a remote server".into(),
            }),
            #[cfg(feature = "autodiff")]
            Some(DispatchDevice::Autodiff(_)) => {
                unreachable!("Autodiff stripped by DispatchDevice::inner")
            }
        }
    }};
}

/// `$devices` as `$variant`'s own devices, or [`ServeError::MixedBackends`].
macro_rules! hosted {
    ($devices:expr, $variant:ident) => {
        $devices
            .into_iter()
            .map(|device| match device {
                DispatchDevice::$variant(device) => Ok(device),
                #[allow(unreachable_patterns)]
                _ => Err(ServeError::MixedBackends),
            })
            .collect::<Result<Vec<_>, ServeError>>()
    };
}

/// Serve `devices` on `transport`, blocking until Ctrl+C or `SIGTERM`.
#[cfg(not(target_family = "wasm"))]
pub fn serve(
    devices: Vec<DispatchDevice>,
    settings: ServerSettings,
    transport: Transport,
) -> Result<(), ServeError> {
    with_backend!(devices, |B, hosted| BackendServer::<B>::new(hosted)
        .with_settings(settings)
        .serve(transport))
}

/// Serve `devices` on `transport` until the returned future is dropped.
#[cfg(not(target_family = "wasm"))]
pub async fn serve_async(
    devices: Vec<DispatchDevice>,
    settings: ServerSettings,
    transport: Transport,
) -> Result<(), ServeError> {
    with_backend!(devices, |B, hosted| BackendServer::<B>::new(hosted)
        .with_settings(settings)
        .serve_async(transport)
        .await)
}

/// Burn Remote's handler for an application's Iroh router on `endpoint`, serving `devices`.
pub fn into_protocol(
    devices: Vec<DispatchDevice>,
    settings: ServerSettings,
    endpoint: &Endpoint,
) -> Result<RemoteProtocol, ServeError> {
    with_backend!(devices, |B, hosted| BackendServer::<B>::new(hosted)
        .with_settings(settings)
        .into_protocol(endpoint))
}