diff --git a/docs/source/_toctree.yml b/docs/source/_toctree.yml index 564426287..2763b9fd2 100644 --- a/docs/source/_toctree.yml +++ b/docs/source/_toctree.yml @@ -27,6 +27,8 @@ title: Simulation vs Production - local: guides/mcp-environment-lifecycle title: MCP Environment Lifecycle + - local: guides/task-api + title: The Task API - local: guides/connecting title: Connecting to Servers - local: guides/runtime-providers diff --git a/docs/source/guides/task-api.md b/docs/source/guides/task-api.md new file mode 100644 index 000000000..540acfd29 --- /dev/null +++ b/docs/source/guides/task-api.md @@ -0,0 +1,273 @@ +# The Task API + +This guide explains the Task API: the optional discovery layer that lets a dataset-backed environment publish the *set of tasks* it can run, so a trainer or evaluator can enumerate them and choose which one each episode uses. + +It exists to answer a common question: OpenEnv's step loop gives you one episode at a time, so how does a training script find out that an environment holds 7,595 test problems, and how does it ask for problem number 12? + +## The Short Answer + +The Task API has two halves: + +- **Server side**: your environment optionally implements five methods — `list_splits()`, `list_tasks()`, `num_tasks()`, `get_task()`, `get_task_range()`. These are described by the `TaskProvider` protocol. +- **HTTP side**: when those methods exist, the environment server automatically exposes them as routes under `/{env_name}/…`. You do not register anything. + +Task discovery is **metadata only**. It never starts an episode. Selecting a task happens the usual way, through `reset()`: + +```python +env.list_splits() # ["train", "test"] +env.num_tasks("test") # 7595 +env.get_task("test", 12) # {"id": "test-12", "index": 12, "split": "test"} +env.reset(split="test", index=12) # <- this is what starts the episode +``` + +The API is shaped to be compatible with ORS/OpenReward task and split conventions, which is why environments imported with `openenv import` get it for free. + +## When You Need It + +Implement the Task API when your environment is **dataset-backed** — when an episode means "run one row of a dataset" and a trainer needs to iterate, shard, or shuffle over those rows. + +You do **not** need it for environments where every episode is generated or self-contained: Echo, Snake, a game, a REPL. Those environments simply omit the methods, and the routes report `501 Not Implemented`. + +## The Protocol + +`TaskProvider` is a `typing.Protocol`, not a base class. You do not inherit from it — you just define the methods on your `Environment` subclass and the server discovers them by name. + +```python +from typing import Any, Optional + +class TaskProvider(Protocol): + def list_splits(self) -> list[Any]: ... + def list_tasks(self, split: str) -> list[Any]: ... + def num_tasks(self, split: str) -> int: ... + def get_task(self, split: str, index: int) -> Any: ... + def get_task_range( + self, + split: str, + start: Optional[int] = None, + stop: Optional[int] = None, + ) -> list[Any]: ... +``` + +Each method may be sync or `async`; the server awaits the result when it is awaitable, so you can back task discovery with async dataset clients. + +A "task spec" is deliberately untyped (`Any`). Return whatever your environment needs — a dict, a Pydantic model, a dataclass. The server converts it to JSON, handling Pydantic models, dataclasses, and plain objects. In practice a positional stub such as `{"id": "test-12", "index": 12, "split": "test"}` is enough, because `reset()` is what actually loads the row. + +### Two rules that are easy to miss + +1. **Task methods must be side-effect-free.** They are discovery, not control. They must not mutate episode state, consume a stream, or advance a cursor. +2. **They must work on a freshly constructed environment.** Each task route builds a short-lived environment instance, calls the one method, and closes it again. Nothing you set up in `reset()` is available, so read configuration in `__init__` (or lazily inside the method) rather than relying on episode state. + +## HTTP Routes + +The routes are registered by `HTTPEnvServer` and appear in the server's OpenAPI schema under the `Task API` tag. All of them are namespaced by environment name except `/list_environments`. + +| Method | Route | Body | Response | +|--------|-------|------|----------| +| `GET` | `/list_environments` | — | `["latex_ocr_env"]` | +| `GET` | `/{env_name}/splits` | — | `[{"name": "train", "type": "train"}, …]` | +| `POST` | `/{env_name}/tasks` | `{"split": "test"}` | `{"tasks": [...], "env_name": "latex_ocr_env"}` | +| `POST` | `/{env_name}/num_tasks` | `{"split": "test"}` | `{"num_tasks": 7595}` | +| `POST` | `/{env_name}/task` | `{"split": "test", "index": 12}` | `{"task": {...}}` | +| `POST` | `/{env_name}/task_range` | `{"split": "test", "start": 0, "stop": 32}` | `{"tasks": [...]}` | + +The request bodies are `ListTasksRequest`, `NumTasksRequest`, `GetTaskRequest`, and `GetTaskRangeRequest`. `start` and `stop` follow Python slice semantics: `start` is inclusive, `stop` is exclusive, and both may be omitted. + +These are HTTP-only. There are no WebSocket message types for task discovery — the `/ws` session protocol stays focused on `reset`, `step`, `state`, and `close`. + +### `env_name` + +`{env_name}` is the name passed to `create_app(..., env_name="latex_ocr_env")`. If you do not pass one, it defaults to the environment factory's class name (for example `LatexOCREnvironment`). Matching is case-insensitive, and any other name returns `404`. Pass `env_name` explicitly for any environment that implements the Task API — the generated URL is part of your public surface, and a class rename should not break a trainer. + +```python +from openenv.core.env_server import create_app + +app = create_app( + LatexOCREnvironment, + LatexOCRAction, + LatexOCRObservation, + env_name="latex_ocr_env", +) +``` + +### Split normalization + +`list_splits()` may return plain strings, dicts, or Pydantic models. The server normalizes each entry to `{"name": ..., "type": ...}`: + +- a dict or model is passed through as-is (after JSON conversion) +- a string becomes `{"name": s, "type": s}` when `s` is `train`, `validation`, or `test` +- any other string becomes `{"name": s, "type": "validation"}` + +So returning `["train", "holdout"]` yields `[{"name": "train", "type": "train"}, {"name": "holdout", "type": "validation"}]`. Return dicts yourself if you need to control the `type` for a custom split name. + +### Error semantics + +| Condition | Status | +|-----------|--------| +| `{env_name}` does not match the server's environment | `404 Not Found` | +| Environment does not define the method, or it raises `NotImplementedError` | `501 Not Implemented` | +| Method raises `IndexError` (index out of range) | `400 Bad Request` | + +Raise `IndexError` from `get_task()` for an out-of-range index so callers get a `400` rather than a `500`. Reserve `NotImplementedError` for capabilities the environment genuinely does not support — for example, random access on a streaming split. + +## Implementing It + +A dataset-backed environment reads its dataset in `__init__`, answers discovery from metadata, and loads the actual row in `reset()`: + +```python +from typing import Any, Optional +from uuid import uuid4 + +from openenv.core.env_server import Environment +from openenv.core.env_server.types import State + + +class LatexOCREnvironment(Environment): + SPLITS = ["train", "test"] + + def __init__(self, dataset_name: str = "unsloth/LaTeX_OCR") -> None: + super().__init__() + self.dataset_name = dataset_name + + # --- Task API: discovery only, no side effects --- + + def list_splits(self) -> list[str]: + return self.SPLITS + + def num_tasks(self, split: str) -> int: + return len(self._load_split(split)) + + def list_tasks(self, split: str) -> list[dict[str, Any]]: + return [ + {"id": f"{split}-{i}", "index": i} + for i in range(self.num_tasks(split)) + ] + + def get_task(self, split: str, index: int) -> dict[str, Any]: + n = self.num_tasks(split) + if index < 0 or index >= n: + raise IndexError(f"index {index} out of range for split {split} (n={n})") + return {"id": f"{split}-{index}", "index": index, "split": split} + + def get_task_range( + self, split: str, start: Optional[int] = None, stop: Optional[int] = None + ) -> list[dict[str, Any]]: + n = self.num_tasks(split) + start = 0 if start is None else start + stop = n if stop is None else min(stop, n) + return [ + {"id": f"{split}-{i}", "index": i, "split": split} + for i in range(start, stop) + ] + + # --- Episode: this is where a task is actually loaded --- + + def reset( + self, + split: str = "test", + index: Optional[int] = None, + seed: Optional[int] = None, + episode_id: Optional[str] = None, + **kwargs: Any, + ) -> LatexOCRObservation: + if split not in self.list_splits(): + raise ValueError(f"unknown split {split!r}; expected {self.list_splits()}") + self._state = State(episode_id=episode_id or str(uuid4()), step_count=0) + row = self._load_row(split, index, seed) + return LatexOCRObservation(image_base64=row["image"], done=False) +``` + +### Selecting a task in `reset()` + +`ResetRequest` allows extra fields, and both the HTTP `/reset` route and the WebSocket `reset` message forward unknown keys to your `reset()` — after filtering them against its signature. So naming the parameters `split` and `index` in `reset()` is all it takes for this to work: + +```python +result = env.reset(split="test", index=12) +``` + +The corresponding wire calls are `POST /reset` with `{"split": "test", "index": 12}`, or a WebSocket `{"type": "reset", "data": {"split": "test", "index": 12}}`. The server filters incoming keys against your `reset()` signature, so a key it does not declare is dropped silently unless the signature also has `**kwargs`. A misspelled parameter therefore shows up as "the environment ignored my selection" rather than an error — worth checking first when a task selection appears not to take effect. + +When `index` is omitted, pick a row from `seed` so episodes stay reproducible. + +### Streaming and very large splits + +`list_tasks()` and `get_task_range()` can be asked for more rows than you want to materialize. Two defenses are worth building in: + +- return positional stubs rather than real rows, and cap how many you generate (`list_tasks()` returning a bounded preview is fine — `num_tasks()` still reports the honest total) +- if a split is streamed sequentially and cannot be randomly accessed, raise `ValueError` or `NotImplementedError` from `reset(index=...)` and say which mode does support indexing + +Report the true count from `num_tasks()` even when `list_tasks()` is truncated. A trainer that shards work by `num_tasks()` needs the real denominator. + +## Consuming the Task API + +The core clients (`EnvClient`, `SyncEnvClient`) do not ship Task API methods, because task specs are environment-specific. Environment clients add thin HTTP helpers alongside the inherited Gym-style `reset()`/`step()`: + +```python +import requests +from urllib.parse import urljoin + +ENV_NAME = "latex_ocr_env" + + +class LatexOCREnv(EnvClient[LatexOCRAction, LatexOCRObservation, State]): + def list_splits(self) -> list[str]: + resp = requests.get(urljoin(self._http_base(), f"{ENV_NAME}/splits"), timeout=30) + resp.raise_for_status() + return [s["name"] for s in resp.json()] + + def num_tasks(self, split: str) -> int: + resp = requests.post( + urljoin(self._http_base(), f"{ENV_NAME}/num_tasks"), + json={"split": split}, + timeout=60, + ) + resp.raise_for_status() + return int(resp.json()["num_tasks"]) +``` + +Note that the Task API is served over HTTP while episodes typically run over the WebSocket `/ws` session. A client that was constructed from a `ws://` URL has to derive the HTTP base URL before calling these routes. + +From a training or evaluation loop, discovery and episode control then compose naturally: + +```python +with LatexOCREnv.from_docker_image("latex-ocr-env:latest") as env: + total = env.num_tasks("test") + for index in range(total): + result = env.reset(split="test", index=index) + prediction = model(result.observation.image_base64) + result = env.step(LatexOCRAction(latex=prediction)) + print(index, result.reward) +``` + +Or with `curl`, against a locally running server: + +```bash +curl http://localhost:8000/list_environments +curl http://localhost:8000/latex_ocr_env/splits +curl -X POST http://localhost:8000/latex_ocr_env/num_tasks \ + -H 'Content-Type: application/json' -d '{"split": "test"}' +curl -X POST http://localhost:8000/latex_ocr_env/task_range \ + -H 'Content-Type: application/json' -d '{"split": "test", "start": 0, "stop": 4}' +``` + +## Imported Environments + +`openenv import` generates wrappers that already implement the Task API by delegating to the source environment's own task and split methods, and whose `reset()` accepts `split` / `index` / `task_spec`. If you are bringing in an ORS/OpenReward or Prime Intellect Verifiers environment, you get task discovery without writing any of the above — see the [CLI reference](../reference/cli.md) for `openenv import`. + +## Checklist + +Before shipping a dataset-backed environment: + +1. All five methods defined, or a deliberate `NotImplementedError` for the ones you cannot support. +2. Every method works on a fresh instance, with no reliance on `reset()` having run. +3. No side effects — no cursor advanced, no stream consumed, no episode state mutated. +4. `env_name` passed explicitly to `create_app()`. +5. `get_task()` raises `IndexError` for out-of-range indices. +6. `num_tasks()` reports the honest total, even if `list_tasks()` truncates. +7. `reset()` accepts `split` and `index` with the same names the trainer will use, and falls back to `seed`-driven selection when `index` is omitted. + +## Related Reading + +- [Concepts](concepts.md) +- [RL Training](rl-integration.md) +- [Core API](../reference/core.md) +- [CLI reference](../reference/cli.md) diff --git a/docs/source/reference/core.md b/docs/source/reference/core.md index bf9de12cf..a8177587a 100644 --- a/docs/source/reference/core.md +++ b/docs/source/reference/core.md @@ -6,6 +6,8 @@ If you are trying to understand when OpenEnv exposes the training loop versus di For a high-level explanation of how MCP-backed environments move through `step()`, `step_async()`, and convenience tool helpers, see the [MCP environment lifecycle](../guides/mcp-environment-lifecycle.md) guide. +For dataset-backed environments that publish enumerable tasks and splits, see the [Task API](../guides/task-api.md) guide. + ## Server ### Environment server primitives @@ -16,6 +18,8 @@ For a high-level explanation of how MCP-backed environments move through `step() [[autodoc]] openenv.core.env_server.interfaces.Transform +[[autodoc]] openenv.core.env_server.interfaces.TaskProvider + [[autodoc]] openenv.core.env_server.interfaces.Environment ### Types @@ -50,6 +54,14 @@ For a high-level explanation of how MCP-backed environments move through `step() [[autodoc]] openenv.core.env_server.types.HealthResponse +[[autodoc]] openenv.core.env_server.types.ListTasksRequest + +[[autodoc]] openenv.core.env_server.types.NumTasksRequest + +[[autodoc]] openenv.core.env_server.types.GetTaskRequest + +[[autodoc]] openenv.core.env_server.types.GetTaskRangeRequest + [[autodoc]] openenv.core.env_server.types.WSResetMessage [[autodoc]] openenv.core.env_server.types.WSStepMessage diff --git a/src/openenv/core/env_server/interfaces.py b/src/openenv/core/env_server/interfaces.py index 09e960651..8081789d9 100644 --- a/src/openenv/core/env_server/interfaces.py +++ b/src/openenv/core/env_server/interfaces.py @@ -78,28 +78,91 @@ def decode( class TaskProvider(Protocol): - """Optional task discovery API for dataset-backed environments. + """ + Optional task discovery API for dataset-backed environments. + + An environment implements this protocol structurally — declare the methods on + an [`~openenv.core.env_server.interfaces.Environment`] subclass, without + inheriting from `TaskProvider`. When the methods are present, + [`~openenv.core.env_server.http_server.HTTPEnvServer`] exposes them as HTTP + routes under `/{env_name}/…`; when they are absent, those routes return + `501 Not Implemented`. Each method may be sync or async. Task provider methods are for metadata/discovery only and should be side-effect-free. They must be callable on a freshly constructed environment instance because HTTP compatibility routes may create a short-lived instance solely for task discovery. + + Selecting a task is not part of this protocol — pass the chosen split and + index to `reset()` instead. See the + [Task API guide](https://huggingface.co/docs/openenv/guides/task-api). + + Examples: + + ```python + env.list_splits() # ["train", "test"] + env.num_tasks("test") # 7595 + env.get_task("test", 12) # {"id": "test-12", "index": 12} + env.reset(split="test", index=12) + ``` """ def list_splits(self) -> list[Any]: - """Return task split descriptors supported by this environment.""" + """ + Return task split descriptors supported by this environment. + + Returns: + `list[Any]`: Split descriptors. Plain strings, dicts, and Pydantic + models are all accepted; the server normalizes each entry to + `{"name": ..., "type": ...}`. + """ ... def list_tasks(self, split: str) -> list[Any]: - """Return all task specs for a split.""" + """ + Return all task specs for a split. + + Args: + split (`str`): + Task split name. + + Returns: + `list[Any]`: Task specs for the split. Environments backed by very + large or streamed splits may return a bounded preview, but + `num_tasks` should still report the true total. + """ ... def num_tasks(self, split: str) -> int: - """Return the number of task specs in a split.""" + """ + Return the number of task specs in a split. + + Args: + split (`str`): + Task split name. + + Returns: + `int`: Number of task specs available in the split. + """ ... def get_task(self, split: str, index: int) -> Any: - """Return one task spec by split and index.""" + """ + Return one task spec by split and index. + + Args: + split (`str`): + Task split name. + index (`int`): + Task index within the split. + + Returns: + `Any`: The task spec at that position. + + Raises: + `IndexError`: If `index` is out of range for the split. The HTTP + route converts this to a `400 Bad Request`. + """ ... def get_task_range( @@ -108,7 +171,20 @@ def get_task_range( start: Optional[int] = None, stop: Optional[int] = None, ) -> list[Any]: - """Return task specs for Python slice-style range bounds.""" + """ + Return task specs for Python slice-style range bounds. + + Args: + split (`str`): + Task split name. + start (`int`, *optional*): + Inclusive start index. Defaults to the beginning of the split. + stop (`int`, *optional*): + Exclusive stop index. Defaults to the end of the split. + + Returns: + `list[Any]`: Task specs in `[start, stop)`. + """ ...