heat_sdk/command/
train.rs

1use burn::{config::Config, module::Module, tensor::backend::Backend};
2
3use crate::client::HeatClient;
4
5#[derive(Debug, Clone)]
6pub struct MultiDevice<B: Backend>(pub Vec<B::Device>);
7
8#[derive(Debug, Clone)]
9pub struct TrainCommandContext<B: Backend> {
10    client: HeatClient,
11    devices: Vec<B::Device>,
12    config: String,
13}
14
15impl<B: Backend> TrainCommandContext<B> {
16    pub fn new(client: HeatClient, devices: Vec<B::Device>, config: String) -> Self {
17        Self {
18            client,
19            devices,
20            config,
21        }
22    }
23
24    pub fn client(&mut self) -> &mut HeatClient {
25        &mut self.client
26    }
27
28    pub fn devices(&mut self) -> &mut Vec<B::Device> {
29        &mut self.devices
30    }
31
32    pub fn config(&self) -> &str {
33        &self.config
34    }
35}
36
37trait FromTrainCommandContext<B: Backend> {
38    fn from_context(context: &TrainCommandContext<B>) -> Self;
39}
40
41impl<B: Backend> FromTrainCommandContext<B> for HeatClient {
42    fn from_context(context: &TrainCommandContext<B>) -> Self {
43        context.client.clone()
44    }
45}
46
47impl<B: Backend> FromTrainCommandContext<B> for MultiDevice<B> {
48    fn from_context(context: &TrainCommandContext<B>) -> Self {
49        MultiDevice(context.devices.clone())
50    }
51}
52
53impl<B: Backend> IntoIterator for MultiDevice<B> {
54    type Item = B::Device;
55    type IntoIter = std::vec::IntoIter<Self::Item>;
56
57    fn into_iter(self) -> Self::IntoIter {
58        self.0.into_iter()
59    }
60}
61
62impl<B: Backend, T: Config> FromTrainCommandContext<B> for T {
63    fn from_context(context: &TrainCommandContext<B>) -> Self {
64        T::load_binary(context.config.as_bytes()).expect("Config should be loaded")
65    }
66}
67
68pub trait TrainCommandHandler<B: Backend, T, M: Module<B>, E: Into<Box<dyn std::error::Error>>> {
69    fn call(self, context: TrainCommandContext<B>) -> Result<M, E>;
70}
71
72impl<F, M, B, E: Into<Box<dyn std::error::Error>>> TrainCommandHandler<B, (), M, E> for F
73where
74    F: Fn() -> Result<M, E>,
75    M: Module<B>,
76    B: Backend,
77{
78    fn call(self, _context: TrainCommandContext<B>) -> Result<M, E> {
79        (self)()
80    }
81}
82
83impl<F, T, M, B, E: Into<Box<dyn std::error::Error>>> TrainCommandHandler<B, (T,), M, E> for F
84where
85    F: Fn(T) -> Result<M, E>,
86    T: FromTrainCommandContext<B>,
87    M: Module<B>,
88    B: Backend,
89{
90    fn call(self, context: TrainCommandContext<B>) -> Result<M, E> {
91        (self)(T::from_context(&context))
92    }
93}
94
95impl<F, T1, T2, M, B, E: Into<Box<dyn std::error::Error>>> TrainCommandHandler<B, (T1, T2), M, E>
96    for F
97where
98    F: Fn(T1, T2) -> Result<M, E>,
99    T1: FromTrainCommandContext<B>,
100    T2: FromTrainCommandContext<B>,
101    M: Module<B>,
102    B: Backend,
103{
104    fn call(self, context: TrainCommandContext<B>) -> Result<M, E> {
105        (self)(T1::from_context(&context), T2::from_context(&context))
106    }
107}
108
109impl<F, T1, T2, T3, M, B, E: Into<Box<dyn std::error::Error>>>
110    TrainCommandHandler<B, (T1, T2, T3), M, E> for F
111where
112    F: Fn(T1, T2, T3) -> Result<M, E>,
113    T1: FromTrainCommandContext<B>,
114    T2: FromTrainCommandContext<B>,
115    T3: FromTrainCommandContext<B>,
116    M: Module<B>,
117    B: Backend,
118{
119    fn call(self, context: TrainCommandContext<B>) -> Result<M, E> {
120        (self)(
121            T1::from_context(&context),
122            T2::from_context(&context),
123            T3::from_context(&context),
124        )
125    }
126}
heat_sdk/command/train.rs

heat_sdk/command/
train.rs
