Create container image and move parakeet fork to vendor dir

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(&parakeet_dir, &parakeet_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(