Merge pull request #61 from awslabs/feat/slime-data-contract

feat(slime): support arbitrary agent payload shapes in the training backend

Author: lliquid

src/agentcore_rl_toolkit/backends/slime/SETUP.md

@@ -112,11 +112,22 @@ ds = load_dataset('openai/gsm8k', 'main', split='train')
 with open('/path/to/gsm8k_tiny.jsonl', 'w') as f:
     for i, row in enumerate(ds):
         if i >= 64: break
+        question = row['question']
         answer = row['answer'].split('####')[-1].strip()
-        f.write(json.dumps({'prompt': row['question'], 'label': answer}) + '\n')
+        # Top-level 'prompt' is read by slime (tokenization, length filter).
+        # 'metadata' is the agent payload verbatim — shape it however the agent expects.
+        f.write(json.dumps({
+            'prompt': question,
+            'metadata': {'prompt': question, 'answer': answer},
+        }) + '\n')
 "
 ```
 
+The agent-visible payload is exactly the contents of ``metadata``, so
+different agents can use different payload shapes (e.g. ``{'task_id': ...}``
+for AppWorld, ``{'repo_uri': ..., 'metadata_uri': ..., ...}`` for
+migration) without any slime-side changes.
+
 ### 3.3 Configure deployment settings
 
 ```bash

src/agentcore_rl_toolkit/backends/slime/examples/math_agent/train.sh

@@ -80,7 +80,6 @@ ray job submit --address="http://127.0.0.1:8265" \
   --tensor-model-parallel-size ${TP_SIZE} \
   --rollout-num-gpus-per-engine ${ROLLOUT_GPUS_PER_ENGINE} \
   --input-key prompt \
-  --label-key label \
   --rollout-batch-size 32 \
   --n-samples-per-prompt 8 \
   --rollout-max-response-len 1024 \

src/agentcore_rl_toolkit/backends/slime/integration/rollout.py

@@ -209,52 +209,20 @@ def _ensure_initialized(args: Namespace):
 
 
 def _sample_to_payload(sample) -> dict:
-    """Convert a slime Sample to an ART invocation payload.
-
-    Extracts all non-None public fields from the Sample into the payload.
-    The agent's @rollout_entrypoint receives this dict as `payload`.
+    """The agent payload is the JSONL row's ``metadata`` dict, verbatim.
+
+    slime's Dataset reads the JSONL row's ``metadata`` field into
+    ``Sample.metadata``; we hand that dict to the agent unchanged. The JSONL's
+    top-level ``prompt`` field is for slime (tokenization, length filtering);
+    the agent's payload shape is entirely defined by whatever the data author
+    put in ``metadata``. A shallow copy isolates the agent's view from
+    downstream mutations to ``Sample.metadata`` (e.g. ``task_metadata``
+    injection in ``_process_one_episode``).
     """
-    payload = {}
-
-    if hasattr(sample, "prompt") and sample.prompt:
-        payload["prompt"] = sample.prompt
-    if hasattr(sample, "label") and sample.label is not None:
-        payload["answer"] = sample.label
-    if hasattr(sample, "metadata") and sample.metadata:
-        payload["metadata"] = sample.metadata
-
-    if hasattr(sample, "to_dict"):
-        for key, value in sample.to_dict().items():
-            if (
-                key not in payload
-                and value is not None
-                and key
-                not in (
-                    "tokens",
-                    "rollout_log_probs",
-                    "loss_mask",
-                    "teacher_log_probs",
-                    "rollout_routed_experts",
-                    "multimodal_inputs",
-                    "multimodal_train_inputs",
-                    "group_index",
-                    "index",
-                    "status",
-                    "session_id",
-                    "spec_info",
-                    "prefix_cache_info",
-                    "response_length",
-                    "response",
-                    "weight_versions",
-                    "remove_sample",
-                    "non_generation_time",
-                    "generate_function_path",
-                    "train_metadata",
-                )
-            ):
-                payload[key] = value
-
-    return payload
+    metadata = getattr(sample, "metadata", None)
+    if isinstance(metadata, dict):
+        return dict(metadata)
+    return {}
 
 
 # ---------------------------------------------------------------------------

tests/test_slime_rollout_payload.py

@@ -0,0 +1,49 @@
+"""Tests for the slime backend's agent-payload conversion.
+
+``_sample_to_payload`` is the contract between slime's ``Sample`` object
+and the agent's ``@rollout_entrypoint`` payload. The rule: the agent
+receives ``Sample.metadata`` verbatim (shallow-copied).
+"""
+
+from types import SimpleNamespace
+
+from agentcore_rl_toolkit.backends.slime.integration.rollout import _sample_to_payload
+
+
+def test_metadata_is_returned_verbatim():
+    """Core contract: the agent payload is Sample.metadata as-is.
+
+    Fields on Sample outside metadata (prompt, label) are slime's own
+    concern and must not leak into the payload.
+    """
+    sample = SimpleNamespace(
+        prompt="slime-side prompt",
+        label="slime-side label",
+        metadata={"task_id": "t1", "answer": "42"},
+    )
+    assert _sample_to_payload(sample) == {"task_id": "t1", "answer": "42"}
+
+
+def test_returned_dict_is_a_shallow_copy():
+    """Mutations to the payload must not leak back into Sample.metadata.
+
+    ``_process_one_episode`` later injects keys into ``Sample.metadata``
+    (e.g. ``task_metadata``); the agent's view must stay stable.
+    """
+    metadata = {"prompt": "hi"}
+    sample = SimpleNamespace(metadata=metadata)
+
+    payload = _sample_to_payload(sample)
+    payload["injected"] = True
+
+    assert "injected" not in metadata
+
+
+def test_missing_or_invalid_metadata_returns_empty_dict():
+    """Defensive fallback when metadata is absent or not a dict."""
+    for sample in [
+        SimpleNamespace(),  # attribute absent
+        SimpleNamespace(metadata=None),  # explicit None
+        SimpleNamespace(metadata="not a dict"),  # wrong type
+    ]:
+        assert _sample_to_payload(sample) == {}