From b26c89ae751bfdb3e508134c0a564aed7050695d Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Mon, 19 Oct 2026 09:25:46 +0000 Subject: [PATCH] Run verl's Harbor agent loop on boat sandboxes Keep the admission gate and the boat quota per running event loop (LoopLocal): both hold asyncio primitives, which bind to the first loop that waits on them. verl's AgentLoopWorker is a Ray async actor, so this stays one per worker process; a process with several loops no longer fails with 'bound to a different event loop'. The Miles gate uses the same helper. HarborAgentLoop takes the provider from HARBOR_ENV_TYPE when the row has none, and refuses a boat task when YETO_HARBOR_MAX_IN_FLIGHT exceeds the boat active-sandbox cap. BoatEnvironment accepts an explicit quota. New CPU test drives HarborAgentLoop.run() through the real Harbor Trial and boat provider on a fake boat account (VM created and deleted, tokens assembled, ABORTED on a failed container start). Co-Authored-By: Claude Opus 5.5 --- tests/test_agentic_harbor.py | 11 +- tests/test_harbor_boat.py | 2 +- tests/test_verl_harbor_boat.py | 147 +++++++++++++++++++++ yeto/rl/adapters/miles/harbor_agent.py | 16 ++- yeto/rl/adapters/verl/harbor_agent_loop.py | 62 ++++++++- yeto/rl/agentic/loop_local.py | 47 +++++++ yeto/rl/harness/harbor/README.md | 16 +++ yeto/rl/harness/harbor/boat_client.py | 21 ++- yeto/rl/harness/harbor/boat_env.py | 11 +- 9 files changed, 306 insertions(+), 27 deletions(-) create mode 100644 tests/test_verl_harbor_boat.py create mode 100644 yeto/rl/agentic/loop_local.py diff --git a/tests/test_agentic_harbor.py b/tests/test_agentic_harbor.py index 99833925a..c3178396b 100644 --- a/tests/test_agentic_harbor.py +++ b/tests/test_agentic_harbor.py @@ -274,7 +274,7 @@ def build_trial_config(metadata, session_url, request_kwargs): monkeypatch.setitem(sys.modules, "harbor_agent_function", example) monkeypatch.setenv("HARBOR_ENV_TYPE", "modal") mod = importlib.import_module("yeto.rl.adapters.miles.harbor_agent") - mod.gate.cache_clear() + mod._gates.clear() return mod, example @@ -379,7 +379,7 @@ def deco(cls): monkeypatch.setitem(sys.modules, "verl.experimental.agent_loop.agent_loop", mod) monkeypatch.delitem(sys.modules, "yeto.rl.adapters.verl.harbor_agent_loop", raising=False) loop_mod = importlib.import_module("yeto.rl.adapters.verl.harbor_agent_loop") - loop_mod._gate.cache_clear() + loop_mod._gates.clear() yield loop_mod, mod sys.modules.pop("yeto.rl.adapters.verl.harbor_agent_loop", None) # bound to the stub @@ -450,11 +450,14 @@ def test_verl_loop_aborts_on_misaligned_tokens_and_infra_errors(verl_loop): bad = [_turn([5, 6], [40]), _turn([5, 6, 77, 1], [41])] # history rewritten (compaction) _FakeTrial.result = _result({"reward": 1.0}) loop = _make_loop(loop_mod, _Recorder(bad), HarborHarness(lambda t, u: {}, trial_cls=_FakeTrial)) - out = _run(loop.run({}, **_ROW)) + + async def run_and_snapshot(): + return await loop.run({}, **_ROW), loop_mod._gate().metrics.snapshot() # the gate is per loop + + out, snap = _run(run_and_snapshot()) assert out.extra_fields[OUTCOME_KEY] == "aborted" and out.extra_fields["outcome_reason"].startswith( "token_alignment") assert (out.response_mask, out.reward_score) == ([0], 0.0) # placeholder: no loss tokens - snap = loop_mod._gate().metrics.snapshot() assert (snap["harness_outcome_ok_total"], snap["harness_outcome_aborted_total"]) == (0, 1) _FakeTrial.result = _result(None, "VerifierTimeoutError") diff --git a/tests/test_harbor_boat.py b/tests/test_harbor_boat.py index ee874995d..5d0a38a95 100644 --- a/tests/test_harbor_boat.py +++ b/tests/test_harbor_boat.py @@ -393,7 +393,7 @@ def test_boat_environment_on_harbor_interface(tmp_path: Path): environment_dir=tmp_path, environment_name="t1", session_id="t1__x__env", trial_paths=TrialPaths(trial_dir=tmp_path / "trial"), task_env_config=EnvironmentConfig(docker_image="img:1", memory_mb=2048, cpus=1), - client=fake) + client=fake, quota=BoatQuota(2, 12)) assert env.type() == "boat" and not env.capabilities.mounted async def go(): diff --git a/tests/test_verl_harbor_boat.py b/tests/test_verl_harbor_boat.py new file mode 100644 index 000000000..841d6ed94 --- /dev/null +++ b/tests/test_verl_harbor_boat.py @@ -0,0 +1,147 @@ +"""S25 N10.8c: verl's HarborAgentLoop on boat sandboxes, end to end on CPU. + +Real pieces: ``HarborAgentLoop.run``, the outer layer, ``build_trial_config`` +(provider ``boat`` -> ``import_path``), Harbor's ``Trial`` and the boat +provider. Fakes: verl's agent-loop module (stub, as in +``test_agentic_harbor``), the boat account (``FakeBoat``) and the token +recorder. Harbor's ``nop`` agent never calls the model, so the recorder +returns preset turns in place of the calls a policy agent would make. +""" + +from __future__ import annotations + +import asyncio +from pathlib import Path + +import pytest +import test_agentic_harbor as agentic +from test_agentic_harbor import _make_loop, _Recorder, _turn +from test_harbor_boat import FakeBoat + +from yeto.rl.agentic import OUTCOME_KEY, HarnessConfigError +from yeto.rl.agentic.loop_local import LoopLocal +from yeto.rl.harness.harbor import boat_client + +pytest.importorskip("harbor") + +verl_loop = agentic.verl_loop # the verl agent-loop stub fixture + +TASK_TOML = """\ +schema_version = "1.1" + +[task] +name = "yeto/boat-smoke" + +[verifier] +timeout_sec = 60.0 + +[agent] +timeout_sec = 60.0 + +[environment] +docker_image = "alpine:3.20" +cpus = 1 +memory_mb = 1024 +""" + + +@pytest.fixture +def task_root(tmp_path: Path) -> Path: + task = tmp_path / "tasks" / "boat-smoke" + for sub in ("environment", "tests", "solution"): + (task / sub).mkdir(parents=True) + (task / "task.toml").write_text(TASK_TOML) + (task / "instruction.md").write_text("Do nothing.\n") + (task / "tests" / "test.sh").write_text("#!/bin/bash\necho 1 > /logs/verifier/reward.txt\n") + (task / "solution" / "solve.sh").write_text("#!/bin/bash\ntrue\n") + return tmp_path + + +def _boat_loop(loop_mod, root: Path, fake: FakeBoat, turns, monkeypatch): + from yeto.rl.harness.harbor import HarborHarness, HarborSettings, boat_env, build_trial_config + + # Harbor writes the trial config to JSON, so the fake account cannot travel as a + # provider kwarg; it replaces the process client instead. + monkeypatch.setattr(boat_env, "process_client", lambda: fake) + settings = HarborSettings(tasks_dir=root / "tasks", trials_dir=root / "trials") + harness = HarborHarness(lambda task, url: build_trial_config(task, url, settings)) + return _make_loop(loop_mod, _Recorder(turns), harness) + + +def _row(provider: str | None = "boat") -> dict: + task = {"task_id": "boat-smoke", "agent": "nop", "trajectory_id": "traj-boat", "wall_clock_s": 120.0} + if provider is not None: + task["provider"] = provider + return {"extra_info": {"agent_task": task}} + + +@pytest.fixture +def fresh_quota(monkeypatch): + monkeypatch.setattr(boat_client, "_process_quota", LoopLocal(boat_client.BoatQuota.from_env)) + + +def test_verl_loop_runs_a_harbor_trial_on_a_boat_sandbox(verl_loop, task_root, fresh_quota, monkeypatch): + loop_mod, _ = verl_loop + fake = FakeBoat(states=["ready"]) + turns = [_turn([5, 6], [40, 41], [-0.1, -0.2]), _turn([5, 6, 40, 41, 90], [42, 2], [-0.3, -0.4])] + loop = _boat_loop(loop_mod, task_root, fake, turns, monkeypatch) + + out = asyncio.run(loop.run({"temperature": 1.0}, **_row())) + + assert fake.created == 1 and fake.deleted == ["bx_1"] and not fake.live # one VM, deleted + assert any("docker run -d --name main" in c and "alpine:3.20" in c for c in fake.commands) + assert fake.machine_type == "small" + assert out.extra_fields[OUTCOME_KEY] == "ok" and out.reward_score == 1.0 + assert out.prompt_ids == [5, 6] and out.response_ids == [40, 41, 90, 42, 2] + assert out.response_mask == [1, 1, 0, 1, 1] + assert out.response_logprobs == [-0.1, -0.2, 0.0, -0.3, -0.4] + + +def test_verl_loop_aborts_when_the_boat_container_fails(verl_loop, task_root, fresh_quota, monkeypatch): + loop_mod, _ = verl_loop + fake = FakeBoat(states=["ready"], fail_on="docker pull") + loop = _boat_loop(loop_mod, task_root, fake, [_turn([5], [6])], monkeypatch) + + async def run_and_snapshot(): + return await loop.run({}, **_row()), loop_mod._gate().metrics.snapshot() + + out, snap = asyncio.run(run_and_snapshot()) + + assert fake.deleted == ["bx_1"] and not fake.live # the failed start still deletes the VM + assert out.extra_fields[OUTCOME_KEY] == "aborted" and out.reward_score == 0.0 + assert out.response_mask == [0] # placeholder: no policy-gradient tokens + assert (snap["harness_outcome_aborted_total"], snap["harness_in_flight"]) == (1, 0) + + +def test_provider_defaults_to_harbor_env_type_and_boat_cap_is_checked(verl_loop, monkeypatch): + loop_mod, _ = verl_loop + monkeypatch.setenv("HARBOR_ENV_TYPE", "Boat") + assert loop_mod.task_from_row(_row(provider=None)).provider == "boat" + assert loop_mod.task_from_row(_row(provider="modal")).provider == "modal" # the row wins + + monkeypatch.setenv("YETO_HARBOR_MAX_IN_FLIGHT", "200") + with pytest.raises(HarnessConfigError, match="active-sandbox cap"): + loop_mod.task_from_row(_row()) + loop_mod.task_from_row(_row(provider="modal")) # the boat cap does not apply to other providers + + +def test_gate_and_quota_are_per_event_loop(verl_loop, fresh_quota): + """Two loops in one process (asyncio.run twice) must not share asyncio primitives.""" + loop_mod, _ = verl_loop + + async def hold_and_contend(): + gate, quota = loop_mod._gate(), boat_client.process_quota() + assert gate is loop_mod._gate() and quota is boat_client.process_quota() # one per loop + # Make both primitives wait, which binds them to this loop. + quota_holders = [asyncio.create_task(quota.acquire()) for _ in range(quota.max_active + 1)] + await asyncio.sleep(0) + for task in quota_holders: + task.cancel() + await asyncio.gather(*quota_holders, return_exceptions=True) + return gate, quota + + first = asyncio.run(hold_and_contend()) + second = asyncio.run(hold_and_contend()) # with process singletons: "bound to a different event loop" + assert first[0] is not second[0] and first[1] is not second[1] + with pytest.raises(RuntimeError): # outside a running loop there is no object to hand out + boat_client.process_quota() diff --git a/yeto/rl/adapters/miles/harbor_agent.py b/yeto/rl/adapters/miles/harbor_agent.py index 0c0caf055..4ebd2ce11 100644 --- a/yeto/rl/adapters/miles/harbor_agent.py +++ b/yeto/rl/adapters/miles/harbor_agent.py @@ -30,7 +30,7 @@ import importlib import os -from functools import lru_cache, partial +from functools import partial from typing import Any from yeto.rl.agentic import ( @@ -41,6 +41,7 @@ Outcome, run_agent_task, ) +from yeto.rl.agentic.loop_local import LoopLocal from yeto.rl.agentic.trust import check_agent_env from yeto.rl.harness.harbor import CODEX_VERSION, HarborHarness, HarborSettings, build_trial_config from yeto.rl.harness.harbor.trial import CUSTOM_PROVIDERS, TrialConfigBuilder @@ -57,10 +58,17 @@ class UntrustedAgentOutcome(RuntimeError): """A sample reached the reward hook without a yeto outcome (the agent function raised).""" -@lru_cache(maxsize=1) +_gates: LoopLocal[AdmissionGate] = LoopLocal( + lambda: AdmissionGate(int(os.environ.get(MAX_IN_FLIGHT_ENV, _DEFAULT_MAX_IN_FLIGHT)))) + + def gate() -> AdmissionGate: - """One gate per rollout-worker process (one sandbox per admitted trajectory).""" - return AdmissionGate(int(os.environ.get(MAX_IN_FLIGHT_ENV, _DEFAULT_MAX_IN_FLIGHT))) + """The gate of the rollout loop (one sandbox per admitted trajectory). + + Miles runs its rollout on one event loop, so this is one gate per worker + process; it is kept per loop because it holds an ``asyncio.Semaphore``. + """ + return _gates.get() def _example() -> Any: diff --git a/yeto/rl/adapters/verl/harbor_agent_loop.py b/yeto/rl/adapters/verl/harbor_agent_loop.py index f4fb79042..ec110a462 100644 --- a/yeto/rl/adapters/verl/harbor_agent_loop.py +++ b/yeto/rl/adapters/verl/harbor_agent_loop.py @@ -21,13 +21,32 @@ It adds no policy-gradient tokens, but it still counts in its GRPO group's reward mean; a group-level drop on verl is not implemented (not verified). -Configuration (worker environment): - YETO_AGENTIC_RECORDER ``module:callable`` returning a TokenRecorder (required) +Configuration (worker environment; pass it to the Ray workers, e.g. through +``ray_kwargs.ray_init.runtime_env.env_vars``): + YETO_AGENTIC_RECORDER ``module:callable`` returning a TokenRecorder (required). + verl builds the agent loop once per sample, so the + callable is called per sample; return a shared object. HARBOR_TASKS_DIR Harbor task directories (required) HARBOR_TRIALS_DIR trial output (default /tmp/harbor_trials) - HARBOR_ENV_KWARGS JSON provider kwargs (default {}) + HARBOR_ENV_TYPE sandbox provider of rows without ``agent_task.provider`` + (``boat``, ``modal``, ``docker``, ...; same name as Miles) + HARBOR_ENV_KWARGS JSON provider kwargs (default {}); boat: ``machine_type`` + (default: smallest that fits the task), ``ttl_s`` HARBOR_TIMEOUT_MULTIPLIER Harbor timeout multiplier (default 1.0) YETO_HARBOR_MAX_IN_FLIGHT sandboxes per worker process (default 64) + YETO_BOAT_MAX_ACTIVE, YETO_BOAT_STARTS_PER_MINUTE + boat caps per worker process (``boat_client``) + +Concurrency: verl runs ``rollout.agent.num_workers`` worker processes, so the +account holds up to ``num_workers * YETO_HARBOR_MAX_IN_FLIGHT`` sandboxes. For +boat, a worker refuses a ``YETO_HARBOR_MAX_IN_FLIGHT`` above +``YETO_BOAT_MAX_ACTIVE``; keeping ``num_workers * YETO_BOAT_MAX_ACTIVE`` within +the plan (box_20: 100 active, 12 starts per minute) is the launcher's job. + +The admission gate and the boat quota hold asyncio primitives and are kept +per event loop (:class:`~yeto.rl.agentic.loop_local.LoopLocal`). verl runs +``AgentLoopWorker`` as a Ray async actor (one loop per process), so in a +worker they are per process. """ from __future__ import annotations @@ -35,7 +54,7 @@ import importlib import os import uuid -from functools import lru_cache, partial +from functools import partial from typing import Any from verl.experimental.agent_loop.agent_loop import AgentLoopBase, AgentLoopMetrics, AgentLoopOutput, register @@ -54,15 +73,27 @@ join_turns, run_agent_task, ) +from yeto.rl.agentic.loop_local import LoopLocal from yeto.rl.harness.harbor import HarborHarness, HarborSettings, build_trial_config +from yeto.rl.harness.harbor.boat_client import check_in_flight_cap AGENT_LOOP_NAME = "yeto_harbor" RECORDER_ENV = "YETO_AGENTIC_RECORDER" +PROVIDER_ENV = "HARBOR_ENV_TYPE" +MAX_IN_FLIGHT_ENV = "YETO_HARBOR_MAX_IN_FLIGHT" +_DEFAULT_MAX_IN_FLIGHT = 64 + + +def _max_in_flight() -> int: + return int(os.environ.get(MAX_IN_FLIGHT_ENV, _DEFAULT_MAX_IN_FLIGHT)) + + +_gates: LoopLocal[AdmissionGate] = LoopLocal(lambda: AdmissionGate(_max_in_flight())) -@lru_cache(maxsize=1) def _gate() -> AdmissionGate: - return AdmissionGate(int(os.environ.get("YETO_HARBOR_MAX_IN_FLIGHT", 64))) + """The admission gate of this worker's event loop (one sandbox per admitted sample).""" + return _gates.get() def _recorder_from_env() -> TokenRecorder: @@ -77,6 +108,23 @@ def _harness_from_env() -> AgentHarness: return HarborHarness(partial(build_trial_config, settings=HarborSettings.from_env())) +def task_from_row(kwargs: dict[str, Any]) -> AgentTask: + """The sample's :class:`AgentTask`; ``HARBOR_ENV_TYPE`` fills a missing provider. + + A boat task is refused (configuration error, fails the batch) when this + worker's gate admits more sandboxes than its boat quota holds. + """ + row = agent_task_row(kwargs) + provider = os.environ.get(PROVIDER_ENV, "").strip().lower() + try: + task = AgentTask.from_mapping(row, defaults={"provider": provider or None}) + if task.provider == "boat": + check_in_flight_cap(_max_in_flight()) + except ValueError as exc: + raise HarnessConfigError(str(exc)) from exc + return task + + def agent_task_row(kwargs: dict[str, Any]) -> dict[str, Any]: """The ``agent_task`` mapping of a verl dataset row (``extra_info.agent_task``).""" extra = kwargs.get("extra_info") or {} @@ -102,7 +150,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: self._harness = self.harness_factory() async def run(self, sampling_params: dict[str, Any], **kwargs: Any) -> AgentLoopOutput: - task = AgentTask.from_mapping(agent_task_row(kwargs)) + task = task_from_row(kwargs) trajectory_id = task.trajectory_id or uuid.uuid4().hex endpoint = await self._recorder.open(trajectory_id, sampling_params) try: diff --git a/yeto/rl/agentic/loop_local.py b/yeto/rl/agentic/loop_local.py new file mode 100644 index 000000000..a2d1d85ed --- /dev/null +++ b/yeto/rl/agentic/loop_local.py @@ -0,0 +1,47 @@ +"""One object per running asyncio event loop (S25 N10.8c). + +The agentic path keeps process-wide guards (the admission gate, the boat +quota) that hold asyncio primitives. An ``asyncio.Semaphore`` or ``Lock`` +binds to the first loop that waits on it; used later from another loop it +raises ``RuntimeError: ... is bound to a different event loop``. + +verl runs each ``AgentLoopWorker`` as a Ray async actor, and Miles runs its +rollout on one background loop, so in a training worker there is one loop +per process and :class:`LoopLocal` gives exactly one object per process. +A process that starts several loops (``asyncio.run`` per call in scripts and +tests, verl's ``auto_await`` outside a running loop) gets one object per +loop instead of a crash. Limits held by such an object are then per loop. +""" + +from __future__ import annotations + +import asyncio +import threading +import weakref +from collections.abc import Callable +from typing import Generic, TypeVar + +T = TypeVar("T") + + +class LoopLocal(Generic[T]): + """Lazily builds ``factory()`` once for each running event loop.""" + + def __init__(self, factory: Callable[[], T]) -> None: + self._factory = factory + self._by_loop: weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, T] = weakref.WeakKeyDictionary() + self._lock = threading.Lock() # loops of different threads may ask at once + + def get(self) -> T: + """The object of the running loop; raises ``RuntimeError`` outside a running loop.""" + loop = asyncio.get_running_loop() + with self._lock: + value = self._by_loop.get(loop) + if value is None: + value = self._by_loop[loop] = self._factory() + return value + + def clear(self) -> None: + """Forget every object (tests, or after the configuration changed).""" + with self._lock: + self._by_loop.clear() diff --git a/yeto/rl/harness/harbor/README.md b/yeto/rl/harness/harbor/README.md index 2e4d3eed6..8295646a2 100644 --- a/yeto/rl/harness/harbor/README.md +++ b/yeto/rl/harness/harbor/README.md @@ -103,6 +103,22 @@ sample is a mask-0 placeholder with reward 0; it still counts in its GRPO group mean. The verl launcher and island entry refuse the Miles agentic flags (`--custom-agent-function-path`, `--use-session-server`, `--tito-model`, ...). +boat on verl: `HARBOR_ENV_TYPE=boat` (or `provider: boat` in the row; the row +wins), `HARBOR_ENV_KWARGS='{"machine_type": "small"}'` (optional). Each of +`rollout.agent.num_workers` worker processes admits up to +`YETO_HARBOR_MAX_IN_FLIGHT` sandboxes and refuses a value above +`YETO_BOAT_MAX_ACTIVE`; keep `num_workers * YETO_BOAT_MAX_ACTIVE` within the +plan. The admission gate and the boat quota are kept per event loop +(`yeto/rl/agentic/loop_local.py`); verl's `AgentLoopWorker` is a Ray async +actor, so that is one per worker process. verl builds the agent loop once per +sample: `YETO_AGENTIC_RECORDER` should return a shared recorder. + +Verified 2026-10-19 (T-N108C): `HarborAgentLoop.run()` (verl agent-loop +module stubbed, configuration from the env variables above) on boat, fix-git +and regex-log: oracle reward 1 on both, nop reward 0 on both, outcome ok; +41 billable s, `boat list` empty afterwards. The recorder returned one +synthetic turn (oracle/nop call no model). A verl training run: not verified. + ## boat sandbox provider (N10.8) `provider: boat` maps to `EnvironmentConfig.import_path = diff --git a/yeto/rl/harness/harbor/boat_client.py b/yeto/rl/harness/harbor/boat_client.py index e88784357..4939617cb 100644 --- a/yeto/rl/harness/harbor/boat_client.py +++ b/yeto/rl/harness/harbor/boat_client.py @@ -11,7 +11,8 @@ Quota (box_20 plan, read from ``GET /limits`` on 2026-10-19): 100 active sandboxes, 12 starts per minute, 60 per hour, 200 per day. :class:`BoatQuota` -holds the first two per process; :func:`check_in_flight_cap` refuses an +holds the first two per process (per event loop, see :func:`process_quota`); +:func:`check_in_flight_cap` refuses an admission gate larger than the active-sandbox cap. The hourly and daily caps are left to the server (a refused create raises :class:`BoatQuotaError`). """ @@ -30,6 +31,8 @@ from pathlib import Path from typing import Protocol +from yeto.rl.agentic.loop_local import LoopLocal + API_HOST = "https://boat.dev/api/v1" API_KEY_ENV = "BOAT_API_KEY" CLI_CONFIG = Path("~/.config/ascii/boat/config.json") @@ -226,7 +229,6 @@ def release(self) -> None: self._active.release() -_process_quota: BoatQuota | None = None _process_client: SdkBoatClient | None = None @@ -239,8 +241,13 @@ def process_client() -> SdkBoatClient: def process_quota() -> BoatQuota: - """The quota shared by every boat sandbox of this process (created on first use).""" - global _process_quota - if _process_quota is None: - _process_quota = BoatQuota.from_env() - return _process_quota + """The quota shared by every boat sandbox of this process's event loop. + + It holds asyncio primitives, so it is kept per running loop + (:class:`~yeto.rl.agentic.loop_local.LoopLocal`); a training worker runs + one loop, so this is one quota per process. Call it inside a running loop. + """ + return _process_quota.get() + + +_process_quota: LoopLocal[BoatQuota] = LoopLocal(BoatQuota.from_env) diff --git a/yeto/rl/harness/harbor/boat_env.py b/yeto/rl/harness/harbor/boat_env.py index de8cb5292..486137dc9 100644 --- a/yeto/rl/harness/harbor/boat_env.py +++ b/yeto/rl/harness/harbor/boat_env.py @@ -24,7 +24,7 @@ from harbor.environments.base import BaseEnvironment, ExecResult from harbor.environments.capabilities import EnvironmentCapabilities -from .boat_client import process_client, process_quota +from .boat_client import BoatQuota, process_client, process_quota from .boat_sandbox import BoatSandbox, ContainerSpec, machine_type_for BOAT_IMPORT_PATH = "yeto.rl.harness.harbor.boat_env:BoatEnvironment" @@ -35,18 +35,21 @@ class BoatEnvironment(BaseEnvironment): Extra kwargs (``EnvironmentConfig.kwargs``): ``machine_type`` (default: the smallest that fits the task), ``ttl_s`` (VM auto-stop, a backstop for a - lost delete), ``client`` (a :class:`~.boat_client.BoatClient`, for tests). + lost delete), ``client`` (a :class:`~.boat_client.BoatClient`) and + ``quota`` (a :class:`~.boat_client.BoatQuota`), both for tests. Without + ``quota`` the environment takes the process quota, so it must be built + inside the running loop (Harbor builds it in ``Trial.create``). """ def __init__(self, *args: Any, machine_type: str | None = None, ttl_s: int | None = None, - client: Any = None, **kwargs: Any) -> None: + client: Any = None, quota: BoatQuota | None = None, **kwargs: Any) -> None: super().__init__(*args, **kwargs) sandbox_kwargs: dict[str, Any] = {} if ttl_s is not None: sandbox_kwargs["ttl_s"] = ttl_s self._sandbox = BoatSandbox( client if client is not None else process_client(), - process_quota(), + quota if quota is not None else process_quota(), name=self.session_id, machine_type=machine_type or machine_type_for(self._effective_memory_mb, self._effective_cpus), **sandbox_kwargs,