Singularity
Singularity Environment class
Full source code
#!/usr/bin/env python3
import logging
import os
import shutil
import subprocess
import tempfile
import uuid
from pathlib import Path
from typing import Any
from pydantic import BaseModel
from minisweagent.exceptions import Submitted
from minisweagent.utils.serialize import recursive_merge
class SingularityEnvironmentConfig(BaseModel):
image: str
cwd: str = "/"
env: dict[str, str] = {}
"""Environment variables to set in the container."""
forward_env: list[str] = []
"""Environment variables to forward to the container."""
timeout: int = 30
"""Timeout for executing commands in the container."""
executable: str = os.getenv("MSWEA_SINGULARITY_EXECUTABLE", "singularity")
"""Path to the singularity executable."""
sandbox_build_retries: int = 3
"""Number of retries for building the sandbox if an error occurs."""
global_args: list[str] = ["--quiet"]
"""Global arguments passed before the subcommand (e.g., --quiet, --debug)."""
exec_args: list[str] = ["--contain", "--cleanenv", "--fakeroot"]
"""Arguments passed to `singularity exec`."""
class SingularityEnvironment:
def __init__(
self, *, config_class: type = SingularityEnvironmentConfig, logger: logging.Logger | None = None, **kwargs
):
"""Singularity environment. See `SingularityEnvironmentConfig` for kwargs."""
self.logger = logger or logging.getLogger("minisweagent.environment")
self.config = config_class(**kwargs)
self.sandbox_dir = self._build_sandbox()
def _build_sandbox(self) -> Path:
# Building the sandbox can fail (very rarely), so we retry it
max_retries = self.config.sandbox_build_retries
for attempt in range(max_retries):
sandbox_dir = Path(tempfile.gettempdir()) / f"minisweagent-{uuid.uuid4().hex[:8]}"
try:
subprocess.run(
[self.config.executable, "build", "--sandbox", sandbox_dir, self.config.image],
check=True,
capture_output=True,
)
break
except subprocess.CalledProcessError as e:
shutil.rmtree(sandbox_dir, ignore_errors=True)
self.logger.error(
f"Error building image {self.config.image}, stdout: {e.stdout}, stderr: {e.stderr} (attempt {attempt + 1}/{max_retries})"
)
if attempt == max_retries - 1:
raise
return sandbox_dir
def get_template_vars(self, **kwargs) -> dict[str, Any]:
return recursive_merge(self.config.model_dump(), kwargs)
def serialize(self) -> dict:
return {
"info": {
"config": {
"environment": self.config.model_dump(mode="json"),
"environment_type": f"{self.__class__.__module__}.{self.__class__.__name__}",
}
}
}
def execute(self, action: dict, cwd: str = "", *, timeout: int | None = None) -> dict[str, Any]:
"""Execute a command in a Singularity container and return the result as a dict."""
command = action.get("command", "")
cmd = [self.config.executable, *self.config.global_args, "exec", *self.config.exec_args]
work_dir = cwd or self.config.cwd
if work_dir and work_dir != "/":
cmd.extend(["--pwd", work_dir])
for key in self.config.forward_env:
if (value := os.getenv(key)) is not None:
cmd.extend(["--env", f"{key}={value}"])
for key, value in self.config.env.items():
cmd.extend(["--env", f"{key}={value}"])
cmd.extend(["--writable", str(self.sandbox_dir), "bash", "-c", command])
try:
result = subprocess.run(
cmd,
text=True,
timeout=timeout or self.config.timeout,
encoding="utf-8",
errors="replace",
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
)
output = {"output": result.stdout, "returncode": result.returncode, "exception_info": ""}
except Exception as e:
raw_output = getattr(e, "output", None)
raw_output = (
raw_output.decode("utf-8", errors="replace") if isinstance(raw_output, bytes) else (raw_output or "")
)
output = {
"output": raw_output,
"returncode": -1,
"exception_info": f"An error occurred while executing the command: {e}",
"extra": {"exception_type": type(e).__name__, "exception": str(e)},
}
self._check_finished(output)
return output
def _check_finished(self, output: dict):
"""Raises Submitted if the output indicates task completion."""
lines = output.get("output", "").lstrip().splitlines(keepends=True)
if lines and lines[0].strip() == "COMPLETE_TASK_AND_SUBMIT_FINAL_OUTPUT" and output["returncode"] == 0:
submission = "".join(lines[1:])
raise Submitted(
{
"role": "exit",
"content": submission,
"extra": {"exit_status": "Submitted", "submission": submission},
}
)
def cleanup(self):
shutil.rmtree(self.sandbox_dir, ignore_errors=True)
def __del__(self):
"""Cleanup sandbox when object is destroyed."""
self.cleanup()
Bases: BaseModel
instance-attribute
image: str
class-attribute
instance-attribute
cwd: str = '/'
class-attribute
instance-attribute
env: dict[str, str] = {}
Environment variables to set in the container.
class-attribute
instance-attribute
forward_env: list[str] = []
Environment variables to forward to the container.
class-attribute
instance-attribute
timeout: int = 30
Timeout for executing commands in the container.
class-attribute
instance-attribute
executable: str = getenv(
"MSWEA_SINGULARITY_EXECUTABLE", "singularity"
)
Path to the singularity executable.
class-attribute
instance-attribute
sandbox_build_retries: int = 3
Number of retries for building the sandbox if an error occurs.
class-attribute
instance-attribute
global_args: list[str] = ['--quiet']
Global arguments passed before the subcommand (e.g., --quiet, --debug).
class-attribute
instance-attribute
exec_args: list[str] = [
"--contain",
"--cleanenv",
"--fakeroot",
]
Arguments passed to singularity exec.
SingularityEnvironment(
*,
config_class: type = SingularityEnvironmentConfig,
logger: Logger | None = None,
**kwargs,
)
Singularity environment. See SingularityEnvironmentConfig for kwargs.
Source code in src/minisweagent/environments/singularity.py
38 39 40 41 42 43 44 | |
instance-attribute
logger = logger or getLogger('minisweagent.environment')
instance-attribute
config = config_class(**kwargs)
instance-attribute
sandbox_dir = _build_sandbox()
get_template_vars(**kwargs) -> dict[str, Any]
Source code in src/minisweagent/environments/singularity.py
67 68 | |
serialize() -> dict
Source code in src/minisweagent/environments/singularity.py
70 71 72 73 74 75 76 77 78 | |
execute(
action: dict,
cwd: str = "",
*,
timeout: int | None = None,
) -> dict[str, Any]
Execute a command in a Singularity container and return the result as a dict.
Source code in src/minisweagent/environments/singularity.py
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 | |
cleanup()
Source code in src/minisweagent/environments/singularity.py
134 135 | |
