Inference Backend

1 The Backend Abstraction

Every LLM call in a rollout goes through a backend — an object turning a prompt into a completion, defined as a structural Protocol (rl_engine/rollout/backends/base.py), not a base class. Two modes coexist:

The agent loop reads supports_token_ids to pick the input type and whether to record tokens. Both return one dataclass, GenerationResult:

# rl_engine/rollout/backends/base.py
@dataclass
class GenerationResult:
    text: str
    token_ids: list[int] | None = None      # SGLang only
    logprobs: list[float] | None = None    # SGLang only
    tool_calls: list[dict] | None = None    # API structured calls
    finish_reason: str = "stop"              # stop | length | abort | tool_calls | error:*
    usage: dict[str, int] | None = None  # API token accounting

class ModelBackend(Protocol):
    async def generate(self, *, input_ids=None, messages=None,
                       tools=None, sampling_params) -> GenerationResult: ...
    def tokenize(self, text) -> list[int]: ...   # API backends raise NotImplementedError
    def decode(self, token_ids) -> str: ...

The backends directory

2 The SGLang Student Backend

SGLangBackend (backends/sglang.py) is the student policy's generation path — token-in / token-out. It reads the exact sampled token IDs and their logprobs from SGLang's native /generate: no re-tokenization, no chat-template round-trip, so the tokens the model emitted are the tokens training sees.

# SGLangBackend.generate — the request
url = f"{self._sglang_url}/generate"
payload = {
    "input_ids": input_ids,
    "sampling_params": sampling_params,
    "return_logprob": True,
}
output = await self._post_fn(url, payload)

# finish_reason == "abort" -> empty result (SGLang killed the request)
# else read meta_info["output_token_logprobs"] as (logprob, token_id) pairs
token_ids = [item[1] for item in meta["output_token_logprobs"]]
logprobs  = [item[0] for item in meta["output_token_logprobs"]]

A post_fn argument overrides the HTTP transport: everything uses slime's pooled post except the recording proxy, which injects its own loop-local httpx client (§4).

Where the URL comes from

The backend is built lazily inside the in-house SWE agent (swe_agent.py) from context.sglang_url, which the Evaluator sets to http://{sglang_router_ip}:{sglang_router_port} (scheduler/evaluator.py). Those router args come from two places:

• Colocated training — slime stands up its own SGLang router alongside the Megatron workers and sets the args automatically. Nothing to configure.
• External / eval — eval.sh passes --rollout-external --rollout-external-engine-addrs plus explicit --sglang-router-ip/--sglang-router-port, from ROLLOUT_EXTERNAL_ENGINE_ADDRS (default 0.0.0.0:9001); the engine is a standalone server (§6).

Weight sync between iterations

In colocated training the SGLang engines share GPUs with the trainer, so updated weights must reach them before the next rollout. The main loop calls trainer.update_weights() before training and after each iteration (rl_engine/train.py), fanning out only to policy workers (roles with sync_to_rollout). Colocated runs hand off via CPU shared memory (dodging CUDA-IPC failures under Ray fractional-GPU placement), otherwise NCCL broadcast; trainer-side mechanics are on the Backends page.

Serving knobs

The installed-scaffold factory (evaluator_factory.py) caps agent context at a fixed 131072 rather than reading MAX_CONTEXT_LEN — that env var applies only to the in-house SWE-agent path.

3 The Teacher Backend

TeacherBackend (backends/teacher.py) wraps an external SGLang server hosting the frozen Qwen3-Coder-30B teacher. Because the teacher speaks a different chat template and tool syntax than the student, it loads its own tokenizer (CPU-only) and resolves its own tool-call format. It never acts in the environment, and any failure returns None without touching the student's trajectory. It plays two roles.

Role A — per-turn action labelling (DAgger / AggreVaTe rollin)

When a turn is teacher-executed, label_turn(history, task_prompt, system_prompt) renders a prompt in the teacher's format (include_reasoning=True, so it sees its own prior chain-of-thought) and parses the reply into a CanonicalAction; the lower-level generate_from_tokens(input_ids) takes a pre-built token sequence instead. The rollin protocols live in dagger_agent.py — see Algorithm Recipes.

Role B — post-rollout force-decode for the K3 loss

After a rollout is collected and finish-filtered, each kept trajectory is force-decoded by the teacher to recover its probability of every student-emitted token — the log p_t term the K3 reverse-KL needs. Orchestration is _apply_distillation_capture (rollout_manager.py), calling score_response_with_topk (distill_capture.py).

Teacher force-decode. Kept trajectories flow right to the frozen teacher, whose per-token chosen + top-K logprobs return along the bottom as K3 loss inputs.

The payload force-decodes the exact response — one greedy step over prompt_ids + response_ids with logprobs on:

# distill_capture.py — score_response_with_topk (payload)
sampling_params = {
    "max_new_tokens": 1, "temperature": 0.0, "top_p": 1.0,
    "top_k": -1, "no_stop_trim": True,
}
payload = {
    "input_ids": prompt_ids + response_ids,   # force-decode the exact response
    "sampling_params": sampling_params,
    "return_logprob": True,
    "logprob_start_len": start_len,           # = prompt_len - 1
    "top_logprobs_num": K,                    # STACX_DISTILL_TOPK (default 20)
}

A persistent teacher failure is not swallowed — it propagates, failing the run loudly rather than silently emitting empty logprobs. The captured fields are pinned onto each Sample (teacher_log_probs, teacher_topk_indices, teacher_topk_logprobs), and loss_mask is overwritten with the action mask so the loss sees exactly the student-emitted tokens.

Teacher & distillation knobs

4 The Recording Proxy

The in-house SWE agent emits tokens natively; the black-box installed scaffolds — OpenHands, Terminus-2 — cannot. The RecordingProxy (proxy/recording_proxy.py) bridges the gap: an OpenAI-compatible façade between scaffold and SGLang that intercepts /v1/chat/completions and reconstructs a token-level trajectory without touching the scaffold. One proxy per task, on an auto-assigned port, torn down at task end.

Token capture. The proxy builds one token series across a task's chat calls.

The per-call cycle

1. Sync the series — _sync_series finds the messages appended since the last turn, skipping the assistant turns the scaffold echoes back (already in the series).
2. Tokenize the increment — new observation/tool messages go in mid-conversation via tokenize_messages_increment with loss_mask = 0; the first call renders the whole prompt.
3. Generate — /generate runs on the exact series; the returned token IDs + logprobs are appended verbatim with loss_mask = 1.
4. Reply — those tokens are decoded back into an OpenAI ChatCompletion (with tool_calls) so the scaffold acts as it would have.

The append-only guard

The scheme assumes an append-only conversation. If a scaffold rewrites or summarizes its history, the stored series no longer matches — so _sync_series checks a per-message signature (_sig: role + content shape) against everything consumed so far. A shrunk or rewritten prefix raises _AppendOnlyViolation: capture is disabled and the request passes through, so the run completes but emits no token data. Semantics are pinned by test_recording_proxy_guard.py:

# the guard, distilled from the tests
_sig(m)  ->  (role, content_shape)         # role or content change is detected

pure append (messages grow)      ->  OK
history shrank  (len drops)      ->  raise _AppendOnlyViolation("shrank")
prefix rewritten (summary swap)  ->  raise _AppendOnlyViolation("rewritten at message i")

Other failures are equally defensive: an SGLang abort or any unexpected exception flips capture_valid off and passes through, while get_result_fields() returns None on an invalid or length-inconsistent series (logprobs, mask, and tokens must align) — dropping the trajectory rather than risking a corrupt sample. The result matches the native SWE-agent shape: one completion region of prompt_token_ids, completion_token_ids, logprobs, and a loss_mask that is 1 on generated tokens (append_generation) and 0 on observations (append_observation).

5 The Installed-Scaffold Path

InstalledAgent (agents/installed.py) runs a scaffold as a black-box CLI tool inside a ROCK container: it owns the container lifecycle, while a per-scaffold InstalledScaffold supplies the install/run commands and log/termination semantics — shell commands in, files out.

The runner

InstalledAgent.run() executes a fixed per-task sequence: start ROCK sandbox → optionally start recording proxy → install scaffold → run it → freeze benchmark artifact → run verifier → parse trajectory → capture logs → classify termination → tear down. The proxy interposition is the switch that turns eval into training:

# installed.py — _maybe_start_proxy (train vs eval, one branch)
if not self.capture_tokens:
    return None                       # eval: scaffold talks to SGLang directly
proxy = RecordingProxy(sglang_url=..., tokenizer_path=hf_checkpoint, ...)
proxy.start()
context.proxy_url = proxy.url          # scaffold is pointed here instead

The scaffold's base URL is then context.proxy_url or context.sglang_url — the proxy when capture is on, SGLang directly when off (as in eval). After teardown, proxy.get_result_fields() pulls the token series onto the AgentResult.

Termination-reason classification

Every run gets a _termination_reason that drives the failure-class summary. The base classifier (InstalledScaffold.termination_reason) keys off the exit code; scaffolds refine the clean exit. OpenHands' (pinned by test_openhands_termination.py) applies this precedence, first match wins:

Terminus-2's is simpler: a clean exit with a mark_task_complete call is finish, else max_turns if the step budget was hit, else finish; nonzero exits defer to the base classes.

Capture vs. plain evaluators

Which variant runs is a build-function choice. SkyRL (skyrl/evaluator_config.py) exposes four: build_evaluators_openhands, …_openhands_capture, …_terminus2, …_terminus2_capture. The _capture variants pass capture_tokens=True — the master switch that turns the proxy on and forces append-only history (raising if APPEND_ONLY=off). The rest leave the proxy off and allow summarization: pure eval.

6 Serving Topologies

The same backends wire up three ways:

7 Quick Reference

Related pages: Rollout Architecture · Task Scheduler · Algorithm Recipes · Sandbox Management · Trainer Backends.