diff --git a/server/backend/linux.py b/server/backend/linux.py index 431381a..32b1dfb 100644 --- a/server/backend/linux.py +++ b/server/backend/linux.py @@ -53,10 +53,12 @@ logger = logging.getLogger("app") from typing import Any, Dict, List, Optional, Union, Iterator import httpx from pathlib import Path +from ..config import get_llama_config _model_cache: Dict[str, LlamaRunner] = {} _default_max_tokens: Optional[int] = None # Use dynamic model-aware limits by default _current_model_path: Optional[str] = None +_current_llama_config: Dict[str, Any] | None = None # Store generated responses for follow-up support (previous_response_id) _responses: Dict[str, ResponsesResponse] = {} @@ -64,13 +66,20 @@ _responses: Dict[str, ResponsesResponse] = {} def get_or_load_model( - model_spec: str, model_cache_path: str | None = None, verbose: bool = True + model_spec: str, + model_cache_path: str | None = None, + verbose: bool = True, ) -> LlamaRunner: """Get model from cache or load it if not cached.""" - global _model_cache, _current_model_path + global _model_cache, _current_model_path, _current_llama_config model_name = model_spec + llama_config = get_llama_config() - if model_cache_path is None and _current_model_path in _model_cache: + if ( + model_cache_path is None + and _current_model_path in _model_cache + and _current_llama_config == llama_config + ): logger.info(f"Model {model_name} already in memory") return _model_cache[_current_model_path] @@ -100,7 +109,10 @@ def get_or_load_model( ) # Check if we need to load a different model - if _current_model_path != model_path_str: + if ( + _current_model_path != model_path_str + or _current_llama_config != llama_config + ): # Proactively clean up any previously loaded runner to release memory if _model_cache: try: @@ -117,11 +129,14 @@ def get_or_load_model( print(f"Loading model: {model_name}") logger.info(f"Loading model: {model_name}") - runner = LlamaRunner(model_path_str, verbose=verbose) + runner = LlamaRunner( + model_path_str, verbose=verbose, llama_config=llama_config + ) runner.load_model() _model_cache[model_path_str] = runner _current_model_path = model_path_str + _current_llama_config = llama_config else: logger.info(f"Model {model_name} already in memory") diff --git a/server/backend/llama_cpp_runner.py b/server/backend/llama_cpp_runner.py index 03ecfd7..696fca8 100644 --- a/server/backend/llama_cpp_runner.py +++ b/server/backend/llama_cpp_runner.py @@ -78,17 +78,20 @@ def _preload_cuda_runtime_libs() -> None: ctypes.CDLL(str(lib_path), mode=ctypes.RTLD_GLOBAL) -def get_model_context_length_gguf(model_path: str) -> int: +def get_model_context_length_gguf( + model_path: str, configured_max_ctx: int | None = None +) -> int: """Extract context length from config.json alongside the GGUF file. Args: model_path: Path to the model directory Returns: - Maximum context length for the model, capped by TILES_LLAMA_CPP_MAX_CTX - or 30000 by default. If model metadata is unavailable, use that cap. + Maximum context length for the model, capped by the configured llama + context or 30000 by default. If model metadata is unavailable, use + that cap. """ - max_ctx = int(os.environ.get("TILES_LLAMA_CPP_MAX_CTX", "30000")) + max_ctx = configured_max_ctx or 30000 config_path = os.path.join(model_path, "config.json") try: with open(config_path) as f: @@ -113,17 +116,6 @@ def get_model_context_length_gguf(model_path: str) -> int: return max_ctx -def _get_env_int(name: str, default: int) -> int: - return int(os.environ.get(name, str(default))) - - -def _get_env_bool(name: str, default: bool) -> bool: - value = os.environ.get(name) - if value is None: - return default - return value.lower() in {"1", "true", "yes", "on"} - - class LlamaRunner: """Direct llama.cpp model runner with streaming and interactive capabilities. @@ -144,7 +136,12 @@ class LlamaRunner: verbose: bool _model_loaded: bool - def __init__(self, model_path: str, verbose: bool = False): + def __init__( + self, + model_path: str, + verbose: bool = False, + llama_config: dict | None = None, + ): """Initialize the runner with a model. Args: @@ -153,6 +150,7 @@ class LlamaRunner: """ self.model_path = Path(model_path) self.model = None + self.llama_config = llama_config or {} # Stop-token state -- populated in _extract_stop_tokens() self._stop_tokens: list[str] | None = None @@ -224,12 +222,19 @@ class LlamaRunner: if self.verbose: print(f"Using GGUF file: {gguf_file}") + configured_context_length = self.llama_config.get("context_length") requested_context_length = get_model_context_length_gguf( - str(self.model_path) + str(self.model_path), configured_context_length ) - n_gpu_layers = _get_env_int("TILES_LLAMA_CPP_N_GPU_LAYERS", 10) - offload_kqv = _get_env_bool("TILES_LLAMA_CPP_OFFLOAD_KQV", True) - n_batch = _get_env_int("TILES_LLAMA_CPP_N_BATCH", 512) + n_gpu_layers = self.llama_config.get("gpu_layers") + if n_gpu_layers is None: + n_gpu_layers = 10 + offload_kqv = self.llama_config.get("offload_kqv") + if offload_kqv is None: + offload_kqv = True + n_batch = self.llama_config.get("batch_size") + if n_batch is None: + n_batch = 512 self._context_length = requested_context_length self.model = Llama( diff --git a/server/backend/mlx.py b/server/backend/mlx.py index 246c597..f1aad7b 100644 --- a/server/backend/mlx.py +++ b/server/backend/mlx.py @@ -48,7 +48,9 @@ _current_model_path: Optional[str] = None def get_or_load_model( - model_spec: str, model_cache_path: str | None = None, verbose: bool = True + model_spec: str, + model_cache_path: str | None = None, + verbose: bool = True, ) -> MLXRunner: """Get model from cache or load it if not cached.""" global _model_cache, _current_model_path diff --git a/server/config.py b/server/config.py index 2575cdc..781b59d 100644 --- a/server/config.py +++ b/server/config.py @@ -1,7 +1,28 @@ -from pathlib import Path import os +import httpx +from pydantic import BaseModel + PORT = 6969 +DAEMON_PORT = 1729 MODEL_ID = "driaforall/mem-agent" MEMORY_PATH = os.path.expanduser("~") + "/tiles_memory" + + +class LlamaConfig(BaseModel): + context_length: int | None = None + gpu_layers: int | None = None + offload_kqv: bool | None = None + batch_size: int | None = None + + +def get_llama_config() -> dict: + try: + response = httpx.get(f"http://127.0.0.1:{DAEMON_PORT}/config", timeout=5) + response.raise_for_status() + config = response.json() + except httpx.HTTPError: + return {} + + return LlamaConfig(**config.get("llama", {})).model_dump(exclude_none=True) diff --git a/server/tests/test_llama_cpp_runner.py b/server/tests/test_llama_cpp_runner.py index de20b74..515bfdb 100644 --- a/server/tests/test_llama_cpp_runner.py +++ b/server/tests/test_llama_cpp_runner.py @@ -9,8 +9,6 @@ from openai_harmony import ( from server.backend.llama_cpp_runner import ( LlamaRunner, - _get_env_bool, - _get_env_int, get_model_context_length_gguf, ) from server.schemas import ToolCallStart @@ -87,41 +85,26 @@ def test_gpt_streaming_rejects_prompt_that_exceeds_context_and_resets_model(): assert model.was_reset -def test_gguf_context_length_caps_at_30000_by_default(tmp_path, monkeypatch): - monkeypatch.delenv("TILES_LLAMA_CPP_MAX_CTX", raising=False) +def test_gguf_context_length_caps_at_30000_by_default(tmp_path): (tmp_path / "config.json").write_text('{"context_length": 131072}') assert get_model_context_length_gguf(str(tmp_path)) == 30000 -def test_gguf_context_length_uses_env_cap(tmp_path, monkeypatch): +def test_gguf_context_length_uses_configured_cap(tmp_path): (tmp_path / "config.json").write_text('{"context_length": 131072}') - monkeypatch.setenv("TILES_LLAMA_CPP_MAX_CTX", "12000") - assert get_model_context_length_gguf(str(tmp_path)) == 12000 + assert get_model_context_length_gguf(str(tmp_path), 12000) == 12000 -def test_gguf_context_length_uses_env_when_config_missing(tmp_path, monkeypatch): - monkeypatch.setenv("TILES_LLAMA_CPP_MAX_CTX", "2048") +def test_gguf_context_length_uses_configured_cap_when_config_missing(tmp_path): + assert get_model_context_length_gguf(str(tmp_path), 2048) == 2048 - assert get_model_context_length_gguf(str(tmp_path)) == 2048 - -def test_gguf_context_length_uses_env_when_config_invalid(tmp_path, monkeypatch): +def test_gguf_context_length_uses_configured_cap_when_config_invalid(tmp_path): (tmp_path / "config.json").write_text("{") - monkeypatch.setenv("TILES_LLAMA_CPP_MAX_CTX", "3072") - - assert get_model_context_length_gguf(str(tmp_path)) == 3072 - - -def test_llama_cpp_env_helpers(monkeypatch): - monkeypatch.setenv("TILES_TEST_INT", "12") - monkeypatch.setenv("TILES_TEST_BOOL", "false") - assert _get_env_int("TILES_TEST_INT", 3) == 12 - assert _get_env_int("TILES_MISSING_INT", 3) == 3 - assert _get_env_bool("TILES_TEST_BOOL", True) is False - assert _get_env_bool("TILES_MISSING_BOOL", True) is True + assert get_model_context_length_gguf(str(tmp_path), 3072) == 3072 class TokenizingFakeModel: diff --git a/tiles/src/daemon.rs b/tiles/src/daemon.rs index 1a2f0a1..db3e8be 100644 --- a/tiles/src/daemon.rs +++ b/tiles/src/daemon.rs @@ -25,7 +25,7 @@ use tokio::sync::oneshot::{self, Receiver, Sender}; use crate::{ core::account::atproto::AtCallbackParams, - utils::config::{ConfigProvider, DefaultProvider, get_model_cache}, + utils::config::{ConfigProvider, DefaultProvider, get_config_json, get_model_cache}, }; struct AppState { @@ -132,6 +132,7 @@ pub async fn start_server(port: Option) -> Result<()> { let shared_state = Arc::new(state); let app = Router::new() .route("/", get(root)) + .route("/config", get(get_config)) .route("/shutdown", get(shutdown)) .route("/model-cache-path", get(get_model_cache_path)) .with_state(shared_state); @@ -220,6 +221,12 @@ async fn get_model_cache_path( } } +async fn get_config(State(_state): State>) -> Result { + get_config_json() + .and_then(|config| serde_json::to_string(&config).map_err(Into::into)) + .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR) +} + async fn stop_server(port: Option) -> Result<()> { let dyn_port = get_port(port); let client = Client::new(); diff --git a/tiles/src/main.rs b/tiles/src/main.rs index 20d3f13..0930b98 100644 --- a/tiles/src/main.rs +++ b/tiles/src/main.rs @@ -14,7 +14,7 @@ use tiles::{ }, daemon::{start_cmd, start_server, stop_cmd}, repl::{self, RunArgs}, - utils::installer, + utils::{config::LlamaConfig, installer}, }; use crate::commands::{set_inference_config_to_run_bg, show_peers, unlink_peer}; @@ -197,6 +197,37 @@ struct RunFlags { // Use PI repl instead of Tiles #[arg(short = 'p', long, hide = true)] pi: bool, + + /// Context window for local llama.cpp inference + #[arg(long)] + context_length: Option, + + /// Number of model layers to offload to GPU for llama.cpp + #[arg(long)] + gpu_layers: Option, + + /// Offload K/Q/V attention operations for llama.cpp + #[arg(long, num_args = 0..=1, default_missing_value = "true")] + offload_kqv: Option, + + /// Prompt processing batch size for llama.cpp + #[arg(long)] + batch_size: Option, +} + +fn llama_config_from_flags(flags: &RunFlags) -> Option { + let config = LlamaConfig { + context_length: flags.context_length, + gpu_layers: flags.gpu_layers, + offload_kqv: flags.offload_kqv, + batch_size: flags.batch_size, + }; + + if config.is_empty() { + None + } else { + Some(config) + } } #[derive(Debug, Args)] @@ -328,6 +359,7 @@ pub async fn main() -> Result<(), Box> { relay_count: cli.flags.relay_count, memory: cli.flags.memory, pi: cli.flags.pi, + llama_config: llama_config_from_flags(&cli.flags), }; commands::run_setup_for_ftue(&run_args) @@ -359,6 +391,7 @@ pub async fn main() -> Result<(), Box> { relay_count: flags.relay_count, memory: flags.memory, pi: flags.pi, + llama_config: llama_config_from_flags(&flags), }; commands::run_setup_for_ftue(&run_args) .await diff --git a/tiles/src/repl.rs b/tiles/src/repl.rs index 2eb29c2..c7887df 100644 --- a/tiles/src/repl.rs +++ b/tiles/src/repl.rs @@ -6,8 +6,8 @@ use crate::core::chats::{ }; use crate::core::storage::db::Dbconn; use crate::utils::config::{ - ConfigProvider, DefaultProvider, create_pi_provider_config, get_inference_config, - get_memory_path, get_model_cache, update_current_model, + ConfigProvider, DefaultProvider, LlamaConfig, create_pi_provider_config, get_inference_config, + get_memory_path, get_model_cache, update_current_model, update_llama_config, }; use crate::utils::hf_model_downloader::*; use anyhow::{Context, Result, anyhow}; @@ -65,6 +65,7 @@ pub struct RunArgs { pub relay_count: u32, pub memory: bool, // Future flags go here pub pi: bool, + pub llama_config: Option, } #[derive(Clone, Debug)] pub struct ChatResponse { @@ -226,23 +227,36 @@ impl From for String { const PY_PORT: u32 = 6969; pub async fn run(run_args: RunArgs, db_conn: &Dbconn) -> Result<()> { - let default_modelfile_path = get_default_modelfile(run_args.memory)?; - let default_modelfile = - tilekit::modelfile::parse_from_file(default_modelfile_path.to_str().unwrap()).unwrap(); - let modelfile_parse_result = if let Some(modelfile_str) = &run_args.modelfile_path { - tilekit::modelfile::parse_from_file(modelfile_str.as_str()) + let (modelfile, default_modelfile) = if let Some(modelfile_str) = &run_args.modelfile_path { + let modelfile = match tilekit::modelfile::parse_from_file(modelfile_str.as_str()) { + Ok(mf) => mf, + Err(_err) => { + println!("Invalid Modelfile"); + return Ok(()); + } + }; + let default_modelfile = get_default_modelfile(run_args.memory) + .ok() + .and_then(|path| tilekit::modelfile::parse_from_file(path.to_str()?).ok()) + .unwrap_or_else(|| modelfile.clone()); + (modelfile, default_modelfile) } else { - Err("NOT PROVIDED".to_string()) + let default_modelfile_path = get_default_modelfile(run_args.memory)?; + let default_modelfile = + match tilekit::modelfile::parse_from_file(default_modelfile_path.to_str().unwrap()) { + Ok(mf) => mf, + Err(_err) => { + println!("Invalid default Modelfile"); + return Ok(()); + } + }; + (default_modelfile.clone(), default_modelfile) }; - let modelfile = match modelfile_parse_result { - Ok(mf) => mf, - Err(err) if err == "NOT PROVIDED" => default_modelfile.clone(), - Err(_err) => { - println!("Invalid Modelfile"); - return Ok(()); - } - }; + if modelfile.from.is_none() { + println!("Invalid Modelfile"); + return Ok(()); + } run_model_with_server(modelfile, default_modelfile, &run_args, db_conn).await } @@ -523,6 +537,9 @@ async fn run_model_with_server( } // loading the model from mem-agent via daemon server let memory_path = get_memory_path().context("Setting/Retrieving memory_path failed")?; + if let Some(llama_config) = &run_args.llama_config { + update_llama_config(llama_config.clone()).context("Failed to update llama config")?; + } match load_model(&modelfile, &default_modelfile, &memory_path, 0).await { Ok(_) => start_repl(&modelfile, run_args, db_conn) .await diff --git a/tiles/src/utils/config.rs b/tiles/src/utils/config.rs index c18da14..4233c79 100644 --- a/tiles/src/utils/config.rs +++ b/tiles/src/utils/config.rs @@ -45,6 +45,23 @@ pub struct InferenceConfig { pub daemon: bool, } +#[derive(Clone, Serialize, Deserialize, Debug, Default, PartialEq)] +pub struct LlamaConfig { + pub context_length: Option, + pub gpu_layers: Option, + pub offload_kqv: Option, + pub batch_size: Option, +} + +impl LlamaConfig { + pub fn is_empty(&self) -> bool { + self.context_length.is_none() + && self.gpu_layers.is_none() + && self.offload_kqv.is_none() + && self.batch_size.is_none() + } +} + #[derive(Serialize, Deserialize, Debug)] struct RootConfig { #[serde(rename = "root-user")] @@ -52,6 +69,7 @@ struct RootConfig { pub data: Option, pub model: Option, pub inference: Option, + pub llama: Option, } #[derive(Serialize, Deserialize, Clone)] @@ -467,19 +485,10 @@ fn do_update_current_model(config: &mut RootConfig, model_name: &str) -> Result< Ok(()) } -fn get_env_u32(name: &str) -> Option { - env::var(name).ok().and_then(|value| value.parse().ok()) -} - -fn get_pi_context_window_with_env(get_env: &F) -> Option -where - F: Fn(&str) -> Option, -{ - get_env("TILES_LLAMA_CPP_MAX_CTX").and_then(|value| value.parse().ok()) -} - fn get_pi_context_window() -> Option { - get_env_u32("TILES_LLAMA_CPP_MAX_CTX") + get_llama_config() + .ok() + .and_then(|config| config.context_length) } fn get_pi_max_tokens(context_window: Option) -> Option { @@ -487,18 +496,14 @@ fn get_pi_max_tokens(context_window: Option) -> Option { } pub fn create_pi_provider_config(model_name: &str, enpoint_base_url: &str) -> Result { - create_pi_provider_config_with_env(model_name, enpoint_base_url, &|name| env::var(name).ok()) + create_pi_provider_config_with_context(model_name, enpoint_base_url, get_pi_context_window()) } -fn create_pi_provider_config_with_env( +fn create_pi_provider_config_with_context( model_name: &str, enpoint_base_url: &str, - get_env: &F, -) -> Result -where - F: Fn(&str) -> Option, -{ - let context_window = get_pi_context_window_with_env(get_env); + context_window: Option, +) -> Result { let max_tokens = get_pi_max_tokens(context_window); let provider_config = PiProviderConfig { api: String::from("openai-responses"), @@ -561,6 +566,37 @@ pub fn update_inference_config(config: InferenceConfig) -> Result<()> { save_root_config(&root_config) } +pub fn get_llama_config() -> Result { + let root_config = get_or_create_root_config()?; + Ok(root_config.llama.unwrap_or_default()) +} + +pub fn get_config_json() -> Result { + let root_config = get_or_create_root_config()?; + serde_json::to_value(root_config).map_err(Into::::into) +} + +pub fn update_llama_config(config: LlamaConfig) -> Result<()> { + let mut root_config = get_or_create_root_config()?; + let mut llama_config = root_config.llama.unwrap_or_default(); + + if config.context_length.is_some() { + llama_config.context_length = config.context_length; + } + if config.gpu_layers.is_some() { + llama_config.gpu_layers = config.gpu_layers; + } + if config.offload_kqv.is_some() { + llama_config.offload_kqv = config.offload_kqv; + } + if config.batch_size.is_some() { + llama_config.batch_size = config.batch_size; + } + + root_config.llama = Some(llama_config); + save_root_config(&root_config) +} + #[cfg(test)] mod tests { @@ -651,14 +687,11 @@ mod tests { } #[test] - fn test_pi_provider_uses_llama_cpp_context_env() { - let config_str = create_pi_provider_config_with_env( + fn test_pi_provider_uses_llama_context_length() { + let config_str = create_pi_provider_config_with_context( "unsloth/gpt-oss-20b-GGUF", "http://127.0.0.1:6969/v1", - &|name| match name { - "TILES_LLAMA_CPP_MAX_CTX" => Some("12000".to_owned()), - _ => None, - }, + Some(12_000), ) .unwrap(); let config: Value = serde_json::from_str(&config_str).unwrap();