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

plato/worlds/README.md

@@ -28,7 +28,8 @@ Create a JSON config file:
   },
   "git_token": "ghp_...",
   "session_id": "local-test-001",
-  "callback_url": ""
+  "otel_url": "",
+  "upload_url": ""
 }
 ```
 

plato/worlds/__init__.py

@@ -52,7 +52,7 @@ from plato.worlds.base import (
     get_world,
     register_world,
 )
-from plato.worlds.config import Agent, AgentConfig, Env, EnvConfig, RunConfig, Secret
+from plato.worlds.config import Agent, AgentConfig, CheckpointConfig, Env, EnvConfig, RunConfig, Secret, StateConfig
 from plato.worlds.runner import run_world
 
 __all__ = [
@@ -66,6 +66,8 @@ __all__ = [
     "get_world",
     # Config
     "RunConfig",
+    "CheckpointConfig",
+    "StateConfig",
     "AgentConfig",
     "Agent",
     "Secret",

plato/worlds/base.py

@@ -3,7 +3,10 @@
 from __future__ import annotations
 
 import logging
+import os
+import subprocess
 from abc import ABC, abstractmethod
+from pathlib import Path
 from typing import TYPE_CHECKING, Any, ClassVar, Generic, TypeVar, get_args, get_origin
 
 from pydantic import BaseModel, Field
@@ -14,13 +17,29 @@ if TYPE_CHECKING:
     from plato.v2.async_.environment import Environment
     from plato.v2.async_.session import Session
 
-from plato.agents.logging import init_logging as _init_chronos_logging
-from plato.agents.logging import log_event as _log_event
-from plato.agents.logging import reset_logging as _reset_chronos_logging
-from plato.agents.logging import span as _span
+from plato.agents.artifacts import (
+    upload_artifact as _upload_artifact_raw,
+)
+from plato.agents.otel import (
+    get_tracer,
+    init_tracing,
+    shutdown_tracing,
+)
+from plato.agents.runner import run_agent as _run_agent_raw
 
 logger = logging.getLogger(__name__)
 
+
+def _get_plato_version() -> str:
+    """Get the installed plato SDK version."""
+    try:
+        from importlib.metadata import version
+
+        return version("plato")
+    except Exception:
+        return "unknown"
+
+
 # Global registry of worlds
 _WORLD_REGISTRY: dict[str, type[BaseWorld]] = {}
 
@@ -107,6 +126,8 @@ class BaseWorld(ABC, Generic[ConfigT]):
         self._step_count: int = 0
         self.plato_session = None
         self._current_step_id: str | None = None
+        self._session_id: str | None = None
+        self._agent_containers: list[str] = []  # Track spawned agent containers for cleanup
 
     @classmethod
     def get_config_class(cls) -> type[RunConfig]:
@@ -166,7 +187,70 @@ class BaseWorld(ABC, Generic[ConfigT]):
 
     async def close(self) -> None:
         """Cleanup resources. Called after run completes."""
-        pass
+        await self._cleanup_agent_containers()
+
+    async def _cleanup_agent_containers(self) -> None:
+        """Stop any agent containers spawned by this world."""
+        import asyncio
+
+        if not self._agent_containers:
+            return
+
+        self.logger.info(f"Stopping {len(self._agent_containers)} agent container(s)...")
+        for container_name in self._agent_containers:
+            try:
+                proc = await asyncio.create_subprocess_exec(
+                    "docker",
+                    "stop",
+                    container_name,
+                    stdout=asyncio.subprocess.DEVNULL,
+                    stderr=asyncio.subprocess.DEVNULL,
+                )
+                await proc.wait()
+                self.logger.debug(f"Stopped container: {container_name}")
+            except Exception as e:
+                self.logger.warning(f"Failed to stop container {container_name}: {e}")
+        self._agent_containers.clear()
+        self.logger.info("Agent containers stopped")
+
+    async def run_agent(
+        self,
+        image: str,
+        config: dict,
+        secrets: dict[str, str],
+        instruction: str,
+        workspace: str | None = None,
+        logs_dir: str | None = None,
+        pull: bool = True,
+    ) -> str:
+        """Run an agent in a Docker container, tracking the container for cleanup.
+
+        This is a wrapper around plato.agents.runner.run_agent that automatically
+        tracks spawned containers so they can be cleaned up when the world closes.
+
+        Args:
+            image: Docker image URI
+            config: Agent configuration dict
+            secrets: Secret values (API keys, etc.)
+            instruction: Task instruction for the agent
+            workspace: Docker volume name for workspace
+            logs_dir: Ignored (kept for backwards compatibility)
+            pull: Whether to pull the image first
+
+        Returns:
+            The container name that was created
+        """
+        container_name = await _run_agent_raw(
+            image=image,
+            config=config,
+            secrets=secrets,
+            instruction=instruction,
+            workspace=workspace,
+            logs_dir=logs_dir,
+            pull=pull,
+        )
+        self._agent_containers.append(container_name)
+        return container_name
 
     async def _connect_plato_session(self) -> None:
         """Connect to Plato session from config.
@@ -195,6 +279,258 @@ class BaseWorld(ABC, Generic[ConfigT]):
             except Exception as e:
                 self.logger.warning(f"Error stopping Plato heartbeat: {e}")
 
+    async def _create_checkpoint(self) -> dict[str, str] | None:
+        """Create a checkpoint snapshot of all environments (excluding configured envs).
+
+        Uses snapshot_store for efficient chunk-based deduplication.
+
+        Returns:
+            Dict mapping environment alias to artifact_id, or None if no session connected.
+        """
+        if not self.plato_session:
+            self.logger.warning("Cannot create checkpoint: Plato session not connected")
+            return None
+
+        exclude_envs = set(self.config.checkpoint.exclude_envs)
+        envs_to_snapshot = [env for env in self.plato_session.envs if env.alias not in exclude_envs]
+
+        if not envs_to_snapshot:
+            self.logger.info("No environments to checkpoint (all excluded)")
+            return {}
+
+        self.logger.info(
+            f"Creating checkpoint for {len(envs_to_snapshot)} environment(s): {[e.alias for e in envs_to_snapshot]}"
+        )
+
+        results: dict[str, str] = {}
+        for env in envs_to_snapshot:
+            try:
+                result = await env.snapshot_store()
+                artifact_id = result.artifact_id
+                results[env.alias] = artifact_id
+
+                # Check for success/error fields (available after SDK regeneration)
+                success = getattr(result, "success", True)
+                error = getattr(result, "error", None)
+
+                if not success or error:
+                    self.logger.error(
+                        f"Checkpoint failed for '{env.alias}': {error or 'unknown error'} (job_id={env.job_id})"
+                    )
+                elif artifact_id:
+                    self.logger.info(f"Checkpoint created for '{env.alias}': {artifact_id}")
+                else:
+                    self.logger.warning(
+                        f"Checkpoint for '{env.alias}' returned empty artifact_id (job_id={env.job_id})"
+                    )
+            except Exception as e:
+                self.logger.error(f"Failed to checkpoint '{env.alias}': {e}")
+
+        return results
+
+    def _init_state_directory(self) -> None:
+        """Initialize the state directory as a git repository.
+
+        Creates the state directory if it doesn't exist and initializes it
+        as a git repository with an initial commit.
+        """
+        if not self.config.state.enabled:
+            return
+
+        state_path = Path(self.config.state.path)
+
+        # Create directory if it doesn't exist
+        if not state_path.exists():
+            state_path.mkdir(parents=True)
+            self.logger.info(f"Created state directory: {state_path}")
+
+        # Check if already a git repo
+        git_dir = state_path / ".git"
+        if git_dir.exists():
+            self.logger.info(f"State directory already initialized: {state_path}")
+            return
+
+        # Initialize git repo
+        try:
+            subprocess.run(
+                ["git", "init"],
+                cwd=state_path,
+                capture_output=True,
+                check=True,
+            )
+            # Create initial commit (even if empty)
+            subprocess.run(
+                ["git", "config", "user.email", "plato@plato.so"],
+                cwd=state_path,
+                capture_output=True,
+                check=True,
+            )
+            subprocess.run(
+                ["git", "config", "user.name", "Plato"],
+                cwd=state_path,
+                capture_output=True,
+                check=True,
+            )
+            # Add all files and create initial commit
+            subprocess.run(
+                ["git", "add", "-A"],
+                cwd=state_path,
+                capture_output=True,
+                check=True,
+            )
+            subprocess.run(
+                ["git", "commit", "--allow-empty", "-m", "Initial state"],
+                cwd=state_path,
+                capture_output=True,
+                check=True,
+            )
+            self.logger.info(f"Initialized git repo in state directory: {state_path}")
+        except subprocess.CalledProcessError as e:
+            self.logger.warning(f"Failed to initialize state git repo: {e.stderr}")
+
+    def _commit_state(self, message: str) -> bool:
+        """Commit current state directory changes.
+
+        Args:
+            message: Commit message
+
+        Returns:
+            True if commit was created (or no changes), False on error.
+        """
+        if not self.config.state.enabled:
+            return True
+
+        state_path = Path(self.config.state.path)
+        if not state_path.exists():
+            return True
+
+        try:
+            # Add all changes
+            subprocess.run(
+                ["git", "add", "-A"],
+                cwd=state_path,
+                capture_output=True,
+                check=True,
+            )
+            # Check if there are changes to commit
+            result = subprocess.run(
+                ["git", "status", "--porcelain"],
+                cwd=state_path,
+                capture_output=True,
+                text=True,
+                check=True,
+            )
+            if not result.stdout.strip():
+                self.logger.debug("No state changes to commit")
+                return True
+
+            # Commit changes
+            subprocess.run(
+                ["git", "commit", "-m", message],
+                cwd=state_path,
+                capture_output=True,
+                check=True,
+            )
+            self.logger.info(f"Committed state changes: {message}")
+            return True
+        except subprocess.CalledProcessError as e:
+            self.logger.warning(f"Failed to commit state: {e.stderr}")
+            return False
+
+    def _create_state_bundle(self) -> bytes | None:
+        """Create a git bundle of the state directory.
+
+        Returns:
+            Bundle bytes if successful, None otherwise.
+        """
+        if not self.config.state.enabled:
+            return None
+
+        state_path = Path(self.config.state.path)
+        if not state_path.exists():
+            return None
+
+        git_dir = state_path / ".git"
+        if not git_dir.exists():
+            self.logger.warning("State directory is not a git repository")
+            return None
+
+        try:
+            # Create bundle to stdout
+            result = subprocess.run(
+                ["git", "bundle", "create", "-", "--all"],
+                cwd=state_path,
+                capture_output=True,
+                check=True,
+            )
+            bundle_data = result.stdout
+            self.logger.info(f"Created state bundle: {len(bundle_data)} bytes")
+            return bundle_data
+        except subprocess.CalledProcessError as e:
+            self.logger.warning(f"Failed to create state bundle: {e.stderr}")
+            return None
+
+    async def _upload_artifact(
+        self,
+        data: bytes,
+        content_type: str = "application/octet-stream",
+    ) -> bool:
+        """Upload an artifact directly to S3.
+
+        Args:
+            data: Raw bytes of the artifact
+            content_type: MIME type of the content
+
+        Returns:
+            True if successful, False otherwise
+        """
+        if not self.config.upload_url:
+            self.logger.warning("Cannot upload artifact: upload_url not set")
+            return False
+        return await _upload_artifact_raw(
+            upload_url=self.config.upload_url,
+            data=data,
+            content_type=content_type,
+        )
+
+    async def _create_and_upload_checkpoint(self) -> tuple[dict[str, str], bool]:
+        """Create a full checkpoint including env snapshots and state bundle.
+
+        This method:
+        1. Commits any pending state changes
+        2. Creates env snapshots using snapshot_store
+        3. Creates and uploads state bundle to S3
+
+        Returns:
+            Tuple of (env_snapshots dict, state_bundle_uploaded bool)
+        """
+        # Commit state changes first
+        self._commit_state(f"Checkpoint at step {self._step_count}")
+
+        # Create env snapshots
+        env_snapshots = await self._create_checkpoint()
+        if env_snapshots is None:
+            env_snapshots = {}
+
+        state_bundle_uploaded = True  # Default to True if state not enabled
+
+        # Create and upload state bundle
+        if self.config.state.enabled:
+            bundle_data = self._create_state_bundle()
+            if bundle_data:
+                success = await self._upload_artifact(
+                    data=bundle_data,
+                    content_type="application/octet-stream",
+                )
+                if success:
+                    self.logger.info(f"Uploaded state bundle at step {self._step_count}")
+                    state_bundle_uploaded = True
+                else:
+                    self.logger.warning(f"Failed to upload state bundle at step {self._step_count}")
+                    state_bundle_uploaded = False
+
+        return env_snapshots, state_bundle_uploaded
+
     def get_env(self, alias: str) -> Environment | None:
         """Get an environment by alias.
 
@@ -236,6 +572,9 @@ class BaseWorld(ABC, Generic[ConfigT]):
         Returns:
             Dict of environment variable name -> value
 
+        Raises:
+            ImportError: If a sim environment is configured but package is not installed.
+
         Example:
             env_vars = self.get_sim_env_vars()
             # Returns: {"AWS_ENDPOINT_URL": "https://...", "GITEA_URL": "https://...", ...}
@@ -263,7 +602,13 @@ class BaseWorld(ABC, Generic[ConfigT]):
                 env_vars.update(sim_vars)
                 self.logger.info(f"{package_name} env vars: {list(sim_vars.keys())}")
             except ImportError:
-                self.logger.debug(f"{package_name} sim package not installed, skipping")
+                raise ImportError(
+                    f"Environment '{env_alias}' is configured but 'plato.sims.{package_name}' "
+                    f"package is not installed.\n\n"
+                    f"Install sims packages:\n"
+                    f'  export INDEX_URL="https://__token__:${{PLATO_API_KEY}}@plato.so/api/v2/pypi/sims/simple/"\n'
+                    f"  uv pip install '.[sims]' --extra-index-url $INDEX_URL"
+                ) from None
             except Exception as e:
                 self.logger.warning(f"Failed to get {package_name} env vars: {e}")
 
@@ -278,6 +623,9 @@ class BaseWorld(ABC, Generic[ConfigT]):
         Returns:
             Markdown string with instructions, or empty string if no sims configured.
 
+        Raises:
+            ImportError: If a sim environment is configured but package is not installed.
+
         Example:
             instructions = self.get_sim_instructions()
             # Returns markdown with LocalStack/Gitea setup instructions
@@ -306,7 +654,13 @@ class BaseWorld(ABC, Generic[ConfigT]):
                         instructions_parts.append(instructions)
                         self.logger.info(f"Added {package_name} instructions to prompt")
             except ImportError:
-                self.logger.debug(f"{package_name} sim package not installed, skipping instructions")
+                raise ImportError(
+                    f"Environment '{env_alias}' is configured but 'plato.sims.{package_name}' "
+                    f"package is not installed.\n\n"
+                    f"Install sims packages:\n"
+                    f'  export INDEX_URL="https://__token__:${{PLATO_API_KEY}}@plato.so/api/v2/pypi/sims/simple/"\n'
+                    f"  uv pip install '.[sims]' --extra-index-url $INDEX_URL"
+                ) from None
             except Exception as e:
                 self.logger.warning(f"Failed to get {package_name} instructions: {e}")
 
@@ -363,74 +717,115 @@ The following services are available for your use:
 
         self.logger.info(f"Starting world '{self.name}'")
 
-        # Initialize the logging singleton for agents to use
-        if config.callback_url and config.session_id:
-            _init_chronos_logging(
-                callback_url=config.callback_url,
-                session_id=config.session_id,
-            )
+        # Initialize state directory (creates git repo if needed)
+        self._init_state_directory()
+
+        # Initialize OTel tracing and session info for artifact uploads
+        if config.session_id:
+            self._session_id = config.session_id
+
+            # Set environment variables for agent runners (which run in Docker)
+            os.environ["SESSION_ID"] = config.session_id
+            if config.otel_url:
+                # For agents in Docker, convert localhost to host.docker.internal
+                # so they can reach the host machine's Chronos instance
+                agent_otel_url = config.otel_url
+                if "localhost" in agent_otel_url or "127.0.0.1" in agent_otel_url:
+                    agent_otel_url = agent_otel_url.replace("localhost", "host.docker.internal")
+                    agent_otel_url = agent_otel_url.replace("127.0.0.1", "host.docker.internal")
+                os.environ["OTEL_EXPORTER_OTLP_ENDPOINT"] = agent_otel_url
+                os.environ["OTEL_EXPORTER_OTLP_PROTOCOL"] = "http/protobuf"
+            if config.upload_url:
+                os.environ["UPLOAD_URL"] = config.upload_url
+
+            # Initialize OTel tracing for the world itself (runs on host, not in Docker)
+            if config.otel_url:
+                logger.debug(f"Initializing OTel tracing with endpoint: {config.otel_url}")
+                init_tracing(
+                    service_name=f"world-{self.name}",
+                    session_id=config.session_id,
+                    otlp_endpoint=config.otel_url,
+                )
+            else:
+                logger.debug("No otel_url in config - OTel tracing disabled")
+
+        # Log version info (goes to OTel after init_tracing)
+        plato_version = _get_plato_version()
+        world_version = self.get_version()
+        self.logger.info(f"World version: {world_version}, Plato SDK version: {plato_version}")
 
         # Connect to Plato session if configured (for heartbeats)
         await self._connect_plato_session()
 
-        # Log session start
-        await _log_event(
-            span_type="session_start",
-            content=f"World '{self.name}' started",
-            source="world",
-            extra={"world_name": self.name, "world_version": self.get_version()},
-        )
+        # Get tracer for spans
+        tracer = get_tracer("plato.world")
 
-        try:
-            # Execute reset with automatic span tracking
-            async with _span("reset", span_type="reset", source="world") as reset_span:
-                reset_span.log(f"Resetting world '{self.name}'")
-                obs = await self.reset()
-                reset_span.set_extra({"observation": obs.model_dump() if hasattr(obs, "model_dump") else str(obs)})
-            self.logger.info(f"World reset complete: {obs}")
-
-            while True:
-                self._step_count += 1
-
-                # Execute step with automatic span tracking
-                # The span automatically sets itself as the current parent,
-                # so agent trajectories will nest under this step
-                async with _span(
-                    f"step_{self._step_count}",
-                    span_type="step",
-                    source="world",
-                ) as step_span:
-                    self._current_step_id = step_span.event_id
-                    step_span.log(f"Step {self._step_count} started")
-                    result = await self.step()
-                    step_span.set_extra(
-                        {
-                            "done": result.done,
-                            "observation": result.observation.model_dump()
-                            if hasattr(result.observation, "model_dump")
-                            else str(result.observation),
-                            "info": result.info,
-                        }
-                    )
+        # Create root session span that encompasses everything
+        # This ensures all child spans share the same trace_id
+        with tracer.start_as_current_span("session") as session_span:
+            session_span.set_attribute("plato.world.name", self.name)
+            session_span.set_attribute("plato.world.version", self.get_version())
+            session_span.set_attribute("plato.session.id", config.session_id)
 
-                self.logger.info(f"Step {self._step_count}: done={result.done}")
+            try:
+                # Execute reset with OTel span
+                with tracer.start_as_current_span("reset") as reset_span:
+                    obs = await self.reset()
+                    obs_data = obs.model_dump() if hasattr(obs, "model_dump") else str(obs)
+                    reset_span.set_attribute("plato.observation", str(obs_data)[:1000])
+                self.logger.info(f"World reset complete: {obs}")
 
-                if result.done:
-                    break
+                while True:
+                    self._step_count += 1
 
-        finally:
-            await self.close()
-            await self._disconnect_plato_session()
+                    # Execute step with OTel span
+                    with tracer.start_as_current_span(f"step_{self._step_count}") as step_span:
+                        step_span.set_attribute("plato.step.number", self._step_count)
 
-            # Log session end
-            await _log_event(
-                span_type="session_end",
-                content=f"World '{self.name}' completed after {self._step_count} steps",
-                source="world",
-                extra={"total_steps": self._step_count},
-            )
+                        # Store span context for nested agent spans
+                        self._current_step_id = format(step_span.get_span_context().span_id, "016x")
 
-            # Reset the logging singleton
-            _reset_chronos_logging()
+                        result = await self.step()
 
-            self.logger.info(f"World '{self.name}' completed after {self._step_count} steps")
+                        step_span.set_attribute("plato.step.done", result.done)
+                        obs_data = (
+                            result.observation.model_dump()
+                            if hasattr(result.observation, "model_dump")
+                            else str(result.observation)
+                        )
+                        step_span.set_attribute("plato.step.observation", str(obs_data)[:1000])
+
+                    self.logger.info(f"Step {self._step_count}: done={result.done}")
+
+                    # Create checkpoint if enabled and interval matches
+                    if self.config.checkpoint.enabled and self._step_count % self.config.checkpoint.interval == 0:
+                        self.logger.info(f"Creating checkpoint after step {self._step_count}")
+                        with tracer.start_as_current_span("checkpoint") as checkpoint_span:
+                            checkpoint_span.set_attribute("plato.checkpoint.step", self._step_count)
+                            env_snapshots, state_bundle_uploaded = await self._create_and_upload_checkpoint()
+
+                            checkpoint_span.set_attribute("plato.checkpoint.success", len(env_snapshots) > 0)
+                            checkpoint_span.set_attribute(
+                                "plato.checkpoint.state_bundle_uploaded", state_bundle_uploaded
+                            )
+
+                            if env_snapshots:
+                                checkpoint_span.set_attribute(
+                                    "plato.checkpoint.environments", list(env_snapshots.keys())
+                                )
+                                checkpoint_span.set_attribute(
+                                    "plato.checkpoint.artifact_ids", list(env_snapshots.values())
+                                )
+
+                    if result.done:
+                        break
+
+            finally:
+                await self.close()
+                await self._disconnect_plato_session()
+
+        # Shutdown OTel tracing and clear session info (outside the span)
+        shutdown_tracing()
+        self._session_id = None
+
+        self.logger.info(f"World '{self.name}' completed after {self._step_count} steps")