test: add configurable memory, cpu and networking limits

Author: martvaha

app/api/base.py

@@ -97,13 +97,7 @@ async def execute_code(
                     continue
 
         # Execute code in Docker container
-        result = await docker_executor.execute(
-            code=request.code,
-            session_id=session_id,
-            lang=request.lang,
-            files=files,
-            timeout=settings.SANDBOX_MAX_EXECUTION_TIME,
-        )
+        result = await docker_executor.execute(code=request.code, session_id=session_id, lang=request.lang, files=files)
 
         # Add a language-specific error message if the stdout is empty
         if not result.get("stdout"):

app/services/docker_executor.py

@@ -34,6 +34,7 @@ class ContainerMetrics:
 @dataclass
 class FileState:
     """Tracks the state of a file for change detection."""
+
     path: Path
     size: int
     mtime: float
@@ -46,21 +47,17 @@ class DockerExecutor:
 
     WORK_DIR = "/mnt/data"  # Working directory will be the same as data mount point
     DATA_MOUNT = "/mnt/data"  # Mount point for session data
-    
+
     # Language-specific execution commands
     LANGUAGE_EXECUTORS = {
         "py": ["python", "-c"],
         "r": ["Rscript", "-e"],
     }
-    
+
     # Language-specific messages
     LANGUAGE_SPECIFIC_MESSAGES = {
-        "py": {
-            "empty_output": "Empty. Make sure to explicitly print() the results in Python"
-        },
-        "r": {
-            "empty_output": "Empty. Make sure to use print() or cat() to display results in R"
-        }
+        "py": {"empty_output": "Empty. Make sure to explicitly print() the results in Python"},
+        "r": {"empty_output": "Empty. Make sure to use print() or cat() to display results in R"},
     }
 
     def __init__(self):
@@ -113,63 +110,56 @@ def _scan_directory(self, directory: Path) -> Dict[str, FileState]:
         Returns a dictionary mapping relative file paths to their FileState objects.
         """
         file_states = {}
-        
+
         if not directory.exists():
             logger.warning(f"Directory {directory} does not exist")
             return file_states
-            
+
         # Walk through the directory recursively
         for root, _, files in os.walk(directory):
             root_path = Path(root)
-            
+
             # Compute relative path from the base directory
             rel_root = root_path.relative_to(directory)
-            
+
             for filename in files:
                 # Skip lock files
-                if filename.endswith('.lock'):
+                if filename.endswith(".lock"):
                     continue
-                    
+
                 file_path = root_path / filename
-                
+
                 # Compute relative path for dictionary key
-                if rel_root == Path('.'):
+                if rel_root == Path("."):
                     rel_path = filename
                 else:
                     rel_path = str(rel_root / filename)
-                
+
                 try:
                     # Get file stats
                     stat = file_path.stat()
                     size = stat.st_size
                     mtime = stat.st_mtime
-                    
+
                     # Calculate MD5 hash for content comparison
                     md5_hash = hashlib.md5(file_path.read_bytes()).hexdigest()
-                    
+
                     # Store file state
-                    file_states[rel_path] = FileState(
-                        path=file_path,
-                        size=size,
-                        mtime=mtime,
-                        md5_hash=md5_hash
-                    )
+                    file_states[rel_path] = FileState(path=file_path, size=size, mtime=mtime, md5_hash=md5_hash)
                     logger.debug(f"Scanned file: {rel_path}, size: {size}, hash: {md5_hash}")
                 except (PermissionError, FileNotFoundError) as e:
                     logger.warning(f"Error scanning file {file_path}: {str(e)}")
                     continue
-        
+
         return file_states
 
-    def _find_changed_files(self, 
-                           before_states: Dict[str, FileState], 
-                           after_states: Dict[str, FileState]) -> Set[str]:
+    def _find_changed_files(self, before_states: Dict[str, FileState], after_states: Dict[str, FileState]) -> Set[str]:
         """
         Compare before and after file states to identify new or modified files.
         Returns a set of relative paths of changed files.
         """
         changed_files = set()
-        
+
         # Find new or modified files
         for rel_path, after_state in after_states.items():
             if rel_path not in before_states:
@@ -179,20 +169,23 @@ def _find_changed_files(self,
             else:
                 before_state = before_states[rel_path]
                 # Check if file was modified (size, hash, or timestamp changed)
-                if (before_state.size != after_state.size or 
-                    before_state.md5_hash != after_state.md5_hash):
-                    logger.info(f"Modified file detected: {rel_path}, before={before_state.size}:{before_state.md5_hash}, after={after_state.size}:{after_state.md5_hash}")
+                if before_state.size != after_state.size or before_state.md5_hash != after_state.md5_hash:
+                    logger.info(
+                        f"Modified file detected: {rel_path}, before={before_state.size}:{before_state.md5_hash}, after={after_state.size}:{after_state.md5_hash}"
+                    )
                     changed_files.add(rel_path)
                 else:
                     logger.info(f"Unchanged file: {rel_path}, size={after_state.size}, hash={after_state.md5_hash}")
-        
+
         # Add debug logs for summarizing scan results
         for rel_path in before_states:
             if rel_path not in after_states:
                 logger.info(f"File deleted: {rel_path}")
-                
-        logger.info(f"Before scan: {len(before_states)} files, After scan: {len(after_states)} files, Changed: {len(changed_files)} files")
-        
+
+        logger.info(
+            f"Before scan: {len(before_states)} files, After scan: {len(after_states)} files, Changed: {len(changed_files)} files"
+        )
+
         return changed_files
 
     async def _update_container_metrics(self, container) -> None:
@@ -274,10 +267,11 @@ async def execute(
         session_id: str,
         lang: Literal["py", "r"],
         files: Optional[List[Dict[str, Any]]] = None,
-        timeout: int = 30,
+        config: Optional[Dict[str, Any]] = None,
     ) -> Dict[str, Any]:
         """Execute code in a Docker container with file management."""
         container = None
+        config = config or {}
 
         try:
             # Ensure Docker client is initialized and valid
@@ -355,14 +349,24 @@ async def execute(
                             logger.error(f"Error checking for image {image_name}: {str(e)}")
                             raise
 
+                    # Get container configuration, with provided config overriding settings
+                    memory_limit_mb = config.get("memory_limit_mb", settings.CONTAINER_MEMORY_LIMIT_MB)
+                    cpu_limit = config.get("cpu_limit", settings.CONTAINER_CPU_LIMIT)
+                    network_enabled = config.get("network_enabled", settings.DOCKER_NETWORK_ENABLED)
+
+                    logger.info(
+                        f"Container config - Memory: {memory_limit_mb}MB, CPU: {cpu_limit}, Network: {network_enabled}"
+                    )
+
                     # Create container config
                     config = {
                         "Image": image_name,
                         "Cmd": ["sleep", "infinity"],
                         "WorkingDir": self.WORK_DIR,
-                        "NetworkDisabled": True,
+                        "NetworkDisabled": not network_enabled,
                         "HostConfig": {
-                            "Memory": 512 * 1024 * 1024,  # 512MB in bytes
+                            "Memory": memory_limit_mb * 1024 * 1024,  # Convert MB to bytes
+                            "NanoCpus": int(cpu_limit * 1e9),  # Convert CPU cores to nano CPUs
                             "Mounts": [
                                 {
                                     "Type": "bind",
@@ -414,11 +418,11 @@ async def execute(
                     # Execute the code with the appropriate interpreter
                     logger.info(f"Code to execute: {code}")
                     logger.info(f"Language: {lang}")
-                    
+
                     # Get the execution command for the specified language
                     exec_cmd = self.LANGUAGE_EXECUTORS.get(lang, self.LANGUAGE_EXECUTORS["py"])
                     logger.info(f"Using execution command: {exec_cmd}")
-                    
+
                     # Execute the code with the appropriate interpreter
                     exec = await container.exec(cmd=[*exec_cmd, code], user="jovyan", stdout=True, stderr=True)
                     # Use raw API call to get output
@@ -450,7 +454,7 @@ async def execute(
                     output_files = []
                     existing_filenames = {file["name"] for file in (files or [])}
                     logger.info(f"Existing filenames: {existing_filenames}")
-                    
+
                     for rel_path in changed_file_paths:
                         file_path = session_path / rel_path
                         if file_path.is_file():
@@ -466,7 +470,7 @@ async def execute(
                             # Use directory structure in filepath if present
                             filepath = f"{session_id}/{rel_path}"
                             filename = Path(rel_path).name
-                            
+
                             file_data = {
                                 "id": file_id,
                                 "session_id": session_id,

app/shared/config.py

@@ -97,6 +97,11 @@ def LANGUAGE_CONTAINERS(self) -> Dict[str, str]:
 
     # Docker execution settings
     MAX_CONCURRENT_CONTAINERS: int = 10  # Maximum number of concurrent Docker containers
+    CONTAINER_MEMORY_LIMIT_MB: int = 512  # Memory limit for Docker containers in MB
+    CONTAINER_CPU_LIMIT: float = 1.0  # CPU limit for Docker containers (number of cores)
+
+    # Docker network settings
+    DOCKER_NETWORK_ENABLED: bool = False  # Whether Docker containers have network access
 
 
 @lru_cache()

tests/main/test_memory_limit.py

@@ -0,0 +1,121 @@
+import pytest
+import logging
+from app.services.docker_executor import docker_executor
+from app.shared.config import get_settings
+from app.utils.generate_id import generate_id
+
+# Setup logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+settings = get_settings()
+
+
+@pytest.fixture(scope="function", autouse=True)
+async def setup_docker():
+    """Setup and teardown Docker for tests."""
+    # Initialize Docker
+    await docker_executor.initialize()
+
+    # Yield control to tests
+    yield
+
+    # Cleanup
+    await docker_executor.close()
+
+
+@pytest.mark.asyncio
+async def test_memory_limit_enforced():
+    """Test that Docker container memory limits are enforced."""
+    # Set a low memory limit (50MB)
+    memory_limit_mb = 50
+    session_id = generate_id()
+
+    # Print test info
+    print(f"\nRunning memory limit test with limit: {memory_limit_mb}MB")
+
+    # Python code that attempts to allocate more memory than the limit
+    code = """
+import numpy as np
+# Try to allocate more memory than the limit
+# Each float64 is 8 bytes, so 13,000,000 elements is about 100MB
+try:
+    # Create a large array (> 50MB)
+    large_array = np.ones(13_000_000, dtype=np.float64)
+    print(f"Array created with shape {large_array.shape} and size {large_array.nbytes / (1024*1024):.2f} MB")
+except MemoryError:
+    print("MemoryError: Memory limit enforced successfully")
+except Exception as e:
+    print(f"Unexpected error: {type(e).__name__}: {e}")
+"""
+
+    # Call docker_executor directly instead of using the API
+    result = await docker_executor.execute(
+        code=code, session_id=session_id, lang="py", config={"memory_limit_mb": memory_limit_mb}
+    )
+
+    # Print the result for debugging
+    print(f"Memory limit test result: {result}")
+
+    # When memory limit is enforced, we could see either:
+    # 1. A status of 'error' with empty stdout/stderr (container killed by OOM)
+    # 2. A MemoryError message explicitly caught
+    # 3. A 'Killed' message in stderr
+
+    # Check if memory limit enforcement was detected
+    memory_error_detected = False
+
+    # Case 1: The container was terminated by OOM with status 'error'
+    if result["status"] == "error" and result["stdout"] == "" and result["stderr"] == "":
+        memory_error_detected = True
+        print("Memory limit enforcement detected: Container was terminated with error status")
+
+    # Case 2: Explicit MemoryError caught by Python
+    elif "MemoryError" in result["stdout"] or "MemoryError" in result["stderr"]:
+        memory_error_detected = True
+        print("Memory limit enforcement detected: MemoryError in output")
+
+    # Case 3: Process was killed due to memory limit
+    elif "Killed" in result["stderr"]:
+        memory_error_detected = True
+        print("Memory limit enforcement detected: Process was killed")
+
+    # Container could also be terminated abnormally
+    elif result["status"] == "error":
+        memory_error_detected = True
+        print("Memory limit enforcement detected: Container terminated abnormally")
+
+    # Assert that memory limit enforcement was detected in some form
+    assert memory_error_detected, f"Memory limit enforcement not detected. Output: {result}"
+
+
+@pytest.mark.asyncio
+async def test_memory_limit_adequate():
+    """Test that adequate memory limits allow execution."""
+    # Set a higher memory limit that should be sufficient
+    memory_limit_mb = 200
+    session_id = generate_id()
+
+    # Print test info
+    print(f"\nRunning memory limit test with adequate limit: {memory_limit_mb}MB")
+
+    # Python code that allocates memory but stays under the limit
+    code = """
+import numpy as np
+# Create an array that should fit within memory limits
+# Each float64 is 8 bytes, so ~13 million elements is about 100MB
+array = np.ones(13_000_000, dtype=np.float64)
+print(f"Array created with shape {array.shape} and size {array.nbytes / (1024*1024):.2f} MB")
+"""
+
+    # Call docker_executor directly instead of using the API
+    result = await docker_executor.execute(
+        code=code, session_id=session_id, lang="py", config={"memory_limit_mb": memory_limit_mb}
+    )
+
+    # Print the result for debugging
+    print(f"Adequate memory test result: {result}")
+
+    # The code should execute successfully
+    assert result["status"] == "ok"
+    assert "Array created with shape" in result["stdout"]
+    assert "MB" in result["stdout"]