Python API

This document provides a comprehensive guide to the pie_client Python library, an asynchronous client designed for interacting with the Pie WebSocket server. It allows you to compile, upload, and manage WebAssembly programs on a remote server.

Prerequisites

Before using the client, ensure you have the necessary libraries installed:

cd client/python
pip install -e .

---

Quick Start

Here’s a complete example demonstrating the typical workflow: connecting to the server, compiling and uploading a Rust program, launching it as an instance, and communicating with it.

main.py

import asyncio
import blake3
from pie_client import PieClient, compile_program, Event

# A simple Rust program that echoes messages and sends back a blob.
RUST_CODE = r"""
use inferlet::{Args, Result};

#[inferlet::main]
async fn main(mut args: Args) -> Result<()> {
    println!("Hello, world!");
    Ok(())
}
"""

async def main():
    # Connect to the Pie server using an async context manager
    async with PieClient("ws://127.0.0.1:8080") as client:
        # Authenticate with the server
        success, msg = await client.authenticate("your-secret-token")
        if not success:
            raise Exception(f"Authentication failed: {msg}")

        # 1. Compile the Rust code to WASM
        # Dependencies can be specified, e.g., ["serde = \"1.0\""]
        wasm_bytes = await compile_program(RUST_CODE, dependencies=[])
        program_hash = blake3.blake3(wasm_bytes).hexdigest()

        # 2. Upload the program if it doesn't already exist on the server
        if not await client.program_exists(program_hash):
            print(f"Program {program_hash} not found. Uploading...")
            await client.upload_program(wasm_bytes)
        else:
            print(f"Program {program_hash} already exists on server.")

        # 3. Launch the program as a new instance
        print("Launching instance...")
        instance = await client.launch_instance(program_hash)
        print(f"Instance launched with ID: {instance.instance_id}")

        # 4. Interact with the instance
        await instance.send("Hello from the client!")

        # 5. Listen for events from the instance
        while True:
            event, message = await instance.recv()
            print(f"Received Event: {event.name}, Message: {message}")

            if event == Event.Blob:
                # The message for a Blob event is the binary data itself
                print(f"Decoded blob content: '{message.decode()}'")
            elif event in [Event.Completed, Event.Aborted, Event.Exception]:
                print("Instance has terminated. Exiting event loop.")
                break

if __name__ == "__main__":
    asyncio.run(main())

---

PieClient Class

The PieClient is the primary interface for communicating with the Pie server. It handles connection, authentication, and program management. It's designed to be used as an asynchronous context manager.

Initialization

class PieClient:
    def __init__(self, server_uri: str):
        """
        Initializes the client.

        Args:
            server_uri (str): The WebSocket URI of the Pie server (e.g., "ws://127.0.0.1:8080").
        """

Usage:

client = PieClient("ws://localhost:8080")

Connection Management

The client is best used with an async with block, which automatically handles connecting and disconnecting.

async def main():
    async with PieClient("ws://127.0.0.1:8080") as client:
        # Client is connected and ready to use
        ...
    # Client is automatically disconnected here

Methods

authenticate

Authenticates the client session. This should be the first call after connecting.

async def authenticate(self, token: str) -> tuple[bool, str]:

• Args:
  • token (str): The authentication token.
• Returns: A tuple (successful, result) where successful is a boolean and result is a message from the server.

upload_program

Uploads a compiled WASM program to the server. The data is sent in chunks.

async def upload_program(self, program_bytes: bytes):

• Args:
  • program_bytes (bytes): The raw binary content of the WASM file.
• Raises: Exception if the upload fails.

program_exists

Checks if a program with a given hash is already available on the server, avoiding unnecessary uploads.

async def program_exists(self, program_hash: str) -> bool:

• Args:
  • program_hash (str): The blake3 hash of the program bytes.
• Returns: True if the program exists, False otherwise.

launch_instance

Requests the server to start a new instance of an uploaded program.

async def launch_instance(self, program_hash: str, arguments: list[str] = None) -> Instance:

• Args:
  • program_hash (str): The hash of the program to launch.
  • arguments (list[str], optional): A list of command-line arguments for the instance.
• Returns: An Instance object to interact with the running program.
• Raises: Exception if the launch fails.

launch_server_instance

Launches a program that acts as a server, listening on a specified port on the host machine.

async def launch_server_instance(self, program_hash: str, port: int, arguments: list[str] = None):

• Args:
  • program_hash (str): The hash of the program to launch.
  • port (int): The host port the instance should listen on.
  • arguments (list[str], optional): A list of command-line arguments.
• Raises: Exception if the launch fails.

---

Instance Class

An Instance object represents a single running program on the server. It is returned by PieClient.launch_instance() and provides methods for direct interaction.

Methods

send

Sends a string message to the instance. This is typically received by the instance via its standard input or a specific function call.

async def send(self, message: str):

• Args:
  • message (str): The string message to send.

upload_blob

Uploads a blob of binary data to the instance. This is useful for providing large data files or inputs to a running program.

async def upload_blob(self, blob_bytes: bytes):

• Args:
  • blob_bytes (bytes): The binary data to upload.

recv

Waits for and receives the next event from the instance. This is a blocking call.

async def recv(self) -> tuple[Event, str | bytes]:

• Returns: A tuple (event, message).
  • event: An Event enum member indicating the event type.
  • message: A str for most events, or bytes if the event is Event.Blob.

terminate

Sends a request to the server to terminate the instance. This is a fire-and-forget operation.

async def terminate(self):

---

Event Enum

The Event enumeration defines the types of events you can receive from an instance via instance.recv().

class Event(Enum):
    Message = 0        # Standard output (stdout) from the instance.
    Completed = 1      # The instance finished successfully (exit code 0).
    Aborted = 2        # The instance terminated with a non-zero exit code.
    Exception = 3      # An internal WASM trap/exception occurred.
    ServerError = 4    # The server encountered an error managing the instance.
    OutOfResources = 5 # The instance was terminated for exceeding resource limits.
    Blob = 6           # The instance sent a binary data blob.

---

compile_program Utility

A helper function that compiles Rust source code into a WASM binary compatible with the Pie server. It runs the Rust compiler (cargo) in a separate process.

danger

This function executes a local cargo command and writes temporary files to your disk. It requires a working Rust environment.

async def compile_program(source: str | Path, dependencies: list[str]) -> bytes:

• Args:
  • source (str | Path): The Rust source code as a string, or a pathlib.Path to a .rs file.
  • dependencies (list[str]): A list of Cargo dependencies, e.g., ['serde = "1.0"'].
• Returns: The compiled WASM binary as bytes.
• Raises: RuntimeError if compilation fails or cargo is not found.

Example:

# Compile from a string
rust_code = 'fn main() { println!("Hello"); }'
wasm_bytes = await compile_program(rust_code, dependencies=[])

# Compile from a file
from pathlib import Path
wasm_bytes_from_file = await compile_program(Path("src/my_program.rs"), dependencies=[])