주식자동매매 AI 프로그램 초기 모델

modules/services/ollama.py

@@ -0,0 +1,114 @@
+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