Kinetic

Run Keras and JAX workloads on cloud TPUs and GPUs with a simple decorator. No infrastructure management required.

import kinetic


@kinetic.run(accelerator="tpu-v5e-1")
def train_model():
  import keras

  model = keras.Sequential([...])
  model.fit(x_train, y_train)
  return model.history.history["loss"][-1]


# Executes on a TPU v5e-1 slice, returns the result locally
final_loss = train_model()

Why Kinetic

• Simple remote execution. A @kinetic.run() decorator runs the function on the accelerator you ask for and returns the result. Nothing else changes about your code.
• Detached jobs. Use func.run_async() for long runs. You get a JobHandle back — poll status, tail logs, collect the result later, or reattach from another machine entirely.
• Data and checkpoint support. Wrap inputs in kinetic.Data(...) to ship local files (or stream from GCS) into the job. Write durable outputs and resumable checkpoints under KINETIC_OUTPUT_DIR.
• Your project travels with the job. Kinetic finds your package root and ships the source with the pickled function. Kinetic then rebuilds sys.path and the working directory on the pod. Multi-module projects and relative-path reads thus operate as they do on your machine. See What Ships to the Pod.

Documentation

Comprehensive documentation is available at: https://kinetic.readthedocs.io

Install

uv pip install keras-kinetic

This installs the @kinetic.run() decorator and the kinetic CLI, which provisions and manages infrastructure.

One-time setup

kinetic init

This detects your local environment, then either joins an existing Kinetic cluster in the project (your own or a teammate's — discovery goes through the shared state bucket) or walks you through creating a new one. It ends by saving a profile that becomes your active context — subsequent commands pick up project, zone, and cluster automatically.

Behind the scenes, the Create path runs kinetic up to enable APIs, provision a GKE cluster with an accelerator node pool, and configure kubectl access. Run kinetic down when you're done.

Recommended first run

python examples/fashion_mnist.py

No environment variables needed — kinetic init set an active profile. Make sure the cluster has a node pool for the script's accelerator (kinetic pool list; kinetic pool add --accelerator tpu-v5litepod-1 if not). The first run takes 5–10 minutes (it builds a container image with your dependencies via Cloud Build). Subsequent runs with unchanged dependencies start in under a minute.

For the full first-run walkthrough, see the Getting Started guide.

Where to go next

Configuration

The recommended way to configure Kinetic is via a profile — the named context that kinetic init creates and kinetic profile ls | use manages. For ad-hoc overrides, every profile field also has a KINETIC_* env-var equivalent (KINETIC_PROJECT, KINETIC_ZONE, KINETIC_CLUSTER, KINETIC_NAMESPACE) and a matching CLI flag.

Precedence is: CLI flag > KINETIC_* env var > active profile > built-in default.

The full surface — every variable, every CLI flag, and the profile model — lives in the Configuration reference.

Contributing

See the Contributing guide.

License

Apache 2.0