Spaces:
Sleeping
Sleeping
File size: 6,242 Bytes
1b35d41 364b897 1b35d41 364b897 1b35d41 364b897 1b35d41 364b897 1b35d41 364b897 1b35d41 364b897 1b35d41 364b897 1b35d41 364b897 1b35d41 364b897 1b35d41 364b897 1b35d41 364b897 1b35d41 364b897 1b35d41 364b897 1b35d41 364b897 1b35d41 364b897 1b35d41 364b897 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 | # Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.
"""Docker sandbox for running tests inside containers.
Falls back to running tests via venv when Docker is not available.
"""
from __future__ import annotations
import shutil
import subprocess
from dataclasses import dataclass
from pathlib import Path
@dataclass
class TestResult:
"""Result of a test execution."""
truncated_log: str = ""
full_log: str | None = None
exit_code: int | None = None
def _is_docker_available() -> bool:
"""Check if Docker CLI exists and daemon is running."""
if not shutil.which("docker"):
return False
try:
subprocess.run(["docker", "info"], capture_output=True, timeout=5, check=True)
return True
except Exception:
return False
def _truncate_log(full_log: str, max_lines: int = 100) -> str:
"""Truncate log to last N lines with line numbers."""
lines = full_log.splitlines(keepends=True)
num_lines = len(lines)
lines_to_show = lines[-max_lines:]
start_line = num_lines - len(lines_to_show) + 1
truncated = "".join(
f"{i}: {line}"
for i, line in enumerate(lines_to_show, start=start_line)
)
lines_before = max(0, num_lines - max_lines)
if lines_before > 0:
truncated = f"({lines_before} lines above)\n" + truncated
return truncated.rstrip()
class DockerSandbox:
"""Build Docker images and run containers with timeout/memory limits.
Falls back to running tests via venv when Docker is not available.
"""
def __init__(self, timeout: int = 600, memory_limit: str = "16g") -> None:
self.timeout = timeout
self.memory_limit = memory_limit
def run_tests(self, image_name: str, workspace_dir: str) -> TestResult:
"""Run tests. Uses Docker if available, otherwise runs via venv."""
if _is_docker_available():
return self._run_tests_docker(image_name, workspace_dir)
else:
return self._run_tests_venv(image_name, workspace_dir)
def _run_tests_docker(self, image_name: str, workspace_dir: str) -> TestResult:
"""Run the container, capture output, truncate to last 100 lines."""
container_name = image_name
try:
command = [
"timeout", "--foreground", "-s", "KILL", str(self.timeout),
"docker", "run",
"-v", f"{workspace_dir}:/work",
f"--memory={self.memory_limit}",
"--name", container_name,
image_name,
]
result = subprocess.run(
command, text=True, stdout=subprocess.PIPE,
stderr=subprocess.STDOUT, check=False,
)
return TestResult(
truncated_log=_truncate_log(result.stdout),
full_log=result.stdout,
exit_code=result.returncode,
)
except Exception as e:
return TestResult(truncated_log=f"Test execution failed: {e}", full_log=None, exit_code=None)
finally:
try:
self.remove_container(container_name)
except Exception:
pass
def _run_tests_venv(self, image_name: str, workspace_dir: str) -> TestResult:
"""Run tests using the venv created during setup (no Docker)."""
try:
venv_python = str(Path(workspace_dir) / ".venv" / "bin" / "python")
if not Path(venv_python).exists():
return TestResult(
truncated_log="No venv found. Setup may have failed.",
full_log=None, exit_code=1,
)
# Find the test script: image_name is like "brmc__shortbus_new"
base_name = image_name.replace("_new", "")
test_script = Path(workspace_dir) / f"test_{base_name}.sh"
if not test_script.exists():
candidates = list(Path(workspace_dir).glob("test_*.sh"))
test_script = candidates[0] if candidates else None
if not test_script or not test_script.exists():
return TestResult(
truncated_log="No test script found in workspace.",
full_log=None, exit_code=1,
)
# Parse the test command from the script
test_cmd = None
with open(test_script) as f:
for line in f:
line = line.strip()
if not line or line.startswith("#") or line.startswith("set "):
continue
test_cmd = line
break
if not test_cmd:
return TestResult(
truncated_log="Test script is empty.",
full_log=None, exit_code=1,
)
# Replace python with venv python
test_cmd = test_cmd.replace("python3 ", f"{venv_python} ").replace("python ", f"{venv_python} ")
print(f"[VENV-TEST] Running: {test_cmd}", flush=True)
result = subprocess.run(
test_cmd, shell=True, cwd=workspace_dir,
text=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
timeout=self.timeout,
)
return TestResult(
truncated_log=_truncate_log(result.stdout),
full_log=result.stdout,
exit_code=result.returncode,
)
except subprocess.TimeoutExpired:
return TestResult(
truncated_log=f"Test execution timed out after {self.timeout}s.",
full_log=None, exit_code=1,
)
except Exception as e:
return TestResult(
truncated_log=f"Test execution failed: {e}",
full_log=None, exit_code=None,
)
@staticmethod
def remove_container(container_name: str) -> None:
"""Force-remove a Docker container."""
subprocess.run(["docker", "rm", "-f", container_name], capture_output=True)
|