use burn_core::data::dataset::DatasetError;
use crate::{
EventProcessorError,
checkpoint::CheckpointerError,
train::{MultiDeviceStepError, WorkerFailure, fmt_worker_failures},
};
#[derive(Debug)]
pub enum TrainingError {
EventProcessor(EventProcessorError),
Dataset(DatasetError),
Workers(Vec<WorkerFailure>),
Checkpoint(CheckpointerError),
}
impl TrainingError {
pub fn is_device_poisoned(&self) -> bool {
match self {
Self::EventProcessor(err) => err.is_device_poisoned(),
_ => false,
}
}
}
impl core::fmt::Display for TrainingError {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
match self {
Self::EventProcessor(err) => write!(f, "{err}"),
Self::Dataset(err) => write!(f, "Dataset error: {err}"),
Self::Workers(failures) => fmt_worker_failures(failures, f),
Self::Checkpoint(err) => write!(f, "Checkpoint error: {err}"),
}
}
}
impl core::error::Error for TrainingError {
fn source(&self) -> Option<&(dyn core::error::Error + 'static)> {
match self {
Self::EventProcessor(err) => Some(err),
Self::Dataset(err) => Some(err),
Self::Workers(_) => None,
Self::Checkpoint(err) => Some(err),
}
}
}
impl From<EventProcessorError> for TrainingError {
fn from(err: EventProcessorError) -> Self {
Self::EventProcessor(err)
}
}
impl From<DatasetError> for TrainingError {
fn from(err: DatasetError) -> Self {
Self::Dataset(err)
}
}
impl From<CheckpointerError> for TrainingError {
fn from(err: CheckpointerError) -> Self {
Self::Checkpoint(err)
}
}
impl From<MultiDeviceStepError> for TrainingError {
fn from(err: MultiDeviceStepError) -> Self {
match err {
MultiDeviceStepError::Dataset(err) => Self::Dataset(err),
MultiDeviceStepError::Workers(failures) => Self::Workers(failures),
}
}
}