Clean url validation (#18)

* Suggest removing the docker, and instead installing the CLI.

* Putting back the profile.

* Getting the API from environment (to catch other things like `HF_HOME`).

* Since we want a URL from the start we can actually use a URL all the
way.

Fixing.

* Fixing the URL handling.

---------

Co-authored-by: Hugo Larcher <[email protected]>

Author: Narsil

src/lib.rs

@@ -14,6 +14,7 @@ pub use crate::requests::TokenizeOptions;
 use chrono::Local;
 use crossterm::ExecutableCommand;
 use log::{debug, error, info, warn, Level, LevelFilter};
+use reqwest::Url;
 use tokenizers::{FromPretrainedParameters, Tokenizer};
 use tokio::sync::broadcast::Sender;
 use tokio::sync::Mutex;
@@ -32,7 +33,7 @@ mod table;
 mod writers;
 
 pub struct RunConfiguration {
-    pub url: String,
+    pub url: Url,
     pub tokenizer_name: String,
     pub profile: Option<String>,
     pub max_vus: u64,
@@ -85,7 +86,7 @@ pub async fn run(mut run_config: RunConfiguration, stop_sender: Sender<()>) -> a
     let tokenizer = Arc::new(tokenizer);
     let backend = OpenAITextGenerationBackend::try_new(
         "".to_string(),
-        run_config.url.clone(),
+        run_config.url,
         run_config.model_name.clone(),
         tokenizer,
         run_config.duration,

src/main.rs

@@ -45,8 +45,8 @@ struct Args {
     warmup: Duration,
     /// The URL of the backend to benchmark. Must be compatible with OpenAI Message API
     #[clap(default_value = "http://localhost:8000", short, long, env)]
-    #[arg(value_parser = parse_url)]
-    url: String,
+    url: Url,
+
     /// Disable console UI
     #[clap(short, long, env)]
     no_console: bool,
@@ -115,13 +115,6 @@ fn parse_duration(s: &str) -> Result<Duration, Error> {
     humantime::parse_duration(s).map_err(|_| Error::new(InvalidValue))
 }
 
-fn parse_url(s: &str) -> Result<String, Error> {
-    match Url::parse(s) {
-        Ok(_) => Ok(s.to_string()),
-        Err(_) => Err(Error::new(InvalidValue)),
-    }
-}
-
 fn parse_key_val(s: &str) -> Result<HashMap<String, String>, Error> {
     let mut key_val_map = HashMap::new();
     let items = s.split(",").collect::<Vec<&str>>();

src/requests.rs

@@ -6,6 +6,7 @@ use log::{debug, error, info, trace, warn};
 use rand_distr::Distribution;
 use rayon::iter::split;
 use rayon::prelude::*;
+use reqwest::Url;
 use reqwest_eventsource::{Error, Event, EventSource};
 use serde::{Deserialize, Serialize};
 use std::cmp::Ordering;
@@ -58,7 +59,7 @@ impl Clone for Box<dyn TextGenerationBackend + Send + Sync> {
 #[derive(Debug, Clone)]
 pub struct OpenAITextGenerationBackend {
     pub api_key: String,
-    pub base_url: String,
+    pub base_url: Url,
     pub model_name: String,
     pub client: reqwest::Client,
     pub tokenizer: Arc<Tokenizer>,
@@ -101,7 +102,7 @@ pub struct OpenAITextGenerationRequest {
 impl OpenAITextGenerationBackend {
     pub fn try_new(
         api_key: String,
-        base_url: String,
+        base_url: Url,
         model_name: String,
         tokenizer: Arc<Tokenizer>,
         timeout: time::Duration,
@@ -128,7 +129,9 @@ impl TextGenerationBackend for OpenAITextGenerationBackend {
         request: Arc<TextGenerationRequest>,
         sender: Sender<TextGenerationAggregatedResponse>,
     ) {
-        let url = format!("{base_url}/v1/chat/completions", base_url = self.base_url);
+        let mut url = self.base_url.clone();
+        url.set_path("/v1/chat/completions");
+        // let url = format!("{base_url}", base_url = self.base_url);
         let mut aggregated_response = TextGenerationAggregatedResponse::new(request.clone());
         let messages = vec![OpenAITextGenerationMessage {
             role: "user".to_string(),
@@ -829,7 +832,7 @@ mod tests {
                 w.write_all(b"data: [DONE]\n\n")
             })
             .create_async().await;
-        let url = s.url();
+        let url = s.url().parse().unwrap();
         let tokenizer = Arc::new(Tokenizer::from_pretrained("gpt2", None).unwrap());
         let backend = OpenAITextGenerationBackend::try_new(
             "".to_string(),
@@ -890,7 +893,7 @@ mod tests {
                 w.write_all(b"data: [DONE]\n\n")
             })
             .create_async().await;
-        let url = s.url();
+        let url = s.url().parse().unwrap();
         let tokenizer = Arc::new(Tokenizer::from_pretrained("gpt2", None).unwrap());
         let backend = OpenAITextGenerationBackend::try_new(
             "".to_string(),
@@ -975,7 +978,7 @@ mod tests {
             .with_chunked_body(|w| w.write_all(b"data: {\"error\": \"Internal server error\"}\n\n"))
             .create_async()
             .await;
-        let url = s.url();
+        let url = s.url().parse().unwrap();
         let tokenizer = Arc::new(Tokenizer::from_pretrained("gpt2", None).unwrap());
         let backend = OpenAITextGenerationBackend::try_new(
             "".to_string(),
@@ -1021,7 +1024,7 @@ mod tests {
             .with_chunked_body(|w| w.write_all(b"this is wrong\n\n"))
             .create_async()
             .await;
-        let url = s.url();
+        let url = s.url().parse().unwrap();
         let tokenizer = Arc::new(Tokenizer::from_pretrained("gpt2", None).unwrap());
         let backend = OpenAITextGenerationBackend::try_new(
             "".to_string(),
@@ -1067,7 +1070,7 @@ mod tests {
             .with_chunked_body(|w| w.write_all(b"data: {\"foo\": \"bar\"}\n\n"))
             .create_async()
             .await;
-        let url = s.url();
+        let url = s.url().parse().unwrap();
         let tokenizer = Arc::new(Tokenizer::from_pretrained("gpt2", None).unwrap());
         let backend = OpenAITextGenerationBackend::try_new(
             "".to_string(),
@@ -1117,7 +1120,7 @@ mod tests {
                 w.write_all(b"data: [DONE]\n\n")
             })
             .create_async().await;
-        let url = s.url();
+        let url = s.url().parse().unwrap();
         let tokenizer = Arc::new(Tokenizer::from_pretrained("gpt2", None).unwrap());
         let backend = OpenAITextGenerationBackend::try_new(
             "".to_string(),

src/scheduler.rs

@@ -232,7 +232,7 @@ mod tests {
                 w.write_all(b"data: [DONE]\n\n")
             })
             .create_async().await;
-        let url = s.url();
+        let url = s.url().parse().unwrap();
         let tokenizer = Arc::new(Tokenizer::from_pretrained("gpt2", None).unwrap());
         let backend = OpenAITextGenerationBackend::try_new(
             "".to_string(),