(WIP) snowflake endpoint POC

.github/workflows/build.yaml

@@ -5,6 +5,7 @@ on:
   push:
     branches:
       - 'main'
+      - 'snowflake-endpoint'
     tags:
       - 'v*'
 

router/src/lib.rs

@@ -14,9 +14,9 @@ use infer::Infer;
 use loader::AdapterLoader;
 use queue::Entry;
 use serde::{Deserialize, Serialize};
-use serde_json::json;
 use utoipa::ToSchema;
 use validation::Validation;
+use std::collections::HashMap;
 
 /// Hub type
 #[derive(Clone, Debug, Deserialize)]
@@ -301,6 +301,19 @@ pub(crate) struct GenerateRequest {
     pub parameters: GenerateParameters,
 }
 
+#[derive(Clone, Debug, Deserialize, ToSchema)]
+pub(crate) struct SnowflakeGenerateRequest {
+    #[schema(example = "data: [[row_index, value]]")]
+    pub data: Vec<(i32, String, GenerateParameters)>,
+    // 'data: [[0, "'{inputs: "blah"}'"], [1, "next row data" ]]
+}
+
+#[derive(Serialize, ToSchema, Clone, Debug, Deserialize)]
+pub(crate) struct SnowflakeGenerateResponse {
+    #[schema(example = "data: [[row_index, value]]")]
+    pub data: Vec<(i32, GenerateResponse)>,
+    // 'data: [[0, "'{inputs: "blah"}'"], [1, "next row data" ]]
+}
 #[derive(Clone, Debug, Deserialize, ToSchema)]
 pub(crate) struct CompatGenerateRequest {
     #[schema(example = "My name is Olivier and I")]
@@ -321,7 +334,28 @@ impl From<CompatGenerateRequest> for GenerateRequest {
     }
 }
 
-#[derive(Debug, Serialize, ToSchema)]
+impl From<SnowflakeGenerateRequest> for GenerateRequest {
+    fn from(req: SnowflakeGenerateRequest) -> Self {
+        let input_row = req.data.into_iter().nth(0).unwrap();
+        Self {
+            inputs: input_row.1,
+            parameters: input_row.2
+        }
+    }
+}
+
+impl From<GenerateResponse> for SnowflakeGenerateResponse {
+    fn from(req: GenerateResponse) -> Self {
+        let mut vec: Vec<(i32, GenerateResponse)> = Vec::new();
+        let data_tuple: (i32, GenerateResponse) = (0, req);
+        vec.push(data_tuple);
+        Self {
+            data: vec,
+        }
+    }
+}
+
+#[derive(Debug, Serialize, ToSchema, Clone, Deserialize)]
 pub struct PrefillToken {
     #[schema(example = 0)]
     id: u32,
@@ -331,7 +365,7 @@ pub struct PrefillToken {
     logprob: f32,
 }
 
-#[derive(Debug, Serialize, ToSchema)]
+#[derive(Debug, Serialize, ToSchema, Clone, Deserialize)]
 pub struct AlternativeToken {
     #[schema(example = 0)]
     id: u32,
@@ -341,7 +375,7 @@ pub struct AlternativeToken {
     logprob: f32,
 }
 
-#[derive(Debug, Serialize, ToSchema)]
+#[derive(Debug, Serialize, ToSchema, Clone, Deserialize)]
 pub struct Token {
     #[schema(example = 0)]
     id: u32,
@@ -356,7 +390,7 @@ pub struct Token {
     alternative_tokens: Option<Vec<AlternativeToken>>,
 }
 
-#[derive(Serialize, ToSchema)]
+#[derive(Serialize, ToSchema, Debug, Clone, Deserialize)]
 #[serde(rename_all(serialize = "snake_case"))]
 pub(crate) enum FinishReason {
     #[schema(rename = "length")]
@@ -368,7 +402,7 @@ pub(crate) enum FinishReason {
     StopSequence,
 }
 
-#[derive(Serialize, ToSchema)]
+#[derive(Serialize, ToSchema, Debug, Clone, Deserialize)]
 pub(crate) struct BestOfSequence {
     #[schema(example = "test")]
     pub generated_text: String,
@@ -382,7 +416,7 @@ pub(crate) struct BestOfSequence {
     pub tokens: Vec<Token>,
 }
 
-#[derive(Serialize, ToSchema)]
+#[derive(Serialize, ToSchema, Clone, Debug, Deserialize)]
 pub(crate) struct Details {
     #[schema(example = "length")]
     pub finish_reason: FinishReason,
@@ -398,7 +432,7 @@ pub(crate) struct Details {
     pub best_of_sequences: Option<Vec<BestOfSequence>>,
 }
 
-#[derive(Serialize, ToSchema)]
+#[derive(Serialize, ToSchema, Clone, Debug, Deserialize)]
 pub(crate) struct GenerateResponse {
     #[schema(example = "test")]
     pub generated_text: String,

router/src/server.rs

@@ -4,10 +4,7 @@ use crate::health::Health;
 use crate::infer::{InferError, InferResponse, InferStreamResponse};
 use crate::validation::ValidationError;
 use crate::{
-    BestOfSequence, ChatCompletionRequest, ChatCompletionResponse, ChatCompletionStreamResponse,
-    CompatGenerateRequest, CompletionRequest, CompletionResponse, CompletionStreamResponse,
-    Details, ErrorResponse, FinishReason, GenerateParameters, GenerateRequest, GenerateResponse,
-    HubModelInfo, Infer, Info, PrefillToken, StreamDetails, StreamResponse, Token, Validation,
+    BestOfSequence, ChatCompletionRequest, ChatCompletionResponse, ChatCompletionStreamResponse, CompatGenerateRequest, CompletionRequest, CompletionResponse, CompletionStreamResponse, Details, ErrorResponse, FinishReason, GenerateParameters, GenerateRequest, GenerateResponse, HubModelInfo, Infer, Info, PrefillToken, SnowflakeGenerateRequest, SnowflakeGenerateResponse, StreamDetails, StreamResponse, Token, Validation
 };
 use axum::extract::Extension;
 use axum::http::{HeaderMap, Method, StatusCode};
@@ -85,6 +82,40 @@ async fn compat_generate(
         Ok((headers, Json(vec![generation.0])).into_response())
     }
 }
+/// Snowflake compatible generation endpoint
+#[utoipa::path(
+post,
+tag = "LoRAX",
+path = "/snowflake/generate",
+request_body = SnowflakeGenerateRequest,
+responses(
+(status = 200, description = "Generated Text", body = SnowflakeGenerateResponse),
+(status = 424, description = "Generation Error", body = ErrorResponse,
+example = json ! ({"error": "Request failed during generation"})),
+(status = 429, description = "Model is overloaded", body = ErrorResponse,
+example = json ! ({"error": "Model is overloaded"})),
+(status = 422, description = "Input validation error", body = ErrorResponse,
+example = json ! ({"error": "Input validation error"})),
+(status = 500, description = "Incomplete generation", body = ErrorResponse,
+example = json ! ({"error": "Incomplete generation"})),
+)
+)]
+#[instrument(skip(infer, req))]
+async fn snowflake_generate(
+    default_return_full_text: Extension<bool>,
+    infer: Extension<Infer>,
+    req_headers: HeaderMap,
+    req: Json<SnowflakeGenerateRequest>,
+) -> Result<(HeaderMap, Json<SnowflakeGenerateResponse>), (StatusCode, Json<ErrorResponse>)> {
+    let req = req.0;
+    let gen_req = GenerateRequest::from(req);
+    let (headers, generation) = generate(infer, req_headers, Json(gen_req.into())).await?;
+    let details = generation.details.clone();
+    let generated_text = generation.generated_text.clone();
+    let response = SnowflakeGenerateResponse::from(GenerateResponse{generated_text, details});
+    Ok((headers, Json(response)))
+}
+
 
 /// OpenAI compatible completions endpoint
 #[utoipa::path(
@@ -945,6 +976,7 @@ pub async fn run(
         .route("/generate", post(generate))
         .route("/generate_stream", post(generate_stream))
         .route("/v1/completions", post(completions_v1))
+        .route("/snowflake/generate", post(snowflake_generate))
         .route("/v1/chat/completions", post(chat_completions_v1))
         // AWS Sagemaker route
         .route("/invocations", post(compat_generate))