diff options
Diffstat (limited to 'makima')
| -rw-r--r-- | makima/Cargo.toml | 2 | ||||
| -rw-r--r-- | makima/Dockerfile | 58 | ||||
| -rw-r--r-- | makima/sh/download-models.sh | 60 | ||||
| -rw-r--r-- | makima/src/bin/server.rs | 32 | ||||
| -rw-r--r-- | makima/src/server/mod.rs | 23 |
5 files changed, 165 insertions, 10 deletions
diff --git a/makima/Cargo.toml b/makima/Cargo.toml index 4d8076e..3368a6e 100644 --- a/makima/Cargo.toml +++ b/makima/Cargo.toml @@ -10,7 +10,7 @@ path = "src/bin/server.rs" [dependencies] # ML/Audio (existing) # Local fork with streaming API exposed (diarize_streaming method) -parakeet-rs = { path = "../parakeet-rs", features = ["sortformer"] } +parakeet-rs = { path = "../vendor/parakeet-rs", features = ["sortformer"] } symphonia = { version = "0.5", features = ["mp3", "aac", "flac", "ogg", "vorbis", "wav", "pcm"] } ort = "2.0.0-rc.10" tokenizers = "0.21" diff --git a/makima/Dockerfile b/makima/Dockerfile new file mode 100644 index 0000000..c8e7a2b --- /dev/null +++ b/makima/Dockerfile @@ -0,0 +1,58 @@ +# Build stage +FROM rust:1.83-bookworm AS builder + +WORKDIR /app + +# Install build dependencies +RUN apt-get update && apt-get install -y \ + pkg-config \ + libssl-dev \ + && rm -rf /var/lib/apt/lists/* + +# Copy workspace files +COPY Cargo.toml Cargo.lock ./ +COPY makima ./makima +COPY vendor ./vendor + +# Build release binary +RUN cargo build --release --package makima --bin makima-server + +# Runtime stage +FROM debian:bookworm-slim + +WORKDIR /app + +# Install runtime dependencies +RUN apt-get update && apt-get install -y \ + ca-certificates \ + libssl3 \ + curl \ + && rm -rf /var/lib/apt/lists/* + +# Copy the binary +COPY --from=builder /app/target/release/makima-server /app/makima-server + +# Copy model download script +COPY makima/sh/download-models.sh /app/download-models.sh +RUN chmod +x /app/download-models.sh + +# Create models directory +RUN mkdir -p /app/models + +# Set default environment variables +ENV PORT=8080 +ENV RUST_LOG=makima=info,tower_http=info +ENV MODELS_DIR=/app/models +ENV PARAKEET_MODEL_DIR=/app/models/parakeet-tdt-0.6b-v3 +ENV PARAKEET_EOU_DIR=/app/models/realtime_eou_120m-v1-onnx +ENV SORTFORMER_MODEL_PATH=/app/models/diarization/diar_streaming_sortformer_4spk-v2.onnx + +EXPOSE 8080 + +# Health check +HEALTHCHECK --interval=30s --timeout=10s --start-period=120s --retries=3 \ + CMD curl -f http://localhost:${PORT}/api/v1/healthcheck || exit 1 + +# Use download script as entrypoint to ensure models exist +ENTRYPOINT ["/app/download-models.sh"] +CMD ["/app/makima-server"] diff --git a/makima/sh/download-models.sh b/makima/sh/download-models.sh new file mode 100644 index 0000000..ddb7454 --- /dev/null +++ b/makima/sh/download-models.sh @@ -0,0 +1,60 @@ +#!/bin/bash +set -e + +MODELS_DIR="${MODELS_DIR:-/app/models}" +MODEL_BASE_URL="${MODEL_BASE_URL:-}" + +# Model directories to check/download +PARAKEET_DIR="parakeet-tdt-0.6b-v3" +EOU_DIR="realtime_eou_120m-v1-onnx" +DIARIZATION_DIR="diarization" + +download_model() { + local name=$1 + local url=$2 + local dest="$MODELS_DIR/$name" + + if [ -d "$dest" ] && [ "$(ls -A $dest 2>/dev/null)" ]; then + echo "Model $name already exists, skipping..." + return 0 + fi + + echo "Downloading $name..." + mkdir -p "$dest" + + # Download and extract tar.gz + curl -L "$url" | tar -xz -C "$dest" --strip-components=1 + + echo "Downloaded $name successfully" +} + +# Check if models exist +check_models_exist() { + [ -d "$MODELS_DIR/$PARAKEET_DIR" ] && \ + [ -d "$MODELS_DIR/$EOU_DIR" ] && \ + [ -f "$MODELS_DIR/$DIARIZATION_DIR/diar_streaming_sortformer_4spk-v2.onnx" ] +} + +if check_models_exist; then + echo "All models present" +else + if [ -z "$MODEL_BASE_URL" ]; then + echo "ERROR: Models not found and MODEL_BASE_URL not set" + echo "Please set MODEL_BASE_URL to the base URL containing model archives:" + echo " - \${MODEL_BASE_URL}/parakeet-tdt-0.6b-v3.tar.gz" + echo " - \${MODEL_BASE_URL}/realtime_eou_120m-v1-onnx.tar.gz" + echo " - \${MODEL_BASE_URL}/diarization.tar.gz" + exit 1 + fi + + mkdir -p "$MODELS_DIR" + + download_model "$PARAKEET_DIR" "${MODEL_BASE_URL}/parakeet-tdt-0.6b-v3.tar.gz" + download_model "$EOU_DIR" "${MODEL_BASE_URL}/realtime_eou_120m-v1-onnx.tar.gz" + download_model "$DIARIZATION_DIR" "${MODEL_BASE_URL}/diarization.tar.gz" + + echo "All models downloaded successfully" +fi + +# Execute the main command +exec "$@" diff --git a/makima/src/bin/server.rs b/makima/src/bin/server.rs index 06b6585..470e295 100644 --- a/makima/src/bin/server.rs +++ b/makima/src/bin/server.rs @@ -8,10 +8,10 @@ use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt}; use makima::server::{run_server, state::AppState}; -/// Default model paths relative to the working directory. -const PARAKEET_MODEL_DIR: &str = "models/parakeet-tdt-0.6b-v3"; -const PARAKEET_EOU_DIR: &str = "models/realtime_eou_120m-v1-onnx"; -const SORTFORMER_MODEL_PATH: &str = "models/diarization/diar_streaming_sortformer_4spk-v2.onnx"; +/// Default model paths (can be overridden via environment variables). +const DEFAULT_PARAKEET_MODEL_DIR: &str = "models/parakeet-tdt-0.6b-v3"; +const DEFAULT_PARAKEET_EOU_DIR: &str = "models/realtime_eou_120m-v1-onnx"; +const DEFAULT_SORTFORMER_MODEL_PATH: &str = "models/diarization/diar_streaming_sortformer_4spk-v2.onnx"; #[tokio::main] async fn main() -> anyhow::Result<()> { @@ -19,22 +19,38 @@ async fn main() -> anyhow::Result<()> { tracing_subscriber::registry() .with( tracing_subscriber::EnvFilter::try_from_default_env() - .unwrap_or_else(|_| "makima=debug,tower_http=debug".into()), + .unwrap_or_else(|_| "makima=info,tower_http=info".into()), ) .with(tracing_subscriber::fmt::layer()) .init(); tracing::info!("Starting Makima Listening API Server"); - tracing::info!("Loading ML models..."); + + // Read configuration from environment + let port = std::env::var("PORT").unwrap_or_else(|_| "8080".to_string()); + let parakeet_dir = std::env::var("PARAKEET_MODEL_DIR") + .unwrap_or_else(|_| DEFAULT_PARAKEET_MODEL_DIR.to_string()); + let parakeet_eou_dir = std::env::var("PARAKEET_EOU_DIR") + .unwrap_or_else(|_| DEFAULT_PARAKEET_EOU_DIR.to_string()); + let sortformer_path = std::env::var("SORTFORMER_MODEL_PATH") + .unwrap_or_else(|_| DEFAULT_SORTFORMER_MODEL_PATH.to_string()); + + tracing::info!( + parakeet = %parakeet_dir, + eou = %parakeet_eou_dir, + sortformer = %sortformer_path, + "Loading ML models..." + ); // Load ML models let state = Arc::new( - AppState::new(PARAKEET_MODEL_DIR, PARAKEET_EOU_DIR, SORTFORMER_MODEL_PATH) + AppState::new(¶keet_dir, ¶keet_eou_dir, &sortformer_path) .map_err(|e| anyhow::anyhow!("Failed to load models: {}", e))?, ); tracing::info!("Models loaded successfully"); // Run the server - run_server(state, "0.0.0.0:8080").await + let addr = format!("0.0.0.0:{}", port); + run_server(state, &addr).await } diff --git a/makima/src/server/mod.rs b/makima/src/server/mod.rs index a6e0525..c509afa 100644 --- a/makima/src/server/mod.rs +++ b/makima/src/server/mod.rs @@ -6,9 +6,12 @@ pub mod openapi; pub mod state; use axum::{ + http::StatusCode, + response::IntoResponse, routing::get, - Router, + Json, Router, }; +use serde::Serialize; use tower_http::cors::{Any, CorsLayer}; use tower_http::trace::TraceLayer; use utoipa::OpenApi; @@ -18,6 +21,23 @@ use crate::server::handlers::listen; use crate::server::openapi::ApiDoc; use crate::server::state::SharedState; +#[derive(Serialize)] +struct HealthResponse { + status: &'static str, + version: &'static str, +} + +/// Health check endpoint for load balancers and orchestrators. +async fn health_check() -> impl IntoResponse { + ( + StatusCode::OK, + Json(HealthResponse { + status: "healthy", + version: env!("CARGO_PKG_VERSION"), + }), + ) +} + /// Create the axum Router with all routes configured. pub fn make_router(state: SharedState) -> Router { // API v1 routes @@ -29,6 +49,7 @@ pub fn make_router(state: SharedState) -> Router { .url("/api-docs/openapi.json", ApiDoc::openapi()); Router::new() + .route("/api/v1/healthcheck", get(health_check)) .nest("/api/v1", api_v1) .merge(swagger) .layer( |