feat: gather task resource info and reserve

Signed-off-by: elisam0 <elisam@nvidia.com>

Author: elisam0

responses_api_agents/anyterminal_agent/app.py

@@ -273,7 +273,9 @@ class AnyTerminalAgentConfig(BaseResponsesAPIAgentConfig):
     )
     tb_agent_timeout: int = 1800
     tb_eval_timeout: int = 300
-    apptainer_memory_limit_mb: int = 32768
+    apptainer_memory_limit_mb: int = 32768  # fallback container memory cap when a task omits memory_mb
+    agent_overhead_mb: int = 2048  # extra container memory on top of the task's memory_mb for the
+    # in-container agent harness
     concurrency: int = 256
 
 
@@ -497,6 +499,13 @@ def _apptainer_exec(
         env: str = "",
         workdir: Optional[str] = None,
     ) -> str:
+        # NOTE: we do NOT hard-enforce the task's cpu/memory envelope on the container. Apptainer
+        # cgroup flags (--cpus/--memory) are unusable here — rootless cgroups don't work under
+        # --fakeroot ("cannot use cgroups - rootless cgroups is not usable in fakeroot mode") — and
+        # per-task `ulimit -v` is unsafe (it caps virtual memory, which crashes torch/CUDA tasks at
+        # import). So the container keeps the generous global virtual-memory cap, and the task's
+        # cpu/memory footprint is only *reserved* in the Ray scheduler (_ray_resource_opts) to avoid
+        # host oversubscription — not enforced as a hard per-task limit.
         pwd_flag = f"--pwd {shlex.quote(workdir)} " if workdir else ""
         cmd = (
             f"apptainer exec --writable-tmpfs --fakeroot --cleanenv --pid --no-mount home,tmp,bind-paths "
@@ -619,8 +628,34 @@ async def responses(self, body: NeMoGymResponseCreateParamsNonStreaming = Body()
             print(f"[{params.task_name}] exception: see {tb_path}", file=sys.stderr)
             raise
 
+    @staticmethod
+    def _ray_resource_opts(params: AnyTerminalInstanceConfig) -> dict:
+        """Reserve the container's cpu/memory footprint in the Ray scheduler so concurrent containers
+        don't oversubscribe the host — the main cause of the compute-starved "productive-but-timed-out"
+        runs. The launcher task mostly awaits the apptainer subprocess, so these reservations are a
+        proxy for the container's real resource use. Since rootless cgroups can't hard-cap the
+        container here, this scheduling reservation is our only resource-isolation lever. Memory
+        reserves the task's memory_mb + agent_overhead_mb (the in-container Hermes harness)."""
+
+        def _f(key):
+            v = params.problem_info.get(key)
+            try:
+                return float(v) if v is not None else None
+            except (TypeError, ValueError):
+                return None
+
+        cpus = _f("cpus")
+        mem_mb = _f("memory_mb")
+        opts: dict = {"num_cpus": cpus if (cpus and cpus > 0) else 1}
+        if mem_mb and mem_mb > 0:
+            opts["memory"] = (int(mem_mb) + params.agent_overhead_mb) * 1024 * 1024
+        gpus = _f("gpus") or 0
+        if gpus > 0:
+            opts["num_gpus"] = gpus
+        return opts
+
     async def _inner_responses(self, params: AnyTerminalInstanceConfig) -> NeMoGymResponse:
-        await _run_remote.remote(params.model_dump())
+        await _run_remote.options(**self._ray_resource_opts(params)).remote(params.model_dump())
 
         persisted = TerminalBenchMetrics.model_validate_json(params.metrics_fpath.read_text())
         mask_sample = bool(persisted.container_timed_out or persisted.agent_timed_out)

responses_api_agents/anyterminal_agent/prepare.py

@@ -63,10 +63,12 @@ def _load_task_config(task_dir: Path) -> dict:
     with open(config_path, "rb") as f:
         config = tomllib.load(f)
 
-    docker_image = (config.get("environment") or {}).get("docker_image", "ubuntu:22.04")
+    env = config.get("environment") or {}
+    docker_image = env.get("docker_image", "ubuntu:22.04")
     problem_statement = instruction_path.read_text() if instruction_path.exists() else ""
     agent_timeout = (config.get("agent") or {}).get("timeout_sec", None)
     verifier_timeout = (config.get("verifier") or {}).get("timeout_sec", None)
+    resources = _parse_env_resources(env)
 
     # Parse last WORKDIR from Dockerfile (if present).
     workdir = None
@@ -82,6 +84,38 @@ def _load_task_config(task_dir: Path) -> dict:
         "agent_timeout_sec": agent_timeout,
         "verifier_timeout_sec": verifier_timeout,
         "workdir": workdir,
+        **resources,
+    }
+
+
+def _mem_to_mb(val) -> int | None:
+    """Normalize a memory/storage spec to MB. Accepts an int (already MB, e.g. ``memory_mb``)
+    or a string with a unit (``"2G"``, ``"512M"``, ``"10G"``)."""
+    if val is None:
+        return None
+    if isinstance(val, (int, float)):
+        return int(val)
+    s = str(val).strip().upper().rstrip("B")
+    units = {"K": 1 / 1024, "M": 1.0, "G": 1024.0, "T": 1024.0 * 1024}
+    if s and s[-1] in units:
+        return int(float(s[:-1]) * units[s[-1]])
+    return int(float(s))  # bare number → assume MB
+
+
+def _parse_env_resources(env: dict) -> dict:
+    """Pull cpu/memory/storage/gpu limits from a task.toml ``[environment]`` table, tolerating
+    both schema v1.0 (``memory = "2G"``) and v1.1 (``memory_mb = 2048``)."""
+    memory_mb = env.get("memory_mb")
+    if memory_mb is None:
+        memory_mb = _mem_to_mb(env.get("memory"))
+    storage_mb = env.get("storage_mb")
+    if storage_mb is None:
+        storage_mb = _mem_to_mb(env.get("storage"))
+    return {
+        "cpus": env.get("cpus"),
+        "memory_mb": memory_mb,
+        "storage_mb": storage_mb,
+        "gpus": env.get("gpus", 0),
     }
 
 
@@ -129,6 +163,10 @@ def _to_gym_row(task_dir: Path, task_cfg: dict) -> dict:
                 if task_cfg.get("verifier_timeout_sec") is not None
                 else None,
                 "workdir": task_cfg.get("workdir"),
+                "cpus": str(task_cfg["cpus"]) if task_cfg.get("cpus") is not None else None,
+                "memory_mb": str(task_cfg["memory_mb"]) if task_cfg.get("memory_mb") is not None else None,
+                "storage_mb": str(task_cfg["storage_mb"]) if task_cfg.get("storage_mb") is not None else None,
+                "gpus": str(task_cfg["gpus"]) if task_cfg.get("gpus") is not None else None,
             },
         },
     }