115 lines
4.4 KiB
Python
115 lines
4.4 KiB
Python
import requests
|
|
import json
|
|
import psutil
|
|
try:
|
|
import pynvml
|
|
except ImportError:
|
|
pynvml = None
|
|
|
|
from modules.config import Config
|
|
|
|
class OllamaManager:
|
|
"""
|
|
Ollama API 세션 관리 및 메모리 누수 방지 래퍼
|
|
- GPU VRAM 사용량 모니터링
|
|
- keep_alive 파라미터를 통한 메모리 관리
|
|
"""
|
|
def __init__(self, model_name=None, base_url=None):
|
|
self.model_name = model_name or Config.OLLAMA_MODEL
|
|
self.base_url = base_url or Config.OLLAMA_API_URL
|
|
self.generate_url = f"{self.base_url}/api/generate"
|
|
|
|
self.gpu_available = False
|
|
try:
|
|
if pynvml:
|
|
pynvml.nvmlInit()
|
|
self.handle = pynvml.nvmlDeviceGetHandleByIndex(0) # 0번 GPU (3070 Ti)
|
|
self.gpu_available = True
|
|
print("✅ [OllamaManager] NVIDIA GPU Monitoring On")
|
|
else:
|
|
print("⚠️ [OllamaManager] 'nvidia-ml-py' not installed. GPU monitoring disabled.")
|
|
except Exception as e:
|
|
print(f"⚠️ [OllamaManager] GPU Init Failed: {e}")
|
|
|
|
def check_vram(self):
|
|
"""현재 GPU VRAM 사용량(GB) 반환"""
|
|
if not self.gpu_available:
|
|
return 0.0
|
|
try:
|
|
info = pynvml.nvmlDeviceGetMemoryInfo(self.handle)
|
|
used_gb = info.used / 1024**3
|
|
return used_gb
|
|
except Exception:
|
|
return 0.0
|
|
|
|
def get_gpu_status(self):
|
|
"""GPU 종합 상태 반환 (온도, 메모리, 사용률, 이름)"""
|
|
if not self.gpu_available:
|
|
return {"name": "N/A", "temp": 0, "vram_used": 0, "vram_total": 0, "load": 0}
|
|
|
|
try:
|
|
# GPU 이름
|
|
name = pynvml.nvmlDeviceGetName(self.handle)
|
|
if isinstance(name, bytes):
|
|
name = name.decode('utf-8')
|
|
|
|
# 온도
|
|
temp = pynvml.nvmlDeviceGetTemperature(self.handle, pynvml.NVML_TEMPERATURE_GPU)
|
|
# 메모리
|
|
mem_info = pynvml.nvmlDeviceGetMemoryInfo(self.handle)
|
|
vram_used = mem_info.used / 1024**3
|
|
vram_total = mem_info.total / 1024**3
|
|
# 사용률
|
|
util = pynvml.nvmlDeviceGetUtilizationRates(self.handle)
|
|
load = util.gpu
|
|
|
|
return {
|
|
"name": name,
|
|
"temp": temp,
|
|
"vram_used": round(vram_used, 1),
|
|
"vram_total": round(vram_total, 1),
|
|
"load": load
|
|
}
|
|
except Exception as e:
|
|
print(f"⚠️ GPU Status Check Failed: {e}")
|
|
return {"name": "N/A", "temp": 0, "vram_used": 0, "vram_total": 0, "load": 0}
|
|
|
|
def request_inference(self, prompt, context_data=None):
|
|
"""
|
|
Ollama에 추론 요청
|
|
:param prompt: 시스템 프롬프트 + 사용자 입력
|
|
:param context_data: (Optional) 이전 대화 컨텍스트
|
|
"""
|
|
# [5070Ti 최적화] VRAM이 14GB 이상이면 모델 언로드 시도 (16GB 중 여유분 확보)
|
|
vram = self.check_vram()
|
|
if vram > 14.0:
|
|
print(f"⚠️ [OllamaManager] High VRAM Usage ({vram:.1f}GB). Requesting unload.")
|
|
try:
|
|
# keep_alive=0으로 설정하여 모델 즉시 언로드
|
|
requests.post(self.generate_url,
|
|
json={"model": self.model_name, "keep_alive": 0}, timeout=5)
|
|
except Exception as e:
|
|
print(f"Warning: Failed to unload model: {e}")
|
|
|
|
payload = {
|
|
"model": self.model_name,
|
|
"prompt": prompt,
|
|
"stream": False,
|
|
"format": "json", # JSON 강제
|
|
"options": {
|
|
"num_ctx": 8192, # [5070Ti 최적화] 컨텍스트 크기 2배 증가 (4096 -> 8192)
|
|
"temperature": 0.2, # 분석 일관성 유지
|
|
"num_gpu": 1, # GPU 사용 명시
|
|
"num_thread": 8 # CPU 스레드 수 (9800X3D 활용)
|
|
},
|
|
"keep_alive": "10m" # [5070Ti 최적화] 10분간 유지 (메모리 여유 있음)
|
|
}
|
|
|
|
try:
|
|
response = requests.post(self.generate_url, json=payload, timeout=180) # 타임아웃 증가
|
|
response.raise_for_status()
|
|
return response.json().get('response')
|
|
except Exception as e:
|
|
print(f"❌ Inference Error: {e}")
|
|
return None
|