From eac7ada1d44b11a6f3547dde67a80dfe6804d69c Mon Sep 17 00:00:00 2001 From: beatgeek Date: Fri, 23 Jan 2026 16:00:02 -0800 Subject: [PATCH 1/2] helion support --- Cargo.lock | 13 ++ Cargo.toml | 2 +- Dockerfile.helion | 32 ++++ Dockerfile.helion.mock | 25 +++ README.md | 63 +++++++ crates/kernelport-backend-helion/Cargo.toml | 12 ++ crates/kernelport-backend-helion/src/lib.rs | 169 +++++++++++++++++ .../kernelport-backend-helion/tests/smoke.rs | 20 ++ crates/kernelport-core/src/artifact.rs | 1 + crates/kernelport-server/Cargo.toml | 1 + crates/kernelport-server/src/cli.rs | 12 ++ crates/kernelport-server/src/main.rs | 30 ++- crates/kernelport-server/src/registry.rs | 27 +++ docker-compose.mock.yml | 21 +++ docker-compose.yml | 32 ++++ docs/deploy/helion-sidecar.yaml | 50 +++++ docs/validation.md | 71 +++++++ scripts/helion/README.md | 7 + scripts/helion/helion_worker.py | 174 ++++++++++++++++++ scripts/helion/mock/helion_worker_mock.py | 148 +++++++++++++++ 20 files changed, 907 insertions(+), 3 deletions(-) create mode 100644 Dockerfile.helion create mode 100644 Dockerfile.helion.mock create mode 100644 crates/kernelport-backend-helion/Cargo.toml create mode 100644 crates/kernelport-backend-helion/src/lib.rs create mode 100644 crates/kernelport-backend-helion/tests/smoke.rs create mode 100644 docker-compose.mock.yml create mode 100644 docker-compose.yml create mode 100644 docs/deploy/helion-sidecar.yaml create mode 100644 scripts/helion/README.md create mode 100644 scripts/helion/helion_worker.py create mode 100644 scripts/helion/mock/helion_worker_mock.py diff --git a/Cargo.lock b/Cargo.lock index 3861191..6488f30 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -627,6 +627,18 @@ version = "1.0.17" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "92ecc6618181def0457392ccd0ee51198e065e016d1d527a7ac1b6dc7c1f09d2" +[[package]] +name = "kernelport-backend-helion" +version = "0.1.0" +dependencies = [ + "anyhow", + "bytes", + "kernelport-core", + "kernelport-proto", + "tokio", + "tonic", +] + [[package]] name = "kernelport-backend-ort" version = "0.0.0" @@ -674,6 +686,7 @@ dependencies = [ "anyhow", "bytes", "clap", + "kernelport-backend-helion", "kernelport-backend-ort", "kernelport-core", "kernelport-proto", diff --git a/Cargo.toml b/Cargo.toml index 08ba4fd..a11ffff 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -7,7 +7,7 @@ members = [ "crates/kernelport-core", "crates/kernelport-runtime", "crates/kernelport-backend-ort", + "crates/kernelport-backend-helion", "crates/kernelport-proto", "crates/kernelport-server", ] - diff --git a/Dockerfile.helion b/Dockerfile.helion new file mode 100644 index 0000000..ca73b02 --- /dev/null +++ b/Dockerfile.helion @@ -0,0 +1,32 @@ +FROM nvidia/cuda:12.2.0-runtime-ubuntu22.04 + +RUN apt-get update \ + && apt-get install -y --no-install-recommends \ + ca-certificates \ + curl \ + python3 \ + python3-venv \ + && rm -rf /var/lib/apt/lists/* + +RUN curl -LsSf https://astral.sh/uv/install.sh | sh + +ENV VIRTUAL_ENV=/opt/venv +ENV PATH="/opt/venv/bin:/root/.cargo/bin:${PATH}" + +RUN uv venv /opt/venv \ + && uv pip install \ + "torch==2.9.*" \ + --index-url https://download.pytorch.org/whl/cu126 \ + && uv pip install \ + helion \ + grpcio \ + grpcio-tools \ + numpy + +WORKDIR /app +COPY scripts/helion /app/scripts/helion +COPY crates/kernelport-proto/src /app/crates/kernelport-proto/src + +EXPOSE 50061 + +CMD ["python", "/app/scripts/helion/helion_worker.py", "--addr", "0.0.0.0:50061", "--device", "cuda"] diff --git a/Dockerfile.helion.mock b/Dockerfile.helion.mock new file mode 100644 index 0000000..49fc2d4 --- /dev/null +++ b/Dockerfile.helion.mock @@ -0,0 +1,25 @@ +FROM python:3.11-slim + +RUN apt-get update \ + && apt-get install -y --no-install-recommends \ + ca-certificates \ + curl \ + && rm -rf /var/lib/apt/lists/* + +RUN curl -LsSf https://astral.sh/uv/install.sh | sh + +ENV PATH="/root/.local/bin:/root/.cargo/bin:${PATH}" + +WORKDIR /app +COPY scripts/helion /app/scripts/helion +COPY crates/kernelport-proto/src /app/crates/kernelport-proto/src + +RUN /root/.local/bin/uv venv /app/.venv \ + && . /app/.venv/bin/activate \ + && /root/.local/bin/uv pip install numpy grpcio grpcio-tools + +ENV PATH="/app/.venv/bin:${PATH}" + +EXPOSE 50061 + +CMD ["python", "/app/scripts/helion/mock/helion_worker_mock.py", "--addr", "0.0.0.0:50061"] diff --git a/README.md b/README.md index d95756c..a6acd26 100644 --- a/README.md +++ b/README.md @@ -53,6 +53,69 @@ This ensures you get: - Framework correctness - No silent performance regressions +### Helion (Experimental) +KernelPort can proxy Helion kernels via a Python sidecar. This keeps the Rust +server lean while letting MLEs author kernels in Helion (higher-level Triton with +autotuning). + +Sidecar flow (v0): +- Run the Helion gRPC worker in Python (see `scripts/helion/helion_worker.py`). +- Start kernelportd with `--backend helion --helion-addr http://127.0.0.1:50061`. +- Send standard KernelPort gRPC requests to `kernelportd`; it forwards to Helion. +- Expect a first-run autotune warm-up (can be minutes depending on kernel/search). + +Planned follow-up: +- Optional in-process Helion embedding (pyo3) for lower per-request latency. + +Helion worker Python deps (uv): +```bash +uv venv .venv +source .venv/bin/activate +uv pip install "torch==2.9.*" --index-url https://download.pytorch.org/whl/cu126 +uv pip install helion grpcio grpcio-tools numpy +``` + +Helion sidecar (Docker) build/run: +```bash +docker build -f Dockerfile.gpu -t kernelport:gpu . +docker build -f Dockerfile.helion -t kernelport-helion:gpu . +docker compose up --build +``` + +Bring it down: +```bash +docker compose down +``` + +GPU pinning and cache: +- Pin GPUs with `CUDA_VISIBLE_DEVICES` in each container (e.g. helion worker uses `0`, kernelport uses `0` or a different GPU). +- Persist Helion autotune artifacts by mounting a volume to the worker cache dir (e.g. set `XDG_CACHE_HOME=/cache` and mount `-v /path/to/cache:/cache`). + +Example request (replace base64 data as needed): +```bash +grpcurl -plaintext -d '{ + "model": "demo", + "inputs": [ + { "name": "x", "dtype": "F16", "shape": [4, 8], "data": "" } + ] +}' localhost:50051 kernelport.v1.InferenceService/Infer +``` + +Generate base64 payload: +```bash +python - <<'PY' +import base64 +import torch +x = torch.randn(4, 8, dtype=torch.float16) +print(base64.b64encode(x.numpy().tobytes()).decode()) +PY +``` + +CPU-only mock Helion (Docker, see `scripts/helion/mock/`): +```bash +docker compose -f docker-compose.mock.yml up --build +``` + --- ### Dynamic Batching diff --git a/crates/kernelport-backend-helion/Cargo.toml b/crates/kernelport-backend-helion/Cargo.toml new file mode 100644 index 0000000..ff4539b --- /dev/null +++ b/crates/kernelport-backend-helion/Cargo.toml @@ -0,0 +1,12 @@ +[package] +name = "kernelport-backend-helion" +version = "0.1.0" +edition = "2021" + +[dependencies] +anyhow = "1" +bytes = "1" +kernelport-core = { path = "../kernelport-core" } +kernelport-proto = { path = "../kernelport-proto" } +tokio = { version = "1", features = ["rt", "macros", "rt-multi-thread"] } +tonic = { version = "0.14", features = ["transport"] } diff --git a/crates/kernelport-backend-helion/src/lib.rs b/crates/kernelport-backend-helion/src/lib.rs new file mode 100644 index 0000000..4d17222 --- /dev/null +++ b/crates/kernelport-backend-helion/src/lib.rs @@ -0,0 +1,169 @@ +use std::time::Duration; + +use anyhow::{bail, Context, Result}; +use bytes::Bytes; +use kernelport_core::{ + Backend, BackendCapabilities, BackendModel, DType, Device, IOName, ModelArtifact, ModelSpec, + Shape, Tensor, TensorSpec, TensorStorage, +}; +use kernelport_proto::kernelport::v1 as pb; +use kernelport_proto::kernelport::v1::inference_service_client::InferenceServiceClient; +use tonic::transport::{Channel, Endpoint}; + +pub struct HelionBackend; + +impl HelionBackend { + pub fn new() -> Self { + Self + } +} + +impl Default for HelionBackend { + fn default() -> Self { + Self::new() + } +} + +pub struct HelionModel { + spec: ModelSpec, + client: InferenceServiceClient, + model: String, + timeout: Duration, +} + +impl Backend for HelionBackend { + type Model = HelionModel; + + fn name(&self) -> &'static str { + "helion" + } + + fn load(&self, artifact: &ModelArtifact, _device: Device) -> Result { + let ModelArtifact::HelionGrpc { addr, model } = artifact else { + bail!("helion backend expects a HelionGrpc artifact"); + }; + + let endpoint = Endpoint::from_shared(addr.clone()) + .context("invalid helion gRPC address")? + .connect_lazy(); + let client = InferenceServiceClient::new(endpoint); + + // v0: assume softmax-like 2D f16 tensors with x->y names and dynamic dims. + let spec = ModelSpec { + inputs: vec![TensorSpec { + name: IOName("x".to_string()), + dtype: DType::F16, + rank: 2, + dims: vec![None, None], + }], + outputs: vec![TensorSpec { + name: IOName("y".to_string()), + dtype: DType::F16, + rank: 2, + dims: vec![None, None], + }], + max_batch: 1, + }; + + Ok(HelionModel { + spec, + client, + model: model.clone(), + timeout: Duration::from_secs(120), + }) + } + + fn capabilities(&self) -> BackendCapabilities { + BackendCapabilities { + supports_dynamic_shapes: true, + prefers_nchw: false, + allows_cuda_graphs: false, + } + } +} + +impl BackendModel for HelionModel { + fn spec(&self) -> &ModelSpec { + &self.spec + } + + fn infer(&mut self, inputs: Vec) -> Result> { + let mut pb_inputs = Vec::with_capacity(inputs.len()); + for (idx, input) in inputs.into_iter().enumerate() { + pb_inputs.push(tensor_to_pb(&format!("input{idx}"), input)?); + } + + let mut request = tonic::Request::new(pb::InferRequest { + model: self.model.clone(), + inputs: pb_inputs, + }); + request.set_timeout(self.timeout); + + let mut client = self.client.clone(); + let response: tonic::Response = tokio::task::block_in_place(|| { + let handle = tokio::runtime::Handle::current(); + handle.block_on(async { client.infer(request).await }) + }) + .context("helion inference request failed")?; + let response = response.into_inner(); + + let mut outputs = Vec::with_capacity(response.outputs.len()); + for output in response.outputs { + outputs.push(pb_to_tensor(output)?); + } + + Ok(outputs) + } +} + +fn tensor_to_pb(name: &str, tensor: Tensor) -> Result { + let data = match tensor.storage { + TensorStorage::CpuBytes(bytes) => bytes, + TensorStorage::CpuPinned(p) => p.bytes, + TensorStorage::CudaDevice(_) => bail!("helion backend only supports CPU tensors"), + }; + + Ok(pb::Tensor { + name: name.to_string(), + dtype: to_proto_dtype(tensor.desc.dtype) as i32, + shape: tensor.desc.shape.0.iter().map(|d| *d as i64).collect(), + data: data.to_vec(), + }) +} + +fn pb_to_tensor(tensor: pb::Tensor) -> Result { + let dtype = parse_dtype(tensor.dtype)?; + let shape: Vec = tensor + .shape + .into_iter() + .map(|d| usize::try_from(d).unwrap_or(0)) + .collect(); + + Ok(Tensor::from_cpu_bytes( + dtype, + Shape::from_slice(&shape), + Bytes::from(tensor.data), + )) +} + +fn parse_dtype(raw: i32) -> Result { + let dtype = pb::DType::try_from(raw).context("unknown dtype enum value")?; + Ok(match dtype { + pb::DType::F32 => DType::F32, + pb::DType::F16 => DType::F16, + pb::DType::I64 => DType::I64, + pb::DType::I32 => DType::I32, + pb::DType::U8 => DType::U8, + pb::DType::DtypeUnspecified => bail!("dtype is unspecified"), + }) +} + +fn to_proto_dtype(dtype: DType) -> pb::DType { + match dtype { + DType::F32 => pb::DType::F32, + DType::F16 => pb::DType::F16, + DType::I64 => pb::DType::I64, + DType::I32 => pb::DType::I32, + DType::U8 => pb::DType::U8, + } +} diff --git a/crates/kernelport-backend-helion/tests/smoke.rs b/crates/kernelport-backend-helion/tests/smoke.rs new file mode 100644 index 0000000..42616f2 --- /dev/null +++ b/crates/kernelport-backend-helion/tests/smoke.rs @@ -0,0 +1,20 @@ +use kernelport_backend_helion::HelionBackend; +use kernelport_core::{Backend, BackendModel, Device, ModelArtifact}; + +#[tokio::test(flavor = "multi_thread")] +async fn loads_helion_spec() { + let backend = HelionBackend::new(); + let artifact = ModelArtifact::HelionGrpc { + addr: "http://127.0.0.1:50061".to_string(), + model: "softmax_two_pass".to_string(), + }; + + let model = backend + .load(&artifact, Device::Cuda { device_id: 0 }) + .expect("load helion model"); + + assert_eq!(model.spec().inputs.len(), 1); + assert_eq!(model.spec().outputs.len(), 1); + assert_eq!(model.spec().inputs[0].name.0, "x"); + assert_eq!(model.spec().outputs[0].name.0, "y"); +} diff --git a/crates/kernelport-core/src/artifact.rs b/crates/kernelport-core/src/artifact.rs index bfdad1b..d7d1b94 100644 --- a/crates/kernelport-core/src/artifact.rs +++ b/crates/kernelport-core/src/artifact.rs @@ -4,4 +4,5 @@ pub enum ModelArtifact { TensorRtEnginePath(std::path::PathBuf), TorchScriptPath(std::path::PathBuf), TfSavedModelDir(std::path::PathBuf), + HelionGrpc { addr: String, model: String }, } diff --git a/crates/kernelport-server/Cargo.toml b/crates/kernelport-server/Cargo.toml index 383cb90..64fc3ee 100644 --- a/crates/kernelport-server/Cargo.toml +++ b/crates/kernelport-server/Cargo.toml @@ -21,6 +21,7 @@ kernelport-core = { path = "../kernelport-core" } kernelport-runtime = { path = "../kernelport-runtime" } kernelport-proto = { path = "../kernelport-proto" } kernelport-backend-ort = { path = "../kernelport-backend-ort" } +kernelport-backend-helion = { path = "../kernelport-backend-helion" } [features] ort-cuda = ["kernelport-backend-ort/cuda"] diff --git a/crates/kernelport-server/src/cli.rs b/crates/kernelport-server/src/cli.rs index 061264f..49b05f8 100644 --- a/crates/kernelport-server/src/cli.rs +++ b/crates/kernelport-server/src/cli.rs @@ -23,8 +23,20 @@ pub enum Command { #[arg(long, default_value = "cpu")] device: String, + /// Backend type (onnx or helion) + #[arg(long, default_value = "onnx")] + backend: String, + /// Path to ONNX model file #[arg(long, default_value = "models/demo.onnx")] model_path: String, + + /// Helion sidecar address (gRPC) + #[arg(long, default_value = "http://127.0.0.1:50061")] + helion_addr: String, + + /// Helion kernel entrypoint name + #[arg(long, default_value = "softmax_two_pass")] + helion_model: String, }, } diff --git a/crates/kernelport-server/src/main.rs b/crates/kernelport-server/src/main.rs index 8941b03..5df533d 100644 --- a/crates/kernelport-server/src/main.rs +++ b/crates/kernelport-server/src/main.rs @@ -22,10 +22,22 @@ async fn main() -> Result<()> { grpc_addr, log, device, + backend, model_path, + helion_addr, + helion_model, } => { let device = parse_device(&device)?; - serve(grpc_addr, log, device, model_path.into()).await + serve( + grpc_addr, + log, + device, + backend, + model_path.into(), + helion_addr, + helion_model, + ) + .await } } } @@ -34,7 +46,10 @@ async fn serve( grpc_addr: String, log: String, device: kernelport_core::Device, + backend: String, model_path: std::path::PathBuf, + helion_addr: String, + helion_model: String, ) -> Result<()> { std::env::set_var("RUST_LOG", &log); tracing_subscriber::fmt() @@ -59,7 +74,18 @@ async fn serve( // Model registry let mut reg = registry::ModelRegistry::new(); - reg.load_onnx("demo", model_path, device).ok(); // allow startup without the file for now + match backend.as_str() { + "onnx" => { + reg.load_onnx("demo", model_path, device).ok(); + } + "helion" => { + reg.load_helion("demo", helion_addr, helion_model, device) + .ok(); + } + other => { + anyhow::bail!("unsupported backend: {other} (expected onnx or helion)"); + } + } let loaded = reg.get("demo"); let worker_model = DemoWorkerModel { loaded }; diff --git a/crates/kernelport-server/src/registry.rs b/crates/kernelport-server/src/registry.rs index ddd3172..6dcd185 100644 --- a/crates/kernelport-server/src/registry.rs +++ b/crates/kernelport-server/src/registry.rs @@ -2,6 +2,7 @@ use std::collections::HashMap; use std::sync::{Arc, Mutex}; use anyhow::Result; +use kernelport_backend_helion::HelionBackend; use kernelport_backend_ort::OrtBackend; use kernelport_core::{Backend, BackendModel, Device, IOName, ModelArtifact, Tensor}; @@ -57,6 +58,32 @@ impl ModelRegistry { Ok(()) } + pub fn load_helion( + &mut self, + name: &str, + addr: String, + model: String, + device: Device, + ) -> Result<()> { + let backend = HelionBackend::new(); + let artifact = ModelArtifact::HelionGrpc { addr, model }; + let model = backend.load(&artifact, device)?; + + let output_names = model + .spec() + .outputs + .iter() + .map(|spec| spec.name.clone()) + .collect(); + let loaded = LoadedModel { + model: Mutex::new(Box::new(model)), + output_names, + }; + + self.models.insert(name.to_string(), Arc::new(loaded)); + Ok(()) + } + pub fn get(&self, name: &str) -> Option> { self.models.get(name).cloned() } diff --git a/docker-compose.mock.yml b/docker-compose.mock.yml new file mode 100644 index 0000000..650f85a --- /dev/null +++ b/docker-compose.mock.yml @@ -0,0 +1,21 @@ +version: "3.8" + +services: + helion-worker: + build: + context: . + dockerfile: Dockerfile.helion.mock + ports: + - "50061:50061" + environment: + KERNELPORT_REPO_ROOT: "/app" + + kernelport: + build: + context: . + dockerfile: Dockerfile.cpu + ports: + - "50051:50051" + depends_on: + - helion-worker + command: ["--backend", "helion", "--device", "cpu", "--helion-addr", "http://helion-worker:50061", "--helion-model", "softmax_mock"] diff --git a/docker-compose.yml b/docker-compose.yml new file mode 100644 index 0000000..ff4224e --- /dev/null +++ b/docker-compose.yml @@ -0,0 +1,32 @@ +version: "3.8" + +services: + helion-worker: + build: + context: . + dockerfile: Dockerfile.helion + ports: + - "50061:50061" + environment: + CUDA_VISIBLE_DEVICES: "0" + XDG_CACHE_HOME: "/cache" + KERNELPORT_REPO_ROOT: "/app" + volumes: + - ./helion-cache:/cache + device_requests: + - driver: nvidia + count: 1 + capabilities: ["gpu"] + command: ["python", "/app/scripts/helion/helion_worker.py", "--addr", "0.0.0.0:50061", "--device", "cuda"] + + kernelport: + build: + context: . + dockerfile: Dockerfile.gpu + ports: + - "50051:50051" + environment: + CUDA_VISIBLE_DEVICES: "0" + depends_on: + - helion-worker + command: ["--backend", "helion", "--device", "cuda:0", "--helion-addr", "http://helion-worker:50061", "--helion-model", "softmax_two_pass"] diff --git a/docs/deploy/helion-sidecar.yaml b/docs/deploy/helion-sidecar.yaml new file mode 100644 index 0000000..b72c7df --- /dev/null +++ b/docs/deploy/helion-sidecar.yaml @@ -0,0 +1,50 @@ +apiVersion: v1 +kind: Pod +metadata: + name: kernelport-helion + labels: + app: kernelport-helion +spec: + containers: + - name: kernelport + image: kernelport:gpu + args: + - "--backend" + - "helion" + - "--device" + - "cuda:0" + - "--helion-addr" + - "http://127.0.0.1:50061" + - "--helion-model" + - "softmax_two_pass" + ports: + - containerPort: 50051 + env: + - name: RUST_LOG + value: info + - name: helion-worker + image: kernelport-helion:gpu + args: + - "python" + - "/app/scripts/helion/helion_worker.py" + - "--addr" + - "0.0.0.0:50061" + - "--device" + - "cuda" + ports: + - containerPort: 50061 + resources: + limits: + nvidia.com/gpu: 1 +--- +apiVersion: v1 +kind: Service +metadata: + name: kernelport-helion +spec: + selector: + app: kernelport-helion + ports: + - name: grpc + port: 50051 + targetPort: 50051 diff --git a/docs/validation.md b/docs/validation.md index deb963a..46fdfd1 100644 --- a/docs/validation.md +++ b/docs/validation.md @@ -54,3 +54,74 @@ Notes: - `data` is raw bytes, base64-encoded (`[1.0,2.0,3.0,4.0]` = `AACAPwAAAEAAAEBAAACAQA==`). - Input/output names must match the model (sample identity model uses `x` -> `y`). - `dtype` is an enum; grpcurl accepts the symbolic name (`F32`). + +## 4) Helion softmax (experimental) + +Prereqs (Python, uv): +```bash +uv venv .venv +source .venv/bin/activate +uv pip install "torch==2.9.*" --index-url https://download.pytorch.org/whl/cu126 +uv pip install helion grpcio grpcio-tools numpy +``` + +Start the Helion worker: +```bash +python scripts/helion/helion_worker.py --addr 0.0.0.0:50061 --device cuda +``` + +Start kernelportd with the Helion backend: +```bash +cargo run -p kernelport-server -- serve \ + --backend helion \ + --device cuda:0 \ + --helion-addr http://127.0.0.1:50061 \ + --helion-model softmax_two_pass +``` + +Generate a base64 payload for a small F16 input: +```bash +python - <<'PY' +import base64 +import torch +x = torch.randn(4, 8, dtype=torch.float16) +print(base64.b64encode(x.numpy().tobytes()).decode()) +PY +``` + +Then call the server (paste the base64 string into `data`): +```bash +grpcurl -plaintext -d '{ + "model": "demo", + "inputs": [ + { "name": "x", "dtype": "F16", "shape": [4, 8], "data": "" } + ] +}' localhost:50051 kernelport.v1.InferenceService/Infer +``` + +## 5) Mock Helion (CPU, Docker) + +Start the CPU mock sidecar + kernelport: +```bash +docker compose -f docker-compose.mock.yml up --build +``` + +Generate a base64 payload: +```bash +python - <<'PY' +import base64 +import numpy as np +x = np.random.randn(4, 8).astype(np.float32) +print(base64.b64encode(x.tobytes()).decode()) +PY +``` + +Then call the server (paste the base64 string into `data`): +```bash +grpcurl -plaintext -d '{ + "model": "demo", + "inputs": [ + { "name": "x", "dtype": "F32", "shape": [4, 8], "data": "" } + ] +}' localhost:50051 kernelport.v1.InferenceService/Infer +``` diff --git a/scripts/helion/README.md b/scripts/helion/README.md new file mode 100644 index 0000000..20d8b4d --- /dev/null +++ b/scripts/helion/README.md @@ -0,0 +1,7 @@ +# Helion Workers + +- `helion_worker.py`: CUDA Helion worker (real kernel, needs GPU). +- `mock/helion_worker_mock.py`: CPU-only mock worker for local testing. + +The mock worker mirrors the gRPC interface but uses NumPy softmax and runs +without CUDA or Helion installed. diff --git a/scripts/helion/helion_worker.py b/scripts/helion/helion_worker.py new file mode 100644 index 0000000..d6596ae --- /dev/null +++ b/scripts/helion/helion_worker.py @@ -0,0 +1,174 @@ +from __future__ import annotations + +import argparse +import os +import sys +import tempfile +from concurrent import futures +from typing import Dict, Tuple + +import helion +import helion.language as hl +import torch +from grpc_tools import protoc +import grpc + + +def find_repo_root() -> str: + env_root = os.environ.get("KERNELPORT_REPO_ROOT") + if env_root: + return env_root + here = os.path.abspath(os.path.dirname(__file__)) + for _ in range(4): + candidate = os.path.join(here, "crates", "kernelport-proto", "src", "inference.proto") + if os.path.exists(candidate): + return here + here = os.path.abspath(os.path.join(here, "..")) + return os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..")) + + +def compile_proto(repo_root: str) -> Tuple[object, object]: + proto_dir = os.path.join(repo_root, "crates", "kernelport-proto", "src") + proto_path = os.path.join(proto_dir, "inference.proto") + tmpdir = tempfile.mkdtemp(prefix="kernelport_proto_") + + args = [ + "grpc_tools.protoc", + f"-I{proto_dir}", + f"--python_out={tmpdir}", + f"--grpc_python_out={tmpdir}", + proto_path, + ] + if protoc.main(args) != 0: + raise RuntimeError("failed to generate grpc stubs") + + pkg_dir = os.path.join(tmpdir, "kernelport", "v1") + os.makedirs(pkg_dir, exist_ok=True) + for path in ( + os.path.join(tmpdir, "kernelport", "__init__.py"), + os.path.join(tmpdir, "kernelport", "v1", "__init__.py"), + ): + if not os.path.exists(path): + with open(path, "w", encoding="utf-8") as handle: + handle.write("") + + sys.path.insert(0, tmpdir) + try: + from kernelport.v1 import inference_pb2, inference_pb2_grpc + except ImportError: + import inference_pb2 # type: ignore + import inference_pb2_grpc # type: ignore + + return inference_pb2, inference_pb2_grpc + + +@helion.kernel(autotune_effort="quick") +def softmax_two_pass(x: torch.Tensor) -> torch.Tensor: + m, n = x.size() + out = torch.empty_like(x) + block_size_m = hl.register_block_size(m) + block_size_n = hl.register_block_size(n) + for tile_m in hl.tile(m, block_size=block_size_m): + mi = hl.full([tile_m], float("-inf"), dtype=torch.float32) + di = hl.zeros([tile_m], dtype=torch.float32) + for tile_n in hl.tile(n, block_size=block_size_n): + values = x[tile_m, tile_n] + local_amax = torch.amax(values, dim=1) + mi_next = torch.maximum(mi, local_amax) + di = di * torch.exp(mi - mi_next) + torch.exp(values - mi_next[:, None]).sum(dim=1) + mi = mi_next + for tile_n in hl.tile(n, block_size=block_size_n): + values = x[tile_m, tile_n] + out[tile_m, tile_n] = torch.exp(values - mi[:, None]) / di[:, None] + return out + + +def dtype_from_pb(dtype: int, pb) -> torch.dtype: + mapping = { + pb.F32: torch.float32, + pb.F16: torch.float16, + pb.I64: torch.int64, + pb.I32: torch.int32, + pb.U8: torch.uint8, + } + if dtype not in mapping: + raise ValueError(f"unsupported dtype: {dtype}") + return mapping[dtype] + + +def pb_from_dtype(dtype: torch.dtype, pb) -> int: + mapping = { + torch.float32: pb.F32, + torch.float16: pb.F16, + torch.int64: pb.I64, + torch.int32: pb.I32, + torch.uint8: pb.U8, + } + if dtype not in mapping: + raise ValueError(f"unsupported dtype: {dtype}") + return mapping[dtype] + + +def tensor_from_pb(tensor, pb) -> torch.Tensor: + dtype = dtype_from_pb(tensor.dtype, pb) + shape = list(tensor.shape) + data = torch.frombuffer(memoryview(tensor.data), dtype=dtype) + return data.reshape(shape) + + +def tensor_to_pb(name: str, tensor: torch.Tensor, pb) -> object: + tensor = tensor.detach().contiguous().cpu() + data = tensor.numpy().tobytes() + return pb.Tensor( + name=name, + dtype=pb_from_dtype(tensor.dtype, pb), + shape=list(tensor.shape), + data=data, + ) + + +class HelionService: + def __init__(self, kernels: Dict[str, object], device: str, pb) -> None: + self.kernels = kernels + self.device = device + self.pb = pb + + def Infer(self, request, context): + if not request.inputs: + context.abort(grpc.StatusCode.INVALID_ARGUMENT, "no inputs provided") + kernel = self.kernels.get(request.model) + if kernel is None: + context.abort(grpc.StatusCode.NOT_FOUND, f"unknown model: {request.model}") + + try: + x = tensor_from_pb(request.inputs[0], self.pb).to(self.device) + y = kernel(x) + return self.pb.InferResponse(outputs=[tensor_to_pb("y", y, self.pb)]) + except Exception as exc: + context.abort(grpc.StatusCode.INTERNAL, f"helion inference failed: {exc}") + + +def main() -> None: + parser = argparse.ArgumentParser(description="Helion gRPC worker") + parser.add_argument("--addr", default="0.0.0.0:50061") + parser.add_argument("--device", default="cuda") + args = parser.parse_args() + + repo_root = find_repo_root() + pb, pb_grpc = compile_proto(repo_root) + + kernels = {"softmax_two_pass": softmax_two_pass} + + server = grpc.server(thread_pool=futures.ThreadPoolExecutor(max_workers=4)) + pb_grpc.add_InferenceServiceServicer_to_server( + HelionService(kernels, args.device, pb), server + ) + + server.add_insecure_port(args.addr) + server.start() + print(f"Helion worker listening on {args.addr}") + server.wait_for_termination() + + +if __name__ == "__main__": + main() diff --git a/scripts/helion/mock/helion_worker_mock.py b/scripts/helion/mock/helion_worker_mock.py new file mode 100644 index 0000000..3d1d076 --- /dev/null +++ b/scripts/helion/mock/helion_worker_mock.py @@ -0,0 +1,148 @@ +from __future__ import annotations + +import argparse +import os +import sys +import tempfile +from concurrent import futures +from typing import Tuple + +import grpc +import numpy as np +from grpc_tools import protoc + + +def find_repo_root() -> str: + env_root = os.environ.get("KERNELPORT_REPO_ROOT") + if env_root: + return env_root + here = os.path.abspath(os.path.dirname(__file__)) + for _ in range(5): + candidate = os.path.join(here, "crates", "kernelport-proto", "src", "inference.proto") + if os.path.exists(candidate): + return here + here = os.path.abspath(os.path.join(here, "..")) + return os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..", "..")) + + +def compile_proto(repo_root: str) -> Tuple[object, object]: + proto_dir = os.path.join(repo_root, "crates", "kernelport-proto", "src") + proto_path = os.path.join(proto_dir, "inference.proto") + tmpdir = tempfile.mkdtemp(prefix="kernelport_proto_") + + args = [ + "grpc_tools.protoc", + f"-I{proto_dir}", + f"--python_out={tmpdir}", + f"--grpc_python_out={tmpdir}", + proto_path, + ] + if protoc.main(args) != 0: + raise RuntimeError("failed to generate grpc stubs") + + pkg_dir = os.path.join(tmpdir, "kernelport", "v1") + os.makedirs(pkg_dir, exist_ok=True) + for path in ( + os.path.join(tmpdir, "kernelport", "__init__.py"), + os.path.join(tmpdir, "kernelport", "v1", "__init__.py"), + ): + if not os.path.exists(path): + with open(path, "w", encoding="utf-8") as handle: + handle.write("") + + sys.path.insert(0, tmpdir) + try: + from kernelport.v1 import inference_pb2, inference_pb2_grpc + except ImportError: + import inference_pb2 # type: ignore + import inference_pb2_grpc # type: ignore + + return inference_pb2, inference_pb2_grpc + + +def dtype_from_pb(dtype: int, pb) -> np.dtype: + mapping = { + pb.F32: np.float32, + pb.F16: np.float16, + pb.I64: np.int64, + pb.I32: np.int32, + pb.U8: np.uint8, + } + if dtype not in mapping: + raise ValueError(f"unsupported dtype: {dtype}") + return mapping[dtype] + + +def pb_from_dtype(dtype: np.dtype, pb) -> int: + mapping = { + np.dtype(np.float32): pb.F32, + np.dtype(np.float16): pb.F16, + np.dtype(np.int64): pb.I64, + np.dtype(np.int32): pb.I32, + np.dtype(np.uint8): pb.U8, + } + if np.dtype(dtype) not in mapping: + raise ValueError(f"unsupported dtype: {dtype}") + return mapping[np.dtype(dtype)] + + +def tensor_from_pb(tensor, pb) -> np.ndarray: + dtype = dtype_from_pb(tensor.dtype, pb) + shape = list(tensor.shape) + data = np.frombuffer(tensor.data, dtype=dtype) + return data.reshape(shape) + + +def tensor_to_pb(name: str, array: np.ndarray, pb) -> object: + array = np.ascontiguousarray(array) + return pb.Tensor( + name=name, + dtype=pb_from_dtype(array.dtype, pb), + shape=list(array.shape), + data=array.tobytes(), + ) + + +def softmax(x: np.ndarray, axis: int = 1) -> np.ndarray: + x_max = np.max(x, axis=axis, keepdims=True) + exp = np.exp(x - x_max) + return exp / np.sum(exp, axis=axis, keepdims=True) + + +class MockHelionService: + def __init__(self, pb) -> None: + self.pb = pb + + def Infer(self, request, context): + if not request.inputs: + context.abort(grpc.StatusCode.INVALID_ARGUMENT, "no inputs provided") + if request.model not in ("softmax_two_pass", "softmax_mock"): + context.abort(grpc.StatusCode.NOT_FOUND, f"unknown model: {request.model}") + + try: + x = tensor_from_pb(request.inputs[0], self.pb) + y = softmax(x, axis=1) + return self.pb.InferResponse(outputs=[tensor_to_pb("y", y, self.pb)]) + except Exception as exc: + context.abort(grpc.StatusCode.INTERNAL, f"mock inference failed: {exc}") + + +def main() -> None: + parser = argparse.ArgumentParser(description="Mock Helion gRPC worker (CPU)") + parser.add_argument("--addr", default="0.0.0.0:50061") + args = parser.parse_args() + + repo_root = find_repo_root() + pb, pb_grpc = compile_proto(repo_root) + + server = grpc.server(thread_pool=futures.ThreadPoolExecutor(max_workers=4)) + pb_grpc.add_InferenceServiceServicer_to_server(MockHelionService(pb), server) + + server.add_insecure_port(args.addr) + server.start() + print(f"Mock Helion worker listening on {args.addr}") + server.wait_for_termination() + + +if __name__ == "__main__": + main() From 0e214cbd4f4e1b6fc0a709dee7da733a9b3439e6 Mon Sep 17 00:00:00 2001 From: beatgeek Date: Sat, 31 Jan 2026 18:25:36 -0800 Subject: [PATCH 2/2] LuxTTS Lambda deploy: manifest, worker, GHCR, deploy workflow - Add LuxTTS model manifest and deployments config - Add LuxTTS gRPC worker (tensor contract) and Dockerfile.luxtts (uv) - Add docker-compose.luxtts.yml override - Add manifest_to_serve_args.py and entrypoint enhancement - Add deploy-lambda.yml: build/push to GHCR, launch Lambda instance - Add Python script tests and pre-commit hooks - Document deployments, Lambda deploy, local venv; add .vscode to gitignore Co-authored-by: Cursor --- .github/workflows/ci.yml | 11 ++ .github/workflows/deploy-lambda.yml | 127 +++++++++++++++++++++ .gitignore | 3 + .pre-commit-config.yaml | 12 ++ Dockerfile.luxtts | 38 +++++++ Makefile | 20 ++-- README.md | 56 ++++++++++ deployments/deployments.yaml | 19 ++++ docker-compose.luxtts.yml | 23 ++++ docs/deployments.md | 55 +++++++++ models/luxtts-manifest.yaml | 11 ++ scripts/entrypoint.sh | 4 +- scripts/luxtts/luxtts_worker.py | 167 ++++++++++++++++++++++++++++ scripts/luxtts/proto_util.py | 57 ++++++++++ scripts/manifest_to_serve_args.py | 57 ++++++++++ tests/test_manifest_serve_args.py | 84 ++++++++++++++ 16 files changed, 735 insertions(+), 9 deletions(-) create mode 100644 .github/workflows/deploy-lambda.yml create mode 100644 Dockerfile.luxtts create mode 100644 deployments/deployments.yaml create mode 100644 docker-compose.luxtts.yml create mode 100644 docs/deployments.md create mode 100644 models/luxtts-manifest.yaml create mode 100644 scripts/luxtts/luxtts_worker.py create mode 100644 scripts/luxtts/proto_util.py create mode 100644 scripts/manifest_to_serve_args.py create mode 100644 tests/test_manifest_serve_args.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 3575880..7fcfcb2 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -46,3 +46,14 @@ jobs: - name: ort identity (cpu) run: cargo test -p kernelport-backend-ort --test identity + + python-scripts: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - name: install PyYAML + run: pip install pyyaml + + - name: manifest_to_serve_args tests + run: python tests/test_manifest_serve_args.py diff --git a/.github/workflows/deploy-lambda.yml b/.github/workflows/deploy-lambda.yml new file mode 100644 index 0000000..88d1128 --- /dev/null +++ b/.github/workflows/deploy-lambda.yml @@ -0,0 +1,127 @@ +# Deploy KernelPort (LuxTTS) to Lambda Cloud: build/push images to GHCR, launch GPU instance. +# Trigger: workflow_dispatch. Requires repository secrets: LAMBDA_CLOUD_API_KEY, HUGGINGFACE_HUB_TOKEN. +# You must add an SSH key in Lambda Cloud (https://cloud.lambda.ai/ssh-keys) and pass its name as input. +name: Deploy to Lambda Cloud + +on: + workflow_dispatch: + inputs: + instance_type_name: + description: "Lambda instance type (e.g. gpu_1x_a100)" + required: true + default: "gpu_1x_a100" + region_name: + description: "Lambda region (e.g. us-tx-1)" + required: true + default: "us-tx-1" + ssh_key_name: + description: "Name of SSH key already added in Lambda Cloud" + required: true + +env: + LAMBDA_API_BASE: "https://cloud.lambdalabs.com/api/v1" + +jobs: + build-push: + runs-on: ubuntu-latest + permissions: + contents: read + packages: write + steps: + - uses: actions/checkout@v4 + + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@v3 + + - name: Log in to GHCR + uses: docker/login-action@v3 + with: + registry: ghcr.io + username: ${{ github.actor }} + password: ${{ secrets.GITHUB_TOKEN }} + + - name: Build and push kernelport (GPU) + uses: docker/build-push-action@v6 + with: + context: . + file: ./Dockerfile.gpu + push: true + tags: ghcr.io/${{ github.repository_owner }}/kernelport:gpu,ghcr.io/${{ github.repository_owner }}/kernelport:gpu-${{ github.sha }} + cache-from: type=gha + cache-to: type=gha,mode=max + + - name: Build and push kernelport-luxtts (GPU) + uses: docker/build-push-action@v6 + with: + context: . + file: ./Dockerfile.luxtts + push: true + tags: ghcr.io/${{ github.repository_owner }}/kernelport-luxtts:gpu,ghcr.io/${{ github.repository_owner }}/kernelport-luxtts:gpu-${{ github.sha }} + cache-from: type=gha + cache-to: type=gha,mode=max + + launch: + needs: build-push + runs-on: ubuntu-latest + steps: + - name: List instance types (optional) + run: | + curl -sSf -u "${{ secrets.LAMBDA_CLOUD_API_KEY }}:" \ + "${{ env.LAMBDA_API_BASE }}/instance-types" | jq -r '.data | keys[]' | head -20 + + - name: Launch instance + id: launch + run: | + body=$(jq -n \ + --arg region "${{ github.event.inputs.region_name }}" \ + --arg type "${{ github.event.inputs.instance_type_name }}" \ + --arg key "${{ github.event.inputs.ssh_key_name }}" \ + '{region_name: $region, instance_type_name: $type, ssh_key_names: [$key]}') + resp=$(curl -sSf -u "${{ secrets.LAMBDA_CLOUD_API_KEY }}:" \ + -X POST "${{ env.LAMBDA_API_BASE }}/instance-operations/launch" \ + -H "Content-Type: application/json" \ + -d "$body") + echo "instance_ids=$(echo "$resp" | jq -r '.data.instance_ids[0]')" >> $GITHUB_OUTPUT + echo "$resp" | jq . + + - name: Poll until instance is active + id: poll + run: | + id="${{ steps.launch.outputs.instance_ids }}" + for i in $(seq 1 30); do + resp=$(curl -sSf -u "${{ secrets.LAMBDA_CLOUD_API_KEY }}:" \ + "${{ env.LAMBDA_API_BASE }}/instances/$id") + status=$(echo "$resp" | jq -r '.data.status') + ip=$(echo "$resp" | jq -r '.data.ip // empty') + echo "Attempt $i: status=$status ip=$ip" + if [ "$status" = "active" ] && [ -n "$ip" ]; then + echo "instance_ip=$ip" >> $GITHUB_OUTPUT + echo "Instance is up at $ip" + exit 0 + fi + sleep 20 + done + echo "Instance did not become active in time" + exit 1 + + - name: Smoke test (grpcurl) + continue-on-error: true + run: | + ip="${{ steps.poll.outputs.instance_ip }}" + if [ -z "$ip" ]; then exit 0; fi + # Instance may not have the stack running yet; user can SSH and run docker compose + if command -v grpcurl >/dev/null 2>&1; then + grpcurl -plaintext -connect-timeout 5 "$ip:50051" list || true + else + echo "grpcurl not installed; skip smoke test. Endpoint: $ip:50051" + fi + + - name: Summary + run: | + echo "## Lambda Cloud instance" >> $GITHUB_STEP_SUMMARY + echo "" >> $GITHUB_STEP_SUMMARY + echo "- **Instance ID**: ${{ steps.launch.outputs.instance_ids }}" >> $GITHUB_STEP_SUMMARY + echo "- **IP**: ${{ steps.poll.outputs.instance_ip }}" >> $GITHUB_STEP_SUMMARY + echo "- **gRPC endpoint**: ${{ steps.poll.outputs.instance_ip }}:50051 (after you deploy the stack on the instance)" >> $GITHUB_STEP_SUMMARY + echo "" >> $GITHUB_STEP_SUMMARY + echo "SSH into the instance and run your stack (e.g. docker compose). Set \`HUGGINGFACE_HUB_TOKEN\` when running the LuxTTS worker." >> $GITHUB_STEP_SUMMARY diff --git a/.gitignore b/.gitignore index 8f0a58b..af3c053 100644 --- a/.gitignore +++ b/.gitignore @@ -17,6 +17,9 @@ target *.onnx !models/identity.onnx +# Editor / IDE +.vscode/ + # RustRover # JetBrains specific template is maintained in a separate JetBrains.gitignore that can # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 6238b31..40f0f11 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -7,3 +7,15 @@ repos: language: system types: [rust] pass_filenames: false + + - id: cargo-test + name: cargo test + entry: cargo test --all + language: system + pass_filenames: false + + - id: python-scripts + name: Python script tests (manifest_to_serve_args) + entry: python tests/test_manifest_serve_args.py + language: system + pass_filenames: false diff --git a/Dockerfile.luxtts b/Dockerfile.luxtts new file mode 100644 index 0000000..45b0ef5 --- /dev/null +++ b/Dockerfile.luxtts @@ -0,0 +1,38 @@ +# LuxTTS gRPC worker (Helion-style sidecar). Loads YatharthS/LuxTTS from HF. +# Uses uv for consistency with Dockerfile.helion. +FROM nvidia/cuda:12.2.0-runtime-ubuntu22.04 + +RUN apt-get update \ + && apt-get install -y --no-install-recommends \ + ca-certificates \ + curl \ + git \ + python3 \ + python3-venv \ + && rm -rf /var/lib/apt/lists/* + +RUN curl -LsSf https://astral.sh/uv/install.sh | sh + +ENV VIRTUAL_ENV=/opt/venv +ENV PATH="/opt/venv/bin:/root/.local/bin:${PATH}" + +WORKDIR /app + +# Clone LuxTTS so zipvoice.luxvoice is available; install its deps with uv +RUN git clone --depth 1 https://github.com/ysharma3501/LuxTTS.git /app/LuxTTS +RUN uv venv /opt/venv \ + && uv pip install -r /app/LuxTTS/requirements.txt + +# gRPC and proto for kernelport InferenceService +RUN uv pip install grpcio grpcio-tools + +COPY scripts/luxtts /app/scripts/luxtts +COPY crates/kernelport-proto/src /app/crates/kernelport-proto/src + +ENV LUXTTS_REPO=/app/LuxTTS +ENV KERNELPORT_REPO_ROOT=/app +ENV PYTHONPATH=/app/LuxTTS:/app + +EXPOSE 50061 + +CMD ["python", "/app/scripts/luxtts/luxtts_worker.py", "--addr", "0.0.0.0:50061", "--device", "cuda"] diff --git a/Makefile b/Makefile index edb303e..7ea9961 100644 --- a/Makefile +++ b/Makefile @@ -1,5 +1,5 @@ .DEFAULT_GOAL := help -.PHONY: fmt clippy test build check run help +.PHONY: fmt clippy test test-python build check run help PROJECT_NAME := kernelport PROJECT_VERSION := $(shell cargo pkgid -p kernelport-server 2>/dev/null | sed -E 's/.*@//') @@ -16,10 +16,13 @@ clippy: test: cargo test --all +test-python: + python3 tests/test_manifest_serve_args.py + build: cargo build --all -check: fmt clippy test build +check: fmt clippy test test-python build run: RUST_LOG=info cargo run -p kernelport-server @@ -28,9 +31,10 @@ help: @printf "%s\n" \ "$(PROJECT_NAME) $(PROJECT_VERSION)" \ "Targets:" \ - " fmt - Run rustfmt" \ - " clippy - Run clippy with warnings denied" \ - " test - Run tests" \ - " build - Build all crates" \ - " check - Run fmt, clippy, test, build" \ - " run - Run kernelport-server" + " fmt - Run rustfmt" \ + " clippy - Run clippy with warnings denied" \ + " test - Run Rust tests" \ + " test-python - Run Python script tests (manifest_to_serve_args)" \ + " build - Build all crates" \ + " check - Run fmt, clippy, test, test-python, build" \ + " run - Run kernelport-server" diff --git a/README.md b/README.md index a6acd26..249e468 100644 --- a/README.md +++ b/README.md @@ -151,6 +151,22 @@ The scheduler can be tuned for: ## Local Development +### Local Python (venv) + +For the Python pieces (e.g. `make test-python`, pre-commit Python hook, `model_fetch.py`, or +Helion/LuxTTS workers), use a venv so your system Python setup doesn't conflict: + +```bash +uv venv .venv +source .venv/bin/activate # or .venv\Scripts\activate on Windows +uv pip install pyyaml +``` + +Then run `make test-python` or `pre-commit run --all-files`; the hook will use the active venv. +To run the Helion worker locally, add: `uv pip install "torch==2.9.*" helion grpcio grpcio-tools numpy` +(see Helion section for PyTorch index). For LuxTTS, use the Docker setup or install from the +LuxTTS repo requirements. + ### CPU (macOS or Linux) - Install Rust components: `rustup component add rustfmt clippy` @@ -244,6 +260,41 @@ See `docs/model-ingestion.md` for the HuggingFace model ingestion plan and Validation steps are in `docs/validation.md`. +### Deployments and secrets + +For deployment options (LuxTTS, Lambda Cloud, GHCR) and required secrets, see +[docs/deployments.md](docs/deployments.md). Summary: + +- **LuxTTS**: Set `HUGGINGFACE_HUB_TOKEN` (or `HF_TOKEN`) when running the LuxTTS worker. +- **Lambda Cloud deploy** (GitHub Actions): Add repository secrets + **LAMBDA_CLOUD_API_KEY** and **HUGGINGFACE_HUB_TOKEN** in + Settings → Secrets and variables → Actions. + +### Deploy to Lambda Cloud + +The GitHub Actions workflow **Deploy to Lambda Cloud** builds the kernelport and LuxTTS +images, pushes them to GHCR, and launches a GPU instance on [Lambda Cloud](https://cloud.lambda.ai/). + +1. **Secrets** (one-time): In the repo go to **Settings → Secrets and variables → Actions**. + Add **LAMBDA_CLOUD_API_KEY** and **HUGGINGFACE_HUB_TOKEN**. + +2. **SSH key** (one-time): In [Lambda Cloud SSH keys](https://cloud.lambda.ai/ssh-keys), add + an SSH key and note its **name** (e.g. `macbook-pro`). The workflow needs this name. + +3. **Run the workflow**: Push your branch, then go to **Actions → Deploy to Lambda Cloud → + Run workflow**. Fill the inputs: + - **instance_type_name**: e.g. `gpu_1x_a100` (see [Lambda instance types](https://cloud.lambda.ai/instances)). + - **region_name**: e.g. `us-tx-1`. + - **ssh_key_name**: the exact name of your SSH key from step 2 (required). + +4. **After the run**: The job summary shows the **instance ID** and **public IP**. SSH into + the instance and run your stack (e.g. pull images from GHCR and run docker compose with + `HUGGINGFACE_HUB_TOKEN`). The gRPC inference endpoint is **`:50051`** once + the stack is running. + +See [.github/workflows/deploy-lambda.yml](.github/workflows/deploy-lambda.yml) and +[docs/deployments.md](docs/deployments.md) for details. + ## Pre-commit (local) Install pre-commit and enable the hook: @@ -253,8 +304,13 @@ pipx install pre-commit pre-commit install ``` +Hooks run: `cargo fmt`, `cargo test --all`, and Python script tests (`manifest_to_serve_args`). +The Python hook requires PyYAML; using a venv is recommended (see [Local Python (venv)](#local-python-venv)). + Run on demand: ```bash pre-commit run --all-files ``` + +Run only Python script tests: `make test-python` (requires PyYAML). diff --git a/deployments/deployments.yaml b/deployments/deployments.yaml new file mode 100644 index 0000000..8c53f76 --- /dev/null +++ b/deployments/deployments.yaml @@ -0,0 +1,19 @@ +# Deployment name -> manifest, backend, optional worker image. +# Used by CI/Lambda and docs to know how to run each deployment. +deployments: + identity-onnx: + manifest: models/sample-manifest.yaml + backend: onnx + worker_image: null + + helion-demo: + manifest: null + backend: helion + helion_model: softmax_two_pass + worker_image: kernelport-helion:gpu + + luxtts: + manifest: models/luxtts-manifest.yaml + backend: helion + helion_model: luxtts + worker_image: kernelport-luxtts:gpu diff --git a/docker-compose.luxtts.yml b/docker-compose.luxtts.yml new file mode 100644 index 0000000..d1510fe --- /dev/null +++ b/docker-compose.luxtts.yml @@ -0,0 +1,23 @@ +# Override to run LuxTTS worker instead of Helion demo worker. +# Usage: docker compose -f docker-compose.yml -f docker-compose.luxtts.yml up --build +# Set HUGGINGFACE_HUB_TOKEN (or HF_TOKEN) so the LuxTTS worker can pull the model. +version: "3.8" + +services: + helion-worker: + build: + context: . + dockerfile: Dockerfile.luxtts + image: kernelport-luxtts:gpu + environment: + CUDA_VISIBLE_DEVICES: "0" + KERNELPORT_REPO_ROOT: "/app" + LUXTTS_REPO: "/app/LuxTTS" + HUGGINGFACE_HUB_TOKEN: ${HUGGINGFACE_HUB_TOKEN:-} + HF_TOKEN: ${HF_TOKEN:-} + command: ["python", "/app/scripts/luxtts/luxtts_worker.py", "--addr", "0.0.0.0:50061", "--device", "cuda", "--model-name", "luxtts"] + + kernelport: + command: ["--backend", "helion", "--device", "cuda:0", "--helion-addr", "http://helion-worker:50061", "--helion-model", "luxtts"] + depends_on: + - helion-worker diff --git a/docs/deployments.md b/docs/deployments.md new file mode 100644 index 0000000..df0bd49 --- /dev/null +++ b/docs/deployments.md @@ -0,0 +1,55 @@ +# Deployments + +This document describes how to run each KernelPort deployment (model + backend + optional worker). + +## Deployment list + +| Name | Manifest | Backend | Worker image | How to run | +|----------------|-----------------------------|---------|------------------------|------------| +| identity-onnx | models/sample-manifest.yaml | onnx | — | `MODEL_MANIFEST_PATH=... model_fetch` then `serve --backend onnx --model-path /models//model.onnx` | +| helion-demo | — | helion | kernelport-helion:gpu | `docker compose up` (default) | +| luxtts | models/luxtts-manifest.yaml | helion | kernelport-luxtts:gpu | `docker compose -f docker-compose.yml -f docker-compose.luxtts.yml up` or Lambda deploy | + +The authoritative mapping is in [deployments/deployments.yaml](../deployments/deployments.yaml). + +## Running the LuxTTS deployment + +1. Set `HUGGINGFACE_HUB_TOKEN` (or `HF_TOKEN`) so the LuxTTS worker can pull the model from Hugging Face. +2. Build and run with the LuxTTS override: + ```bash + docker compose -f docker-compose.yml -f docker-compose.luxtts.yml up --build + ``` +3. Inference endpoint: gRPC at `localhost:50051`, method `kernelport.v1.InferenceService/Infer`. See [LuxTTS tensor contract](#luxtts-tensor-contract) for input/output tensors. + +## Lambda Cloud deploy + +The GitHub Actions workflow [.github/workflows/deploy-lambda.yml](../.github/workflows/deploy-lambda.yml) can create a Lambda Cloud GPU instance, build and push images to GHCR, and run the LuxTTS stack (or another deployment) via cloud-init. + +Required repository secrets: + +- **LAMBDA_CLOUD_API_KEY** — Lambda Cloud API key for launching instances. +- **HUGGINGFACE_HUB_TOKEN** — Hugging Face token for pulling models (e.g. LuxTTS); passed to the LuxTTS worker at runtime. + +Trigger: `workflow_dispatch` (manual) or push to a branch you configure. See the workflow file for inputs (e.g. instance type). + +## LuxTTS tensor contract + +Inputs (all optional except `text` and `prompt_audio`; omitted use LuxTTS defaults): + +| Name | Dtype | Shape | Description | +|-----------------|-------|-------|-------------| +| text | U8 | [N] | UTF-8 bytes of the text to synthesize. | +| prompt_audio | U8 | [M] | Raw reference audio file bytes (WAV or MP3). | +| rms | F32 | [1] | Loudness (default 0.01). | +| t_shift | F32 | [1] | Sampling param (default 0.9). | +| num_steps | I32 | [1] | Sampling steps (default 4). | +| speed | F32 | [1] | Playback speed (default 1.0). | +| return_smooth | I32 | [1] | 0 = false, 1 = true (default 0). | +| ref_duration | I32 | [1] | Reference duration in seconds (default 5). | + +Outputs: + +| Name | Dtype | Shape | Description | +|--------------|-------|-------|-------------| +| audio | F32 | [S] | Mono waveform at 48 kHz. | +| sample_rate | I32 | [1] | 48000. | diff --git a/models/luxtts-manifest.yaml b/models/luxtts-manifest.yaml new file mode 100644 index 0000000..544e4db --- /dev/null +++ b/models/luxtts-manifest.yaml @@ -0,0 +1,11 @@ +id: luxtts +source: + kind: huggingface + repo: YatharthS/LuxTTS + revision: main +format: python +backend: helion +cache: + dir: /models +runtime: + device: cuda:0 diff --git a/scripts/entrypoint.sh b/scripts/entrypoint.sh index f481511..1979fef 100644 --- a/scripts/entrypoint.sh +++ b/scripts/entrypoint.sh @@ -1,8 +1,10 @@ #!/usr/bin/env sh set -eu +SERVE_ARGS="" if [ -n "${MODEL_MANIFEST_PATH:-}" ]; then python3 /app/scripts/model_fetch.py "$MODEL_MANIFEST_PATH" + SERVE_ARGS=$(python3 /app/scripts/manifest_to_serve_args.py "$MODEL_MANIFEST_PATH") fi -exec /app/kernelportd serve "$@" +exec /app/kernelportd serve $SERVE_ARGS "$@" diff --git a/scripts/luxtts/luxtts_worker.py b/scripts/luxtts/luxtts_worker.py new file mode 100644 index 0000000..b223ccd --- /dev/null +++ b/scripts/luxtts/luxtts_worker.py @@ -0,0 +1,167 @@ +""" +LuxTTS gRPC worker: implements kernelport.v1.InferenceService for LuxTTS. + +Input tensors (LuxTTS contract): text (U8), prompt_audio (U8), optional rms, t_shift, +num_steps, speed, return_smooth, ref_duration. Output: audio (F32), sample_rate (I32). +""" +from __future__ import annotations + +import argparse +import os +import struct +import sys +import tempfile +from concurrent import futures +from typing import Any, Dict, Optional + +import grpc +import numpy as np + +# Ensure script dir is on path for proto_util +_script_dir = os.path.dirname(os.path.abspath(__file__)) +if _script_dir not in sys.path: + sys.path.insert(0, _script_dir) +from proto_util import compile_proto, find_repo_root + +# LuxTTS repo is expected at LUXTTS_REPO or /app/LuxTTS in Docker +_luxtts_repo = os.environ.get("LUXTTS_REPO", "/app/LuxTTS") +if os.path.isdir(_luxtts_repo): + sys.path.insert(0, _luxtts_repo) + +# After path is set, import LuxTTS +try: + from zipvoice.luxvoice import LuxTTS +except ImportError as e: + LuxTTS = None # type: ignore + _import_error = e +else: + _import_error = None + +# DType enum values from inference.proto +DTYPE_F32 = 1 +DTYPE_I32 = 4 +DTYPE_U8 = 5 + + +def _inputs_by_name(request: Any) -> Dict[str, Any]: + """Build a dict of input name -> tensor (name, dtype, shape, data).""" + out = {} + for t in request.inputs: + out[t.name] = {"dtype": t.dtype, "shape": list(t.shape), "data": t.data} + return out + + +def _scalar_f32(data: bytes) -> float: + if len(data) < 4: + return 0.0 + return struct.unpack(" int: + if len(data) < 4: + return 0 + return struct.unpack(" Any: + return pb_module.Tensor(name=name, dtype=dtype, shape=shape, data=data) + + +class LuxTTSService: + def __init__(self, model_name: str, device: str, pb_module: Any) -> None: + if LuxTTS is None: + raise RuntimeError(f"LuxTTS not available: {_import_error}") + self.pb = pb_module + self.device = device + self.model_name = model_name + token = os.environ.get("HUGGINGFACE_HUB_TOKEN") or os.environ.get("HF_TOKEN") + self.lux_tts = LuxTTS("YatharthS/LuxTTS", device=device, token=token) + + def Infer(self, request: Any, context: grpc.ServicerContext) -> Any: + if request.model != self.model_name: + context.abort(grpc.StatusCode.NOT_FOUND, f"unknown model: {request.model}") + + inputs = _inputs_by_name(request) + if "text" not in inputs or "prompt_audio" not in inputs: + context.abort( + grpc.StatusCode.INVALID_ARGUMENT, + "inputs must include 'text' (U8) and 'prompt_audio' (U8)", + ) + + try: + text = inputs["text"]["data"].decode("utf-8") + except Exception as e: + context.abort(grpc.StatusCode.INVALID_ARGUMENT, f"text must be UTF-8: {e}") + + prompt_audio_bytes = inputs["prompt_audio"]["data"] + rms = _scalar_f32(inputs["rms"]["data"]) if "rms" in inputs else 0.01 + t_shift = _scalar_f32(inputs["t_shift"]["data"]) if "t_shift" in inputs else 0.9 + num_steps = _scalar_i32(inputs["num_steps"]["data"]) if "num_steps" in inputs else 4 + speed = _scalar_f32(inputs["speed"]["data"]) if "speed" in inputs else 1.0 + return_smooth = bool(_scalar_i32(inputs["return_smooth"]["data"])) if "return_smooth" in inputs else False + ref_duration = _scalar_i32(inputs["ref_duration"]["data"]) if "ref_duration" in inputs else 5 + + with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f: + f.write(prompt_audio_bytes) + prompt_path = f.name + try: + encoded_prompt = self.lux_tts.encode_prompt( + prompt_path, duration=ref_duration, rms=rms + ) + final_wav = self.lux_tts.generate_speech( + text, + encoded_prompt, + num_steps=num_steps, + t_shift=t_shift, + speed=speed, + return_smooth=return_smooth, + ) + finally: + try: + os.unlink(prompt_path) + except OSError: + pass + + # LuxTTS returns tensor; ensure numpy float32 1D + if hasattr(final_wav, "numpy"): + audio_np = final_wav.numpy().squeeze().astype(np.float32) + else: + audio_np = np.asarray(final_wav, dtype=np.float32).squeeze() + if audio_np.ndim != 1: + audio_np = audio_np.ravel() + audio_bytes = audio_np.tobytes() + sample_rate = 48000 + sr_bytes = np.array([sample_rate], dtype=np.int32).tobytes() + + return self.pb.InferResponse( + outputs=[ + _tensor_pb(self.pb, "audio", DTYPE_F32, [len(audio_np)], audio_bytes), + _tensor_pb(self.pb, "sample_rate", DTYPE_I32, [1], sr_bytes), + ], + queued_us=0, + batched_us=0, + backend_us=0, + ) + + +def main() -> None: + parser = argparse.ArgumentParser(description="LuxTTS gRPC worker") + parser.add_argument("--addr", default="0.0.0.0:50061") + parser.add_argument("--device", default="cuda") + parser.add_argument("--model-name", default="luxtts") + args = parser.parse_args() + + repo_root = find_repo_root() + pb_mod, pb_grpc_mod = compile_proto(repo_root) + + service = LuxTTSService(args.model_name, args.device, pb_mod) + server = grpc.server(thread_pool=futures.ThreadPoolExecutor(max_workers=4)) + pb_grpc_mod.add_InferenceServiceServicer_to_server(service, server) + server.add_insecure_port(args.addr) + server.start() + print(f"LuxTTS worker listening on {args.addr}") + server.wait_for_termination() + + +if __name__ == "__main__": + main() diff --git a/scripts/luxtts/proto_util.py b/scripts/luxtts/proto_util.py new file mode 100644 index 0000000..acd5029 --- /dev/null +++ b/scripts/luxtts/proto_util.py @@ -0,0 +1,57 @@ +"""Shared proto compilation for workers that use kernelport inference proto.""" +from __future__ import annotations + +import os +import sys +import tempfile +from typing import Tuple + +from grpc_tools import protoc + + +def find_repo_root() -> str: + env_root = os.environ.get("KERNELPORT_REPO_ROOT") + if env_root: + return env_root + here = os.path.abspath(os.path.dirname(__file__)) + for _ in range(5): + candidate = os.path.join(here, "crates", "kernelport-proto", "src", "inference.proto") + if os.path.exists(candidate): + return here + here = os.path.abspath(os.path.join(here, "..")) + return os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..")) + + +def compile_proto(repo_root: str) -> Tuple[object, object]: + proto_dir = os.path.join(repo_root, "crates", "kernelport-proto", "src") + proto_path = os.path.join(proto_dir, "inference.proto") + tmpdir = tempfile.mkdtemp(prefix="kernelport_proto_") + + args = [ + "grpc_tools.protoc", + f"-I{proto_dir}", + f"--python_out={tmpdir}", + f"--grpc_python_out={tmpdir}", + proto_path, + ] + if protoc.main(args) != 0: + raise RuntimeError("failed to generate grpc stubs") + + pkg_dir = os.path.join(tmpdir, "kernelport", "v1") + os.makedirs(pkg_dir, exist_ok=True) + for path in ( + os.path.join(tmpdir, "kernelport", "__init__.py"), + os.path.join(tmpdir, "kernelport", "v1", "__init__.py"), + ): + if not os.path.exists(path): + with open(path, "w", encoding="utf-8") as handle: + handle.write("") + + sys.path.insert(0, tmpdir) + try: + from kernelport.v1 import inference_pb2, inference_pb2_grpc + except ImportError: + import inference_pb2 # type: ignore + import inference_pb2_grpc # type: ignore + + return inference_pb2, inference_pb2_grpc diff --git a/scripts/manifest_to_serve_args.py b/scripts/manifest_to_serve_args.py new file mode 100644 index 0000000..a104940 --- /dev/null +++ b/scripts/manifest_to_serve_args.py @@ -0,0 +1,57 @@ +#!/usr/bin/env python3 +""" +Read a model manifest and print kernelportd serve args (--backend, --model-path or +--helion-model, etc.). Used by entrypoint.sh when MODEL_MANIFEST_PATH is set. +""" +from __future__ import annotations + +import os +import sys +from pathlib import Path + +import yaml + + +def load_manifest(path: Path) -> dict: + with path.open("r", encoding="utf-8") as f: + return yaml.safe_load(f) + + +def main() -> int: + if len(sys.argv) != 2: + print("usage: manifest_to_serve_args.py ", file=sys.stderr) + return 2 + + manifest_path = Path(sys.argv[1]).resolve() + manifest = load_manifest(manifest_path) + + backend = manifest.get("backend") + if not backend: + backend = "onnx" if manifest.get("format") == "onnx" else "helion" + model_id = manifest.get("id", "model") + cache_dir = manifest.get("cache", {}).get("dir", "/models") + + args = ["--backend", backend] + + if backend == "onnx": + files = manifest.get("files", ["model.onnx"]) + model_file = files[0] if files else "model.onnx" + model_path = Path(cache_dir) / model_id / model_file + args.extend(["--model-path", str(model_path)]) + elif backend == "helion": + helion_model = manifest.get("helion_model", model_id) + helion_addr = os.environ.get("HELION_ADDR", "http://127.0.0.1:50061") + args.extend(["--helion-addr", helion_addr, "--helion-model", helion_model]) + else: + print(f"unknown backend: {backend}", file=sys.stderr) + return 1 + + device = manifest.get("runtime", {}).get("device", "cpu") + args.extend(["--device", device]) + + print(" ".join(args)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/test_manifest_serve_args.py b/tests/test_manifest_serve_args.py new file mode 100644 index 0000000..cf3c7fd --- /dev/null +++ b/tests/test_manifest_serve_args.py @@ -0,0 +1,84 @@ +#!/usr/bin/env python3 +""" +Test manifest_to_serve_args.py output for sample and LuxTTS manifests. +Run from repo root: python3 tests/test_manifest_serve_args.py +Requires PyYAML: pip3 install pyyaml or uv pip install pyyaml +No pytest required. +""" +from __future__ import annotations + +import os +import subprocess +import sys +from pathlib import Path + +# manifest_to_serve_args.py imports yaml; fail fast with a clear message +try: + import yaml # noqa: F401 +except ModuleNotFoundError: + print( + "PyYAML is required. Run: pip3 install pyyaml or uv pip install pyyaml", + file=sys.stderr, + ) + sys.exit(1) + + +def repo_root() -> Path: + root = Path(__file__).resolve().parent.parent + if not (root / "scripts" / "manifest_to_serve_args.py").exists(): + raise RuntimeError(f"repo root not found (expected {root})") + return root + + +def run_manifest_to_serve_args(manifest_path: Path) -> str: + script = repo_root() / "scripts" / "manifest_to_serve_args.py" + result = subprocess.run( + [sys.executable, str(script), str(manifest_path)], + capture_output=True, + text=True, + cwd=str(repo_root()), + ) + if result.returncode != 0: + raise AssertionError( + f"manifest_to_serve_args failed: {result.stderr or result.stdout}" + ) + return result.stdout.strip() + + +def test_sample_manifest_onnx() -> None: + root = repo_root() + manifest = root / "models" / "sample-manifest.yaml" + if not manifest.exists(): + raise FileNotFoundError(f"manifest not found: {manifest}") + out = run_manifest_to_serve_args(manifest) + assert "--backend onnx" in out, f"expected --backend onnx in {out!r}" + assert "--model-path" in out, f"expected --model-path in {out!r}" + assert "/models/" in out, f"expected /models/ in {out!r}" + + +def test_luxtts_manifest_helion() -> None: + root = repo_root() + manifest = root / "models" / "luxtts-manifest.yaml" + if not manifest.exists(): + raise FileNotFoundError(f"manifest not found: {manifest}") + out = run_manifest_to_serve_args(manifest) + assert "--backend helion" in out, f"expected --backend helion in {out!r}" + assert "--helion-model luxtts" in out, f"expected --helion-model luxtts in {out!r}" + assert "--helion-addr" in out, f"expected --helion-addr in {out!r}" + + +def main() -> int: + root = repo_root() + os.chdir(root) + try: + test_sample_manifest_onnx() + test_luxtts_manifest_helion() + print("manifest_to_serve_args tests passed") + return 0 + except Exception as e: + print(f"FAIL: {e}", file=sys.stderr) + return 1 + + +if __name__ == "__main__": + sys.exit(main())