Compare commits
2 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 3e71dbc70b | |||
| 23160a0957 |
44
tests/test_resource_limits.py
Normal file
44
tests/test_resource_limits.py
Normal file
@@ -0,0 +1,44 @@
|
|||||||
|
"""
|
||||||
|
Tests for resource limits (#755).
|
||||||
|
"""
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
from tools.resource_limits import ResourceLimiter, ResourceLimits, ResourceResult, ResourceViolation
|
||||||
|
|
||||||
|
|
||||||
|
class TestResourceLimiter:
|
||||||
|
def test_successful_execution(self):
|
||||||
|
limiter = ResourceLimiter(ResourceLimits(memory_mb=2048, timeout_seconds=10))
|
||||||
|
result = limiter.execute("echo hello")
|
||||||
|
assert result.success is True
|
||||||
|
assert result.exit_code == 0
|
||||||
|
assert "hello" in result.stdout
|
||||||
|
assert result.violation == ResourceViolation.NONE
|
||||||
|
|
||||||
|
def test_timeout_violation(self):
|
||||||
|
limiter = ResourceLimiter(ResourceLimits(timeout_seconds=1))
|
||||||
|
result = limiter.execute("sleep 10")
|
||||||
|
assert result.success is False
|
||||||
|
assert result.violation == ResourceViolation.TIME
|
||||||
|
assert result.killed is True
|
||||||
|
|
||||||
|
def test_failed_command(self):
|
||||||
|
limiter = ResourceLimiter()
|
||||||
|
result = limiter.execute("exit 1")
|
||||||
|
assert result.success is False
|
||||||
|
assert result.exit_code == 1
|
||||||
|
|
||||||
|
def test_resource_report(self):
|
||||||
|
from tools.resource_limits import format_resource_report
|
||||||
|
result = ResourceResult(
|
||||||
|
success=True, stdout="", stderr="", exit_code=0,
|
||||||
|
violation=ResourceViolation.NONE, violation_message="",
|
||||||
|
memory_used_mb=100, cpu_time_seconds=0.5, wall_time_seconds=1.0,
|
||||||
|
)
|
||||||
|
report = format_resource_report(result)
|
||||||
|
assert "Exit code: 0" in report
|
||||||
|
assert "100MB" in report
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
pytest.main([__file__])
|
||||||
249
tools/resource_limits.py
Normal file
249
tools/resource_limits.py
Normal file
@@ -0,0 +1,249 @@
|
|||||||
|
"""
|
||||||
|
Terminal Sandbox Resource Limits — CPU, memory, time.
|
||||||
|
|
||||||
|
Provides resource limits for agent terminal commands to prevent
|
||||||
|
OOM kills, runaway processes, and excessive resource consumption.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import logging
|
||||||
|
import os
|
||||||
|
import signal
|
||||||
|
import subprocess
|
||||||
|
import time
|
||||||
|
from dataclasses import dataclass, field
|
||||||
|
from typing import Optional, Dict, Any
|
||||||
|
from enum import Enum
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
class ResourceViolation(Enum):
|
||||||
|
"""Types of resource violations."""
|
||||||
|
MEMORY = "memory"
|
||||||
|
CPU = "cpu"
|
||||||
|
TIME = "time"
|
||||||
|
NONE = "none"
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class ResourceLimits:
|
||||||
|
"""Resource limits for a subprocess."""
|
||||||
|
memory_mb: int = 2048 # 2GB default
|
||||||
|
cpu_percent: int = 80 # 80% of one core
|
||||||
|
timeout_seconds: int = 300 # 5 minutes
|
||||||
|
kill_timeout: int = 10 # SIGKILL after 10s if SIGTERM fails
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class ResourceResult:
|
||||||
|
"""Result of a resource-limited execution."""
|
||||||
|
success: bool
|
||||||
|
stdout: str
|
||||||
|
stderr: str
|
||||||
|
exit_code: int
|
||||||
|
violation: ResourceViolation
|
||||||
|
violation_message: str
|
||||||
|
memory_used_mb: float
|
||||||
|
cpu_time_seconds: float
|
||||||
|
wall_time_seconds: float
|
||||||
|
killed: bool = False
|
||||||
|
|
||||||
|
|
||||||
|
class ResourceLimiter:
|
||||||
|
"""Apply resource limits to subprocess execution."""
|
||||||
|
|
||||||
|
def __init__(self, limits: Optional[ResourceLimits] = None):
|
||||||
|
self.limits = limits or ResourceLimits()
|
||||||
|
|
||||||
|
def _get_resource_rlimit(self) -> Dict[str, Any]:
|
||||||
|
"""Get resource limits for subprocess (Unix only)."""
|
||||||
|
import resource
|
||||||
|
|
||||||
|
rlimit = {}
|
||||||
|
|
||||||
|
# Memory limit (RSS)
|
||||||
|
if self.limits.memory_mb > 0:
|
||||||
|
mem_bytes = self.limits.memory_mb * 1024 * 1024
|
||||||
|
rlimit[resource.RLIMIT_AS] = (mem_bytes, mem_bytes)
|
||||||
|
|
||||||
|
# CPU time limit
|
||||||
|
if self.limits.timeout_seconds > 0:
|
||||||
|
rlimit[resource.RLIMIT_CPU] = (self.limits.timeout_seconds, self.limits.timeout_seconds)
|
||||||
|
|
||||||
|
return rlimit
|
||||||
|
|
||||||
|
def _check_resource_usage(self, process: subprocess.Popen) -> Dict[str, float]:
|
||||||
|
"""Check resource usage of a process (Unix only)."""
|
||||||
|
try:
|
||||||
|
import resource
|
||||||
|
usage = resource.getrusage(resource.RUSAGE_CHILDREN)
|
||||||
|
return {
|
||||||
|
"user_time": usage.ru_utime,
|
||||||
|
"system_time": usage.ru_stime,
|
||||||
|
"max_rss_mb": usage.ru_maxrss / 1024, # KB to MB
|
||||||
|
}
|
||||||
|
except:
|
||||||
|
return {"user_time": 0, "system_time": 0, "max_rss_mb": 0}
|
||||||
|
|
||||||
|
def execute(self, command: str, **kwargs) -> ResourceResult:
|
||||||
|
"""
|
||||||
|
Execute a command with resource limits.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
command: Shell command to execute
|
||||||
|
**kwargs: Additional subprocess arguments
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
ResourceResult with execution details
|
||||||
|
"""
|
||||||
|
start_time = time.time()
|
||||||
|
|
||||||
|
# Try to use resource limits (Unix only)
|
||||||
|
preexec_fn = None
|
||||||
|
try:
|
||||||
|
import resource
|
||||||
|
rlimit = self._get_resource_rlimit()
|
||||||
|
|
||||||
|
def set_limits():
|
||||||
|
for res, limits in rlimit.items():
|
||||||
|
resource.setrlimit(res, limits)
|
||||||
|
|
||||||
|
preexec_fn = set_limits
|
||||||
|
except ImportError:
|
||||||
|
logger.debug("resource module not available, skipping limits")
|
||||||
|
|
||||||
|
try:
|
||||||
|
# Execute with timeout
|
||||||
|
result = subprocess.run(
|
||||||
|
command,
|
||||||
|
shell=True,
|
||||||
|
capture_output=True,
|
||||||
|
text=True,
|
||||||
|
timeout=self.limits.timeout_seconds,
|
||||||
|
preexec_fn=preexec_fn,
|
||||||
|
**kwargs,
|
||||||
|
)
|
||||||
|
|
||||||
|
wall_time = time.time() - start_time
|
||||||
|
usage = self._check_resource_usage(result)
|
||||||
|
|
||||||
|
# Check for violations
|
||||||
|
violation = ResourceViolation.NONE
|
||||||
|
violation_message = ""
|
||||||
|
|
||||||
|
# Check memory (if we can get it)
|
||||||
|
if usage["max_rss_mb"] > self.limits.memory_mb:
|
||||||
|
violation = ResourceViolation.MEMORY
|
||||||
|
violation_message = f"Memory limit exceeded: {usage['max_rss_mb']:.0f}MB > {self.limits.memory_mb}MB"
|
||||||
|
|
||||||
|
return ResourceResult(
|
||||||
|
success=result.returncode == 0,
|
||||||
|
stdout=result.stdout,
|
||||||
|
stderr=result.stderr,
|
||||||
|
exit_code=result.returncode,
|
||||||
|
violation=violation,
|
||||||
|
violation_message=violation_message,
|
||||||
|
memory_used_mb=usage["max_rss_mb"],
|
||||||
|
cpu_time_seconds=usage["user_time"] + usage["system_time"],
|
||||||
|
wall_time_seconds=wall_time,
|
||||||
|
)
|
||||||
|
|
||||||
|
except subprocess.TimeoutExpired as e:
|
||||||
|
wall_time = time.time() - start_time
|
||||||
|
|
||||||
|
# Try to kill gracefully
|
||||||
|
if hasattr(e, 'process') and e.process:
|
||||||
|
try:
|
||||||
|
e.process.terminate()
|
||||||
|
time.sleep(self.limits.kill_timeout)
|
||||||
|
if e.process.poll() is None:
|
||||||
|
e.process.kill()
|
||||||
|
except:
|
||||||
|
pass
|
||||||
|
|
||||||
|
return ResourceResult(
|
||||||
|
success=False,
|
||||||
|
stdout=e.stdout.decode() if e.stdout else "",
|
||||||
|
stderr=e.stderr.decode() if e.stderr else "",
|
||||||
|
exit_code=-1,
|
||||||
|
violation=ResourceViolation.TIME,
|
||||||
|
violation_message=f"Timeout after {self.limits.timeout_seconds}s",
|
||||||
|
memory_used_mb=0,
|
||||||
|
cpu_time_seconds=0,
|
||||||
|
wall_time_seconds=wall_time,
|
||||||
|
killed=True,
|
||||||
|
)
|
||||||
|
|
||||||
|
except MemoryError:
|
||||||
|
wall_time = time.time() - start_time
|
||||||
|
return ResourceResult(
|
||||||
|
success=False,
|
||||||
|
stdout="",
|
||||||
|
stderr=f"Memory limit exceeded ({self.limits.memory_mb}MB)",
|
||||||
|
exit_code=-1,
|
||||||
|
violation=ResourceViolation.MEMORY,
|
||||||
|
violation_message=f"Memory limit exceeded: {self.limits.memory_mb}MB",
|
||||||
|
memory_used_mb=self.limits.memory_mb,
|
||||||
|
cpu_time_seconds=0,
|
||||||
|
wall_time_seconds=wall_time,
|
||||||
|
killed=True,
|
||||||
|
)
|
||||||
|
|
||||||
|
except Exception as e:
|
||||||
|
wall_time = time.time() - start_time
|
||||||
|
return ResourceResult(
|
||||||
|
success=False,
|
||||||
|
stdout="",
|
||||||
|
stderr=str(e),
|
||||||
|
exit_code=-1,
|
||||||
|
violation=ResourceViolation.NONE,
|
||||||
|
violation_message=f"Execution error: {e}",
|
||||||
|
memory_used_mb=0,
|
||||||
|
cpu_time_seconds=0,
|
||||||
|
wall_time_seconds=wall_time,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def format_resource_report(result: ResourceResult) -> str:
|
||||||
|
"""Format resource usage as a report string."""
|
||||||
|
lines = [
|
||||||
|
f"Exit code: {result.exit_code}",
|
||||||
|
f"Wall time: {result.wall_time_seconds:.2f}s",
|
||||||
|
f"CPU time: {result.cpu_time_seconds:.2f}s",
|
||||||
|
f"Memory: {result.memory_used_mb:.0f}MB",
|
||||||
|
]
|
||||||
|
|
||||||
|
if result.violation != ResourceViolation.NONE:
|
||||||
|
lines.append(f"⚠️ Violation: {result.violation_message}")
|
||||||
|
|
||||||
|
if result.killed:
|
||||||
|
lines.append("🔴 Process killed")
|
||||||
|
|
||||||
|
return " | ".join(lines)
|
||||||
|
|
||||||
|
|
||||||
|
def execute_with_limits(
|
||||||
|
command: str,
|
||||||
|
memory_mb: int = 2048,
|
||||||
|
cpu_percent: int = 80,
|
||||||
|
timeout_seconds: int = 300,
|
||||||
|
) -> ResourceResult:
|
||||||
|
"""
|
||||||
|
Convenience function to execute with resource limits.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
command: Shell command
|
||||||
|
memory_mb: Memory limit in MB
|
||||||
|
cpu_percent: CPU limit as percent of one core
|
||||||
|
timeout_seconds: Timeout in seconds
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
ResourceResult
|
||||||
|
"""
|
||||||
|
limits = ResourceLimits(
|
||||||
|
memory_mb=memory_mb,
|
||||||
|
cpu_percent=cpu_percent,
|
||||||
|
timeout_seconds=timeout_seconds,
|
||||||
|
)
|
||||||
|
limiter = ResourceLimiter(limits)
|
||||||
|
return limiter.execute(command)
|
||||||
Reference in New Issue
Block a user