Skip to content

Commit

Permalink
Dust Apps: Improve resilience (#4552)
Browse files Browse the repository at this point in the history 
* Always retry model errors at least once

* bump data source retrieval max topk to 1024

* Bump parallelism of model execution in dust apps maps
  • Loading branch information
@spolu
spolu authored Apr 4, 2024
1 parent 3cb6a54 commit 43e726f
Show file tree
Hide file tree
Showing 7 changed files with 95 additions and 39 deletions.
2 changes: 1 addition & 1 deletion core/src/data_sources/data_source.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1310,7 +1310,7 @@ impl DataSource {
Ok(document)
}

const MAX_TOP_K_SEARCH: usize = 128;
const MAX_TOP_K_SEARCH: usize = 1024;

pub async fn search(
&self,
Expand Down
10 changes: 7 additions & 3 deletions core/src/providers/ai21.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -131,17 +131,21 @@ impl AI21LLM {
Err(ModelError {
message: format!("Ai21APIError: {}", error.detail),
retryable: Some(ModelErrorRetryOptions {
sleep: Duration::from_millis(2000),
sleep: Duration::from_millis(500),
factor: 2,
retries: 8,
retries: 1,
}),
})
}
_ => {
let error: Error = serde_json::from_slice(c)?;
Err(ModelError {
message: format!("Ai21APIError: {}", error.detail),
retryable: None,
retryable: Some(ModelErrorRetryOptions {
sleep: Duration::from_millis(500),
factor: 1,
retries: 1,
}),
})
}
}?;
Expand Down
20 changes: 16 additions & 4 deletions core/src/providers/anthropic.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@ use crate::providers::embedder::{Embedder, EmbedderVector};
use crate::providers::llm::{
ChatMessage, ChatMessageRole, LLMChatGeneration, LLMGeneration, Tokens, LLM,
};
use crate::providers::provider::{ModelError, Provider, ProviderID};
use crate::providers::provider::{ModelError, ModelErrorRetryOptions, Provider, ProviderID};
use crate::providers::tiktoken::tiktoken::anthropic_base_singleton;
use crate::run::Credentials;
use crate::utils;
Expand Down Expand Up @@ -280,7 +280,11 @@ impl AnthropicLLM {
let error: Error = serde_json::from_slice(c)?;
Err(ModelError {
message: format!("Anthropic API Error: {}", error.to_string()),
retryable: None,
retryable: Some(ModelErrorRetryOptions {
sleep: Duration::from_millis(500),
factor: 1,
retries: 1,
}),
})
}
}?;
Expand Down Expand Up @@ -537,7 +541,11 @@ impl AnthropicLLM {
"Anthropic API Error: {}",
event.error.to_string()
),
retryable: None,
retryable: Some(ModelErrorRetryOptions {
sleep: Duration::from_millis(500),
factor: 1,
retries: 1,
}),
})?;
break 'stream;
}
Expand Down Expand Up @@ -750,7 +758,11 @@ impl AnthropicLLM {
let error: Error = serde_json::from_slice(c)?;
Err(ModelError {
message: format!("Anthropic API Error: {}", error.to_string()),
retryable: None,
retryable: Some(ModelErrorRetryOptions {
sleep: Duration::from_millis(500),
factor: 1,
retries: 1,
}),
})
}
}?;
Expand Down
34 changes: 23 additions & 11 deletions core/src/providers/cohere.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -58,17 +58,21 @@ async fn api_encode(api_key: &str, text: &str) -> Result<Vec<usize>> {
Err(ModelError {
message: format!("CohereAPIError: {}", error.message),
retryable: Some(ModelErrorRetryOptions {
sleep: Duration::from_millis(2000),
sleep: Duration::from_millis(500),
factor: 2,
retries: 8,
retries: 3,
}),
})
}
_ => {
let error: Error = serde_json::from_slice(c)?;
Err(ModelError {
message: format!("CohereAPIError: {}", error.message),
retryable: None,
retryable: Some(ModelErrorRetryOptions {
sleep: Duration::from_millis(500),
factor: 1,
retries: 1,
}),
})
}
}?;
Expand Down Expand Up @@ -105,17 +109,21 @@ async fn api_decode(api_key: &str, tokens: Vec<usize>) -> Result<String> {
Err(ModelError {
message: format!("CohereAPIError: {}", error.message),
retryable: Some(ModelErrorRetryOptions {
sleep: Duration::from_millis(2000),
sleep: Duration::from_millis(500),
factor: 2,
retries: 8,
retries: 3,
}),
})
}
_ => {
let error: Error = serde_json::from_slice(c)?;
Err(ModelError {
message: format!("CohereAPIError: {}", error.message),
retryable: None,
retryable: Some(ModelErrorRetryOptions {
sleep: Duration::from_millis(500),
factor: 1,
retries: 1,
}),
})
}
}?;
Expand Down Expand Up @@ -226,9 +234,9 @@ impl CohereLLM {
Err(ModelError {
message: format!("CohereAPIError: {}", error.message),
retryable: Some(ModelErrorRetryOptions {
sleep: Duration::from_millis(2000),
sleep: Duration::from_millis(500),
factor: 2,
retries: 8,
retries: 3,
}),
})
}
Expand Down Expand Up @@ -450,17 +458,21 @@ impl CohereEmbedder {
Err(ModelError {
message: format!("CohereAPIError: {}", error.message),
retryable: Some(ModelErrorRetryOptions {
sleep: Duration::from_millis(2000),
sleep: Duration::from_millis(500),
factor: 2,
retries: 8,
retries: 3,
}),
})
}
_ => {
let error: Error = serde_json::from_slice(c)?;
Err(ModelError {
message: format!("CohereAPIError: {}", error.message),
retryable: None,
retryable: Some(ModelErrorRetryOptions {
sleep: Duration::from_millis(500),
factor: 1,
retries: 1,
}),
})
}
}?;
Expand Down
18 changes: 13 additions & 5 deletions core/src/providers/mistral.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -490,14 +490,18 @@ impl MistralAILLM {
true => Err(ModelError {
message: error.message(),
retryable: Some(ModelErrorRetryOptions {
sleep: Duration::from_millis(100),
sleep: Duration::from_millis(500),
factor: 2,
retries: 3,
}),
})?,
false => Err(ModelError {
message: error.message(),
retryable: None,
retryable: Some(ModelErrorRetryOptions {
sleep: Duration::from_millis(500),
factor: 1,
retries: 1,
}),
})?,
}
break 'stream;
Expand Down Expand Up @@ -710,14 +714,18 @@ impl MistralAILLM {
true => Err(ModelError {
message: error.message(),
retryable: Some(ModelErrorRetryOptions {
sleep: Duration::from_millis(2000),
sleep: Duration::from_millis(500),
factor: 2,
retries: 8,
retries: 3,
}),
}),
false => Err(ModelError {
message: error.message(),
retryable: None,
retryable: Some(ModelErrorRetryOptions {
sleep: Duration::from_millis(500),
factor: 1,
retries: 1,
}),
}),
}
}
Expand Down
46 changes: 33 additions & 13 deletions core/src/providers/openai.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -273,14 +273,18 @@ pub async fn streamed_completion(
true => Err(ModelError {
message: error.message(),
retryable: Some(ModelErrorRetryOptions {
sleep: Duration::from_millis(100),
sleep: Duration::from_millis(500),
factor: 2,
retries: 3,
}),
})?,
false => Err(ModelError {
message: error.message(),
retryable: None,
retryable: Some(ModelErrorRetryOptions {
sleep: Duration::from_millis(500),
factor: 1,
retries: 1,
}),
})?,
}
break 'stream;
Expand Down Expand Up @@ -500,14 +504,18 @@ pub async fn completion(
true => Err(ModelError {
message: error.message(),
retryable: Some(ModelErrorRetryOptions {
sleep: Duration::from_millis(2000),
sleep: Duration::from_millis(500),
factor: 2,
retries: 8,
retries: 3,
}),
}),
false => Err(ModelError {
message: error.message(),
retryable: None,
retryable: Some(ModelErrorRetryOptions {
sleep: Duration::from_millis(500),
factor: 1,
retries: 1,
}),
}),
}
}
Expand Down Expand Up @@ -651,14 +659,18 @@ pub async fn streamed_chat_completion(
true => Err(ModelError {
message: error.message(),
retryable: Some(ModelErrorRetryOptions {
sleep: Duration::from_millis(100),
sleep: Duration::from_millis(500),
factor: 2,
retries: 3,
}),
})?,
false => Err(ModelError {
message: error.message(),
retryable: None,
retryable: Some(ModelErrorRetryOptions {
sleep: Duration::from_millis(500),
factor: 1,
retries: 1,
}),
})?,
}
break 'stream;
Expand Down Expand Up @@ -982,14 +994,18 @@ pub async fn chat_completion(
true => Err(ModelError {
message: error.message(),
retryable: Some(ModelErrorRetryOptions {
sleep: Duration::from_millis(2000),
sleep: Duration::from_millis(500),
factor: 2,
retries: 8,
retries: 3,
}),
}),
false => Err(ModelError {
message: error.message(),
retryable: None,
retryable: Some(ModelErrorRetryOptions {
sleep: Duration::from_millis(500),
factor: 1,
retries: 1,
}),
}),
}
}
Expand Down Expand Up @@ -1078,14 +1094,18 @@ pub async fn embed(
true => Err(ModelError {
message: error.message(),
retryable: Some(ModelErrorRetryOptions {
sleep: Duration::from_millis(2000),
sleep: Duration::from_millis(500),
factor: 2,
retries: 8,
retries: 3,
}),
}),
false => Err(ModelError {
message: error.message(),
retryable: None,
retryable: Some(ModelErrorRetryOptions {
sleep: Duration::from_millis(500),
factor: 1,
retries: 1,
}),
}),
}
}
Expand Down
4 changes: 2 additions & 2 deletions core/src/run.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -47,8 +47,8 @@ impl RunConfig {
BlockType::Data => 64,
BlockType::DataSource => 8,
BlockType::Code => 64,
BlockType::LLM => 8,
BlockType::Chat => 8,
BlockType::LLM => 32,
BlockType::Chat => 32,
BlockType::Map => 64,
BlockType::Reduce => 64,
BlockType::Search => 8,
Expand Down

0 comments on commit 43e726f

Please sign in to comment.