isa-model 0.0.1__py3-none-any.whl

isa_model/core/resource_manager.py

@@ -0,0 +1,202 @@
+import time
+import threading
+import logging
+from typing import Dict, List, Any, Optional
+import psutil
+try:
+    import torch
+    import nvidia_smi
+    HAS_GPU = torch.cuda.is_available()
+except ImportError:
+    HAS_GPU = False
+
+logger = logging.getLogger(__name__)
+
+class ResourceManager:
+    """
+    Monitors system resources and manages model loading/unloading
+    to prevent resource exhaustion.
+    """
+    
+    def __init__(self, model_registry, monitoring_interval=30):
+        self.registry = model_registry
+        self.monitoring_interval = monitoring_interval  # seconds
+        
+        self.max_memory_percent = 90  # Maximum memory usage percentage
+        self.max_gpu_memory_percent = 90  # Maximum GPU memory usage percentage
+        
+        self._stop_event = threading.Event()
+        self._monitor_thread = None
+        
+        # Track resource usage over time
+        self.resource_history = {
+            "timestamps": [],
+            "cpu_percent": [],
+            "memory_percent": [],
+            "gpu_utilization": [],
+            "gpu_memory_percent": []
+        }
+        
+        # Initialize GPU monitoring if available
+        if HAS_GPU:
+            try:
+                nvidia_smi.nvmlInit()
+                self.gpu_count = torch.cuda.device_count()
+                logger.info(f"Initialized GPU monitoring with {self.gpu_count} devices")
+            except Exception as e:
+                logger.warning(f"Failed to initialize NVIDIA SMI: {str(e)}")
+                self.gpu_count = 0
+        else:
+            self.gpu_count = 0
+            
+        logger.info("Initialized ResourceManager")
+    
+    def start_monitoring(self):
+        """Start the resource monitoring thread"""
+        if self._monitor_thread is not None and self._monitor_thread.is_alive():
+            logger.warning("Resource monitoring already running")
+            return
+            
+        self._stop_event.clear()
+        self._monitor_thread = threading.Thread(
+            target=self._monitor_resources,
+            daemon=True
+        )
+        self._monitor_thread.start()
+        logger.info("Started resource monitoring thread")
+    
+    def stop_monitoring(self):
+        """Stop the resource monitoring thread"""
+        if self._monitor_thread is not None:
+            self._stop_event.set()
+            self._monitor_thread.join(timeout=5)
+            self._monitor_thread = None
+            logger.info("Stopped resource monitoring thread")
+    
+    def _monitor_resources(self):
+        """Monitor system resources in a loop"""
+        while not self._stop_event.is_set():
+            try:
+                # Get current resource usage
+                cpu_percent = psutil.cpu_percent(interval=1)
+                memory = psutil.virtual_memory()
+                memory_percent = memory.percent
+                
+                # GPU monitoring
+                gpu_utilization = 0
+                gpu_memory_percent = 0
+                
+                if self.gpu_count > 0:
+                    gpu_util_sum = 0
+                    gpu_mem_percent_sum = 0
+                    
+                    for i in range(self.gpu_count):
+                        try:
+                            handle = nvidia_smi.nvmlDeviceGetHandleByIndex(i)
+                            util = nvidia_smi.nvmlDeviceGetUtilizationRates(handle)
+                            mem_info = nvidia_smi.nvmlDeviceGetMemoryInfo(handle)
+                            
+                            gpu_util_sum += util.gpu
+                            gpu_mem_percent = (mem_info.used / mem_info.total) * 100
+                            gpu_mem_percent_sum += gpu_mem_percent
+                        except Exception as e:
+                            logger.error(f"Error getting GPU {i} stats: {str(e)}")
+                    
+                    if self.gpu_count > 0:
+                        gpu_utilization = gpu_util_sum / self.gpu_count
+                        gpu_memory_percent = gpu_mem_percent_sum / self.gpu_count
+                
+                # Record history (keep last 60 samples)
+                now = time.time()
+                self.resource_history["timestamps"].append(now)
+                self.resource_history["cpu_percent"].append(cpu_percent)
+                self.resource_history["memory_percent"].append(memory_percent)
+                self.resource_history["gpu_utilization"].append(gpu_utilization)
+                self.resource_history["gpu_memory_percent"].append(gpu_memory_percent)
+                
+                # Trim history to last 60 samples
+                max_history = 60
+                if len(self.resource_history["timestamps"]) > max_history:
+                    for key in self.resource_history:
+                        self.resource_history[key] = self.resource_history[key][-max_history:]
+                
+                # Check if we need to unload models
+                self._check_resource_constraints(memory_percent, gpu_memory_percent)
+                
+                # Log current usage
+                logger.debug(
+                    f"Resource usage - CPU: {cpu_percent:.1f}%, Memory: {memory_percent:.1f}%, "
+                    f"GPU: {gpu_utilization:.1f}%, GPU Memory: {gpu_memory_percent:.1f}%"
+                )
+                
+                # Wait for next check
+                self._stop_event.wait(self.monitoring_interval)
+                
+            except Exception as e:
+                logger.error(f"Error in resource monitoring: {str(e)}")
+                # Wait a bit before retrying
+                self._stop_event.wait(5)
+    
+    def _check_resource_constraints(self, memory_percent, gpu_memory_percent):
+        """Check if we need to unload models due to resource constraints"""
+        # Check memory usage
+        if memory_percent > self.max_memory_percent:
+            logger.warning(
+                f"Memory usage ({memory_percent:.1f}%) exceeds threshold ({self.max_memory_percent}%). "
+                "Unloading least used model."
+            )
+            # This would trigger model unloading
+            # self.registry._evict_least_used_model()
+        
+        # Check GPU memory usage
+        if HAS_GPU and gpu_memory_percent > self.max_gpu_memory_percent:
+            logger.warning(
+                f"GPU memory usage ({gpu_memory_percent:.1f}%) exceeds threshold ({self.max_gpu_memory_percent}%). "
+                "Unloading least used model."
+            )
+            # This would trigger model unloading
+            # self.registry._evict_least_used_model()
+    
+    def get_resource_usage(self) -> Dict[str, Any]:
+        """Get current resource usage stats"""
+        try:
+            cpu_percent = psutil.cpu_percent(interval=0.1)
+            memory = psutil.virtual_memory()
+            memory_percent = memory.percent
+            
+            result = {
+                "cpu_percent": cpu_percent,
+                "memory_total_gb": memory.total / (1024**3),
+                "memory_available_gb": memory.available / (1024**3),
+                "memory_percent": memory_percent,
+                "gpus": []
+            }
+            
+            # GPU stats
+            if HAS_GPU:
+                for i in range(self.gpu_count):
+                    try:
+                        handle = nvidia_smi.nvmlDeviceGetHandleByIndex(i)
+                        name = nvidia_smi.nvmlDeviceGetName(handle)
+                        util = nvidia_smi.nvmlDeviceGetUtilizationRates(handle)
+                        mem_info = nvidia_smi.nvmlDeviceGetMemoryInfo(handle)
+                        temp = nvidia_smi.nvmlDeviceGetTemperature(
+                            handle, nvidia_smi.NVML_TEMPERATURE_GPU
+                        )
+                        
+                        result["gpus"].append({
+                            "index": i,
+                            "name": name,
+                            "utilization_percent": util.gpu,
+                            "memory_total_gb": mem_info.total / (1024**3),
+                            "memory_used_gb": mem_info.used / (1024**3),
+                            "memory_percent": (mem_info.used / mem_info.total) * 100,
+                            "temperature_c": temp
+                        })
+                    except Exception as e:
+                        logger.error(f"Error getting GPU {i} stats: {str(e)}")
+            
+            return result
+        except Exception as e:
+            logger.error(f"Error getting resource usage: {str(e)}")
+            return {"error": str(e)}

isa_model/deployment/gpu_fp16_ds8/models/deepseek_r1/1/model.py

@@ -0,0 +1,120 @@
+import json
+import numpy as np
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import os 
+import triton_python_backend_utils as pb_utils
+
+class TritonPythonModel:
+    def initialize(self, args):
+        """初始化模型"""
+        self.model_config = json.loads(args['model_config'])
+        
+        # --- START: CORRECTED PATH LOGIC ---
+
+        # model_repository 是父目录, e.g., /models/deepseek_r1
+        model_repository = args['model_repository']
+        # model_version 是版本号, e.g., '1'
+        model_version = args['model_version']
+        
+        # 将它们拼接成指向模型文件的确切路径
+        model_path = os.path.join(model_repository, model_version)
+
+        print(f"Loading model from specific version path: {model_path}")
+
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            model_path,  # 从正确的版本目录加载
+            trust_remote_code=True
+        )
+        
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_path,  # 从正确的版本目录加载
+            torch_dtype=torch.bfloat16, 
+            device_map="gpu",           
+            trust_remote_code=True
+        )
+        
+        # --- END: CORRECTED PATH LOGIC ---
+        
+        # ... (您代码的其余部分保持不变) ...
+        output_config = pb_utils.get_output_config_by_name(
+            self.model_config, "OUTPUT_TEXT"
+        )
+        self.output_dtype = pb_utils.triton_string_to_numpy(
+            output_config['data_type']
+        )
+        
+        self.generation_config = {
+            'max_new_tokens': 512,
+            'temperature': 0.7,
+            'do_sample': True,
+            'top_p': 0.9,
+            'repetition_penalty': 1.1,
+            'pad_token_id': self.tokenizer.eos_token_id
+        }
+        
+        print("Model loaded successfully!")
+    
+    def execute(self, requests):
+        """执行推理"""
+        responses = []
+        
+        for request in requests:
+            # 获取输入文本
+            input_text = pb_utils.get_input_tensor_by_name(
+                request, "INPUT_TEXT"
+            ).as_numpy()
+            
+            # 解码输入文本
+            input_texts = [text.decode('utf-8') for text in input_text.flatten()]
+            
+            # 批量推理
+            output_texts = []
+            for text in input_texts:
+                try:
+                    # 编码输入
+                    inputs = self.tokenizer.encode(
+                        text, 
+                        return_tensors="pt"
+                    ).to(self.model.device)
+                    
+                    # 生成响应
+                    with torch.no_grad():
+                        outputs = self.model.generate(
+                            inputs,
+                            **self.generation_config
+                        )
+                    
+                    # 解码输出
+                    response = self.tokenizer.decode(
+                        outputs[0][inputs.shape[-1]:], 
+                        skip_special_tokens=True
+                    )
+                    
+                    output_texts.append(response)
+                    
+                except Exception as e:
+                    print(f"Error processing text: {e}")
+                    output_texts.append(f"Error: {str(e)}")
+            
+            # 准备输出
+            output_texts_np = np.array(
+                [[text.encode('utf-8')] for text in output_texts], 
+                dtype=object
+            )
+            
+            output_tensor = pb_utils.Tensor(
+                "OUTPUT_TEXT", 
+                output_texts_np.astype(self.output_dtype)
+            )
+            
+            response = pb_utils.InferenceResponse(
+                output_tensors=[output_tensor]
+            )
+            responses.append(response)
+        
+        return responses
+    
+    def finalize(self):
+        """清理资源"""
+        print("Cleaning up...")

isa_model/deployment/gpu_fp16_ds8/scripts/download_model.py

@@ -0,0 +1,18 @@
+from huggingface_hub import snapshot_download
+import os
+
+model_name = 'deepseek-ai/DeepSeek-R1-0528-Qwen3-8B'
+# 定义Triton模型仓库中该模型的版本路径
+local_model_path = os.path.join("models", "deepseek_r1", "1")
+
+print(f"开始下载模型 '{model_name}' 到 '{local_model_path}'...")
+
+# 使用 snapshot_download 下载整个模型仓库
+# 它会下载所有文件，包括.safetensors权重文件
+snapshot_download(
+    repo_id=model_name,
+    local_dir=local_model_path,
+    local_dir_use_symlinks=False,  
+)
+
+print("模型所有文件下载完成!")

isa_model/deployment/gpu_int8_ds8/app/server.py

@@ -0,0 +1,66 @@
+import os
+from fastapi import FastAPI
+from pydantic import BaseModel
+from contextlib import asynccontextmanager
+from pathlib import Path
+from threading import Thread
+from transformers import AutoTokenizer
+from tensorrt_llm.runtime import ModelRunner
+
+# --- 全局变量 ---
+ENGINE_PATH = "/app/built_engine/deepseek_engine"
+TOKENIZER_PATH = "/app/hf_model" # 我们需要原始HF模型中的tokenizer
+runner = None
+tokenizer = None
+
+# --- FastAPI生命周期事件 ---
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    global runner, tokenizer
+    print("--- 正在加载模型引擎和Tokenizer... ---")
+    tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_PATH, trust_remote_code=True)
+    runner = ModelRunner.from_dir(engine_dir=ENGINE_PATH, rank=0, stream=True)
+    print("--- ✅ 模型加载完毕，服务准备就绪 ---")
+    yield
+    print("--- 正在清理资源... ---")
+    runner = None
+    tokenizer = None
+
+app = FastAPI(lifespan=lifespan)
+
+# --- API请求和响应模型 ---
+class GenerateRequest(BaseModel):
+    prompt: str
+    max_new_tokens: int = 256
+    temperature: float = 0.7
+
+class GenerateResponse(BaseModel):
+    text: str
+
+# --- API端点 ---
+@app.post("/generate", response_model=GenerateResponse)
+async def generate(request: GenerateRequest):
+    print(f"收到请求: {request.prompt}")
+    
+    # 准备输入
+    input_ids = tokenizer.encode(request.prompt, return_tensors="pt").to("cuda")
+    
+    # 执行推理
+    output_ids = runner.generate(
+        input_ids,
+        max_new_tokens=request.max_new_tokens,
+        temperature=request.temperature,
+        eos_token_id=tokenizer.eos_token_id,
+        pad_token_id=tokenizer.pad_token_id,
+    )
+    
+    # 清理并解码输出
+    # output_ids[0] 的形状是 [beam_width, seq_length]
+    generated_text = tokenizer.decode(output_ids[0, 0, len(input_ids[0]):], skip_special_tokens=True)
+    
+    print(f"生成响应: {generated_text}")
+    return GenerateResponse(text=generated_text)
+
+@app.get("/health")
+async def health_check():
+    return {"status": "ok" if runner is not None else "loading"}

isa_model/deployment/gpu_int8_ds8/scripts/test_client.py

@@ -0,0 +1,43 @@
+import requests
+import json
+
+# --- 配置 ---
+TRITON_SERVER_URL = "http://localhost:8000"
+MODEL_NAME = "deepseek_trtllm"
+PROMPT = "请给我讲一个关于人工智能的笑话。"
+MAX_TOKENS = 256
+STREAM = False
+# ----------------------------------------------------
+
+def main():
+    """向Triton服务器发送请求并打印结果。"""
+    url = f"{TRITON_SERVER_URL}/v2/models/{MODEL_NAME}/generate"
+    payload = {
+        "text_input": PROMPT,
+        "max_new_tokens": MAX_TOKENS,
+        "temperature": 0.7,
+        "stream": STREAM
+    }
+    print(f"Sending request to: {url}")
+    print(f"Payload: {json.dumps(payload, indent=2, ensure_ascii=False)}")
+    print("-" * 30)
+
+    try:
+        response = requests.post(url, json=payload, headers={"Accept": "application/json"})
+        response.raise_for_status()
+        response_data = response.json()
+        generated_text = response_data.get('text_output', 'Error: "text_output" key not found.')
+        
+        print("✅ Request successful!")
+        print("-" * 30)
+        print("Prompt:", PROMPT)
+        print("\nGenerated Text:", generated_text)
+
+    except requests.exceptions.RequestException as e:
+        print(f"❌ Error making request to Triton server: {e}")
+        if e.response:
+            print(f"Response Status Code: {e.response.status_code}")
+            print(f"Response Body: {e.response.text}")
+
+if __name__ == '__main__':
+    main()

isa_model/deployment/gpu_int8_ds8/scripts/test_client_os.py

@@ -0,0 +1,35 @@
+import requests
+import json
+
+PROMPT = "请给我讲一个关于人工智能的笑话。"
+API_URL = "http://localhost:8000/generate"
+
+def main():
+    payload = {
+        "prompt": PROMPT,
+        "max_new_tokens": 100
+    }
+    
+    print(f"Sending request to: {API_URL}")
+    print(f"Payload: {json.dumps(payload, ensure_ascii=False)}")
+    print("-" * 30)
+
+    try:
+        response = requests.post(API_URL, json=payload)
+        response.raise_for_status()
+        
+        response_data = response.json()
+        generated_text = response_data.get('text')
+        
+        print("✅ Request successful!")
+        print("-" * 30)
+        print("Prompt:", PROMPT)
+        print("\nGenerated Text:", generated_text)
+
+    except requests.exceptions.RequestException as e:
+        print(f"❌ Error making request: {e}")
+        if e.response:
+            print(f"Response Body: {e.response.text}")
+
+if __name__ == '__main__':
+    main()

isa_model/inference/__init__.py

@@ -0,0 +1,11 @@
+"""
+Inference module for isA_Model
+
+File: isa_model/inference/__init__.py
+This module provides the main inference components for the IsA Model system.
+"""
+
+from .ai_factory import AIFactory
+from .base import ModelType, Capability, RoutingStrategy
+
+__all__ = ["AIFactory", "ModelType", "Capability", "RoutingStrategy"]