plato-sdk-v2 2.3.0__py3-none-any.whl → 2.4.1__py3-none-any.whl

plato/agents/runner.py

@@ -1,16 +1,23 @@
-"""Agent runner - run agents in Docker containers."""
+"""Agent runner - run agents in Docker containers.
+
+Agents emit their own OTel spans for trajectory events. This runner:
+1. Runs agents in Docker containers
+2. Streams stdout/stderr for logging
+3. Passes OTel environment variables for trace context propagation
+4. Uploads artifacts to S3 when complete
+"""
 
 from __future__ import annotations
 
 import asyncio
+import base64
 import json
 import logging
 import os
 import platform
-import tempfile
-from pathlib import Path
+import uuid
 
-from plato.agents.logging import log_event, span, upload_artifacts
+from opentelemetry import trace
 
 logger = logging.getLogger(__name__)
 
@@ -20,10 +27,10 @@ async def run_agent(
     config: dict,
     secrets: dict[str, str],
     instruction: str,
-    workspace: str,
+    workspace: str | None = None,
     logs_dir: str | None = None,
     pull: bool = True,
-) -> None:
+) -> str:
     """Run an agent in a Docker container.
 
     Args:
@@ -31,131 +38,228 @@ async def run_agent(
         config: Agent configuration dict
         secrets: Secret values (API keys, etc.)
         instruction: Task instruction for the agent
-        workspace: Host directory to mount as /workspace
-        logs_dir: Host directory for logs (temp dir if None)
+        workspace: Docker volume name for workspace (created if None)
+        logs_dir: Ignored (kept for backwards compatibility)
         pull: Whether to pull the image first
+
+    Returns:
+        The container name that was created (for cleanup purposes)
+
+    Note: Agents handle their own OTel tracing. This runner only passes
+    the trace context (TRACEPARENT) so agent spans link to the parent step.
+
+    Note: This uses Docker volumes (not bind mounts) for DIND compatibility.
+    The workspace parameter should be a Docker volume name.
     """
-    logs_dir = logs_dir or tempfile.mkdtemp(prefix="agent_logs_")
-    agent_name = image.split("/")[-1].split(":")[0]
-
-    async with span(agent_name, span_type="agent", source="agent") as agent_span:
-        agent_span.log(f"Starting agent: {agent_name} ({image})")
-
-        # Pull image if requested
-        if pull:
-            agent_span.log(f"Pulling image: {image}")
-            pull_proc = await asyncio.create_subprocess_exec(
-                "docker",
-                "pull",
-                image,
-                stdout=asyncio.subprocess.PIPE,
-                stderr=asyncio.subprocess.STDOUT,
-            )
-            await pull_proc.wait()
-
-        # Setup
-        os.makedirs(os.path.join(logs_dir, "agent"), exist_ok=True)
-        config_file = tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False)
-        json.dump(config, config_file)
-        config_file.close()
-
-        try:
-            # Build docker command
-            docker_cmd = ["docker", "run", "--rm"]
-
-            # Determine if we need host networking:
-            # - Required on Linux without iptables for connectivity
-            # - Skip on macOS where --network=host doesn't work properly
-            use_host_network = False
-            is_macos = platform.system() == "Darwin"
-
-            if not is_macos:
-                try:
-                    proc = await asyncio.create_subprocess_exec(
-                        "iptables",
-                        "-L",
-                        "-n",
-                        stdout=asyncio.subprocess.DEVNULL,
-                        stderr=asyncio.subprocess.DEVNULL,
-                    )
-                    await proc.wait()
-                    has_iptables = proc.returncode == 0
-                except (FileNotFoundError, PermissionError):
-                    has_iptables = False
-
-                use_host_network = not has_iptables
-
-            if use_host_network:
-                docker_cmd.extend(["--network=host", "--add-host=localhost:127.0.0.1"])
+    # Get session info from environment variables
+    session_id = os.environ.get("SESSION_ID")
+    otel_url = os.environ.get("OTEL_EXPORTER_OTLP_ENDPOINT")
+    upload_url = os.environ.get("UPLOAD_URL")
+
+    # Pull image if requested
+    if pull:
+        pull_proc = await asyncio.create_subprocess_exec(
+            "docker",
+            "pull",
+            image,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.STDOUT,
+        )
+        await pull_proc.wait()
+
+    # Encode config as base64 to pass via environment variable
+    # This avoids file mount issues in Docker-in-Docker scenarios
+    config_json = json.dumps(config)
+    config_b64 = base64.b64encode(config_json.encode()).decode()
+
+    # Generate a unique container name for inspection
+    container_name = f"agent-{uuid.uuid4().hex[:8]}"
+
+    # Use WORKSPACE_VOLUME env var if set (for DIND compatibility)
+    # Otherwise create a new volume
+    workspace_volume = os.environ.get("WORKSPACE_VOLUME") or workspace or f"workspace-{uuid.uuid4().hex[:8]}"
+    if not os.environ.get("WORKSPACE_VOLUME") and not workspace:
+        await asyncio.create_subprocess_exec(
+            "docker",
+            "volume",
+            "create",
+            workspace_volume,
+            stdout=asyncio.subprocess.DEVNULL,
+            stderr=asyncio.subprocess.DEVNULL,
+        )
+
+    # Create logs volume
+    logs_volume = f"logs-{uuid.uuid4().hex[:8]}"
+    await asyncio.create_subprocess_exec(
+        "docker",
+        "volume",
+        "create",
+        logs_volume,
+        stdout=asyncio.subprocess.DEVNULL,
+        stderr=asyncio.subprocess.DEVNULL,
+    )
+
+    try:
+        # Build docker command
+        docker_cmd = ["docker", "run", "--rm", "--privileged", "--name", container_name]
+
+        # Determine if we need host networking
+        use_host_network = False
+        is_macos = platform.system() == "Darwin"
+
+        if not is_macos:
+            try:
+                proc = await asyncio.create_subprocess_exec(
+                    "iptables",
+                    "-L",
+                    "-n",
+                    stdout=asyncio.subprocess.DEVNULL,
+                    stderr=asyncio.subprocess.DEVNULL,
+                )
+                await proc.wait()
+                has_iptables = proc.returncode == 0
+            except (FileNotFoundError, PermissionError):
+                has_iptables = False
+
+            use_host_network = not has_iptables
 
+        if use_host_network:
+            docker_cmd.extend(["--network=host", "--add-host=localhost:127.0.0.1"])
+
+        # Use Docker volumes instead of bind mounts for DIND compatibility
+        docker_cmd.extend(
+            [
+                "-v",
+                f"{workspace_volume}:/workspace",
+                "-v",
+                f"{logs_volume}:/logs",
+                "-v",
+                "/var/run/docker.sock:/var/run/docker.sock",
+                "-w",
+                "/workspace",
+                "-e",
+                f"AGENT_CONFIG_B64={config_b64}",
+            ]
+        )
+
+        # Pass session info to agent
+        if otel_url:
+            traces_endpoint = f"{otel_url.rstrip('/')}/v1/traces"
+            docker_cmd.extend(["-e", f"OTEL_EXPORTER_OTLP_ENDPOINT={otel_url}"])
+            docker_cmd.extend(["-e", f"OTEL_EXPORTER_OTLP_TRACES_ENDPOINT={traces_endpoint}"])
+            docker_cmd.extend(["-e", "OTEL_EXPORTER_OTLP_PROTOCOL=http/protobuf"])
+        if session_id:
+            docker_cmd.extend(["-e", f"SESSION_ID={session_id}"])
+        if upload_url:
+            docker_cmd.extend(["-e", f"UPLOAD_URL={upload_url}"])
+
+        # Pass trace context to agent for parent linking
+        # Agent spans will be children of the current step span
+        current_span = trace.get_current_span()
+        span_context = current_span.get_span_context()
+        if span_context.is_valid:
+            trace_id = format(span_context.trace_id, "032x")
+            span_id = format(span_context.span_id, "016x")
+            # W3C Trace Context format for TRACEPARENT
+            traceparent = f"00-{trace_id}-{span_id}-01"
             docker_cmd.extend(
                 [
-                    "-v",
-                    f"{workspace}:/workspace",
-                    "-v",
-                    f"{logs_dir}:/logs",
-                    "-v",
-                    f"{config_file.name}:/config.json:ro",
-                    "-w",
-                    "/workspace",
+                    "-e",
+                    f"TRACEPARENT={traceparent}",
+                    "-e",
+                    f"OTEL_TRACE_ID={trace_id}",
+                    "-e",
+                    f"OTEL_PARENT_SPAN_ID={span_id}",
                 ]
             )
 
-            for key, value in secrets.items():
-                docker_cmd.extend(["-e", f"{key.upper()}={value}"])
+        for key, value in secrets.items():
+            docker_cmd.extend(["-e", f"{key.upper()}={value}"])
 
-            docker_cmd.append(image)
+        docker_cmd.append(image)
 
-            # Pass instruction via CLI arg (agents expect --instruction flag)
-            docker_cmd.extend(["--instruction", instruction])
+        # Pass instruction via CLI arg
+        docker_cmd.extend(["--instruction", instruction])
 
-            # Run container and stream output
-            process = await asyncio.create_subprocess_exec(
-                *docker_cmd,
-                stdout=asyncio.subprocess.PIPE,
-                stderr=asyncio.subprocess.STDOUT,
-            )
+        logger.info(f"Starting container: {container_name}")
+
+        # Run container - agents emit their own OTel spans
+        # Use large limit to handle agents that output long lines (e.g., JSON with file contents)
+        process = await asyncio.create_subprocess_exec(
+            *docker_cmd,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.STDOUT,
+            limit=100 * 1024 * 1024,  # 100MB buffer limit
+        )
+
+        # Get and print container IP in background
+        async def print_container_ip():
+            await asyncio.sleep(3)  # Wait for container to start
+            try:
+                inspect_proc = await asyncio.create_subprocess_exec(
+                    "docker",
+                    "inspect",
+                    "-f",
+                    "{{.NetworkSettings.IPAddress}}",
+                    container_name,
+                    stdout=asyncio.subprocess.PIPE,
+                    stderr=asyncio.subprocess.PIPE,
+                )
+                stdout, _ = await inspect_proc.communicate()
+                container_ip = stdout.decode().strip()
+                if container_ip:
+                    logger.info("=" * 50)
+                    logger.info(f"Container: {container_name}")
+                    logger.info(f"Container IP: {container_ip}")
+                    logger.info(f"noVNC: http://{container_ip}:6080")
+                    logger.info("=" * 50)
+            except Exception:
+                pass
+
+        asyncio.create_task(print_container_ip())
+
+        # Stream and capture output for error reporting using chunked reads to handle large lines
+        output_lines: list[str] = []
+        assert process.stdout is not None
+        buffer = ""
+        while True:
+            try:
+                chunk = await process.stdout.read(65536)
+            except Exception:
+                break
+            if not chunk:
+                break
+            buffer += chunk.decode(errors="replace")
+
+            while "\n" in buffer:
+                line, buffer = buffer.split("\n", 1)
+                output_lines.append(line)
+                # Print agent output in real-time
+                print(f"[agent] {line}")
+
+        # Handle any remaining content in buffer
+        if buffer.strip():
+            output_lines.append(buffer)
+            print(f"[agent] {buffer}")
+
+        await process.wait()
+
+        exit_code = process.returncode or 0
+        if exit_code != 0:
+            error_context = "\n".join(output_lines[-50:]) if output_lines else "No output captured"
+            raise RuntimeError(f"Agent failed with exit code {exit_code}\n\nAgent output:\n{error_context}")
+
+    finally:
+        # Clean up volumes
+        await asyncio.create_subprocess_exec(
+            "docker",
+            "volume",
+            "rm",
+            "-f",
+            logs_volume,
+            stdout=asyncio.subprocess.DEVNULL,
+            stderr=asyncio.subprocess.DEVNULL,
+        )
+        # Note: workspace_volume is not cleaned up as it may be shared
 
-            # Stream output line by line
-            assert process.stdout is not None
-            while True:
-                line = await process.stdout.readline()
-                if not line:
-                    break
-                logger.info(f"[agent] {line.decode().rstrip()}")
-
-            await process.wait()
-
-            if process.returncode != 0:
-                raise RuntimeError(f"Agent failed with exit code {process.returncode}")
-
-            agent_span.log("Agent completed successfully")
-
-        finally:
-            os.unlink(config_file.name)
-
-            # Load trajectory and add to span
-            trajectory_path = Path(logs_dir) / "agent" / "trajectory.json"
-            if trajectory_path.exists():
-                try:
-                    with open(trajectory_path) as f:
-                        trajectory = json.load(f)
-                    if isinstance(trajectory, dict) and "schema_version" in trajectory:
-                        # Add agent image to trajectory
-                        agent_data = trajectory.get("agent", {})
-                        extra = agent_data.get("extra") or {}
-                        extra["image"] = image
-                        agent_data["extra"] = extra
-                        trajectory["agent"] = agent_data
-
-                        # Log trajectory as separate event
-                        await log_event(
-                            span_type="trajectory",
-                            log_type="atif",
-                            extra=trajectory,
-                            source="agent",
-                        )
-                except Exception as e:
-                    logger.warning(f"Failed to load trajectory: {e}")
-
-            await upload_artifacts(logs_dir)
+    return container_name

plato/chronos/models/__init__.py

@@ -10,10 +10,18 @@ from pydantic import AwareDatetime, BaseModel, ConfigDict, Field
 
 
 class AgentConfig(BaseModel):
+    """Agent config - supports multiple formats.
+
+    New format: agent + version (version optional, defaults to latest)
+    Legacy format: agent_id (public_id)
+    """
+
     model_config = ConfigDict(
         extra="allow",
     )
-    agent_id: Annotated[str, Field(title="Agent Id")]
+    agent: Annotated[str | None, Field(title="Agent")] = None
+    version: Annotated[str | None, Field(title="Version")] = None
+    agent_id: Annotated[str | None, Field(title="Agent Id")] = None  # backwards compat
     config: Annotated[dict[str, Any] | None, Field(title="Config")] = {}