Run ML workloads on cloud TPUs and GPUs

Run ML workloads on cloud TPUs and GPUs#

Run any Python function on a cloud TPU or GPU with one decorator. Kinetic creates the infrastructure, builds the container image, ships your code, and returns the result.

Get started

How it works

import kinetic


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

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


final_loss = train_model()  # runs on a 4-chip TPU v5e slice

Start here#

Read these three pages in order. They take about 30 minutes. When something does not work, see Troubleshooting, and for short answers see the FAQ.

1. Getting Started

Install Kinetic, run kinetic init, and run a Keras job on a TPU.

Getting Started

2. How Kinetic Works

The vocabulary, the job lifecycle, and where your code, data, and results go.

How Kinetic Works

3. Examples

Runnable scripts, from a first run to multi-host LLM fine-tuning.

Examples

What happens on every job#

Package. Kinetic serializes your function and archives your project source. Data(...) arguments upload one time, keyed by content.
Build. Kinetic builds a container image with the packages from your requirements.txt or pyproject.toml, and caches the image. Later runs with the same dependencies skip this step.
Schedule. Kinetic creates a Kubernetes Job on your GKE cluster, or a LeaderWorkerSet for a multi-host TPU slice. The autoscaler starts an accelerator node in the matching node pool.
Run. The pod runs your function with KINETIC_OUTPUT_DIR set. The logs stream to your terminal.
Collect. The pod uploads the return value to Cloud Storage. Kinetic downloads the value and deletes the job resources. Files that you wrote under KINETIC_OUTPUT_DIR stay.