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()- 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. Switch to
@kinetic.submit()for long runs. You get aJobHandleback — 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 underKINETIC_OUTPUT_DIR.
Comprehensive documentation is available at: https://kinetic.readthedocs.io
uv pip install keras-kineticThis installs both the decorator and the kinetic CLI.
If nobody on your team has provisioned a Kinetic cluster yet, run:
kinetic upThis enables the required GCP APIs, creates an Artifact Registry
repository, provisions a GKE cluster with an accelerator node pool,
and configures local Docker / kubectl access. Run kinetic down
when you're finished to tear everything back down.
export KINETIC_PROJECT="your-gcp-project-id"
python examples/fashion_mnist.pyThe first run takes ~5 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.
| Question | Where to look |
|---|---|
| How do I get my first job running? | Getting Started |
When should I use submit() instead of run()? |
Detached Jobs |
| How do I ship data and persist outputs? | Data and Checkpointing |
| Bundled vs prebuilt vs custom image — which one? | Execution Modes |
| Something's broken; where do I start? | Troubleshooting |
Kinetic reads KINETIC_PROJECT (required), KINETIC_ZONE,
KINETIC_CLUSTER, and a handful of other environment variables. The
short version:
export KINETIC_PROJECT="your-project-id" # required
export KINETIC_ZONE="us-central1-a" # optional
export KINETIC_CLUSTER="kinetic-cluster" # optionalThe full surface — every variable, every CLI flag, and how the precedence rules combine them — lives in the Configuration reference.
See the Contributing guide.