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-kinetic[cli]"The base keras-kinetic package installs the @kinetic.run()
decorator. The [cli] extra adds the dependencies the kinetic CLI
needs to provision and manage infrastructure. Drop the [cli] extra
only if you just need to submit jobs against an already-provisioned
cluster.
kinetic initThis 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
local Docker / kubectl access. Run kinetic down when you're done.
python examples/fashion_mnist.pyNo environment variables needed — kinetic init set an active
profile. The 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 |
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.
See the Contributing guide.