plato-sdk-v2 2.2.4__py3-none-any.whl → 2.3.3__py3-none-any.whl

plato/v2/async_/environment.py

@@ -7,6 +7,7 @@ from typing import TYPE_CHECKING
 
 from plato._generated.api.v2 import jobs
 from plato._generated.models import (
+    AppApiV2SchemasSessionCreateSnapshotRequest,
     ConnectRoutingInfoResult,
     CreateCheckpointRequest,
     CreateSnapshotResult,
@@ -136,6 +137,36 @@ class Environment:
             x_api_key=self._api_key,
         )
 
+    async def snapshot_store(
+        self,
+        override_service: str | None = None,
+        override_version: str | None = None,
+        override_dataset: str | None = None,
+    ) -> CreateSnapshotResult:
+        """Create a snapshot-store snapshot of this environment.
+
+        Uses the snapshot-store pipeline for chunk-based deduplication and
+        efficient storage. This is the preferred method for new base snapshots.
+
+        Args:
+            override_service: Override simulator/service name in artifact metadata.
+            override_version: Override version/git_hash in artifact metadata.
+            override_dataset: Override dataset name in artifact metadata.
+
+        Returns:
+            CreateSnapshotResult with artifact_id.
+        """
+        return await jobs.snapshot_store.asyncio(
+            client=self._http,
+            job_id=self.job_id,
+            body=AppApiV2SchemasSessionCreateSnapshotRequest(
+                override_service=override_service,
+                override_version=override_version,
+                override_dataset=override_dataset,
+            ),
+            x_api_key=self._api_key,
+        )
+
     async def close(self) -> None:
         """Close this environment."""
         await jobs.close.asyncio(

plato/v2/async_/session.py

@@ -34,6 +34,7 @@ from plato._generated.api.v2.sessions import reset as sessions_reset
 from plato._generated.api.v2.sessions import set_date as sessions_set_date
 from plato._generated.api.v2.sessions import setup_sandbox as sessions_setup_sandbox
 from plato._generated.api.v2.sessions import snapshot as sessions_snapshot
+from plato._generated.api.v2.sessions import snapshot_store as sessions_snapshot_store
 from plato._generated.api.v2.sessions import state as sessions_state
 from plato._generated.api.v2.sessions import wait_for_ready as sessions_wait_for_ready
 from plato._generated.models import (
@@ -622,6 +623,38 @@ class Session:
             x_api_key=self._api_key,
         )
 
+    async def snapshot_store(
+        self,
+        override_service: str | None = None,
+        override_version: str | None = None,
+        override_dataset: str | None = None,
+    ) -> AppApiV2SchemasSessionCreateSnapshotResponse:
+        """Create a snapshot-store snapshot of all environments in the session.
+
+        Uses the snapshot-store pipeline for chunk-based deduplication and
+        efficient storage. This is the preferred method for new base snapshots.
+
+        Args:
+            override_service: Override simulator/service name in artifact metadata.
+            override_version: Override version/git_hash in artifact metadata.
+            override_dataset: Override dataset name in artifact metadata.
+
+        Returns:
+            Snapshot response with info per job_id.
+        """
+        self._check_closed()
+
+        return await sessions_snapshot_store.asyncio(
+            client=self._http,
+            session_id=self.session_id,
+            body=AppApiV2SchemasSessionCreateSnapshotRequest(
+                override_service=override_service,
+                override_version=override_version,
+                override_dataset=override_dataset,
+            ),
+            x_api_key=self._api_key,
+        )
+
     async def disk_snapshot(
         self,
         override_service: str | None = None,

plato/v2/sync/environment.py

@@ -7,6 +7,7 @@ from typing import TYPE_CHECKING
 
 from plato._generated.api.v2 import jobs
 from plato._generated.models import (
+    AppApiV2SchemasSessionCreateSnapshotRequest,
     CreateCheckpointRequest,
     CreateSnapshotResult,
     ExecuteCommandRequest,
@@ -135,6 +136,36 @@ class Environment:
             x_api_key=self._api_key,
         )
 
+    def snapshot_store(
+        self,
+        override_service: str | None = None,
+        override_version: str | None = None,
+        override_dataset: str | None = None,
+    ) -> CreateSnapshotResult:
+        """Create a snapshot-store snapshot of this environment.
+
+        Uses the snapshot-store pipeline for chunk-based deduplication and
+        efficient storage. This is the preferred method for new base snapshots.
+
+        Args:
+            override_service: Override simulator/service name in artifact metadata.
+            override_version: Override version/git_hash in artifact metadata.
+            override_dataset: Override dataset name in artifact metadata.
+
+        Returns:
+            CreateSnapshotResult with artifact_id.
+        """
+        return jobs.snapshot_store.sync(
+            client=self._http,
+            job_id=self.job_id,
+            body=AppApiV2SchemasSessionCreateSnapshotRequest(
+                override_service=override_service,
+                override_version=override_version,
+                override_dataset=override_dataset,
+            ),
+            x_api_key=self._api_key,
+        )
+
     def close(self) -> None:
         """Close this environment."""
         jobs.close.sync(

plato/v2/sync/session.py

@@ -32,6 +32,7 @@ from plato._generated.api.v2.sessions import reset as sessions_reset
 from plato._generated.api.v2.sessions import set_date as sessions_set_date
 from plato._generated.api.v2.sessions import setup_sandbox as sessions_setup_sandbox
 from plato._generated.api.v2.sessions import snapshot as sessions_snapshot
+from plato._generated.api.v2.sessions import snapshot_store as sessions_snapshot_store
 from plato._generated.api.v2.sessions import state as sessions_state
 from plato._generated.api.v2.sessions import wait_for_ready as sessions_wait_for_ready
 from plato._generated.models import (
@@ -595,6 +596,38 @@ class Session:
             x_api_key=self._api_key,
         )
 
+    def snapshot_store(
+        self,
+        override_service: str | None = None,
+        override_version: str | None = None,
+        override_dataset: str | None = None,
+    ) -> AppApiV2SchemasSessionCreateSnapshotResponse:
+        """Create a snapshot-store snapshot of all environments in the session.
+
+        Uses the snapshot-store pipeline for chunk-based deduplication and
+        efficient storage. This is the preferred method for new base snapshots.
+
+        Args:
+            override_service: Override simulator/service name in artifact metadata.
+            override_version: Override version/git_hash in artifact metadata.
+            override_dataset: Override dataset name in artifact metadata.
+
+        Returns:
+            Snapshot response with info per job_id.
+        """
+        self._check_closed()
+
+        return sessions_snapshot_store.sync(
+            client=self._http,
+            session_id=self.session_id,
+            body=AppApiV2SchemasSessionCreateSnapshotRequest(
+                override_service=override_service,
+                override_version=override_version,
+                override_dataset=override_dataset,
+            ),
+            x_api_key=self._api_key,
+        )
+
     def disk_snapshot(
         self,
         override_service: str | None = None,

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,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/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: