Something went wrong. Try again.
a tiny pet with an even tinier brain, living in your browser
Something went wrong. Try again.
4.5 kB · 115 lines
Rust
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116//! Required env: SUBTOKENLING_LISTEN, SUBTOKENLING_DB, SUBTOKENLING_MODEL_DIR,//! SUBTOKENLING_MODEL_ID, SUBTOKENLING_ALLOWED_ORIGIN, ORT_DYLIB_PATH.//! Optional: SUBTOKENLING_MODEL_DTYPE=q4|q8|fp32 (default q4); SUBTOKENLING_THREADS=4;//! SUBTOKENLING_GPU_DEVICE (logical CUDA index after CUDA_VISIBLE_DEVICES,//! requires a CUDA-enabled ORT library); SUBTOKENLING_MAX_PETS=64;//! SUBTOKENLING_WEATHER_SEED=""; SUBTOKENLING_FEATURES="[]" (JSON feature names).//! No state, access keys, prompts or generated text are logged.use anyhow::{Context, Result, ensure};use axum::http::{HeaderValue, Method, header};use std::{ env, path::Path, sync::{Arc, Mutex},};use tokenling_server::{ brain::Brain, host::{Host, Worker}, store::Store,};use tower_http::{cors::CorsLayer, set_header::SetResponseHeaderLayer};fn required(name: &str) -> Result<String> { env::var(name).with_context(|| format!("set {name}"))}#[tokio::main]async fn main() -> Result<()> { let listen: std::net::SocketAddr = required("SUBTOKENLING_LISTEN")?.parse()?; let db = required("SUBTOKENLING_DB")?; let model_dir = required("SUBTOKENLING_MODEL_DIR")?; let model_id = required("SUBTOKENLING_MODEL_ID")?; let dtype = env::var("SUBTOKENLING_MODEL_DTYPE") .unwrap_or_else(|_| "q4".into()) .parse()?; let origins: Vec<HeaderValue> = required("SUBTOKENLING_ALLOWED_ORIGIN")? .split(',') .map(|s| s.trim().parse()) .collect::<std::result::Result<_, _>>()?; ensure!( !origins.is_empty() && origins.iter().all(|s| s.as_bytes() != b"*"), "configure explicit allowed browser origins" ); let threads: usize = env::var("SUBTOKENLING_THREADS") .unwrap_or_else(|_| "4".into()) .parse()?; ensure!(threads > 0, "SUBTOKENLING_THREADS must be positive"); let max_pets: usize = env::var("SUBTOKENLING_MAX_PETS") .unwrap_or_else(|_| "64".into()) .parse()?; ensure!(max_pets > 0, "SUBTOKENLING_MAX_PETS must be positive"); let store = Store::open(Path::new(&db))?; let foreign = store.foreign_brains(&model_id)?; ensure!( foreign == 0, "{foreign} hosted pet(s) run a brain other than SUBTOKENLING_MODEL_ID; keep their model, \ or have owners bring them local before changing it (no pet was changed)" ); ort::init_from(required("ORT_DYLIB_PATH")?)? .with_name("tokenling") .commit(); let cpu = Worker::new( "cpu", Brain::load(Path::new(&model_dir), dtype, threads, None)?, ); let gpu = if let Ok(device) = env::var("SUBTOKENLING_GPU_DEVICE") { match Brain::load(Path::new(&model_dir), dtype, 1, Some(device.parse()?)) { Ok(brain) => Some(Worker::new("gpu", brain)), Err(error) => { eprintln!("optional CUDA worker unavailable: {error}"); None } } } else { None }; let host = Arc::new(Host { db: Mutex::new(store), model_id, features: serde_json::from_str( &env::var("SUBTOKENLING_FEATURES").unwrap_or_else(|_| "[]".into()), )?, max_pets, weather_seed: env::var("SUBTOKENLING_WEATHER_SEED").unwrap_or_default(), cpu, gpu, jobs: Mutex::new(std::collections::HashMap::new()), notify: tokio::sync::Notify::new(), queue: Mutex::new(Default::default()), }); let router = host .clone() .router() .layer( CorsLayer::new() .allow_origin(origins) .allow_methods([Method::GET, Method::POST, Method::DELETE]) .allow_headers([header::AUTHORIZATION, header::CONTENT_TYPE]), ) .layer(SetResponseHeaderLayer::if_not_present( header::CACHE_CONTROL, HeaderValue::from_static("no-store"), )) .layer(SetResponseHeaderLayer::if_not_present( header::HeaderName::from_static("cross-origin-resource-policy"), HeaderValue::from_static("cross-origin"), )); let listener = tokio::net::TcpListener::bind(listen).await?; let scheduler = tokio::spawn(host.run()); eprintln!("subtokenling host listening on {listen}"); axum::serve(listener, router) .with_graceful_shutdown(async { let _ = tokio::signal::ctrl_c().await; }) .await?; scheduler.abort(); Ok(())}