Merge pull request #7 from modelscope/add/pool

[Feature] Add sandbox pools

Author: Yunnglin

examples/sandbox_usage_examples.py

@@ -13,7 +13,7 @@ async def direct_sandbox_example():
 
     # Create Docker sandbox configuration
     config = DockerSandboxConfig(
-        image='python-sandbox',
+        image='python:3.11-slim',
         timeout=30,
         memory_limit='512m',
         cpu_limit=1.0,

ms_enclave/sandbox/manager/__init__.py

@@ -1,11 +1,12 @@
 """Sandbox manager implementations."""
 
-from .base import SandboxManager
+from .base import SandboxManager, SandboxManagerFactory
 from .http_manager import HttpSandboxManager
 from .local_manager import LocalSandboxManager
 
 __all__ = [
     'SandboxManager',
+    'SandboxManagerFactory',
     'LocalSandboxManager',
     'HttpSandboxManager',
 ]

ms_enclave/sandbox/manager/base.py

@@ -1,9 +1,19 @@
 """Base sandbox manager interface."""
 
+import asyncio
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
-
-from ..model import SandboxConfig, SandboxInfo, SandboxStatus, SandboxType, ToolResult
+from collections import deque
+from typing import TYPE_CHECKING, Any, Deque, Dict, List, Optional, Union
+
+from ..model import (
+    SandboxConfig,
+    SandboxInfo,
+    SandboxManagerConfig,
+    SandboxManagerType,
+    SandboxStatus,
+    SandboxType,
+    ToolResult,
+)
 
 if TYPE_CHECKING:
     from ..boxes import Sandbox
@@ -12,10 +22,18 @@
 class SandboxManager(ABC):
     """Abstract base class for sandbox managers."""
 
-    def __init__(self):
-        """Initialize the sandbox manager."""
+    def __init__(self, config: Optional[SandboxManagerConfig] = None, **kwargs):
+        """Initialize the sandbox manager.
+
+        Args:
+            config: Sandbox manager configuration
+        """
+        self.config = config or SandboxManagerConfig()
         self._running = False
         self._sandboxes: Dict[str, 'Sandbox'] = {}
+        self._sandbox_pool: Deque[str] = deque()
+        self._pool_lock = asyncio.Lock()
+        self._pool_initialized = False
 
     @abstractmethod
     async def start(self) -> None:
@@ -144,6 +162,48 @@ async def cleanup_all_sandboxes(self) -> None:
         """Clean up all sandboxes."""
         pass
 
+    @abstractmethod
+    async def initialize_pool(
+        self,
+        pool_size: Optional[int] = None,
+        sandbox_type: Optional[SandboxType] = None,
+        config: Optional[Union[SandboxConfig, Dict]] = None
+    ) -> List[str]:
+        """Initialize sandbox pool.
+
+        Args:
+            pool_size: Number of sandboxes in pool (uses config if not provided)
+            sandbox_type: Type of sandbox to create
+            config: Sandbox configuration (uses config.sandbox_config if not provided)
+
+        Returns:
+            List of created sandbox IDs
+        """
+        pass
+
+    @abstractmethod
+    async def execute_tool_in_pool(
+        self, tool_name: str, parameters: Dict[str, Any], timeout: Optional[float] = None
+    ) -> ToolResult:
+        """Execute tool using an available sandbox from the pool.
+
+        Uses FIFO queue to get an idle sandbox, marks it as busy during execution,
+        then returns it to the pool as idle.
+
+        Args:
+            tool_name: Tool name to execute
+            parameters: Tool parameters
+            timeout: Optional timeout for waiting for available sandbox
+
+        Returns:
+            Tool execution result
+
+        Raises:
+            ValueError: If pool is empty or no sandbox available
+            TimeoutError: If timeout waiting for available sandbox
+        """
+        pass
+
     # Context manager support
     async def __aenter__(self):
         """Async context manager entry."""
@@ -153,3 +213,81 @@ async def __aenter__(self):
     async def __aexit__(self, exc_type, exc_val, exc_tb):
         """Async context manager exit."""
         await self.stop()
+
+
+class SandboxManagerFactory:
+    """Factory for creating sandbox managers."""
+
+    _registry: Dict[SandboxManagerType, type] = {}
+
+    @classmethod
+    def register(cls, manager_type: SandboxManagerType, manager_class: type) -> None:
+        """Register a sandbox manager class.
+
+        Args:
+            manager_type: Manager type to register
+            manager_class: Manager class to register
+        """
+        cls._registry[manager_type] = manager_class
+
+    @classmethod
+    def create_manager(
+        cls,
+        manager_type: Optional[SandboxManagerType] = None,
+        config: Optional[SandboxManagerConfig] = None,
+        **kwargs
+    ) -> SandboxManager:
+        """Create a sandbox manager instance.
+
+        Args:
+            manager_type: Type of manager to create
+            config: Manager configuration
+            **kwargs: Additional manager-specific parameters
+
+        Returns:
+            Sandbox manager instance
+
+        Raises:
+            ValueError: If manager type is not registered
+        """
+        if not manager_type:
+            api_url = config.base_url if config else kwargs.get('base_url')
+            if api_url:
+                manager_type = SandboxManagerType.HTTP
+            else:
+                manager_type = SandboxManagerType.LOCAL
+
+        if manager_type not in cls._registry:
+            raise ValueError(
+                f"Sandbox manager type '{manager_type}' not registered. "
+                f'Available types: {list(cls._registry.keys())}'
+            )
+
+        manager_class = cls._registry[manager_type]
+        return manager_class(config=config, **kwargs)
+
+    @classmethod
+    def get_registered_types(cls) -> List[SandboxManagerType]:
+        """Get list of registered manager types.
+
+        Returns:
+            List of registered manager types
+        """
+        return list(cls._registry.keys())
+
+
+def register_manager(manager_type: SandboxManagerType):
+    """Decorator to register a sandbox manager class.
+
+    Args:
+        manager_type: Manager type to register
+
+    Returns:
+        Decorator function
+    """
+
+    def decorator(manager_class: type) -> type:
+        SandboxManagerFactory.register(manager_type, manager_class)
+        return manager_class
+
+    return decorator