feat: add benchmarks and improve inference module with Docker support

README.md

@@ -216,6 +216,20 @@ Typical downstream evaluation reads:
 - how many deadlines were missed
 - how much important work remained unfinished
 
+## Benchmarks
+
+Verified self-contained inference run using:
+
+1. `qwen/qwen3.5-9b`
+
+Results:
+
+| Preset   | Success | Steps | Score   |
+| -------- | ------- | ----- | ------- |
+| `easy`   | `true`  | `11`  | `0.952` |
+| `medium` | `true`  | `20`  | `0.945` |
+| `hard`   | `true`  | `45`  | `0.652` |
+
 ## Local Development
 
 Validate the environment:

inference.py

@@ -1,8 +1,12 @@
 from __future__ import annotations
 
+import asyncio
 import json
 import os
+import subprocess
+from contextlib import asynccontextmanager
 from dataclasses import dataclass
+from pathlib import Path
 
 from openai import OpenAI
 from workflow_arena import WorkflowArenaAction, WorkflowArenaEnv
@@ -12,7 +16,6 @@ from workflow_arena.models import (
     WorkflowArenaObservation,
     WorkflowTaskView,
 )
-from workflow_arena.presets import get_preset_config
 
 BENCHMARK = "WorkflowArena"
 PRESETS = [
@@ -20,7 +23,9 @@ PRESETS = [
     DifficultyPreset.MEDIUM,
     DifficultyPreset.HARD,
 ]
-DEFAULT_BASE_URL = os.getenv("WORKFLOW_ARENA_BASE_URL", "http://localhost:8000")
+PROJECT_DIR = Path(__file__).resolve().parent
+IMAGE_NAME = "workflow-arena-inference:latest"
+DOCKERFILE_PATH = PROJECT_DIR / "server" / "Dockerfile"
 TEMPERATURE = 0.0
 MAX_STEPS = 256
 
@@ -190,7 +195,54 @@ class EpisodeResult:
     rewards: list[float]
 
 
-def run_episode(
+def ensure_local_image() -> None:
+    try:
+        inspect_result = subprocess.run(
+            ["docker", "image", "inspect", IMAGE_NAME],
+            cwd=PROJECT_DIR,
+            stdout=subprocess.DEVNULL,
+            stderr=subprocess.DEVNULL,
+            check=False,
+        )
+    except OSError as exc:
+        raise RuntimeError(f"Failed to execute docker: {exc}") from exc
+
+    if inspect_result.returncode == 0:
+        return
+
+    try:
+        build_result = subprocess.run(
+            ["docker", "build", "-t", IMAGE_NAME, "-f", str(DOCKERFILE_PATH), "."],
+            cwd=PROJECT_DIR,
+            capture_output=True,
+            text=True,
+            check=False,
+        )
+    except OSError as exc:
+        raise RuntimeError(f"Failed to execute docker build: {exc}") from exc
+
+    if build_result.returncode != 0:
+        raise RuntimeError(
+            "Failed to build Docker image for inference.\n"
+            f"Command: docker build -t {IMAGE_NAME} -f {DOCKERFILE_PATH} .\n"
+            f"Exit code: {build_result.returncode}\n"
+            f"Stdout: {build_result.stdout}\n"
+            f"Stderr: {build_result.stderr}"
+        )
+
+
+@asynccontextmanager
+async def managed_env():
+    ensure_local_image()
+    env = await WorkflowArenaEnv.from_docker_image(IMAGE_NAME)
+    try:
+        yield env
+    finally:
+        await env.close()
+
+
+async def run_episode(
+    env,
     client: OpenAI | None,
     model_name: str,
     preset: DifficultyPreset,
@@ -203,56 +255,53 @@ def run_episode(
 
     log_start(task=preset.value, env=BENCHMARK, model=model_name)
 
-    with WorkflowArenaEnv(base_url=DEFAULT_BASE_URL).sync() as env:
-        preset_config = get_preset_config(preset)
-        result = env.reset(
-            seed=seed,
-            preset=preset.value,
-            worker_count=preset_config.worker_count,
-        )
-        observation = result.observation
+    result = await env.reset(
+        seed=seed,
+        preset=preset.value,
+    )
+    observation = result.observation
 
-        while not observation.done and steps_taken < MAX_STEPS:
-            try:
-                if client is None:
-                    action = heuristic_action(observation)
-                else:
-                    action = get_model_action(client, model_name, observation)
-            except (
-                Exception
-            ):  # pragma: no cover - network/model failures are expected sometimes
+    while not observation.done and steps_taken < MAX_STEPS:
+        try:
+            if client is None:
                 action = heuristic_action(observation)
+            else:
+                action = get_model_action(client, model_name, observation)
+        except (
+            Exception
+        ):  # pragma: no cover - network/model failures are expected sometimes
+            action = heuristic_action(observation)
+
+        try:
+            result = await env.step(action)
+        except (
+            Exception
+        ):  # pragma: no cover - preserve log format and continue safely
+            action = heuristic_action(observation)
+            result = await env.step(action)
 
-            try:
-                result = env.step(action)
-            except (
-                Exception
-            ):  # pragma: no cover - preserve log format and continue safely
-                action = heuristic_action(observation)
-                result = env.step(action)
-
-            observation = result.observation
-            reward = float(result.reward or 0.0)
-            rewards.append(reward)
-            steps_taken += 1
-            log_step(
-                step=steps_taken,
-                action=action_to_log_string(action),
-                reward=reward,
-                done=bool(result.done),
-                error=observation.validation_error,
-            )
+        observation = result.observation
+        reward = float(result.reward or 0.0)
+        rewards.append(reward)
+        steps_taken += 1
+        log_step(
+            step=steps_taken,
+            action=action_to_log_string(action),
+            reward=reward,
+            done=bool(result.done),
+            error=observation.validation_error,
+        )
 
-        success = is_success(observation)
-        score = compute_score(observation) if observation.done else 0.0
-        log_end(success=success, steps=steps_taken, score=score, rewards=rewards)
+    success = is_success(observation)
+    score = compute_score(observation) if observation.done else 0.0
+    log_end(success=success, steps=steps_taken, score=score, rewards=rewards)
 
     return EpisodeResult(
         success=success, steps=steps_taken, score=score, rewards=rewards
     )
 
 
-def main() -> None:
+async def main() -> None:
     api_base_url = os.environ["API_BASE_URL"]
     model_name = os.environ["MODEL_NAME"]
     api_key = os.getenv("HF_TOKEN") or os.getenv("OPENAI_API_KEY")
@@ -261,14 +310,16 @@ def main() -> None:
 
     client = OpenAI(base_url=api_base_url, api_key=api_key)
 
-    for index, preset in enumerate(PRESETS):
-        run_episode(
-            client=client,
-            model_name=model_name,
-            preset=preset,
-            seed=100 + index,
-        )
+    async with managed_env() as env:
+        for index, preset in enumerate(PRESETS):
+            await run_episode(
+                env=env,
+                client=client,
+                model_name=model_name,
+                preset=preset,
+                seed=100 + index,
+            )
 
 
 if __name__ == "__main__":
-    main()
+    asyncio.run(main())