diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000..3cc6e61 --- /dev/null +++ b/.dockerignore @@ -0,0 +1,10 @@ +# Keep the Docker build context small — the repo root is ~1.8G otherwise. +# The Dockerfiles (deploy/Dockerfile.base, .nvidia) only COPY constraints.txt, +# serve_requirements.txt, requirements-nvidia.txt and src/ — nothing below is +# needed in the build context. +.venv/ +.git/ +node_modules/ +**/__pycache__/ +**/*.pyc +*.log diff --git a/.gitignore b/.gitignore index 92a7c35..72491bd 100644 --- a/.gitignore +++ b/.gitignore @@ -179,6 +179,42 @@ cython_debug/ # Helm **/charts/*.tgz +# Terraform +# Local .terraform directories +**/.terraform/* + +# .tfstate files +*.tfstate +*.tfstate.* + +# Crash log files +crash.log +crash.*.log + +# Exclude all .tfvars files, which are likely to contain sensitive data, such as +# password, private keys, and other secrets. These should not be part of version +# control as they are data points which are potentially sensitive and subject +# to change depending on the environment. +*.tfvars +*.tfvars.json + +# Ignore override files as they are usually used to override resources locally and so +# are not checked in +override.tf +override.tf.json +*_override.tf +*_override.tf.json + +# Include override files you do wish to add to version control using negated pattern +# !example_override.tf + +# Include tfplan files to ignore the plan output of command: terraform plan -out=tfplan +# example: *tfplan* + +# Ignore CLI configuration files +.terraformrc +terraform.rc + # Others .trash docs diff --git a/config/client_aux.yaml b/config/client_aux.yaml index 1ad2f08..fac9d02 100644 --- a/config/client_aux.yaml +++ b/config/client_aux.yaml @@ -15,17 +15,11 @@ senders: args: sample_rate: 16000 frame_duration: 0.030 - text: - name: text - topic: question - args: - sample_rate: 16000 - frame_duration: 0.030 hooks: - text: - name: text - topics: [question, answer] + token: + name: token + topics: [token] audio: name: audio topics: [audio] @@ -48,10 +42,34 @@ modules: # block_duration: ${senders.audio.args.frame_duration} # tag: # name: tag + stt: + name: stt + args: + language: "en" + block_duration: 0.030 + logging: INFO + tag: + name: tag + logging: INFO + emo: + name: emo + args: + block_duration: 0.030 + eag: + name: eag + qag: + name: qag rag: name: rag args: language: en tone: formal + response_format: paragraph + max_length: 1024 + logging: INFO tts: name: tts + logging: INFO + gesture: + name: gesture + logging: INFO diff --git a/config/client_text.yaml b/config/client_text.yaml index 5771e18..49f735e 100644 --- a/config/client_text.yaml +++ b/config/client_text.yaml @@ -8,9 +8,9 @@ senders: topic: question hooks: - text: - name: text - topics: [rag_response] + token: + name: token + topics: [token] modules: rag: diff --git a/constraints.txt b/constraints.txt new file mode 100644 index 0000000..835369d --- /dev/null +++ b/constraints.txt @@ -0,0 +1,12 @@ +# Global pip constraints for all HuRI images. Pass to EVERY `pip install` via +# `-c /app/constraints.txt` so a later requirements file cannot silently override +# these versions (which is exactly how protobuf got clobbered before). +# +# protobuf: Ray Serve 2.55's ray/serve/_private/config.py:_proto_to_dict() reads +# FieldDescriptor.label, an attribute protobuf's upb (C) backend removed at >=5.0. +# With protobuf 5/6/7 every Serve replica crashes in __init__ with +# AttributeError: 'FieldDescriptor' object has no attribute 'label' +# 4.25.8 is the floor onnxruntime requires (>=4.25.8) and stays <5 for +# google-api-core / googleapis-common-protos / proto-plus / opentelemetry-proto / +# google-cloud-storage, so it satisfies the whole dependency graph. +protobuf==4.25.8 diff --git a/deploy/Dockerfile.amd b/deploy/Dockerfile.amd index 786094b..5d29277 100644 --- a/deploy/Dockerfile.amd +++ b/deploy/Dockerfile.amd @@ -38,8 +38,11 @@ ENV LD_LIBRARY_PATH="${ROCM_PATH}/lib:${LD_LIBRARY_PATH}" USER ray +# constraints.txt pins protobuf==4.25.8 so dep-resolving installs below can't +# bump protobuf past 4.x (breaks Ray Serve's _proto_to_dict). See the file. +COPY constraints.txt /app COPY serve_requirements.txt /app -RUN pip install --no-cache-dir -r serve_requirements.txt +RUN pip install --no-cache-dir -c /app/constraints.txt -r serve_requirements.txt # 1. AMD's PyTorch built for ROCm (NOT the PyPI one — it's built for ROCm 6.2 and will silently break) ARG ROCM_VERSION=7.2 @@ -69,11 +72,11 @@ RUN apt-get update && apt-get install -y unzip curl \ USER ray # 4. faster-whisper -RUN pip install --no-cache-dir faster-whisper +RUN pip install --no-cache-dir -c /app/constraints.txt faster-whisper # 5. RAG / LLM extras (httpx, qdrant-client, sentence-transformers, …) # Installed last so the ROCm torch wheel installed above is the resolved one. COPY requirements-amd.txt /app -RUN pip install --no-cache-dir -r requirements-amd.txt +RUN pip install --no-cache-dir -c /app/constraints.txt -r requirements-amd.txt COPY src /app/src diff --git a/deploy/Dockerfile.base b/deploy/Dockerfile.base index f46ac35..0095aaa 100644 --- a/deploy/Dockerfile.base +++ b/deploy/Dockerfile.base @@ -10,7 +10,10 @@ build-essential \ USER ray +# constraints.txt pins protobuf==4.25.8 so serve_requirements.txt's deps can't +# bump protobuf past 4.x (breaks Ray Serve's _proto_to_dict). See the file. +COPY constraints.txt /app COPY serve_requirements.txt /app -RUN pip install --no-cache-dir -r serve_requirements.txt +RUN pip install --no-cache-dir -c /app/constraints.txt -r serve_requirements.txt COPY src /app/src diff --git a/deploy/Dockerfile.nvidia b/deploy/Dockerfile.nvidia index f0e6e4d..90fa791 100644 --- a/deploy/Dockerfile.nvidia +++ b/deploy/Dockerfile.nvidia @@ -4,14 +4,18 @@ WORKDIR /app # Full CUDA 12.1 dependency stack (CosyVoice2, faster-whisper, TensorRT, …). # PyTorch cu121 wheels live on the PyTorch index; TensorRT wheels on the NGC index. +# constraints.txt pins protobuf==4.25.8 (see the file). Passed to every pip +# install below so serve_requirements.txt's deps (qdrant-client, onnxruntime — +# no protobuf upper bound) can't bump protobuf past 4.x and break Ray Serve. +COPY constraints.txt /app COPY requirements-nvidia.txt /app -RUN pip install --no-cache-dir \ +RUN pip install --no-cache-dir -c /app/constraints.txt \ --extra-index-url https://download.pytorch.org/whl/cu121 \ --extra-index-url https://pypi.ngc.nvidia.com \ -r requirements-nvidia.txt COPY serve_requirements.txt /app -RUN pip install --no-cache-dir -r serve_requirements.txt +RUN pip install --no-cache-dir -c /app/constraints.txt -r serve_requirements.txt USER root diff --git a/deploy/examples/cloud/.terraform.lock.hcl b/deploy/examples/cloud/.terraform.lock.hcl new file mode 100644 index 0000000..ede0f64 --- /dev/null +++ b/deploy/examples/cloud/.terraform.lock.hcl @@ -0,0 +1,61 @@ +# This file is maintained automatically by "terraform init". +# Manual edits may be lost in future updates. + +provider "registry.terraform.io/hashicorp/google" { + version = "7.37.0" + hashes = [ + "h1:Q5CT8lbZiCt7zS1LKo33hPIMQcWC3eUIjugf8XXVb/U=", + "zh:07373454562956f666099852978eed017a26f514535dc56409427662199d4a76", + "zh:168a3c56e35e215e40eb5bbfa4fd96690190eace4dddfd3994f765b6278cf1d8", + "zh:1af3aa69b8144699a21692bb497efd79153b4f25a5260a03bde508950fe9a10f", + "zh:274434b34327142196ae046fc397a4737c4826b7933f33227c72ad551f47e321", + "zh:28177baac1533e051d4ef7a7d44f468b1ade85c8eb25127141ab20c63bb38e6b", + "zh:35c50d8aefca5141732627eb4d469a0243de75da89fbbc93128672daabf3951c", + "zh:66b0ebf74e5d785a3426b948ee0922c2024568d953d2fd0dff6fa1cb1673365c", + "zh:8490723876b95e2c0cc87ead3dd83bc0b68af8028f6f1965469966826b3d8fd4", + "zh:9d76b3474f60c2803db3e4dfd41322ad6d47ca22da7f59eb5c26377452b5347c", + "zh:abfb9ed4282491b84761d206c20840db7b2b7749c279e3ab860e982c201eae27", + "zh:e3f9e865e0cffdb7e461c79813164315bb41949200f7b279aa81cc0f731aa0c8", + "zh:f569b65999264a9416862bca5cd2a6177d94ccb0424f3a4ef424428912b9cb3c", + ] +} + +provider "registry.terraform.io/hashicorp/helm" { + version = "3.2.0" + hashes = [ + "h1:ad8XdNRv+trfXgmu0tdI1YeFp7WSO5g9PI3QFOxRR8Y=", + "zh:2f1d55bf4e6a9c2629dfd3b162a05632f2e251bf6083d8613f38af3d51cea553", + "zh:33b378f15d39c2050a9272f6d5e8437f162972c93244af6df6de54ba2b0a416c", + "zh:501d7e7b3d42f5b40a6e5e979d5a6a4d67eba0fb37d786e2c3d7186e742dd557", + "zh:52e430da694bad4a06d049aa574d4a7a2b4e11c47d7bab637131068cc1160593", + "zh:77fb8ecaa27f4218177917bb3865551b058b92193408fa20366c45a529f0ac94", + "zh:b114d6ea5ab4486dc26b83cd595f0af820b3d66d0d1cf81975dfe7baa84419e2", + "zh:b6729f6f32fab90945e3cb4ba0268b73262dad0d14c1d71514ac767bac644595", + "zh:b9ff1756e698e1d3bdb0c2605ee31794f15d8275e2f8817a2bf66f3272ed9362", + "zh:be4791496afea715783f0efb65d21c99d3dc84ddb94ba4868f3afee2d50e71ef", + "zh:c8f66bd76991d7521a805cce3000a5226b1b18baf1d6b63a03d38e09208a564b", + "zh:e2bdd80b8956307f7307f90ffab0e91169bb83babeb540c712239dbc5c09916f", + "zh:f569b65999264a9416862bca5cd2a6177d94ccb0424f3a4ef424428912b9cb3c", + "zh:fcda8394edc50d2a2132534edf8655cb25dc76f55937e859b9ffb5f787430f37", + ] +} + +provider "registry.terraform.io/hashicorp/kubernetes" { + version = "3.2.0" + hashes = [ + "h1:0rX4RXNq0jqXm0yMzrtZd/1W9QrDO0EOC+zkfDNxAfQ=", + "zh:2e33acc20154d96ce5b3ab6d5fa0407403759a0852c63276baed0bbef4dbf1d4", + "zh:38721fee7fa1857414942040291d930b3c3f2a979845a2ff66289a73ad9f17ff", + "zh:38ec0cb28383f0e50065a98e215f40edbfb93c202ad3140afd1590a93f1965aa", + "zh:74adb0c844e49cecda97869022c3dfd7929e532ccf7cf9dff6eee87255fa0a54", + "zh:843f70f10b296eaa9f847629ceb4c4b6439cabf13be722d81afadf45ced859fc", + "zh:b004807190af53c2f5236847e7b6ec3e086beba268c8a027fb70aaf93b75f2ae", + "zh:b8459ab1fe6d7cceef20b9c47cb5e94e041fd92aca808ff80381073bfa440042", + "zh:c8529299f1d92d20ec65e80f4d5d134f21e699a1ead40a63bd66163cbca05960", + "zh:c98a3dd8884abf3dc0f6d4eb9129e2d1849089ef17a55b90a73159d06c74dc70", + "zh:d1bc7c68bb1a2abbcda03bc37f13b61a149f4f86d194849400d614873445f749", + "zh:d97f761796c306fb15d77ada4ccb2413a721f0831a7e91e436d9db1ce8ee9e26", + "zh:eecf88b257b564a87a9b5c56a6bc10a5340c4ecec6ffd75f3a179e9e64e72107", + "zh:f569b65999264a9416862bca5cd2a6177d94ccb0424f3a4ef424428912b9cb3c", + ] +} diff --git a/deploy/examples/cloud/artifact_registry.tf b/deploy/examples/cloud/artifact_registry.tf new file mode 100644 index 0000000..a3c7e0f --- /dev/null +++ b/deploy/examples/cloud/artifact_registry.tf @@ -0,0 +1,34 @@ +# Private Artifact Registry holding the website backend image. +# +# This repository ALREADY EXISTS — Terraform should adopt it rather than try to +# recreate it. After `terraform init`, import it once: +# +# terraform import google_artifact_registry_repository.website \ +# projects/${PROJECT}/locations/${LOCATION}/repositories/${REPO_ID} +# +# (LOCATION = var.website_image_location, REPO_ID = var.website_image_repo.) +# The image itself is built and pushed out-of-band from +# HuRI_website_demo/backend/Dockerfile. +resource "google_artifact_registry_repository" "website" { + location = var.website_image_location + repository_id = var.website_image_repo + format = "DOCKER" + description = "HuRI website backend images" + + # Don't let an apply delete an existing shared registry. + lifecycle { + prevent_destroy = true + } +} + +data "google_project" "current" {} + +# GKE nodes pull with the node pool's service account (the default Compute +# Engine SA here). Grant it read access to the repo so the kubelet can pull. +resource "google_artifact_registry_repository_iam_member" "website_puller" { + project = var.project_id + location = google_artifact_registry_repository.website.location + repository = google_artifact_registry_repository.website.repository_id + role = "roles/artifactregistry.reader" + member = "serviceAccount:${data.google_project.current.number}-compute@developer.gserviceaccount.com" +} diff --git a/deploy/examples/cloud/backend.tf b/deploy/examples/cloud/backend.tf new file mode 100644 index 0000000..7fa131f --- /dev/null +++ b/deploy/examples/cloud/backend.tf @@ -0,0 +1,16 @@ +# Remote Terraform state in GCS (versioned, IAM-restricted). Holds the cluster +# state — and, because Terraform decrypts Secret Manager values at apply time, +# secret material too. Keep the bucket private. +# +# The bucket must EXIST before `terraform init` (chicken-and-egg: backends can't +# create their own bucket and can't interpolate variables). It is supplied at +# init time, so only the static prefix is committed here: +# +# terraform init -backend-config="bucket=" +# +# See gcp_steps.md §0 for creating the bucket. +terraform { + backend "gcs" { + prefix = "huri-cloud" + } +} diff --git a/deploy/examples/cloud/embedding.tf b/deploy/examples/cloud/embedding.tf new file mode 100644 index 0000000..269520e --- /dev/null +++ b/deploy/examples/cloud/embedding.tf @@ -0,0 +1,153 @@ +# Embedding server — serves bge-large-en-v1.5 over an OpenAI-compatible +# /v1/embeddings endpoint, in-cluster, replacing the old external +# embedding.huri.lan host. +# +# Why in-cluster and not a managed GCP/Vertex endpoint: +# RAG retrieval embeds the query and searches Qdrant by vector similarity, so +# query vectors MUST come from the SAME model that produced the stored vectors +# at ingestion time. Those were created with the bge-large-en-v1.5 GGUF +# (Q4_K_M, 1024-dim). Vertex AI only serves its own embedding models (different +# vectors and dimensions), which would be incompatible without re-ingesting the +# whole corpus. So we run the exact same GGUF here via llama.cpp's server, whose +# /v1/embeddings path and request shape match RemoteEmbedder (ingestion.py) and +# RAGHandle._embed (rag.py). +# +# The RAGHandle (values-gcp.yaml user_config) points embedding_url at +# http://embedding.huri.svc.cluster.local:8080. llama.cpp ignores the request's +# "model" field and uses the loaded GGUF, so embedding_model is informational. + +resource "kubernetes_deployment_v1" "embedding" { + metadata { + name = "embedding" + namespace = "huri" + labels = { app = "embedding" } + } + + spec { + replicas = 1 + selector { + match_labels = { app = "embedding" } + } + # The model cache is a ReadWriteOnce PVC, so the old pod must release it + # before the new one mounts it — a rolling update would deadlock. + strategy { + type = "Recreate" + } + template { + metadata { + labels = { app = "embedding" } + } + spec { + # CPU-only — bge-large (~335M params, ~200MB at Q4_K_M) runs fine on the + # system pool. The GPU pool is reserved for the Ray workers. + node_selector = { + "huri.io/system" = "true" + } + + container { + name = "embedding" + # Pinned by digest for reproducibility; the tag is kept for readability. + # Re-resolve with: docker buildx imagetools inspect ghcr.io/ggml-org/llama.cpp:server + # IfNotPresent (also the default for a non-:latest tag) reuses the image + # from the node cache across restarts; only re-pulled if the node is replaced. + image = "ghcr.io/ggml-org/llama.cpp:server@sha256:9e2ab5775c2e9cb52ac611e506ea5b873247033aeba54d71f20d71bf7f4e219d" + # -hf pulls the GGUF from Hugging Face on startup into HF_HOME. + # --pooling cls matches how bge models are meant to be pooled. + # -c 512 covers bge's 512-token context (ingestion chunks stay under + # this — see the rag-ingestion-chunk-size note). + args = [ + "-hf", var.embedding_hf_repo, + "--embedding", + "--pooling", "cls", + "--host", "0.0.0.0", + "--port", "8080", + "-c", "512", + ] + + env { + name = "HF_HOME" + value = "/models" + } + + port { + container_port = 8080 + } + + volume_mount { + name = "models" + mount_path = "/models" + } + + resources { + requests = { cpu = "250m", memory = "1Gi" } + limits = { cpu = "2", memory = "2Gi" } + } + + # The model is downloaded at startup, so give it generous startup slack + # before the readiness gate kicks in. + readiness_probe { + http_get { + path = "/health" + port = 8080 + } + initial_delay_seconds = 30 + period_seconds = 10 + failure_threshold = 30 + } + } + + # Persistent model cache so the GGUF is downloaded once and survives pod + # restarts/reschedules (emptyDir would re-pull it every time). + volume { + name = "models" + persistent_volume_claim { + claim_name = kubernetes_persistent_volume_claim_v1.embedding_models.metadata[0].name + } + } + } + } + } + + # The huri namespace is created by the qdrant release; reuse it. + depends_on = [helm_release.qdrant] +} + +resource "kubernetes_persistent_volume_claim_v1" "embedding_models" { + metadata { + name = "embedding-models" + namespace = "huri" + } + # standard-rwo binds WaitForFirstConsumer, so the PVC stays Pending until the + # pod mounts it. Don't block apply waiting for a bind that can't happen yet. + wait_until_bound = false + + spec { + access_modes = ["ReadWriteOnce"] + storage_class_name = "standard-rwo" + resources { + requests = { storage = "1Gi" } # models shouldn't be that big + } + } + + # The huri namespace is created by the qdrant release; reuse it. + depends_on = [helm_release.qdrant] +} + +resource "kubernetes_service_v1" "embedding" { + metadata { + name = "embedding" + namespace = "huri" + } + spec { + selector = { app = "embedding" } + port { + port = 8080 + target_port = 8080 + } + type = "ClusterIP" + } + + # The huri namespace is created by the qdrant release; reuse it. Without this + # the service races ahead of namespace creation ("namespaces huri not found"). + depends_on = [helm_release.qdrant] +} diff --git a/deploy/examples/cloud/gke.tf b/deploy/examples/cloud/gke.tf new file mode 100644 index 0000000..c2fe20c --- /dev/null +++ b/deploy/examples/cloud/gke.tf @@ -0,0 +1,161 @@ +resource "google_container_cluster" "primary" { + name = var.cluster_name + # Zonal (single zone) rather than regional: a regional cluster replicates every + # node pool across all 3 zones, tripling SSD/CPU/GPU usage and blowing past + # free-trial quotas. Zonal keeps one node per pool. + location = var.zone + + # Allow `terraform destroy` to delete the cluster. The provider defaults this + # to true, which blocks teardown — fine to disable for a throwaway trial stack. + deletion_protection = false + + # We can't create a cluster with no node pool defined, but we want to only use + # separately managed node pools. So we create the smallest possible default + # node pool and immediately delete it. + remove_default_node_pool = true + initial_node_count = 1 + + # This default pool is created (then removed) during cluster creation, so it + # must also fit the free-trial SSD quota. Without this it uses GKE's default + # 100 GB pd-balanced (SSD) disk. pd-standard keeps it off SSD_TOTAL_GB. + node_config { + disk_type = "pd-standard" + disk_size_gb = 50 + } + + network = google_compute_network.vpc.name + subnetwork = google_compute_subnetwork.subnet.name + + ip_allocation_policy { + cluster_secondary_range_name = "pods" + services_secondary_range_name = "services" + } + + release_channel { + channel = "REGULAR" + } + + # Enable workload identity + workload_identity_config { + workload_pool = "${var.project_id}.svc.id.goog" + } +} + +# General purpose node pool for Ray Head and other system components +resource "google_container_node_pool" "system_nodes" { + name = "system-pool" + location = var.zone + cluster = google_container_cluster.primary.name + node_count = 1 + + management { + auto_repair = true + auto_upgrade = true + } + + node_config { + machine_type = "e2-standard-4" + + # pd-standard (HDD) boot disk counts against DISKS_TOTAL_GB, not the much + # smaller SSD_TOTAL_GB quota that pd-balanced/pd-ssd consume. Keeps the + # free trial from tripping its 250 GB SSD limit. + disk_type = "pd-standard" + disk_size_gb = 50 + + # GKE reserves the node-role.kubernetes.io/ prefix and rejects it as a node + # label, so we use a custom key to pin system pods (Ray head, Qdrant, + # embedding, LiteLLM) here. Must match the nodeSelectors in values-gcp.yaml, + # qdrant.tf, embedding.tf and litellm.tf. + labels = { + "huri.io/system" = "true" + } + + # Removed taint to allow init jobs and other system pods to schedule easily + + workload_metadata_config { + mode = "GKE_METADATA" + } + } +} + +# GPU node pool for NVIDIA workers +resource "google_container_node_pool" "gpu_nodes" { + name = "gpu-pool" + location = var.zone + cluster = google_container_cluster.primary.name + + # Layer ③ — GKE provisions/removes GPU VMs as pending GPU pods appear/clear. + # Only fires once KubeRay (layer ②, ray.enableInTreeAutoscaling in values) + # marks worker pods Pending for lack of GPU. This is a zonal pool + # (location = zone), so these counts are the actual node counts (no per-zone + # multiplier). + initial_node_count = var.gpu_min_nodes + + # Spread the pool across the region's zones (when gpu_node_zones is set) so the + # autoscaler isn't trapped in one zone's GPU stock. Empty list -> nodes stay in + # the cluster zone (var.zone). All listed zones must be within var.region. + node_locations = length(var.gpu_node_zones) > 0 ? var.gpu_node_zones : null + + autoscaling { + min_node_count = var.gpu_min_nodes + max_node_count = var.gpu_max_nodes + + # ANY = provision in whichever listed zone currently has GPU capacity, rather + # than the default BALANCED round-robin that keeps hitting the exhausted zone. + # This is the knob that actually dodges RESOURCE_POOL_EXHAUSTED. + location_policy = "ANY" + } + + management { + auto_repair = true + auto_upgrade = true + } + + node_config { + # machine_type and gpu_type must be compatible: N1 for T4, G2 for L4, + # A2 for A100 (see variables.tf). guest_accelerator takes the *accelerator* + # type (nvidia-tesla-t4, ...), never a machine type. + machine_type = var.gpu_machine_type + + guest_accelerator { + type = var.gpu_type + count = var.gpu_count + } + + # G2 (L4) and A2 (A100) machine types reject pd-standard (HDD) boot disks — + # unlike the N1/e2 system pools, they only accept pd-balanced/pd-ssd/hyperdisk. + # So this pool must use SSD-backed pd-balanced (the cheapest G2-compatible + # option) and does count against SSD_TOTAL_GB: 150 GB * gpu_max_nodes. + # + # 150 GB (not 50): the huri:nvidia serve image is large, and a 50 GB disk + # leaves only ~17 GB allocatable ephemeral-storage after GKE reservations — + # kubelet then evicts the Ray worker for ephemeral-storage pressure before it + # finishes init. 150 GB gives ~110 GB ephemeral headroom. + disk_type = "pd-balanced" + disk_size_gb = 150 + + # Automatically install NVIDIA drivers + # This is recommended for GKE + # metadata = { + # "install-nvidia-driver" = "true" + # } + + labels = { + gpu = "nvidia" + } + + oauth_scopes = [ + "https://www.googleapis.com/auth/cloud-platform" + ] + + workload_metadata_config { + mode = "GKE_METADATA" + } + } + + lifecycle { + # The autoscaler owns the running node count; don't let Terraform reset it + # back to the initial count on every apply. + ignore_changes = [node_count] + } +} diff --git a/deploy/examples/cloud/litellm-config.yaml b/deploy/examples/cloud/litellm-config.yaml new file mode 100644 index 0000000..c957a91 --- /dev/null +++ b/deploy/examples/cloud/litellm-config.yaml @@ -0,0 +1,31 @@ +# LiteLLM proxy config (rendered by templatefile() from litellm.tf). +# +# Exposes one OpenAI-compatible model alias, "huri-fast", that forwards to +# Anthropic's OpenAI-compatible API. The RAGHandle calls +# {llm_url}/v1/chat/completions with model="huri-fast" (llm_provider: vllm), +# and LiteLLM proxies to Anthropic (Claude). +# +# The Anthropic API key is injected via the MISTRAL_API_KEY env var (name kept +# for compatibility) from the litellm-secrets Kubernetes Secret — never stored +# in this file or in git. + +model_list: + - model_name: huri-fast + litellm_params: + # openai/ provider = generic OpenAI-compatible endpoint; api_base points + # it at Anthropic's OpenAI-compatible /v1 API instead of api.openai.com. + # LiteLLM POSTs to {api_base}/chat/completions, so the /v1 is required. + model: openai/${model} + api_base: https://api.anthropic.com/v1 + api_key: os.environ/MISTRAL_API_KEY + +litellm_settings: + # Drop params a given model doesn't support instead of erroring, so the + # RAGHandle's OpenAI-style request body is always accepted. + drop_params: true + +general_settings: + # Internal ClusterIP service — no virtual-key auth. Access is restricted at + # the network layer (ClusterIP, same namespace), so requests need no bearer + # token and the Mistral key stays solely in the Secret. + master_key: null diff --git a/deploy/examples/cloud/litellm.tf b/deploy/examples/cloud/litellm.tf new file mode 100644 index 0000000..79fb0fe --- /dev/null +++ b/deploy/examples/cloud/litellm.tf @@ -0,0 +1,134 @@ +# LiteLLM proxy — a fast, OpenAI-compatible LLM endpoint backed by Mistral. +# +# The RAGHandle (values-gcp.yaml user_config) points llm_url at +# http://litellm.huri.svc.cluster.local:4000 with llm_provider: vllm and +# llm_model: huri-fast. LiteLLM exposes /v1/chat/completions (with streaming), +# which is exactly the path the handler builds, and forwards to Mistral's +# OpenAI-compatible API. + +# Mistral API key, read from GCP Secret Manager (secrets.tf), never committed. +resource "kubernetes_secret_v1" "litellm" { + metadata { + name = "litellm-secrets" + namespace = "huri" + } + data = { + # trimspace() guards against a trailing newline/CRLF in the Secret Manager + # value (e.g. created with `echo` instead of `printf`). LiteLLM puts this + # key straight into the outbound Authorization header, and aiohttp rejects + # any header value containing \r or \n ("header injection"), which surfaces + # as a 500 back to the RAG caller. + MISTRAL_API_KEY = trimspace(data.google_secret_manager_secret_version.mistral_api_key.secret_data) + } + + # The huri namespace is created by the qdrant release; reuse it. + depends_on = [helm_release.qdrant] +} + +# model_list config, with the Mistral model name templated in. +resource "kubernetes_config_map_v1" "litellm" { + metadata { + name = "litellm-config" + namespace = "huri" + } + data = { + "config.yaml" = templatefile("${path.module}/litellm-config.yaml", { + model = var.llm_model + }) + } + + depends_on = [helm_release.qdrant] +} + +resource "kubernetes_deployment_v1" "litellm" { + metadata { + name = "litellm" + namespace = "huri" + labels = { app = "litellm" } + } + + spec { + replicas = 1 + selector { + match_labels = { app = "litellm" } + } + template { + metadata { + labels = { app = "litellm" } + } + spec { + # CPU-only proxy — keep it off the GPU pool. + node_selector = { + "huri.io/system" = "true" + } + + container { + name = "litellm" + # Pinned by digest for reproducibility; the tag is kept for readability. + # Re-resolve with: docker buildx imagetools inspect ghcr.io/berriai/litellm:main-stable + image = "ghcr.io/berriai/litellm:main-stable@sha256:8f3517476b8293ccc1c3b0da7940d8dabbfe64b013454113f5a3cb68265c439b" + args = ["--config", "/etc/litellm/config.yaml", "--port", "4000"] + + port { + container_port = 4000 + } + + env_from { + secret_ref { + name = kubernetes_secret_v1.litellm.metadata[0].name + } + } + + volume_mount { + name = "config" + mount_path = "/etc/litellm" + read_only = true + } + + resources { + # The litellm:main-stable image loads many provider SDKs at startup + # and OOMs under a 1Gi cap before it can serve /health. 2Gi is the + # observed safe ceiling for the proxy. + requests = { cpu = "100m", memory = "512Mi" } + limits = { cpu = "500m", memory = "2Gi" } + } + + readiness_probe { + http_get { + path = "/health/liveliness" + port = 4000 + } + initial_delay_seconds = 10 + period_seconds = 10 + } + } + + volume { + name = "config" + config_map { + name = kubernetes_config_map_v1.litellm.metadata[0].name + } + } + } + } + } +} + +resource "kubernetes_service_v1" "litellm" { + metadata { + name = "litellm" + namespace = "huri" + } + spec { + selector = { app = "litellm" } + port { + port = 4000 + target_port = 4000 + } + type = "ClusterIP" + } + + # The huri namespace is created by the qdrant release; reuse it. Without this + # the service races ahead of namespace creation ("namespaces huri not found"). + depends_on = [helm_release.qdrant] +} diff --git a/deploy/examples/cloud/main.tf b/deploy/examples/cloud/main.tf new file mode 100644 index 0000000..cce8c6b --- /dev/null +++ b/deploy/examples/cloud/main.tf @@ -0,0 +1,58 @@ +provider "google" { + project = var.project_id + region = var.region +} + +# Get GKE cluster credentials +data "google_client_config" "default" {} + +data "google_container_cluster" "primary" { + name = google_container_cluster.primary.name + location = var.zone + depends_on = [google_container_cluster.primary] +} + +provider "kubernetes" { + host = "https://${data.google_container_cluster.primary.endpoint}" + token = data.google_client_config.default.access_token + cluster_ca_certificate = base64decode(data.google_container_cluster.primary.master_auth[0].cluster_ca_certificate) +} + +provider "helm" { + kubernetes = { + host = "https://${data.google_container_cluster.primary.endpoint}" + token = data.google_client_config.default.access_token + cluster_ca_certificate = base64decode(data.google_container_cluster.primary.master_auth[0].cluster_ca_certificate) + } +} + +# Install KubeRay Operator +resource "helm_release" "kuberay_operator" { + name = "kuberay-operator" + repository = "https://ray-project.github.io/kuberay-helm/" + chart = "kuberay-operator" + namespace = "ray-system" + create_namespace = true +} + +# Install HuRI App +resource "helm_release" "huri" { + name = "huri" + chart = "../../../helm" + namespace = "huri" + create_namespace = true + + values = [ + file("${path.module}/values-gcp.yaml") + ] + + depends_on = [ + helm_release.kuberay_operator, + google_container_node_pool.gpu_nodes, + google_container_node_pool.system_nodes, + # RAGHandle reads these at startup via values-gcp.yaml user_config. + helm_release.qdrant, + kubernetes_service_v1.litellm, + kubernetes_service_v1.embedding + ] +} diff --git a/deploy/examples/cloud/outputs.tf b/deploy/examples/cloud/outputs.tf new file mode 100644 index 0000000..be3490b --- /dev/null +++ b/deploy/examples/cloud/outputs.tf @@ -0,0 +1,29 @@ +output "kubernetes_cluster_name" { + value = google_container_cluster.primary.name + description = "GKE Cluster Name" +} + +output "kubernetes_cluster_host" { + value = google_container_cluster.primary.endpoint + description = "GKE Cluster Host" +} + +output "region" { + value = var.region + description = "GCP region" +} + +output "project_id" { + value = var.project_id + description = "GCP Project ID" +} + +output "qdrant_url" { + value = "http://qdrant.huri.svc.cluster.local:6333" + description = "In-cluster Qdrant endpoint (matches values-gcp.yaml)" +} + +output "llm_url" { + value = "http://litellm.huri.svc.cluster.local:4000" + description = "In-cluster LiteLLM (Gemini) endpoint (matches values-gcp.yaml)" +} diff --git a/deploy/examples/cloud/qdrant.tf b/deploy/examples/cloud/qdrant.tf new file mode 100644 index 0000000..9577198 --- /dev/null +++ b/deploy/examples/cloud/qdrant.tf @@ -0,0 +1,41 @@ +# Qdrant vector database — backs the HuRI RAG store. +# +# Deployed in-cluster so the RAGHandle reaches it over the cluster network at +# http://qdrant.huri.svc.cluster.local:6333 (see values-gcp.yaml user_config), +# replacing the old external qdrant.pommier.lan host. make_qdrant_client() +# derives port 6333 for non-https URLs, so no app change is needed. + +resource "helm_release" "qdrant" { + name = "qdrant" + repository = "https://qdrant.github.io/qdrant-helm" + chart = "qdrant" + # Pinned so the chart (and the Qdrant image it carries) don't float on apply. + # Chart 1.18.2 → appVersion v1.18.1. Bump deliberately, not implicitly. + version = "1.18.2" + namespace = "huri" + create_namespace = true + + # Single-node, persistent. The chart's StatefulSet exposes a ClusterIP + # Service named "qdrant" on the REST (6333) and gRPC (6334) ports. + values = [ + yamlencode({ + replicaCount = 1 + persistence = { + size = "10Gi" + storageClassName = "standard-rwo" + } + service = { + type = "ClusterIP" + } + # Keep Qdrant on the system pool (CPU only); the GPU pool is reserved + # for the Ray workers. + nodeSelector = { + "huri.io/system" = "true" + } + resources = { + requests = { cpu = "250m", memory = "512Mi" } + limits = { cpu = "1", memory = "2Gi" } + } + }) + ] +} diff --git a/deploy/examples/cloud/secrets.tf b/deploy/examples/cloud/secrets.tf new file mode 100644 index 0000000..9110634 --- /dev/null +++ b/deploy/examples/cloud/secrets.tf @@ -0,0 +1,12 @@ +# Secrets come from GCP Secret Manager, not from variables or env vars. Create +# the secret and add a version OUT OF BAND before applying (see README §1); +# Terraform only *reads* the latest version here, so the plaintext key is never +# a Terraform input and never touches git. +# +# Caveat: the decrypted value still lands in Terraform state, so keep state in a +# secured remote backend (GCS). The identity running Terraform needs +# roles/secretmanager.secretAccessor on this secret. +data "google_secret_manager_secret_version" "mistral_api_key" { + secret = "mistral_api_key" + # version defaults to "latest". +} diff --git a/deploy/examples/cloud/terraform.tfvars.example b/deploy/examples/cloud/terraform.tfvars.example new file mode 100644 index 0000000..93300ea --- /dev/null +++ b/deploy/examples/cloud/terraform.tfvars.example @@ -0,0 +1,40 @@ +project_id = "your-gcp-project-id" +region = "europe-west4" +# Zone for the (zonal) GKE cluster + node pools. Must be within `region`. +# A zonal cluster keeps SSD/CPU/GPU usage within free-trial quotas. +zone = "europe-west4-a" +cluster_name = "huri-ray-cluster" + +# GPU accelerator + the machine that hosts it (the two must be compatible): +# T4 -> gpu_type = "nvidia-tesla-t4", gpu_machine_type = "n1-standard-4" +# L4 -> gpu_type = "nvidia-l4", gpu_machine_type = "g2-standard-8" +# A100 40GB -> gpu_type = "nvidia-tesla-a100", gpu_machine_type = "a2-highgpu-1g" +# A100 80GB -> gpu_type = "nvidia-tesla-a100", gpu_machine_type = "a2-ultragpu-1g" +# A2 machines bundle their A100(s); gpu_count must match the machine (1 here). +gpu_type = "nvidia-tesla-v100" +gpu_machine_type = "n1-standard-8" +gpu_count = 1 + +# Multi-zone GPU pool: span every zone in `region` so the autoscaler (with +# location_policy=ANY, set in gke.tf) provisions the GPU node wherever stock is +# available instead of hard-failing on RESOURCE_POOL_EXHAUSTED in one zone. All +# zones must be within `region`. NOTE: gpu_min/max_nodes are PER ZONE — keep +# gpu_min_nodes=0 (default) or you hold one warm node in EACH listed zone. +# Leave empty/unset for single-zone behavior (pool stays in `zone`). +gpu_node_zones = ["europe-west4-a", "europe-west4-b", "europe-west4-c"] + +# Fast LLM (LiteLLM -> Mistral). Override the model if you like; any fast model. +llm_model = "ministral-14b-latest" + +# Website image registry (the "GCR"). HuRI owns the Artifact Registry repo; the +# website backend image is built/pushed into it and deployed by the website's own +# Terraform (HuRI_website_demo/deploy), which is also where Authelia + the website +# Deployment now live. +website_image_location = "europe-west4" +website_image_repo = "huri" + +# ----------------------------------------------------------------------------- +# Secrets are NOT set here and are NOT Terraform variables. The Mistral API key +# lives in GCP Secret Manager (secret id "mistral_api_key") and is read by +# secrets.tf. Create it once, out of band, before applying — see README §1. +# ----------------------------------------------------------------------------- diff --git a/deploy/examples/cloud/values-gcp.yaml b/deploy/examples/cloud/values-gcp.yaml new file mode 100644 index 0000000..5ae4817 --- /dev/null +++ b/deploy/examples/cloud/values-gcp.yaml @@ -0,0 +1,392 @@ +# GCP Specific Overrides for HuRI Ray Cluster + +image: + # Image for the Ray HEAD and any worker group without its own `image:` (CPU + # actors land here). Built from deploy/Dockerfile.base and pushed to the + # Artifact Registry BEFORE `terraform apply` — see ../../../gcp_steps.md §1c. + # The GPU worker group below overrides this with the nvidia image. + repository: europe-west4-docker.pkg.dev/hurigcp/huri/huri + tag: base-2.55.1@sha256:514da8e891d7f7fccd723ebd422db11ca12c55de153d6aa3b68fb010cf405689 + pullPolicy: IfNotPresent + +# The KubeRay operator is installed as its OWN Helm release +# (helm_release.kuberay_operator in main.tf), so this chart must not also bundle +# it as a sub-chart — otherwise the operator's cluster-scoped RBAC (e.g. +# raycronjob-editor-role) collides with the existing release's Helm ownership. +kuberay: + install: false + +ray: + # RayService.spec.rayClusterConfig.rayVersion — MUST be a non-null string or the + # KubeRay CRD rejects the object under Helm's server-side apply ("rayVersion: + # Invalid value: null"). Match the Ray version baked into the images + # (deploy/Dockerfile.base/.nvidia -> FROM rayproject/ray:2.55.1). + version: "2.55.1" + # Layer â‘¡ — let the Ray (KubeRay) autoscaler add GPU worker pods when Serve + # asks for replicas that don't fit on the current pods. Pairs with the + # workerGroups maxReplicas below and the GKE node-pool autoscaling in gke.tf + # (layer â‘¢). Leave unset/false in the local single-GPU values to stay pinned. + enableInTreeAutoscaling: true + # First-bring-up ceiling (ray.io/initializing-timeout). Must exceed a cold GPU + # node provision + driver install + nvidia image pull, or the RayService hits a + # TERMINAL failure state and stops recreating the cluster (see rayservice.yaml). + # Set to 24h (not 40m): during an L4 stockout the gpu-pool can sit "GCE out of + # resources" for hours. A short timeout makes the RayService self-destruct and + # tear down the head + the *standing Pending GPU worker* — which is the only + # thing that keeps the autoscaler asking GCE for capacity. A long timeout keeps + # that request alive so the node is grabbed the moment stock returns. Lower it + # back to ~40m once you're off scarce/on-demand GPU capacity. + initializingTimeout: "24h" + # Give a cold GPU worker time to come up before KubeRay declares the cluster + # unhealthy and recreates it. On GCP the gpu-pool scales 0→1, the node installs + # drivers, and the large nvidia image is pulled — easily several minutes. With + # the default (~5 min) thresholds the Serve deployments are still UPDATING when + # the controller gives up, so it deletes the RayCluster (head + worker + # terminate together) and the cycle repeats. 1800s/1200s comfortably covers a + # cold node provision + image pull. Paired with ray.io/initializing-timeout + # (20m) on the RayService for the very first bring-up. + serviceUnhealthySecondThreshold: 1800 + deploymentUnhealthySecondThreshold: 1200 + serveConfig: | + proxy_location: EveryNode + http_options: + host: 0.0.0.0 + port: 8000 + applications: + - name: huri-app + route_prefix: / + import_path: src.app:app + runtime_env: + env_vars: + RAY_COLOR_PREFIX: "1" + HURI_GESTURE_CONTEXT_SEC: "2.0" + HURI_GESTURE_MIN_CHUNK_SEC: "0.5" + HURI_VOICE_TRANSCRIPT: "You are a helpful assistant.<|endofprompt|>Instinct creates its own oppressors and bids us rise up against them." + HURI_STT_MODEL_PATH: /models/whisper/Systran/faster-whisper-base + # Concurrent transcription slots on the single STT replica. Each + # client's audio cache is per-session (in the STT module), so this + # only affects throughput, never isolation. Raise for more + # simultaneous speakers; it shares the one GPU slice below. + HURI_STT_NUM_WORKERS: "2" + # ────────────────────────────────────────────────────────────────── + # Autoscaling (layer â‘ : Ray Serve replicas). + # The autoscaling_config blocks below let Serve add/remove model + # replicas based on in-flight requests: + # desired = ceil(total_ongoing_requests / target_ongoing_requests) + # clamped to [min_replicas, max_replicas]. min_replicas:1 keeps a warm + # replica so the baseline never cold-starts. + # + # End-to-end GPU autoscaling is now wired across all three layers: + # â‘  these autoscaling_config blocks (Serve replicas) + # â‘¡ ray.enableInTreeAutoscaling + workerGroups maxReplicas (KubeRay + # scales GPU worker pods) + # â‘¢ gke.tf gpu_nodes autoscaling { } (GKE provisions GPU VMs) + # Flow: load rises → Serve wants replicas → no free GPU → KubeRay adds a + # worker pod → pod Pending → GKE adds a GPU node → replica runs. + # The remaining ceiling is GPU quota + budget (and per-zone node counts + # for the regional pool). Cold-start is minutes (VM + driver + model + # load), so min_replicas:1 stays warm; for bursts on the existing GPU, + # HURI_STT_NUM_WORKERS adds in-replica concurrency without waiting for a + # node. + # + # GPU PACKING NOTE — num_gpus below is a *logical* Ray fraction, not a + # memory reservation; it only controls how many actors Ray packs per + # card. Sized for the V100 16GB (n1-standard-8): each model's real + # footprint is a few GB (Whisper base <1GB, CosyVoice3 ~4GB, EMAGE ~3GB), + # far under 16GB, so the fractions are deliberately small to pack many + # replicas onto ONE card before autoscaling buys a second: + # baseline (1 each): STT 0.2 + TTS 0.1 + Gesture 0.1 = 0.4 + # one full V100 e.g.: STT*3 + TTS*2 + Gesture*2 = 1.0 + # Shrink further for more density — the binding limit is logical packing, + # not memory. (On a 16GB T4 you'd raise these back toward ~0.5/0.2/0.2.) + # ────────────────────────────────────────────────────────────────── + deployments: + # Stateful WebSocket ingress + per-session routing — single coordinator, + # not autoscaled. + - name: HuRI + num_replicas: 1 + ray_actor_options: + num_cpus: 1 + num_gpus: 0 + - name: STT + # Client isolation is per-session (audio cache lives in the STT + # module); this actor is stateless per call. HURI_STT_NUM_WORKERS + # gives in-replica concurrency; replicas add cross-actor parallelism + # and pack onto the A100 at 0.2 each (see packing note above). + max_ongoing_requests: 8 + autoscaling_config: + min_replicas: 1 + max_replicas: 4 # 0.2 GPU each → up to ~4 share one A100 + target_ongoing_requests: 4 + ray_actor_options: + num_cpus: 1 + num_gpus: 0.2 + resources: {"GPU_TYPE_NVIDIA": 0.2} + - name: RAGHandle + # CPU/network-bound (num_gpus: 0) — these replicas schedule on the + # existing nodes without needing GPU, so this one genuinely scales + # within the current cluster. Work is async LLM/embedding I/O. + max_ongoing_requests: 16 + autoscaling_config: + min_replicas: 1 + max_replicas: 4 + target_ongoing_requests: 8 + ray_actor_options: + num_cpus: 1 + num_gpus: 0 + # Removed AMD specific resource + user_config: + # In-cluster services provisioned by Terraform (qdrant.tf, litellm.tf). + qdrant_url: "http://qdrant.huri.svc.cluster.local:6333" + llm_url: "http://litellm.huri.svc.cluster.local:4000" + # In-cluster embedding server (embedding.tf): llama.cpp serving the + # same bge-large-en-v1.5 GGUF used at ingestion, so query vectors + # match the stored ones. /v1/embeddings is appended by the handler. + embedding_url: "http://embedding.huri.svc.cluster.local:8080" + embedding_model: "bge-large-en-v1.5-gguf-Q4_K_M" + # LiteLLM exposes an OpenAI-compatible /v1/chat/completions; "vllm" + # sends no auth header (the proxy is cluster-internal). "huri-fast" + # is the LiteLLM alias that forwards to Gemini. + llm_provider: "vllm" + llm_model: "huri-fast" + verify_ssl: true + - name: TTS + # Streaming synth runs per-utterance in a thread (run_in_executor), + # so one replica already serves several sessions; replicas pack onto + # the A100 at 0.1 each for more concurrent utterances. + autoscaling_config: + min_replicas: 1 + max_replicas: 4 + target_ongoing_requests: 8 + ray_actor_options: + num_cpus: 1 + num_gpus: 0.1 # logical GPU fraction for Ray packing (see note above) + resources: {"GPU_TYPE_NVIDIA": 0.1} + - name: GestureGeneration + # Sync infer() is auto-threaded by Serve (concurrent within a + # replica). Replicas pack onto the A100 at 0.1 each. + autoscaling_config: + min_replicas: 1 + max_replicas: 4 + target_ongoing_requests: 4 + ray_actor_options: + num_cpus: 1 + num_gpus: 0.1 + resources: {"GPU_TYPE_NVIDIA": 0.1} + +head: + serviceType: ClusterIP + # Ray autoscaler-v2 (enableInTreeAutoscaling) refuses to start unless it can + # detect the head's CPU, via either the num-cpus rayStartParam or a CPU + # resource limit. Without this the autoscaler sidecar crashloops with + # "failed to detect `CPU` resources for group headgroup", the head never goes + # ready, and the RayService hits its 20m init timeout -> terminal failure. + # num-cpus: "0" also keeps Ray tasks off the head (they belong on the GPU worker). + rayStartParams: + num-cpus: "0" + resources: + requests: + cpu: "1" + memory: "2Gi" + limits: + cpu: "2" + memory: "4Gi" + nodeSelector: + huri.io/system: "true" + tolerations: + - key: huri.io/system + operator: Exists + effect: NoSchedule + +workerGroups: + - groupName: gpu-nvidia + # Private Artifact Registry (the "GCR"): {location}-docker.pkg.dev/{project}/{repo}/huri:{tag} + image: europe-west4-docker.pkg.dev/hurigcp/huri/huri:nvidia-2.55.1@sha256:c68efb655adbe7024f389996256922293127e55bf9008667e949364c95f0d9c2 + replicas: 1 + minReplicas: 1 + # One V100 node (gpu_max_nodes = 1) hosts one worker pod (nvidia.com/gpu: 1), + # so maxReplicas = 1. The autoscaler packs multiple Serve model actors onto + # this single card via fractional GPU requests, not multiple pods. + maxReplicas: 1 + mountVoiceAssets: true + # cudaCacheHostPath is typically not used in GKE unless you have a specific reason + nodeSelector: + gpu: nvidia + # The gpu-pool nodes carry the GKE GPU taint nvidia.com/gpu=present:NoSchedule. + # GKE *usually* auto-injects a matching toleration for pods that request + # nvidia.com/gpu, but the KubeRay worker pod must not rely on that admission + # path: without this toleration the worker stays Pending on a healthy GPU node + # (node provisions Ready, pod never lands, KubeRay tears the cluster down, + # autoscaler reclaims the idle node — an endless up/down loop). + tolerations: + - key: nvidia.com/gpu + operator: Exists + effect: NoSchedule + # GKE doesn't use the 'nvidia' runtimeClassName by default for GPU nodes + # The drivers are handled at the node level. + runtimeClassName: null + # Ray's KubeRay autoscaler parses this field as: + # json.loads(resources_string[1:-1].replace("\\", "")) + # i.e. it strips an OUTER pair of double-quotes and removes backslash escapes + # before json.loads. So the stored CR value must be a quote-wrapped, escaped + # JSON string: "{\"GPU_TYPE_NVIDIA\":1}" (see Ray's ray-cluster.complete.yaml). + # The chart wraps this in single quotes (squote), so we write the outer + # double-quotes + escaped inner quotes here. WITHOUT the outer quote pair the + # autoscaler strips the braces instead, json.loads("GPU_TYPE_NVIDIA":1) fails + # with "Extra data", and the head's autoscaler sidecar crashloops. + customResources: '"{\"GPU_TYPE_NVIDIA\":1}"' + rayStartParams: + num-gpus: "1" + # Pin Ray's logical CPU count to match the pod requests below. Each model + # actor requests num_cpus:1, so this — not the GPU fraction — is what caps + # how many replicas pack onto the card. + num-cpus: "6" + # PVCs mounted into this (V100) worker: whisper (STT) + cosytts (TTS) + + # emotion (EMO prosody). Their download Jobs are pinned to this same gpu node + # (see models.* above) so the zonal RWO PVCs bind in the V100's zone. emage + # stays unmounted (disabled). + mountedModels: + - whisper + - cosytts + - emotion + resources: + # Sized for n1-standard-8 (8 vCPU / 30 GB / 1*V100 16GB), the machine the + # tfvars provisions. GKE leaves only ~7.5 allocatable vCPU and ~26Gi + # allocatable RAM on this node, so requests must stay under that or the + # autoscaler reports "Insufficient cpu/memory" and never scales the GPU + # pool up. Leaves headroom for the kubelet/system daemons and the raylet. + limits: + cpu: "7" + memory: "24Gi" + nvidia.com/gpu: "1" + requests: + cpu: "6" + memory: "14Gi" + nvidia.com/gpu: "1" + + # NOTE: the AMD worker group is intentionally omitted on GCP. Leaving it in + # (even at replicas: 0) makes Ray's autoscaler-v2 validate it and crash with + # "failed to detect `CPU` resources for group gpu-amd", because it has no + # num-cpus rayStartParam or CPU resource limit. There's no AMD GPU on GCP, so + # we simply don't define the group. + +# Voice models (STT/TTS) mount on the gpu-nvidia (V100) worker. Their PVCs use +# standard-rwo — a ZONAL ReadWriteOnce PD that can only attach to a node in the +# zone it was provisioned in. The gpu-pool is pinned to us-central1-c (gpu_node_zones +# in terraform.tfvars) so the V100 lands there with the already-bound PVCs, while the +# system pool is pinned to us-central1-a. So the download Jobs MUST run in the V100's +# zone, i.e. ON the gpu +# node itself (nodeSelector gpu: nvidia), so WaitForFirstConsumer binds each PVC +# where the worker can mount it. That node carries nvidia.com/gpu=present:NoSchedule +# and the Job requests no GPU, so GKE won't auto-inject a toleration — we add it. +models: + cosytts: + enabled: true + nodeSelector: + gpu: nvidia + tolerations: + - key: nvidia.com/gpu + operator: Exists + effect: NoSchedule + pvc: + storageClassName: "standard-rwo" + size: 20Gi + accessModes: + - ReadWriteOnce + mountPath: /models/cosytts + forceDownload: false + modelSource: + type: modelscope + modelId: FunAudioLLM/Fun-CosyVoice3-0.5B-2512 + env: + HURI_MODEL_PATH: /models/cosytts/FunAudioLLM/Fun-CosyVoice3-0.5B-2512 + HURI_COSY_DIR: /app/cosyvoice + whisper: + enabled: true + nodeSelector: + gpu: nvidia + tolerations: + - key: nvidia.com/gpu + operator: Exists + effect: NoSchedule + pvc: + storageClassName: "standard-rwo" + size: 2Gi + accessModes: + - ReadWriteOnce + mountPath: /models/whisper + modelSource: + type: huggingface + repoId: Systran/faster-whisper-base + env: + HURI_STT_MODEL_PATH: /models/whisper/Systran/faster-whisper-base + # Prosody / emotion model for the EMO module (voice -> emotion). Enabling the + # full-prosody voice pipeline (emo + eag + qag) in the website's huri_config.yaml + # made EMO load superb/hubert-large-superb-er (~1.2GB HF audio-classification + # model). EMO is a PLAIN Module — it runs inside the gpu-nvidia worker process + # and loads the weights with transformers `from_pretrained()`, so without a + # cached copy every session would re-download from HuggingFace at init. We cache + # it on a PVC here (same zonal-RWO, gpu-node-pinned download-Job pattern as + # whisper/cosytts) and point the module at the mount via huri_config.yaml + # (emo.args.model_name: /models/emotion/superb/hubert-large-superb-er). Unlike + # whisper/STT, EMO reads no env var for its path — hence the module arg, not `env`. + emotion: + enabled: true + nodeSelector: + gpu: nvidia + tolerations: + - key: nvidia.com/gpu + operator: Exists + effect: NoSchedule + pvc: + storageClassName: "standard-rwo" + size: 3Gi + accessModes: + - ReadWriteOnce + mountPath: /models/emotion + modelSource: + type: huggingface + repoId: superb/hubert-large-superb-er + # emage (gesture) weights stay DISABLED on GCP: GestureGeneration is already + # HEALTHY without them. Leave off until it's actually needed. + emage: + nodeSelector: + huri.io/system: "true" + pvc: + storageClassName: "standard-rwo" + +# TTS reference voice sample. The PVC (WaitForFirstConsumer) binds when the V100 +# worker mounts it (mountVoiceAssets: true); populate it AFTER the worker is +# Running with: kubectl cp voice.wav :/assets/voice.wav -n huri +voiceAssets: + enabled: true + pvc: + storageClassName: "standard-rwo" + size: 100Mi + accessModes: + - ReadWriteOnce + mountPath: /assets + env: + HURI_VOICE_SAMPLE_PATH: /assets/voice.wav + +# The raw Ray Serve websocket (huri-serve-svc:8000) is INTENTIONALLY not exposed +# by an Ingress. It has no authentication of its own — the website backend is the +# auth layer (OIDC relying party that pins the Authelia `sub` as the RAG +# partition key). Publishing this socket would let anyone connect to /session +# directly, bypass Authelia, and read/poison any user's RAG documents by setting +# an arbitrary user_id. The website reaches HuRI privately over the in-cluster +# ClusterIP (ws://huri-serve-svc.huri.svc.cluster.local:8000/session); only the +# website (app.huri.pommier.dev) and Authelia (auth.huri.pommier.dev) — both in +# the website-demo Terraform (HuRI_website_demo/deploy) — are public. +ingress: + enabled: false + +# Ray Dashboard (port 8265) ingress. Kept DISABLED on GCP: there is no nginx +# ingress controller in this cluster, and the dashboard has no auth of its own +# (same reasoning as the serve socket above — exposing it publicly is unsafe). +# The chart's ingress-dashboard.yaml template dereferences this block, so it must +# exist even when disabled. Reach the dashboard ad hoc with a port-forward: +# kubectl port-forward -n huri svc/huri-head-svc 8265:8265 +dashboard: + ingress: + enabled: false diff --git a/deploy/examples/cloud/variables.tf b/deploy/examples/cloud/variables.tf new file mode 100644 index 0000000..e389e58 --- /dev/null +++ b/deploy/examples/cloud/variables.tf @@ -0,0 +1,110 @@ +variable "project_id" { + description = "The GCP project ID" + type = string +} + +variable "region" { + description = "The GCP region. Used for regional resources (VPC, Artifact Registry). The GKE cluster itself is zonal (see var.zone)." + type = string + default = "us-central1" +} + +variable "zone" { + description = "The GCP zone the GKE cluster and its node pools run in. A zonal cluster places one node per pool (vs. one per zone for a regional cluster), which keeps SSD/CPU/GPU usage within free-trial quotas. Must be a zone within var.region." + type = string + default = "us-central1-a" +} + +variable "cluster_name" { + description = "The name of the GKE cluster" + type = string + default = "huri-ray-cluster" +} + +variable "gpu_type" { + description = "GPU *accelerator* type attached to each node — NOT a machine type. E.g. nvidia-tesla-t4, nvidia-l4, nvidia-tesla-a100." + type = string + default = "nvidia-tesla-t4" +} + +variable "gpu_machine_type" { + description = "Compute Engine *machine* type that hosts the GPU. Must be compatible with gpu_type: N1 (e.g. n1-standard-4) for T4, G2 (e.g. g2-standard-8) for L4, A2 (e.g. a2-highgpu-1g) for A100." + type = string + default = "n1-standard-4" +} + +variable "gpu_count" { + description = "Number of GPUs per node" + type = number + default = 1 +} + +variable "gpu_min_nodes" { + description = "Minimum GPU nodes per zone in the autoscaling pool. Keep >=1 to hold a warm GPU (model cold-start is minutes)." + type = number + default = 1 +} + +variable "gpu_max_nodes" { + description = "Maximum GPU nodes per zone the autoscaler may provision under load. Bounded by GPU quota and budget." + type = number + default = 3 +} + +variable "gpu_node_zones" { + description = <<-EOT + Zones the GPU node pool may place nodes in (all must be within var.region). + Spanning every zone in the region + location_policy=ANY lets the autoscaler + provision the GPU node in whichever zone currently has stock, instead of hard + -failing on RESOURCE_POOL_EXHAUSTED in a single zone. NOTE: gpu_min/max_nodes + are PER ZONE, so keep gpu_min_nodes=0 here or you hold one warm node *per zone*. + Empty list = single-zone behavior (pool stays in var.zone). + EOT + type = list(string) + default = [] +} + +# --------------------------------------------------------------------------- +# Fast LLM (LiteLLM -> Mistral) +# --------------------------------------------------------------------------- + +# NOTE: the Mistral API key is NOT a Terraform variable. It is read from GCP +# Secret Manager (secret id "mistral_api_key") in secrets.tf. Populate it out of +# band per the README before applying. + +variable "llm_model" { + description = "Mistral model the LiteLLM 'huri-fast' alias forwards to (via Mistral's OpenAI-compatible API). Any fast model works." + type = string + default = "ministral-14b-latest" +} + +# --------------------------------------------------------------------------- +# Embedding server (llama.cpp -> bge-large-en-v1.5 GGUF) +# --------------------------------------------------------------------------- + +variable "embedding_hf_repo" { + description = "Hugging Face GGUF repo[:quant] llama.cpp pulls for the embedding server. MUST be the same bge-large-en-v1.5 quant used at ingestion, or query vectors won't match the stored ones." + type = string + default = "CompendiumLabs/bge-large-en-v1.5-gguf:Q4_K_M" +} + +# --------------------------------------------------------------------------- +# Website image registry (the "GCR") +# +# HuRI owns the Artifact Registry that holds the website backend image, but is +# otherwise agnostic to the website itself. The website's own Terraform +# (HuRI_website_demo/deploy) builds the image reference from these same values +# and deploys it. Only the registry location/id are needed here. +# --------------------------------------------------------------------------- + +variable "website_image_location" { + description = "Artifact Registry location (region) for the website image repo, e.g. us-central1." + type = string + default = "us-central1" +} + +variable "website_image_repo" { + description = "Artifact Registry repository id holding the website image (imported, not created)." + type = string + default = "huri" +} diff --git a/deploy/examples/cloud/vpc.tf b/deploy/examples/cloud/vpc.tf new file mode 100644 index 0000000..fda867d --- /dev/null +++ b/deploy/examples/cloud/vpc.tf @@ -0,0 +1,21 @@ +resource "google_compute_network" "vpc" { + name = "${var.cluster_name}-vpc" + auto_create_subnetworks = false +} + +resource "google_compute_subnetwork" "subnet" { + name = "${var.cluster_name}-subnet" + region = var.region + network = google_compute_network.vpc.name + ip_cidr_range = "10.0.0.0/16" + + secondary_ip_range { + range_name = "pods" + ip_cidr_range = "10.1.0.0/16" + } + + secondary_ip_range { + range_name = "services" + ip_cidr_range = "10.2.0.0/16" + } +} diff --git a/deploy/examples/local_nvidia_amd/Chart.lock b/helm/Chart.lock similarity index 100% rename from deploy/examples/local_nvidia_amd/Chart.lock rename to helm/Chart.lock diff --git a/deploy/examples/local_nvidia_amd/Chart.yaml b/helm/Chart.yaml similarity index 100% rename from deploy/examples/local_nvidia_amd/Chart.yaml rename to helm/Chart.yaml diff --git a/deploy/examples/local_nvidia_amd/templates/_helpers.tpl b/helm/templates/_helpers.tpl similarity index 100% rename from deploy/examples/local_nvidia_amd/templates/_helpers.tpl rename to helm/templates/_helpers.tpl diff --git a/deploy/examples/local_nvidia_amd/templates/cosytts-model-init-job.yaml b/helm/templates/cosytts-model-init-job.yaml similarity index 91% rename from deploy/examples/local_nvidia_amd/templates/cosytts-model-init-job.yaml rename to helm/templates/cosytts-model-init-job.yaml index b0e3bc7..8d8ed91 100644 --- a/deploy/examples/local_nvidia_amd/templates/cosytts-model-init-job.yaml +++ b/helm/templates/cosytts-model-init-job.yaml @@ -49,6 +49,13 @@ spec: nodeSelector: {{- toYaml . | nindent 8 }} {{- end }} + {{- with $model.tolerations }} + # Lets the download Job schedule onto a tainted node (e.g. a GPU node with + # nvidia.com/gpu=present:NoSchedule) so the PVC binds in that node's zone — + # required when the mounting worker is a GPU pod on a zonal RWO PVC. + tolerations: + {{- toYaml . | nindent 8 }} + {{- end }} volumes: - name: models persistentVolumeClaim: diff --git a/deploy/examples/local_nvidia_amd/templates/emage-model-init-job.yaml b/helm/templates/emage-model-init-job.yaml similarity index 100% rename from deploy/examples/local_nvidia_amd/templates/emage-model-init-job.yaml rename to helm/templates/emage-model-init-job.yaml diff --git a/helm/templates/emotion-model-init-job.yaml b/helm/templates/emotion-model-init-job.yaml new file mode 100644 index 0000000..2c8b84c --- /dev/null +++ b/helm/templates/emotion-model-init-job.yaml @@ -0,0 +1,95 @@ +{{- if .Values.models.emotion.enabled }} +{{- $model := .Values.models.emotion }} +{{- $pvcName := printf "%s-emotion-models" (include "huri.fullname" .) }} +{{- if not (lookup "v1" "PersistentVolumeClaim" .Release.Namespace $pvcName) }} +--- +apiVersion: v1 +kind: PersistentVolumeClaim +metadata: + name: {{ $pvcName }} + labels: + {{- include "huri.labels" . | nindent 4 }} + annotations: + "helm.sh/hook": pre-install,pre-upgrade + "helm.sh/hook-weight": "-10" + "helm.sh/resource-policy": keep +spec: + accessModes: + {{- toYaml $model.pvc.accessModes | nindent 4 }} + resources: + requests: + storage: {{ $model.pvc.size }} + {{- if $model.pvc.storageClassName }} + storageClassName: {{ $model.pvc.storageClassName }} + {{- end }} +{{- end }} +--- +# Runs only on first install (not on upgrade) — model is already on the PVC. +apiVersion: batch/v1 +kind: Job +metadata: + name: {{ include "huri.fullname" . }}-emotion-init + labels: + {{- include "huri.labels" . | nindent 4 }} + annotations: + "helm.sh/hook": pre-install,pre-upgrade + "helm.sh/hook-weight": "-5" + "helm.sh/hook-delete-policy": hook-succeeded,before-hook-creation +spec: + backoffLimit: 3 + template: + metadata: + labels: + {{- include "huri.selectorLabels" . | nindent 8 }} + spec: + restartPolicy: OnFailure + {{- with $model.nodeSelector }} + nodeSelector: + {{- toYaml . | nindent 8 }} + {{- end }} + {{- with $model.tolerations }} + # Lets the download Job schedule onto a tainted node (e.g. a GPU node with + # nvidia.com/gpu=present:NoSchedule) so the PVC binds in that node's zone — + # required when the mounting worker is a GPU pod on a zonal RWO PVC. + tolerations: + {{- toYaml . | nindent 8 }} + {{- end }} + volumes: + - name: models + persistentVolumeClaim: + claimName: {{ include "huri.fullname" . }}-emotion-models + containers: + - name: emotion-downloader + image: python:3.11-slim + command: ["/bin/sh", "-c"] + args: + - | + set -e + MODEL_DIR="{{ $model.mountPath }}/{{ $model.modelSource.repoId }}" + # config.json is the reliable presence marker for a HF transformers + # repo (the weights file name varies: pytorch_model.bin / *.safetensors). + if [ -f "$MODEL_DIR/config.json" ]; then + echo "Model already present at $MODEL_DIR — skipping download." + exit 0 + fi + echo "Downloading {{ $model.modelSource.repoId }} into $MODEL_DIR …" + pip install --quiet huggingface_hub + python - <<'PYEOF' + from huggingface_hub import snapshot_download + snapshot_download( + "{{ $model.modelSource.repoId }}", + local_dir="{{ $model.mountPath }}/{{ $model.modelSource.repoId }}", + ) + PYEOF + echo "Download complete." + volumeMounts: + - name: models + mountPath: {{ $model.mountPath }} + resources: + requests: + cpu: "500m" + memory: "512Mi" + limits: + cpu: "2" + memory: "2Gi" +{{- end }} diff --git a/deploy/examples/local_nvidia_amd/templates/head-dashboard-svc.yaml b/helm/templates/head-dashboard-svc.yaml similarity index 100% rename from deploy/examples/local_nvidia_amd/templates/head-dashboard-svc.yaml rename to helm/templates/head-dashboard-svc.yaml diff --git a/deploy/examples/local_nvidia_amd/templates/ingress-dashboard.yaml b/helm/templates/ingress-dashboard.yaml similarity index 100% rename from deploy/examples/local_nvidia_amd/templates/ingress-dashboard.yaml rename to helm/templates/ingress-dashboard.yaml diff --git a/deploy/examples/local_nvidia_amd/templates/ingress.yaml b/helm/templates/ingress.yaml similarity index 100% rename from deploy/examples/local_nvidia_amd/templates/ingress.yaml rename to helm/templates/ingress.yaml diff --git a/deploy/examples/local_nvidia_amd/templates/rayservice.yaml b/helm/templates/rayservice.yaml similarity index 80% rename from deploy/examples/local_nvidia_amd/templates/rayservice.yaml rename to helm/templates/rayservice.yaml index 1922d2e..3b8aaf5 100644 --- a/deploy/examples/local_nvidia_amd/templates/rayservice.yaml +++ b/helm/templates/rayservice.yaml @@ -5,12 +5,36 @@ metadata: labels: {{- include "huri.labels" . | nindent 4 }} annotations: - ray.io/initializing-timeout: "20m" + # Hard ceiling for the FIRST bring-up: if the Serve apps aren't ready within + # this window the RayService enters a TERMINAL failure state, deletes the + # RayCluster, and refuses to recreate it until the object is recreated. Cold + # GPU bring-up on GCP (gpu-pool 0→1 scale-up + driver install + pulling the + # large nvidia image) can exceed the old 20m, which is exactly what tripped + # the terminal failure. Override via ray.initializingTimeout. + ray.io/initializing-timeout: {{ .Values.ray.initializingTimeout | default "20m" | quote }} spec: + {{- if .Values.ray.serviceUnhealthySecondThreshold }} + # How long the Serve *service* may stay unhealthy before the controller + # recreates the RayCluster. Cold GPU bring-up (node provision + driver install + # + large nvidia image pull) can take several minutes; the default is too short + # and tears the cluster down mid-bootstrap (head + worker terminate together, + # then the idle GPU node is reclaimed — an endless up/down loop). + serviceUnhealthySecondThreshold: {{ .Values.ray.serviceUnhealthySecondThreshold }} + {{- end }} + {{- if .Values.ray.deploymentUnhealthySecondThreshold }} + # Same idea, per individual Serve deployment (STT/TTS/Gesture wait on the GPU). + deploymentUnhealthySecondThreshold: {{ .Values.ray.deploymentUnhealthySecondThreshold }} + {{- end }} serveConfigV2: | {{ .Values.ray.serveConfig | indent 4 }} rayClusterConfig: rayVersion: {{ .Values.ray.version | quote }} + {{- if .Values.ray.enableInTreeAutoscaling }} + # Layer ② — the Ray autoscaler turns Serve replica demand into worker-pod + # demand, scaling workerGroupSpecs between their min/maxReplicas. Off unless + # explicitly enabled in values, so single-GPU deployments stay pinned. + enableInTreeAutoscaling: true + {{- end }} headGroupSpec: serviceType: {{ .Values.head.serviceType }} diff --git a/deploy/examples/local_nvidia_amd/templates/voice-assets-pvc.yaml b/helm/templates/voice-assets-pvc.yaml similarity index 100% rename from deploy/examples/local_nvidia_amd/templates/voice-assets-pvc.yaml rename to helm/templates/voice-assets-pvc.yaml diff --git a/deploy/examples/local_nvidia_amd/templates/whisper-model-init-job.yaml b/helm/templates/whisper-model-init-job.yaml similarity index 88% rename from deploy/examples/local_nvidia_amd/templates/whisper-model-init-job.yaml rename to helm/templates/whisper-model-init-job.yaml index 879d908..e25572f 100644 --- a/deploy/examples/local_nvidia_amd/templates/whisper-model-init-job.yaml +++ b/helm/templates/whisper-model-init-job.yaml @@ -47,6 +47,13 @@ spec: nodeSelector: {{- toYaml . | nindent 8 }} {{- end }} + {{- with $model.tolerations }} + # Lets the download Job schedule onto a tainted node (e.g. a GPU node with + # nvidia.com/gpu=present:NoSchedule) so the PVC binds in that node's zone — + # required when the mounting worker is a GPU pod on a zonal RWO PVC. + tolerations: + {{- toYaml . | nindent 8 }} + {{- end }} volumes: - name: models persistentVolumeClaim: diff --git a/requirements-nvidia.txt b/requirements-nvidia.txt index 98a8950..9ff2e43 100644 --- a/requirements-nvidia.txt +++ b/requirements-nvidia.txt @@ -20,7 +20,7 @@ x-transformers==2.11.24 einops==0.8.2 tiktoken==0.13.0 # cosyvoice/tokenizer pyarrow==18.1.0 # imported by cli paths via dataset utils? actually only dataset/processor — can drop -protobuf==4.25 +protobuf==4.25.8 # keep <5 so Ray Serve's _proto_to_dict works; 4.25.8 satisfies onnxruntime>=4.25.8 (also enforced by constraints.txt) pydantic==2.7.0 # transitive (transformers/fastapi), but pinning avoids drift regex==2025.11.3 tqdm==4.67.3 diff --git a/requirements.txt b/requirements.txt index 9f337e0..eadd691 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,24 +1,25 @@ -black -isort -mypy -flake8 -flake8-toml-config -pytest +black==26.3.1 +flake8==7.3.0 +flake8-toml-config==1.0.0 +isort==8.0.1 +mypy==1.20.2 +pytest==9.0.3 +pytest-asyncio==1.4.0 # server -numpy -ray[serve] -webrtcvad -faster-whisper -qdrant-client -sentence-transformers -pypdf -semantic_chunker +numpy==1.26.4 +ray[serve]==2.55.1 +webrtcvad==2.0.10 +faster-whisper==1.2.1 +qdrant-client==1.18.0 +sentence-transformers==3.2.1 +pypdf==5.1.0 +semantic_chunker==0.2.0 # client -sounddevice -websockets -omegaconf -prompt-toolkit +sounddevice==0.5.5 +websockets==16.0 +omegaconf==2.3.0 +prompt-toolkit==3.0.52 diff --git a/src/core/client.py b/src/core/client.py index 3786b4e..9df62d0 100644 --- a/src/core/client.py +++ b/src/core/client.py @@ -149,9 +149,17 @@ async def _receive_loop(self, ws: websockets.ClientConnection): data = event["data"] for hook in self.hooks[topic]: - if not isinstance(data, bytes): - data = hook.input_type(**data) - asyncio.create_task(hook.hook(data)) + # Hydrate via from_wire (not a bare **data splat) so events + # with nested EventData fields — e.g. RAGQuestion.transcript / + # .emotion — are rebuilt as dataclasses, mirroring the server's + # EventDataFactory. Build a fresh instance per hook so a topic + # with several hooks doesn't re-splat an already-built event. + hook_data = ( + data + if isinstance(data, bytes) + else hook.input_type.from_wire(data) + ) + asyncio.create_task(hook.hook(hook_data)) except (asyncio.CancelledError, websockets.ConnectionClosedOK): pass diff --git a/src/core/events.py b/src/core/events.py index 43176d2..862760b 100644 --- a/src/core/events.py +++ b/src/core/events.py @@ -2,6 +2,7 @@ import logging from collections import defaultdict from dataclasses import dataclass +from typing import Any, Mapping import numpy as np @@ -15,7 +16,17 @@ class EventData: """An event data must be derived from this class, and use @dataclass decorator. Or they can be bytes.""" - pass + @classmethod + def from_wire(cls, data: Mapping[str, Any]) -> "EventData": + """Build an event from a JSON payload sent by a client. + + Default: keyword-splat the payload onto the dataclass. Events whose fields + are nested dataclasses (or that accept a simpler external shape than the + in-pipeline one) override this — e.g. RAGQuestion accepts a bare + ``{"text": ...}`` typed question. In-process producers construct the + dataclass directly and never go through this path. + """ + return cls(**data) class EventGraph: diff --git a/src/core/huri.py b/src/core/huri.py index 66feb3f..6b92e12 100644 --- a/src/core/huri.py +++ b/src/core/huri.py @@ -141,5 +141,9 @@ async def receive_loop(session: Session, ws: WebSocket): await fin() except Exception: import traceback - print(f"[HuRI] finalize failed for {type(module).__name__}:\n{traceback.format_exc()}") - self.clients.pop(session_id, None) \ No newline at end of file + + print( + f"[HuRI] finalize failed for {type(module).__name__}:\n" + f"{traceback.format_exc()}" + ) + self.clients.pop(session_id, None) diff --git a/src/interfaces/cli_interface.py b/src/interfaces/cli_interface.py index 39cebb0..3cf2282 100644 --- a/src/interfaces/cli_interface.py +++ b/src/interfaces/cli_interface.py @@ -12,9 +12,9 @@ from src.core.client import ClientHook, ClientSender from src.core.interface import Interface -from src.modules.rag.events import RAGResult -from src.modules.speech_to_text.events import Sentence -from src.modules.text_to_speech.events import Audio +from src.modules.rag.events import RAGQuestion, RAGResult +from src.modules.speech_to_text.events import Transcript +from src.modules.text_to_speech.events import Audio, Token class AudioSender(ClientSender[bytes]): @@ -46,8 +46,8 @@ def callback(indata: np.ndarray, frames, time, status): await self.send(ws, chunk.tobytes()) -class TextSender(ClientSender[Sentence]): - output_type = Sentence +class TextSender(ClientSender[RAGQuestion]): + output_type = RAGQuestion def __init__(self, **kwargs): super().__init__(**kwargs) @@ -61,7 +61,7 @@ async def input_loop(self, ws): text = await session.prompt_async(">> ") if text == "\\exit": return - await self.send(ws, Sentence(text)) + await self.send(ws, RAGQuestion(Transcript(text, True), None)) except (EOFError, KeyboardInterrupt): pass @@ -172,6 +172,26 @@ async def hook(self, data: RAGResult): print("<<", data.answer) +class TokenHook(ClientHook[Token]): + """Print the RAG's streamed reply token-by-token (topic ``token``). + + This is the live text output of the pipeline (RAG -> Token). The website + backend consumes the same stream; the CLI just echoes each delta inline and + emits a newline on the end-of-utterance marker. + """ + + input_type = Token + + def __init__(self, **kwargs): + super().__init__(**kwargs) + + async def hook(self, data: Token): + if data.end: + print() + else: + print(data.text, end="", flush=True) + + class CLIInterface(Interface): def __init__(self): super().__init__(singletton=None) @@ -180,7 +200,7 @@ def get_senders(self) -> Dict[str, Type[ClientSender]]: return {"audio": AudioSender, "text": TextSender} def get_hooks(self) -> Dict[str, Type[ClientHook]]: - return {"audio": AudioHook, "text": TextHook} + return {"audio": AudioHook, "text": TextHook, "token": TokenHook} cli_interface = CLIInterface() diff --git a/src/modules/emotion/prosody_analysis.py b/src/modules/emotion/prosody_analysis.py index cd50b3d..a5ea9c9 100644 --- a/src/modules/emotion/prosody_analysis.py +++ b/src/modules/emotion/prosody_analysis.py @@ -2,8 +2,6 @@ from typing import List, Optional import numpy as np -import torch -from transformers import AutoModelForAudioClassification, Wav2Vec2FeatureExtractor from src.core.module import Module from src.modules.speech_to_text.events import Voice @@ -37,6 +35,11 @@ def __init__( ): super().__init__() + from transformers import ( + AutoModelForAudioClassification, + Wav2Vec2FeatureExtractor, + ) + self.model = AutoModelForAudioClassification.from_pretrained(model_name) self.feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(model_name) @@ -51,6 +54,8 @@ def __init__( self.lock: asyncio.Lock = asyncio.Lock() def _predict_emotion(self, audio_np: np.ndarray): + import torch + inputs = self.feature_extractor( audio_np, sampling_rate=self.sample_rate, return_tensors="pt", padding=True ) @@ -70,25 +75,47 @@ def _predict_emotion(self, audio_np: np.ndarray): } async def process(self, voice: Voice) -> Optional[Emotion]: + # End-of-utterance marker. MIC emits Voice(None) exactly ONCE per + # utterance, and this call MUST always produce an Emotion(end=True): it is + # what makes EAG finalize and, in turn, unblocks QAG (with use_emotion=True + # QAG holds the entire question until this emotion lands). The old code + # routed the marker through the same `running` / sliding-window guard as + # speech frames, so whenever an inference was mid-flight — or a late frame + # flipped the shared `self.silence` back to False before it was read — the + # single end marker was silently dropped and the whole voice turn hung. + # Handle it on its own path so it can never be swallowed. if voice.data is None: - self.silence = True - else: - self.silence = False async with self.lock: - self.buffer.append(voice.data) - + self.silence = True + tail = self.buffer + self.buffer = [] + # Read the tail so the final emotion reflects the actual utterance; + # fall back to a short zero buffer for a very short turn so the + # feature extractor still gets valid input and EAG gets scores. + audio = ( + np.concatenate(tail, axis=0) + if tail + else np.zeros(self.sample_rate // 10, dtype=np.float32) + ) + emotion_result = await asyncio.to_thread( + self._predict_emotion, audio_np=audio + ) + return Emotion( + emotion_result["label"], + emotion_result["confidence"], + emotion_result["scores"], + True, + ) + + # Speech frame: accumulate, and once a full analysis window is buffered run + # one interim inference (end=False). Only one inference runs at a time; + # extra frames just buffer until it finishes. async with self.lock: - if self.running: + self.silence = False + self.buffer.append(voice.data) + if self.running or len(self.buffer) < self.window_size: return None self.running = True - - async with self.lock: - buffer_size = len(self.buffer) - if buffer_size == 0 or ( - self.silence is False and buffer_size < self.window_size - ): - self.running = False - return None processing_chunks = self.buffer[: self.window_size] processing_audio = np.concatenate(processing_chunks, axis=0) @@ -105,5 +132,5 @@ async def process(self, voice: Voice) -> Optional[Emotion]: emotion_result["label"], emotion_result["confidence"], emotion_result["scores"], - self.silence, + False, ) diff --git a/src/modules/factory.py b/src/modules/factory.py index fb0f58c..315dd6e 100644 --- a/src/modules/factory.py +++ b/src/modules/factory.py @@ -37,7 +37,7 @@ def create(self, topic: str, data: Mapping[str, Any] | bytes) -> EventData | byt else: if issubclass(event_cls, EventData): - return event_cls(**data) + return event_cls.from_wire(data) else: raise RuntimeError(f"mismatched event data type: \ {event_cls} is not derived from EventData but should be.") diff --git a/src/modules/gesture/emage/configuration.py b/src/modules/gesture/emage/configuration.py index ad97076..4afd5a3 100644 --- a/src/modules/gesture/emage/configuration.py +++ b/src/modules/gesture/emage/configuration.py @@ -1,3 +1,5 @@ +from typing import Any, Dict, cast + from omegaconf import OmegaConf from transformers import PretrainedConfig @@ -7,7 +9,9 @@ class EmageAudioConfig(PretrainedConfig): def __init__(self, config_obj=None, **kwargs): if config_obj is not None: - cfg_dict = OmegaConf.to_container(config_obj, resolve=True) + cfg_dict = cast( + Dict[str, Any], OmegaConf.to_container(config_obj, resolve=True) + ) kwargs.update(cfg_dict) super().__init__(**kwargs) @@ -17,7 +21,9 @@ class EmageVQVAEConvConfig(PretrainedConfig): def __init__(self, config_obj=None, **kwargs): if config_obj is not None: - cfg_dict = OmegaConf.to_container(config_obj, resolve=True) + cfg_dict = cast( + Dict[str, Any], OmegaConf.to_container(config_obj, resolve=True) + ) kwargs.update(cfg_dict) super().__init__(**kwargs) @@ -27,6 +33,8 @@ class EmageVAEConvConfig(PretrainedConfig): def __init__(self, config_obj=None, **kwargs): if config_obj is not None: - cfg_dict = OmegaConf.to_container(config_obj, resolve=True) + cfg_dict = cast( + Dict[str, Any], OmegaConf.to_container(config_obj, resolve=True) + ) kwargs.update(cfg_dict) super().__init__(**kwargs) diff --git a/src/modules/gesture/emage/modeling.py b/src/modules/gesture/emage/modeling.py index 3fa1da9..f58b967 100644 --- a/src/modules/gesture/emage/modeling.py +++ b/src/modules/gesture/emage/modeling.py @@ -14,7 +14,6 @@ VQEncoderV5, VQEncoderV6, WavEncoder, - axis_angle_to_matrix, axis_angle_to_rotation_6d, recover_from_mask_ts, rotation_6d_to_axis_angle, diff --git a/src/modules/gesture/emage/processing.py b/src/modules/gesture/emage/processing.py index 91266b0..b40c34f 100644 --- a/src/modules/gesture/emage/processing.py +++ b/src/modules/gesture/emage/processing.py @@ -1,4 +1,5 @@ import math +from typing import Optional import torch import torch.nn as nn @@ -132,7 +133,7 @@ def recover_from_mask_ts(selected_motion: torch.Tensor, mask: list) -> torch.Ten dtype = selected_motion.dtype mask_arr = torch.tensor(mask, dtype=torch.bool, device=device) j = len(mask_arr) - sum_mask = mask_arr.sum().item() + sum_mask = int(mask_arr.sum().item()) c_channels = selected_motion.shape[-1] // sum_mask new_shape = selected_motion.shape[:-1] + (sum_mask, c_channels) selected_motion = selected_motion.reshape(new_shape) @@ -270,7 +271,7 @@ def __init__(self, args): input_size = args.vae_length n_resblk = 2 if input_size == channels[0]: - layers = [] + layers: list[nn.Module] = [] else: layers = [nn.Conv1d(input_size, channels[0], 3, 1, 1)] for i in range(n_resblk): @@ -326,7 +327,7 @@ def __init__( self.bn2 = norm_layer(planes) self.act2 = act_layer(inplace=True) if downsample is not None: - self.downsample = nn.Sequential( + self.downsample: Optional[nn.Sequential] = nn.Sequential( nn.Conv1d( inplanes, planes, @@ -401,6 +402,7 @@ class PeriodicPositionalEncoding(nn.Module): def __init__(self, d_model, dropout=0.1, period=15, max_seq_len=60): super().__init__() self.dropout = nn.Dropout(p=dropout) + self.pe: torch.Tensor pe = torch.zeros(period, d_model) position = torch.arange(0, period, dtype=torch.float).unsqueeze(1) div_term = torch.exp( diff --git a/src/modules/gesture/gesture.py b/src/modules/gesture/gesture.py index 554b8f4..a39704a 100644 --- a/src/modules/gesture/gesture.py +++ b/src/modules/gesture/gesture.py @@ -83,24 +83,24 @@ def __init__( print("[Gesture] loading face_vq...") face_vq = EmageVQVAEConv.from_pretrained(hf_repo, subfolder="emage_vq/face").to( - self.device + self.device # type: ignore[arg-type] ) print("[Gesture] loading upper_vq...") upper_vq = EmageVQVAEConv.from_pretrained( hf_repo, subfolder="emage_vq/upper" - ).to(self.device) + ).to(self.device) # type: ignore[arg-type] print("[Gesture] loading lower_vq...") lower_vq = EmageVQVAEConv.from_pretrained( hf_repo, subfolder="emage_vq/lower" - ).to(self.device) + ).to(self.device) # type: ignore[arg-type] print("[Gesture] loading hands_vq...") hands_vq = EmageVQVAEConv.from_pretrained( hf_repo, subfolder="emage_vq/hands" - ).to(self.device) + ).to(self.device) # type: ignore[arg-type] print("[Gesture] loading global_ae...") global_ae = EmageVAEConv.from_pretrained( hf_repo, subfolder="emage_vq/global" - ).to(self.device) + ).to(self.device) # type: ignore[arg-type] self.motion_vq = EmageVQModel( face_model=face_vq, @@ -112,7 +112,9 @@ def __init__( self.motion_vq.eval() print("[Gesture] loading EmageAudioModel...") - self.model = EmageAudioModel.from_pretrained(hf_repo).to(self.device) + self.model = EmageAudioModel.from_pretrained(hf_repo).to( + self.device # type: ignore[arg-type] + ) self.model.eval() self._warmup() @@ -167,12 +169,12 @@ def _warmup(self) -> None: torch.cuda.synchronize(self.device) print( f"[Gesture] warmup pass {pass_idx} {s:.2f}s " - f"({n} samples @ {_WARMUP_SRC_SR} Hz) \ -in {time.time() - ts:.2f}s", + f"({n} samples @ {_WARMUP_SRC_SR} Hz) " + f"in {time.time() - ts:.2f}s", ) print( - f"[Gesture] warmup done ({len(secs)} shapes x2) \ -in {time.time() - t0:.2f}s", + f"[Gesture] warmup done ({len(secs)} shapes x2) " + f"in {time.time() - t0:.2f}s", ) except Exception as e: # noqa: BLE001 — warmup is an optimisation, never fatal print(f"[Gesture] WARNING warmup failed: {e!r}") @@ -339,7 +341,9 @@ def _end_utterance(self) -> None: self._buf_start = 0 self._emitted = 0 - async def process(self, audio: Audio) -> AsyncGenerator[Motion, None]: + async def process( # type: ignore[override] + self, audio: Audio + ) -> AsyncGenerator[Motion, None]: # Each chunk arrives as its own process() task on the shared per-session # instance, so serialise under a lock to keep the buffer ordered. async with self._lock: diff --git a/src/modules/modules.py b/src/modules/modules.py index 66b627d..55f4978 100644 --- a/src/modules/modules.py +++ b/src/modules/modules.py @@ -19,7 +19,6 @@ def get_modules() -> Dict[str, Type[Module]]: "stt": STT, "tag": TAG, "emo": EMO, - "rag": RAG, "eag": EAG, "qag": QAG, "rag": RAG, diff --git a/src/modules/rag/events.py b/src/modules/rag/events.py index 24f883f..0873376 100644 --- a/src/modules/rag/events.py +++ b/src/modules/rag/events.py @@ -1,5 +1,5 @@ from dataclasses import dataclass, field -from typing import Optional +from typing import Any, Mapping, Optional from src.core.events import EventData from src.modules.emotion.events import Emotion @@ -27,4 +27,34 @@ class RAGQuestion(EventData): """Fully aggregated question to send to the RAG.""" transcript: Transcript - emotion: Optional[Emotion] + emotion: Optional[Emotion] = None + + @classmethod + def from_wire(cls, data: Mapping[str, Any]) -> "RAGQuestion": + """Accept a typed question straight off the wire. + + In-pipeline, QAG builds this from a Transcript + Emotion. External clients + (the website backend, the CLI text sender) inject a typed question as a + bare ``{"text": ...}`` — no live emotion — which we wrap into a final + Transcript. The nested branch also hydrates the dataclasses (JSON gives + plain dicts) so a fully-shaped payload round-trips too. + """ + if "text" in data: + return cls(Transcript(text=data["text"], end=True), None) + + # Thomas: it's nasty, but ragquestion is a subclass and doesn't work well with + # ray for my end, i get RAGQuestion.__init__() got an unexpected keyword + # argument 'text' if there is a agnostic way, we can fix this, but until now + # it's like this + + transcript = data.get("transcript") + if isinstance(transcript, Mapping): + transcript = Transcript(**transcript) + if not isinstance(transcript, Transcript): + raise ValueError( + "RAGQuestion.from_wire needs a 'text' or 'transcript' field" + ) + emotion = data.get("emotion") + if isinstance(emotion, Mapping): + emotion = Emotion(**emotion) + return cls(transcript, emotion) diff --git a/src/modules/rag/ingestion.py b/src/modules/rag/ingestion.py index 2a825b3..3ddb32b 100644 --- a/src/modules/rag/ingestion.py +++ b/src/modules/rag/ingestion.py @@ -185,8 +185,8 @@ def ingest_chunks( def chunk_strat(text: str, args, model: Any) -> list[str] | Any: """Pick the right chunking strategy based on args.""" if args.chunking == "semantic": - # Thomas: I need to import here, bceause it takes too much time earlier, - # or use a jupyter notebook to do it instead + # Thomas: I need to import here, bceause it takes too much time earlier, or + # use a jupyter notebook to do it instead from .semantic_chunker import SemanticChunker chunker = SemanticChunker( @@ -518,7 +518,7 @@ def main(): try: from .qdrant_utils import make_qdrant_client except ImportError: - from qdrant_utils import make_qdrant_client + from qdrant_utils import make_qdrant_client # type: ignore[no-redef] client = make_qdrant_client(args.qdrant_url, verify_ssl) # Lazy-load the model only if the command needs embeddings. @@ -528,7 +528,7 @@ def main(): if needs_embeddings: if args.embedding_url: print( - f"Embedding remotely via {args.embedding_url}" + f"Embedding remotely via {args.embedding_url} " f"(model={args.embedding_model})" ) model = RemoteEmbedder(args.embedding_url, args.embedding_model) diff --git a/src/modules/rag/memory_inspect.py b/src/modules/rag/memory_inspect.py index f912b16..e5ac964 100644 --- a/src/modules/rag/memory_inspect.py +++ b/src/modules/rag/memory_inspect.py @@ -4,13 +4,14 @@ python -m src.modules.rag.memory.memory_inspect python -m src.modules.rag.memory.memory_inspect --user-id """ + import argparse from datetime import datetime, timedelta try: from .qdrant_utils import make_qdrant_client except ImportError: - from qdrant_utils import make_qdrant_client + from qdrant_utils import make_qdrant_client # type: ignore[no-redef] HALF_LIFE_DAYS = 5.0 DELETE_BELOW, CONSOLIDATE_BELOW = 0.05, 0.30 @@ -19,14 +20,17 @@ def strength(payload: dict, at: datetime | None = None) -> float: """Query-independent strength: recency * importance (matches maintenance).""" at = at or datetime.now() - imp = payload.get("importance", 3) + imp: int = payload.get("importance", 3) half = max(HALF_LIFE_DAYS * (imp / 5.0), 0.5) try: - last = datetime.fromisoformat(payload.get("last_accessed") or payload["created_at"]) + last = datetime.fromisoformat( + payload.get("last_accessed") or payload["created_at"] + ) age = (at - last).total_seconds() / 86400.0 except Exception: age = 0.0 - return (0.5 ** (max(age, 0) / half)) * (imp / 10.0) + recency: float = 0.5 ** (max(age, 0) / half) + return recency * (imp / 10.0) def fate(s: float) -> str: @@ -47,8 +51,13 @@ def main(): qdrant = make_qdrant_client(args.qdrant_url) points, offset = [], None while True: - batch, offset = qdrant.scroll(collection_name=args.collection, limit=200, - offset=offset, with_payload=True, with_vectors=False) + batch, offset = qdrant.scroll( + collection_name=args.collection, + limit=200, + offset=offset, + with_payload=True, + with_vectors=False, + ) points.extend(batch) if offset is None: break @@ -56,33 +65,48 @@ def main(): now = datetime.now() rows = [] for p in points: - pl = p.payload + pl = p.payload or {} if pl.get("type") == "maintenance_marker": print(f"[marker] last maintenance run: {pl.get('last_run')}\n") continue if args.user_id and pl.get("_user_id") != args.user_id: continue s_now = strength(pl, now) - rows.append({ - "text": pl.get("text", "")[:60].replace("\n", " "), - "type": pl.get("type", "?"), - "imp": pl.get("importance", "?"), - "acc": pl.get("access_count", 0), - "age_d": round((now - datetime.fromisoformat( - pl.get("last_accessed") or pl["created_at"])).total_seconds() / 86400, 1), - "now": round(s_now, 3), - "+5d": round(strength(pl, now + timedelta(days=5)), 3), - "+15d": round(strength(pl, now + timedelta(days=15)), 3), - "fate": fate(s_now), - }) + rows.append( + { + "text": pl.get("text", "")[:60].replace("\n", " "), + "type": pl.get("type", "?"), + "imp": pl.get("importance", "?"), + "acc": pl.get("access_count", 0), + "age_d": round( + ( + now + - datetime.fromisoformat( + pl.get("last_accessed") or pl["created_at"] + ) + ).total_seconds() + / 86400, + 1, + ), + "now": round(s_now, 3), + "+5d": round(strength(pl, now + timedelta(days=5)), 3), + "+15d": round(strength(pl, now + timedelta(days=15)), 3), + "fate": fate(s_now), + } + ) rows.sort(key=lambda r: r["now"], reverse=True) - hdr = f"{'strength':>8} {'+5d':>6} {'+15d':>6} {'imp':>3} {'acc':>3} {'age':>5} {'fate':<12} text" + hdr = ( + f"{'strength':>8} {'+5d':>6} {'+15d':>6} {'imp':>3} {'acc':>3} " + f"{'age':>5} {'fate':<12} text" + ) print(hdr) print("-" * len(hdr)) for r in rows: - print(f"{r['now']:>8} {r['+5d']:>6} {r['+15d']:>6} {r['imp']:>3} " - f"{r['acc']:>3} {r['age_d']:>5} {r['fate']:<12} {r['text']}") + print( + f"{r['now']:>8} {r['+5d']:>6} {r['+15d']:>6} {r['imp']:>3} " + f"{r['acc']:>3} {r['age_d']:>5} {r['fate']:<12} {r['text']}" + ) print(f"\n{len(rows)} memories") diff --git a/src/modules/rag/memory_maintenance.py b/src/modules/rag/memory_maintenance.py index d20cc5a..052ca80 100644 --- a/src/modules/rag/memory_maintenance.py +++ b/src/modules/rag/memory_maintenance.py @@ -6,19 +6,19 @@ dead ones are deleted. Decay itself is computed lazily at query time in rag.py; this job only prunes and compresses. """ + import argparse -import json import uuid from collections import defaultdict from datetime import datetime import httpx -from qdrant_client.models import Distance, PointIdsList, PointStruct, VectorParams +from qdrant_client.models import PointIdsList, PointStruct try: from .qdrant_utils import make_qdrant_client except ImportError: - from qdrant_utils import make_qdrant_client + from qdrant_utils import make_qdrant_client # type: ignore[no-redef] DELETE_BELOW = 0.05 CONSOLIDATE_BELOW = 0.30 @@ -27,31 +27,39 @@ def strength(payload: dict) -> float: """Query-independent strength: recency * importance (no relevance term).""" - importance = payload.get("importance", 3) + importance: int = payload.get("importance", 3) half_life = max(HALF_LIFE_DAYS * (importance / 5.0), 0.5) try: - last = datetime.fromisoformat(payload.get("last_accessed") or payload["created_at"]) + last = datetime.fromisoformat( + payload.get("last_accessed") or payload["created_at"] + ) age_days = (datetime.now() - last).total_seconds() / 86400.0 except Exception: age_days = 0.0 - recency = 0.5 ** (age_days / half_life) + recency: float = 0.5 ** (age_days / half_life) return recency * (importance / 10.0) def embed(client: httpx.Client, url: str, model: str, text: str) -> list[float]: r = client.post(f"{url}/v1/embeddings", json={"model": model, "input": text}) r.raise_for_status() - return r.json()["data"][0]["embedding"] + embedding: list[float] = r.json()["data"][0]["embedding"] + return embedding def llm(client: httpx.Client, url: str, model: str, prompt: str) -> str: - r = client.post(f"{url}/api/chat", json={ - "model": model, "stream": False, - "messages": [{"role": "user", "content": prompt}], - "options": {"num_predict": 300}, - }) + r = client.post( + f"{url}/api/chat", + json={ + "model": model, + "stream": False, + "messages": [{"role": "user", "content": prompt}], + "options": {"num_predict": 300}, + }, + ) r.raise_for_status() - return r.json()["message"]["content"] + content: str = r.json()["message"]["content"] + return content def main(): @@ -69,8 +77,13 @@ def main(): points, offset = [], None while True: - batch, offset = qdrant.scroll(collection_name=args.collection, limit=200, - offset=offset, with_payload=True, with_vectors=False) + batch, offset = qdrant.scroll( + collection_name=args.collection, + limit=200, + offset=offset, + with_payload=True, + with_vectors=False, + ) points.extend(batch) if offset is None: break @@ -78,12 +91,16 @@ def main(): to_delete, weak_by_user = [], defaultdict(list) for p in points: - s = strength(p.payload) + payload = p.payload or {} + s = strength(payload) if s < DELETE_BELOW: to_delete.append(p) elif s < CONSOLIDATE_BELOW: - weak_by_user[p.payload.get("_user_id", "anonymous")].append(p) - print(f"delete: {len(to_delete)}, consolidate candidates: {sum(map(len, weak_by_user.values()))}") + weak_by_user[payload.get("_user_id", "anonymous")].append(p) + print( + f"delete: {len(to_delete)}, " + f"consolidate candidates: {sum(map(len, weak_by_user.values()))}" + ) if args.dry_run: return @@ -91,29 +108,47 @@ def main(): for user, weak in weak_by_user.items(): if len(weak) < 3: continue # not worth merging yet; keep decaying - texts = [p.payload["text"] for p in weak] - merged = llm(http, args.ollama_url, args.llm_model, - "These are old memories about conversations with the same person. " - "Merge them into a single 3-5 sentence memory keeping only durable " - "facts, preferences and recurring themes. Drop one-off small talk.\n\n" - + "\n---\n".join(texts)).strip() + texts = [(p.payload or {})["text"] for p in weak] + merged = llm( + http, + args.ollama_url, + args.llm_model, + "These are old memories about conversations with the same person. " + "Merge them into a single 3-5 sentence memory keeping only durable " + "facts, preferences and recurring themes. Drop one-off small talk.\n\n" + + "\n---\n".join(texts), + ).strip() vec = embed(http, args.ollama_url, args.embedding_model, merged) now = datetime.now().isoformat() - imp = min(max(p.payload.get("importance", 3) for p in weak) + 1, 10) - qdrant.upsert(collection_name=args.collection, points=[PointStruct( - id=str(uuid.uuid4()), vector=vec, payload={ - "text": merged, "_user_id": user, "type": "conversation_consolidated", - "created_at": now, "last_accessed": now, "access_count": 0, - "importance": imp, - })]) + imp = min(max((p.payload or {}).get("importance", 3) for p in weak) + 1, 10) + qdrant.upsert( + collection_name=args.collection, + points=[ + PointStruct( + id=str(uuid.uuid4()), + vector=vec, + payload={ + "text": merged, + "_user_id": user, + "type": "conversation_consolidated", + "created_at": now, + "last_accessed": now, + "access_count": 0, + "importance": imp, + }, + ) + ], + ) to_delete.extend(weak) print(f"[{user}] consolidated {len(weak)} → 1 (importance={imp})") if to_delete: - qdrant.delete(collection_name=args.collection, - points_selector=PointIdsList(points=[p.id for p in to_delete])) + qdrant.delete( + collection_name=args.collection, + points_selector=PointIdsList(points=[p.id for p in to_delete]), + ) print(f"deleted {len(to_delete)} points") if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/src/modules/rag/question_aggregator.py b/src/modules/rag/question_aggregator.py index 8fd864b..2f07f05 100644 --- a/src/modules/rag/question_aggregator.py +++ b/src/modules/rag/question_aggregator.py @@ -37,11 +37,20 @@ async def process(self, partial_question: PartialQuestion) -> Optional[RAGQuesti if partial_question.transcript is not None: self.current_transcript = partial_question.transcript - if self.current_transcript is not None: - if self.use_emotion: - if self.current_emotion is not None: - return RAGQuestion(self.current_transcript, self.current_emotion) - else: - return RAGQuestion(self.current_transcript, None) - - return None + if self.current_transcript is None: + return None + # With emotion enabled, hold the question back until its emotion lands. + if self.use_emotion and self.current_emotion is None: + return None + + emotion = self.current_emotion if self.use_emotion else None + question = RAGQuestion(self.current_transcript, emotion) + + # Clear the aggregation state after firing. Otherwise a later partial + # carrying only one half (the next turn's emotion-only update, or a + # stray transcript) re-emits this same question with stale data — which + # made the avatar answer the previous turn again. TAG/EAG already reset + # their own state on emit; QAG must do the same. + self.current_transcript = None + self.current_emotion = None + return question diff --git a/src/modules/rag/rag.py b/src/modules/rag/rag.py index ab09957..23fddee 100644 --- a/src/modules/rag/rag.py +++ b/src/modules/rag/rag.py @@ -1,15 +1,21 @@ +import asyncio import json import os import traceback import uuid -import asyncio from dataclasses import dataclass, field from datetime import datetime from typing import Any, AsyncGenerator import httpx from pydantic import BaseModel -from qdrant_client.models import FieldCondition, Filter, MatchValue +from qdrant_client.models import ( + FieldCondition, + Filter, + MatchValue, + PointIdsList, + PointStruct, +) from ray import serve from ray.serve import handle @@ -19,6 +25,12 @@ from .events import RAGQuestion from .qdrant_utils import make_qdrant_client +# Reserved _user_id for documents visible to EVERY user (e.g. the HuRI project +# overview). Retrieval matches the querying user's own id OR this shared id, so +# one ingested copy is reachable by all sessions. Ingest global docs with +# `ingestion.py --user-id __shared__ ...`. Keep in sync with any ingestion. +SHARED_USER_ID = "__shared__" + # Default character persona. Overridable per session via the `persona` key in the # client config's module args, or globally via HURI_RAG_DEFAULT_PERSONA in the # Serve app runtime_env.env_vars (see deploy values.yaml) — no rebuild needed. @@ -43,7 +55,7 @@ class RAGDeploymentConfig(BaseModel): verify_ssl: bool = True top_k: int = 5 score_threshold: float = 0.5 - + memory_collection: str = "conversations" memory_top_k: int = 3 memory_half_life_days: float = 5.0 @@ -72,7 +84,9 @@ class RAGHandle: """Stateless RAG processor. Streams LLM tokens to the caller.""" _MAINTENANCE_MARKER_ID = "00000000-0000-0000-0000-00000000feed" + def __init__(self, **kwargs): + self._maintenance_task: asyncio.Task | None = None self._cfg = RAGDeploymentConfig(**kwargs) self._apply_config() @@ -87,7 +101,8 @@ def _apply_config(self) -> None: print(f"[RAGHandle] Connected to Qdrant at {cfg.qdrant_url}") self._embed_client = httpx.AsyncClient(timeout=30.0, verify=cfg.verify_ssl) self._llm_client = httpx.AsyncClient(timeout=120.0, verify=cfg.verify_ssl) - if not hasattr(self, "_maintenance_task") or self._maintenance_task.done(): + task = self._maintenance_task + if task is None or task.done(): self._maintenance_task = asyncio.get_event_loop().create_task( self._maintenance_loop() ) @@ -114,7 +129,8 @@ async def _embed(self, text: str) -> list[float]: f"Embedding non-JSON response from {url}: {resp.text[:1000]}" ) from e try: - return payload["data"][0]["embedding"] + embedding: list[float] = payload["data"][0]["embedding"] + return embedding except (KeyError, IndexError, TypeError) as e: raise RuntimeError( f"Embedding unexpected schema from {url}: {str(payload)[:1000]}" @@ -144,10 +160,11 @@ def _get_profile(self, collection: str, _user_id: str) -> list[str]: ) except Exception: return [] - return [p.payload.get("text", "") for p in points if p.payload.get("text")] + return [t for p in points if (t := (p.payload or {}).get("text", ""))] def _ensure_memory_collection(self, vector_size: int) -> None: from qdrant_client.models import Distance, VectorParams + names = [c.name for c in self._qdrant.get_collections().collections] if self._cfg.memory_collection not in names: self._qdrant.create_collection( @@ -155,26 +172,34 @@ def _ensure_memory_collection(self, vector_size: int) -> None: vectors_config=VectorParams(size=vector_size, distance=Distance.COSINE), ) - async def _llm_complete(self, system_prompt: str, user_prompt: str, max_tokens: int = 300) -> str: + async def _llm_complete( + self, system_prompt: str, user_prompt: str, max_tokens: int = 300 + ) -> str: """Non-streamed convenience wrapper over _llm_stream.""" parts = [] - async for d in self._llm_stream(system_prompt, user_prompt, {"max_length": max_tokens}, None): + async for d in self._llm_stream( + system_prompt, user_prompt, {"max_length": max_tokens}, None + ): parts.append(d) return "".join(parts) def _memory_strength(self, payload: dict, relevance: float) -> float: - importance = payload.get("importance", 3) + importance: int = payload.get("importance", 3) half_life = max(self._cfg.memory_half_life_days * (importance / 5.0), 0.5) try: - last = datetime.fromisoformat(payload.get("last_accessed") or payload["created_at"]) + last = datetime.fromisoformat( + payload.get("last_accessed") or payload["created_at"] + ) age_days = (datetime.now() - last).total_seconds() / 86400.0 except Exception: age_days = 0.0 - recency = 0.5 ** (age_days / half_life) + recency: float = 0.5 ** (age_days / half_life) cfg = self._cfg - return (cfg.memory_w_relevance * relevance - + cfg.memory_w_recency * recency - + cfg.memory_w_importance * (importance / 10.0)) + return ( + cfg.memory_w_relevance * relevance + + cfg.memory_w_recency * recency + + cfg.memory_w_importance * (importance / 10.0) + ) def _search_memories(self, query_vector: list[float], _user_id: str) -> list[str]: """Retrieve, re-rank (relevance+recency+importance), reinforce, return texts.""" @@ -183,16 +208,26 @@ def _search_memories(self, query_vector: list[float], _user_id: str) -> list[str collection_name=self._cfg.memory_collection, query=query_vector, query_filter=Filter( - must=[FieldCondition(key="_user_id", match=MatchValue(value=_user_id))], - must_not=[FieldCondition(key="type", match=MatchValue(value="maintenance_marker"))], + must=[ + FieldCondition(key="_user_id", match=MatchValue(value=_user_id)) + ], + must_not=[ + FieldCondition( + key="type", match=MatchValue(value="maintenance_marker") + ) + ], ), limit=10, - score_threshold=0.2, # permissive; real filtering is the re-rank + score_threshold=0.2, # permissive; real filtering is the re-rank ).points except Exception: - return [] # collection missing / qdrant down → just no memories + return [] # collection missing / qdrant down → just no memories - scored = sorted(hits, key=lambda p: self._memory_strength(p.payload, p.score), reverse=True) + scored = sorted( + hits, + key=lambda p: self._memory_strength(p.payload or {}, p.score), + reverse=True, + ) top = scored[: self._cfg.memory_top_k] # MemoryBank-style reinforcement: recalled memories decay slower. @@ -201,13 +236,15 @@ def _search_memories(self, query_vector: list[float], _user_id: str) -> list[str try: self._qdrant.set_payload( collection_name=self._cfg.memory_collection, - payload={"last_accessed": now, - "access_count": p.payload.get("access_count", 0) + 1}, + payload={ + "last_accessed": now, + "access_count": (p.payload or {}).get("access_count", 0) + 1, + }, points=[p.id], ) except Exception: pass - return [p.payload.get("text", "") for p in top if p.payload.get("text")] + return [t for p in top if (t := (p.payload or {}).get("text", ""))] async def save_conversation(self, _user_id: str, history: list) -> None: """Summarize a finished session into one memory point. Called at disconnect.""" @@ -223,12 +260,21 @@ async def save_conversation(self, _user_id: str, history: list) -> None: ) try: raw = await self._llm_complete("You are a memory summarizer.", prompt) - cleaned = raw.strip().removeprefix("```json").removeprefix("```").removesuffix("```").strip() + cleaned = ( + raw.strip() + .removeprefix("```json") + .removeprefix("```") + .removesuffix("```") + .strip() + ) data = json.loads(cleaned) summary = str(data["summary"]) importance = max(1, min(int(data["importance"]), 10)) except Exception: - print(f"[RAG] memory summarization failed, storing raw tail:\n{traceback.format_exc()}") + print( + "[RAG] memory summarization failed, storing raw tail:\n" + f"{traceback.format_exc()}" + ) summary, importance = transcript[-500:], 3 if importance <= 1: @@ -240,23 +286,38 @@ async def save_conversation(self, _user_id: str, history: list) -> None: now = datetime.now().isoformat() self._qdrant.upsert( collection_name=self._cfg.memory_collection, - points=[PointStruct(id=str(uuid.uuid4()), vector=vector, payload={ - "text": summary, "_user_id": _user_id, "type": "conversation", - "created_at": now, "last_accessed": now, - "access_count": 0, "importance": importance, - })], + points=[ + PointStruct( + id=str(uuid.uuid4()), + vector=vector, + payload={ + "text": summary, + "_user_id": _user_id, + "type": "conversation", + "created_at": now, + "last_accessed": now, + "access_count": 0, + "importance": importance, + }, + ) + ], + ) + print( + f"[RAG] Saved conversation memory (importance={importance}): " + f"{summary[:80]}..." ) - print(f"[RAG] Saved conversation memory (importance={importance}): {summary[:80]}...") - def _last_maintenance(self) -> datetime | None: try: pts = self._qdrant.retrieve( collection_name=self._cfg.memory_collection, - ids=[self._MAINTENANCE_MARKER_ID], with_payload=True, with_vectors=False, + ids=[self._MAINTENANCE_MARKER_ID], + with_payload=True, + with_vectors=False, ) if pts: - return datetime.fromisoformat(pts[0].payload["last_run"]) + payload = pts[0].payload or {} + return datetime.fromisoformat(payload["last_run"]) except Exception: pass return None @@ -266,23 +327,30 @@ def _mark_maintenance_done(self, vector_size: int) -> None: # retrieval automatically (zero vector never scores) but be explicit anyway. self._qdrant.upsert( collection_name=self._cfg.memory_collection, - points=[PointStruct( - id=self._MAINTENANCE_MARKER_ID, - vector=[0.0] * vector_size, - payload={"type": "maintenance_marker", - "last_run": datetime.now().isoformat()}, - )], + points=[ + PointStruct( + id=self._MAINTENANCE_MARKER_ID, + vector=[0.0] * vector_size, + payload={ + "type": "maintenance_marker", + "last_run": datetime.now().isoformat(), + }, + ) + ], ) async def _maintenance_loop(self) -> None: import asyncio + check_secs = self._cfg.memory_maintenance_check_hours * 3600 while True: try: last = self._last_maintenance() - due = (last is None or - (datetime.now() - last).total_seconds() - >= self._cfg.memory_maintenance_days * 86400) + due = ( + last is None + or (datetime.now() - last).total_seconds() + >= self._cfg.memory_maintenance_days * 86400 + ) if due: print("[RAG] Running memory maintenance...") await self._run_maintenance() @@ -293,6 +361,7 @@ async def _maintenance_loop(self) -> None: async def _run_maintenance(self) -> None: """Decay-based pruning + consolidation. Same logic as memory_maintenance.py.""" from collections import defaultdict + DELETE_BELOW, CONSOLIDATE_BELOW = 0.05, 0.30 # scroll everything @@ -300,65 +369,92 @@ async def _run_maintenance(self) -> None: try: while True: batch, offset = self._qdrant.scroll( - collection_name=self._cfg.memory_collection, limit=200, - offset=offset, with_payload=True, with_vectors=False) + collection_name=self._cfg.memory_collection, + limit=200, + offset=offset, + with_payload=True, + with_vectors=False, + ) points.extend(batch) if offset is None: break except Exception: - return # collection doesn't exist yet — nothing to do + return # collection doesn't exist yet — nothing to do def base_strength(payload: dict) -> float: # query-independent: recency * importance - imp = payload.get("importance", 3) + imp: int = payload.get("importance", 3) half = max(self._cfg.memory_half_life_days * (imp / 5.0), 0.5) try: - last = datetime.fromisoformat(payload.get("last_accessed") or payload["created_at"]) + last = datetime.fromisoformat( + payload.get("last_accessed") or payload["created_at"] + ) age = (datetime.now() - last).total_seconds() / 86400.0 except Exception: age = 0.0 - return (0.5 ** (age / half)) * (imp / 10.0) + recency: float = 0.5 ** (age / half) + return recency * (imp / 10.0) to_delete, weak_by_user = [], defaultdict(list) vector_size = None for p in points: - if p.payload.get("type") == "maintenance_marker": + payload = p.payload or {} + if payload.get("type") == "maintenance_marker": continue - s = base_strength(p.payload) + s = base_strength(payload) if s < DELETE_BELOW: to_delete.append(p) elif s < CONSOLIDATE_BELOW: - weak_by_user[p.payload.get("_user_id", "anonymous")].append(p) + weak_by_user[payload.get("_user_id", "anonymous")].append(p) for user, weak in weak_by_user.items(): if len(weak) < 3: continue - texts = [p.payload["text"] for p in weak] - merged = (await self._llm_complete( - "You are a memory consolidator.", - "These are old memories about conversations with the same person. " - "Merge them into a single 3-5 sentence memory keeping only durable " - "facts, preferences and recurring themes. Drop one-off small talk.\n\n" - + "\n---\n".join(texts))).strip() + texts = [(p.payload or {})["text"] for p in weak] + merged = ( + await self._llm_complete( + "You are a memory consolidator.", + "These are old memories about conversations with the " + "same person. Merge them into a single 3-5 sentence " + "memory keeping only durable facts, preferences and " + "recurring themes. Drop one-off small talk.\n\n" + + "\n---\n".join(texts), + ) + ).strip() if not merged: continue vec = await self._embed(merged) vector_size = len(vec) now = datetime.now().isoformat() - imp = min(max(p.payload.get("importance", 3) for p in weak) + 1, 10) - self._qdrant.upsert(collection_name=self._cfg.memory_collection, - points=[PointStruct(id=str(uuid.uuid4()), vector=vec, payload={ - "text": merged, "_user_id": user, - "type": "conversation_consolidated", - "created_at": now, "last_accessed": now, - "access_count": 0, "importance": imp, - })]) + imp = min( + max((p.payload or {}).get("importance", 3) for p in weak) + 1, 10 + ) + self._qdrant.upsert( + collection_name=self._cfg.memory_collection, + points=[ + PointStruct( + id=str(uuid.uuid4()), + vector=vec, + payload={ + "text": merged, + "_user_id": user, + "type": "conversation_consolidated", + "created_at": now, + "last_accessed": now, + "access_count": 0, + "importance": imp, + }, + ) + ], + ) to_delete.extend(weak) print(f"[RAG] Consolidated {len(weak)} memories → 1 for user {user}") if to_delete: - self._qdrant.delete(collection_name=self._cfg.memory_collection, - points_selector=PointIdsList(points=[p.id for p in to_delete])) + self._qdrant.delete( + collection_name=self._cfg.memory_collection, + points_selector=PointIdsList(points=[p.id for p in to_delete]), + ) print(f"[RAG] Deleted {len(to_delete)} decayed memories") if vector_size is None: @@ -376,11 +472,22 @@ def _search( ) -> list[dict]: qdrant_filter: Any = None if filters: - conditions: Any = [ - FieldCondition(key=k, match=MatchValue(value=v)) - for k, v in filters.items() - ] - qdrant_filter = Filter(must=conditions) + must: Any = [] + should: Any = None + for k, v in filters.items(): + if k == "_user_id": + # Match the querying user's own docs OR the shared/global + # partition, so "all users" docs (ingested under + # SHARED_USER_ID) are retrieved alongside personal ones. + should = [ + FieldCondition(key=k, match=MatchValue(value=v)), + FieldCondition(key=k, match=MatchValue(value=SHARED_USER_ID)), + ] + else: + must.append(FieldCondition(key=k, match=MatchValue(value=v))) + # With `should`, Qdrant requires >=1 of the OR conditions to match; + # any other filters stay as `must` (AND). + qdrant_filter = Filter(must=must or None, should=should) try: results = qdrant.query_points( @@ -424,17 +531,18 @@ def _build_prompt( parts.append(f"Use a {preferences['tone']} tone.") if preferences.get("response_format") == "bullet_points": parts.append("Format your answer as bullet points.") - elif preferences.get("response_format") == "short": - parts.append("Keep your answer to 2-3 sentences maximum.") if preferences.get("extra_instructions"): parts.append(preferences["extra_instructions"]) + parts.append("Keep your answer to 2-3 sentences maximum. This is a discussion.") parts.append( "Use the context and memories in the user's message to inform your " "answers when relevant, but always answer in character. If you have " "relevant memories of past conversations, use them naturally. Only if " "you genuinely know nothing relevant, improvise in character rather " - "than admitting you lack information or breaking character. " ) + "than admitting you lack information or breaking character. " + ) + system_prompt = " ".join(parts) memory_block = "" @@ -447,10 +555,10 @@ def _build_prompt( if not chunks: user_prompt = ( - memory_block - + "No relevant context was found.\n\n" + memory_block + "No relevant context was found.\n\n" f"Question: {question}\n\n" - "Answer based on your memories above if relevant, otherwise general knowledge." + "Answer based on your memories above if relevant, otherwise " + "general knowledge." ) else: context_parts = [] @@ -462,13 +570,20 @@ def _build_prompt( ) context_block = "\n\n".join(context_parts) user_prompt = ( - memory_block - + f"Context:\n{context_block}\n\n" + memory_block + f"Context:\n{context_block}\n\n" f"Question: {question}\n\n" "Answer based on the context and your memories above. " "Don't speak about the sources, just use them to answer." ) + # Final line the model reads before generating — the highest-compliance + # slot for a formatting rule. open-mistral-nemo skips the persona-level + # no-Ah/Oh rule often enough that we restate it right at the tail. + user_prompt += ( + '\n\nStart your answer straight on the substance — do not open with ' + '"Ah" or "Oh".' + ) + return system_prompt, user_prompt async def _stream_ollama( @@ -617,7 +732,7 @@ async def stream(self, query: RAGQuery) -> AsyncGenerator[str, None]: class RAG(ModuleWithHandle, ModuleWithId): """RAG Module — streams LLM tokens. - input: question (Sentence) + input: question (RAGQuestion) output: token (Token) """ @@ -643,8 +758,9 @@ def __init__( print( f"[RAG] Initialized with user_id={_user_id}, language={language}, " - f"tone={tone}, response_format={response_format}, max_length={max_length}, " - f"temperature={temperature}, max_history_turns={max_history_turns}" + f"tone={tone}, response_format={response_format}, " + f"max_length={max_length}, temperature={temperature}, " + f"max_history_turns={max_history_turns}" ) self.preferences = { @@ -664,7 +780,9 @@ def __init__( self._max_history_turns = max_history_turns self.history: list[dict] = [] - async def process(self, data: RAGQuestion) -> AsyncGenerator[Token, None]: # type: ignore[override] + async def process( # type: ignore[override] + self, data: RAGQuestion + ) -> AsyncGenerator[Token, None]: """ Called when a "question" event arrives through the event bus. Packages _user_id + question, sends to the stateless RAGHandle. @@ -680,7 +798,7 @@ async def process(self, data: RAGQuestion) -> AsyncGenerator[Token, None]: # ty parts: list[str] = [] stream = self._handle.options(stream=True).stream.remote(query) - async for delta in stream: + async for delta in stream: # type: ignore[union-attr] parts.append(delta) yield Token(text=delta, end=False) yield Token(text="", end=True) @@ -701,7 +819,7 @@ def _record_turn(self, question: str, answer: str) -> None: def update_preferences(self, new_preferences: dict): self.preferences.update(new_preferences) - + async def finalize(self) -> None: """Called when the session ends — persist this conversation as a memory.""" if not self.history: diff --git a/src/modules/speech_to_text/speech_to_text.py b/src/modules/speech_to_text/speech_to_text.py index 7ce1f0f..a9a1b58 100644 --- a/src/modules/speech_to_text/speech_to_text.py +++ b/src/modules/speech_to_text/speech_to_text.py @@ -1,5 +1,6 @@ import asyncio import os +from concurrent.futures import ThreadPoolExecutor from typing import List, Optional import numpy as np @@ -11,9 +12,10 @@ from .events import Transcript, Voice _MODEL_PATH = os.environ.get("HURI_STT_MODEL_PATH", "base") +_NUM_WORKERS = int(os.environ.get("HURI_STT_NUM_WORKERS", "2")) -@serve.deployment(name="STT") +@serve.deployment(name="STT", max_ongoing_requests=8) class STTDeployment: """faster-whisper model wrapper. @@ -36,16 +38,39 @@ def __init__( model: str = _MODEL_PATH, device: str = "auto", compute_type: str = "auto", + num_workers: int = _NUM_WORKERS, ): from faster_whisper import WhisperModel + # num_workers lets CTranslate2 service several transcriptions at once on + # this single replica — each "worker" is an independent inference slot + # over the shared (read-only) weights. Combined with the thread pool + # below, N sessions are transcribed concurrently while staying fully + # independent: no per-call client state is ever held here (the sliding + # window lives in the per-session STT module). self.model_faster = WhisperModel( model, device=device, compute_type=compute_type, + num_workers=num_workers, + ) + # Run the *blocking* faster-whisper call off the actor's asyncio loop. + # transcribe() used to be an async method that called the synchronous + # model inline, blocking the replica's event loop for the whole inference + # — serialising every session on the replica and stalling Serve health + # checks. Offloading to a thread pool (sized to num_workers) keeps the + # loop free to accept other sessions' requests while inference runs. + self._executor = ThreadPoolExecutor( + max_workers=num_workers, thread_name_prefix="stt-transcribe" ) async def transcribe(self, audio: np.ndarray, language: str = "en") -> str: + loop = asyncio.get_running_loop() + return await loop.run_in_executor( + self._executor, self._transcribe_sync, audio, language + ) + + def _transcribe_sync(self, audio: np.ndarray, language: str) -> str: segments, _ = self.model_faster.transcribe( audio, language=language, @@ -103,6 +128,13 @@ def __init__( self.running = False self.lock: asyncio.Lock = asyncio.Lock() + # Whether we've transcribed some non-empty speech during the CURRENT + # utterance. Guards the end-of-turn reset (in process()) against a stray + # noise blip that transcribes to "" being mistaken for a finished turn. + # It is cleared when the turn ends, so the session handles continuous + # back-to-back turns rather than latching shut after the first one. + self._heard_speech: bool = False + async def process(self, voice: Voice) -> Optional[Transcript]: if voice.data is None: self.silence = True @@ -133,7 +165,27 @@ async def process(self, voice: Voice) -> Optional[Transcript]: processed_size = self.window_size - self.step_size async with self.lock: - self.buffer = self.buffer[processed_size:] self.running = False + # Track that we've heard actual speech this turn (guards the reset + # below against a stray noise blip that transcribes to ""). + if current_text: + self._heard_speech = True + + if self.silence and self._heard_speech: + # End of a real utterance. This Transcript carries end=True, so + # TAG/QAG emit the finished question downstream. Reset per-turn + # state so the NEXT utterance is fully independent: drop the + # residual sliding-window frames (they would otherwise be + # re-transcribed as a phantom continuation of this turn) and clear + # the speech flag. The client mutes the mic while the avatar + # responds (half-duplex), so its own TTS can't echo back and open + # a spurious turn — which is what the old single-turn latch was + # working around. + self.buffer = [] + self._heard_speech = False + else: + # Mid-utterance: slide the window forward, keeping the overlap. + self.buffer = self.buffer[processed_size:] + return Transcript(current_text, self.silence) diff --git a/src/modules/text_to_speech/text_to_speech.py b/src/modules/text_to_speech/text_to_speech.py index 95d2099..63e1b46 100644 --- a/src/modules/text_to_speech/text_to_speech.py +++ b/src/modules/text_to_speech/text_to_speech.py @@ -84,7 +84,39 @@ def __init__( voice_sample_transcript = raw voice_sample_transcript = _normalize_transcript(voice_sample_transcript) - self.model = CosyVoice3(model_dir=model_path, load_trt=False) + # fp16 is the whole point of running on a bandwidth-rich GPU (e.g. V100): + # it routes CosyVoice's LM/flow/vocoder matmuls through the fp16 tensor + # cores and halves memory traffic. CosyVoice3 defaults fp16=False (fp32), + # which leaves that advantage entirely unused. Default ON when a CUDA + # device is present; override with HURI_TTS_FP16 (fp16 needs CUDA, so it + # is forced off on CPU regardless). + import torch + + _fp16_env = os.environ.get("HURI_TTS_FP16") + if _fp16_env is not None: + fp16 = _fp16_env.strip().lower() in ("1", "true", "yes", "on") + else: + fp16 = torch.cuda.is_available() + if fp16 and not torch.cuda.is_available(): + print("[TTS] fp16 requested but no CUDA device; falling back to fp32") + fp16 = False + + # TensorRT accelerates the flow-matching estimator (often the single + # biggest TTS speedup) but builds a device-specific engine at startup + # (~minutes) and needs validation per GPU arch. OFF by default; flip + # HURI_TTS_TRT=1 to experiment once fp16 is confirmed working. + load_trt = os.environ.get("HURI_TTS_TRT", "").strip().lower() in ( + "1", + "true", + "yes", + "on", + ) + + print( + f"[TTS] loading CosyVoice3 (fp16={fp16}, load_trt={load_trt}) " + f"from {model_path!r}" + ) + self.model = CosyVoice3(model_dir=model_path, load_trt=load_trt, fp16=fp16) self.sample_rate: int = self.model.sample_rate self.prompt_speech = voice_sample_path @@ -92,6 +124,50 @@ def __init__( self._text_queues: dict[str, queue.Queue] = {} + # Pay CosyVoice's first-synth costs now, at deploy time, rather than on + # the first user utterance (see _warmup). + self._warmup() + + def _warmup(self) -> None: + """Run one throwaway synthesis so the first real utterance is hot. + + The first CosyVoice call pays one-time costs — CUDA context init, + cuBLAS/cuDNN algorithm selection, the LM/flow/vocoder first-call + compiles, and the caching allocator's first growth — that otherwise land + on the first user utterance and stall it for seconds. Draining a dummy + synth through the *real* fp16 path (same prompt, stream=True) pays them + upfront, mirroring the Gesture module's warmup. + + Best-effort: never fatal. The reference sample (voice.wav) is uploaded to + its PVC *after* the worker starts, so on a brand-new volume it may be + absent on first boot — warmup is skipped then and runs on the next + restart once the PVC holds the file. + """ + import time + + if not os.path.isfile(self.prompt_speech): + print( + f"[TTS] warmup skipped: reference sample {self.prompt_speech!r} " + "not present yet (upload voice.wav, then restart to warm)" + ) + return + try: + t0 = time.time() + + def _dummy_text(): + yield "Hello, this is a warm up." + + audio_iter = self.model.inference_zero_shot( + _dummy_text(), + self.prompt_text, + self.prompt_speech, + stream=True, + ) + chunks = sum(1 for _ in audio_iter) + print(f"[TTS] warmup done ({chunks} chunks) in {time.time() - t0:.2f}s") + except Exception as e: # noqa: BLE001 — warmup is an optimisation, never fatal + print(f"[TTS] WARNING warmup failed: {e!r}") + async def get_sample_rate(self) -> int: return self.sample_rate @@ -172,7 +248,9 @@ def __init__(self, _handle: handle.DeploymentHandle): # and silently drop trailing words). self._push_lock = asyncio.Lock() - async def process(self, token: Token) -> AsyncGenerator[Audio, None]: + async def process( # type: ignore[override] + self, token: Token + ) -> AsyncGenerator[Audio, None]: # Acquire BEFORE any await so lock-acquisition order matches token order. # Setup + push happen under the lock; only the first token of an # utterance goes on to drain/yield audio (outside the lock, so pushes of @@ -227,7 +305,7 @@ async def _drain_audio(self, session_id: str, audio_q: asyncio.Queue) -> None: response = self._handle.options(stream=True).stream_audio.remote(session_id) count = 0 pts = 0.0 - async for audio in response: # type: ignore[attr-defined] + async for audio in response: # type: ignore[union-attr] count += 1 audio.pts = pts pts += audio.data.shape[0] / audio.sample_rate diff --git a/src/modules/utils/sender.py b/src/modules/utils/sender.py index db5f633..434f457 100644 --- a/src/modules/utils/sender.py +++ b/src/modules/utils/sender.py @@ -17,6 +17,11 @@ class Sender(Module): """Sender Module Send output data to the client. + This data must be JSON serialisable, like a dataclass. + Audio wire format: [4B sample_rate uint32][1B end][8B pts float64][float32 PCM]. + Motion wire format: [8B pts float64][4B fps uint32][4B n_frames uint32] + [poses float32 n*165][expressions float32 n*100] + [trans float32 n*3]. This data must be JSON serialisable, like a dataclass. diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/core/__init__.py b/tests/core/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/core/test_events.py b/tests/core/test_events.py new file mode 100644 index 0000000..c4485b1 --- /dev/null +++ b/tests/core/test_events.py @@ -0,0 +1,230 @@ +import asyncio +import logging +from dataclasses import dataclass + +import numpy as np +import pytest + +from src.core.events import EventData, EventGraph, _summarize + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +class DummyModule: + input_type = None + output_type = None + + async def process(self, data): + raise NotImplementedError + + +@dataclass +class DummyEvent(EventData): + value: int + + +@dataclass +class ArrayEvent(EventData): + data: np.ndarray + + +# --------------------------------------------------------------------------- +# register() +# --------------------------------------------------------------------------- + + +def test_register(): + graph = EventGraph() + + module = DummyModule() + module.input_type = "input" + + graph.register(module) + + assert module in graph.subscribers["input"] + + +# --------------------------------------------------------------------------- +# publish() +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_publish_calls_subscriber(monkeypatch): + graph = EventGraph() + + called = asyncio.Event() + + class Module(DummyModule): + input_type = "input" + + async def process(self, data): + called.set() + return None + + module = Module() + graph.register(module) + + await graph.publish("input", DummyEvent(1)) + + await asyncio.wait_for(called.wait(), timeout=1) + + +# --------------------------------------------------------------------------- +# coroutine output +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_coroutine_output_published(): + graph = EventGraph() + + received = asyncio.Event() + + class Producer(DummyModule): + input_type = "start" + output_type = "next" + + async def process(self, data): + return DummyEvent(42) + + class Consumer(DummyModule): + input_type = "next" + + async def process(self, data): + assert data.value == 42 + received.set() + + graph.register(Producer()) + graph.register(Consumer()) + + await graph.publish("start", DummyEvent(0)) + + await asyncio.wait_for(received.wait(), timeout=1) + + +# --------------------------------------------------------------------------- +# async generator output +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_async_generator_output(): + graph = EventGraph() + + values = [] + + class Producer(DummyModule): + input_type = "start" + output_type = "stream" + + async def process(self, data): + yield DummyEvent(1) + yield DummyEvent(2) + yield DummyEvent(3) + + class Consumer(DummyModule): + input_type = "stream" + + async def process(self, data): + values.append(data.value) + + graph.register(Producer()) + graph.register(Consumer()) + + await graph.publish("start", DummyEvent(0)) + + await asyncio.sleep(0.1) + + assert values == [1, 2, 3] + + +# --------------------------------------------------------------------------- +# None outputs ignored +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_none_output_not_published(): + graph = EventGraph() + + called = False + + class Producer(DummyModule): + input_type = "start" + output_type = "next" + + async def process(self, data): + return None + + class Consumer(DummyModule): + input_type = "next" + + async def process(self, data): + nonlocal called + called = True + + graph.register(Producer()) + graph.register(Consumer()) + + await graph.publish("start", DummyEvent(0)) + await asyncio.sleep(0.05) + + assert not called + + +# --------------------------------------------------------------------------- +# recursive routing +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_recursive_routing(): + graph = EventGraph() + + done = asyncio.Event() + + class First(DummyModule): + input_type = "a" + output_type = "b" + + async def process(self, data): + return DummyEvent(data.value + 1) + + class Second(DummyModule): + input_type = "b" + + async def process(self, data): + assert data.value == 2 + done.set() + + graph.register(First()) + graph.register(Second()) + + await graph.publish("a", DummyEvent(1)) + + await asyncio.wait_for(done.wait(), timeout=1) + + +# --------------------------------------------------------------------------- +# summarize() +# --------------------------------------------------------------------------- + + +def test_summarize_numpy(): + event = ArrayEvent(np.zeros((3, 4), dtype=np.float32)) + + text = _summarize(event) + + assert "shape=(3, 4)" in text + assert "float32" in text + + +def test_summarize_regular(): + event = DummyEvent(123) + + text = _summarize(event) + + assert "DummyEvent" in text