Ok inital commit for server side need to run inference on the specific model once downloaded

to verify that the model is compatible with our infrastructure.

crates/lib/deprecated/config/src/bad_urls.rs

@@ -3,11 +3,16 @@
 use crate::is_bad_download_url::is_bad_download_url;
 
 /// Reports if the URL is inappropriate for downloading as a TTS model, vocoder, etc.
+/// DEPRECATE:
 pub fn is_bad_tts_model_download_url(url: &str) -> anyhow::Result<bool> {
   // TODO: Just use this.
   is_bad_download_url(url)
 }
 
+pub fn is_bad_model_weights_download_url(url: &str) -> anyhow::Result<bool> {
+  is_bad_download_url(url)
+}
+
 #[cfg(test)]
 mod tests {
   use crate::bad_urls::is_bad_tts_model_download_url;

crates/service/web/storyteller_web/src/http_server/endpoints/weights/drain_multipart_weights_request.rs

@@ -1,100 +0,0 @@
-use actix_multipart::Multipart;
-use actix_web::web::BytesMut;
-use futures::TryStreamExt;
-use log::warn;
-
-use errors::AnyhowResult;
-use tokens::tokens::weights::ModelWeightToken;
-
-use crate::http_server::web_utils::read_multipart_field_bytes::{checked_read_multipart_bytes, read_multipart_field_as_text};
-
-pub struct WeightUploadRequest {
-  pub uuid_idempotency_token: Option<String>,
-  pub model_weights_token: Option<String>,
-  pub file_name: Option<String>,
-  pub file_bytes: Option<BytesMut>,
-  pub media_source: MediaSource,
-}
-
-/// Where the frontend tells us the file came from.
-#[derive(Eq, PartialEq, Clone, Copy)]
-pub enum MediaSource {
-  Unknown,
-  UserFile,
-  /// Eg. the web audio API in Javascript.
-  UserDeviceApi,
-}
-
-/// Pull common parts out of multipart media HTTP requests, typically for handling file uploads.
-pub async fn drain_multipart_request(mut multipart_payload: Multipart) -> AnyhowResult<WeightUploadRequest> {
-  let mut uuid_idempotency_token = None;
-  let mut model_weight_token = None;
-  let mut file_bytes = None;
-  let mut file_name = None;
-  let mut media_source = None;
-
-  while let Ok(Some(mut field)) = multipart_payload.try_next().await {
-    let mut field_name = None;
-    let mut field_filename = None;
-
-    if let content_disposition = field.content_disposition() {
-      field_name = content_disposition.get_name()
-          .map(|s| s.to_string());
-      field_filename = content_disposition.get_filename() // NB: Only used for the file bytes.
-          .map(|s| s.to_string());
-    }
-
-    match field_name.as_deref() {
-      Some("uuid_idempotency_token") => {
-        uuid_idempotency_token = read_multipart_field_as_text(&mut field).await
-            .map_err(|e| {
-              warn!("Error reading idempotency token: {:}", e);
-              e
-            })?;
-      },
-      Some("dataset_token") => {
-        model_weight_token = read_multipart_field_as_text(&mut field).await
-            .map_err(|e| {
-              warn!("Error reading dataset token: {:}", e);
-              e
-            })?;
-      },
-      Some("file") => {
-        file_name = field_filename.clone();
-        file_bytes = checked_read_multipart_bytes(&mut field).await
-            .map_err(|e| {
-              warn!("Error reading audio upload: {:}", e);
-              e
-            })?;
-      },
-      Some("source") => {
-        media_source = read_multipart_field_as_text(&mut field).await
-            .map_err(|e| {
-              warn!("Error reading source: {:}", e);
-              e
-            })?;
-      },
-      _ => continue,
-    }
-  }
-
-  let media_source = match media_source.as_deref() {
-    Some("device") => MediaSource::UserDeviceApi,
-    Some("file") => MediaSource::UserFile,
-    _ => {
-      if file_name.as_deref() == Some("blob") {
-        MediaSource::UserDeviceApi
-      } else {
-        MediaSource::Unknown
-      }
-    },
-  };
-
-  Ok(UploadSampleRequest {
-    uuid_idempotency_token,
-    dataset_token,
-    file_name,
-    file_bytes,
-    media_source,
-  })
-}

crates/service/web/storyteller_web/src/http_server/endpoints/weights/upload_weight.rs

@@ -2,10 +2,10 @@ use std::collections::HashSet;
 use std::sync::Arc;
 
 use actix_multipart::Multipart;
-use actix_web::{HttpRequest, HttpResponse, web};
+use actix_web::{ HttpRequest, HttpResponse, web };
 use enums::by_table::model_weights::weights_category::WeightsCategory;
 use enums::by_table::model_weights::weights_types::WeightsType;
-use mysql_queries::queries::model_weights::create_weight::{self, create_weight};
+use mysql_queries::queries::model_weights::create_weight::{ self, create_weight };
 use once_cell::sync::Lazy;
 
 use tokens::tokens::model_weights::ModelWeightToken;
@@ -13,259 +13,180 @@ use tokens::tokens::model_weights::ModelWeightToken;
 use crate::http_server::endpoints::media_uploads::common::upload_error::UploadError;
 use crate::server_state::ServerState;
 
-use log::{error, info, warn};
+use log::{ error, info, warn };
 
 use buckets::public::weight_uploads::original_file::WeightUploadOriginalFilePath;
 use enums::by_table::media_uploads::media_upload_source::MediaUploadSource;
 use enums::by_table::media_uploads::media_upload_type::MediaUploadType;
-use enums::common::visibility::Visibility;
+
 use hashing::sha256::sha256_hash_bytes::sha256_hash_bytes;
 use http_server_common::request::get_request_ip::get_request_ip;
 use media::decode_basic_audio_info::decode_basic_audio_bytes_info;
 use mimetypes::mimetype_for_bytes::get_mimetype_for_bytes;
 use mysql_queries::queries::media_uploads::get_media_upload_by_uuid::get_media_upload_by_uuid_with_connection;
-use mysql_queries::queries::media_uploads::insert_media_upload::{Args, insert_media_upload};
+use mysql_queries::queries::media_uploads::insert_media_upload::{ Args, insert_media_upload };
 
-use crate::http_server::endpoints::media_uploads::common::drain_multipart_request::{drain_multipart_request, MediaSource};
+use crate::http_server::endpoints::media_uploads::common::drain_multipart_request::{
+    drain_multipart_request,
+    MediaSource,
+};
 use crate::validations::validate_idempotency_token_format::validate_idempotency_token_format;
 use utoipa::ToSchema;
 
+use config::bad_urls::is_bad_model_weights_download_url;
+
+use enums::by_table::model_weights::{
+    weights_category::WeightsCategory,
+    weights_types::WeightsType,
+};
+use enums::common::visibility::Visibility;
+use utoipa::ToSchema;
+
+#[derive(Deserialize, ToSchema)]
+pub struct UploadModelWeightsRequest {
+    idempotency_token: String,
+    title: String,
+    download_url: String,
+    model_weights_type: WeightsType,
+    model_weights_category: WeightsCategory,
+    creator_set_visibility: Option<Visibility>,
+}
+
+/// Tell the frontend how to deal with the download queue.
+#[derive(Serialize, ToSchema)]
+pub enum DownloadJobType {
+    /// Modern shared download type (everything will use this in the
+    /// future, and this endpoint will die.)
+    #[serde(rename = "generic")]
+    Generic,
+}
+
+#[derive(Serialize, ToSchema)]
+pub struct UploadModelWeightsSuccessResponse {
+    pub success: bool,
+    /// This is how frontend clients can request the job execution status.
+    pub job_token: String,
+    /// This is a transitional field to tell the frontend how to process the job checking.
+    pub job_type: DownloadJobType,
+}
+
+#[derive(Debug)]
+pub enum UploadModelWeightsError {
+    BadInput(String),
+    MustBeLoggedIn,
+    ServerError,
+    RateLimited,
+}
+
+impl ResponseError for UploadModelWeightsError {
+    fn status_code(&self) -> StatusCode {
+        match *self {
+            UploadModelWeightsError::BadInput(_) => StatusCode::BAD_REQUEST,
+            UploadModelWeightsError::MustBeLoggedIn => StatusCode::UNAUTHORIZED,
+            UploadModelWeightsError::ServerError => StatusCode::INTERNAL_SERVER_ERROR,
+            UploadModelWeightsError::RateLimited => StatusCode::TOO_MANY_REQUESTS,
+        }
+    }
+
+    fn error_response(&self) -> HttpResponse {
+        let error_reason = match self {
+            UploadModelWeightsError::BadInput(reason) => reason.to_string(),
+            UploadModelWeightsError::MustBeLoggedIn => "user must be logged in".to_string(),
+            UploadModelWeightsError::ServerError => "server error".to_string(),
+            UploadModelWeightsError::RateLimited => "rate limited".to_string(),
+        };
+
+        to_simple_json_error(&error_reason, self.status_code())
+    }
+}
+
 pub async fn upload_weights_handler(
     http_request: HttpRequest,
     server_state: web::Data<Arc<ServerState>>,
-    mut multipart_payload: Multipart,
-  ) -> Result<HttpResponse, UploadError> {
-  
-    let response = handle_weight_upload(
-      &http_request,
-      &server_state,
-      multipart_payload,
-      &ALLOWED_MIME_TYPES).await?;
-    
-    // send a body with the required payload along side the multipart upload 
-  
-    let weight_upload_token = response.to_model_token();
-  
-    let response: UploadWeightsSuccessResponse = UploadWeightsSuccessResponse {
-      success: true,
-      upload_token: weight_upload_token,
-    };
-  
-    let body = serde_json::to_string(&response)
-        .map_err(|e| UploadError::ServerError)?;
-  
-    return Ok(HttpResponse::Ok()
-        .content_type("application/json")
-        .body(body));
-  }
-  
-  pub enum SuccessCase {
-    WeightAlreadyUploaded {
-      existing_upload_token: ModelWeightToken,
-    },
-    WeightSuccessfullyUploaded {
-      upload_token: ModelWeightToken,
+    mut multipart_payload: Multipart
+) -> Result<HttpResponse, UploadError> {
+    // Rate limiter check
+    if
+        let Err(_err) = server_state.redis_rate_limiters.model_upload.rate_limit_request(
+            &http_request
+        )
+    {
+        return Err(UploadModelWeightsError::RateLimited);
     }
-  }
-  
-  impl SuccessCase {
-    pub fn to_model_token(self) -> ModelWeightToken {
-      match self {
-        SuccessCase::WeightAlreadyUploaded { existing_upload_token } => existing_upload_token,
-        SuccessCase::WeightSuccessfullyUploaded { upload_token } => upload_token,
-      }
-    }
-  }
-  
-  pub async fn handle_weight_upload(
-    http_request: &HttpRequest,
-    server_state: &web::Data<Arc<ServerState>>,
-    mut multipart_payload: Multipart,
-    allowed_mimetypes: &HashSet<&'static str>,
-  ) -> Result<SuccessCase, UploadError> {
-  
-    let mut mysql_connection = server_state.mysql_pool
-        .acquire()
-        .await
-        .map_err(|err| {
-          error!("MySql pool error: {:?}", err);
-          UploadError::ServerError
-        })?;
-  
-    // ==================== READ SESSION ==================== //
-  
-    let maybe_user_session = server_state
-        .session_checker
-        .maybe_get_user_session_from_connection(http_request, &mut mysql_connection)
-        .await
+
+    let maybe_user_session = server_state.session_checker
+        .maybe_get_user_session(&http_request, &server_state.mysql_pool).await
         .map_err(|e| {
-          error!("Session checker error: {:?}", e);
-          UploadError::ServerError
+            warn!("Session checker error: {:?}", e);
+            UploadModelWeightsError::ServerError
         })?;
-  
-    // ==================== RATE LIMIT ==================== //
-  
-    let rate_limiter = match maybe_user_session {
-      None => &server_state.redis_rate_limiters.logged_out,
-      Some(ref user) => {
-        if user.is_banned {
-          return Err(UploadError::NotAuthorized);
+
+    let user_session = match maybe_user_session {
+        Some(session) => session,
+        None => {
+            warn!("not logged in");
+            return Err(UploadModelWeightsError::MustBeLoggedIn);
         }
-        &server_state.redis_rate_limiters.logged_in
-      },
     };
-  
-    if let Err(_err) = rate_limiter.rate_limit_request(&http_request) {
-      return Err(UploadError::RateLimited);
+
+    if let Err(reason) = validate_idempotency_token(&request.idempotency_token) {
+        return Err(UploadModelWeightsError::BadInput(reason));
     }
-  
-    if let Err(_err) = server_state.redis_rate_limiters.model_upload.rate_limit_request(&http_request) {
-      return Err(UploadError::RateLimited);
+
+    if let Err(reason) = validate_model_title(&request.title) {
+        return Err(UploadModelWeightsError::BadInput(reason));
     }
-  
-    // ==================== READ MULTIPART REQUEST ==================== //
-  
-    let upload_weights_request = drain_multipart_request(multipart_payload)
-        .await
-        .map_err(|e| {
-          // TODO: Error handling could be nicer.
-          UploadError::BadInput("bad request".to_string())
-        })?;
-  
-    let uuid_idempotency_token = upload_weights_request.uuid_idempotency_token
-        .ok_or(UploadError::BadInput("no uuid".to_string()))?;
-  
-    let maybe_existing_upload =
-        get_media_upload_by_uuid_with_connection(&uuid_idempotency_token, &mut mysql_connection)
-            .await;
-  
-    match maybe_existing_upload {
-      Err(err) => {
-        error!("Error checking for previous upload: {:?}", err);
-        return Err(UploadError::ServerError);
-      }
-      Ok(Some(upload)) => {
-        // File already uploaded and frontend didn't protect us.
-        return Ok(SuccessCase::WeightAlreadyUploaded {
-          existing_upload_token: upload.token,
-        });
-      }
-      Ok(None) => {
-        // Proceed.
-      }
+
+    match is_bad_model_weights_download_url(&download_url) {
+        Ok(false) => {} // Ok case
+        Ok(true) => {
+            return Err(UploadModelWeightsError::BadInput("Bad model download URL".to_string()));
+        }
+        Err(err) => {
+            warn!("Error parsing url: {:?}", err);
+            return Err(UploadModelWeightsError::BadInput("Bad model download URL".to_string()));
+        }
     }
-  
-    if let Err(reason) = validate_idempotency_token_format(&uuid_idempotency_token) {
-      return Err(UploadError::BadInput(reason));
-    }
-  
-    let creator_set_visibility = maybe_user_session
-        .as_ref()
-        .map(|user_session| user_session.preferred_tts_result_visibility) // TODO: We need a new type of visibility control.
-        .unwrap_or(Visibility::default());
-  
-    let ip_address = get_request_ip(&http_request);
-  
-    let maybe_user_token = maybe_user_session
-        .map(|session| session.get_strongly_typed_user_token());
-  
-    let maybe_file_size_bytes = upload_weights_request.file_bytes
-        .as_ref()
-        .map(|bytes| bytes.len());
-  
-    info!("Upload maybe filesize: {:?}", maybe_file_size_bytes);
-  
-    let maybe_mimetype = upload_weights_request.file_bytes
-        .as_ref()
-        .map(|bytes| get_mimetype_for_bytes(bytes.as_ref()))
-        .flatten();
-      
-  
-    let bytes = match upload_weights_request.file_bytes {
-      None => return Err(UploadError::BadInput("missing file contents".to_string())),
-      Some(bytes) => bytes,
+
+    let supported_model_weights_type = request.model_weights_type;
+
+    let job_token;
+    let download_job_type;
+
+    let (download_job_token, _record_id) = insert_generic_download_job(
+        InsertGenericDownloadJobArgs {
+            uuid_idempotency_token: &uuid,
+            download_type: GenericDownloadType::Vits,
+            download_url: &download_url,
+            title: &title,
+            creator_user_token: &user_session.user_token,
+            creator_ip_address: &ip_address,
+            creator_set_visibility: visibility, // TODO: Creator set preference
+            mysql_pool: &server_state.mysql_pool,
+        }
+    ).await.map_err(|err| {
+        warn!("New generic download creation DB error (for generic model): {:?}", err);
+        UploadModelWeightsError::ServerError
+    })?;
+
+    job_token = download_job_token.to_string();
+    download_job_type = DownloadJobType::Generic;
+
+    let response = UploadModelWeightsSuccessResponse {
+        success: true,
+        job_token: job_token.to_string(),
+        job_type: download_job_type,
     };
-  
-    let hash = sha256_hash_bytes(&bytes)
-        .map_err(|io_error| {
-          error!("Problem hashing bytes: {:?}", io_error);
-          UploadError::ServerError
-        })?;
-  
-    let file_size_bytes = bytes.len();
-  
-    let mut maybe_duration_millis = None;
-    let mut maybe_codec_name = None;
-    let mut media_upload_type = None;
-  
-    if let Some(mimetype) = maybe_mimetype.as_deref() {
-  
-      if !allowed_mimetypes.contains(mimetype) {
-        // NB: Don't let our error message inject malicious strings
-        let filtered_mimetype = mimetype
-            .chars()
-            .filter(|c| c.is_ascii())
-            .filter(|c| c.is_alphanumeric() || *c == '/')
-            .collect::<String>();
-        return Err(UploadError::BadInput(format!("unpermitted mime type: {}", &filtered_mimetype)));
-      }
-      media_upload_type = match mimetype {
-        "binary/octet-stream" /* .bin */ => Some(MediaUploadType::Binary),
-        _ => None,
-      };
+
+    let body = serde_json::to_string(&response).map_err(|_e| UploadModelWeightsError::ServerError)?;
+
+    Ok(HttpResponse::Ok().content_type("application/json").body(body))
+}
+
+fn validate_idempotency_token(token: &str) -> Result<(), String> {
+    if token.len() != 36 {
+        return Err("idempotency token should be 36 characters".to_string());
     }
-  
-    // This is fine for now we just don't expose the path!
-    let public_upload_path = WeightUploadOriginalFilePath::generate_new();
-  
-    info!("Uploading media to bucket path: {}", public_upload_path.get_full_object_path_str());
-  
-    // need to upload based on visibility 
-    server_state.public_bucket_client.upload_file_with_content_type(
-      public_upload_path.get_full_object_path_str(),
-      bytes.as_ref(),
-      mime_type)
-        .await
-        .map_err(|e| {
-          warn!("Upload media bytes to bucket error: {:?}", e);
-          UploadError::ServerError
-        })?;
-  
-    let token = create_weight(CreateModelWeightsArgs {
-      token: &ModelWeightToken::generate(),
-      weights_type: WeightsType::LoRA, // to do parse multi part for this
-      weights_category: WeightsCategory::ImageGeneration,
-      title: upload_weights_request.file_name.as_deref(),
-      maybe_thumbnail_token: &ThumbnailToken::generate(),
-      description_markdown: "",
-      description_rendered_html: "",
-      creator_user_token: maybe_user_token.as_str(),
-      creator_ip_address: &ip_address,
-      creator_set_visibility,
-      maybe_last_update_user_token: "",
-      original_download_url: "",
-      original_filename: upload_weights_request.file_name.as_deref(),
-      file_size_bytes: file_size_bytes as i32,
-      file_checksum_sha2: hash,
-      private_bucket_hash: private_bucket_hash,
-      maybe_private_bucket_prefix: None,
-      maybe_private_bucket_extension: None,
-      cached_user_ratings_total_count: 0,
-      cached_user_ratings_positive_count: 0,
-      cached_user_ratings_negative_count: 0,
-      maybe_cached_user_ratings_ratio: None,
-      version: 0,
-      mysql_pool: &server_state.mysql_pool,
-    }).await
-        .map_err(|e| {
-          error!("Error creating weights upload: {:?}", e);
-          UploadError::ServerError
-        })?;
-  
-    info!("new weights upload id: {} token: {:?}", record_id, &token);
-  
-    Ok(SuccessCase::WeightSuccessfullyUploaded {
-      upload_token: token,
-    })
-  }
-  
+
+    Ok(())
+}