chore: move errors to atoma-types #51

Cargo.toml

@@ -10,7 +10,7 @@ members = [
     "atoma-node",
     "atoma-json-rpc",
     "atoma-storage",
-    "atoma-types"
+    "atoma-types",
 ]
 
 [workspace.package]
@@ -36,10 +36,13 @@ ethers = "2.0.14"
 futures = "0.3.30"
 futures-util = "0.3.30"
 hf-hub = "0.3.2"
-image = { version = "0.25.0", default-features = false, features = ["jpeg", "png"] }
+image = { version = "0.25.0", default-features = false, features = [
+    "jpeg",
+    "png",
+] }
 serde = "1.0.197"
 serde_json = "1.0.114"
-sui-sdk = { git = "https://github.com/mystenlabs/sui", package = "sui-sdk"}
+sui-sdk = { git = "https://github.com/mystenlabs/sui", package = "sui-sdk" }
 # solana-client = "1.18.9"
 # solana-sdk = "1.18.8"
 rand = "0.8.5"

atoma-inference/src/apis/hugging_face.rs

@@ -1,6 +1,7 @@
 use std::path::PathBuf;
 
 use async_trait::async_trait;
+use atoma_types::ApiError;
 use hf_hub::{
     api::sync::{Api, ApiBuilder},
     Repo, RepoType,
@@ -9,7 +10,7 @@ use tracing::error;
 
 use crate::models::ModelId;
 
-use super::{ApiError, ApiTrait};
+use super::ApiTrait;
 
 #[async_trait]
 impl ApiTrait for Api {

atoma-inference/src/apis/mod.rs

@@ -1,26 +1,11 @@
 pub mod hugging_face;
-use hf_hub::api::sync::ApiError as HuggingFaceError;
 
 use std::path::PathBuf;
 
-use thiserror::Error;
+use atoma_types::ApiError;
 
 use crate::models::ModelId;
 
-#[derive(Debug, Error)]
-pub enum ApiError {
-    #[error("Api Error: `{0}`")]
-    ApiError(String),
-    #[error("HuggingFace API error: `{0}`")]
-    HuggingFaceError(HuggingFaceError),
-}
-
-impl From<HuggingFaceError> for ApiError {
-    fn from(error: HuggingFaceError) -> Self {
-        Self::HuggingFaceError(error)
-    }
-}
-
 pub trait ApiTrait: Send {
     fn fetch(&self, model_id: ModelId, revision: String) -> Result<Vec<PathBuf>, ApiError>;
     fn create(api_key: String, cache_dir: PathBuf) -> Result<Self, ApiError>

atoma-inference/src/main.rs

@@ -1,8 +1,5 @@
-use atoma_inference::{
-    jrpc_server,
-    models::config::ModelsConfig,
-    service::{ModelService, ModelServiceError},
-};
+use atoma_inference::{jrpc_server, models::config::ModelsConfig, service::ModelService};
+use atoma_types::ModelServiceError;
 use ed25519_consensus::SigningKey as PrivateKey;
 
 #[tokio::main]

atoma-inference/src/model_thread.rs

@@ -1,56 +1,26 @@
-use std::{
-    collections::HashMap, fmt::Debug, path::PathBuf, str::FromStr, sync::mpsc, thread::JoinHandle,
-};
+use std::{collections::HashMap, path::PathBuf, str::FromStr, sync::mpsc, thread::JoinHandle};
 
-use atoma_types::{Request, Response};
+use atoma_types::{ModelThreadError, Request, Response};
 use ed25519_consensus::VerificationKey as PublicKey;
 use futures::stream::FuturesUnordered;
-use thiserror::Error;
-use tokio::sync::oneshot::{self, error::RecvError};
+use tokio::sync::oneshot;
 use tracing::{debug, error, info, warn};
 
-use crate::{
-    apis::ApiError,
-    models::{
-        candle::{
-            falcon::FalconModel, llama::LlamaModel, mamba::MambaModel, mistral::MistralModel,
-            mixtral::MixtralModel, quantized::QuantizedModel, stable_diffusion::StableDiffusion,
-        },
-        config::{ModelConfig, ModelsConfig},
-        types::ModelType,
-        ModelError, ModelId, ModelTrait,
+use crate::models::{
+    candle::{
+        falcon::FalconModel, llama::LlamaModel, mamba::MambaModel, mistral::MistralModel,
+        mixtral::MixtralModel, quantized::QuantizedModel, stable_diffusion::StableDiffusion,
     },
+    config::{ModelConfig, ModelsConfig},
+    types::ModelType,
+    ModelId, ModelTrait,
 };
 
 pub struct ModelThreadCommand {
     pub(crate) request: Request,
     pub(crate) sender: oneshot::Sender<Response>,
 }
 
-#[derive(Debug, Error)]
-pub enum ModelThreadError {
-    #[error("Model thread shutdown: `{0}`")]
-    ApiError(ApiError),
-    #[error("Model thread shutdown: `{0}`")]
-    ModelError(ModelError),
-    #[error("Core thread shutdown: `{0}`")]
-    Shutdown(RecvError),
-    #[error("Serde error: `{0}`")]
-    SerdeError(#[from] serde_json::Error),
-}
-
-impl From<ModelError> for ModelThreadError {
-    fn from(error: ModelError) -> Self {
-        Self::ModelError(error)
-    }
-}
-
-impl From<ApiError> for ModelThreadError {
-    fn from(error: ApiError) -> Self {
-        Self::ApiError(error)
-    }
-}
-
 pub struct ModelThreadHandle {
     sender: mpsc::Sender<ModelThreadCommand>,
     join_handle: std::thread::JoinHandle<Result<(), ModelThreadError>>,

atoma-inference/src/models/candle/falcon.rs

@@ -1,5 +1,6 @@
 use std::{path::PathBuf, str::FromStr, time::Instant};
 
+use atoma_types::ModelError;
 use candle::{DType, Device, Tensor};
 use candle_nn::VarBuilder;
 use candle_transformers::{
@@ -15,7 +16,7 @@ use crate::models::{
     candle::hub_load_safetensors,
     config::ModelConfig,
     types::{LlmLoadData, ModelType, TextModelInput, TextModelOutput},
-    ModelError, ModelTrait,
+    ModelTrait,
 };
 
 use super::device;

atoma-inference/src/models/candle/llama.rs

@@ -1,5 +1,6 @@
 use std::{path::PathBuf, str::FromStr, time::Instant};
 
+use atoma_types::ModelError;
 use candle::{DType, Device, Tensor};
 use candle_nn::VarBuilder;
 use candle_transformers::{
@@ -16,7 +17,7 @@ use crate::models::{
     config::ModelConfig,
     token_output_stream::TokenOutputStream,
     types::{LlmLoadData, ModelType, TextModelInput, TextModelOutput},
-    ModelError, ModelTrait,
+    ModelTrait,
 };
 
 use super::{device, hub_load_safetensors};

atoma-inference/src/models/candle/mamba.rs

@@ -1,5 +1,6 @@
 use std::{path::PathBuf, str::FromStr, time::Instant};
 
+use atoma_types::{bail, ModelError};
 use candle::{DType, Device, Tensor};
 use candle_nn::VarBuilder;
 use candle_transformers::{
@@ -11,15 +12,12 @@ use hf_hub::{api::sync::ApiBuilder, Repo, RepoType};
 use tokenizers::Tokenizer;
 use tracing::info;
 
-use crate::{
-    bail,
-    models::{
-        candle::device,
-        config::ModelConfig,
-        token_output_stream::TokenOutputStream,
-        types::{LlmLoadData, ModelType, TextModelInput, TextModelOutput},
-        ModelError, ModelTrait,
-    },
+use crate::models::{
+    candle::device,
+    config::ModelConfig,
+    token_output_stream::TokenOutputStream,
+    types::{LlmLoadData, ModelType, TextModelInput, TextModelOutput},
+    ModelTrait,
 };
 
 pub struct MambaModel {

atoma-inference/src/models/candle/mistral.rs

@@ -1,5 +1,6 @@
 use std::str::FromStr;
 
+use atoma_types::{bail, ModelError};
 use candle::{DType, Device, Tensor};
 use candle_nn::VarBuilder;
 use candle_transformers::{
@@ -11,14 +12,11 @@ use hf_hub::{api::sync::ApiBuilder, Repo, RepoType};
 use tokenizers::Tokenizer;
 use tracing::info;
 
-use crate::{
-    bail,
-    models::{
-        candle::{device, hub_load_safetensors},
-        token_output_stream::TokenOutputStream,
-        types::{LlmLoadData, ModelType, TextModelInput, TextModelOutput},
-        ModelError, ModelTrait,
-    },
+use crate::models::{
+    candle::{device, hub_load_safetensors},
+    token_output_stream::TokenOutputStream,
+    types::{LlmLoadData, ModelType, TextModelInput, TextModelOutput},
+    ModelTrait,
 };
 
 pub struct MistralModel {

atoma-inference/src/models/candle/mixtral.rs

@@ -1,3 +1,4 @@
+use atoma_types::{bail, ModelError};
 use candle::{DType, Device, Tensor};
 use candle_nn::VarBuilder;
 use candle_transformers::{
@@ -9,14 +10,11 @@ use hf_hub::{api::sync::ApiBuilder, Repo, RepoType};
 use tokenizers::Tokenizer;
 use tracing::info;
 
-use crate::{
-    bail,
-    models::{
-        candle::{device, hub_load_safetensors},
-        token_output_stream::TokenOutputStream,
-        types::{LlmLoadData, ModelType, TextModelInput, TextModelOutput},
-        ModelError, ModelTrait,
-    },
+use crate::models::{
+    candle::{device, hub_load_safetensors},
+    token_output_stream::TokenOutputStream,
+    types::{LlmLoadData, ModelType, TextModelInput, TextModelOutput},
+    ModelTrait,
 };
 
 pub struct MixtralModel {

atoma-inference/src/models/candle/mod.rs

@@ -1,15 +1,12 @@
 use std::{fs::File, io::Write, path::PathBuf};
 
+use atoma_types::{bail, ModelError};
 use candle::{
     utils::{cuda_is_available, metal_is_available},
     DType, Device, Tensor,
 };
 use tracing::info;
 
-use crate::bail;
-
-use super::ModelError;
-
 pub mod falcon;
 pub mod llama;
 pub mod mamba;

atoma-inference/src/models/candle/quantized.rs

@@ -1,5 +1,6 @@
 use std::{path::PathBuf, str::FromStr};
 
+use atoma_types::ModelError;
 use candle::{
     quantized::{ggml_file, gguf_file},
     DType, Device, Tensor,
@@ -17,7 +18,7 @@ use crate::models::{
     config::ModelConfig,
     token_output_stream::TokenOutputStream,
     types::{LlmLoadData, ModelType, TextModelInput, TextModelOutput},
-    ModelError, ModelTrait,
+    ModelTrait,
 };
 use candle_transformers::models::quantized_llama as model;
 

atoma-inference/src/models/candle/stable_diffusion.rs

@@ -1,5 +1,6 @@
 use std::{path::PathBuf, str::FromStr, time::Instant};
 
+use atoma_types::{bail, ModelError};
 use candle_transformers::models::stable_diffusion::{
     self, clip::ClipTextTransformer, unet_2d::UNet2DConditionModel, vae::AutoEncoderKL,
     StableDiffusionConfig,
@@ -11,11 +12,8 @@ use serde::Deserialize;
 use tokenizers::Tokenizer;
 use tracing::{debug, info};
 
-use crate::{
-    bail,
-    models::{
-        candle::save_image, config::ModelConfig, types::ModelType, ModelError, ModelId, ModelTrait,
-    },
+use crate::models::{
+    candle::save_image, config::ModelConfig, types::ModelType, ModelId, ModelTrait,
 };
 
 use super::{convert_to_image, device, save_tensor_to_file};

atoma-inference/src/models/mod.rs

@@ -1,11 +1,9 @@
 use std::path::PathBuf;
 
-use ::candle::{DTypeParseError, Error as CandleError};
+use self::{config::ModelConfig, types::ModelType};
+use atoma_types::ModelError;
 use ed25519_consensus::VerificationKey as PublicKey;
 use serde::{de::DeserializeOwned, Serialize};
-use thiserror::Error;
-
-use self::{config::ModelConfig, types::ModelType};
 
 pub mod candle;
 pub mod config;
@@ -45,40 +43,3 @@ pub trait Response: Send + 'static {
 
     fn from_model_output(model_output: Self::ModelOutput) -> Self;
 }
-
-#[derive(Debug, Error)]
-pub enum ModelError {
-    #[error("Deserialize error: `{0}`")]
-    DeserializeError(#[from] serde_json::Error),
-    #[error("{0}")]
-    Msg(String),
-    #[error("Candle error: `{0}`")]
-    CandleError(#[from] CandleError),
-    #[error("Config error: `{0}`")]
-    Config(String),
-    #[error("Image error: `{0}`")]
-    ImageError(#[from] image::ImageError),
-    #[error("Io error: `{0}`")]
-    IoError(#[from] std::io::Error),
-    #[error("Error: `{0}`")]
-    BoxedError(#[from] Box<dyn std::error::Error + Send + Sync>),
-    #[error("ApiError error: `{0}`")]
-    ApiError(#[from] hf_hub::api::sync::ApiError),
-    #[error("DTypeParseError: `{0}`")]
-    DTypeParseError(#[from] DTypeParseError),
-    #[error("Invalid model type: `{0}`")]
-    InvalidModelType(String),
-}
-
-#[macro_export]
-macro_rules! bail {
-    ($msg:literal $(,)?) => {
-        return Err(ModelError::Msg(format!($msg).into()))
-    };
-    ($err:expr $(,)?) => {
-        return Err(ModelError::Msg(format!($err).into()).bt())
-    };
-    ($fmt:expr, $($arg:tt)*) => {
-        return Err(ModelError::Msg(format!($fmt, $($arg)*).into()).bt())
-    };
-}

atoma-inference/src/models/token_output_stream.rs

@@ -1,4 +1,4 @@
-use crate::{bail, models::ModelError};
+use atoma_types::{bail, ModelError};
 
 /// This is a wrapper around a tokenizer to ensure that tokens can be returned to the user in a
 /// streaming way rather than having to wait for the full decoding.

atoma-inference/src/models/types.rs

@@ -1,12 +1,13 @@
 use std::{fmt::Display, path::PathBuf, str::FromStr};
 
+use atoma_types::ModelError;
 use candle::{DType, Device};
 use ed25519_consensus::VerificationKey as PublicKey;
 use serde::{Deserialize, Serialize};
 
 use crate::models::{ModelId, Request, Response};
 
-use super::{candle::stable_diffusion::StableDiffusionInput, ModelError};
+use super::candle::stable_diffusion::StableDiffusionInput;
 
 pub type NodeId = PublicKey;