How Google AX Uses Kubernetes and Snapshots to Run AI Agents

| | 9 min read

A long-running AI agent can spend minutes setting up a workspace, make a few tool calls, then sit idle while it waits for a person or another service. A conventional container scheduler can restart a process, but that does not automatically preserve the agent’s live memory, open work, and filesystem state. Google AX is designed to manage this different kind of workload: one that is stateful, bursty, and expensive to keep running continuously.

AX gives operators a declarative, Kubernetes-shaped way to describe agent tasks. It relies on Agent Substrate underneath for sandboxed execution and snapshot-based suspension and resumption. The distinction matters: AX is an orchestration layer for agent work, not a new agent reasoning framework, and not a replacement for Kubernetes.

In brief

  • AX defines and manages tasks, workspaces, and model configuration through resources and a CLI.
  • Agent Substrate runs tasks in isolated sandboxes and can snapshot an actor so its worker can be released while it is idle.
  • Kubernetes still provides the cluster foundation. AX adds an agent-oriented lifecycle above it.
  • Snapshot restore can preserve process and filesystem progress, but it is not a promise of zero latency, zero cost, or lossless recovery from every failure.
  • The project is in active development. Check the current specifications and test your workload before treating it as production infrastructure.

Why AI agents need a different runtime

A web service usually handles a request, returns a response, and can often serve the next request in a fresh process. A batch job has a defined start and finish. An agent may behave differently: it can inspect a repository, call tools, wait for approval, keep conversational or task state in memory, and resume later. It may also create child tasks or retry work.

Leaving every agent process running while it waits can waste worker capacity. Killing it and starting from scratch can discard useful state or force the agent to repeat setup. Kubernetes provides general-purpose workload primitives and controllers, but an agent system still needs to define how agent state is stored, how isolated execution is provisioned, and when an idle agent should be restored.

AX’s design addresses that operational gap by representing work as small, isolated tasks and delegating their execution lifecycle to Agent Substrate. The project’s README describes a high-throughput design goal, but that goal should not be confused with independently demonstrated capacity for billions of real tasks in a customer cluster.

How Google AX fits with Kubernetes

Think of the system as three layers. Kubernetes manages the cluster and runs AX’s control-plane components. AX exposes agent-oriented resources and a CLI that feels familiar to Kubernetes users. Agent Substrate handles sandboxed actors, worker assignment, routing, and lifecycle operations such as suspend and resume.

Layer Responsibility
Kubernetes Provides the cluster, nodes, and infrastructure on which the system runs.
AX Describes agent tasks, their workspaces and model configuration, and provides commands to manage them.
Agent Substrate Runs isolated actors, assigns them to workers, and manages snapshot-based lifecycle transitions.

This is why “Kubernetes for AI agents” is a useful analogy but an incomplete description. AX borrows the declarative workflow, yet the snapshot and actor lifecycle come from the substrate beneath it. Kubernetes itself already supports multiple workload patterns, including Deployments, StatefulSets, and Jobs; AX adds abstractions tailored to agent execution rather than replacing those primitives.

For a broader look at the runtime boundary around agent work, see AI agent evaluation in production. Evaluation tells you whether an agent’s work is good; orchestration determines where and how that work runs.

Google AX Tasks, Workspaces, and Models

AX’s core resources separate the job from the material and configuration it needs:

  • Task: the unit of isolated execution. It can specify a container image, command, environment variables, resource requests and limits, and workspace bindings.
  • Workspace: prepared files and tools that the task can use. A workspace can describe Git repositories, MCP servers, and skills. Multiple tasks can reuse a workspace definition.
  • Model: a reusable provider and model configuration. Credentials can be referenced from a Kubernetes Secret rather than copied into each task manifest.

AX also includes gateway and networking capabilities for reaching running tasks. These building blocks do not dictate how an agent plans or reasons. Your chosen runner or application still supplies that behavior.

AX’s concepts guide explains the resource model, while its manifest reference documents the current fields. Because the project marks its APIs as early-stage, treat the repository’s current examples as authoritative for the version you install.

A Google AX manifest and lifecycle

This example follows the repository’s current sample shape. It binds a task to a Git-backed workspace and enables debugging so an operator can use ax ssh. It assumes AX and Agent Substrate are already installed and configured in the target cluster.

apiVersion: ax.io/v1alpha1
kind: Task
metadata:
  name: task123
  atespace: default
spec:
  image: "gcr.io/ax-substrate/ate-images/ax-task-runner@sha256:3a0dea6ad8b55278685db58aca6e37dc4ba04056831d45bef3aaeafdca43cac6"
  resources:
    requests:
      cpu: "500m"
      memory: "1Gi"
    limits:
      cpu: "2"
      memory: "4Gi"
  workspaces:
    - name: default-workspace
      path: "/workspace"
  debug: true
---
apiVersion: ax.io/v1alpha1
kind: Workspace
metadata:
  name: default-workspace
  atespace: default
spec:
  git:
    - name: origin
      repo: "https://github.com/chalk/chalk.git"
      branch: "main"

The image digest and repository above come from the project’s example. They make the resource shape concrete, but they are not a general-purpose agent configuration. For a real task, use an image and command appropriate to your runner, and configure any required model resource and credentials. AX must be deployed with a reachable Agent Substrate control API before it can schedule the task.

Once the manifest is saved as task.yaml, the lifecycle commands look like this:

ax apply -f task.yaml
ax watch task task123
ax ssh task123 -- ls -al /workspace
ax suspend task task123
ax resume task task123

ax apply submits the resources, ax watch follows task status, and ax ssh provides access when debugging is enabled. Suspending asks the runtime to checkpoint and release the actor; resuming asks it to restore the suspended task. The exact runtime behavior depends on the configured sandbox and snapshot support.

For example, a task that runs an agent process against a code repository could be suspended after the agent has analyzed the code and before it starts an expensive test suite. On resume, the process can continue from the captured state rather than repeating all of its setup, assuming the snapshot is compatible and the suspend completed successfully.

What Google AX snapshots preserve, and what they do not

Agent Substrate describes actor snapshots that preserve memory and durable data, allowing the actor to resume on a worker later. For microVM-backed actors, the architecture describes memory snapshot restoration using demand paging. The runtime can therefore restore more than a fresh container image and a mounted disk: it can restore the actor’s execution state.

That capability has conditions. A snapshot represents a point in time. Work performed after the last successful durable snapshot may be lost if the actor crashes before another snapshot completes. A changed sandbox template can also make a captured memory image incompatible, in which case the runtime may restore durable data and boot a fresh guest instead. Treat snapshot cadence, compatibility, and failure recovery as design questions, not implementation details.

Substrate advertises sub-500-millisecond resume operations in its README. That is a project performance claim, not a universal latency guarantee for every image, snapshot size, worker pool, or cluster load. Restore still takes time, and practical latency should be measured with your own workload.

Suspending idle agents can reduce compute, not every cost

Suspension can free a worker that an idle actor would otherwise occupy. That can improve utilization when many agents alternate between active work and waiting. It does not mean the entire deployment costs nothing while an actor is idle. The cluster, control plane, snapshot storage, networking, and any warm capacity can still incur costs. Model API charges are separate, and suspending a process does not refund tokens it already used.

AX’s own materials call out the risk of agents burning money in a loop. However, the project roadmap lists token and timeout budgets, approval policies, governance, and audit controls among areas to stabilize or develop. Operators should verify what guardrails exist in the version they deploy and enforce hard spend limits at the model-provider or application layer where necessary.

We explore a related state-management question in our guide to LangChain MCP Adapter’s stateless pause-and-resume flow. The key distinction is that protocol-level continuation and process-level snapshot restoration solve different parts of the continuity problem.

Where AX may fit, and where to be cautious

AX is most interesting when tasks need isolation, repeatable workspace preparation, and a way to release compute during idle periods without always discarding process state. That can describe coding agents, tool-using workers, or long-running agent sessions that wait on external events.

It is less compelling for a short, stateless model call or a predictable batch job that already fits a simpler execution model. Adding a controller, sandbox runtime, snapshot storage, and operational surface creates complexity. A good evaluation should compare that complexity with the cost of keeping agents warm or rebuilding their state after restarts.

The repository currently warns that AX and several of its features are in heavy development, with breaking changes likely before a stable release. Its roadmap includes stabilizing resource specifications, evolving the actor architecture, and expanding governance and observability. That makes AX worth experimenting with, while making production adoption a deliberate infrastructure decision.

The practical takeaway

AX combines familiar Kubernetes-style declarations with an agent-specific lifecycle built on Agent Substrate. Its compelling idea is not that agents become magically cheap or impossible to interrupt. It is that a platform can treat an agent as a stateful actor, snapshot it when idle, and restore it when work resumes.

If you are evaluating AX, start with the official AX repository, its roadmap, and the Agent Substrate architecture. Test suspend, resume, crash recovery, worker contention, and total cost on a representative task before relying on the headline performance or scale claims.

The post How Google AX Uses Kubernetes and Snapshots to Run AI Agents appeared first on Alpesh Kumar.

Subscribe to Our Newsletter

We don’t spam! Read our privacy policy for more info.