burn_train/checkpoint/
async_checkpoint.rs

1use super::{Checkpoint, Checkpointer, CheckpointerError};
2use crate::Interrupter;
3use std::sync::mpsc;
4
5enum Message<R> {
6    Restore(usize, mpsc::SyncSender<Result<R, CheckpointerError>>),
7    Save(usize, R),
8    Delete(usize),
9    Interrupter(Interrupter),
10    End,
11}
12
13struct CheckpointerThread<C, R> {
14    checkpointer: C,
15    receiver: mpsc::Receiver<Message<R>>,
16    interrupter: Option<Interrupter>,
17}
18
19impl<C, R> CheckpointerThread<C, R>
20where
21    C: Checkpointer<R>,
22    R: Checkpoint,
23{
24    fn new(checkpointer: C, receiver: mpsc::Receiver<Message<R>>) -> Self {
25        Self {
26            checkpointer,
27            receiver,
28            interrupter: None,
29        }
30    }
31
32    fn run(mut self) {
33        while let Ok(item) = self.receiver.recv() {
34            match item {
35                Message::Restore(epoch, callback) => {
36                    let record = self.checkpointer.restore(epoch);
37                    if let Err(err) = callback.send(record) {
38                        self.fail(
39                            "Error when sending response through callback channel",
40                            CheckpointerError::Unknown(err.to_string()),
41                        );
42                    }
43                }
44                Message::Save(epoch, state) => {
45                    if let Err(err) = self.checkpointer.save(epoch, state) {
46                        self.fail("Error when saving the state", err);
47                    }
48                }
49                Message::Delete(epoch) => {
50                    if let Err(err) = self.checkpointer.delete(epoch) {
51                        self.fail("Error when deleting the state", err);
52                    }
53                }
54                Message::Interrupter(interrupter) => {
55                    self.interrupter = Some(interrupter);
56                }
57                Message::End => {
58                    return;
59                }
60            };
61        }
62    }
63
64    /// Interrupt training with `err`, or panic when there is no interrupter to report it to.
65    fn fail(&self, context: &str, err: CheckpointerError) {
66        match &self.interrupter {
67            Some(interrupter) => interrupter.fail(err),
68            None => panic!("{context}: {err}"),
69        }
70    }
71}
72
73/// Async checkpointer.
74pub struct AsyncCheckpointer<R> {
75    sender: mpsc::SyncSender<Message<R>>,
76    handler: Option<std::thread::JoinHandle<()>>,
77}
78
79impl<R> AsyncCheckpointer<R>
80where
81    R: Checkpoint,
82{
83    /// Create a new async checkpointer.
84    ///
85    /// # Arguments
86    ///
87    /// * `checkpointer` - The checkpointer.
88    ///
89    /// # Returns
90    ///
91    /// The async checkpointer.
92    pub fn new<C>(checkpointer: C) -> Self
93    where
94        C: Checkpointer<R> + Send + 'static,
95    {
96        // Only on checkpoint can be done in advance.
97        let (sender, receiver) = mpsc::sync_channel(0);
98        let thread = CheckpointerThread::new(checkpointer, receiver);
99        let handler = Some(std::thread::spawn(move || thread.run()));
100
101        Self { sender, handler }
102    }
103
104    /// Assign a handle used to interrupt training in case of checkpointing error.
105    pub fn with_interrupter(self, interrupter: Interrupter) -> Self {
106        self.sender
107            .send(Message::Interrupter(interrupter))
108            .expect("Can send message to checkpointer thread.");
109        self
110    }
111}
112
113impl<R> Checkpointer<R> for AsyncCheckpointer<R>
114where
115    R: Checkpoint,
116{
117    fn save(&self, epoch: usize, record: R) -> Result<(), CheckpointerError> {
118        self.sender
119            .send(Message::Save(epoch, record))
120            .expect("Can send message to checkpointer thread.");
121
122        Ok(())
123    }
124
125    fn restore(&self, epoch: usize) -> Result<R, CheckpointerError> {
126        let (sender, receiver) = mpsc::sync_channel(1);
127        self.sender
128            .send(Message::Restore(epoch, sender))
129            .map_err(|e| CheckpointerError::Unknown(e.to_string()))?;
130
131        if let Ok(record) = receiver.recv() {
132            return record;
133        };
134
135        Err(CheckpointerError::Unknown("Channel error.".to_string()))
136    }
137
138    fn delete(&self, epoch: usize) -> Result<(), CheckpointerError> {
139        self.sender
140            .send(Message::Delete(epoch))
141            .map_err(|e| CheckpointerError::Unknown(e.to_string()))?;
142
143        Ok(())
144    }
145}
146
147impl<E> Drop for AsyncCheckpointer<E> {
148    fn drop(&mut self) {
149        self.sender
150            .send(Message::End)
151            .expect("Can send the end message to the checkpointer thread.");
152        let handler = self.handler.take();
153
154        if let Some(handler) = handler {
155            handler
156                .join()
157                .expect("The checkpointer thread should stop.");
158        }
159    }
160}
burn_train/checkpoint/async_checkpoint.rs

burn_train/checkpoint/
async_checkpoint.rs
