Skip to content

C
coderai
Commit f524118d authored by Stefy Lanza (nextime / spora )'s avatar Stefy Lanza (nextime / spora )
Browse files
Options

Parallel loading

parent 0cf7c3c7
Show whitespace changes
Inline Side-by-side
Showing with 505 additions and 162 deletions
README.md
View file @ f524118d
# CoderAI
![CoderAI](CoderAI.gif)
An OpenAI-compatible API server with web administration dashboard, supporting multiple GPU backends: NVIDIA (CUDA), AMD (Vulkan), and Intel (Vulkan). Configuration-driven architecture with per-model settings and full multi-modal support.
## Features
......@@ -72,6 +74,20 @@ Built-in multi-step pipelines callable from the API or web UI:
---
## Quick Start
```bash
git clone git@git.nexlab.net:nexlab/coderai.git
cd coderai
./build.sh all # build all backends (recommended)
source venv_all/bin/activate
python coderai # starts on http://127.0.0.1:8776
```
That's it. Open `http://127.0.0.1:8776/admin` and log in with `admin` / `admin`.
---
## Installation
### Prerequisites
......@@ -81,7 +97,7 @@ Built-in multi-step pipelines callable from the API or web UI:
- For AMD/Intel GPUs (Vulkan): Vulkan drivers and SDK
- For CPU-only: No additional requirements
### Quick Install with Build Script
### Build Script
```bash
git clone git@git.nexlab.net:nexlab/coderai.git
......@@ -159,16 +175,16 @@ python coderai --config /path/to/cfg # Custom config directory
python coderai --debug # Debug mode
```
Server starts on `http://0.0.0.0:8000`.
Server starts on `http://127.0.0.1:8776` by default.
### Access Points
| URL | Description |
|---|---|
| `http://localhost:8000/admin` | Admin dashboard |
| `http://localhost:8000/chat` | Web Studio (generation UI) |
| `http://localhost:8000/v1/*` | OpenAI-compatible API |
| `http://localhost:8000/docs` | Interactive API docs |
| `http://127.0.0.1:8776/admin` | Admin dashboard |
| `http://127.0.0.1:8776/chat` | Web Studio (generation UI) |
| `http://127.0.0.1:8776/v1/*` | OpenAI-compatible API |
| `http://127.0.0.1:8776/docs` | Interactive API docs |
Default credentials: `admin` / `admin` (prompted to change on first login).
......@@ -191,7 +207,7 @@ Config files live in `~/.coderai/` (or `--config` path):
```json
{
"server": { "host": "0.0.0.0", "port": 8000 },
"server": { "host": "127.0.0.1", "port": 8776 },
"backend": { "type": "auto" },
"models": { "default_load_mode": "ondemand" },
"offload": { "load_in_4bit": false, "flash_attention": false },
......@@ -399,10 +415,23 @@ MAX_JOBS=4 pip install flash-attn --no-build-isolation
---
## Developer
**Stefy Lanza** <stefy@nexlab.net>
## License
GNU General Public License v3.0 — see [LICENSE.md](LICENSE.md).
This program is free software: you can redistribute it and/or modify it under the terms of the GNU General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version.
## Donations
If you find CoderAI useful, consider supporting development:
- **Bitcoin (BTC):** `bc1qcpt2uutqkz4456j5r78rjm3gwq03h5fpwmcc5u`
- **Ethereum (ETH):** `0xdA6dAb526515b5cb556d20269207D43fcc760E51`
## Contributing
Merge requests welcome.
......
codai/admin/routes.py
View file @ f524118d
......@@ -1188,15 +1188,22 @@ async def api_model_disable(request: Request, username: str = Depends(require_ad
"""Remove a model from models.json (keeps it cached locally)."""
if config_manager is None:
raise HTTPException(status_code=503, detail="Config manager not initialized")
import os as _os
data = await request.json()
path = data.get("path") or data.get("model_id", "")
# Also match by bare filename so entries stored without full path are caught
fname = _os.path.basename(path) if (_os.sep in path or "/" in path) else ""
def _matches(m_entry) -> bool:
key = m_entry if isinstance(m_entry, str) else m_entry.get("path", m_entry.get("id", ""))
return key == path or (fname and _os.path.basename(key) == fname)
changed = False
for cat in ("text_models", "image_models", "audio_models",
"gguf_models", "tts_models", "vision_models", "video_models",
"audio_gen_models", "embedding_models"):
lst = config_manager.models_data.get(cat, [])
new_lst = [m for m in lst
if (m if isinstance(m, str) else m.get("path", m.get("id", ""))) != path]
new_lst = [m for m in lst if not _matches(m)]
if len(new_lst) != len(lst):
config_manager.models_data[cat] = new_lst
changed = True
......@@ -1391,13 +1398,23 @@ async def api_model_configure(request: Request, username: str = Depends(require_
if not model_types:
model_types = ["text_models"]
# Remove from all categories (handles type changes)
# Remove from all categories (handles type changes and quant switches)
# paths_to_remove: the new path + the original path before a quant change
import os as _os
paths_to_remove = {path}
orig_path = (data.get("orig_path") or "").strip()
if orig_path and orig_path != path:
paths_to_remove.add(orig_path)
# Also match by bare filename so entries stored without full path are caught
fnames_to_remove = {_os.path.basename(p) for p in paths_to_remove if _os.sep in p or "/" in p}
def _should_remove(m_entry) -> bool:
key = m_entry if isinstance(m_entry, str) else m_entry.get("path", m_entry.get("id", ""))
return key in paths_to_remove or (fnames_to_remove and _os.path.basename(key) in fnames_to_remove)
for cat in valid | {"gguf_models"}:
lst = config_manager.models_data.get(cat, [])
config_manager.models_data[cat] = [
m for m in lst
if (m if isinstance(m, str) else m.get("path", m.get("id", ""))) != path
]
config_manager.models_data[cat] = [m for m in lst if not _should_remove(m)]
# Auto-estimate used_vram_gb from file size if not provided
used_vram_gb = data.get("used_vram_gb")
......@@ -1418,7 +1435,8 @@ async def api_model_configure(request: Request, username: str = Depends(require_
for key in ("alias", "backend", "load_mode", "n_gpu_layers", "n_ctx",
"max_gpu_percent", "manual_ram_gb", "load_in_4bit", "load_in_8bit",
"flash_attention", "no_ram", "offload_strategy", "offload_dir",
"system_prompt", "parser", "tools_closer_prompt", "grammar_guided"):
"system_prompt", "parser", "tools_closer_prompt", "grammar_guided",
"max_instances", "preload_all_instances"):
if key in data:
entry[key] = data[key]
......
codai/admin/templates/models.html
View file @ f524118d
......@@ -338,6 +338,7 @@
</div>
<div class="modal-body">
<input type="hidden" id="cfg-path">
<input type="hidden" id="cfg-orig-path">
<input type="hidden" id="cfg-orig-type">
<!-- identity -->
......@@ -386,7 +387,7 @@
</div>
<div class="form-row" style="margin:0">
<label class="form-label">Load mode</label>
<select id="cfg-load-mode" class="form-input">
<select id="cfg-load-mode" class="form-input" onchange="_updatePreloadAllVisibility()">
<option value="load">Load (pre-load in VRAM)</option>
<option value="on-request">On-request (load when needed)</option>
</select>
......@@ -398,6 +399,20 @@
<input type="number" id="cfg-used-vram" class="form-input" min="0" step="0.1" placeholder="auto-estimate">
<span class="form-hint" id="cfg-used-vram-hint" style="font-size:11px;color:var(--text-3)"></span>
</div>
<div class="form-row" style="margin:0">
<label class="form-label">Max instances
<span class="muted" title="How many copies of this model can be loaded simultaneously in VRAM to serve parallel requests. Requires enough free VRAM for each additional copy."> (?)</span>
</label>
<input type="number" id="cfg-max-instances" class="form-input" min="1" step="1" value="1" oninput="_updatePreloadAllVisibility()">
<span style="font-size:11px;color:var(--text-3)">Parallel copies in VRAM</span>
</div>
</div>
<div id="cfg-preload-all-row" style="display:none;margin-top:.5rem">
<label style="display:flex;align-items:center;gap:.5rem;font-size:13px;cursor:pointer">
<input type="checkbox" id="cfg-preload-all-instances">
Load all instances at boot
<span class="muted" style="font-size:11px" title="Pre-load all N copies of this model into VRAM at server startup instead of loading each copy on demand.">(load all ${max_instances} copies on startup)</span>
</label>
</div>
<!-- inference -->
......@@ -1354,6 +1369,7 @@ function openCfgModal(idx){
document.getElementById('cfg-modal-title').textContent = m.in_config ? 'Configure model' : 'Add to CoderAI';
document.getElementById('cfg-id-label').textContent = m.label;
document.getElementById('cfg-path').value = m.path;
document.getElementById('cfg-orig-path').value = m.path; // frozen at open time
document.getElementById('cfg-orig-type').value = m.defaultType;
// Quantization selector for grouped GGUF models
......@@ -1416,11 +1432,15 @@ function openCfgModal(idx){
// Used VRAM
const usedVram = s.used_vram_gb != null ? s.used_vram_gb : null;
document.getElementById('cfg-used-vram').value = usedVram != null ? usedVram : '';
// Show estimate hint from file size (GGUF: ~1.1x file size; HF: from size_gb)
const estVram = _estimateVram(m);
// Show estimate hint (weights × 1.15 + KV cache ~0.5 MB per 1 K ctx)
const nCtxForEst = s.n_ctx || 2048;
const estVram = _estimateVram(m, nCtxForEst);
document.getElementById('cfg-used-vram-hint').textContent = estVram ? `Estimated: ~${estVram.toFixed(1)} GB` : '';
document.getElementById('cfg-gpu-layers').value = s.n_gpu_layers !== undefined ? s.n_gpu_layers : -1;
document.getElementById('cfg-n-ctx').value = s.n_ctx || 2048;
document.getElementById('cfg-n-ctx').value = nCtxForEst;
document.getElementById('cfg-max-instances').value = s.max_instances != null ? s.max_instances : 1;
document.getElementById('cfg-preload-all-instances').checked = !!s.preload_all_instances;
_updatePreloadAllVisibility();
document.getElementById('cfg-max-gpu').value = s.max_gpu_percent != null ? s.max_gpu_percent : '';
document.getElementById('cfg-ram-gb').value = s.manual_ram_gb != null ? s.manual_ram_gb : '';
document.getElementById('cfg-4bit').checked = !!s.load_in_4bit;
......@@ -1439,10 +1459,19 @@ function openCfgModal(idx){
openModal('cfg-modal');
}
function _estimateVram(m) {
// Estimate VRAM from file size: GGUF ~1.1x, HF safetensors ~1.2x
if (m.size_gb) return m.size_gb * (m.cacheType === 'gguf' ? 1.1 : 1.2);
return null;
function _updatePreloadAllVisibility() {
const loadMode = document.getElementById('cfg-load-mode').value;
const maxInst = parseInt(document.getElementById('cfg-max-instances').value) || 1;
const row = document.getElementById('cfg-preload-all-row');
row.style.display = (loadMode === 'load' && maxInst > 1) ? '' : 'none';
}
function _estimateVram(m, nCtx) {
// Mirror the server-side formula: weights × 1.15 + KV cache (0.5 MB per 1 K ctx)
if (!m.size_gb) return null;
const weights = m.size_gb * (m.cacheType === 'gguf' ? 1.0 : 1.1);
const kvCacheGb = ((nCtx || 2048) / 1024) * 0.5 / 1024;
return weights * 1.15 + kvCacheGb;
}
async function saveModelConfig(){
......@@ -1451,8 +1480,10 @@ async function saveModelConfig(){
const ramGb = parseFloat(document.getElementById('cfg-ram-gb').value);
const usedVram = parseFloat(document.getElementById('cfg-used-vram').value);
const {primaryType, model_types, video_subtypes} = _getCheckedTypes();
const origPath = document.getElementById('cfg-orig-path').value;
const data = {
path,
orig_path: (origPath && origPath !== path) ? origPath : undefined,
model_type: primaryType,
model_types: model_types,
video_subtypes: video_subtypes.length ? video_subtypes : undefined,
......@@ -1460,6 +1491,8 @@ async function saveModelConfig(){
backend: document.getElementById('cfg-backend').value,
load_mode: document.getElementById('cfg-load-mode').value,
used_vram_gb: isNaN(usedVram) ? null : usedVram,
max_instances: parseInt(document.getElementById('cfg-max-instances').value) || 1,
preload_all_instances: document.getElementById('cfg-preload-all-instances').checked,
n_gpu_layers: parseInt(document.getElementById('cfg-gpu-layers').value) || -1,
n_ctx: parseInt(document.getElementById('cfg-n-ctx').value) || 2048,
max_gpu_percent: isNaN(maxGpu) ? null : maxGpu,
......@@ -1557,14 +1590,20 @@ async function unloadModel(idx){
async function disableModel(idx){
const m = _localModels[idx];
if(!confirm('Remove this model from CoderAI config? It will stay in the local cache.')) return;
// Collect all paths that are marked in_config for this model (covers quant groups)
const paths = m.ggufGroup
? m.ggufGroup.filter(f=>f.in_config).map(f=>f.path)
: (m.in_config ? [m.path] : []);
if(!paths.length) paths.push(m.path); // fallback
try{
const r = await fetch('/admin/api/model-disable',{
// Disable each configured path in the group
for(const p of paths){
await fetch('/admin/api/model-disable',{
method:'POST', headers:{'Content-Type':'application/json'},
body: JSON.stringify({path: m.path})
body: JSON.stringify({path: p})
});
const d = await r.json();
if(d.success) loadCachedModels();
else alert('Error: '+(d.detail||'Unknown'));
}
loadCachedModels();
}catch(e){ alert('Error: '+e.message); }
}
......
codai/api/text.py
View file @ f524118d
......@@ -293,13 +293,22 @@ async def chat_completions(request: ChatCompletionRequest, http_request: Request
if model_info.get('error'):
raise HTTPException(status_code=404, detail=model_info['error'])
# Try to get the appropriate model (request_model handles VRAM cleanup)
# Acquire the least-busy instance (increments ref-count; released on response completion)
_model_key = model_info.get('model_key')
_instance_idx = None
_acq = multi_model_manager.acquire_model_instance(_model_key) if _model_key else None
if _acq:
_instance_idx, mm = _acq
else:
mm = multi_model_manager.get_model_for_request(requested_model)
def _release_instance():
if _instance_idx is not None and _model_key:
multi_model_manager.release_model_instance(_model_key, _instance_idx)
if mm is None:
# Model not loaded - try to use default
_release_instance()
if model_manager.backend is not None:
# Fallback to legacy model_manager
current_manager = model_manager
else:
raise HTTPException(status_code=503, detail="Model not loaded")
......@@ -817,8 +826,15 @@ async def chat_completions(request: ChatCompletionRequest, http_request: Request
yield f"data: {json.dumps({'choices': [{'delta': {'content': second_pass_result}, 'finish_reason': 'stop'}]})}\n\n"
yield "data: [DONE]\n\n"
async def _raw_stream_with_release():
try:
async for chunk in raw_stream_generate():
yield chunk
finally:
_release_instance()
from fastapi.responses import StreamingResponse
return StreamingResponse(raw_stream_generate(), media_type="text/event-stream")
return StreamingResponse(_raw_stream_with_release(), media_type="text/event-stream")
# Non-streaming path (already implemented above)
# First pass: generate until reasoning close tag
......@@ -1127,9 +1143,9 @@ async def chat_completions(request: ChatCompletionRequest, http_request: Request
return JSONResponse(content=formatted_response, headers=headers)
if request.stream:
from fastapi.responses import StreamingResponse
return StreamingResponse(
stream_chat_response(
async def _managed_stream():
try:
async for chunk in stream_chat_response(
messages_dict,
response_model_name,
request.max_tokens,
......@@ -1140,10 +1156,15 @@ async def chat_completions(request: ChatCompletionRequest, http_request: Request
current_manager,
tool_parser,
request.response_format,
),
media_type="text/event-stream",
)
):
yield chunk
finally:
_release_instance()
from fastapi.responses import StreamingResponse
return StreamingResponse(_managed_stream(), media_type="text/event-stream")
else:
try:
return await generate_chat_response(
messages_dict,
response_model_name,
......@@ -1157,6 +1178,8 @@ async def chat_completions(request: ChatCompletionRequest, http_request: Request
request.response_format,
force_reasoning_args,
)
finally:
_release_instance()
async def stream_chat_response(
messages: List[Dict],
......@@ -1701,13 +1724,22 @@ async def completions(request: CompletionRequest):
if model_info.get('error'):
raise HTTPException(status_code=404, detail=model_info['error'])
# Try to get the appropriate model (request_model handles VRAM cleanup)
# Acquire the least-busy instance (increments ref-count; released on response completion)
_model_key = model_info.get('model_key')
_instance_idx = None
_acq = multi_model_manager.acquire_model_instance(_model_key) if _model_key else None
if _acq:
_instance_idx, mm = _acq
else:
mm = multi_model_manager.get_model_for_request(requested_model)
def _release_instance():
if _instance_idx is not None and _model_key:
multi_model_manager.release_model_instance(_model_key, _instance_idx)
if mm is None:
# Model not loaded - try to use default
_release_instance()
if model_manager.backend is not None:
# Fallback to legacy model_manager
current_manager = model_manager
else:
raise HTTPException(status_code=503, detail="Model not loaded")
......@@ -1720,9 +1752,9 @@ async def completions(request: CompletionRequest):
stop_sequences = [request.stop] if isinstance(request.stop, str) else request.stop
if request.stream:
from fastapi.responses import StreamingResponse
return StreamingResponse(
stream_completion_response(
async def _managed_completion_stream():
try:
async for chunk in stream_completion_response(
prompts[0],
request.model,
request.max_tokens,
......@@ -1730,10 +1762,15 @@ async def completions(request: CompletionRequest):
request.top_p,
stop_sequences,
current_manager,
),
media_type="text/event-stream",
)
):
yield chunk
finally:
_release_instance()
from fastapi.responses import StreamingResponse
return StreamingResponse(_managed_completion_stream(), media_type="text/event-stream")
else:
try:
return await generate_completion_response(
prompts[0],
request.model,
......@@ -1743,6 +1780,8 @@ async def completions(request: CompletionRequest):
stop_sequences,
current_manager,
)
finally:
_release_instance()
async def stream_completion_response(
prompt: str,
......
codai/config.py
View file @ f524118d
......@@ -25,8 +25,8 @@ from dataclasses import dataclass, field
@dataclass
class ServerConfig:
"""Server configuration."""
host: str = "0.0.0.0"
port: int = 8000
host: str = "127.0.0.1"
port: int = 8776
https: bool = False
https_key_path: Optional[str] = None
https_cert_path: Optional[str] = None
......@@ -49,6 +49,7 @@ class ModelsConfig:
default_load_mode: str = "ondemand"
hf_cache_dir: Optional[str] = None
gguf_cache_dir: Optional[str] = None
max_model_instances: int = 1 # max concurrent instances per model (global default; overridable per-model via "max_instances")
@dataclass
......@@ -150,8 +151,8 @@ class ConfigManager:
default_config = {
"version": "1.0",
"server": {
"host": "0.0.0.0",
"port": 8000,
"host": "127.0.0.1",
"port": 8776,
"https": False,
"https_key_path": None,
"https_cert_path": None
......
codai/main.py
View file @ f524118d
......@@ -234,6 +234,7 @@ def main():
load_mode = config.models.default_load_mode or "ondemand"
set_load_mode(load_mode)
multi_model_manager.set_load_mode(load_mode)
multi_model_manager._global_max_instances = config.models.max_model_instances
print(f"\nLoad mode: {load_mode}")
if load_mode == "ondemand":
......@@ -463,6 +464,18 @@ def main():
if mm is not None and mm.backend is not None:
multi_model_manager.active_in_vram = mid
print(f" Loaded: {mid}")
# Load additional instances if preload_all_instances is set
if isinstance(m, dict) and m.get("preload_all_instances"):
max_inst = m.get("max_instances", 1)
if isinstance(max_inst, int) and max_inst > 1:
for i in range(max_inst - 1):
multi_model_manager._pending_new_instance.add(mid)
mm2 = multi_model_manager._load_model_by_name(mid)
if mm2 is not None and mm2.backend is not None:
print(f" Loaded extra instance {i+2}/{max_inst}: {mid}")
else:
print(f" Warning: extra instance {i+2}/{max_inst} of '{mid}' failed to load")
break
else:
print(f" Warning: {mid} failed to load")
elif mtype == "audio" and mid in multi_model_manager.whisper_servers:
......
codai/models/manager.py
View file @ f524118d
......@@ -395,6 +395,86 @@ class WhisperServerManager:
self.stop()
class ModelInstancePool:
"""
Ref-counted pool of loaded instances for a single model key.
When max_instances > 1 and all existing instances are handling at least one
active request, the manager can load additional instances (up to the limit)
to serve concurrent requests in parallel. Each caller must call release()
when it finishes using the instance so the ref-count stays accurate.
"""
def __init__(self, max_instances: int = 1):
self.instances: list = []
self.ref_counts: list = []
self.max_instances: int = max_instances
self._lock = threading.Lock()
@property
def count(self) -> int:
return len(self.instances)
@property
def primary(self):
return self.instances[0] if self.instances else None
def add(self, model_obj) -> int:
"""Register a new instance; return its index."""
with self._lock:
idx = len(self.instances)
self.instances.append(model_obj)
self.ref_counts.append(0)
return idx
def acquire(self):
"""Return (idx, instance) of the least-busy instance, incrementing its ref-count."""
with self._lock:
if not self.instances:
return None
idx = min(range(len(self.instances)), key=lambda i: self.ref_counts[i])
self.ref_counts[idx] += 1
return idx, self.instances[idx]
def release(self, idx: int) -> None:
with self._lock:
if 0 <= idx < len(self.ref_counts):
self.ref_counts[idx] = max(0, self.ref_counts[idx] - 1)
def least_busy_instance(self):
"""Return the instance with the lowest ref-count without incrementing."""
with self._lock:
if not self.instances:
return None
idx = min(range(len(self.instances)), key=lambda i: self.ref_counts[i])
return self.instances[idx]
def all_busy(self) -> bool:
"""True if every instance is handling at least one request."""
with self._lock:
return bool(self.ref_counts) and all(r > 0 for r in self.ref_counts)
def can_add(self) -> bool:
return len(self.instances) < self.max_instances
def needs_new_instance(self) -> bool:
"""True when all instances are busy and we are below max_instances."""
return self.all_busy() and self.can_add()
def cleanup_all(self) -> None:
with self._lock:
for obj in self.instances:
try:
if hasattr(obj, 'cleanup'):
obj.cleanup()
elif hasattr(obj, 'to'):
obj.to('cpu')
except Exception:
pass
self.instances.clear()
self.ref_counts.clear()
class MultiModelManager:
"""
Manages multiple models: main text model, audio transcription, and image generation.
......@@ -422,6 +502,9 @@ class MultiModelManager:
self.model_backend_types: Dict[str, str] = {}
self.tool_breaker = FuzzyToolBreaker(threshold=3) # Circuit breaker for repetitive tool calls
self._load_lock = threading.Lock() # Prevents duplicate on-demand model loads
self.model_pools: Dict[str, ModelInstancePool] = {} # per-key instance pools
self._pending_new_instance: set = set() # keys awaiting a second+ instance load
self._global_max_instances: int = 1 # set from config at startup
@property
def image_model(self) -> Optional[str]:
......@@ -546,13 +629,14 @@ class MultiModelManager:
return None
# Fast path: already loaded (checked without lock for performance)
if self.default_model in self.models:
return self.models[self.default_model]
if self.default_model in self.models and self.default_model not in self._pending_new_instance:
return self._get_least_busy_instance(self.default_model)
with self._load_lock:
# Re-check inside the lock to avoid duplicate loads from concurrent requests
if self.default_model in self.models:
return self.models[self.default_model]
if self.default_model in self.models and self.default_model not in self._pending_new_instance:
return self._get_least_busy_instance(self.default_model)
self._pending_new_instance.discard(self.default_model)
config = self.config.get(self.default_model, {})
backend_type = self.model_backend_types.get(self.default_model, "auto")
......@@ -590,7 +674,7 @@ class MultiModelManager:
print(f"Loading default model on demand: {self.default_model}")
model_manager.load_model(self.default_model, backend_type=backend_type, **kwargs)
self.models[self.default_model] = model_manager
self.add_model(self.default_model, model_manager)
self.current_model_key = self.default_model
print(f"Model loaded successfully: {self.default_model}")
return model_manager
......@@ -600,13 +684,14 @@ class MultiModelManager:
def _load_model_by_name(self, model_name: str):
"""Load a model by name on demand (thread-safe)."""
if model_name in self.models:
return self.models[model_name]
if model_name in self.models and model_name not in self._pending_new_instance:
return self._get_least_busy_instance(model_name)
with self._load_lock:
# Re-check inside lock to prevent duplicate loads
if model_name in self.models:
return self.models[model_name]
# Re-check inside lock to prevent duplicate loads, but honour pending new instance.
if model_name in self.models and model_name not in self._pending_new_instance:
return self._get_least_busy_instance(model_name)
self._pending_new_instance.discard(model_name)
config = self.config.get(model_name, {})
backend_type = self.model_backend_types.get(model_name, "auto")
......@@ -642,11 +727,15 @@ class MultiModelManager:
if hasattr(global_args, 'max_gpu_percent'):
kwargs['max_gpu_percent'] = global_args.max_gpu_percent
print(f"Loading model on demand: {model_name}")
pool = self.model_pools.get(model_name)
inst_num = pool.count + 1 if pool else 1
print(f"Loading model on demand: {model_name}"
+ (f" (instance {inst_num})" if inst_num > 1 else ""))
model_manager.load_model(model_name, backend_type=backend_type, **kwargs)
self.models[model_name] = model_manager
self.add_model(model_name, model_manager)
self.current_model_key = model_name
print(f"Model loaded successfully: {model_name}")
print(f"Model loaded successfully: {model_name}"
+ (f" (instance {inst_num})" if inst_num > 1 else ""))
return model_manager
except Exception as e:
print(f"Error loading model {model_name}: {e}")
......@@ -942,32 +1031,25 @@ class MultiModelManager:
# Handle empty or "default" model names
if not requested_model or requested_model == "default":
if self.default_model:
# Check if already loaded
if self.default_model in self.models:
self.current_model_key = self.default_model
return self.models[self.default_model]
# Model not loaded yet - try to load it
return self._get_least_busy_instance(self.default_model)
return self._load_default_model()
return None
# Handle "audio" alias
if requested_model == "audio":
if self.audio_models:
first_audio = self.audio_models[0]
key = f"audio:{first_audio}"
key = f"audio:{self.audio_models[0]}"
if key in self.models:
self.current_model_key = key
return self.models[key]
return self._get_least_busy_instance(key)
return None
# Handle "image" alias
if requested_model == "image":
if self.image_models:
first_image = self.image_models[0]
key = f"image:{first_image}"
key = f"image:{self.image_models[0]}"
if key in self.models:
self.current_model_key = key
return self.models[key]
return self._get_least_busy_instance(key)
return None
# Handle "tts" alias
......@@ -975,53 +1057,40 @@ class MultiModelManager:
if self.tts_model:
key = f"tts:{self.tts_model}"
if key in self.models:
self.current_model_key = key
return self.models[key]
return self._get_least_busy_instance(key)
return None
# Handle prefixed models
if requested_model.startswith("audio:"):
audio_name = requested_model[6:]
key = f"audio:{audio_name}"
key = f"audio:{requested_model[6:]}"
if key in self.models:
self.current_model_key = key
return self.models[key]
return self._get_least_busy_instance(key)
return None
if requested_model.startswith("tts:"):
tts_name = requested_model[4:]
key = f"tts:{tts_name}"
key = f"tts:{requested_model[4:]}"
if key in self.models:
self.current_model_key = key
return self.models[key]
return self._get_least_busy_instance(key)
return None
if requested_model.startswith("vision:") or requested_model.startswith("image:"):
if requested_model.startswith("vision:"):
image_name = requested_model[7:]
else:
image_name = requested_model[6:]
image_name = requested_model[7:] if requested_model.startswith("vision:") else requested_model[6:]
key = f"image:{image_name}"
if key in self.models:
self.current_model_key = key
return self.models[key]
return self._get_least_busy_instance(key)
return None
# Check if it's the default model
if self.default_model and (requested_model == self.default_model or
requested_model.endswith(self.default_model.split("/")[-1])):
# Check if already loaded
if self.default_model in self.models:
self.current_model_key = self.default_model
return self.models[self.default_model]
# Try to load the default model
return self._get_least_busy_instance(self.default_model)
return self._load_default_model()
# Check if any loaded model matches
for key, model in self.models.items():
for key in self.models:
if requested_model in key or key.endswith(requested_model.split("/")[-1]):
self.current_model_key = key
return model
return self._get_least_busy_instance(key)
# Validate the model is allowed before attempting to load it.
# This prevents loading arbitrary models not registered via command line.
......@@ -1292,12 +1361,71 @@ class MultiModelManager:
return free / 1e9
except Exception:
pass
# AMD GPU via sysfs (covers Vulkan/ROCm on Linux)
try:
import glob
for total_path in sorted(glob.glob('/sys/class/drm/card*/device/mem_info_vram_total')):
used_path = total_path.replace('vram_total', 'vram_used')
with open(total_path) as f:
total = int(f.read().strip())
with open(used_path) as f:
used = int(f.read().strip())
return (total - used) / 1e9
except Exception:
pass
return 999.0 # Unknown — assume enough
def _get_model_used_vram_gb(self, model_key: str) -> float:
"""Return the configured used_vram_gb for a model, or 0 if unknown."""
def _get_model_used_vram_gb(self, model_key: str, resolved_name: str = None) -> float:
"""Return VRAM requirement in GB for a model.
Uses configured used_vram_gb when set; otherwise estimates from file size
for local model files (GGUF, GGML, whisper .bin, etc.) — weights + KV cache
+ 15% buffer. Returns 0 when the requirement cannot be determined.
"""
cfg = self.config.get(model_key, {})
return float(cfg.get("used_vram_gb") or 0)
explicit = cfg.get("used_vram_gb")
if explicit:
return float(explicit)
# Build a list of candidate paths to check
candidates = []
if resolved_name:
candidates.append(resolved_name)
# model_key is often the path for text GGUF models
candidates.append(model_key)
# Strip type prefix (e.g. "audio:/path/to/model.bin")
if ":" in model_key:
candidates.append(model_key.split(":", 1)[1])
# Config may store the path explicitly
for field in ("path", "model_path", "model"):
v = cfg.get(field)
if v:
candidates.append(v)
import os
local_exts = {'.gguf', '.ggml', '.bin', '.pt', '.safetensors'}
for path in candidates:
if not path:
continue
_, ext = os.path.splitext(path)
if ext.lower() not in local_exts:
continue
try:
size_bytes = os.path.getsize(path)
weights_gb = size_bytes / 1e9
# KV cache: 2 bytes × 2 (K+V) × n_layers × n_ctx × head_dim
# We don't know the architecture, so estimate from n_ctx alone.
# Empirically ~0.5 MB per 1 K tokens covers most small models.
n_ctx = cfg.get("n_ctx") or 2048
kv_cache_gb = (n_ctx / 1024) * 0.5 / 1024 # 0.5 MB per 1 K ctx → GB
# 15% overhead for activations, graph buffers, scratch space
return weights_gb * 1.15 + kv_cache_gb
except OSError:
continue
return 0
def _evict_models_for_vram(self, needed_gb: float):
"""Unload loaded models (LRU first) until we have at least needed_gb free VRAM."""
......@@ -1509,7 +1637,7 @@ class MultiModelManager:
'already_loaded': True,
}
# Not loaded — check VRAM and evict if needed
needed_gb = self._get_model_used_vram_gb(model_key)
needed_gb = self._get_model_used_vram_gb(model_key, resolved_name)
if needed_gb > 0:
free_gb = self._get_free_vram_gb()
if free_gb < needed_gb:
......@@ -1634,11 +1762,33 @@ class MultiModelManager:
}
# =====================================================================
# ONDEMAND MODE (default): Only one model in memory at a time.
# Fully unload the current model before loading the new one.
# ONDEMAND MODE (default): One model in memory at a time — but if a model
# has max_instances > 1 and all instances are busy, load another copy.
# =====================================================================
if existing_model is not None:
# Already loaded and it's the only model - return it
pool = self.model_pools.get(model_key)
if pool is None:
# Pool wasn't created yet (model registered outside add_model); adopt it now.
pool = ModelInstancePool(max_instances=self._get_max_instances(model_key))
pool.add(existing_model)
self.model_pools[model_key] = pool
if pool.needs_new_instance():
needed_gb = self._get_model_used_vram_gb(model_key, resolved_name)
free_gb = self._get_free_vram_gb()
if needed_gb > 0 and free_gb >= needed_gb:
print(f"Ondemand: all {pool.count} instance(s) of '{model_key}' busy — "
f"loading instance {pool.count + 1}/{pool.max_instances} "
f"(need {needed_gb:.1f} GB, have {free_gb:.1f} GB free)")
self._pending_new_instance.add(model_key)
return {
'model_key': model_key,
'model_name': resolved_name,
'model_object': None,
'config': self.config.get(model_key, {}),
'already_loaded': False,
}
self.current_model_key = model_key
self.active_in_vram = model_key
return {
......@@ -1663,6 +1813,13 @@ class MultiModelManager:
loaded_canonical = "legacy_model_manager"
if loaded_canonical and loaded_canonical != model_key:
needed_gb = self._get_model_used_vram_gb(model_key, resolved_name)
free_gb = self._get_free_vram_gb()
if needed_gb > 0 and free_gb >= needed_gb:
print(f"Ondemand mode - keeping '{loaded_canonical}' in VRAM alongside new model "
f"(need {needed_gb:.1f} GB, have {free_gb:.1f} GB free)")
else:
print(f"Ondemand mode - model switch detected:")
print(f" Requested: '{model_key}' (resolved: '{resolved_name}')")
print(f" Currently loaded: '{loaded_canonical}'")
......@@ -1760,6 +1917,8 @@ class MultiModelManager:
self.current_model_key = None
self.active_in_vram = None
self.models_in_vram = set()
self.model_pools.clear()
self._pending_new_instance.clear()
# Force garbage collection
for _ in range(3):
......@@ -1779,12 +1938,57 @@ class MultiModelManager:
time.sleep(1)
print("=== FULL VRAM CLEANUP: Complete ===")
def _get_max_instances(self, model_key: str) -> int:
"""Per-model max_instances, falling back to the global default."""
per_model = self.config.get(model_key, {}).get("max_instances")
if per_model is not None:
return int(per_model)
return self._global_max_instances
def _get_least_busy_instance(self, key: str):
"""Return the least-busy instance for key without incrementing its ref-count."""
pool = self.model_pools.get(key)
if pool and pool.count > 0:
obj = pool.least_busy_instance()
if obj is not None:
self.current_model_key = key
return obj
obj = self.models.get(key)
if obj is not None:
self.current_model_key = key
return obj
def add_model(self, key: str, manager):
"""Add a model (ModelManager, diffusers pipeline, sd.cpp model, etc.) for a specific key."""
self.models[key] = manager
pool = self.model_pools.get(key)
if pool is None:
pool = ModelInstancePool(max_instances=self._get_max_instances(key))
self.model_pools[key] = pool
pool.add(manager)
self.models[key] = pool.primary # backward compat: always the first instance
self.active_in_vram = key
self.models_in_vram.add(key)
def acquire_model_instance(self, model_key: str):
"""Acquire the least-busy instance, incrementing its ref-count.
Returns (instance_idx, model_obj) or None if no instance is loaded.
Callers MUST call release_model_instance() when done.
"""
pool = self.model_pools.get(model_key)
if pool and pool.count > 0:
return pool.acquire()
obj = self.models.get(model_key)
if obj is not None:
return 0, obj
return None
def release_model_instance(self, model_key: str, instance_idx: int) -> None:
"""Release a previously acquired instance, decrementing its ref-count."""
pool = self.model_pools.get(model_key)
if pool:
pool.release(instance_idx)
def get_model(self, key: str) -> Optional[ModelManager]:
"""Get a model manager by key."""
return self.models.get(key)
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment