synth-ai 0.2.9.dev4__py3-none-any.whl → 0.2.9.dev6__py3-none-any.whl

examples/swe/task_app/hosted/storage/__init__.py

@@ -0,0 +1,5 @@
+"""Storage module for Modal Volume operations."""
+
+from .volume import VolumeStorage, storage
+
+__all__ = ["VolumeStorage", "storage"]

examples/swe/task_app/hosted/storage/volume.py

@@ -0,0 +1,211 @@
+from __future__ import annotations
+
+import gzip
+import hashlib
+import json
+import os
+import tarfile
+import tempfile
+from datetime import datetime
+from pathlib import Path
+from typing import Any
+
+
+class VolumeStorage:
+    """Helpers for Modal Volume storage operations."""
+
+    def __init__(self, base_path: str = "/data/state") -> None:
+        self.base_path = Path(base_path)
+
+    def get_snapshot_path(
+        self,
+        rl_run_id: str,
+        kind: str,
+        snapshot_id: str,
+    ) -> Path:
+        """Build the path for a snapshot file."""
+        # Use first 2 chars of snapshot_id for sharding
+        shard1 = snapshot_id[:2] if len(snapshot_id) >= 2 else "00"
+        shard2 = snapshot_id[2:4] if len(snapshot_id) >= 4 else "00"
+
+        return (
+            self.base_path / "runs" / rl_run_id / kind / shard1 / shard2 / f"{snapshot_id}.tar.gz"
+        )
+
+    def get_index_path(self, rl_run_id: str) -> Path:
+        """Get the index file path for a run."""
+        return self.base_path / "runs" / rl_run_id / "index" / "meta.jsonl"
+
+    def write_snapshot_atomic(
+        self,
+        path: Path,
+        archive_bytes: bytes,
+    ) -> None:
+        """Atomically write a snapshot archive to disk."""
+        # Ensure parent directory exists
+        path.parent.mkdir(parents=True, exist_ok=True)
+
+        # Write to temp file first
+        tmp_path = path.with_suffix(".tmp")
+        with open(tmp_path, "wb") as f:
+            f.write(archive_bytes)
+            f.flush()
+            os.fsync(f.fileno())
+
+        # Atomic rename
+        os.replace(tmp_path, path)
+
+    def create_archive(
+        self,
+        state_dict: dict[str, Any],
+        meta: dict[str, Any],
+    ) -> bytes:
+        """Create a tar.gz archive with state and metadata."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            tmppath = Path(tmpdir)
+
+            # Write state.json
+            state_path = tmppath / "state.json"
+            with open(state_path, "w") as f:
+                json.dump(state_dict, f, sort_keys=True, indent=2)
+
+            # Write meta.json
+            meta_path = tmppath / "meta.json"
+            with open(meta_path, "w") as f:
+                json.dump(meta, f, sort_keys=True, indent=2)
+
+            # Create tar archive
+            tar_path = tmppath / "archive.tar"
+            with tarfile.open(tar_path, "w") as tar:
+                tar.add(state_path, arcname="state.json")
+                tar.add(meta_path, arcname="meta.json")
+
+            # Compress with gzip
+            with open(tar_path, "rb") as f:
+                tar_bytes = f.read()
+
+            compressed = gzip.compress(tar_bytes, compresslevel=6)
+
+            return compressed
+
+    def extract_archive(self, archive_bytes: bytes) -> tuple[dict[str, Any], dict[str, Any]]:
+        """Extract state and metadata from a tar.gz archive."""
+        # Decompress
+        tar_bytes = gzip.decompress(archive_bytes)
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            tmppath = Path(tmpdir)
+
+            # Write tar bytes to temp file
+            tar_path = tmppath / "archive.tar"
+            with open(tar_path, "wb") as f:
+                f.write(tar_bytes)
+
+            # Extract tar
+            with tarfile.open(tar_path, "r") as tar:
+                tar.extractall(tmppath)
+
+            # Read state and meta
+            with open(tmppath / "state.json") as f:
+                state = json.load(f)
+
+            with open(tmppath / "meta.json") as f:
+                meta = json.load(f)
+
+            return state, meta
+
+    def compute_snapshot_id(self, archive_bytes: bytes) -> str:
+        """Compute content-addressed snapshot ID."""
+        return hashlib.sha256(archive_bytes).hexdigest()
+
+    def save_snapshot(
+        self,
+        rl_run_id: str,
+        kind: str,
+        state_dict: dict[str, Any],
+        config: dict[str, Any] | None = None,
+        parent_snapshot_id: str | None = None,
+    ) -> tuple[str, str, int]:
+        """Save a snapshot and return (snapshot_id, path, size)."""
+        # Build metadata
+        meta = {
+            "kind": kind,
+            "rl_run_id": rl_run_id,
+            "schema_version": "1.0",
+            "created_at": datetime.utcnow().isoformat(),
+        }
+
+        if parent_snapshot_id:
+            meta["parent_snapshot_id"] = parent_snapshot_id
+
+        if config:
+            config_str = json.dumps(config, sort_keys=True)
+            meta["config_hash"] = hashlib.sha256(config_str.encode()).hexdigest()
+
+        # Create archive
+        archive_bytes = self.create_archive(state_dict, meta)
+
+        # Compute snapshot ID
+        snapshot_id = self.compute_snapshot_id(archive_bytes)
+        meta["snapshot_id"] = snapshot_id
+
+        # Recreate archive with snapshot_id in metadata
+        archive_bytes = self.create_archive(state_dict, meta)
+
+        # Get path and write
+        path = self.get_snapshot_path(rl_run_id, kind, snapshot_id)
+        self.write_snapshot_atomic(path, archive_bytes)
+
+        # Append to index
+        self.append_to_index(rl_run_id, meta)
+
+        return snapshot_id, str(path), len(archive_bytes)
+
+    def load_snapshot(
+        self,
+        rl_run_id: str,
+        kind: str,
+        snapshot_id: str,
+    ) -> tuple[dict[str, Any], dict[str, Any]]:
+        """Load a snapshot and return (state_dict, meta)."""
+        path = self.get_snapshot_path(rl_run_id, kind, snapshot_id)
+
+        if not path.exists():
+            raise FileNotFoundError(f"Snapshot not found: {path}")
+
+        with open(path, "rb") as f:
+            archive_bytes = f.read()
+
+        state, meta = self.extract_archive(archive_bytes)
+        return state, meta
+
+    def append_to_index(
+        self,
+        rl_run_id: str,
+        meta: dict[str, Any],
+    ) -> None:
+        """Append metadata to the run's index file."""
+        index_path = self.get_index_path(rl_run_id)
+        index_path.parent.mkdir(parents=True, exist_ok=True)
+
+        with open(index_path, "a") as f:
+            f.write(json.dumps(meta) + "\n")
+
+    def read_index(self, rl_run_id: str) -> list[dict[str, Any]]:
+        """Read all entries from a run's index file."""
+        index_path = self.get_index_path(rl_run_id)
+
+        if not index_path.exists():
+            return []
+
+        entries = []
+        with open(index_path) as f:
+            for line in f:
+                if line.strip():
+                    entries.append(json.loads(line))
+
+        return entries
+
+
+# Global storage instance
+storage = VolumeStorage()

examples/swe/task_app/hosted/test_agents.py

@@ -0,0 +1,161 @@
+#!/usr/bin/env python3
+"""
+Smoke test for Wordle and Sokoban ReAct agents using the hosted service.
+
+Prereqs:
+- Run the service: python examples/swe/task_app/hosted/main.py
+- Run an OpenAI-compatible inference server (e.g., Flash/vLLM) at VLLM_BASE_URL
+  that serves model "gpt-5-nano" or adjust MODEL below.
+
+This script will:
+- Create a Wordle/Sokoban env
+- Create corresponding *-react policy with tools
+- Ask the policy for tool_calls via /policy/step (which calls the model)
+- Apply tool_calls to the env via /env/step
+"""
+
+import asyncio
+import os
+
+import httpx
+
+BASE_URL = os.environ.get("SYNTH_ENVS_HOSTED_URL", "http://localhost:8000")
+INFER_URL = os.environ.get("VLLM_BASE_URL", "http://localhost:8001")
+MODEL = os.environ.get("MODEL", "gpt-5-nano")
+
+
+async def run_wordle(rounds: int = 3) -> None:
+    async with httpx.AsyncClient() as client:
+        # Create env
+        resp = await client.post(
+            f"{BASE_URL}/env/create",
+            json={
+                "env_name": "Wordle",
+                "config": {"word_length": 5, "max_guesses": 6},
+                "seed": 0,
+                "rl_run_id": "agents-smoke",
+            },
+        )
+        resp.raise_for_status()
+        data = resp.json()
+        env_id = data["env_id"]
+        obs = data["observation"]
+        print("Wordle env created:", env_id)
+
+        # Create policy
+        resp = await client.post(
+            f"{BASE_URL}/policy/create",
+            json={
+                "policy_name": "wordle-react",
+                "config": {
+                    "inference_url": INFER_URL,
+                    "model": MODEL,
+                    "use_tools": True,
+                    "word_length": 5,
+                    "max_guesses": 6,
+                },
+                "rl_run_id": "agents-smoke",
+                "bound_env_id": env_id,
+            },
+        )
+        resp.raise_for_status()
+        policy_id = resp.json()["policy_id"]
+        print("Wordle policy:", policy_id)
+
+        # Loop a few rounds
+        for i in range(rounds):
+            print(f"[Wordle] Round {i + 1}")
+            step_req = {"policy_id": policy_id, "observation": obs, "dry_run": False}
+            resp = await client.post(f"{BASE_URL}/policy/step", json=step_req)
+            resp.raise_for_status()
+            step_out = resp.json()
+            tool_calls = step_out.get("tool_calls", [])
+            print(" tool_calls:", tool_calls)
+            if not tool_calls:
+                break
+            resp = await client.post(
+                f"{BASE_URL}/env/step",
+                json={"env_id": env_id, "tool_calls": tool_calls},
+            )
+            resp.raise_for_status()
+            env_step = resp.json()
+            obs = env_step["observation"]
+            print(" done:", env_step.get("done"), "reward:", env_step.get("reward"))
+            if env_step.get("done"):
+                break
+
+
+async def run_sokoban(rounds: int = 3) -> None:
+    async with httpx.AsyncClient() as client:
+        # Create env (no initial_state provided; relies on env default)
+        resp = await client.post(
+            f"{BASE_URL}/env/create",
+            json={
+                "env_name": "Sokoban",
+                "config": {"difficulty": "easy"},
+                "seed": 0,
+                "rl_run_id": "agents-smoke",
+            },
+        )
+        if resp.status_code != 200:
+            print("Sokoban create failed:", resp.status_code, resp.text)
+            return
+        data = resp.json()
+        env_id = data["env_id"]
+        obs = data["observation"]
+        print("Sokoban env created:", env_id)
+
+        resp = await client.post(
+            f"{BASE_URL}/policy/create",
+            json={
+                "policy_name": "sokoban-react",
+                "config": {
+                    "inference_url": INFER_URL,
+                    "model": MODEL,
+                    "use_tools": True,
+                },
+                "rl_run_id": "agents-smoke",
+                "bound_env_id": env_id,
+            },
+        )
+        if resp.status_code != 200:
+            print("Sokoban policy create failed:", resp.status_code, resp.text)
+            return
+        policy_id = resp.json()["policy_id"]
+        print("Sokoban policy:", policy_id)
+
+        for i in range(rounds):
+            print(f"[Sokoban] Round {i + 1}")
+            step_req = {"policy_id": policy_id, "observation": obs, "dry_run": False}
+            resp = await client.post(f"{BASE_URL}/policy/step", json=step_req)
+            if resp.status_code != 200:
+                print(" policy step failed:", resp.status_code, resp.text)
+                break
+            step_out = resp.json()
+            tool_calls = step_out.get("tool_calls", [])
+            print(" tool_calls:", tool_calls)
+            if not tool_calls:
+                break
+            resp = await client.post(
+                f"{BASE_URL}/env/step",
+                json={"env_id": env_id, "tool_calls": tool_calls},
+            )
+            if resp.status_code != 200:
+                print(" env step failed:", resp.status_code, resp.text)
+                break
+            env_step = resp.json()
+            obs = env_step["observation"]
+            print(" done:", env_step.get("done"), "reward:", env_step.get("reward"))
+            if env_step.get("done"):
+                break
+
+
+async def main():
+    print("Testing Wordle agent with model:", MODEL)
+    await run_wordle(rounds=3)
+    print("\nTesting Sokoban agent with model:", MODEL)
+    await run_sokoban(rounds=3)
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

examples/swe/task_app/hosted/test_service.py

@@ -0,0 +1,137 @@
+#!/usr/bin/env python3
+"""
+Simple test script for the GRPO Synth Envs Hosted Service.
+
+Run this after starting the service with:
+    python main.py
+"""
+
+import asyncio
+import json
+
+import httpx
+
+
+async def test_service():
+    """Test basic service functionality."""
+    base_url = "http://localhost:8000"
+
+    async with httpx.AsyncClient() as client:
+        # Test 1: Service info
+        print("1. Testing /info endpoint...")
+        response = await client.get(f"{base_url}/info")
+        assert response.status_code == 200
+        info = response.json()
+        print(f"   Service info: {json.dumps(info, indent=2)}")
+
+        # Test 2: Health check
+        print("\n2. Testing /health endpoint...")
+        response = await client.get(f"{base_url}/health")
+        assert response.status_code == 200
+        print(f"   Health: {response.json()}")
+
+        # Test 3: Create environment
+        print("\n3. Creating environment...")
+        response = await client.post(
+            f"{base_url}/env/create",
+            json={
+                "env_name": "crafter",
+                "config": {},
+                "seed": 42,
+                "rl_run_id": "test-run-001",
+            },
+        )
+        if response.status_code != 200:
+            print(f"   Error: {response.status_code} - {response.text}")
+            return
+        env_data = response.json()
+        env_id = env_data["env_id"]
+        print(f"   Created env: {env_id}")
+        print(f"   Initial observation keys: {list(env_data['observation'].keys())}")
+
+        # Test 4: Create policy
+        print("\n4. Creating policy...")
+        response = await client.post(
+            f"{base_url}/policy/create",
+            json={
+                "policy_name": "crafter-react",
+                "config": {
+                    "inference_url": "http://localhost:8001",
+                    "model": "test-model",
+                },
+                "rl_run_id": "test-run-001",
+                "bound_env_id": env_id,
+            },
+        )
+        if response.status_code != 200:
+            print(f"   Error: {response.status_code} - {response.text}")
+            return
+        policy_data = response.json()
+        policy_id = policy_data["policy_id"]
+        print(f"   Created policy: {policy_id}")
+
+        # Test 5: Environment step with dummy tool calls
+        print("\n5. Testing environment step...")
+        response = await client.post(
+            f"{base_url}/env/step",
+            json={
+                "env_id": env_id,
+                "tool_calls": [{"tool": "interact", "args": {"action": "move_left"}}],
+            },
+        )
+        if response.status_code != 200:
+            print(f"   Error: {response.status_code} - {response.text}")
+        else:
+            step_data = response.json()
+            print(f"   Step result - done: {step_data['done']}, reward: {step_data.get('reward')}")
+
+        # Test 6: Environment snapshot
+        print("\n6. Creating environment snapshot...")
+        response = await client.post(f"{base_url}/env/snapshot", json={"env_id": env_id})
+        if response.status_code != 200:
+            print(f"   Error: {response.status_code} - {response.text}")
+        else:
+            snapshot_data = response.json()
+            print(f"   Snapshot ID: {snapshot_data['snapshot_id']}")
+            print(f"   Size: {snapshot_data['size']} bytes")
+
+        # Test 7: Policy snapshot
+        print("\n7. Creating policy snapshot...")
+        response = await client.post(f"{base_url}/policy/snapshot", json={"policy_id": policy_id})
+        if response.status_code != 200:
+            print(f"   Error: {response.status_code} - {response.text}")
+        else:
+            snapshot_data = response.json()
+            print(f"   Snapshot ID: {snapshot_data['snapshot_id']}")
+            print(f"   Size: {snapshot_data['size']} bytes")
+
+        # Test 8: Run status
+        print("\n8. Testing run status...")
+        response = await client.get(f"{base_url}/run/status/test-run-001")
+        if response.status_code != 200:
+            print(f"   Error: {response.status_code} - {response.text}")
+        else:
+            status_data = response.json()
+            print(f"   Run status: {status_data['status']}")
+
+        # Test 9: Terminate environment
+        print("\n9. Terminating environment...")
+        response = await client.post(f"{base_url}/env/terminate", json={"env_id": env_id})
+        if response.status_code != 200:
+            print(f"   Error: {response.status_code} - {response.text}")
+        else:
+            print(f"   Environment terminated: {response.json()['ok']}")
+
+        # Test 10: Terminate policy
+        print("\n10. Terminating policy...")
+        response = await client.post(f"{base_url}/policy/terminate", json={"policy_id": policy_id})
+        if response.status_code != 200:
+            print(f"   Error: {response.status_code} - {response.text}")
+        else:
+            print(f"   Policy terminated: {response.json()['ok']}")
+
+        print("\n✅ All basic tests completed!")
+
+
+if __name__ == "__main__":
+    asyncio.run(test_service())

examples/swe/task_app/hosted/utils.py

@@ -0,0 +1,62 @@
+"""Utility functions for the task service."""
+
+from typing import Any
+
+import numpy as np
+
+
+def convert_numpy_to_python(obj: Any) -> Any:
+    """
+    Recursively convert numpy types to Python native types for JSON serialization.
+
+    Args:
+        obj: Object that may contain numpy types
+
+    Returns:
+        Object with numpy types converted to Python native types
+    """
+    if isinstance(obj, np.integer):
+        return int(obj)
+    elif isinstance(obj, np.floating):
+        return float(obj)
+    elif isinstance(obj, np.ndarray):
+        return obj.tolist()
+    elif isinstance(obj, dict):
+        return {key: convert_numpy_to_python(value) for key, value in obj.items()}
+    elif isinstance(obj, list | tuple):
+        return [convert_numpy_to_python(item) for item in obj]
+    else:
+        return obj
+
+
+def sanitize_observation(observation: dict[str, Any]) -> dict[str, Any]:
+    """
+    Sanitize observation data for JSON serialization.
+
+    Converts numpy types and removes non-serializable objects.
+
+    Args:
+        observation: Raw observation from environment
+
+    Returns:
+        Sanitized observation safe for JSON serialization
+    """
+    if not isinstance(observation, dict):
+        return observation
+
+    sanitized = {}
+    for key, value in observation.items():
+        # Skip non-serializable keys or convert them
+        if key in ["semantic_map", "world_material_map", "observation_image"]:
+            # These are likely numpy arrays - convert to lists or skip
+            if isinstance(value, np.ndarray):
+                # For large arrays, we might want to skip or compress
+                # For now, skip them as they're likely debug info
+                continue
+        elif key == "player_position" and isinstance(value, tuple):
+            # Convert tuple with potential numpy types
+            sanitized[key] = [convert_numpy_to_python(v) for v in value]
+        else:
+            sanitized[key] = convert_numpy_to_python(value)
+
+    return sanitized

examples/vlm/README.md

@@ -0,0 +1,68 @@
+# Crafter VLM Pipeline
+
+This folder captures the reference workflow for fine-tuning Crafter policies with
+multimodal (text + image) prompts. It stitches together the new image-aware tracing
+plumbing with lightweight utilities for dataset curation and training.
+
+## Quick Start
+
+1. **Verify image capture**  
+   ```
+   uv run python examples/vlm/crafter_image_only_agent.py --seed 7 --steps 5
+   ```
+   This writes PNG frames to `examples/vlm/output/frames/` and produces a JSONL preview
+   of OpenAI-style image-only user messages.
+
+2. **Collect traced rollouts**  
+   Use the Crafter task app (or your existing pipeline) with tracing enabled. The new
+   tracing schema automatically records `observation_image_base64` and stores image parts
+   in LM call records.
+
+3. **Export multimodal SFT rows**  
+   ```
+   uv run python examples/warming_up_to_rl/export_trace_sft.py \
+       --db traces/v3/synth_ai.db \
+       --output examples/vlm/output/crafter_traces_full.jsonl
+   ```
+   The exporter now emits `metadata.has_image`, `metadata.user_has_image`, and
+   `metadata.assistant_has_image` flags per turn.
+
+4. **Filter to image-rich turns**  
+   ```
+   uv run python examples/vlm/filter_image_rows.py \
+       --input examples/vlm/output/crafter_traces_full.jsonl \
+       --output examples/vlm/output/crafter_vlm_dataset.jsonl
+   ```
+
+5. **(Optional) Split validation or augment**, then upload using the standard CLI:
+   ```
+   uv run python examples/warming_up_to_rl/run_fft_and_save.py \
+       --toml examples/vlm/configs/crafter_vlm_gpt4o.toml \
+       --data examples/vlm/output/crafter_vlm_dataset.jsonl
+   ```
+
+## Config & Utilities
+
+| File | Purpose |
+| --- | --- |
+| `configs/crafter_vlm_gpt4o.toml` | Sample Synth job targeting an image-capable model (`openai/gpt-4o-mini`). Set `job.data` or pass `--data` explicitly. |
+| `crafter_image_only_agent.py` | Captures frames and builds image-only prompts for sanity checks. |
+| `filter_image_rows.py` | Extracts rows with image parts from exported JSONL datasets. |
+
+## Notes & Next Steps
+
+- The training config assumes full-finetuning (`mode = "sft_offline"`). Adjust the
+  model id, hardware, or hyperparameters to match available infrastructure.
+- Dataset rows emitted by `export_trace_sft.py` already contain OpenAI multimodal
+  content parts like:
+  ```json
+  {
+    "role": "user",
+    "content": [
+      {"type": "text", "text": "..."},
+      {"type": "image_url", "image_url": {"url": "data:image/png;base64,..." }}
+    ]
+  }
+  ```
+- See `PROPOSAL.md` for a deeper dive into outstanding work (longer rollouts,
+  richer multimodal augmentations, evaluation ideas).

examples/vlm/configs/crafter_vlm_gpt4o.toml

@@ -0,0 +1,44 @@
+[job]
+model = "openai/gpt-4o-mini-2024-07-18"
+modalities = ["text", "image"]
+# data = "examples/vlm/output/crafter_vlm_dataset.jsonl"
+description = "Crafter VLM SFT (text + image prompts)"
+
+[compute]
+gpu_type = "A100"
+gpu_count = 1
+nodes = 1
+
+[data]
+topology = {}
+# validation_path = "examples/vlm/output/crafter_vlm_dataset.val.jsonl"
+
+[training]
+mode = "sft_offline"
+use_qlora = false
+
+[training.validation]
+enabled = true
+evaluation_strategy = "steps"
+eval_steps = 50
+save_best_model_at_end = true
+metric_for_best_model = "val.loss"
+greater_is_better = false
+
+[hyperparameters]
+n_epochs = 1
+train_kind = "fft"
+per_device_batch = 1
+gradient_accumulation_steps = 32
+sequence_length = 4096
+learning_rate = 1e-5
+warmup_ratio = 0.03
+weight_decay = 0.01
+
+[hyperparameters.parallelism]
+use_deepspeed = true
+deepspeed_stage = 2
+fsdp = false
+bf16 = true
+fp16 = false
+activation_checkpointing = true