isa-model 0.3.4__py3-none-any.whl → 0.3.6__py3-none-any.whl

isa_model/core/models/model_manager.py

@@ -0,0 +1,399 @@
+from typing import Dict, Optional, List, Any
+import logging
+from pathlib import Path
+from datetime import datetime
+from huggingface_hub import hf_hub_download, snapshot_download
+from huggingface_hub.errors import HfHubHTTPError
+from .model_storage import ModelStorage, LocalModelStorage
+from .model_repo import ModelRegistry, ModelType, ModelCapability
+from .model_billing_tracker import ModelBillingTracker, ModelOperationType
+from ..pricing_manager import PricingManager
+from ..config import ConfigManager
+
+logger = logging.getLogger(__name__)
+
+class ModelManager:
+    """
+    Model lifecycle management service
+    
+    Handles the complete model lifecycle:
+    - Model registration and metadata management
+    - Model downloads, versions, and caching  
+    - Cost tracking and billing across all operations
+    - Integration with model training, evaluation, and deployment
+    """
+    
+    
+    def __init__(self, 
+                 storage: Optional[ModelStorage] = None,
+                 registry: Optional[ModelRegistry] = None,
+                 billing_tracker: Optional[ModelBillingTracker] = None,
+                 pricing_manager: Optional[PricingManager] = None,
+                 config_manager: Optional[ConfigManager] = None):
+        self.storage = storage or LocalModelStorage()
+        self.registry = registry or ModelRegistry()
+        self.billing_tracker = billing_tracker or ModelBillingTracker(model_registry=self.registry)
+        self.pricing_manager = pricing_manager or PricingManager()
+        self.config_manager = config_manager or ConfigManager()
+    
+    def get_model_pricing(self, provider: str, model_name: str) -> Dict[str, float]:
+        """获取模型定价信息"""
+        pricing = self.pricing_manager.get_model_pricing(provider, model_name)
+        if pricing:
+            return {"input": pricing.input_cost, "output": pricing.output_cost}
+        return {"input": 0.0, "output": 0.0}
+    
+    def calculate_cost(self, provider: str, model_name: str, input_tokens: int, output_tokens: int) -> float:
+        """计算请求成本"""
+        return self.pricing_manager.calculate_cost(
+            provider=provider,
+            model_name=model_name,
+            input_units=input_tokens,
+            output_units=output_tokens
+        )
+    
+    def get_cheapest_model(self, provider: str, model_type: str = "llm") -> Optional[str]:
+        """获取最便宜的模型"""
+        result = self.pricing_manager.get_cheapest_model(
+            provider=provider, 
+            unit_type="token",
+            min_input_units=1000  # Assume 1K tokens for comparison
+        )
+        return result["model_name"] if result else None
+
+    async def get_model(self, 
+                       model_id: str, 
+                       repo_id: str,
+                       model_type: ModelType,
+                       capabilities: List[ModelCapability],
+                       revision: Optional[str] = None,
+                       force_download: bool = False) -> Optional[Path]:
+        """
+        Get model files, downloading if necessary
+        
+        Args:
+            model_id: Unique identifier for the model
+            repo_id: Hugging Face repository ID
+            model_type: Type of model (LLM, embedding, etc.)
+            capabilities: List of model capabilities
+            revision: Specific model version/tag
+            force_download: Force re-download even if cached
+            
+        Returns:
+            Path to the model files or None if failed
+        """
+        # Check if model is already downloaded
+        if not force_download:
+            model_path = await self.storage.load_model(model_id)
+            if model_path:
+                logger.info(f"Using cached model {model_id}")
+                return model_path
+        
+        try:
+            # Download model files
+            logger.info(f"Downloading model {model_id} from {repo_id}")
+            model_dir = Path(f"./models/temp/{model_id}")
+            model_dir.mkdir(parents=True, exist_ok=True)
+            
+            snapshot_download(
+                repo_id=repo_id,
+                revision=revision,
+                local_dir=model_dir,
+                local_dir_use_symlinks=False
+            )
+            
+            # Save model and metadata
+            metadata = {
+                "repo_id": repo_id,
+                "revision": revision,
+                "downloaded_at": str(Path(model_dir).stat().st_mtime)
+            }
+            
+            # Register model
+            self.registry.register_model(
+                model_id=model_id,
+                model_type=model_type,
+                capabilities=capabilities,
+                metadata=metadata
+            )
+            
+            # Save model files
+            await self.storage.save_model(model_id, str(model_dir), metadata)
+            
+            return await self.storage.load_model(model_id)
+            
+        except HfHubHTTPError as e:
+            logger.error(f"Failed to download model {model_id}: {e}")
+            return None
+        except Exception as e:
+            logger.error(f"Unexpected error downloading model {model_id}: {e}")
+            return None
+    
+    async def list_models(self) -> List[Dict[str, Any]]:
+        """List all downloaded models with their metadata"""
+        models = await self.storage.list_models()
+        return [
+            {
+                "model_id": model_id,
+                **metadata,
+                **(self.registry.get_model_info(model_id) or {})
+            }
+            for model_id, metadata in models.items()
+        ]
+    
+    async def remove_model(self, model_id: str) -> bool:
+        """Remove a model and its metadata"""
+        try:
+            # Remove from storage
+            storage_success = await self.storage.delete_model(model_id)
+            
+            # Unregister from registry
+            registry_success = self.registry.unregister_model(model_id)
+            
+            return storage_success and registry_success
+            
+        except Exception as e:
+            logger.error(f"Failed to remove model {model_id}: {e}")
+            return False
+    
+    async def get_model_info(self, model_id: str) -> Optional[Dict[str, Any]]:
+        """Get information about a specific model"""
+        storage_info = await self.storage.get_metadata(model_id)
+        registry_info = self.registry.get_model_info(model_id)
+        
+        if not storage_info and not registry_info:
+            return None
+            
+        return {
+            **(storage_info or {}),
+            **(registry_info or {})
+        }
+    
+    async def update_model(self, 
+                          model_id: str, 
+                          repo_id: str,
+                          model_type: ModelType,
+                          capabilities: List[ModelCapability],
+                          revision: Optional[str] = None) -> bool:
+        """Update a model to a new version"""
+        try:
+            return bool(await self.get_model(
+                model_id=model_id,
+                repo_id=repo_id,
+                model_type=model_type,
+                capabilities=capabilities,
+                revision=revision,
+                force_download=True
+            ))
+        except Exception as e:
+            logger.error(f"Failed to update model {model_id}: {e}")
+            return False 
+    
+    # === MODEL LIFECYCLE MANAGEMENT ===
+    
+    async def register_model_for_lifecycle(
+        self,
+        model_id: str,
+        model_type: ModelType,
+        capabilities: List[ModelCapability],
+        provider: str = "custom",
+        provider_model_name: Optional[str] = None,
+        metadata: Optional[Dict[str, Any]] = None
+    ) -> bool:
+        """
+        Register a model for lifecycle management
+        
+        Args:
+            model_id: Unique identifier for the model
+            model_type: Type of model (LLM, embedding, etc.)
+            capabilities: List of model capabilities
+            provider: Provider name for billing
+            provider_model_name: Provider-specific model name for pricing
+            metadata: Additional metadata
+            
+        Returns:
+            True if registration successful
+        """
+        try:
+            # Prepare metadata with billing info
+            full_metadata = metadata or {}
+            full_metadata.update({
+                "provider": provider,
+                "provider_model_name": provider_model_name or model_id,
+                "registered_for_lifecycle": True,
+                "lifecycle_stage": "registered"
+            })
+            
+            # Register in model registry
+            success = self.registry.register_model(
+                model_id=model_id,
+                model_type=model_type,
+                capabilities=capabilities,
+                metadata=full_metadata
+            )
+            
+            if success:
+                # Track registration operation
+                self.billing_tracker.track_model_usage(
+                    model_id=model_id,
+                    operation_type=ModelOperationType.STORAGE,
+                    provider=provider,
+                    service_type="model_management",
+                    operation="register_model",
+                    metadata={"stage": "registration"}
+                )
+                
+                logger.info(f"Successfully registered model {model_id} for lifecycle management")
+                
+            return success
+            
+        except Exception as e:
+            logger.error(f"Failed to register model {model_id} for lifecycle: {e}")
+            return False
+    
+    def track_model_usage(
+        self,
+        model_id: str,
+        operation_type: ModelOperationType,
+        provider: str,
+        service_type: str,
+        operation: str,
+        input_tokens: Optional[int] = None,
+        output_tokens: Optional[int] = None,
+        input_units: Optional[float] = None,
+        output_units: Optional[float] = None,
+        metadata: Optional[Dict[str, Any]] = None
+    ):
+        """
+        Track model usage and costs
+        
+        This method should be called by:
+        - Training services when training a model
+        - Evaluation services when evaluating a model  
+        - Deployment services when deploying a model
+        - Inference services when using a model for inference
+        """
+        return self.billing_tracker.track_model_usage(
+            model_id=model_id,
+            operation_type=operation_type,
+            provider=provider,
+            service_type=service_type,
+            operation=operation,
+            input_tokens=input_tokens,
+            output_tokens=output_tokens,
+            input_units=input_units,
+            output_units=output_units,
+            metadata=metadata
+        )
+    
+    async def update_model_stage(
+        self,
+        model_id: str,
+        new_stage: str,
+        metadata: Optional[Dict[str, Any]] = None
+    ) -> bool:
+        """
+        Update model lifecycle stage
+        
+        Args:
+            model_id: Model identifier
+            new_stage: New lifecycle stage (training, evaluation, deployment, production, retired)
+            metadata: Additional metadata for this stage
+            
+        Returns:
+            True if update successful
+        """
+        try:
+            # Get current model info
+            model_info = self.registry.get_model_info(model_id)
+            if not model_info:
+                logger.error(f"Model {model_id} not found in registry")
+                return False
+            
+            # Update metadata with new stage
+            current_metadata = model_info.get("metadata", {})
+            current_metadata.update({
+                "lifecycle_stage": new_stage,
+                "stage_updated_at": str(datetime.now()),
+                **(metadata or {})
+            })
+            
+            # Update in registry
+            success = self.registry.register_model(
+                model_id=model_id,
+                model_type=ModelType(model_info["type"]),
+                capabilities=[ModelCapability(cap) for cap in model_info["capabilities"]],
+                metadata=current_metadata
+            )
+            
+            if success:
+                logger.info(f"Updated model {model_id} to stage: {new_stage}")
+            
+            return success
+            
+        except Exception as e:
+            logger.error(f"Failed to update model {model_id} stage: {e}")
+            return False
+    
+    def get_model_lifecycle_summary(self, model_id: str) -> Optional[Dict[str, Any]]:
+        """
+        Get complete lifecycle summary for a model including costs
+        
+        Returns:
+            Dictionary with model info, lifecycle stage, and billing summary
+        """
+        try:
+            # Get model info from registry
+            model_info = self.registry.get_model_info(model_id)
+            if not model_info:
+                return None
+            
+            # Get billing summary from tracker
+            billing_summary = self.billing_tracker.get_model_usage_summary(model_id)
+            
+            return {
+                "model_id": model_id,
+                "model_info": model_info,
+                "billing_summary": billing_summary,
+                "current_stage": model_info.get("metadata", {}).get("lifecycle_stage", "unknown")
+            }
+            
+        except Exception as e:
+            logger.error(f"Failed to get lifecycle summary for {model_id}: {e}")
+            return None
+    
+    def list_models_by_stage(self, stage: str) -> List[Dict[str, Any]]:
+        """
+        List all models in a specific lifecycle stage
+        
+        Args:
+            stage: Lifecycle stage to filter by
+            
+        Returns:
+            List of model dictionaries
+        """
+        try:
+            all_models = self.registry.list_models()
+            stage_models = []
+            
+            for model_id, model_info in all_models.items():
+                current_stage = model_info.get("metadata", {}).get("lifecycle_stage")
+                if current_stage == stage:
+                    stage_models.append({
+                        "model_id": model_id,
+                        **model_info
+                    })
+            
+            return stage_models
+            
+        except Exception as e:
+            logger.error(f"Failed to list models by stage {stage}: {e}")
+            return []
+    
+    def get_billing_summary_by_operation(self, operation_type: ModelOperationType) -> Dict[str, Any]:
+        """Get billing summary for a specific operation type"""
+        return self.billing_tracker.get_operation_summary(operation_type)
+    
+    def print_model_costs(self, model_id: str):
+        """Print cost summary for a specific model"""
+        self.billing_tracker.print_model_summary(model_id)