From 127b606b7f7e547a324d5a2e4946313755a8528c Mon Sep 17 00:00:00 2001 From: Satrajit Ghosh Date: Sat, 7 Mar 2026 17:21:29 -0500 Subject: [PATCH 1/2] docs: add KIVA modular feature plan --- docs/kiva-modular-feature-plan.md | 619 ++++++++++++++++++++++++++++++ 1 file changed, 619 insertions(+) create mode 100644 docs/kiva-modular-feature-plan.md diff --git a/docs/kiva-modular-feature-plan.md b/docs/kiva-modular-feature-plan.md new file mode 100644 index 0000000..7f24cd4 --- /dev/null +++ b/docs/kiva-modular-feature-plan.md @@ -0,0 +1,619 @@ +# KIVA Modular Feature Implementation Plan + +Status: Draft + +Date: 2026-03-07 + +Branch: `codex/rebuild-spec` + +## 1. Goal + +This document analyzes the requested KIVA features and turns them into a modular implementation plan that can be reused across flows. + +Constraint: + +- audiobook training is only one KIVA flow +- shared capabilities must live in reusable platform modules +- flow-specific behavior should be configured through activity schemas, resources, and flow definitions rather than embedded in one-off runtime code + +## 2. Requested Feature Themes + +The requested items cluster into six workstreams: + +1. Reading-state grounding +2. Pedagogical control and response policy +3. Engagement detection and motivational support +4. Runtime latency and synchronization +5. Analytics, mastery, and cross-session measurement +6. Privacy, validation, and reporting + +## 3. Existing System Baseline + +Relevant capabilities already present: + +- KIVA uses advanced flows with `reading_context` and user state in [`src/server/activities/vocab-tutoring/flow_config.json`](/Users/satra/software/sensein/riverst/src/server/activities/vocab-tutoring/flow_config.json) +- indexable reading content already exists in activity resources and `get_reading_context` flow handling in [`src/server/bot/flows/handlers.py`](/Users/satra/software/sensein/riverst/src/server/bot/flows/handlers.py) +- audiobook playback already exposes word-level timing in [`src/client/react/src/pages/AudioPlayerPage.tsx`](/Users/satra/software/sensein/riverst/src/client/react/src/pages/AudioPlayerPage.tsx) +- session artifacts already include transcripts, audio, and metrics via [`src/server/bot/core/event_manager.py`](/Users/satra/software/sensein/riverst/src/server/bot/core/event_manager.py) and [`src/server/bot/monitoring/metrics_logger.py`](/Users/satra/software/sensein/riverst/src/server/bot/monitoring/metrics_logger.py) +- subtitle rendering is already present in the runtime UI in [`src/client/react/src/components/avatarInteraction/AvatarInteractionContent.tsx`](/Users/satra/software/sensein/riverst/src/client/react/src/components/avatarInteraction/AvatarInteractionContent.tsx) + +Gaps relative to the request: + +- no unified reading-position state model shared across audiobook playback and dialogue flow execution +- no reusable policy engine for single-question prompts, wait time, minimum expectations, or contingent follow-ups +- no explicit disengagement detector +- current metrics are infrastructure-focused, not pedagogy-focused +- no defined cross-session mastery or retention layer +- no validation protocol against human-coded samples + +## 4. Design Principles + +### 4.1 Shared capability first + +Every new feature should be implemented as: + +- a reusable state contract +- a reusable runtime processor or service +- a reusable analytics event schema +- optional per-flow configuration + +Audiobook tutoring should consume these modules, not define them. + +### 4.2 Flow configuration over prompt sprawl + +Behavioral rules such as: + +- minimum vocabulary targets +- minimum student response length +- single-question constraints +- disengagement handling + +should be encoded as structured policy in activity config and flow metadata, then enforced by shared runtime components. + +### 4.3 Event-first analytics + +Pedagogical and engagement metrics should be derived from normalized events instead of fragile transcript-only post-processing. + +### 4.4 Testable contracts + +Each subsystem should be introduced with: + +- schema validation +- unit tests +- replay tests using captured sessions or synthetic traces + +## 5. Proposed Modular Architecture + +## 5.1 Shared modules + +### Module A: Reading State Service + +Purpose: + +- maintain canonical reading position at word, sentence, paragraph, and section level +- expose current location to prompts, dialogue flow, analytics, and UI + +Responsibilities: + +- accept reading progression events from audiobook playback and dialogue flow +- persist current location in session state +- resolve current text span and neighboring context +- detect pauses, rereads, and regressions + +Consumers: + +- audiobook flow +- vocabulary tutoring flow +- future comprehension flows +- analytics layer + +### Module B: Instruction Policy Engine + +Purpose: + +- enforce teaching and response expectations consistently across KIVA flows + +Responsibilities: + +- enforce single-question prompts +- enforce minimum instructional targets per session +- enforce minimum response expectations +- decide when modeling, sentence starters, or adaptive prompts are needed +- maintain stable instructional routines across sessions +- ensure follow-up prompts are contingent on prior student response + +Consumers: + +- all KIVA instructional flows + +### Module C: Engagement Monitor + +Purpose: + +- detect disengagement and guide short motivational or pause behavior + +Responsibilities: + +- classify silence, refusal, minimal response, repeated non-response, and participation drop-off +- emit engagement events +- trigger affect acknowledgment and brief scaffolding +- offer pause or stop options when thresholds are met + +Consumers: + +- KIVA tutoring flows +- future support-oriented flows + +### Module D: Runtime Synchronization and Latency Monitor + +Purpose: + +- reduce perceived and measured lag across ASR, subtitles, avatar response, and animation + +Responsibilities: + +- instrument timestamps across speech input, transcription, subtitle render, LLM response, TTS start, TTS end, and avatar animation start +- expose latency metrics per turn and per session +- align subtitles and animation with audio playback + +Consumers: + +- all real-time avatar flows +- performance dashboards + +### Module E: Pedagogical Analytics Layer + +Purpose: + +- convert session events into interpretable educational and product metrics + +Responsibilities: + +- track pedagogical behaviors such as question type, scaffolding, wait time, motivational prompts, and contingent follow-ups +- track student engagement signals such as response length, response latency, participation trends, and silence +- compute learning metrics such as exposures per word, instructional effort, time to mastery, and retention +- support cross-session growth and mastery trajectories + +Consumers: + +- session review UI +- research exports +- product iteration workflows + +### Module F: Privacy and Evaluation Layer + +Purpose: + +- govern storage, documentation, de-identification, and validation quality + +Responsibilities: + +- document data usage and retention +- support privacy-preserving storage decisions +- maintain validation datasets and agreement checks against human-coded samples +- track confidence and quality of automated measures + +Consumers: + +- engineering +- research +- compliance review + +## 5.2 Flow-specific adapters + +### Audiobook flow adapter + +Responsibilities: + +- generate word-, sentence-, and paragraph-level reading events from the audiobook player +- update shared reading state +- consume policy engine outputs for grounding prompts in current text + +### Vocabulary tutoring flow adapter + +Responsibilities: + +- request current reading state from the shared service +- ground vocabulary and comprehension prompts in current text position +- log exposures, mastery attempts, and scaffolding events + +### Future KIVA flow adapters + +Examples: + +- comprehension check flow +- retell flow +- second-language tutoring flow + +These should reuse the same shared services and only differ in policy configuration and content schemas. + +## 6. Feature Analysis and Implementation Approach + +## 6.1 Reading-state grounding + +Requested features covered: + +- automatic reading position at sentence, paragraph, and word level +- persistent reading state shared across dialogue, prompts, and feedback +- current-text grounding for questions and follow-ups +- logging of progression, pauses, and rereads + +Implementation approach: + +1. Extend reading content schema to include stable IDs for: + - word + - sentence + - paragraph + - chapter or section +2. Introduce a session-scoped `reading_state` object in flow and session state. +3. Emit normalized reading events: + - `word_entered` + - `sentence_completed` + - `paragraph_completed` + - `pause_detected` + - `reread_detected` + - `position_set` +4. Build a `ReadingStateService` that derives current position and recent history. +5. Expose a shared handler to flows such as `get_reading_position` and `get_current_text_window`. + +Key design choice: + +- do not infer progress only from dialogue transcripts +- prefer direct player/runtime events when available +- fall back to transcript inference only as a secondary heuristic + +## 6.2 Pedagogical control and instructional consistency + +Requested features covered: + +- adaptive wait time after child responses +- single-question prompts +- contingent follow-up prompts +- minimum instructional expectations per session +- minimum response expectations per student response +- adaptive support such as modeling or sentence starters +- consistent instructional routines across sessions + +Implementation approach: + +1. Define a reusable `instruction_policy` schema in activity configs. +2. Add policy fields such as: + - `single_question_only` + - `min_target_words_per_session` + - `min_student_words_per_response` + - `adaptive_wait_time_ms` + - `allowed_scaffolds` + - `routine_template` +3. Introduce a `PolicyEngine` processor between transcript interpretation and prompt generation. +4. Have the flow manager consult policy outputs before advancing nodes or issuing prompts. +5. Log policy interventions as explicit analytics events. + +Key design choice: + +- enforce constraints structurally where possible instead of relying on large system prompts alone + +## 6.3 Engagement and motivational support + +Requested features covered: + +- disengagement detection +- affect acknowledgment +- motivational scaffolding +- graceful pause or stopping option + +Implementation approach: + +1. Define disengagement signals: + - refusal phrases + - silence thresholds + - repeated minimal responses + - declining participation trend +2. Build an `EngagementMonitor` that consumes transcript and latency events. +3. Produce engagement states such as: + - `engaged` + - `hesitant` + - `disengaged` + - `paused` +4. Map each state to configurable recovery behaviors. +5. Require flows to explicitly support `pause` and `stop` transitions. + +Key design choice: + +- the engagement monitor should recommend actions +- the flow adapter decides whether and how to surface those actions based on activity policy + +## 6.4 Runtime synchronization and latency + +Requested features covered: + +- lower latency from child speech to transcription, subtitle render, and avatar response +- synchronize audio, subtitles, and avatar animation +- monitor latency as a core performance metric + +Implementation approach: + +1. Add per-turn timestamp instrumentation to runtime processors. +2. Define a latency event model: + - speech end + - ASR final + - subtitle rendered + - LLM first token + - TTS start + - audio playback start + - avatar animation start +3. Compute turn-level latency breakdowns. +4. Refactor subtitle and avatar animation consumers to use shared timing cues where available. +5. Surface latency summary in session metrics. + +Key design choice: + +- latency must be tracked both as an engineering metric and as a pedagogical quality metric because long lag degrades instructional flow + +## 6.5 Analytics and learning measurement + +Requested features covered: + +- session-level interaction data with timestamps and reading position +- pedagogical behavior tracking +- student engagement tracking +- learning and mastery metrics +- cross-session growth and retention +- interpretable session-level summaries +- analytics that inform instructional improvement + +Implementation approach: + +1. Define a normalized `session_event` schema covering: + - reading events + - dialogue events + - policy interventions + - engagement events + - teaching actions + - mastery updates +2. Store raw events and derived metrics separately. +3. Compute derived outputs: + - exposures per word + - number of scaffold attempts + - question-type distribution + - average wait time + - response-length distribution + - participation trend + - time to mastery + - retained/not retained across sessions +4. Generate session summary narratives from these structured metrics. +5. Add longitudinal summaries at the learner level. + +Key design choice: + +- summaries should cite structured signals rather than free-form LLM impression alone + +## 6.6 Validation, privacy, and documentation + +Requested features covered: + +- validate automated metrics against human-coded samples +- privacy-preserving data storage +- clear documentation of data use + +Implementation approach: + +1. Define a human-coding rubric for: + - question type + - scaffolding + - affect support + - disengagement + - mastery event +2. Build a gold-sample evaluation dataset. +3. Compare automated output to coded samples on a scheduled basis. +4. Document: + - data collected + - retention policy + - de-identification strategy + - intended analytic uses +5. Store PII-sensitive and research-sensitive artifacts behind clear boundaries. + +## 7. Data Contract Additions + +## 7.1 Session state additions + +Add to shared session/flow state: + +```json +{ + "reading_state": { + "resource_id": "the_tale_of_peter_rabbit", + "chapter_id": 3, + "paragraph_id": "p12", + "sentence_id": "s48", + "word_id": "w391", + "position_updated_at": "2026-03-07T14:00:00Z", + "recent_positions": [], + "pause_count": 0, + "reread_count": 0 + }, + "instruction_state": { + "target_words_taught": [], + "instructional_steps_completed": [], + "single_question_compliance": true, + "wait_time_events": [] + }, + "engagement_state": { + "current_status": "engaged", + "minimal_response_streak": 0, + "silence_events": 0, + "refusal_events": 0 + }, + "mastery_state": { + "word_exposures": {}, + "word_mastery_status": {}, + "retention_history": [] + } +} +``` + +## 7.2 Event schema additions + +Add normalized events such as: + +- `reading_position_updated` +- `reading_pause_detected` +- `reading_reread_detected` +- `question_prompted` +- `wait_time_applied` +- `student_response_evaluated` +- `scaffold_prompted` +- `motivation_prompted` +- `pause_offered` +- `latency_measured` +- `word_exposure_recorded` +- `mastery_status_changed` + +## 8. Implementation Phases + +## Phase 0: Contracts and instrumentation design + +Deliverables: + +- shared schemas for reading state, instruction policy, engagement events, and session events +- test fixtures and synthetic traces +- design review with KIVA stakeholders + +TDD: + +- schema validation tests +- event serialization tests + +## Phase 1: Shared Reading State Service + +Deliverables: + +- reading state model and persistence +- audiobook event emitter +- flow handlers for current reading location +- pause and reread detection + +TDD: + +- reading position update tests +- pause/reread detection tests +- integration tests from audiobook playback events into flow-visible state + +## Phase 2: Instruction Policy Engine + +Deliverables: + +- configurable instruction policy schema +- enforcement of single-question prompts +- adaptive wait-time controller +- minimum instructional and response expectation tracking +- contingent follow-up enforcement hooks + +TDD: + +- policy rule evaluation tests +- turn replay tests for contingent follow-up behavior +- flow tests for minimum target coverage + +## Phase 3: Engagement Monitor + +Deliverables: + +- disengagement classifier +- motivational scaffold templates +- pause/stop flow hooks + +TDD: + +- silence and minimal-response detection tests +- refusal classification tests +- flow tests for pause/stop transitions + +## Phase 4: Runtime latency and synchronization + +Deliverables: + +- end-to-end timestamp instrumentation +- subtitle/audio/avatar sync improvements +- latency dashboards in metrics summary + +TDD: + +- processor timestamp propagation tests +- subtitle timing tests +- latency budget regression tests + +## Phase 5: Pedagogical analytics and mastery + +Deliverables: + +- normalized event log +- derived pedagogy and engagement metrics +- mastery and cross-session growth tracking +- session summary generator + +TDD: + +- metric derivation tests +- longitudinal aggregation tests +- summary generation tests on fixed event fixtures + +## Phase 6: Validation, privacy, and reporting + +Deliverables: + +- human-coded validation dataset +- agreement evaluation scripts +- privacy and data-use documentation +- instructional analytics views + +TDD: + +- evaluation pipeline tests +- privacy redaction tests +- report-generation tests + +## 9. Recommended GitHub Work Breakdown + +The work should be tracked as one project with nine issues: + +1. KIVA shared contracts and event schemas +2. Shared reading state service and audiobook instrumentation +3. Instruction policy engine for KIVA flows +4. Engagement monitor and pause/stop handling +5. Runtime latency instrumentation and synchronization +6. Pedagogical analytics event pipeline +7. Mastery, growth, and cross-session tracking +8. Validation against human-coded samples +9. Privacy, data documentation, and analytics surfacing + +Rationale: + +- this is small enough for a single board +- the issues are modular and map to reusable system capabilities +- audiobook-specific work is contained inside issue 2 and does not dominate the architecture + +## 10. Risks + +- prompt-only enforcement will drift; policy must be implemented structurally +- word-level reading position may vary across audiobook assets; content contracts must be normalized first +- latency fixes can regress if instrumentation is added without performance budgets +- cross-session growth tracking is only useful if learner identity and session linkage are defined cleanly +- privacy work cannot be deferred until after analytics collection is already in place + +## 11. Acceptance Criteria + +This initiative is complete when: + +- KIVA flows can access canonical reading position without ad hoc inference +- audiobook and tutoring flows share the same reading-state service +- single-question and contingent follow-up behavior is enforced by shared policy +- disengagement detection triggers configurable motivational or pause behavior +- session logs capture reading, pedagogy, engagement, mastery, and latency events +- session summaries are derived from interpretable structured metrics +- cross-session retention metrics are available for supported learners +- automated metrics have documented validation against human-coded samples +- privacy and data-use documentation is committed in the repo + +## 12. Recommended Next Step + +Create the board and issues immediately from the work breakdown in section 9, then implement phase 0 first with schema and replay-test scaffolding before changing any KIVA prompt logic. From af77c2bff5e6073809d6c237c5108e3bd0ed6685 Mon Sep 17 00:00:00 2001 From: Satrajit Ghosh Date: Sat, 7 Mar 2026 20:42:21 -0500 Subject: [PATCH 2/2] Add CPU and GPU deployment modes --- README.md | 18 +++++ docker-compose.gpu.yaml | 8 +++ docker-compose.yaml | 4 ++ docs/gpu-cpu-deployment-plan.md | 43 ++++++++++++ notes/first_steps_to_deploy.md | 6 +- src/server/Dockerfile | 12 +++- src/server/README.md | 17 ++++- src/server/bot/processors/video/processor.py | 18 ++++- src/server/bot/utils/device_utils.py | 57 ++++++++++++---- src/server/env.example | 1 + src/server/requirements.gpu.txt | 4 ++ src/server/tests/test_device_utils.py | 71 ++++++++++++++++++++ 12 files changed, 239 insertions(+), 20 deletions(-) create mode 100644 docker-compose.gpu.yaml create mode 100644 docs/gpu-cpu-deployment-plan.md create mode 100644 src/server/requirements.gpu.txt create mode 100644 src/server/tests/test_device_utils.py diff --git a/README.md b/README.md index 0cd5cae..bfaa32a 100644 --- a/README.md +++ b/README.md @@ -105,6 +105,14 @@ More detailed instructions for setting up the [client](src/client/react/README.m docker compose up --build ``` +This default Docker setup builds the server for CPU / non-GPU deployment. + +To build with GPU-oriented extras and request a GPU-enabled container runtime: + +```bash +docker compose -f docker-compose.yaml -f docker-compose.gpu.yaml up --build +``` + #### 3b. Run manually in two different tabs of your terminal (recommended) - **Start the server:** @@ -117,6 +125,16 @@ docker compose up --build python main.py ``` + For a GPU-oriented install on Linux, use: + ```bash + pip install -r requirements.gpu.txt + ``` + + To force local inference onto CPU even on a machine with an accelerator, set: + ```bash + export RIVERST_COMPUTE_DEVICE=cpu + ``` + - **Start the client:** ```bash cd src/client/react diff --git a/docker-compose.gpu.yaml b/docker-compose.gpu.yaml new file mode 100644 index 0000000..0f2fbb0 --- /dev/null +++ b/docker-compose.gpu.yaml @@ -0,0 +1,8 @@ +services: + server: + build: + args: + RIVERST_DEPLOYMENT_TARGET: gpu + environment: + RIVERST_COMPUTE_DEVICE: auto + gpus: all diff --git a/docker-compose.yaml b/docker-compose.yaml index 6afb22e..4118083 100644 --- a/docker-compose.yaml +++ b/docker-compose.yaml @@ -5,11 +5,15 @@ services: build: context: ./src/server dockerfile: Dockerfile + args: + RIVERST_DEPLOYMENT_TARGET: ${RIVERST_DEPLOYMENT_TARGET:-cpu} container_name: riverst_server ports: - "7860:7860" env_file: - ./src/server/.env + environment: + RIVERST_COMPUTE_DEVICE: ${RIVERST_COMPUTE_DEVICE:-auto} volumes: - ./src/server/sessions:/app/sessions diff --git a/docs/gpu-cpu-deployment-plan.md b/docs/gpu-cpu-deployment-plan.md new file mode 100644 index 0000000..e4b54c9 --- /dev/null +++ b/docs/gpu-cpu-deployment-plan.md @@ -0,0 +1,43 @@ +# GPU and Non-GPU Deployment Plan + +## Goal + +Allow Riverst to deploy in either of these modes without code changes: + +- [x] CPU / non-GPU deployment +- [x] GPU-capable deployment + +## Findings + +After reviewing the repository, no server function strictly requires a GPU to execute. The GPU-sensitive paths are all local-model or local-inference paths where acceleration improves latency or throughput. + +### GPU-sensitive functions and classes + +- [x] `src/server/bot/utils/device_utils.py:get_best_device` +- [x] `src/server/bot/processors/audio/resampling_helper.py:AudioResamplingHelper._torchaudio_resample` +- [x] `src/server/bot/processors/speech/lipsync_processor.py:predict_phonemes_from_waveform` +- [x] `src/server/bot/processors/speech/lipsync_processor.py:load_cupe_model` +- [x] `src/server/bot/processors/speech/lipsync_processor.py:LipsyncProcessor.__init__` +- [x] `src/server/bot/processors/speech/lipsync_processor.py:LipsyncProcessor._warm_up` +- [x] `src/server/bot/processors/speech/lipsync_processor.py:LipsyncProcessor._preprocess_audio` +- [x] `src/server/bot/processors/speech/lipsync_processor.py:LipsyncProcessor._run_lipsync` +- [x] `src/server/bot/core/component_factory.py:BotComponentFactory._build_stt_service` +- [x] `src/server/bot/core/component_factory.py:BotComponentFactory._build_tts_service` +- [x] `src/server/bot/processors/video/processor.py:VideoProcessor.__init__` +- [x] `src/server/bot/processors/video/processor.py:VideoProcessor._run_pose_in_background` +- [x] `src/server/bot/processors/audio/analyzer.py:AudioAnalyzer.analyze_audio` + +## Decisions + +- [x] Add a runtime device policy via `RIVERST_COMPUTE_DEVICE=auto|cpu` +- [x] Add a Docker build target via `RIVERST_DEPLOYMENT_TARGET=cpu|gpu` +- [x] Add a GPU compose override for `gpus: all` +- [x] Keep CPU deployments functional when ONNX export dependencies are unavailable +- [x] Add a small test suite for the runtime device policy + +## Implementation Notes + +- CPU deployments now default to `requirements.txt` +- GPU deployments install `requirements.gpu.txt` +- Runtime device selection is centralized in `bot.utils.device_utils` +- Pose processing falls back to the PyTorch YOLO model if ONNX export is unavailable diff --git a/notes/first_steps_to_deploy.md b/notes/first_steps_to_deploy.md index 1d7f5fb..abee934 100644 --- a/notes/first_steps_to_deploy.md +++ b/notes/first_steps_to_deploy.md @@ -6,7 +6,8 @@ This guide explains step-by-step how to deploy the Riverst project on AWS using ## 1. Create an AWS EC2 instance -- Choose `g4dn.xlarge` as the instance type or more powerful (you may need a more powerful gpu if you want to run models locally without experiencing any unpleasant delay). +- For GPU deployment, choose `g4dn.xlarge` as the instance type or more powerful. +- For non-GPU deployment, choose a general-purpose Ubuntu instance sized for your traffic and disable accelerator usage with `RIVERST_COMPUTE_DEVICE=cpu`. - Use a Linux (Ubuntu) machine. - Set the storage volume to **64 GB** (or larger). - Open these ports in the security group: @@ -37,6 +38,8 @@ ssh -i your-key.pem ubuntu@your-elastic-ip ## 4. Install NVIDIA drivers +Only required for GPU deployment. + ```bash sudo apt update sudo apt install ubuntu-drivers-common @@ -160,6 +163,7 @@ cd riverst/src/server pip install -r requirements.txt cp .env.example .env # Edit `.env` to configure your settings following .env.example +# For CPU-only deployment add: RIVERST_COMPUTE_DEVICE=cpu # Edit the src/server/config/authorized_users.json # [In a tmux tab] diff --git a/src/server/Dockerfile b/src/server/Dockerfile index 081408b..e173d6f 100644 --- a/src/server/Dockerfile +++ b/src/server/Dockerfile @@ -4,6 +4,9 @@ FROM python:3.12 # Set working directory WORKDIR /app +ARG RIVERST_DEPLOYMENT_TARGET=cpu +ENV RIVERST_DEPLOYMENT_TARGET=${RIVERST_DEPLOYMENT_TARGET} + # System packages and Rust installation RUN apt-get update && apt-get install -y --no-install-recommends \ git build-essential g++ curl \ @@ -14,8 +17,13 @@ RUN apt-get update && apt-get install -y --no-install-recommends \ # Copy application code COPY . . -# Install any needed packages specified in requirements.txt -RUN pip install --no-cache-dir -r requirements.txt +# Install dependencies for the selected deployment target +RUN case "$RIVERST_DEPLOYMENT_TARGET" in \ + cpu) requirements_file="requirements.txt" ;; \ + gpu) requirements_file="requirements.gpu.txt" ;; \ + *) echo "Unsupported RIVERST_DEPLOYMENT_TARGET: $RIVERST_DEPLOYMENT_TARGET" >&2; exit 1 ;; \ + esac \ + && pip install --no-cache-dir -r "$requirements_file" # Expose port diff --git a/src/server/README.md b/src/server/README.md index a5610b3..0dc8fed 100644 --- a/src/server/README.md +++ b/src/server/README.md @@ -5,8 +5,8 @@ 1. Set up and activate your virtual environment: ```bash -conda create -n riverst python=3.11 -y -conda activate riverst +python -m venv .venv +source .venv/bin/activate ``` 2. Install python dependencies @@ -14,9 +14,15 @@ conda activate riverst pip install -r requirements.txt ``` +For a GPU-oriented Linux deployment, install: +```bash +pip install -r requirements.gpu.txt +``` + 3. Copy `env.example` to `.env` and configure params: - Set your `OPENAI_API_KEY` for LLM and TTS services - Set your env variables (you can follow the instructions in the `.env.example` file) + - Set `RIVERST_COMPUTE_DEVICE=cpu` if you want to disable GPU/MPS usage at runtime **Note**: Not all API KEYS are strictly required. Only if you want to use a remote service, you need to expose the corresponding API KEY @@ -56,3 +62,10 @@ python main.py docker build --no-cache -t fastapi-server . docker run -p 7860:7860 --env-file .env fastapi-server ``` + +To build a GPU-oriented image instead: + +```bash +docker build --no-cache --build-arg RIVERST_DEPLOYMENT_TARGET=gpu -t fastapi-server . +docker run --gpus all -p 7860:7860 --env-file .env fastapi-server +``` diff --git a/src/server/bot/processors/video/processor.py b/src/server/bot/processors/video/processor.py index f2f0234..5175545 100644 --- a/src/server/bot/processors/video/processor.py +++ b/src/server/bot/processors/video/processor.py @@ -41,9 +41,7 @@ def __init__( if self.enable_pose: logger.info("Initializing YOLO pose model...") - yolo_model = YOLO("yolo11n-pose.pt") - yolo_model.export(format="onnx") # Creates 'yolo11n-pose.onnx' - self.pose_inferencer = YOLO("yolo11n-pose.onnx") + self.pose_inferencer = self._load_pose_inferencer() dummy_img = np.random.randint( 0, 255, (camera_out_height, camera_out_width, 3), dtype=np.uint8 @@ -51,6 +49,20 @@ def __init__( _ = self.pose_inferencer(dummy_img) logger.info("YOLO warmed up.") + def _load_pose_inferencer(self): + """Prefer ONNX when available, but keep CPU deployments functional without it.""" + yolo_model = YOLO("yolo11n-pose.pt") + try: + yolo_model.export(format="onnx") # Creates 'yolo11n-pose.onnx' + logger.info("Using ONNX-backed YOLO pose inferencer.") + return YOLO("yolo11n-pose.onnx") + except Exception as exc: + logger.warning( + "ONNX pose export unavailable, falling back to PyTorch YOLO pose model: " + f"{exc}" + ) + return yolo_model + async def _run_pose_in_background(self, img: np.ndarray) -> None: if self._pose_lock.locked(): return diff --git a/src/server/bot/utils/device_utils.py b/src/server/bot/utils/device_utils.py index 77a3d7d..00a331e 100644 --- a/src/server/bot/utils/device_utils.py +++ b/src/server/bot/utils/device_utils.py @@ -1,27 +1,60 @@ """Device detection utilities for PyTorch.""" +import os + import torch from torch import device as TorchDevice +COMPUTE_DEVICE_ENV_VAR = "RIVERST_COMPUTE_DEVICE" +DEPLOYMENT_TARGET_ENV_VAR = "RIVERST_DEPLOYMENT_TARGET" +DEFAULT_DEVICE_ORDER = ("cuda", "mps", "cpu") +VALID_COMPUTE_DEVICE_POLICIES = {"auto", "cpu"} -def get_best_device(options=["cuda", "mps", "cpu"]) -> TorchDevice: - """Returns the "best" available torch device according to the following strategy: - 1. Use CUDA if available. - 2. If not, use MPS (Metal Performance Shaders) if available. - 3. Otherwise, fall back to CPU. +def get_compute_device_policy() -> str: + """Return the configured runtime compute device policy. - Returns: - torch.device: The best available torch device ('cuda', 'mps', or 'cpu'). + `auto` prefers accelerated backends when available. + `cpu` forces all local model execution onto CPU. """ - if torch.cuda.is_available() and "cuda" in options: + policy = os.getenv(COMPUTE_DEVICE_ENV_VAR, "auto").strip().lower() + if policy not in VALID_COMPUTE_DEVICE_POLICIES: + raise ValueError( + f"Invalid {COMPUTE_DEVICE_ENV_VAR}={policy!r}. " + f"Expected one of {sorted(VALID_COMPUTE_DEVICE_POLICIES)}." + ) + return policy + + +def get_deployment_target() -> str: + """Return the configured deployment target used by packaging/docs.""" + return os.getenv(DEPLOYMENT_TARGET_ENV_VAR, "cpu").strip().lower() + + +def get_best_device(options=None) -> TorchDevice: + """Return the best available torch device for the current runtime policy.""" + requested_options = list(options or DEFAULT_DEVICE_ORDER) + policy = get_compute_device_policy() + + if policy == "cpu": + if "cpu" not in requested_options: + raise ValueError( + "CPU-only runtime policy requested, but 'cpu' is not in the " + f"allowed device options: {requested_options}." + ) + return torch.device("cpu") + + if torch.cuda.is_available() and "cuda" in requested_options: return torch.device("cuda") - elif ( + + if ( hasattr(torch.backends, "mps") and torch.backends.mps.is_available() - and "mps" in options + and "mps" in requested_options ): return torch.device("mps") - else: - # Fallback to CPU + + if "cpu" in requested_options: return torch.device("cpu") + + raise ValueError(f"No supported device found for options: {requested_options}.") diff --git a/src/server/env.example b/src/server/env.example index 75130c6..2771ec7 100644 --- a/src/server/env.example +++ b/src/server/env.example @@ -3,6 +3,7 @@ # ------------------------- OPENAI_API_KEY="" # OpenAI API key for language models HF_TOKEN="" # Hugging Face access token for downloading models (required for many models, including ASR Whisper) +RIVERST_COMPUTE_DEVICE="auto" # Runtime device policy for local models: "auto" prefers accelerators, "cpu" disables GPU/MPS usage # ------------------------- # TURN Server Configuration (for WebRTC) - You can ignore these for local testing diff --git a/src/server/requirements.gpu.txt b/src/server/requirements.gpu.txt new file mode 100644 index 0000000..52463dc --- /dev/null +++ b/src/server/requirements.gpu.txt @@ -0,0 +1,4 @@ +-r requirements.txt +onnx>=1.12.0 +onnxslim>=0.1.67 +onnxruntime-gpu diff --git a/src/server/tests/test_device_utils.py b/src/server/tests/test_device_utils.py new file mode 100644 index 0000000..46dcfd4 --- /dev/null +++ b/src/server/tests/test_device_utils.py @@ -0,0 +1,71 @@ +import os +import sys +import unittest +import types +import importlib +from unittest import mock + + +def make_fake_torch(cuda_available=False, mps_available=False): + torch_module = types.ModuleType("torch") + torch_module.device = lambda name: name + torch_module.cuda = types.SimpleNamespace(is_available=lambda: cuda_available) + torch_module.backends = types.SimpleNamespace( + mps=types.SimpleNamespace(is_available=lambda: mps_available) + ) + return torch_module + + +class DeviceUtilsTest(unittest.TestCase): + def setUp(self): + self.modules_backup = { + "torch": sys.modules.get("torch"), + "bot.utils.device_utils": sys.modules.get("bot.utils.device_utils"), + "bot.utils": sys.modules.get("bot.utils"), + "bot": sys.modules.get("bot"), + } + sys.modules["torch"] = make_fake_torch() + self.device_utils = importlib.import_module("bot.utils.device_utils") + + def tearDown(self): + for module_name, module in self.modules_backup.items(): + if module is None: + sys.modules.pop(module_name, None) + else: + sys.modules[module_name] = module + + def test_default_policy_is_auto(self): + with mock.patch.dict(os.environ, {}, clear=False): + os.environ.pop(self.device_utils.COMPUTE_DEVICE_ENV_VAR, None) + self.assertEqual(self.device_utils.get_compute_device_policy(), "auto") + + def test_cpu_policy_forces_cpu(self): + with mock.patch.dict( + os.environ, + {self.device_utils.COMPUTE_DEVICE_ENV_VAR: "cpu"}, + clear=False, + ): + device = self.device_utils.get_best_device(options=["cuda", "mps", "cpu"]) + self.assertEqual(str(device), "cpu") + + def test_invalid_policy_raises(self): + with mock.patch.dict( + os.environ, + {self.device_utils.COMPUTE_DEVICE_ENV_VAR: "invalid"}, + clear=False, + ): + with self.assertRaises(ValueError): + self.device_utils.get_compute_device_policy() + + def test_cpu_policy_requires_cpu_option(self): + with mock.patch.dict( + os.environ, + {self.device_utils.COMPUTE_DEVICE_ENV_VAR: "cpu"}, + clear=False, + ): + with self.assertRaises(ValueError): + self.device_utils.get_best_device(options=["cuda"]) + + +if __name__ == "__main__": + unittest.main()