Evals and RL environments

import os
import time
import uuid
from concurrent.futures import ThreadPoolExecutor

from boat_sdk import ApiClient, Configuration, exec_command, wait_until_ready
from boat_sdk.api.boat_api import BoatApi
from boat_sdk.exceptions import ApiException
from boat_sdk.models import CreateSandboxRequest

STARTS_PER_MINUTE = 12      # $20 plan. See https://docs.boat.dev/limits
STARTS_PER_HOUR = 60        # $20 plan. The hour limit is the tighter one for long runs.
TEMPLATE = "task-env-v1"

config = Configuration(host="https://boat.dev/api/v1", access_token=os.environ["BOAT_API_KEY"])


def start(boat: BoatApi) -> str:
    request = CreateSandboxRequest(var_from=TEMPLATE, type="small", ttl_seconds=3600)
    key = str(uuid.uuid4())          # the same key on a retry never makes a second sandbox
    for attempt in range(6):
        try:
            return boat.create(idempotency_key=key, create_sandbox_request=request).sandbox.id
        except ApiException as exc:
            if exc.status not in (429, 503):
                raise
            time.sleep(min(60, 2 ** attempt * 5))
    raise RuntimeError("could not start a sandbox")


def episode(boat: BoatApi, n: int) -> tuple[int, int, str]:
    sandbox_id = start(boat)
    try:
        wait_until_ready(boat, sandbox_id)
        result = exec_command(boat, sandbox_id, f"python3 eval-tasks/run.py --episode {n}", timeout_seconds=600)
        out = result.actual_instance
        return n, out.exit_code, out.stdout
    finally:
        boat.delete_sandbox(sandbox_id=sandbox_id, x_ascii_confirm_delete=sandbox_id)


with ApiClient(config) as client:
    boat = BoatApi(client)
    with ThreadPoolExecutor(max_workers=50) as pool:
        futures = []
        for n in range(100):
            futures.append(pool.submit(episode, boat, n))
            # pace the starts under both the minute and the hour limit
            time.sleep(max(60 / STARTS_PER_MINUTE, 3600 / STARTS_PER_HOUR))
        for future in futures:
            n, code, stdout = future.result()
            print(n, code, stdout.strip()[-200:])