plato-sdk-v2 2.3.1__tar.gz → 2.3.2__tar.gz

{plato_sdk_v2-2.3.1 → plato_sdk_v2-2.3.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: plato-sdk-v2
-Version: 2.3.1
+Version: 2.3.2
 Summary: Python SDK for the Plato API
 Author-email: Plato <support@plato.so>
 License-Expression: MIT
@@ -44,6 +44,7 @@ Description-Content-Type: text/markdown
 
 # Plato Python SDK
 
+
 Python SDK for the Plato platform. Uses [Harbor](https://harborframework.com) for agent execution.
 
 ## Installation

{plato_sdk_v2-2.3.1 → plato_sdk_v2-2.3.2}/README.md

@@ -1,5 +1,6 @@
 # Plato Python SDK
 
+
 Python SDK for the Plato platform. Uses [Harbor](https://harborframework.com) for agent execution.
 
 ## Installation

plato_sdk_v2-2.3.2/VERSION

@@ -0,0 +1 @@
+2.3.2

{plato_sdk_v2-2.3.1 → plato_sdk_v2-2.3.2}/plato/agents/__init__.py

@@ -85,6 +85,8 @@ __all__ = [
     "span",
     "log_event",
     "upload_artifacts",
+    "upload_artifact",
+    "upload_checkpoint",
     "reset_logging",
 ]
 
@@ -102,7 +104,9 @@ from plato.agents.logging import (
     log_event,
     reset_logging,
     span,
+    upload_artifact,
     upload_artifacts,
+    upload_checkpoint,
 )
 from plato.agents.runner import run_agent
 from plato.agents.trajectory import (

{plato_sdk_v2-2.3.1 → plato_sdk_v2-2.3.2}/plato/agents/logging.py

@@ -399,3 +399,117 @@ async def upload_artifacts(dir_path: str) -> str | None:
     except Exception as e:
         logger.warning(f"Failed to upload artifacts: {e}")
         return None
+
+
+# =============================================================================
+# Artifact Upload (Generic)
+# =============================================================================
+
+
+async def upload_artifact(
+    data: bytes,
+    artifact_type: str,
+    filename: str | None = None,
+    extra: dict[str, Any] | None = None,
+) -> dict[str, Any] | None:
+    """Upload an artifact to Chronos.
+
+    Artifacts are stored in S3 and linked to the session in the database.
+
+    Args:
+        data: Raw bytes of the artifact
+        artifact_type: Type of artifact (e.g., "state", "logs", "trajectory")
+        filename: Optional filename for the artifact
+        extra: Optional extra data to store with the artifact
+
+    Returns:
+        Dict with artifact_id and s3_url if successful, None otherwise.
+    """
+    chronos = _ChronosLogger._instance
+    if not chronos or not chronos.enabled:
+        return None
+
+    try:
+        data_base64 = base64.b64encode(data).decode("utf-8")
+        logger.info(f"Uploading artifact: type={artifact_type}, size={len(data)} bytes")
+    except Exception as e:
+        logger.warning(f"Failed to encode artifact: {e}")
+        return None
+
+    try:
+        async with httpx.AsyncClient(timeout=120.0) as client:
+            response = await client.post(
+                f"{chronos.callback_url}/artifact",
+                json={
+                    "session_id": chronos.session_id,
+                    "artifact_type": artifact_type,
+                    "data_base64": data_base64,
+                    "filename": filename,
+                    "extra": extra or {},
+                },
+            )
+            if response.status_code == 200:
+                result = response.json()
+                logger.info(f"Uploaded artifact: {result}")
+                return result
+            else:
+                logger.warning(f"Failed to upload artifact: {response.status_code} {response.text}")
+                return None
+    except Exception as e:
+        logger.warning(f"Failed to upload artifact: {e}")
+        return None
+
+
+# =============================================================================
+# Checkpoint Upload
+# =============================================================================
+
+
+async def upload_checkpoint(
+    step_number: int,
+    env_snapshots: dict[str, str],
+    state_artifact_id: str | None = None,
+    extra: dict[str, Any] | None = None,
+) -> dict[str, Any] | None:
+    """Upload checkpoint data to Chronos.
+
+    A checkpoint includes:
+    - Environment snapshots (artifact IDs per env alias)
+    - State artifact (git bundle of /state directory)
+    - Extra data (step number, timestamp, etc.)
+
+    Args:
+        step_number: The step number when this checkpoint was created
+        env_snapshots: Dict mapping env alias to artifact_id
+        state_artifact_id: Artifact ID of the state bundle (from upload_artifact)
+        extra: Optional additional data
+
+    Returns:
+        Dict with checkpoint_id if successful, None otherwise.
+    """
+    chronos = _ChronosLogger._instance
+    if not chronos or not chronos.enabled:
+        return None
+
+    try:
+        async with httpx.AsyncClient(timeout=60.0) as client:
+            response = await client.post(
+                f"{chronos.callback_url}/checkpoint",
+                json={
+                    "session_id": chronos.session_id,
+                    "step_number": step_number,
+                    "env_snapshots": env_snapshots,
+                    "state_artifact_id": state_artifact_id,
+                    "extra": extra or {},
+                },
+            )
+            if response.status_code == 200:
+                result = response.json()
+                logger.info(f"Uploaded checkpoint: step={step_number}, checkpoint_id={result.get('checkpoint_id')}")
+                return result
+            else:
+                logger.warning(f"Failed to upload checkpoint: {response.status_code} {response.text}")
+                return None
+    except Exception as e:
+        logger.warning(f"Failed to upload checkpoint: {e}")
+        return None

{plato_sdk_v2-2.3.1 → plato_sdk_v2-2.3.2}/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_sdk_v2-2.3.1 → plato_sdk_v2-2.3.2}/plato/worlds/base.py

@@ -3,7 +3,9 @@
 from __future__ import annotations
 
 import logging
+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
@@ -18,6 +20,8 @@ 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.logging import upload_artifact as _upload_artifact
+from plato.agents.logging import upload_checkpoint as _upload_checkpoint
 
 logger = logging.getLogger(__name__)
 
@@ -195,6 +199,248 @@ 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 _create_and_upload_checkpoint(self) -> dict[str, Any] | None:
+        """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 as an artifact
+        4. Calls the checkpoint endpoint with all data
+
+        Returns:
+            Checkpoint result dict if successful, None otherwise.
+        """
+        # 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 = {}
+
+        # Create and upload state bundle
+        state_artifact_id: str | None = None
+        if self.config.state.enabled:
+            bundle_data = self._create_state_bundle()
+            if bundle_data:
+                result = await _upload_artifact(
+                    data=bundle_data,
+                    artifact_type="state",
+                    filename=f"state_step_{self._step_count}.bundle",
+                    extra={
+                        "step_number": self._step_count,
+                        "state_path": self.config.state.path,
+                    },
+                )
+                if result:
+                    state_artifact_id = result.get("artifact_id")
+                    self.logger.info(f"Uploaded state artifact: {state_artifact_id}")
+
+        # Upload checkpoint with all data
+        checkpoint_result = await _upload_checkpoint(
+            step_number=self._step_count,
+            env_snapshots=env_snapshots,
+            state_artifact_id=state_artifact_id,
+            extra={
+                "world_name": self.name,
+                "world_version": self.get_version(),
+            },
+        )
+
+        return checkpoint_result
+
     def get_env(self, alias: str) -> Environment | None:
         """Get an environment by alias.
 
@@ -236,6 +482,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 +512,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 +533,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 +564,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,6 +627,9 @@ The following services are available for your use:
 
         self.logger.info(f"Starting world '{self.name}'")
 
+        # Initialize state directory (creates git repo if needed)
+        self._init_state_directory()
+
         # Initialize the logging singleton for agents to use
         if config.callback_url and config.session_id:
             _init_chronos_logging(
@@ -415,6 +682,13 @@ The following services are available for your use:
 
                 self.logger.info(f"Step {self._step_count}: done={result.done}")
 
+                # Create checkpoint if enabled and interval matches
+                # Note: The checkpoint event is created by the callback endpoint,
+                # so we don't need a span wrapper here (would create duplicates)
+                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}")
+                    await self._create_and_upload_checkpoint()
+
                 if result.done:
                     break
 

{plato_sdk_v2-2.3.1 → plato_sdk_v2-2.3.2}/plato/worlds/config.py

@@ -72,6 +72,36 @@ class Env:
         self.required = required
 
 
+class StateConfig(BaseModel):
+    """Configuration for world state persistence.
+
+    The state directory is a git-tracked directory that persists across checkpoints.
+    At each checkpoint, the state directory is git bundled and uploaded as an artifact.
+    On restore, bootstrap.sh downloads and unbundles the state before the world starts.
+
+    Attributes:
+        enabled: Whether to enable state persistence (default: True).
+        path: Path to the state directory (default: /state).
+    """
+
+    enabled: bool = True
+    path: str = "/state"
+
+
+class CheckpointConfig(BaseModel):
+    """Configuration for automatic checkpointing during world execution.
+
+    Attributes:
+        enabled: Whether to enable automatic checkpoints after steps.
+        interval: Create checkpoint every N steps (default: 1 = every step).
+        exclude_envs: Environment aliases to exclude from checkpoints (default: ["runtime"]).
+    """
+
+    enabled: bool = True
+    interval: int = 1
+    exclude_envs: list[str] = Field(default_factory=lambda: ["runtime"])
+
+
 class RunConfig(BaseModel):
     """Base configuration for running a world.
 
@@ -98,6 +128,7 @@ class RunConfig(BaseModel):
         session_id: Unique Chronos session identifier
         callback_url: Callback URL for status updates
         plato_session: Serialized Plato session for connecting to existing VM session
+        checkpoint: Configuration for automatic checkpoints after steps
     """
 
     session_id: str = ""
@@ -108,6 +139,12 @@ class RunConfig(BaseModel):
     # This is the output of Session.dump() - used to restore session with Session.load()
     plato_session: SerializedSession | None = None
 
+    # Checkpoint configuration for automatic snapshots after steps
+    checkpoint: CheckpointConfig = Field(default_factory=CheckpointConfig)
+
+    # State persistence configuration
+    state: StateConfig = Field(default_factory=StateConfig)
+
     model_config = {"extra": "allow"}
 
     @classmethod
@@ -145,7 +182,7 @@ class RunConfig(BaseModel):
         envs = []
 
         # Skip runtime fields
-        runtime_fields = {"session_id", "callback_url", "all_secrets", "plato_session"}
+        runtime_fields = {"session_id", "callback_url", "all_secrets", "plato_session", "checkpoint", "state"}
 
         for field_name, prop_schema in properties.items():
             if field_name in runtime_fields: