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.

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 local Docker / 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. 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.

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