From 826e28b1989fef4229e444d7e7a1b762409f53d8 Mon Sep 17 00:00:00 2001 From: Tomsawyerhu <46986832+Tomsawyerhu@users.noreply.github.com> Date: Tue, 2 Jun 2026 19:02:58 +0800 Subject: [PATCH] Update container.py --- src/programbench/container.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/src/programbench/container.py b/src/programbench/container.py index 4c24e57..2e49d5b 100644 --- a/src/programbench/container.py +++ b/src/programbench/container.py @@ -35,6 +35,14 @@ def __init__( self.cpus = cpus self._name = f"programbench-{uuid.uuid4().hex[:12]}" run_args = list(run_args or []) + memory = os.environ.get("PROGRAMBENCH_DOCKER_MEMORY", "20g").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]) env_dict = {"PYTEST_XDIST_AUTO_NUM_WORKERS": str(cpus), **(env or {})} env_args: list[str] = [] for key, value in env_dict.items():