Clement Tian on Paperspace
133 lines
5.1 KiB
Python
133 lines
5.1 KiB
Python
"""GPU model manager — ensures only one model subprocess runs at a time."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
import os
|
|
import subprocess
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
import httpx
|
|
|
|
ROOT_DIR = Path(__file__).parent.parent
|
|
|
|
ACESTEP_PORT = 8001
|
|
FOUNDATION1_PORT = 8002
|
|
HEALTH_POLL_INTERVAL = 2.0
|
|
HEALTH_TIMEOUT_SECONDS = 180 # model loading can be slow on first run
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def _python_exe(venv_dir: Path) -> Path:
|
|
"""Return the Python executable for a given venv directory."""
|
|
win = venv_dir / "Scripts" / "python.exe"
|
|
if win.exists():
|
|
return win
|
|
return venv_dir / "bin" / "python"
|
|
|
|
|
|
class ModelNotActiveError(Exception):
|
|
def __init__(self, requested: str, active: Optional[str]) -> None:
|
|
self.requested = requested
|
|
self.active = active
|
|
super().__init__(f"Model '{requested}' is not loaded (active: {active!r})")
|
|
|
|
|
|
class ModelManager:
|
|
def __init__(self) -> None:
|
|
self.active_model: Optional[str] = None
|
|
self._lock = asyncio.Lock()
|
|
self._process: Optional[subprocess.Popen] = None
|
|
self.acestep_port: int = ACESTEP_PORT
|
|
self.foundation1_port: int = FOUNDATION1_PORT
|
|
|
|
async def load(self, model: str) -> None:
|
|
"""Load a model, unloading the currently active one first if needed."""
|
|
async with self._lock:
|
|
if self.active_model == model:
|
|
logger.info("Model '%s' is already active.", model)
|
|
return
|
|
if self.active_model is not None:
|
|
logger.info("Unloading model '%s'...", self.active_model)
|
|
await asyncio.get_event_loop().run_in_executor(None, self._stop_process)
|
|
logger.info("Starting model '%s'...", model)
|
|
await asyncio.get_event_loop().run_in_executor(None, self._start_process, model)
|
|
self.active_model = model
|
|
logger.info("Model '%s' is ready.", model)
|
|
|
|
async def unload(self) -> None:
|
|
"""Unload the active model."""
|
|
async with self._lock:
|
|
if self.active_model is None:
|
|
return
|
|
logger.info("Unloading model '%s'...", self.active_model)
|
|
await asyncio.get_event_loop().run_in_executor(None, self._stop_process)
|
|
|
|
def _stop_process(self) -> None:
|
|
if self._process is None:
|
|
return
|
|
try:
|
|
self._process.terminate()
|
|
self._process.wait(timeout=30)
|
|
except subprocess.TimeoutExpired:
|
|
logger.warning("Process did not terminate gracefully, killing.")
|
|
self._process.kill()
|
|
self._process.wait()
|
|
except Exception as exc:
|
|
logger.error("Error stopping process: %s", exc)
|
|
finally:
|
|
self._process = None
|
|
self.active_model = None
|
|
|
|
def _start_process(self, model: str) -> None:
|
|
if model == "fullsong":
|
|
self._process = self._launch_acestep()
|
|
self._wait_healthy(f"http://127.0.0.1:{self.acestep_port}/health")
|
|
elif model == "clip":
|
|
self._process = self._launch_foundation1()
|
|
self._wait_healthy(f"http://127.0.0.1:{self.foundation1_port}/health")
|
|
elif model == "separator":
|
|
pass # no persistent process; previous model already stopped by load(), VRAM is free
|
|
else:
|
|
raise ValueError(f"Unknown model: '{model}'. Must be 'fullsong', 'clip', or 'separator'.")
|
|
|
|
def _launch_acestep(self) -> subprocess.Popen:
|
|
venv_python = _python_exe(ROOT_DIR / "ace-step" / ".venv")
|
|
cmd = [str(venv_python), "-c", "from acestep.api_server import main; main()"]
|
|
env = {**os.environ, "ACESTEP_API_PORT": str(self.acestep_port)}
|
|
env.setdefault("ACESTEP_LM_MODEL_PATH", "acestep-5Hz-lm-0.6B")
|
|
logger.info("Launching ACE-Step on port %d: %s", self.acestep_port, " ".join(cmd))
|
|
return subprocess.Popen(cmd, cwd=str(ROOT_DIR / "ace-step"), env=env)
|
|
|
|
def _launch_foundation1(self) -> subprocess.Popen:
|
|
venv_python = _python_exe(ROOT_DIR / "foundation1" / ".venv")
|
|
server_script = str(ROOT_DIR / "foundation1_server" / "server.py")
|
|
cmd = [str(venv_python), server_script]
|
|
env = {**os.environ, "FOUNDATION1_SERVER_PORT": str(self.foundation1_port)}
|
|
logger.info("Launching Foundation-1 server on port %d: %s", self.foundation1_port, " ".join(cmd))
|
|
return subprocess.Popen(cmd, cwd=str(ROOT_DIR), env=env)
|
|
|
|
def _wait_healthy(self, url: str) -> None:
|
|
deadline = time.monotonic() + HEALTH_TIMEOUT_SECONDS
|
|
logger.info("Waiting for service at %s ...", url)
|
|
while time.monotonic() < deadline:
|
|
try:
|
|
with httpx.Client(timeout=5.0) as client:
|
|
resp = client.get(url)
|
|
if resp.status_code == 200:
|
|
logger.info("Service healthy: %s", url)
|
|
return
|
|
except Exception:
|
|
pass
|
|
time.sleep(HEALTH_POLL_INTERVAL)
|
|
raise TimeoutError(
|
|
f"Service at {url} did not become healthy within {HEALTH_TIMEOUT_SECONDS}s"
|
|
)
|
|
|
|
|
|
model_manager = ModelManager()
|