diff --git a/src/programbench/container.py b/src/programbench/container.py index 4c24e57..3477825 100644 --- a/src/programbench/container.py +++ b/src/programbench/container.py @@ -9,6 +9,7 @@ import uuid from pathlib import Path from typing import Any +import os from programbench.constants import DOCKER_CP_TIMEOUT, DOCKER_RUN_TIMEOUT @@ -35,6 +36,18 @@ def __init__( self.cpus = cpus self._name = f"programbench-{uuid.uuid4().hex[:12]}" run_args = list(run_args or []) + storage_size = os.environ.get("PROGRAMBENCH_DOCKER_STORAGE_SIZE", "50G").strip() + has_storage_opt_arg = any(arg == "--storage-opt" or arg.startswith("--storage-opt=") for arg in run_args) + memory = os.environ.get("PROGRAMBENCH_DOCKER_MEMORY", "60g").strip() + memory_swap = os.environ.get("PROGRAMBENCH_DOCKER_MEMORY_SWAP", memory).strip() + has_memory_arg = any(arg in {"--memory", "-m"} or arg.startswith("--memory=") for arg in run_args) + has_memory_swap_arg = any(arg == "--memory-swap" or arg.startswith("--memory-swap=") for arg in run_args) + if memory and not has_memory_arg: + run_args.extend(["--memory", memory]) + if memory_swap and not has_memory_swap_arg: + run_args.extend(["--memory-swap", memory_swap]) + if storage_size and not has_storage_opt_arg: + run_args.extend(["--storage-opt", f"size={storage_size}"]) env_dict = {"PYTEST_XDIST_AUTO_NUM_WORKERS": str(cpus), **(env or {})} env_args: list[str] = [] for key, value in env_dict.items():