feat(training): train on a bundle with slime through its OpenAI adapter over the OpenEnv server

.github/workflows/slime.yml

@@ -0,0 +1,81 @@
+name: slime contract
+
+# agent_env.training.slime on slime itself (tst/slime): the helper's generate function plays rows the writer wrote
+# through slime's own OpenAI adapter, with slime's own arguments, in front of a stand-in for SGLang's /generate, and
+# slime's dynamic-sampling filter decides each group; on slime 540ea91 and on its main branch. slime isn't on PyPI: its
+# checkout is installed with its requirements over the dev environment, torch the CPU build, and SGLang's and
+# Megatron-LM's checkouts go on PYTHONPATH, at the SGLang tag and the Megatron-LM commit slime's Dockerfile builds on,
+# with the transformers that SGLang tag pins. It runs on pull requests that touch the helper, the OpenEnv server it plays
+# or the tier, weekly, and on demand; it isn't a required check.
+
+on:
+  pull_request:
+    paths:
+      - "src/agent_env/training/**"
+      - "src/agent_env/openenv/**"
+      - "tst/slime/**"
+      - "tst/unit/training/**"
+      - "pyproject.toml"
+      - "uv.lock"
+      - ".github/workflows/slime.yml"
+  schedule:
+    - cron: "41 8 * * 1"
+  workflow_dispatch:
+
+permissions:
+  contents: read
+
+concurrency:
+  group: slime-${{ github.event_name == 'pull_request' && github.ref || github.sha }}
+  cancel-in-progress: ${{ github.event_name == 'pull_request' }}
+
+env:
+  PIP_DISABLE_PIP_VERSION_CHECK: "1"
+
+jobs:
+  generate:
+    name: generate (slime ${{ matrix.slime }})
+    runs-on: ubuntu-latest
+    timeout-minutes: 40
+    strategy:
+      fail-fast: false
+      matrix:
+        slime: ["540ea919e45b8b929198abce008a103be9cfc779", "main"]
+    steps:
+      - uses: actions/checkout@11d5960a326750d5838078e36cf38b85af677262 # v4.4.0
+      - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
+        with:
+          python-version: "3.12"
+      - name: Install from uv.lock with the slime extra
+        run: |
+          python -m pip install --quiet -r .github/ci-requirements.txt
+          uv sync --locked --extra dev --extra slime --python "$(command -v python)"
+      - name: Check out slime, SGLang and Megatron-LM
+        env:
+          SLIME_REF: ${{ matrix.slime }}
+        run: |
+          checkout() { git init --quiet "$RUNNER_TEMP/$1" && git -C "$RUNNER_TEMP/$1" fetch --quiet --depth 1 "$2" "$3" \
+            && git -C "$RUNNER_TEMP/$1" checkout --quiet FETCH_HEAD && echo "$1 $3 $(git -C "$RUNNER_TEMP/$1" rev-parse HEAD)"; }
+          checkout slime https://github.com/THUDM/slime "$SLIME_REF"
+          sglang=$(grep -oP '^ARG SGLANG_IMAGE_TAG=\K\S+' "$RUNNER_TEMP/slime/docker/Dockerfile" | sed -E 's/-cu[0-9]+$//')
+          megatron=$(grep -oP '^ARG MEGATRON_COMMIT=\K\S+' "$RUNNER_TEMP/slime/docker/Dockerfile")
+          checkout sglang https://github.com/sgl-project/sglang "$sglang"
+          checkout Megatron-LM https://github.com/NVIDIA/Megatron-LM "$megatron"
+          echo "PYTHONPATH=$RUNNER_TEMP/sglang/python:$RUNNER_TEMP/Megatron-LM" >> "$GITHUB_ENV"
+      - name: Install slime, its requirements and CPU torch
+        # mcp stays below 2, as the framework pins it; the rest beyond slime's requirements are what SGLang's tool-call
+        # parser and slime's argument parser import on CPU.
+        run: |
+          echo 'mcp>=1.25,<2' > "$RUNNER_TEMP/constraints.txt"
+          transformers=$(grep -oP '"\Ktransformers==[^"]+' "$RUNNER_TEMP/sglang/python/pyproject.toml" | head -1)
+          uv pip install --python .venv/bin/python -e "$RUNNER_TEMP/slime" --no-deps
+          uv pip install --python .venv/bin/python -r "$RUNNER_TEMP/slime/requirements.txt" torch "$transformers" \
+            pybase64 triton msgspec pyzmq IPython partial-json-parser xgrammar sentencepiece \
+            -c "$RUNNER_TEMP/constraints.txt" \
+            --extra-index-url https://download.pytorch.org/whl/cpu --index-strategy unsafe-best-match
+          .venv/bin/python -c "from importlib.metadata import version; print('torch', version('torch'), 'transformers', version('transformers'), 'mcp', version('mcp'))"
+      - name: slime tier
+        run: .venv/bin/python -m pytest tst/slime -rA --timeout=900 --junitxml=junit-slime.xml
+      - name: No skips
+        if: always()
+        run: .venv/bin/python .github/scripts/check_skip_policy.py junit-slime.xml

Makefile

@@ -1,4 +1,4 @@
-.PHONY: help install lint unit-test int-test int-test-fast int-test-slow installer-test trl-test nemo-gym-test verl-test skyrl-test miles-test clean-install-test test clean
+.PHONY: help install lint unit-test int-test int-test-fast int-test-slow installer-test trl-test nemo-gym-test verl-test skyrl-test miles-test slime-test clean-install-test test clean
 
 VENV := .venv
 VENV_BIN := $(abspath $(VENV))/bin
@@ -91,6 +91,12 @@ skyrl-test: ## Play rows through the SkyRL env and generator on SkyRL itself (ne
 miles-test: ## Play rows through Miles' rollout and session server (needs Miles, SGLang and Megatron-LM on PYTHONPATH, CPU torch, and the Hub)
 	$(PYTHON) -m pytest tst/miles --timeout=900 $(PYTEST_ARGS)
 
+# The slime tier plays rows through the slime helper's generate function on slime's own OpenAI adapter, with slime's own
+# arguments and Qwen3's tokenizer from the Hub. slime isn't on PyPI: install a checkout of it, its requirements and CPU
+# torch, and put SGLang's and Megatron-LM's checkouts on PYTHONPATH first, as .github/workflows/slime.yml does.
+slime-test: ## Play rows through the slime helper on slime's OpenAI adapter (needs slime, SGLang and Megatron-LM, CPU torch, and the Hub)
+	$(PYTHON) -m pytest tst/slime --timeout=900 $(PYTEST_ARGS)
+
 # The clean-install gate CI runs, on the same Python. The script gives every step an allowlisted
 # environment, so nothing in this shell reaches the install it checks. A failed run keeps its folder.
 clean-install-test: ## Build both distributions, install them clean from public PyPI and run hello twice (needs python3.11, uv, git)

pyproject.toml

@@ -97,6 +97,11 @@ skyrl = [
 miles = [
     "agentenv-framework[openenv]",
 ]
+# agent_env.training.slime: slime's trainers on a bundle served as an OpenEnv env; installed into slime's own environment,
+# whose slime, SGLang and torch are the trainer's to pin. It needs nothing beyond the openenv extra.
+slime = [
+    "agentenv-framework[openenv]",
+]
 # The AWS store backends (S3, Secrets Manager, DynamoDB, ECR), selected by impl pointer in config.toml.
 aws = [
     "boto3>=1.35.0",

src/agent_env/training/__init__.py

@@ -1,4 +1,4 @@
 """Training on AgentEnv tasks with other projects' trainers, one module per trainer, each behind an extra of its own:
 ``agent_env.training.trl`` for TRL, ``agent_env.training.nemo_gym`` for NeMo Gym and the trainers that take its
-environments, NeMo RL among them, ``agent_env.training.verl`` for verl, ``agent_env.training.skyrl`` for SkyRL, and
-``agent_env.training.miles`` for Miles."""
+environments, NeMo RL among them, ``agent_env.training.verl`` for verl, ``agent_env.training.skyrl`` for SkyRL,
+``agent_env.training.miles`` for Miles, and ``agent_env.training.slime`` for slime."""

src/agent_env/training/_chat.py

@@ -0,0 +1,164 @@
+"""What the Miles and slime helpers share: a server's tasks as chat rows, and a row's task played as one episode by a
+model chatting at a trainer's OpenAI chat endpoint, a session of the trainer's that records the model's turns for
+training. A row holds the task's system prompt and first prompt as ``messages``, and in ``metadata`` its episode, the
+tools its toolset offers and the server it plays on with its settings, so the trainer's rollout processes need no
+settings of their own."""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Iterator
+from pathlib import Path
+from typing import Any
+
+import click
+import httpx
+
+from agent_env.training._openenv import Session, Shared, function_tool, key, messages, tasks_and_tools, token_in
+
+
+@click.command()
+@click.argument("url")
+@click.argument("out", type=click.Path(file_okay=False))
+@click.option("--tool-timeout", type=float, default=60.0, show_default=True,
+              help="One tool call's deadline in seconds; the server ends a call that misses it, and the episode.")
+@click.option("--reset-timeout", type=float, default=600.0, show_default=True,
+              help="How long an episode may take to deploy, and each turn to be ended, in seconds; also how long a "
+                   "rollout waits for room on the server.")
+@click.option("--model-timeout", type=float, default=600.0, show_default=True,
+              help="How long a rollout waits for each of the model's replies, in seconds; an episode whose "
+                   "model misses it fails, and its group is left out.")
+@click.option("--max-turns", type=int, default=4, show_default=True,
+              help="The model turns a rollout may take; the episode is ended after the last.")
+@click.option("--token-var", help="The environment variable that holds the token the server asks for, read now and in "
+                                  "each of the trainer's rollout processes.")
+def command(url: str, out: str, tool_timeout: float, reset_timeout: float, model_timeout: float, max_turns: int,
+         token_var: str | None) -> None:
+    """Writes the tasks of the OpenEnv server at URL as OUT/<split>.jsonl, a row per task, for the trainer's
+    --prompt-data with --input-key messages and --metadata-key metadata."""
+    folder = write_rows(url, out, tool_timeout=tool_timeout, reset_timeout=reset_timeout, model_timeout=model_timeout,
+                        max_turns=max_turns, token_var=token_var)
+    click.echo(f"wrote {', '.join(sorted(path.name for path in folder.iterdir()))} to {folder}")
+
+
+
+def write_rows(url: str, out: str | Path, *, tool_timeout: float = 60.0, reset_timeout: float = 600.0,
+               model_timeout: float = 600.0, max_turns: int = 4, token_var: str | None = None) -> Path:
+    """The tasks of the OpenEnv server at ``url`` as ``out/<split>.jsonl``, a row per task of each split: its system
+    prompt and first prompt as ``messages``, and as ``metadata`` its split, index and task id, the tools its toolset
+    offers, and where and how ``play`` plays it: the server, its deadlines and the model's in seconds, the model turns a
+    rollout may take, and the environment variable holding the token a server given one asks for, read in each of the
+    trainer's rollout processes."""
+    url = url.rstrip("/")
+    _, tasks, probed = tasks_and_tools(url, local="serve the bundle with agent_env.openenv.build_app first",
+                                       timeout=reset_timeout, token=token_in(token_var))
+    tools = {each: [function_tool(tool) for tool in found] for each, found in probed.items()}
+    server = {"url": url, "tool_timeout": tool_timeout, "reset_timeout": reset_timeout, "model_timeout": model_timeout,
+              "max_turns": max_turns, **({"token_var": token_var} if token_var else {})}
+    out = Path(out)
+    out.mkdir(parents=True, exist_ok=True)
+    for split, described in tasks.items():
+        (out / f"{split}.jsonl").write_text("".join(json.dumps({
+            "messages": messages(task["system_prompt"], task["prompt"]),
+            "metadata": {"split": split, "index": task["index"], "task_id": task["id"], "tools": tools[key(task)],
+                         "agentenv": server},
+        }) + "\n" for task in described))
+    return out
+
+
+class _Episode(Session):
+    """One sample's episode on the server its row names."""
+
+    offered_by = "its row lists"
+
+    def __init__(self, shared: Shared):
+        super().__init__()
+        self._shared = shared
+
+
+async def play(base_url: str, prompt: list[dict[str, Any]], request_kwargs: dict[str, Any] | None = None,
+               metadata: dict[str, Any] | None = None, **kwargs: Any) -> dict[str, Any]:
+    """Plays the row's task as one episode: the model chats at ``base_url``, a session of the trainer's that records its
+    turns, with the row's tools; each call is a call on the episode, and the model's last reply, or the turn cap, ends
+    the turn. A call whose arguments aren't a JSON object is answered with why, as the model's error, and the episode
+    goes on. Returns the sample's metadata: the episode's reward and scores, or, when the env or the session failed it,
+    the model missed the row's ``model_timeout`` or the episode ended without a reward, a reward of 0 and why, which
+    ``keep_scored_groups`` leaves out of training; and the calls made and refused. A failed episode isn't played again:
+    the session already holds its turns."""
+    row = metadata or {}
+    server, tools = row["agentenv"], row["tools"]
+    episode = _Episode(Shared(server["url"], frozenset(tool["function"]["name"] for tool in tools),
+                              server["tool_timeout"], server["reset_timeout"], patience=server["reset_timeout"],
+                              waits_for_server=True, token=token_in(server.get("token_var"))))
+    counts = {"agentenv_tool_calls": 0, "agentenv_invalid_calls": 0}
+    try:
+        await episode._start(row["split"], int(row["index"]), row["task_id"])
+        reply = await _converse(episode, base_url, list(prompt), tools, request_kwargs or {}, server, counts)
+        rollout = episode._rollout
+        if not (rollout.done or rollout.failure):
+            await episode._finish(rollout, reply)
+        if rollout.failure is None and rollout.reward is None:
+            rollout.failure = "the episode ended without a reward"
+        if rollout.failure is not None:
+            return {"reward": 0.0, "agentenv_failure": rollout.failure, **counts}
+        return {"reward": float(rollout.reward), "agentenv_scores": rollout.scores or {}, **counts}
+    except Exception as e:
+        return {"reward": 0.0, "agentenv_failure": f"{type(e).__name__}: {e}", **counts}
+    finally:
+        await episode._close()
+
+
+async def reward(args: Any, samples: Any, **kwargs: Any) -> float | list[float]:
+    """Each sample's episode reward, which ``play`` put in its metadata."""
+    if isinstance(samples, list):
+        return [float(sample.metadata.get("reward", 0.0)) for sample in samples]
+    return float(samples.metadata.get("reward", 0.0))
+
+
+def keep_scored_groups(args: Any, samples: list[Any], **kwargs: Any) -> bool:
+    """Whether to train on a group: not when an episode of it failed in the env, its episode never returned (a child
+    process that died mid-episode leaves the session's turns with no reward), or the trainer aborted a sample of it,
+    none of them the policy's doing."""
+    return not any(sample.metadata.get("agentenv_failure") or "reward" not in sample.metadata
+                   or getattr(sample.status, "name", None) == "ABORTED" for sample in _samples(samples))
+
+
+async def _converse(episode: _Episode, base_url: str, messages: list[dict[str, Any]], tools: list[dict[str, Any]],
+                    request_kwargs: dict[str, Any], server: dict[str, Any], counts: dict[str, int]) -> str:
+    """The model's turns, until it replies without a call, the episode ends, or the server's ``max_turns`` pass; its
+    last reply."""
+    async with httpx.AsyncClient(timeout=server["model_timeout"]) as client:
+        for _ in range(server["max_turns"]):
+            response = await client.post(f"{base_url}/v1/chat/completions", json={
+                "model": "default", "messages": messages, "tools": tools, **request_kwargs})
+            response.raise_for_status()
+            message = response.json()["choices"][0]["message"]
+            # Verbatim: the session server matches each message a request replays against the one it recorded.
+            messages.append(message)
+            if not (calls := message.get("tool_calls")):
+                return message.get("content") or ""
+            for call in calls:
+                messages.append({"role": "tool", "tool_call_id": call.get("id", ""),
+                                 "content": await _answer(episode, call, counts)})
+            if episode._rollout.done or episode._rollout.failure:
+                return ""
+    return ""
+
+
+async def _answer(episode: _Episode, call: dict[str, Any], counts: dict[str, int]) -> str:
+    name, arguments = call["function"]["name"], call["function"].get("arguments")
+    try:
+        arguments = json.loads(arguments) if isinstance(arguments, str) else arguments
+    except json.JSONDecodeError as e:
+        counts["agentenv_invalid_calls"] += 1
+        return f"{name} wasn't called: its arguments aren't JSON ({e})"
+    if not isinstance(arguments, dict):
+        counts["agentenv_invalid_calls"] += 1
+        return f"{name} wasn't called: its arguments must be a JSON object, not {type(arguments).__name__}"
+    counts["agentenv_tool_calls"] += 1
+    return await episode._step(name, arguments)
+
+
+def _samples(group: list[Any]) -> Iterator[Any]:
+    for sample in group:
+        yield from sample if isinstance(sample, list) else [sample]

src/agent_env/training/miles/__main__.py

@@ -1,33 +1,6 @@
 """``python -m agent_env.training.miles URL OUT``: writes a bundle's tasks as Miles' prompt data."""
 
-import click
-
-from agent_env.training.miles.prepare import write_rows
-
-
-@click.command()
-@click.argument("url")
-@click.argument("out", type=click.Path(file_okay=False))
-@click.option("--tool-timeout", type=float, default=60.0, show_default=True,
-              help="One tool call's deadline in seconds; the server ends a call that misses it, and the episode.")
-@click.option("--reset-timeout", type=float, default=600.0, show_default=True,
-              help="How long an episode may take to deploy, and each turn to be ended, in seconds; also how long a "
-                   "rollout waits for room on the server.")
-@click.option("--model-timeout", type=float, default=600.0, show_default=True,
-              help="How long the agent function waits for each of the model's replies, in seconds; an episode whose "
-                   "model misses it fails, and its group is left out.")
-@click.option("--max-turns", type=int, default=4, show_default=True,
-              help="The model turns a rollout may take; the episode is ended after the last.")
-@click.option("--token-var", help="The environment variable that holds the token the server asks for, read now and in "
-                                  "each of Miles' rollout processes.")
-def main(url: str, out: str, tool_timeout: float, reset_timeout: float, model_timeout: float, max_turns: int,
-         token_var: str | None) -> None:
-    """Writes the tasks of the OpenEnv server at URL as OUT/<split>.jsonl, a row per task, for Miles' --prompt-data
-    with --input-key messages and --metadata-key metadata."""
-    folder = write_rows(url, out, tool_timeout=tool_timeout, reset_timeout=reset_timeout, model_timeout=model_timeout,
-                        max_turns=max_turns, token_var=token_var)
-    click.echo(f"wrote {', '.join(sorted(path.name for path in folder.iterdir()))} to {folder}")
-
+from agent_env.training._chat import command as main
 
 if __name__ == "__main__":
     main()