From 8e14780bf4e26f29452a6acf89ee64489a646ea2 Mon Sep 17 00:00:00 2001 From: Karol Damaszke Date: Mon, 4 Mar 2024 12:17:24 +0100 Subject: [PATCH] Wait 2sec once shard is ready to improve stability (#92) (#94) Co-authored-by: mswiniarsk <156412439+mswiniarsk@users.noreply.github.com> --- launcher/src/main.rs | 3 ++- router/src/main.rs | 4 +++- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/launcher/src/main.rs b/launcher/src/main.rs index 510dd33f..bbb4c8df 100644 --- a/launcher/src/main.rs +++ b/launcher/src/main.rs @@ -463,7 +463,7 @@ fn shard_manager( envs.push(("MAX_TOTAL_TOKENS".into(), max_total_tokens.to_string().into())); // Torch Distributed Env vars - if world_size == 1 { + if world_size == 1 { envs.push(("RANK".into(), rank.to_string().into())); } envs.push(("WORLD_SIZE".into(), world_size.to_string().into())); @@ -603,6 +603,7 @@ fn shard_manager( // Shard is ready if uds.exists() && !ready { tracing::info!("Shard ready in {:?}", start_time.elapsed()); + sleep(Duration::from_millis(2000)); status_sender.send(ShardStatus::Ready).unwrap(); ready = true; } else if !ready && wait_time.elapsed() > Duration::from_secs(10) { diff --git a/router/src/main.rs b/router/src/main.rs index c4406ab6..fbe0bca5 100644 --- a/router/src/main.rs +++ b/router/src/main.rs @@ -142,7 +142,9 @@ fn main() -> Result<(), RouterError> { // This will only be used to validate payloads let local_path = Path::new(&tokenizer_name); let local_model = local_path.exists() && local_path.is_dir(); - let skip_tokenizer_in_tgi = env::var("SKIP_TOKENIZER_IN_TGI").ok().map_or(false, |value| value.to_lowercase() == "true"); + let skip_tokenizer_in_tgi = env::var("SKIP_TOKENIZER_IN_TGI") + .ok() + .map_or(false, |value| value.to_lowercase() == "true"); let tokenizer = if skip_tokenizer_in_tgi { None } else if local_model {