isa-model 0.1.0__py3-none-any.whl → 0.1.1__py3-none-any.whl

isa_model/__init__.py

@@ -2,4 +2,4 @@
 isA_Model - A simple interface for AI model integration
 """
 
-__version__ = "0.1.0" 
+__version__ = "0.1.1" 

isa_model/core/model_registry.py

@@ -3,6 +3,9 @@ from enum import Enum
 import logging
 from pathlib import Path
 import json
+import sqlite3
+from datetime import datetime
+import threading
 
 logger = logging.getLogger(__name__)
 
@@ -29,27 +32,45 @@ class ModelType(str, Enum):
     VISION = "vision"
 
 class ModelRegistry:
-    """Registry for model metadata and capabilities"""
+    """SQLite-based registry for model metadata and capabilities"""
     
-    def __init__(self, registry_file: str = "./models/model_registry.json"):
-        self.registry_file = Path(registry_file)
-        self.registry: Dict[str, Dict[str, Any]] = {}
-        self._load_registry()
+    def __init__(self, db_path: str = "./models/model_registry.db"):
+        self.db_path = Path(db_path)
+        self.db_path.parent.mkdir(parents=True, exist_ok=True)
+        self._lock = threading.Lock()
+        self._initialize_database()
     
-    def _load_registry(self):
-        """Load model registry from file"""
-        if self.registry_file.exists():
-            with open(self.registry_file, 'r') as f:
-                self.registry = json.load(f)
-        else:
-            self.registry = {}
-            self._save_registry()
-    
-    def _save_registry(self):
-        """Save model registry to file"""
-        self.registry_file.parent.mkdir(parents=True, exist_ok=True)
-        with open(self.registry_file, 'w') as f:
-            json.dump(self.registry, f, indent=2)
+    def _initialize_database(self):
+        """Initialize SQLite database with required tables"""
+        with sqlite3.connect(self.db_path) as conn:
+            conn.execute("""
+                CREATE TABLE IF NOT EXISTS models (
+                    model_id TEXT PRIMARY KEY,
+                    model_type TEXT NOT NULL,
+                    metadata TEXT,
+                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+                )
+            """)
+            
+            conn.execute("""
+                CREATE TABLE IF NOT EXISTS model_capabilities (
+                    model_id TEXT,
+                    capability TEXT,
+                    PRIMARY KEY (model_id, capability),
+                    FOREIGN KEY (model_id) REFERENCES models(model_id) ON DELETE CASCADE
+                )
+            """)
+            
+            conn.execute("""
+                CREATE INDEX IF NOT EXISTS idx_model_type ON models(model_type)
+            """)
+            
+            conn.execute("""
+                CREATE INDEX IF NOT EXISTS idx_capability ON model_capabilities(capability)
+            """)
+            
+            conn.commit()
     
     def register_model(self, 
                       model_id: str,
@@ -58,14 +79,30 @@ class ModelRegistry:
                       metadata: Dict[str, Any]) -> bool:
         """Register a model with its capabilities and metadata"""
         try:
-            self.registry[model_id] = {
-                "type": model_type,
-                "capabilities": [cap.value for cap in capabilities],
-                "metadata": metadata
-            }
-            self._save_registry()
+            with self._lock:
+                with sqlite3.connect(self.db_path) as conn:
+                    # Insert or update model
+                    conn.execute("""
+                        INSERT OR REPLACE INTO models 
+                        (model_id, model_type, metadata, updated_at)
+                        VALUES (?, ?, ?, CURRENT_TIMESTAMP)
+                    """, (model_id, model_type.value, json.dumps(metadata)))
+                    
+                    # Clear existing capabilities
+                    conn.execute("DELETE FROM model_capabilities WHERE model_id = ?", (model_id,))
+                    
+                    # Insert new capabilities
+                    for capability in capabilities:
+                        conn.execute("""
+                            INSERT INTO model_capabilities (model_id, capability)
+                            VALUES (?, ?)
+                        """, (model_id, capability.value))
+                    
+                    conn.commit()
+                    
             logger.info(f"Registered model {model_id}")
             return True
+            
         except Exception as e:
             logger.error(f"Failed to register model {model_id}: {e}")
             return False
@@ -73,43 +110,233 @@ class ModelRegistry:
     def unregister_model(self, model_id: str) -> bool:
         """Unregister a model"""
         try:
-            if model_id in self.registry:
-                del self.registry[model_id]
-                self._save_registry()
-                logger.info(f"Unregistered model {model_id}")
-                return True
-            return False
+            with self._lock:
+                with sqlite3.connect(self.db_path) as conn:
+                    cursor = conn.execute("DELETE FROM models WHERE model_id = ?", (model_id,))
+                    conn.commit()
+                    
+                    if cursor.rowcount > 0:
+                        logger.info(f"Unregistered model {model_id}")
+                        return True
+                    return False
+                    
         except Exception as e:
             logger.error(f"Failed to unregister model {model_id}: {e}")
             return False
     
     def get_model_info(self, model_id: str) -> Optional[Dict[str, Any]]:
         """Get model information"""
-        return self.registry.get(model_id)
+        try:
+            with sqlite3.connect(self.db_path) as conn:
+                conn.row_factory = sqlite3.Row
+                
+                # Get model info
+                model_row = conn.execute("""
+                    SELECT model_id, model_type, metadata, created_at, updated_at
+                    FROM models WHERE model_id = ?
+                """, (model_id,)).fetchone()
+                
+                if not model_row:
+                    return None
+                
+                # Get capabilities
+                capabilities = conn.execute("""
+                    SELECT capability FROM model_capabilities WHERE model_id = ?
+                """, (model_id,)).fetchall()
+                
+                model_info = {
+                    "model_id": model_row["model_id"],
+                    "type": model_row["model_type"],
+                    "capabilities": [cap["capability"] for cap in capabilities],
+                    "metadata": json.loads(model_row["metadata"]) if model_row["metadata"] else {},
+                    "created_at": model_row["created_at"],
+                    "updated_at": model_row["updated_at"]
+                }
+                
+                return model_info
+                
+        except Exception as e:
+            logger.error(f"Failed to get model info for {model_id}: {e}")
+            return None
     
     def get_models_by_type(self, model_type: ModelType) -> Dict[str, Dict[str, Any]]:
         """Get all models of a specific type"""
-        return {
-            model_id: info
-            for model_id, info in self.registry.items()
-            if info["type"] == model_type
-        }
+        try:
+            with sqlite3.connect(self.db_path) as conn:
+                conn.row_factory = sqlite3.Row
+                
+                models = conn.execute("""
+                    SELECT model_id, model_type, metadata, created_at, updated_at
+                    FROM models WHERE model_type = ?
+                """, (model_type.value,)).fetchall()
+                
+                result = {}
+                for model in models:
+                    model_id = model["model_id"]
+                    
+                    # Get capabilities for this model
+                    capabilities = conn.execute("""
+                        SELECT capability FROM model_capabilities WHERE model_id = ?
+                    """, (model_id,)).fetchall()
+                    
+                    result[model_id] = {
+                        "type": model["model_type"],
+                        "capabilities": [cap["capability"] for cap in capabilities],
+                        "metadata": json.loads(model["metadata"]) if model["metadata"] else {},
+                        "created_at": model["created_at"],
+                        "updated_at": model["updated_at"]
+                    }
+                
+                return result
+                
+        except Exception as e:
+            logger.error(f"Failed to get models by type {model_type}: {e}")
+            return {}
     
     def get_models_by_capability(self, capability: ModelCapability) -> Dict[str, Dict[str, Any]]:
         """Get all models with a specific capability"""
-        return {
-            model_id: info
-            for model_id, info in self.registry.items()
-            if capability.value in info["capabilities"]
-        }
+        try:
+            with sqlite3.connect(self.db_path) as conn:
+                conn.row_factory = sqlite3.Row
+                
+                models = conn.execute("""
+                    SELECT DISTINCT m.model_id, m.model_type, m.metadata, m.created_at, m.updated_at
+                    FROM models m
+                    JOIN model_capabilities mc ON m.model_id = mc.model_id
+                    WHERE mc.capability = ?
+                """, (capability.value,)).fetchall()
+                
+                result = {}
+                for model in models:
+                    model_id = model["model_id"]
+                    
+                    # Get all capabilities for this model
+                    capabilities = conn.execute("""
+                        SELECT capability FROM model_capabilities WHERE model_id = ?
+                    """, (model_id,)).fetchall()
+                    
+                    result[model_id] = {
+                        "type": model["model_type"],
+                        "capabilities": [cap["capability"] for cap in capabilities],
+                        "metadata": json.loads(model["metadata"]) if model["metadata"] else {},
+                        "created_at": model["created_at"],
+                        "updated_at": model["updated_at"]
+                    }
+                
+                return result
+                
+        except Exception as e:
+            logger.error(f"Failed to get models by capability {capability}: {e}")
+            return {}
     
     def has_capability(self, model_id: str, capability: ModelCapability) -> bool:
         """Check if a model has a specific capability"""
-        model_info = self.get_model_info(model_id)
-        if not model_info:
+        try:
+            with sqlite3.connect(self.db_path) as conn:
+                result = conn.execute("""
+                    SELECT 1 FROM model_capabilities 
+                    WHERE model_id = ? AND capability = ?
+                """, (model_id, capability.value)).fetchone()
+                
+                return result is not None
+                
+        except Exception as e:
+            logger.error(f"Failed to check capability for {model_id}: {e}")
             return False
-        return capability.value in model_info["capabilities"]
     
     def list_models(self) -> Dict[str, Dict[str, Any]]:
         """List all registered models"""
-        return self.registry 
+        try:
+            with sqlite3.connect(self.db_path) as conn:
+                conn.row_factory = sqlite3.Row
+                
+                models = conn.execute("""
+                    SELECT model_id, model_type, metadata, created_at, updated_at
+                    FROM models ORDER BY created_at DESC
+                """).fetchall()
+                
+                result = {}
+                for model in models:
+                    model_id = model["model_id"]
+                    
+                    # Get capabilities for this model
+                    capabilities = conn.execute("""
+                        SELECT capability FROM model_capabilities WHERE model_id = ?
+                    """, (model_id,)).fetchall()
+                    
+                    result[model_id] = {
+                        "type": model["model_type"],
+                        "capabilities": [cap["capability"] for cap in capabilities],
+                        "metadata": json.loads(model["metadata"]) if model["metadata"] else {},
+                        "created_at": model["created_at"],
+                        "updated_at": model["updated_at"]
+                    }
+                
+                return result
+                
+        except Exception as e:
+            logger.error(f"Failed to list models: {e}")
+            return {}
+    
+    def get_stats(self) -> Dict[str, Any]:
+        """Get registry statistics"""
+        try:
+            with sqlite3.connect(self.db_path) as conn:
+                # Count total models
+                total_models = conn.execute("SELECT COUNT(*) FROM models").fetchone()[0]
+                
+                # Count by type
+                type_counts = dict(conn.execute("""
+                    SELECT model_type, COUNT(*) FROM models GROUP BY model_type
+                """).fetchall())
+                
+                # Count by capability
+                capability_counts = dict(conn.execute("""
+                    SELECT capability, COUNT(*) FROM model_capabilities GROUP BY capability
+                """).fetchall())
+                
+                return {
+                    "total_models": total_models,
+                    "models_by_type": type_counts,
+                    "models_by_capability": capability_counts
+                }
+                
+        except Exception as e:
+            logger.error(f"Failed to get stats: {e}")
+            return {}
+    
+    def search_models(self, query: str) -> Dict[str, Dict[str, Any]]:
+        """Search models by name or metadata"""
+        try:
+            with sqlite3.connect(self.db_path) as conn:
+                conn.row_factory = sqlite3.Row
+                
+                models = conn.execute("""
+                    SELECT model_id, model_type, metadata, created_at, updated_at
+                    FROM models 
+                    WHERE model_id LIKE ? OR metadata LIKE ?
+                    ORDER BY created_at DESC
+                """, (f"%{query}%", f"%{query}%")).fetchall()
+                
+                result = {}
+                for model in models:
+                    model_id = model["model_id"]
+                    
+                    # Get capabilities for this model
+                    capabilities = conn.execute("""
+                        SELECT capability FROM model_capabilities WHERE model_id = ?
+                    """, (model_id,)).fetchall()
+                    
+                    result[model_id] = {
+                        "type": model["model_type"],
+                        "capabilities": [cap["capability"] for cap in capabilities],
+                        "metadata": json.loads(model["metadata"]) if model["metadata"] else {},
+                        "created_at": model["created_at"],
+                        "updated_at": model["updated_at"]
+                    }
+                
+                return result
+                
+        except Exception as e:
+            logger.error(f"Failed to search models with query '{query}': {e}")
+            return {} 

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()