plato-sdk-v2 2.8.7__py3-none-any.whl → 2.9.0__py3-none-any.whl

plato/v2/sync/environment.py

@@ -39,11 +39,17 @@ class Environment:
         job_id: str,
         alias: str,
         artifact_id: str | None = None,
+        simulator: str | None = None,
+        status: str | None = None,
+        public_url: str | None = None,
     ):
         self._session = session
         self.job_id = job_id
         self.alias = alias
         self.artifact_id = artifact_id
+        self.simulator = simulator
+        self.status = status
+        self.public_url = public_url
 
     @property
     def _http(self):

plato/v2/sync/sandbox.py

@@ -53,7 +53,6 @@ from plato._generated.models import (
 )
 from plato.v2.async_.flow_executor import FlowExecutor
 from plato.v2.models import SandboxState
-from plato.v2.sync.client import Plato
 from plato.v2.types import Env, EnvFromArtifact, EnvFromResource, EnvFromSimulator, SimConfigCompute
 
 logger = logging.getLogger(__name__)
@@ -207,6 +206,8 @@ def _run_ssh_command(
 def _start_heartbeat_process(session_id: str, api_key: str) -> int | None:
     """Start a background process that sends heartbeats.
 
+    Uses only stdlib (urllib) to work on any machine without dependencies.
+
     Returns:
         PID of the background process, or None if failed.
     """
@@ -217,10 +218,12 @@ def _start_heartbeat_process(session_id: str, api_key: str) -> int | None:
         base_url = base_url[:-4]
     base_url = base_url.rstrip("/")
 
+    # Use only stdlib - no external dependencies
     heartbeat_script = f'''
 import time
-import os
-import httpx
+import json
+import urllib.request
+import urllib.error
 from datetime import datetime
 
 session_id = "{session_id}"
@@ -235,23 +238,27 @@ def log(msg):
         f.flush()
 
 log(f"Heartbeat process started for session {{session_id}}")
+log(f"URL: {{base_url}}/api/v2/sessions/{{session_id}}/heartbeat")
 
 heartbeat_count = 0
 while True:
     heartbeat_count += 1
     try:
         url = f"{{base_url}}/api/v2/sessions/{{session_id}}/heartbeat"
-        with httpx.Client(timeout=30) as client:
-            response = client.post(
-                url,
-                headers={{"X-API-Key": api_key}},
-            )
-            if response.status_code == 200:
-                result = response.json()
-                success = result.get("success", False)
-                log(f"Heartbeat #{{heartbeat_count}}: status={{response.status_code}}, success={{success}}")
-            else:
-                log(f"Heartbeat #{{heartbeat_count}}: status={{response.status_code}}, url={{url}}, body={{response.text[:500]}}")
+        req = urllib.request.Request(
+            url,
+            method="POST",
+            headers={{"X-API-Key": api_key, "Content-Type": "application/json"}},
+            data=b"{{}}",
+        )
+        with urllib.request.urlopen(req, timeout=30) as resp:
+            status = resp.status
+            body = resp.read().decode("utf-8")
+            result = json.loads(body)
+            success = result.get("success", False)
+            log(f"Heartbeat #{{heartbeat_count}}: status={{status}}, success={{success}}")
+    except urllib.error.HTTPError as e:
+        log(f"Heartbeat #{{heartbeat_count}}: HTTP {{e.code}} - {{e.reason}}")
     except Exception as e:
         log(f"Heartbeat #{{heartbeat_count}} EXCEPTION: {{type(e).__name__}}: {{e}}")
     time.sleep(30)
@@ -290,6 +297,77 @@ class SyncResult(BaseModel):
     bytes_synced: int
 
 
+class SSHConfigInfo(BaseModel):
+    """SSH config information for connecting to a job."""
+
+    config_content: str
+    private_key_path: str
+    job_id: str
+    gateway_host: str
+
+
+def _generate_temp_ssh_key_pair() -> tuple[str, str]:
+    """Generate a temporary SSH key pair.
+
+    Returns:
+        Tuple of (public_key_content, private_key_path).
+    """
+    # Create temp directory for keys
+    temp_dir = tempfile.mkdtemp(prefix="plato_ssh_")
+    private_key_path = os.path.join(temp_dir, "id_ed25519")
+
+    # Generate key pair
+    subprocess.run(
+        [
+            "ssh-keygen",
+            "-t",
+            "ed25519",
+            "-f",
+            private_key_path,
+            "-N",
+            "",
+            "-q",
+        ],
+        check=True,
+    )
+
+    # Read public key
+    public_key = Path(f"{private_key_path}.pub").read_text().strip()
+
+    return public_key, private_key_path
+
+
+def _generate_ssh_config_content(job_id: str, private_key_path: str) -> str:
+    """Generate SSH config content for a job.
+
+    Args:
+        job_id: The job ID for routing.
+        private_key_path: Path to private key.
+
+    Returns:
+        SSH config content as a string.
+    """
+    gateway_host = os.getenv("PLATO_GATEWAY_HOST", "gateway.plato.so")
+
+    # SNI format: {job_id}--{port}.{gateway_host}
+    ssh_port = 22
+    sni = f"{job_id}--{ssh_port}.{gateway_host}"
+
+    config_content = f"""# Plato SSH Config for job: {job_id}
+# Generated dynamically for -J/--job-id option
+
+Host sandbox
+    HostName {job_id}
+    User root
+    IdentityFile {private_key_path}
+    StrictHostKeyChecking no
+    UserKnownHostsFile /dev/null
+    LogLevel ERROR
+    ProxyCommand openssl s_client -quiet -connect {gateway_host}:443 -servername {sni} 2>/dev/null
+"""
+    return config_content
+
+
 # =============================================================================
 # TUNNEL
 # =============================================================================
@@ -569,24 +647,26 @@ class SandboxClient:
         env_config: EnvFromSimulator | EnvFromArtifact | EnvFromResource
 
         if mode == "artifact" and artifact_id:
-            self.console.print(f"Starting from artifact: {artifact_id}")
+            self.console.print(f"[cyan]Mode:[/cyan] artifact ({artifact_id})")
             env_config = Env.artifact(artifact_id)
         elif mode == "simulator" and simulator_name:
-            self.console.print(f"Starting from simulator: {simulator_name}")
+            self.console.print(f"[cyan]Mode:[/cyan] simulator ({simulator_name}:{tag})")
             env_config = Env.simulator(simulator_name, tag=tag, dataset=dataset)
-        elif mode == "blank" and simulator_name:
-            self.console.print("Starting from blank")
+        elif mode == "blank":
+            # Use provided simulator_name or default to "sandbox"
+            sim_name = simulator_name or "sandbox"
+            self.console.print(f"[cyan]Mode:[/cyan] blank VM ({sim_name})")
+            self.console.print(f"[dim]  cpus={cpus}, memory={memory}MB, disk={disk}MB[/dim]")
             sim_config = SimConfigCompute(
                 cpus=cpus, memory=memory, disk=disk, app_port=app_port, plato_messaging_port=messaging_port
             )
-            env_config = Env.resource(simulator_name, sim_config)
+            env_config = Env.resource(sim_name, sim_config)
         elif mode == "config":
-            self.console.print("Starting from config")
+            self.console.print("[cyan]Mode:[/cyan] config (plato-config.yml)")
             # read plato-config.yml
             plato_config_path = self.working_dir / "plato-config.yml"
             with open(plato_config_path, "rb") as f:
                 plato_config = yaml.safe_load(f)
-            self.console.print(f"plato-config: {plato_config}")
             plato_config_model = PlatoConfig.model_validate(plato_config)
             dataset_config = plato_config_model.datasets[dataset]
             simulator_name = plato_config_model.service
@@ -609,24 +689,67 @@ class SandboxClient:
                 plato_messaging_port=config_messaging_port,
             )
             env_config = Env.resource(simulator_name, sim_config)
-            self.console.print(f"env_config: {env_config}")
         else:
             raise ValueError(f"Invalid mode '{mode}' or missing required parameter")
 
-        # Use Plato SDK to create session (handles create, wait, network)
-        self.console.print(f"Creating session and waiting for VM to become ready (timeout={timeout}s)...")
-        plato = Plato(api_key=self.api_key, http_client=self._http)
-        session = plato.sessions.create(
-            envs=[env_config],
-            connect_network=connect_network,
+        # Track total time
+        total_start = time.time()
+
+        # The sessions.create handles: create -> wait_for_ready -> connect_network
+        # But we want to show progress, so we'll do it step by step
+        from plato._generated.api.v2.sessions import make as sessions_make
+        from plato._generated.api.v2.sessions import wait_for_ready as sessions_wait_for_ready
+        from plato._generated.models import CreateSessionFromEnvs, Envs, RunSessionSource
+
+        # Step 1: Create session
+        self.console.print("[yellow]Creating session...[/yellow]")
+        step_start = time.time()
+        request_body = CreateSessionFromEnvs(
+            envs=[Envs(root=env_config)],
             timeout=timeout,
+            source=RunSessionSource.SDK,
+        )
+        response = sessions_make.sync(
+            client=self._http,
+            body=request_body,
+            x_api_key=self.api_key,
         )
-        self.console.print(f"session: {session}")
-        session_id = session.session_id
-        job_id = session.envs[0].job_id if session.envs else None
+        session_id = response.session_id
+        elapsed = time.time() - step_start
+        self.console.print(f"[green]Session created:[/green] {session_id} [dim]({elapsed:.1f}s)[/dim]")
+
+        # Check if any envs failed to create
+        if response.envs:
+            for env_result in response.envs:
+                if not env_result.success:
+                    raise RuntimeError(f"Failed to create environment: {env_result.error}")
+                if env_result.job_id:
+                    self.console.print(f"[dim]  Job: {env_result.job_id}[/dim]")
+        else:
+            raise RuntimeError("No environments created in session")
+
+        # Step 2: Wait for VM
+        self.console.print("[yellow]Waiting for VM to start...[/yellow]")
+        step_start = time.time()
+        ready_response = sessions_wait_for_ready.sync(
+            client=self._http,
+            session_id=session_id,
+            timeout=timeout,
+            x_api_key=self.api_key,
+        )
+        if not ready_response.ready:
+            errors = []
+            if ready_response.results:
+                for jid, result in ready_response.results.items():
+                    if not result.ready:
+                        errors.append(f"{jid}: {result.error or 'Unknown error'}")
+            raise RuntimeError(f"VM failed to start: {', '.join(errors) if errors else 'timeout'}")
+
+        job_id = response.envs[0].job_id if response.envs else None
         if not job_id:
             raise ValueError("No job ID found")
-        self.console.print(f"job_id: {job_id}")
+        elapsed = time.time() - step_start
+        self.console.print(f"[green]VM ready:[/green] {job_id} [dim]({elapsed:.1f}s)[/dim]")
 
         # For artifact mode, we need to get simulator_name from session details BEFORE generating public URL
         # Note: get_session_details returns a dict, not a Pydantic model
@@ -670,10 +793,14 @@ class SandboxClient:
                         else:
                             url = f"{url}?{target_param}"
                     public_url = url
+            elapsed = time.time() - step_start
+            self.console.print(f"[green]Public URL:[/green] {public_url} [dim]({elapsed:.1f}s)[/dim]")
         except Exception as e:
-            raise ValueError(f"Error getting public URL: {e}") from e
+            self.console.print(f"[dim]Public URL not available: {e}[/dim]")
 
         # Setup SSH
+        self.console.print("[yellow]Setting up SSH...[/yellow]")
+        step_start = time.time()
         ssh_config_path = None
         try:
             public_key, private_key_path = _generate_ssh_key_pair(session_id[:8], Path(self.working_dir))
@@ -688,12 +815,24 @@ class SandboxClient:
 
             if add_response.success:
                 ssh_config_path = _generate_ssh_config(job_id, private_key_path, Path(self.working_dir))
+                elapsed = time.time() - step_start
+                self.console.print(
+                    f"[green]SSH configured:[/green] ssh -F .plato/ssh_config sandbox [dim]({elapsed:.1f}s)[/dim]"
+                )
+            else:
+                self.console.print("[dim]SSH key upload failed[/dim]")
         except Exception as e:
-            logger.warning(f"SSH setup failed: {e}")
+            self.console.print(f"[dim]SSH setup failed: {e}[/dim]")
 
         # Start heartbeat
-        heartbeat_pid = None
+        self.console.print("[yellow]Starting heartbeat...[/yellow]")
+        step_start = time.time()
         heartbeat_pid = _start_heartbeat_process(session_id, self.api_key)
+        elapsed = time.time() - step_start
+        if heartbeat_pid:
+            self.console.print(f"[green]Heartbeat started[/green] (pid={heartbeat_pid}) [dim]({elapsed:.1f}s)[/dim]")
+        else:
+            self.console.print("[dim]Heartbeat failed to start[/dim]")
 
         # Convert absolute paths to relative for state storage
         def _to_relative(abs_path: str | None) -> str | None:
@@ -718,6 +857,7 @@ class SandboxClient:
             heartbeat_pid=heartbeat_pid,
             simulator_name=simulator_name,
             dataset=dataset,
+            network_connected=connect_network,
         )
         if mode == "artifact":
             sandbox_state.artifact_id = artifact_id
@@ -740,6 +880,10 @@ class SandboxClient:
         with open(self.working_dir / self.PLATO_DIR / "state.json", "w") as f:
             json.dump(sandbox_state.model_dump(), f)
 
+        total_elapsed = time.time() - total_start
+        self.console.print("")
+        self.console.print(f"[bold green]Sandbox ready![/bold green] [dim](total: {total_elapsed:.1f}s)[/dim]")
+
         return sandbox_state
 
     # CHECKED
@@ -784,7 +928,7 @@ class SandboxClient:
             plato_config_model = PlatoConfig.model_validate(plato_config)
             dataset_config = plato_config_model.datasets[dataset]
             # Convert dataset config back to dict for YAML serialization
-            dataset_dict = dataset_config.model_dump(exclude_none=True, by_alias=True)
+            dataset_dict = dataset_config.model_dump(exclude_none=True, by_alias=True, mode="json")
             checkpoint_request.plato_config = yaml.dump(dataset_dict, default_flow_style=False)
 
             dataset_compute = dataset_config.compute
@@ -839,7 +983,7 @@ class SandboxClient:
 
         # Convert AppApiV2SchemasArtifactSimConfigDataset to AppSchemasBuildModelsSimConfigDataset
         # They have compatible fields but different nested types
-        dataset_config_dict = dataset_config.model_dump(exclude_none=True)
+        dataset_config_dict = dataset_config.model_dump(exclude_none=True, mode="json")
 
         _ = start_worker.sync(
             client=self._http,
@@ -1011,6 +1155,60 @@ class SandboxClient:
             bind_address=bind_address,
         )
 
+    def get_ssh_config_for_job(self, job_id: str) -> SSHConfigInfo:
+        """Get SSH config for connecting to a specific job.
+
+        Generates a temporary SSH key pair, adds the public key to the VM,
+        and returns an SSH config that routes through the Plato gateway.
+
+        Args:
+            job_id: The job public ID to connect to.
+
+        Returns:
+            SSHConfigInfo with the config content, private key path, and metadata.
+
+        Note:
+            The caller is responsible for cleaning up the temporary key files
+            after the SSH session ends. The private key is stored in a temp
+            directory that should be removed when done.
+        """
+        from plato._generated.api.v2.jobs import execute as jobs_execute
+        from plato._generated.models import ExecuteCommandRequest
+
+        gateway_host = os.getenv("PLATO_GATEWAY_HOST", "gateway.plato.so")
+
+        # Generate temp SSH key pair
+        public_key, private_key_path = _generate_temp_ssh_key_pair()
+
+        # Add public key to the VM via execute
+        add_key_cmd = f'mkdir -p /root/.ssh && chmod 700 /root/.ssh && echo "{public_key}" >> /root/.ssh/authorized_keys && chmod 600 /root/.ssh/authorized_keys'
+        try:
+            jobs_execute.sync(
+                client=self._http,
+                job_id=job_id,
+                body=ExecuteCommandRequest(command=add_key_cmd, timeout=30),
+                x_api_key=self.api_key,
+            )
+        except Exception as e:
+            # Clean up temp key on failure
+            try:
+                import shutil
+
+                shutil.rmtree(os.path.dirname(private_key_path))
+            except Exception:
+                pass
+            raise RuntimeError(f"Failed to add SSH key to job {job_id}: {e}") from e
+
+        # Generate SSH config
+        config_content = _generate_ssh_config_content(job_id, private_key_path)
+
+        return SSHConfigInfo(
+            config_content=config_content,
+            private_key_path=private_key_path,
+            job_id=job_id,
+            gateway_host=gateway_host,
+        )
+
     def run_audit_ui(
         self,
         job_id: str | None = None,

plato/v2/sync/session.py

@@ -438,6 +438,8 @@ class Session:
                     job_id=ctx.job_id,
                     alias=ctx.alias,
                     artifact_id=ctx.artifact_id,
+                    simulator=ctx.simulator,
+                    status="running",  # Environments are running after from_envs completes
                 )
                 for ctx in env_contexts
             ]
@@ -640,6 +642,7 @@ class Session:
         override_service: str | None = None,
         override_version: str | None = None,
         override_dataset: str | None = None,
+        target: str | None = None,
     ) -> CreateDiskSnapshotResponse:
         """Create a disk-only snapshot of all environments in the session.
 
@@ -647,10 +650,13 @@ class Session:
         will do a fresh boot with the preserved disk state. This is faster to
         create and smaller to store than full snapshots.
 
+        Uses snapshot-store backend for chunk-based deduplication and efficient storage.
+
         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.
+            target: Target domain for routing (e.g., "sims.plato.so").
 
         Returns:
             CreateDiskSnapshotResponse with artifact_id per job_id.
@@ -664,6 +670,7 @@ class Session:
                 override_service=override_service,
                 override_version=override_version,
                 override_dataset=override_dataset,
+                target=target,
             ),
             x_api_key=self._api_key,
         )
@@ -793,6 +800,8 @@ class Session:
             job_id=job_id,
             alias=env.alias,
             artifact_id=response.env.artifact_id,
+            simulator=getattr(env, "simulator", None),
+            status="running",  # Newly added environments are running
         )
 
         logger.info(f"Added job {job_id} (alias={env.alias}) to session {self.session_id}")

plato/v2/types.py

@@ -30,17 +30,27 @@ class Env:
     def simulator(
         simulator: str,
         *,
-        tag: str = "latest",
+        tag: str | None = None,
+        version: str | None = None,
         dataset: str | None = None,
         alias: str | None = None,
+        restore_memory: bool = True,
     ) -> EnvFromSimulator:
-        """Create env from simulator with tag.
+        """Create env from simulator with tag or version lookup.
+
+        Supports string formats:
+        - "simname" -> uses 'latest' tag
+        - "simname:tag" -> uses specified tag
+        - "simname:version@dataset" -> uses version (git_hash) lookup
 
         Args:
-            simulator: Simulator name, or "env:tag" format (e.g., "espocrm:staging")
-            tag: Artifact tag (default: "latest"). Ignored if simulator contains ":"
+            simulator: Simulator name with optional tag/version format
+            tag: Artifact tag. If neither tag nor version provided, defaults to 'latest'.
+            version: Artifact version (git_hash). If provided, looks up by version instead of tag.
             dataset: Dataset name (e.g., "base", "blank"). If not specified, uses default.
             alias: Custom name for this environment
+            restore_memory: If True (default), resume from memory snapshot.
+                           If False, do a fresh boot with disk state only (for disk snapshots).
 
         Returns:
             EnvFromSimulator
@@ -48,21 +58,33 @@ class Env:
         Examples:
             >>> Env.simulator("espocrm")  # -> uses "latest" tag
             >>> Env.simulator("espocrm:staging")  # -> uses "staging" tag
-            >>> Env.simulator("espocrm", tag="staging")  # -> uses "staging" tag
-            >>> Env.simulator("gitea", dataset="blank")  # -> uses "blank" dataset
+            >>> Env.simulator("espocrm:v1@base")  # -> uses version "v1", dataset "base"
+            >>> Env.simulator("espocrm:v1@base", restore_memory=False)  # disk snapshot
         """
-        # Support "env:tag" format
-        if ":" in simulator:
+        sim_name = simulator
+
+        # Parse "simname:version@dataset" format (version lookup)
+        if "@" in simulator:
+            parts = simulator.split("@", 1)
+            dataset = dataset or parts[1]
+            if ":" in parts[0]:
+                sim_name, version = parts[0].split(":", 1)
+            else:
+                sim_name = parts[0]
+        # Parse "simname:tag" format (tag lookup)
+        elif ":" in simulator:
             sim_name, tag = simulator.split(":", 1)
-        else:
-            sim_name = simulator
 
-        # Build kwargs, only including dataset if specified
+        # Build kwargs
         kwargs: dict[str, Any] = {
             "simulator": sim_name,
-            "tag": tag,
             "alias": alias,
+            "restore_memory": restore_memory,
         }
+        if tag is not None:
+            kwargs["tag"] = tag
+        if version is not None:
+            kwargs["version"] = version
         if dataset is not None:
             kwargs["dataset"] = dataset
 
@@ -94,27 +116,36 @@ class Env:
     @staticmethod
     def resource(
         simulator: str,
-        sim_config: SimConfigCompute,
+        sim_config: SimConfigCompute | None = None,
         *,
         alias: str | None = None,
+        docker_image_url: str | None = None,
+        upload_rootfs: bool | None = True,
     ) -> EnvFromResource:
         """Create env from resource specification (blank VM).
 
         Args:
             simulator: Simulator/service name
-            sim_config: Resource configuration (CPUs, memory, disk)
+            sim_config: Resource configuration (CPUs, memory, disk). If None, uses defaults.
             alias: Custom name for this environment
+            docker_image_url: Custom Docker image URL (ECR). If not set, uses default.
+            upload_rootfs: Upload rootfs to S3/snapshot-store if not cached. Set False for one-off VMs.
 
         Returns:
             EnvFromResource
 
-        Example:
+        Examples:
             >>> Env.resource("redis", SimConfigCompute(cpus=4, memory=8192, disk=20000))
+            >>> Env.resource("agent", docker_image_url="123.dkr.ecr.us-west-1.amazonaws.com/my-image:v1")
         """
+        if sim_config is None:
+            sim_config = SimConfigCompute()
         return EnvFromResource(
             simulator=simulator,
             sim_config=sim_config,
             alias=alias,
+            docker_image_url=docker_image_url,
+            upload_rootfs=upload_rootfs,
         )
 
 

plato/worlds/__init__.py

@@ -46,23 +46,19 @@ from plato._generated.models import (
 from plato.worlds.base import (
     BaseWorld,
     ConfigT,
-    Observation,
-    StepResult,
     get_registered_worlds,
     get_world,
     register_world,
 )
 from plato.worlds.config import (
-    Agent,
     AgentConfig,
     CheckpointConfig,
-    Env,
     EnvConfig,
-    EnvList,
     RunConfig,
-    Secret,
     StateConfig,
 )
+from plato.worlds.markers import Agent, Env, EnvList, Secret
+from plato.worlds.models import Observation, StepResult
 from plato.worlds.runner import run_world
 
 __all__ = [