From 8052d7aba8ab85b8f5530b017a03b9769c9cf7b6 Mon Sep 17 00:00:00 2001 From: CI-DEV <154627941+IlumCI@users.noreply.github.com> Date: Mon, 18 Aug 2025 19:45:45 +0300 Subject: [PATCH] Update multi_model_gpu_manager.py --- swarms/structs/multi_model_gpu_manager.py | 139 +++++++++++++++++++--- 1 file changed, 121 insertions(+), 18 deletions(-) diff --git a/swarms/structs/multi_model_gpu_manager.py b/swarms/structs/multi_model_gpu_manager.py index 8a945e82..1761a7c9 100644 --- a/swarms/structs/multi_model_gpu_manager.py +++ b/swarms/structs/multi_model_gpu_manager.py @@ -43,6 +43,8 @@ except ImportError: "Transformers package not found. HuggingFace models will not be supported." ) +from swarms.security import SwarmShieldIntegration, ShieldConfig + class ModelType(Enum): """Enum defining supported model types.""" @@ -298,35 +300,53 @@ class ModelGrid: allocation_strategy: GPUAllocationStrategy = GPUAllocationStrategy.MEMORY_OPTIMIZED, memory_buffer: float = 0.5, # GB buffer to leave on each GPU max_cpu_models: int = 0, # Maximum models to keep on CPU if no GPU space + max_workers: int = 4, use_multiprocessing: bool = True, + enable_logging: bool = True, log_level: str = "INFO", + shield_config: Optional[ShieldConfig] = None, + enable_security: bool = True, + security_level: str = "standard", ): """ - Initialize the model manager. + Initialize the ModelGrid. Args: allocation_strategy: Strategy for allocating models to GPUs memory_buffer: Memory buffer to leave on each GPU (in GB) max_cpu_models: Maximum number of models to keep on CPU if no GPU space + max_workers: Maximum number of worker processes use_multiprocessing: Whether to use multiprocessing for model execution + enable_logging: Whether to enable detailed logging log_level: Logging level + shield_config (ShieldConfig, optional): Security configuration for SwarmShield integration. Defaults to None. + enable_security (bool, optional): Whether to enable SwarmShield security features. Defaults to True. + security_level (str, optional): Pre-defined security level. Options: "basic", "standard", "enhanced", "maximum". Defaults to "standard". """ # Set log level - logger.remove() - logger.add(sys.stderr, level=log_level) - logger.add( - "gpu_model_manager.log", - rotation="100 MB", - retention="1 week", - level=log_level, - ) + if enable_logging: + logger.remove() + logger.add(sys.stderr, level=log_level) + logger.add( + "gpu_model_manager.log", + rotation="100 MB", + retention="1 week", + level=log_level, + ) - self.models: Dict[str, ModelMetadata] = {} - self.gpu_manager = GPUManager() self.allocation_strategy = allocation_strategy self.memory_buffer = memory_buffer self.max_cpu_models = max_cpu_models + self.max_workers = max_workers self.use_multiprocessing = use_multiprocessing + self.enable_logging = enable_logging + + # Initialize SwarmShield integration + self._initialize_swarm_shield(shield_config, enable_security, security_level) + + # Initialize components + self.gpu_manager = GPUManager() + self.models: Dict[str, ModelMetadata] = {} # Initialize locks and queues for multiprocessing self.manager = ( @@ -338,15 +358,98 @@ class ModelGrid: self.result_queues: Dict[str, Any] = ( self.manager.dict() if use_multiprocessing else {} ) - self.model_locks: Dict[str, Any] = {} - - logger.info( - f"ModelGrid initialized with {len(self.gpu_manager.gpus)} GPUs" - ) - logger.info( - f"Using allocation strategy: {allocation_strategy.value}" + self.model_locks: Dict[str, Any] = ( + self.manager.dict() if use_multiprocessing else {} ) + if self.enable_logging: + logger.info( + f"ModelGrid initialized with {len(self.gpu_manager.gpus)} GPUs" + ) + logger.info( + f"Using allocation strategy: {allocation_strategy.value}" + ) + + def _initialize_swarm_shield( + self, + shield_config: Optional[ShieldConfig] = None, + enable_security: bool = True, + security_level: str = "standard" + ) -> None: + """Initialize SwarmShield integration for security features.""" + self.enable_security = enable_security + self.security_level = security_level + + if enable_security: + if shield_config is None: + shield_config = ShieldConfig.get_security_level(security_level) + + self.swarm_shield = SwarmShieldIntegration(shield_config) + else: + self.swarm_shield = None + + # Security methods + def validate_task_with_shield(self, task: str) -> str: + """Validate and sanitize task input using SwarmShield.""" + if self.swarm_shield: + return self.swarm_shield.validate_and_protect_input(task) + return task + + def validate_agent_config_with_shield(self, agent_config: dict) -> dict: + """Validate agent configuration using SwarmShield.""" + if self.swarm_shield: + return self.swarm_shield.validate_and_protect_input(str(agent_config)) + return agent_config + + def process_agent_communication_with_shield(self, message: str, agent_name: str) -> str: + """Process agent communication through SwarmShield security.""" + if self.swarm_shield: + return self.swarm_shield.process_agent_communication(message, agent_name) + return message + + def check_rate_limit_with_shield(self, agent_name: str) -> bool: + """Check rate limits for an agent using SwarmShield.""" + if self.swarm_shield: + return self.swarm_shield.check_rate_limit(agent_name) + return True + + def add_secure_message(self, message: str, agent_name: str) -> None: + """Add a message to secure conversation history.""" + if self.swarm_shield: + self.swarm_shield.add_secure_message(message, agent_name) + + def get_secure_messages(self) -> List[dict]: + """Get secure conversation messages.""" + if self.swarm_shield: + return self.swarm_shield.get_secure_messages() + return [] + + def get_security_stats(self) -> dict: + """Get security statistics and metrics.""" + if self.swarm_shield: + return self.swarm_shield.get_security_stats() + return {"security_enabled": False} + + def update_shield_config(self, new_config: ShieldConfig) -> None: + """Update SwarmShield configuration.""" + if self.swarm_shield: + self.swarm_shield.update_config(new_config) + + def enable_security(self) -> None: + """Enable SwarmShield security features.""" + if not self.swarm_shield: + self._initialize_swarm_shield(enable_security=True, security_level=self.security_level) + + def disable_security(self) -> None: + """Disable SwarmShield security features.""" + self.swarm_shield = None + self.enable_security = False + + def cleanup_security(self) -> None: + """Clean up SwarmShield resources.""" + if self.swarm_shield: + self.swarm_shield.cleanup() + def add_model( self, model_name: str,