From d4adff31aa8e6d05a4eb91b6acce9ae70ac49b6d Mon Sep 17 00:00:00 2001 From: Simo Lin Date: Mon, 26 Jan 2026 02:43:29 -0500 Subject: [PATCH] update wasm endpoint (#17748) --- sgl-model-gateway/Cargo.toml | 6 - sgl-model-gateway/src/wasm/mod.rs | 9 ++ sgl-model-gateway/src/wasm/route.rs | 228 ++++++++++++++++++++++++++++ 3 files changed, 237 insertions(+), 6 deletions(-) create mode 100644 sgl-model-gateway/src/wasm/mod.rs create mode 100644 sgl-model-gateway/src/wasm/route.rs diff --git a/sgl-model-gateway/Cargo.toml b/sgl-model-gateway/Cargo.toml index 9bc796579..a7a3947aa 100644 --- a/sgl-model-gateway/Cargo.toml +++ b/sgl-model-gateway/Cargo.toml @@ -69,11 +69,8 @@ k8s-openapi = { version = "0.25.0", features = ["v1_33"] } metrics = "0.24.2" metrics-exporter-prometheus = "0.17.0" uuid = { version = "1.10", features = ["v4", "serde"] } -ulid = "1.2.1" parking_lot = "0.12.4" thiserror = "2.0.12" - - regex = "1.10" memchr = "2.7" # SIMD-optimized byte pattern searching url = "2.5.4" @@ -134,9 +131,6 @@ deadpool-redis = "0.18.0" # wasm dependencies sha2 = "0.10" wasmtime = { version = "38.0", features = ["component-model", "async"] } -wasmtime-wasi = "38.0" -async-channel = "2.5" -aho-corasick = "1.1.4" [build-dependencies] tonic-prost-build = "0.14.2" diff --git a/sgl-model-gateway/src/wasm/mod.rs b/sgl-model-gateway/src/wasm/mod.rs new file mode 100644 index 000000000..5567222cf --- /dev/null +++ b/sgl-model-gateway/src/wasm/mod.rs @@ -0,0 +1,9 @@ +//! WebAssembly (WASM) module support for SGL Model Gateway +//! +//! This module re-exports the smg-wasm crate and provides HTTP API routes. + +// Re-export everything from smg-wasm +pub use smg_wasm::*; + +// Local HTTP API routes (depends on app-specific types) +pub mod route; diff --git a/sgl-model-gateway/src/wasm/route.rs b/sgl-model-gateway/src/wasm/route.rs new file mode 100644 index 000000000..5e4bbb5f8 --- /dev/null +++ b/sgl-model-gateway/src/wasm/route.rs @@ -0,0 +1,228 @@ +//! WASM HTTP API Routes +//! +//! Provides REST API endpoints for managing WASM modules: +//! - POST /wasm - Add modules +//! - DELETE /wasm/:uuid - Remove a module +//! - GET /wasm - List all modules with metrics + +use std::{sync::Arc, time::Duration}; + +use axum::{ + extract::{Json, Path, State}, + http::StatusCode, + response::{IntoResponse, Response}, +}; +use uuid::Uuid; + +use crate::{ + core::{job_queue::Job, steps::WasmModuleConfigRequest}, + server::AppState, + wasm::{ + WasmMetrics, WasmModuleAddRequest, WasmModuleAddResponse, WasmModuleAddResult, + WasmModuleListResponse, + }, +}; + +/// Wait for job completion by polling job status +/// Returns the job result message if successful +async fn wait_for_job_completion( + job_queue: &crate::core::job_queue::JobQueue, + status_key: &str, + timeout_duration: Duration, +) -> Result { + let start = std::time::Instant::now(); + let mut poll_interval = Duration::from_millis(100); + let max_poll_interval = Duration::from_millis(2000); + let poll_backoff = Duration::from_millis(200); + + loop { + if start.elapsed() > timeout_duration { + return Err(format!("Job timeout after {}s", timeout_duration.as_secs())); + } + + if let Some(job_status) = job_queue.get_status(status_key) { + match job_status.status.as_str() { + "pending" | "processing" => { + tokio::time::sleep(poll_interval).await; + poll_interval = (poll_interval + poll_backoff).min(max_poll_interval); + continue; + } + "failed" => { + let error_msg = job_status + .message + .unwrap_or_else(|| "Unknown error".to_string()); + job_queue.remove_status(status_key); + return Err(error_msg); + } + _ => { + // Should not happen, but handle gracefully + job_queue.remove_status(status_key); + return Err("Unexpected job status".to_string()); + } + } + } else { + // Job completed successfully (status was removed by record_job_completion) + // We need to get the result from the job execution + // Since job queue removes status on success, we can't get the result here + // We'll need to query the wasm manager to find the module by name + // For now, return a success message and let caller extract UUID from manager + return Ok("Job completed successfully".to_string()); + } + } +} + +pub async fn add_wasm_module( + State(state): State>, + Json(config): Json, +) -> Response { + let Some(_) = state.context.wasm_manager.as_ref() else { + return StatusCode::INTERNAL_SERVER_ERROR.into_response(); + }; + + let Some(job_queue) = state.context.worker_job_queue.get() else { + return StatusCode::INTERNAL_SERVER_ERROR.into_response(); + }; + + let mut status = StatusCode::OK; + let mut modules = config.modules.clone(); + + for module in modules.iter_mut() { + let wasm_config = WasmModuleConfigRequest { + descriptor: module.clone(), + }; + + let job = Job::AddWasmModule { + config: Box::new(wasm_config), + }; + + let worker_url = job.worker_url().to_string(); + + // Submit job to queue + match job_queue.submit(job).await { + Ok(_) => { + // Wait for job completion (timeout: 5 minutes) + let timeout = Duration::from_secs(300); + match wait_for_job_completion(job_queue, &worker_url, timeout).await { + Ok(_) => { + // Job completed successfully, but we need to get the UUID + // Since job queue removes status on success, we need to query + // the workflow engine or wasm manager to get the UUID + // For now, let's try to get it from the wasm manager by name + if let Some(wasm_manager) = state.context.wasm_manager.as_ref() { + if let Ok(all_modules) = wasm_manager.get_modules() { + if let Some(registered_module) = all_modules + .iter() + .find(|m| m.module_meta.name == module.name) + { + module.add_result = Some(WasmModuleAddResult::Success( + registered_module.module_uuid, + )); + } else { + module.add_result = Some(WasmModuleAddResult::Error( + "Module registered but UUID not found".to_string(), + )); + status = StatusCode::BAD_REQUEST; + } + } else { + module.add_result = Some(WasmModuleAddResult::Error( + "Failed to query registered modules".to_string(), + )); + status = StatusCode::BAD_REQUEST; + } + } else { + module.add_result = Some(WasmModuleAddResult::Error( + "WASM manager not available".to_string(), + )); + status = StatusCode::BAD_REQUEST; + } + } + Err(e) => { + module.add_result = Some(WasmModuleAddResult::Error(e)); + status = StatusCode::BAD_REQUEST; + } + } + } + Err(e) => { + module.add_result = Some(WasmModuleAddResult::Error(format!( + "Failed to submit job: {}", + e + ))); + status = StatusCode::BAD_REQUEST; + } + } + } + + let response = WasmModuleAddResponse { modules }; + (status, Json(response)).into_response() +} + +pub async fn remove_wasm_module( + State(state): State>, + Path(module_uuid_str): Path, +) -> Response { + let Ok(module_uuid) = Uuid::parse_str(&module_uuid_str) else { + return StatusCode::BAD_REQUEST.into_response(); + }; + + let Some(_) = state.context.wasm_manager.as_ref() else { + return StatusCode::INTERNAL_SERVER_ERROR.into_response(); + }; + + let Some(job_queue) = state.context.worker_job_queue.get() else { + return StatusCode::INTERNAL_SERVER_ERROR.into_response(); + }; + + use crate::core::steps::WasmModuleRemovalRequest; + + let removal_request = WasmModuleRemovalRequest::new(module_uuid); + + let job = Job::RemoveWasmModule { + request: Box::new(removal_request), + }; + + let worker_url = job.worker_url().to_string(); + + // Submit job to queue + match job_queue.submit(job).await { + Ok(_) => { + // Wait for job completion (timeout: 1 minute) + let timeout = Duration::from_secs(60); + match wait_for_job_completion(job_queue, &worker_url, timeout).await { + Ok(_) => (StatusCode::OK, "Module removed successfully").into_response(), + Err(e) => (StatusCode::BAD_REQUEST, e).into_response(), + } + } + Err(e) => ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("Failed to submit job: {}", e), + ) + .into_response(), + } +} + +pub async fn list_wasm_modules(State(state): State>) -> Response { + let Some(wasm_manager) = state.context.wasm_manager.as_ref() else { + return StatusCode::INTERNAL_SERVER_ERROR.into_response(); + }; + let modules = wasm_manager.get_modules(); + if let Ok(modules) = modules { + let (total, success, failed, total_time_ms, max_time_ms) = wasm_manager.get_metrics(); + let average_execution_time_ms = if total > 0 { + Some(total_time_ms as f64 / total as f64) + } else { + None + }; + let metrics = WasmMetrics { + total_executions: total, + successful_executions: success, + failed_executions: failed, + total_execution_time_ms: total_time_ms, + max_execution_time_ms: max_time_ms, + average_execution_time_ms, + }; + let response = WasmModuleListResponse { modules, metrics }; + (StatusCode::OK, Json(response)).into_response() + } else { + StatusCode::INTERNAL_SERVER_ERROR.into_response() + } +}