Skip to content

Commit

Permalink
added implementation that requires new cli argument
Browse files Browse the repository at this point in the history
  • Loading branch information
Edwinhr716 committed Jul 25, 2024
1 parent 1d1b1ef commit c27075d
Show file tree
Hide file tree
Showing 3 changed files with 21 additions and 0 deletions.
6 changes: 6 additions & 0 deletions launcher/src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1162,6 +1162,7 @@ fn spawn_webserver(
max_input_tokens: usize,
max_total_tokens: usize,
max_batch_prefill_tokens: u32,
startup_time: u64,
shutdown: Arc<AtomicBool>,
shutdown_receiver: &mpsc::Receiver<()>,
) -> Result<Child, LauncherError> {
Expand Down Expand Up @@ -1199,6 +1200,8 @@ fn spawn_webserver(
format!("{}-0", args.shard_uds_path),
"--tokenizer-name".to_string(),
args.model_id,
"--startup-time".to_string(),
startup_time.to_string(),
];

// Grammar support
Expand Down Expand Up @@ -1341,6 +1344,7 @@ fn terminate(process_name: &str, mut process: Child, timeout: Duration) -> io::R
fn main() -> Result<(), LauncherError> {
// Pattern match configuration
let args: Args = Args::parse();
let start_time = Instant::now();

// Filter events with LOG_LEVEL
let varname = "LOG_LEVEL";
Expand Down Expand Up @@ -1622,12 +1626,14 @@ fn main() -> Result<(), LauncherError> {
return Ok(());
}

let download_time = start_time.elapsed().as_secs();
let mut webserver = spawn_webserver(
num_shard,
args,
max_input_tokens,
max_total_tokens,
max_batch_prefill_tokens,
download_time,
shutdown.clone(),
&shutdown_receiver,
)
Expand Down
6 changes: 6 additions & 0 deletions router/src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -87,6 +87,8 @@ struct Args {
disable_grammar_support: bool,
#[clap(default_value = "4", long, env)]
max_client_batch_size: usize,
#[clap(long, env)]
startup_time: u64,
}

#[derive(Debug, Subcommand)]
Expand Down Expand Up @@ -129,6 +131,7 @@ async fn main() -> Result<(), RouterError> {
disable_grammar_support,
max_client_batch_size,
command,
startup_time,
} = args;

let print_schema_command = match command {
Expand Down Expand Up @@ -378,6 +381,8 @@ async fn main() -> Result<(), RouterError> {
}
};

tracing::info!("start time of the model is {startup_time}");

// Run server
server::run(
master_shard_uds_path,
Expand Down Expand Up @@ -409,6 +414,7 @@ async fn main() -> Result<(), RouterError> {
disable_grammar_support,
max_client_batch_size,
print_schema_command,
startup_time,
)
.await?;
Ok(())
Expand Down
9 changes: 9 additions & 0 deletions router/src/server.rs
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,7 @@ use tower_http::cors::{AllowOrigin, CorsLayer};
use tracing::{info_span, instrument, Instrument};
use utoipa::OpenApi;
use utoipa_swagger_ui::SwaggerUi;
use tokio::time::Duration;

/// Generate tokens if `stream == false` or a stream of token if `stream == true`
#[utoipa::path(
Expand Down Expand Up @@ -1433,6 +1434,7 @@ pub async fn run(
grammar_support: bool,
max_client_batch_size: usize,
print_schema_command: bool,
start_time: u64,
) -> Result<(), WebServerError> {
// OpenAPI documentation
#[derive(OpenApi)]
Expand Down Expand Up @@ -1512,6 +1514,7 @@ pub async fn run(
)
)]
struct ApiDoc;
let length_time = Instant::now();

// Create state
if print_schema_command {
Expand Down Expand Up @@ -1892,6 +1895,12 @@ pub async fn run(
.layer(cors_layer);

tracing::info!("Connected");
let total_time = length_time.elapsed() + Duration::from_secs(start_time);
tracing::info!("total time for router to boot up and connect to model server {:?}", length_time.elapsed());
tracing::info!("the total time in secs of boot time is {:?}", total_time);
metrics::gauge!("tgi_model_load_time").set(total_time.as_secs_f64());



if ngrok {
#[cfg(feature = "ngrok")]
Expand Down

0 comments on commit c27075d

Please sign in to comment.