diff --git a/README.md b/README.md
index a72061a2..08b0db80 100644
--- a/README.md
+++ b/README.md
@@ -67,7 +67,7 @@ new one. It ends by saving a profile that becomes your active context
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.
+`kubectl` access. Run `kinetic down` when you're done.
## Recommended first run
@@ -76,8 +76,10 @@ 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
+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
@@ -91,7 +93,8 @@ guide.
| How do I get my first job running? | [Getting Started](https://kinetic.readthedocs.io/en/latest/getting_started.html) |
| When should I use `run_async()` instead of `run()`? | [Detached Jobs](https://kinetic.readthedocs.io/en/latest/guides/async_jobs.html) |
| How do I ship data and persist outputs? | [Data](https://kinetic.readthedocs.io/en/latest/guides/data.html) and [Checkpointing](https://kinetic.readthedocs.io/en/latest/guides/checkpointing.html) |
-| Bundled vs prebuilt vs custom image — which one? | [Execution Modes](https://kinetic.readthedocs.io/en/latest/guides/execution_modes.html) |
+| How does Kinetic work end to end? | [How Kinetic Works](https://kinetic.readthedocs.io/en/latest/concepts.html) |
+| Bundled vs prebuilt vs custom image — which one? | [Container Images](https://kinetic.readthedocs.io/en/latest/guides/containers.html) |
| Something's broken; where do I start? | [Troubleshooting](https://kinetic.readthedocs.io/en/latest/troubleshooting.html) |
## Configuration
diff --git a/docs/accelerators.md b/docs/accelerators.md
index 640359e9..22b51a21 100644
--- a/docs/accelerators.md
+++ b/docs/accelerators.md
@@ -1,73 +1,140 @@
-# Accelerator Support
+# Accelerators
+
+The `accelerator` argument of `@kinetic.run()` and the `--accelerator`
+flag of `kinetic pool add` take the same strings. This page lists every
+name that Kinetic accepts, and the topology behind each TPU name.
:::{important}
-Each accelerator and topology requires setting up its own node pool as a prerequisite.
+A job runs only on a node pool with the same accelerator type and, for
+TPUs, the same topology. Add one node pool for each accelerator that you
+use. See [Clusters and Node Pools](guides/clusters.md).
:::
+Run `kinetic accelerators` to print the names, the chip counts, and the
+topologies (TPUs) or the counts and machine types (GPUs) from your
+installed version. This page adds the machine type and the host count for
+each TPU slice. Add `--live` to mark the accelerators that have a node
+pool on your cluster.
+
+## Name formats
+
+| Format | Examples | Meaning |
+| ------ | -------- | ------- |
+| CPU | `cpu` | A CPU-only node. |
+| GPU name | `gpu-l4`, `l4`, `gpu:l4` | One GPU of that type. |
+| GPU name with count | `gpu-a100x4`, `a100x4` | That many GPUs on one node. |
+| TPU name with chip count | `tpu-v5litepod-8`, `v5litepod-8`, `tpu:v5litepod-8` | A slice with that many chips. |
+| TPU name with topology | `tpu-v5litepod-2x4`, `tpu-v4-2x2x2` | The same slice, named by topology. |
+| TPU name only | `tpu-v5litepod`, `tpu-v6e` | The default chip count of that type: `v5litepod-4`, `v6e-8`, `v5p-8`, `v4-4`, `v3-4`. |
+| Count only | `gpu:4`, `tpu:8` | The most capable type that supports that count. GPUs: `h100`, then `a100-80gb`, `a100`, `l4`, `v100`, `t4`, `p100`, `p4`. TPUs: `v6e`, then `v5p`, `v5litepod`, `v4`, `v3`. |
+| Spot suffix | `gpu-l4:spot` | The same as `spot=True`. |
+
+The `gpu-` and `tpu-` prefixes are optional. `v5e` is an alias for
+`v5litepod`. Names are not case-sensitive.
+
## TPUs
-| Type | Configurations |
-| -------------- | ----------------------------------------------------------------------------------------------------------------------------- |
-| TPU v6e | `v6e-8`, `v6e-16` |
-| TPU v5p | `v5p-8`, `v5p-16`, `v5p-32` |
-| TPU v5 Litepod | `v5litepod-1`, `v5litepod-4`, `v5litepod-8`, `v5litepod-16`, `v5litepod-32`, `v5litepod-64`, `v5litepod-128`, `v5litepod-256` |
-| TPU v4 | `v4-4`, `v4-8`, `v4-16`, `v4-32`, `v4-64`, `v4-128`, `v4-256`, `v4-512`, `v4-1024`, `v4-2048`, `v4-4096` |
-| TPU v3 | `v3-4`, `v3-16`, `v3-32`, `v3-64`, `v3-128`, `v3-256`, `v3-512`, `v3-1024`, `v3-2048` |
+The **Hosts** column decides the backend. A slice with more than one
+host is multi-host: Kinetic runs it on the `pathways` backend, one pod per
+host, and startup takes longer. See
+[Distributed Training](guides/distributed_training.md).
+
+| Type | Name | Topology | Machine type | Hosts |
+| ---- | ---- | -------- | ------------ | ----- |
+| TPU v6e | `v6e-8` | 2x4 | `ct6e-standard-4t` | 2 |
+| | `v6e-16` | 4x4 | `ct6e-standard-4t` | 4 |
+| TPU v5p | `v5p-8` | 2x2x2 | `ct5p-hightpu-4t` | 2 |
+| | `v5p-16` | 2x2x4 | `ct5p-hightpu-4t` | 4 |
+| | `v5p-32` | 2x4x4 | `ct5p-hightpu-4t` | 8 |
+| TPU v5e (`v5litepod`) | `v5litepod-1` | 1x1 | `ct5lp-hightpu-1t` | 1 |
+| | `v5litepod-4` | 2x2 | `ct5lp-hightpu-4t` | 1 |
+| | `v5litepod-8` | 2x4 | `ct5lp-hightpu-8t` | 1 |
+| | `v5litepod-16` | 4x4 | `ct5lp-hightpu-4t` | 4 |
+| | `v5litepod-32` | 4x8 | `ct5lp-hightpu-4t` | 8 |
+| | `v5litepod-64` | 8x8 | `ct5lp-hightpu-4t` | 16 |
+| | `v5litepod-128` | 8x16 | `ct5lp-hightpu-4t` | 32 |
+| | `v5litepod-256` | 16x16 | `ct5lp-hightpu-4t` | 64 |
+| TPU v4 | `v4-4` | 2x2x1 | `ct4p-hightpu-4t` | 1 |
+| | `v4-8` | 2x2x2 | `ct4p-hightpu-4t` | 2 |
+| | `v4-16` | 2x2x4 | `ct4p-hightpu-4t` | 4 |
+| | `v4-32` | 2x4x4 | `ct4p-hightpu-4t` | 8 |
+| | `v4-64` | 4x4x4 | `ct4p-hightpu-4t` | 16 |
+| | `v4-128` | 4x4x8 | `ct4p-hightpu-4t` | 32 |
+| | `v4-256` | 4x8x8 | `ct4p-hightpu-4t` | 64 |
+| | `v4-512` | 8x8x8 | `ct4p-hightpu-4t` | 128 |
+| | `v4-1024` | 8x8x16 | `ct4p-hightpu-4t` | 256 |
+| | `v4-2048` | 8x16x16 | `ct4p-hightpu-4t` | 512 |
+| | `v4-4096` | 16x16x16 | `ct4p-hightpu-4t` | 1024 |
+| TPU v3 | `v3-4` | 2x2 | `ct3-hightpu-4t` | 1 |
+| | `v3-16` | 4x4 | `ct3p-hightpu-4t` | 4 |
+| | `v3-32` | 4x8 | `ct3p-hightpu-4t` | 8 |
+| | `v3-64` | 8x8 | `ct3p-hightpu-4t` | 16 |
+| | `v3-128` | 8x16 | `ct3p-hightpu-4t` | 32 |
+| | `v3-256` | 16x16 | `ct3p-hightpu-4t` | 64 |
+| | `v3-512` | 16x32 | `ct3p-hightpu-4t` | 128 |
+| | `v3-1024` | 32x32 | `ct3p-hightpu-4t` | 256 |
+| | `v3-2048` | 32x64 | `ct3p-hightpu-4t` | 512 |
+
+For a first TPU job, use a single-host slice: `v5litepod-1`,
+`v5litepod-4`, or `v5litepod-8`.
## GPUs
-| Type | Aliases | Multi-GPU Counts |
-| ---------------- | ------------------------------- | ---------------- |
-| NVIDIA H100 | `h100`, `nvidia-h100-80gb` | 1, 2, 4, 8 |
-| NVIDIA A100 80GB | `a100-80gb`, `nvidia-a100-80gb` | 1, 2, 4, 8, 16 |
-| NVIDIA A100 | `a100`, `nvidia-tesla-a100` | 1, 2, 4, 8, 16 |
-| NVIDIA L4 | `l4`, `nvidia-l4` | 1, 2, 4, 8 |
-| NVIDIA V100 | `v100`, `nvidia-tesla-v100` | 1, 2, 4, 8 |
-| NVIDIA T4 | `t4`, `nvidia-tesla-t4` | 1, 2, 4 |
-| NVIDIA P100 | `p100`, `nvidia-tesla-p100` | 1, 2, 4 |
-| NVIDIA P4 | `p4`, `nvidia-tesla-p4` | 1, 2, 4 |
+| Type | Names | GPU counts | Machine types |
+| ---- | ----- | ---------- | ------------- |
+| NVIDIA H100 80GB | `h100`, `nvidia-h100-80gb` | 1, 2, 4, 8 | `a3-highgpu-{1,2,4,8}g` |
+| NVIDIA A100 80GB | `a100-80gb`, `nvidia-a100-80gb` | 1, 2, 4, 8, 16 | `a2-ultragpu-{1,2,4,8,16}g` |
+| NVIDIA A100 40GB | `a100`, `nvidia-tesla-a100` | 1, 2, 4, 8, 16 | `a2-highgpu-{1,2,4,8}g`, `a2-megagpu-16g` |
+| NVIDIA L4 | `l4`, `nvidia-l4` | 1, 2, 4, 8 | `g2-standard-{4,24,48,96}` |
+| NVIDIA V100 | `v100`, `nvidia-tesla-v100` | 1, 2, 4, 8 | `n1-standard-{8,16,32,64}` |
+| NVIDIA T4 | `t4`, `nvidia-tesla-t4` | 1, 2, 4 | `n1-standard-{4,8,16}` |
+| NVIDIA P100 | `p100`, `nvidia-tesla-p100` | 1, 2, 4 | `n1-standard-{4,8,16}` |
+| NVIDIA P4 | `p4`, `nvidia-tesla-p4` | 1, 2, 4 | `n1-standard-{4,8,16}` |
-For multi-GPU configurations on GKE, append the count: `a100x4`, `l4x2`, etc.
+Append `xN` for more than one GPU on one node: `a100x4`, `l4x2`. Every
+GPU job is single-host.
## CPU
-Use `accelerator="cpu"` to run on a CPU-only node (no accelerator attached).
+`accelerator="cpu"` runs the job on a CPU node without an accelerator.
+Every cluster can run CPU jobs. Use `cpu` to test a script before you
+request hardware.
-## Capacity Reservations
+## Capacity
-Newer accelerators (TPU v6e, H100) can have limited on-demand availability. If `kinetic pool add` fails to provision nodes, use a GCP capacity reservation to guarantee hardware:
+Newer accelerators, such as TPU v6e and H100, can have no on-demand
+capacity in a zone. The job then stays `PENDING`. A capacity reservation
+guarantees the hardware for a node pool:
```bash
-kinetic pool add --accelerator tpu-v6e-8 --reservation my-v6e-reservation --project your-project-id
+kinetic pool add --accelerator tpu-v6e-16 --reservation my-v6e-reservation
```
-:::{seealso}
-See the [Capacity Reservations](guides/reservations.md) guide for details.
-:::
+See [Capacity Reservations](guides/reservations.md).
## Related pages
::::{grid} 1 1 2 2
:gutter: 3
-:::{grid-item-card} {octicon}`graph;1em` Cost Optimization
-:link: guides/cost_optimization
+:::{grid-item-card} {octicon}`server;1em` Clusters and Node Pools
+:link: guides/clusters
:link-type: doc
-Spot vs reservations vs on-demand for each accelerator class.
+Add a node pool for each accelerator that you use.
:::
:::{grid-item-card} {octicon}`cpu;1em` Distributed Training
:link: guides/distributed_training
:link-type: doc
-Which TPU topologies are multi-host and need Pathways.
+What changes when a slice spans more than one host.
:::
-:::{grid-item-card} {octicon}`server;1em` Multiple Clusters
-:link: guides/clusters
+:::{grid-item-card} {octicon}`graph;1em` Cost Optimization
+:link: guides/cost_optimization
:link-type: doc
-When to put accelerator pools in separate clusters.
+Spot VMs, reservations, and on-demand capacity.
:::
::::
diff --git a/docs/architecture.md b/docs/architecture.md
index 7a90d9ba..1298a5ec 100644
--- a/docs/architecture.md
+++ b/docs/architecture.md
@@ -1,25 +1,120 @@
# Architecture Overview
-Kinetic automates the process of running Python functions on Google Cloud
-Platform (GCP) accelerators. It handles packaging, infrastructure provisioning,
-and execution management to provide a seamless experience for remote workloads.
+This page is for contributors. It maps the execution pipeline to the
+modules in the `kinetic` package. For the user-level model, see
+[How Kinetic Works](concepts.md).
-## Execution Lifecycle
+## Modules
-When a function decorated with `@kinetic.run()` is executed (either directly
-for a synchronous run, or via `run_async()` for a detached job), the system
-follows these steps:
+```text
+kinetic/
+├── core/ # @run decorator, RemoteCallable, accelerator registry and parser
+├── backend/ # JobContext, GKE and Pathways backends, k8s helpers, log streaming
+├── data/ # Data class, content hashing, data references
+├── infra/ # Container image building and caching (Cloud Build)
+├── runner/ # remote_runner.py — the entry point inside the pod
+├── utils/ # Packager (payload, context.zip, packaging plan) and Cloud Storage helpers
+├── jobs.py # JobHandle, attach(), list_jobs()
+├── collections.py # run_async_map(), BatchHandle, attach_batch()
+├── debug.py # debugpy attach and port-forward helpers
+├── cli/ # The `kinetic` command
+│ ├── commands/ # init, up, down, status, config, pool, jobs, build-image, profile, accelerators
+│ ├── infra/ # Pulumi program, stack and state management, post-deploy steps
+│ ├── profiles.py # Profile store and the resolve_infra() precedence chain
+│ └── options.py # Shared --project/--zone/--cluster options
+├── credentials.py # gcloud, ADC, and kubeconfig checks and setup
+└── constants.py # Zone and region helpers, default names
+```
+
+## Execution lifecycle
+
+A call to a decorated function, direct or through `run_async()`, goes
+through these steps:
:::{container} kinetic-steps
-1. **Context Resolution**: Kinetic aggregates function parameters, environment variables, and local configurations into a unified `JobContext`.
-2. **Credential Validation**: The system verifies active `gcloud` and `kubectl` credentials, performing automatic configuration where necessary to ensure access to GCP services.
-3. **Artifact Preparation**:
- * **Data Dependencies**: Local data paths are hashed and uploaded to Google Cloud Storage (GCS) if they are not already present in the content-addressed cache.
- * **Function Serialization**: Kinetic serializes the decorated function and its closure with `cloudpickle`. Kinetic registers the modules under the package root for serialization by value. The remote pod thus does not import your first-party code to deserialize the job.
- * **Project Packaging**: Kinetic compresses the package root into a ZIP archive. To find the package root, Kinetic first walks up out of any Python package. Kinetic then walks up to the nearest directory with a project marker. The markers are `pyproject.toml`, `requirements.txt`, `setup.py`, `setup.cfg`, and `.git`. `KINETIC_PACKAGE_ROOT` replaces this search. Kinetic excludes the paths that the Data API controls, the default exclusion list, and the `.kineticignore` patterns.
- * **Packaging Plan**: Kinetic writes a packaging plan into the archive at `.kinetic/plan.json`. The plan holds the `sys.path` entries and the client working directory, both relative to the package root. The runner reads the plan to build its `sys.path` and to change its working directory. See [What Ships to the Pod](guides/packaging.md).
-4. **Container Image Management**: Kinetic generates a hash of project dependencies (e.g., `requirements.txt` or `pyproject.toml`). If a corresponding image does not exist in Artifact Registry, Kinetic initiates a Cloud Build job to create it.
-5. **Job Submission**: Based on the requested accelerator type, Kinetic submits a Kubernetes Job (for GKE) or a LeaderWorkerSet (for multi-host Pathways) to the target cluster.
-6. **Remote Execution**: The remote pod pulls the container image, retrieves the serialized artifacts, mounts the required data volumes, and executes the function.
-7. **Result Retrieval**: Upon completion, the function's return value is retrieved from GCS, deserialized, and returned to the local Python process.
+1. **Context resolution.** `resolve_infra()` in `cli/profiles.py`
+ resolves the project, zone, cluster, and namespace from the decorator
+ arguments, the `KINETIC_*` environment variables, the active profile,
+ and the defaults, in that order. `JobContext.from_params()` in
+ `backend/execution.py` collects every other job setting into one
+ mutable `JobContext`.
+2. **Credential validation.** `credentials.py` verifies the `gcloud`
+ login, Application Default Credentials, and the `kubeconfig` entry
+ for the cluster, and configures them when it can.
+3. **Artifact preparation.** `_prepare_artifacts()` in
+ `backend/execution.py`:
+ - Uploads each `Data` object to the content-addressed cache in the
+ jobs bucket, and replaces the object with a data reference.
+ - Registers each module under the package root for serialization by
+ value, and serializes the function, its arguments, and the captured
+ environment variables with `cloudpickle` into `payload.pkl`.
+ - Resolves the package root: up out of every `__init__.py` package,
+ then up to the nearest `pyproject.toml`, `requirements.txt`,
+ `setup.py`, `setup.cfg`, or `.git`. `KINETIC_PACKAGE_ROOT` replaces
+ the search. Archives the root into `context.zip`, without the
+ `Data` paths, the default exclusions, and the `.kineticignore`
+ matches.
+ - Writes a packaging plan into the archive at `.kinetic/plan.json`.
+ The plan holds the client `sys.path` entries and the client working
+ directory, relative to the package root. See
+ [What Ships to the Pod](guides/packaging.md).
+4. **Container image and upload, in parallel.**
+ `infra/container_builder.py` hashes the base image, the accelerator
+ category, the Kinetic version, the filtered dependency file, the
+ runner script, and the Dockerfile template. If Artifact Registry has
+ no image with that tag, it runs Cloud Build. Prebuilt mode resolves a
+ base image instead; custom image mode uses the URI as given.
+5. At the same time, `utils/storage.py` uploads `payload.pkl`,
+ `context.zip`, and, in prebuilt mode, `requirements.txt` to
+ `gs://{jobs bucket}/{job_id}/`. The client records the SHA-256 hash of
+ the payload and of the archive.
+6. **Submission.** `GKEBackend` creates a Kubernetes Job for a
+ single-host accelerator. `PathwaysBackend` creates a LeaderWorkerSet
+ for a multi-host TPU slice. `kinetic.run()` selects the backend from
+ `TpuConfig.num_nodes` unless `backend=` is set. Both backends pass
+ the artifact hashes in the pod specification.
+7. **Remote execution.** `runner/remote_runner.py` downloads both
+ artifacts, verifies the hashes, extracts the archive, rebuilds
+ `sys.path` and the working directory from the plan, installs the
+ dependency file in prebuilt mode, resolves the `Data` references and
+ volumes, applies the captured environment variables, sets
+ `KINETIC_OUTPUT_DIR`, and calls the function.
+8. **Result retrieval.** The runner writes a result payload to
+ `gs://{jobs bucket}/{job_id}/result.pkl`. `JobHandle.result()`
+ downloads it, returns the value or raises the exception with the
+ remote traceback, and, by default, deletes the Kubernetes resource
+ and, when the job succeeded, the job artifacts.
:::
+
+## Result payload
+
+```python
+{
+ "success": bool,
+ "result": Any, # if success is True
+ "exception": Exception, # if success is False
+ "traceback": str, # if success is False
+}
+```
+
+When one of its own startup phases fails, the runner writes a failure
+payload with a `phase` field. When the return value cannot be pickled,
+the runner writes a payload with the flag `serialization_failed`. See
+[Troubleshooting](troubleshooting.md).
+
+## Backend selection
+
+- CPU, GPU, and single-host TPU → `GKEBackend` (Kubernetes Job).
+- Multi-host TPU (`TpuConfig.num_nodes > 1`, one node per host) →
+ `PathwaysBackend` (LeaderWorkerSet).
+- An explicit `backend=` argument overrides the selection.
+
+## Infrastructure state
+
+The CLI keeps three layers of state: the in-memory `InfraConfig`, the
+Pulumi stack in the state bucket `gs://{project}-kinetic-state`, and the
+Google Cloud resources. Each `(project, cluster)` pair has its own stack,
+named `{project}-{cluster}`. All stack operations go through
+`cli/infra/state.py`: `load_state()`, `apply_update()`, and
+`apply_destroy()`. See `AGENTS.md` in the repository for the conventions
+that contributors follow.
diff --git a/docs/concepts.md b/docs/concepts.md
new file mode 100644
index 00000000..72dbfa33
--- /dev/null
+++ b/docs/concepts.md
@@ -0,0 +1,271 @@
+# How Kinetic Works
+
+This page explains the model behind Kinetic. It describes the pieces on
+your machine and in your Google Cloud project, what a job is, what
+travels with a job, and what comes back. Read this page after your
+[first run](getting_started.md) and before the guides. Each section names
+the guide that covers the details.
+
+## The one-sentence model
+
+You write a Python function. You decorate the function with
+`@kinetic.run(accelerator=...)`. When you call the function, Kinetic runs
+the function on that accelerator in your Google Cloud project and returns
+the return value to your local process.
+
+```python
+import kinetic
+
+
+@kinetic.run(accelerator="tpu-v5litepod-4")
+def train(learning_rate):
+ import keras # heavy imports go inside the function
+
+ ...
+ return final_loss
+
+
+loss = train(0.001) # runs on a 4-chip TPU v5e slice, returns a float
+```
+
+Nothing else about your code changes. Kinetic handles the packaging, the
+container image, the cluster, and the transfer of the result.
+
+## The pieces
+
+Kinetic has three parts. Two parts are on your machine, and one part is
+in your Google Cloud project.
+
+**On your machine**
+
+- The `kinetic` Python package. It gives you `@kinetic.run()`,
+ `kinetic.Data`, and the job APIs.
+- The `kinetic` command-line tool. It creates and manages the cloud
+ resources, and it inspects jobs.
+- The **active profile**. A profile is a saved set of four values:
+ project, zone, cluster, and Kubernetes namespace. `kinetic init` creates
+ the first profile. Both the package and the command-line tool read the
+ active profile, so you do not repeat those values in code or in
+ commands. See [Profiles](guides/profiles.md).
+
+**In your Google Cloud project** (created by `kinetic up`, one set per
+cluster, plus one state bucket per project)
+
+- A **GKE cluster**. Default name: `kinetic-cluster`.
+- One or more **accelerator node pools**. Each node pool holds VMs of
+ one accelerator type. A node pool scales to zero when no job runs. A
+ job can run only on a node pool with the same accelerator type. The
+ cluster also has a small default pool of CPU nodes for system pods and
+ for `accelerator="cpu"` jobs. See
+ [Clusters and Node Pools](guides/clusters.md).
+- An **Artifact Registry repository** for the container images that
+ Kinetic builds.
+- Two **Cloud Storage buckets**: a jobs bucket for job artifacts and
+ outputs, and a builds bucket for Cloud Build. Both buckets delete
+ objects that are older than 30 days.
+- A **Cloud NAT gateway**, which gives the private cluster nodes access
+ to the internet.
+- A shared **state bucket**, one per project, where the command-line tool
+ keeps the infrastructure state. Teammates with access to this bucket
+ see the same clusters.
+
+**Inside a job**
+
+- The **pod**. A Kubernetes pod that runs one container on one node of
+ the node pool. The pod downloads your code, runs your function, and
+ uploads the result. A multi-host TPU job uses one pod per host.
+- The **backend**. Kinetic has two: the `gke` backend runs a single-host
+ job as a Kubernetes Job, and the `pathways` backend runs a multi-host
+ TPU slice as a LeaderWorkerSet, one pod per host. Kinetic selects the
+ backend from the accelerator. See
+ [Distributed Training](guides/distributed_training.md).
+
+## The decorator
+
+`@kinetic.run()` takes every setting that a job needs. The parameters
+fall into five groups:
+
+| Group | Parameters | Notes |
+| ----- | ---------- | ----- |
+| Hardware | `accelerator`, `spot`, `backend` | `accelerator` names the hardware and the slice size, for example `"tpu-v5litepod-4"`, `"gpu-l4"`, `"gpu-a100x4"`, or `"cpu"`. `spot=True` requests Spot capacity, and needs a Spot node pool. Kinetic selects `backend` for you. |
+| Where | `project`, `zone`, `cluster`, `namespace` | Leave these unset. Kinetic reads them from the active profile. Set one only for a one-off override. |
+| Inputs | `volumes`, `capture_env_vars` | `volumes` mounts `kinetic.Data(...)` at fixed paths. `capture_env_vars` copies named local environment variables into the pod. Function arguments are the other input path. |
+| Outputs | `output_dir` | The Cloud Storage location that the pod sees as `KINETIC_OUTPUT_DIR`. Defaults to a per-job prefix in the jobs bucket. |
+| Advanced | `container_image`, `base_image_repo`, `debug` | `container_image` changes how Kinetic produces the image. `debug=True` attaches a debugger. |
+
+See the [API reference](api.rst) for every parameter and its default.
+
+## Three ways to call a decorated function
+
+| Call | Returns | Use it when |
+| ---- | ------- | ----------- |
+| `train(...)` — a **blocking call** | The return value of the function | The job is short, or you iterate interactively. The call blocks until the job ends and streams the logs. |
+| `train.run_async(...)` — a **detached job** | A `JobHandle`, as soon as Kinetic submitted the job | The job runs for more than a few minutes. You poll status, tail logs, and collect the result later, from any machine. See [Detached Jobs](guides/async_jobs.md). |
+| `train.run_async_map(inputs, ...)` — a **batch** | A `BatchHandle`, as soon as Kinetic started the submissions | You run the same function over many inputs, for example a hyperparameter sweep. See [Batched Jobs](guides/batched_jobs.md). |
+
+Submission includes the packaging and, when needed, the image build.
+
+`kinetic.attach(job_id)` rebuilds a `JobHandle` from a job ID such as
+`job-3f9a1c2b`, and
+`kinetic.list_jobs()` lists the jobs on the cluster. The `kinetic jobs`
+command group offers the same operations from the shell.
+
+## What travels with a job
+
+Kinetic uploads three kinds of thing to the jobs bucket for every job,
+under `gs://{jobs bucket}/{job_id}/`:
+
+1. **Your function**, in `payload.pkl`. Kinetic serializes the function
+ with `cloudpickle`, together with its arguments and the environment
+ variables that you capture. Kinetic serializes the modules of your own
+ project by value, so the pod does not import them to load the job.
+2. **Your project source**, in `context.zip`. Kinetic starts at the
+ **entry directory**, the directory of the file that defines the
+ decorated function. Kinetic then finds the **package root**: the
+ nearest directory at or above the entry directory that holds a
+ `pyproject.toml`, a `requirements.txt`, a `setup.py`, a `setup.cfg`,
+ or a `.git` entry. Kinetic archives that directory. On the pod, the
+ runner extracts the archive into a **workspace**, rebuilds `sys.path`,
+ and changes to the workspace directory that matches your working
+ directory. Imports and relative paths therefore behave as they do on
+ your machine. See [What Ships to the Pod](guides/packaging.md).
+3. **Your data**, if you wrap a path in `kinetic.Data(...)`. Kinetic
+ uploads local data one time, keyed by a content hash, and gives your
+ function a plain filesystem path on the pod. See
+ [Working with Data](guides/data.md).
+
+Kinetic reads one more file: your **dependency file**, a
+`requirements.txt` or a `pyproject.toml`. That file decides which
+packages the image contains. A `pip install` in your local shell does not
+carry over. See [Dependencies](guides/dependencies.md).
+
+## The container image
+
+The pod runs a container image. By default, Kinetic builds that image for
+you with Cloud Build. The image starts from a Python image that matches
+your local Python version. The image contains the accelerator runtime
+(JAX with `libtpu` or CUDA), Keras, Kinetic, and the packages from your
+dependency file. Kinetic tags the image with a hash of those inputs and
+stores it in the Artifact Registry repository of the cluster.
+
+Two consequences follow:
+
+- The first run with a given dependency set takes about 5 to 10 minutes,
+ because Cloud Build runs.
+- Every later run with the same dependency file skips the build. The
+ pod starts in less than 1 minute while a node is still running, or
+ after the 2 to 5 minutes that a new node needs. A change to the
+ dependency file causes a new build.
+
+Kinetic also has other image modes, for a slow build step or for system
+libraries. You do not need them at first. See
+[Container Images](guides/containers.md).
+
+## The life of a job
+
+Every call goes through the same five phases:
+
+:::{container} kinetic-steps
+1. **Package.** Kinetic resolves the profile, serializes the function,
+ archives the package root, uploads the `Data` objects, and uploads
+ the artifacts to `gs://{jobs bucket}/{job_id}/`.
+2. **Build.** Kinetic reuses a cached image or runs Cloud Build.
+3. **Schedule.** Kinetic creates a Kubernetes Job on the cluster, or a
+ LeaderWorkerSet for a multi-host TPU slice. The cluster autoscaler
+ starts a node in the matching node pool. The job is `PENDING`.
+4. **Run.** The pod downloads the artifacts, extracts the source,
+ resolves the `Data` paths, sets `KINETIC_OUTPUT_DIR`, and calls your
+ function. The job is `RUNNING`. Kinetic streams the pod log to your
+ terminal.
+5. **Collect.** The pod uploads the return value (or the exception) to
+ `gs://{jobs bucket}/{job_id}/result.pkl`. The job is `SUCCEEDED` or
+ `FAILED`. Kinetic downloads the value and returns it, or raises the
+ exception with the remote traceback.
+:::
+
+The pod exits as soon as your function returns. After a blocking call or
+a `result()` call collects the result, Kinetic deletes the Kubernetes Job
+and, on success, the uploaded artifacts. A detached job that nobody
+collects keeps its artifacts until you call `result()` or `cleanup()`, or
+until the 30-day rule of the jobs bucket deletes them. Kubernetes deletes
+a finished Job resource 10 minutes after it ends, without a call from
+you. Kinetic never deletes what you wrote under `KINETIC_OUTPUT_DIR` as
+part of job cleanup, but the 30-day rule of the jobs bucket applies.
+
+## Where the results go
+
+A job produces three kinds of output. Each kind has a different channel:
+
+| Output | Channel | Size |
+| ------ | ------- | ---- |
+| The return value | `result.pkl` in the jobs bucket, then your local process | Small: a metric, a dict, a path |
+| Files that you keep | `KINETIC_OUTPUT_DIR`, a per-job Cloud Storage prefix | Any size: exported models, evaluation results |
+| Checkpoints | A stable subdirectory under `KINETIC_OUTPUT_DIR`, with `output_dir=` set to a fixed location | Any size; a restart reads them |
+
+Write everything that you want to keep under `KINETIC_OUTPUT_DIR`. The
+pod discards its filesystem, including `/tmp`, when it ends. The default
+output directory is different for every job. To resume a job from its
+checkpoints, set `output_dir=` to a fixed location. See
+[Outputs and Checkpoints](guides/checkpointing.md).
+
+## Where the settings come from
+
+Every job needs a project, a zone, a cluster, and a namespace. Kinetic
+resolves each value in this order, and the first value wins:
+
+1. The decorator argument (`project=`) or the CLI flag (`--project`).
+2. The environment variable (`KINETIC_PROJECT`).
+3. The active profile.
+4. The built-in default: `us-central1-a`, `kinetic-cluster`, and
+ `default`. The project has no default. Without a profile or an
+ override, Kinetic raises `Project must be specified`.
+
+The profile is the layer that you set one time. The environment variable
+and the flag are for a one-off override.
+`kinetic config` prints each resolved value and its source. See
+[Configuration](configuration.md).
+
+## What it costs
+
+You pay for accelerator nodes only while a job runs. Each accelerator
+node pool scales to zero after about 10 minutes without a job. Three
+costs continue while the cluster exists. The first is the GKE control
+plane, which Google Cloud covers for one cluster with a monthly credit.
+The second is one small `e2-standard-4` node in the default pool. The
+third is the Cloud NAT gateway. Cloud Build charges for each
+image build. When you no longer need a cluster, `kinetic down` deletes
+it. See [Cost Optimization](guides/cost_optimization.md).
+
+## Where to go next
+
+::::{grid} 1 1 2 2
+:gutter: 3
+
+:::{grid-item-card} {octicon}`database;1em` Working with Data
+:link: guides/data
+:link-type: doc
+
+Ship local files and read Cloud Storage data from your function.
+:::
+
+:::{grid-item-card} {octicon}`history;1em` Outputs and Checkpoints
+:link: guides/checkpointing
+:link-type: doc
+
+`KINETIC_OUTPUT_DIR`, retention, and resumable training.
+:::
+
+:::{grid-item-card} {octicon}`clock;1em` Detached Jobs
+:link: guides/async_jobs
+:link-type: doc
+
+Submit, monitor, reattach, and clean up long jobs.
+:::
+
+:::{grid-item-card} {octicon}`server;1em` Clusters and Node Pools
+:link: guides/clusters
+:link-type: doc
+
+Add accelerator pools, share a cluster, and run more than one cluster.
+:::
+::::
diff --git a/docs/configuration.md b/docs/configuration.md
index 61a88962..9dfbf2c5 100644
--- a/docs/configuration.md
+++ b/docs/configuration.md
@@ -1,149 +1,183 @@
# Configuration
-Kinetic uses environment variables, decorator arguments, CLI flags, and
-optionally [named profiles](guides/profiles.md) for configuration. This
-page is the source of truth for what each one does, what the defaults
-are, and how they come together when they disagree.
-
-:::{tip}
-If you work with more than one cluster or project, consider saving
-those combinations as [profiles](guides/profiles.md) — they remove the
-need to re-export `KINETIC_*` env vars each time you switch.
-:::
+This page is the reference for every setting that Kinetic reads: where
+each setting can come from, which source wins, and what the default is.
+For the everyday workflow, see [Profiles](guides/profiles.md).
+
+## The four sources
+
+Kinetic reads a setting from up to four sources. The first source that
+has a value wins:
+
+1. **A decorator argument or a CLI flag.** For example,
+ `@kinetic.run(project="p")` or `kinetic status --project p`. Use these
+ for a one-off override in one call.
+2. **A `KINETIC_*` environment variable.** For example,
+ `KINETIC_PROJECT=p`. Use these for one shell session, or in CI.
+3. **The active profile.** A profile stores the project, the zone, the
+ cluster, and the namespace. `kinetic init` creates the first profile.
+ This is the source that you set one time.
+4. **The built-in default.**
+
+Profiles cover only the four target settings. The other settings in this
+page come from a decorator argument, a flag, an environment variable, or
+the default.
+
+## Precedence table
+
+| Setting | Decorator argument | CLI flag | Environment variable | Active profile | Built-in default |
+| ------- | ------------------ | -------- | -------------------- | -------------- | ---------------- |
+| Project | `project=` | `--project` | `KINETIC_PROJECT` | `project` | `GOOGLE_CLOUD_PROJECT` (Python API only), else required |
+| Zone | `zone=` | `--zone` | `KINETIC_ZONE` | `zone` | `us-central1-a` |
+| Cluster | `cluster=` | `--cluster` | `KINETIC_CLUSTER` | `cluster` | `kinetic-cluster` |
+| Namespace | `namespace=` | `--namespace` (`kinetic jobs list`, `kinetic init`, `kinetic up`, `kinetic profile create`) | `KINETIC_NAMESPACE` | `namespace` | `default` |
+| Output directory | `output_dir=` | _(none)_ | `KINETIC_OUTPUT_DIR` | _(none)_ | `gs://{jobs bucket}/outputs/{job_id}` |
+| Base image repository | `base_image_repo=` | `kinetic build-image --repo` | `KINETIC_BASE_IMAGE_REPO` | _(none)_ | `kinetic` |
+| Reservation | _(none)_ | `kinetic pool add --reservation` | `KINETIC_RESERVATION` | _(none)_ | _(unset)_ |
+| Profile selection | _(none)_ | `kinetic --profile` | `KINETIC_PROFILE` | the stored `current` | _(none)_ |
+
+Read each row from left to right. A decorator argument beats a flag,
+which beats an environment variable, which beats the profile, which beats
+the default. For example:
+
+```python
+@kinetic.run(accelerator="tpu-v5litepod-4", project="explicit-project")
+def train(): ...
+```
+
+This job runs in `explicit-project`, even if `KINETIC_PROJECT` and the
+active profile name a different project.
+
+A reservation is a property of a node pool, not of a job. You bind a
+reservation to a pool with `kinetic pool add`, and every job that lands
+on that pool uses it. Jobs select pools through `accelerator=`.
## Environment variables
-| Variable | Used by | Default | Description |
-| ---------------------------- | ------------------------- | -------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------ |
-| `KINETIC_PROJECT` | CLI + decorators | _(required)_ | GCP project ID. Falls back to `GOOGLE_CLOUD_PROJECT` if unset. |
-| `KINETIC_ZONE` | CLI + decorators | `us-central1-a` | GCP zone for jobs and clusters. |
-| `KINETIC_CLUSTER` | CLI + decorators | `kinetic-cluster` | GKE cluster name. |
-| `KINETIC_NAMESPACE` | CLI + decorators | `default` | Kubernetes namespace. |
-| `KINETIC_BASE_IMAGE_REPO` | Decorator (prebuilt mode) | `kinetic` | Repo for prebuilt base images. See [Execution Modes](guides/execution_modes.md). |
-| `KINETIC_OUTPUT_DIR` | Library (submit) + remote pod | `gs://{bucket}/outputs/{job_id}` | Per-job durable artifact prefix. Kinetic reads this variable at submit time and sets it in the pod. Kinetic has no equivalent CLI flag. See [Checkpointing](guides/checkpointing.md). |
-| `KINETIC_RESERVATION` | `kinetic pool add` | _(unset)_ | GCP capacity reservation to consume. Pool-level config, not a per-job setting. |
-| `KINETIC_LOG_LEVEL` | Library | `INFO` | `DEBUG`, `INFO`, `WARNING`, `ERROR`, `FATAL`. |
-| `KINETIC_DEBUG_WAIT_TIMEOUT` | Library + remote pod | `600` | Seconds the remote pod waits for a debugger client to attach when `debug=True`. Applies on both sides (local `debug_attach()` and the pod's debugpy server). |
-| `KINETIC_PACKAGE_ROOT` | Library (submit) | _(auto-detected)_ | Set this variable to the directory that Kinetic puts into `context.zip`. The directory must exist. The directory must also contain the directory that defines your function. If one of these conditions is not true, Kinetic raises a `ValueError` at submit time. See [What Ships to the Pod](guides/packaging.md). |
-| `KINETIC_NO_DEFAULT_EXCLUDES`| Library (submit) | _(unset)_ | Set this variable to `1` to turn the default exclusions off. Kinetic then puts `.venv`, `node_modules`, and the cache directories into `context.zip`. Kinetic always excludes `.git` and `__pycache__`. |
-| `KINETIC_CONTEXT_SIZE_WARN_MB`| Library (submit) | `100` | Warning threshold in megabytes for `context.zip`. Above the threshold, Kinetic logs a warning and lists the five largest files. Set the value to `0` to turn the warning off. |
-| `KINETIC_PAYLOAD_SIZE_WARN_MB`| Library (submit) | `50` | Warning threshold in megabytes for `payload.pkl`. Above the threshold, Kinetic logs a warning about the arguments and the module-level globals that it captured by value. Set the value to `0` to turn the warning off. |
-
-Set them in your shell profile (`~/.bashrc`, `~/.zshrc`) so they
-persist across sessions:
+### Target selection
-```bash
-export KINETIC_PROJECT="my-gcp-project-id"
-export KINETIC_ZONE="us-central1-a"
-```
+These four variables have profile equivalents. Set them only for a
+one-off override.
-## Precedence
+| Variable | Default | Description |
+| -------- | ------- | ----------- |
+| `KINETIC_PROJECT` | _(none)_ | Google Cloud project ID. Required when no profile is active. The Python API also accepts `GOOGLE_CLOUD_PROJECT`, after the profile. The CLI does not read `GOOGLE_CLOUD_PROJECT`. |
+| `KINETIC_ZONE` | `us-central1-a` | Zone of the cluster. |
+| `KINETIC_CLUSTER` | `kinetic-cluster` | Name of the cluster. |
+| `KINETIC_NAMESPACE` | `default` | Kubernetes namespace for jobs. |
-When the same setting can come from multiple sources, the highest one
-wins:
+### Profile selection
-| Setting | Decorator arg | CLI flag | Env var | Active [profile](guides/profiles.md) | Built-in default |
-| --------------- | ------------------ | -------------------------------- | ----------------------------------------------- | ------------------------------------ | -------------------------------- |
-| Project | `project=` | `--project` | `KINETIC_PROJECT` (then `GOOGLE_CLOUD_PROJECT`) | `project` | _(required)_ |
-| Zone | `zone=` | `--zone` | `KINETIC_ZONE` | `zone` | `us-central1-a` |
-| Cluster | `cluster=` | `--cluster` | `KINETIC_CLUSTER` | `cluster` | `kinetic-cluster` |
-| Namespace | `namespace=` | `--namespace` | `KINETIC_NAMESPACE` | `namespace` | `default` |
-| Output dir | `output_dir=` | _(n/a)_ | `KINETIC_OUTPUT_DIR` | _(n/a)_ | `gs://{bucket}/outputs/{job_id}` |
-| Base image repo | `base_image_repo=` | `kinetic build-image --repo` | `KINETIC_BASE_IMAGE_REPO` | _(n/a)_ | `kinetic` |
-| Reservation\* | _(n/a)_ | `kinetic pool add --reservation` | `KINETIC_RESERVATION` | _(n/a)_ | _(unset)_ |
+| Variable | Default | Description |
+| -------- | ------- | ----------- |
+| `KINETIC_PROFILE` | _(unset)_ | Name of the profile to use for this process. Overrides the stored active profile. Both the CLI and the Python API read it. |
+| `KINETIC_PROFILES_FILE` | `~/.kinetic/profiles.json` | Path of the profile store. |
-\* Reservation is a node-pool-level setting, not a per-job one. You bind
-a reservation to a pool when you create the pool with `kinetic pool add`,
-and any job that lands on that pool consumes it. Because of that there is
-no decorator argument; jobs select pools indirectly via `accelerator=`.
+### Job behavior
-Read left to right: a decorator argument always beats a CLI flag, which
-beats an env var, which beats a profile field, which beats the built-in
-default. Concretely:
+| Variable | Default | Description |
+| -------- | ------- | ----------- |
+| `KINETIC_OUTPUT_DIR` | `gs://{jobs bucket}/outputs/{job_id}` | On your machine before a submit: the output directory for the job. In the pod: the resolved value, always set. See [Outputs and Checkpoints](guides/checkpointing.md). |
+| `KINETIC_BASE_IMAGE_REPO` | `kinetic` | Repository for prebuilt base images. Used only with `container_image="prebuilt"`. See [Container Images](guides/containers.md). |
+| `KINETIC_NO_TTY_DEBUG` | _(unset)_ | Set to `1` to permit a blocking call with `debug=True` when `stdin` is not a terminal. See [Interactive Debugging](guides/debugging.md). |
-```python
-@kinetic.run(accelerator="tpu-v6e-8", project="explicit-project")
-def train(): ...
-```
+### Packaging
+
+| Variable | Default | Description |
+| -------- | ------- | ----------- |
+| `KINETIC_PACKAGE_ROOT` | _(detected)_ | The directory that Kinetic archives into `context.zip`. The directory must exist, and it must be the directory of your function or a parent of it. Otherwise Kinetic raises a `ValueError` at submit time. See [What Ships to the Pod](guides/packaging.md). |
+| `KINETIC_NO_DEFAULT_EXCLUDES` | _(unset)_ | Set to `1` to turn off the default exclusions. Kinetic then archives `.venv`, `node_modules`, and the cache directories. `.git` and `__pycache__` stay excluded. |
+| `KINETIC_CONTEXT_SIZE_WARN_MB` | `100` | Size of `context.zip`, in megabytes, above which Kinetic logs a warning and lists the five largest files. `0` turns the warning off. |
+| `KINETIC_PAYLOAD_SIZE_WARN_MB` | `50` | Size of `payload.pkl`, in megabytes, above which Kinetic logs a warning about the arguments and globals that it captured by value. `0` turns the warning off. |
-uses `explicit-project` even if `KINETIC_PROJECT` is set to something
-else.
+### CLI only
+
+| Variable | Default | Description |
+| -------- | ------- | ----------- |
+| `KINETIC_RESERVATION` | _(unset)_ | Capacity reservation for `kinetic pool add`. See [Capacity Reservations](guides/reservations.md). |
+| `KINETIC_FORCE_DESTROY` | `true` | Whether `kinetic down` empties the buckets before it deletes them. `kinetic up --no-force-destroy` stores `false` in the stack. |
+| `KINETIC_LOG_LEVEL` | `INFO` | Log level of the `kinetic` package: `DEBUG`, `INFO`, `WARNING`, `ERROR`, or `FATAL`. |
+
+To make an environment variable persist, put the `export` line in your
+shell profile (`~/.bashrc` or `~/.zshrc`). For the target settings,
+prefer a profile.
## Logging
-Kinetic uses `absl-py` for logging. Set `KINETIC_LOG_LEVEL` to control
-verbosity:
+Kinetic logs with `absl-py`. `KINETIC_LOG_LEVEL` sets the level:
-- **DEBUG** — packaging details, dependency hashing, build pipeline,
- GKE submission.
-- **INFO** — major lifecycle milestones (default).
-- **WARNING / ERROR / FATAL** — only the named severity and above.
+- `DEBUG` — packaging details, dependency hashing, the build pipeline,
+ and the Kubernetes submission.
+- `INFO` — the main lifecycle events. This is the default.
+- `WARNING`, `ERROR`, `FATAL` — that level and above only.
```bash
export KINETIC_LOG_LEVEL=DEBUG
```
-## Pulumi state
+## See the resolved values
-Kinetic stores its Pulumi state in a Google Cloud Storage bucket
-derived from the GCP project: `gs://{project}-kinetic-state`. The
-bucket is created on first use (idempotent), with **versioning
-enabled** and **uniform bucket-level access**, no public ACL.
-Multiple clusters in one project share the bucket but get separate
-stacks (named `{project}-{cluster}`), so a team running against the
-same GCP project automatically converges on one authoritative state.
+`kinetic config` prints the active profile and, for the project, the
+zone, the cluster, the namespace, and the output directory, the resolved
+value and its source (`KINETIC_*`, `profile`, or `default`). The command
+also prints the state bucket of the project. Run it first when a setting
+does not take effect.
-### IAM
+```bash
+kinetic config
+```
+
+`kinetic config` cannot see a CLI flag or a decorator argument, because
+those apply to one call. The other variables in this page do not appear
+in the output. Inspect them with `env | grep KINETIC_`.
-Kinetic uses Application Default Credentials, the same auth path as
-`gcloud`. The first admin to run `kinetic up` for a project needs
-`roles/storage.admin` so the state bucket can be created. Every other
-team member only needs `roles/storage.objectAdmin` on the bucket to
-read and write state.
+## Infrastructure state
-## Where to look
+The `kinetic` CLI stores its Pulumi state in a Cloud Storage bucket named
+`gs://{project}-kinetic-state`. The first `kinetic up` in a project
+creates the bucket, with versioning and uniform bucket-level access, and
+without a public ACL. All clusters in the project share the bucket. Each
+cluster has its own stack, named `{project}-{cluster}`. A team that works
+in one project therefore sees one authoritative state.
-If a setting isn't behaving the way you expect, `kinetic config` prints
-the resolved value of the most common variables (project, zone,
-cluster, namespace, output dir, and the per-project Pulumi state
-bucket) and where each came from (env var,
-[profile](guides/profiles.md), or default). Run it before reaching
-for `kinetic init`'s troubleshoot path. Variables that aren't shown there
-(`KINETIC_BASE_IMAGE_REPO`, `KINETIC_RESERVATION`, `KINETIC_LOG_LEVEL`,
-`KINETIC_DEBUG_WAIT_TIMEOUT`) can be inspected with `env | grep
-KINETIC_`.
+### IAM
+
+Kinetic uses Application Default Credentials, the same login path as
+`gcloud`. The first person who runs `kinetic up` in a project needs
+`roles/storage.admin`, because that run creates the state bucket. Every
+other team member needs `roles/storage.objectAdmin` on the bucket to
+read and write the state.
## Related pages
::::{grid} 1 1 2 2
:gutter: 3
-:::{grid-item-card} {octicon}`rocket;1em` Getting Started
-:link: getting_started
+:::{grid-item-card} {octicon}`stack;1em` Profiles
+:link: guides/profiles
:link-type: doc
-Sets the canonical `KINETIC_PROJECT` once.
+The everyday way to set the project, zone, cluster, and namespace.
:::
-:::{grid-item-card} {octicon}`stack;1em` Profiles
-:link: guides/profiles
+:::{grid-item-card} {octicon}`server;1em` Clusters and Node Pools
+:link: guides/clusters
:link-type: doc
-Named bundles for project/zone/cluster/namespace; the ergonomic
-alternative to re-exporting env vars when you target multiple clusters.
+What `kinetic up` creates, and how a team shares it.
:::
:::{grid-item-card} {octicon}`terminal;1em` CLI Reference
:link: cli
:link-type: doc
-Generated reference for every flag.
+Generated reference for every command and flag.
:::
:::{grid-item-card} {octicon}`bug;1em` Troubleshooting
:link: troubleshooting
:link-type: doc
-What to check when a setting doesn't take effect.
+What to check when a setting does not take effect.
:::
::::
diff --git a/docs/contributing.md b/docs/contributing.md
index 10ea5f4c..fa71c7fb 100644
--- a/docs/contributing.md
+++ b/docs/contributing.md
@@ -1,24 +1,26 @@
# Contributing
-We'd love to accept your patches and contributions to this project.
+We welcome your patches and contributions to this project. This page
+explains the legal requirements, the development setup, and the review
+process.
## Before you begin
### Sign our Contributor License Agreement
-Contributions to this project must be accompanied by a
+Every contribution to this project needs a
[Contributor License Agreement](https://cla.developers.google.com/about) (CLA).
-You (or your employer) retain the copyright to your contribution; this simply
-gives us permission to use and redistribute your contributions as part of the
-project.
+You, or your employer, keep the copyright to your contribution. The CLA
+gives us the permission to use and to redistribute your contribution as
+part of the project.
:::{note}
-If you or your current employer have already signed the Google CLA (even if it
-was for a different project), you probably don't need to do it again.
+If you or your current employer already signed the Google CLA, for this
+project or for a different one, you do not need to sign it again.
:::
-Visit to see your current agreements or to
-sign a new one.
+Visit to see your current agreements
+or to sign a new one.
### Review our community guidelines
@@ -43,25 +45,26 @@ This project follows
### Code quality and testing
-Before submitting a pull request, please ensure your changes pass linting and unit tests.
+Before you open a pull request, make sure that your changes pass the
+linter and the unit tests.
-- **Linting:** We use [Ruff](https://docs.astral.sh/ruff/) for linting and formatting. Run it with:
+- **Lint:** We use [Ruff](https://docs.astral.sh/ruff/) for linting and formatting. Run it with:
```bash
ruff check .
```
-- **Unit tests:** We use [Pytest](https://docs.pytest.org/) for unit tests. Run them with:
+- **Unit tests:** We use [pytest](https://docs.pytest.org/) to run the unit tests. Run them with:
```bash
pytest
```
-- **E2E tests:** End-to-end tests run real workloads against a GKE cluster. They live in `tests/e2e/` and are skipped by default unless explicitly enabled.
+- **End-to-end tests:** These tests run real workloads on a GKE cluster. They live in `tests/e2e/`. pytest skips them unless you set `E2E_TESTS=1`.
**Prerequisites:**
- - A GCP project with a provisioned GKE cluster.
- - Google Cloud SDK authenticated (`gcloud auth login` and `gcloud auth application-default login`)
- - GKE credentials configured: `gcloud container clusters get-credentials --zone --project `
- - Test dependencies installed: `uv pip install -e ".[test]"`
+ - A Google Cloud project with a Kinetic cluster (`kinetic init`).
+ - A Google Cloud login (`gcloud auth login` and `gcloud auth application-default login`).
+ - `kubectl` credentials for the cluster. `kinetic init` and `kinetic up` configure them.
+ - The test dependencies: `uv pip install -e ".[test]"`.
**Required environment variables:**
@@ -85,57 +88,52 @@ Before submitting a pull request, please ensure your changes pass linting and un
```
:::{tip}
- Drop `-n auto` to run tests serially to make it easier to debug.
+ Remove `-n auto` to run the tests one at a time. Serial runs are easier to debug.
:::
### Submitting changes
-1. Fork the repository
-2. Create a feature branch: `git checkout -b feature/amazing-feature`
-3. Commit your changes: `git commit -m 'Add amazing feature'`
-4. Push to the branch: `git push origin feature/amazing-feature`
-5. Open a Pull Request
+1. Fork the repository.
+2. Create a feature branch: `git checkout -b feature/my-feature`.
+3. Commit your changes: `git commit -m 'Add my feature'`.
+4. Push the branch: `git push origin feature/my-feature`.
+5. Open a pull request.
### Code reviews
-All submissions, including submissions by project members, require review. We
-use GitHub pull requests for this purpose. Consult
-[GitHub Help](https://help.github.com/articles/about-pull-requests/) for more
-information on using pull requests.
+Every submission needs a review, including a submission from a project
+member. We use GitHub pull requests for reviews. See
+[GitHub Help](https://help.github.com/articles/about-pull-requests/) for
+more information about pull requests.
## Documentation Contribution Process
-```sh
-# Install docs libraries:
-pip install .[docs]
+Install the documentation dependencies, then build and serve the site
+locally:
-# Build and serve docs locally:
+```bash
+uv pip install -e ".[docs]"
sphinx-autobuild docs /tmp/docs
```
+The pages are MyST Markdown under `docs/`. Follow the style of the
+existing pages: short sentences in the active voice, present tense,
+imperative steps, and no contractions.
+
## Releases
-To release a new version of the package to PyPI, follow these steps:
+Only maintainers release Kinetic. The process is in
+[RELEASE_PROCESS.md](https://github.com/keras-team/kinetic/blob/main/RELEASE_PROCESS.md)
+in the repository. In short:
:::{container} kinetic-steps
-1. **Install the release dependencies.**
- ```bash
- uv pip install -e ".[release]"
- ```
-2. **Bump the version** in the following files:
- - [pyproject.toml](../pyproject.toml)
- - [version.py](../kinetic/version.py)
-3. **Build the source distribution and wheel.**
- ```bash
- python3 -m build
- ```
-4. **Upload the packages to PyPI** using `twine`. To avoid `twine` hanging while waiting for interactive input, provide your credentials via environment variables (e.g. using an API token) or a `~/.pypirc` file:
- ```bash
- TWINE_USERNAME=__token__ TWINE_PASSWORD=pypi-... python3 -m twine upload dist/*
- ```
-5. **Create a new release on GitHub** using the `gh` CLI, e.g.:
- ```bash
- gh release create 0.0.2
- ```
+1. **Bump the version** in `pyproject.toml` and `kinetic/version.py`
+ through a pull request.
+2. **Create a release branch** named after the version, for example
+ `r0.0.5`.
+3. **Create a GitHub release** from that branch at
+ .
+4. **Wait for the publish workflow.** The release tag starts the
+ `publish_to_pypi` GitHub Actions workflow, which uploads the package
+ to PyPI. Do not upload with `twine` yourself.
:::
-
diff --git a/docs/examples.md b/docs/examples.md
index 63e3ddec..555cd408 100644
--- a/docs/examples.md
+++ b/docs/examples.md
@@ -7,6 +7,7 @@ examples/fashion_mnist
examples/simple_demo
examples/example_async_jobs
examples/example_data_api
+examples/example_fuse
examples/example_checkpoint
examples/example_keras_checkpoint
examples/example_collections
@@ -17,27 +18,40 @@ examples/gemma3_sft_demo
examples/tunix_sft
```
-A catalog of runnable example scripts using Kinetic. Every example below is rendered directly on this site and is also available as a raw Python script in the GitHub repository.
+This page is a catalog of the runnable example scripts in the
+repository. Each example renders on this site and is also a Python file
+in the `examples/` directory of the
+[GitHub repository](https://github.com/keras-team/kinetic/tree/main/examples).
-Examples are grouped into three tiers:
+The examples come in three tiers:
-- {bdg-success}`Quickstart` — your first run. Minimal setup, sensible defaults.
-- {bdg-secondary}`Core` — the everyday product surface: async jobs, data,
- checkpoints, parallel sweeps.
-- {bdg-secondary}`Advanced` — multi-host Pathways jobs, LLM fine-tuning,
- anything that needs special quota or external credentials.
+- {bdg-success}`Quickstart` — the first run, with the defaults.
+- {bdg-secondary}`Core` — the everyday features: detached jobs, data,
+ checkpoints, and parallel sweeps.
+- {bdg-secondary}`Advanced` — multi-host Pathways jobs and LLM
+ fine-tuning. These need special quota or external credentials.
-To run any example: clone the repo, install Kinetic, set `KINETIC_PROJECT`,
-and `python examples/.py`.
+To run an example, clone the repository, install Kinetic, make sure that
+`kinetic init` has saved an active profile, and run the script:
```bash
git clone https://github.com/keras-team/kinetic.git
cd kinetic
uv pip install -e .
-export KINETIC_PROJECT="your-project-id"
+kinetic init # skip this step if you already have an active profile
python examples/fashion_mnist.py
```
+Each example names its accelerator in the decorator. Your cluster needs a
+node pool with that accelerator, or you change the `accelerator=` value.
+See [Getting Started](getting_started.md).
+
+The LLM examples import packages that the default image does not have,
+such as `keras-hub`, `tunix`, and `wandb`. Before you run one of them,
+put a `requirements.txt` with those packages next to the script, or in
+the `examples/` directory. Kinetic reads that file and builds an image
+with the packages. See [Dependencies](guides/dependencies.md).
+
## Quickstart
::::{grid} 1 2 2 3
@@ -50,9 +64,9 @@ python examples/fashion_mnist.py
:class-body: sd-fs-6
:class-title: sd-fs-5
-The first thing to run after `kinetic up`. A small Keras classifier on
-Fashion-MNIST that confirms your cluster can schedule a TPU pod and
-stream a real result back to your shell.
+The first script to run after `kinetic init`. A small Keras classifier
+on Fashion-MNIST that shows that your cluster can schedule a TPU pod and
+return a result to your shell.
+++
@@ -60,15 +74,15 @@ stream a real result back to your shell.
{bdg-secondary}`TPU`
:::
-:::{grid-item-card} Keras + JAX smoke test
+:::{grid-item-card} Keras + JAX on a CPU node
:link: examples/simple_demo.md
:class-card: sd-shadow-sm
:class-body: sd-fs-6
:class-title: sd-fs-5
-The cheapest sanity check there is. Keras-on-JAX on a CPU node — no
-accelerator quota needed, useful for verifying your install before you
-ask for hardware.
+The smallest possible check. Keras on JAX on a CPU node, without
+accelerator quota. Use it to test your install before you request
+hardware.
+++
@@ -90,9 +104,9 @@ ask for hardware.
:class-body: sd-fs-6
:class-title: sd-fs-5
-Walks through every part of the detached-job API end-to-end: `run_async()`,
-`status()`/`tail()`/`result()`, reattach from another shell with
-`kinetic.attach()`, and enumerate jobs with `list_jobs()`.
+The full detached-job API: `run_async()`, `status()`, `tail()`,
+`result()`, a reattach from another shell with `kinetic.attach()`, and
+`list_jobs()`.
+++
@@ -106,9 +120,9 @@ Walks through every part of the detached-job API end-to-end: `run_async()`,
:class-body: sd-fs-6
:class-title: sd-fs-5
-Wrap a local directory in `kinetic.Data(...)` and let it land as a
-plain filesystem path on the remote — your training code doesn't have
-to know whether the bytes started on your laptop or in GCS.
+Wrap a local directory in `kinetic.Data(...)`. The function receives a
+plain filesystem path, and does not know whether the bytes started on
+your laptop or in Cloud Storage.
+++
@@ -116,15 +130,30 @@ to know whether the bytes started on your laptop or in GCS.
{bdg-secondary}`GCS`
:::
+:::{grid-item-card} Mount data with FUSE
+:link: examples/example_fuse.md
+:class-card: sd-shadow-sm
+:class-body: sd-fs-6
+:class-title: sd-fs-5
+
+`Data(..., fuse=True)` for volumes, single files, several mounts in one
+job, and a mix of mounted and downloaded data.
+
++++
+
+{bdg-secondary}`Data`
+{bdg-secondary}`FUSE`
+:::
+
:::{grid-item-card} Resumable JAX training with Orbax
:link: examples/example_checkpoint.md
:class-card: sd-shadow-sm
:class-body: sd-fs-6
:class-title: sd-fs-5
-JAX training that picks up where it left off. Writes Orbax checkpoints
-to `KINETIC_OUTPUT_DIR` and proves the resume path by relaunching the
-same function and seeing it skip already-completed steps.
+JAX training that continues where it stopped. Writes Orbax checkpoints
+to `KINETIC_OUTPUT_DIR`, and shows the resume path when you run the same
+function again.
+++
@@ -139,9 +168,8 @@ same function and seeing it skip already-completed steps.
:class-body: sd-fs-6
:class-title: sd-fs-5
-Auto-resumable Keras training. Round-trips `model.get_weights()` through
-Orbax so a restarted job picks up at the right step without any custom
-save/load code.
+The same pattern for Keras. Round-trips `model.get_weights()` through
+Orbax, so a restarted job continues at the right step.
+++
@@ -156,9 +184,8 @@ save/load code.
:class-body: sd-fs-6
:class-title: sd-fs-5
-Fan out a grid of jobs with `run_async_map()`, batch submissions to keep
-the cluster happy, and gather results — including how to handle the
-job that inevitably fails halfway through.
+Fan out a grid of jobs with `run_async_map()`, limit the concurrency,
+collect the results, and handle the job that fails.
+++
@@ -166,15 +193,14 @@ job that inevitably fails halfway through.
{bdg-secondary}`Parallel`
:::
-:::{grid-item-card} Mix accelerators in one driver
+:::{grid-item-card} Mix accelerators in one script
:link: examples/example_gke.md
:class-card: sd-shadow-sm
:class-body: sd-fs-6
:class-title: sd-fs-5
-One driver script that successively schedules work on CPU, TPU, and
-GPU pools — handy for verifying which hardware your cluster will
-actually serve.
+One script that runs work on a CPU pool, a TPU pool, and a GPU pool in
+turn. Useful to check which hardware your cluster serves.
+++
@@ -195,8 +221,8 @@ actually serve.
:class-body: sd-fs-6
:class-title: sd-fs-5
-The reference for scaling beyond a single TPU host. A short JAX program
-that verifies cross-host collectives are actually wired up before you
+The reference for a slice with more than one TPU host. A short JAX
+program that checks that the cross-host collectives work before you
trust them with a real workload.
+++
@@ -212,9 +238,10 @@ trust them with a real workload.
:class-body: sd-fs-6
:class-title: sd-fs-5
-End-to-end SFT of Gemma 2B with LoRA across multiple TPU hosts. The
-realistic LLM workload to model your own fine-tuning runs after — pulls
-weights from Kaggle and runs on Pathways.
+Supervised fine-tuning of Gemma 2B with LoRA and the Keras
+`DataParallel` distribution. Pulls the weights from Kaggle, and forces
+the `pathways` backend on a single-host slice to exercise the multi-host
+code path.
+++
@@ -229,9 +256,9 @@ weights from Kaggle and runs on Pathways.
:class-body: sd-fs-6
:class-title: sd-fs-5
-Compact Gemma 3 1B SFT on a single TPU. A good baseline for getting an
-LLM workload running before scaling out to Pathways, and a worked
-example of forwarding Kaggle credentials into the remote pod.
+Compact Gemma 3 1B fine-tuning on one TPU. A good baseline before you
+scale to Pathways, and a worked example of Kaggle credentials in the
+pod.
+++
@@ -239,15 +266,14 @@ example of forwarding Kaggle credentials into the remote pod.
{bdg-secondary}`TPU`
:::
-:::{grid-item-card} Tunix SFT Example
+:::{grid-item-card} Tunix SFT
:link: examples/tunix_sft.md
:class-card: sd-shadow-sm
:class-body: sd-fs-6
:class-title: sd-fs-5
-SFT of Gemma 3 with LoRA/QLoRA on TPU v5litepod. Demonstrates how to
-run the Tunix SFT script on a remote cluster with environment variable
-capture for credentials.
+Supervised fine-tuning of Gemma 3 with LoRA and QLoRA through Tunix on
+a TPU v5e slice, with credentials forwarded through `capture_env_vars`.
+++
@@ -257,9 +283,20 @@ capture for credentials.
:::
::::
-## Related pages
-
-- [Getting Started](getting_started.md): your first run, end-to-end.
-- [Keras Training](examples/keras_training.md): patterns for Keras users.
-- [LLM Fine-tuning](examples/llm_finetuning.md): extended walkthrough using the
- Gemma examples.
+## Tutorials
+
+The pages in the **Examples & Tutorials** section of the sidebar are
+longer walkthroughs:
+
+- [Training Keras Models](examples/keras_training.md) — patterns for an
+ existing Keras script.
+- [Native JAX Training](examples/jax_training.md) — JAX loops,
+ single-host parallelism, and multi-host slices.
+- [PyTorch Training](examples/pytorch_training.md) — PyTorch on GPU
+ nodes.
+- [Fine-tuning Gemma 4 on TPU](examples/gemma4_finetuning.md) — a
+ complete LoRA fine-tune with inference.
+- [Fine-tuning LLMs](examples/llm_finetuning.md) — Keras Hub, Kaggle
+ credentials, and LoRA.
+- [Running vLLM on TPU](guides/vllm_tpu.md) — vLLM inference on a TPU
+ slice.
diff --git a/docs/examples/example_fuse.md b/docs/examples/example_fuse.md
new file mode 100644
index 00000000..fc417fed
--- /dev/null
+++ b/docs/examples/example_fuse.md
@@ -0,0 +1,5 @@
+# Mount data with FUSE
+
+```{literalinclude} ../../examples/example_fuse.py
+:language: python
+```
diff --git a/docs/examples/gemma4_finetuning.md b/docs/examples/gemma4_finetuning.md
index 241410c3..cecde358 100644
--- a/docs/examples/gemma4_finetuning.md
+++ b/docs/examples/gemma4_finetuning.md
@@ -1,44 +1,61 @@
# Fine-tuning Gemma 4 on TPU
-This guide walks through fine-tuning [Gemma 4 Instruct 26B](https://www.kaggle.com/models/keras/gemma4) on a TPU slice using Kinetic. You will use Low-Rank Adaptation (LoRA) to reduce memory requirements, save the adapted weights to GCS, and run inference with the fine-tuned model, all from your local machine.
-
-The model used here is `gemma4_instruct_26b_a4b`, a Mixture of Experts (MoE) architecture with 26B total parameters and 4B active parameters per forward pass. All 26B weights load into memory (~52 GB in bfloat16), so a v5litepod-8 (8 chips × 16 GB = 128 GB HBM) is the minimum required configuration. A self-contained script combining both steps is available at [`examples/gemma4_finetuning.py`](../../examples/gemma4_finetuning/gemma4_finetuning.py).
-
-## Prerequisites
-
-Before starting, you need:
-
-- A GCP project with billing enabled.
-- A Kinetic cluster provisioned (`kinetic up`). See the [Getting Started](../getting_started.md) guide.
-- A **v5litepod-8 TPU node pool** in your cluster. Run `kinetic status` to check what pools you have. If you need to add one:
+This tutorial fine-tunes
+[Gemma 4 Instruct 26B](https://www.kaggle.com/models/keras/gemma4) on
+one TPU v5e host with Kinetic. You use Low-Rank Adaptation (LoRA) to
+reduce the memory that training needs. You save the adapted weights to
+Cloud Storage under `KINETIC_OUTPUT_DIR`. Then you run inference with
+the fine-tuned model in a second job. Every step runs from your local
+machine.
+
+The model is `gemma4_instruct_26b_a4b`, a Mixture of Experts (MoE)
+architecture with 26B total parameters and 4B active parameters per
+forward pass. All 26B weights load into memory, about 52 GB in
+bfloat16. A `tpu-v5litepod-8` slice (8 chips × 16 GB = 128 GB of HBM)
+is the smallest configuration that fits. That slice is one host, so the
+job runs as a plain Kubernetes Job on the cluster. The complete script,
+with both steps, is
+[`examples/gemma4_finetuning/gemma4_finetuning.py`](https://github.com/keras-team/kinetic/blob/main/examples/gemma4_finetuning/gemma4_finetuning.py).
+
+## Before you start
+
+You need these things:
+
+- A Google Cloud project with billing enabled.
+- Kinetic installed, and `kinetic init` complete. `kinetic init` creates
+ the active profile that supplies the project, the zone, and the
+ cluster to every command and every job. See
+ [Getting Started](../getting_started.md).
+- A `tpu-v5litepod-8` node pool in your cluster. `kinetic pool list`
+ shows the pools of the cluster. If no `v5litepod` pool with 8 chips
+ exists, add one:
```bash
- kinetic pool add --accelerator tpu-v5litepod-8 --project your-project-id
+ kinetic pool add --accelerator tpu-v5litepod-8
```
-- A Kaggle account with [Gemma 4 access accepted](https://www.kaggle.com/models/keras/gemma4).
-- `KAGGLE_USERNAME` and `KAGGLE_KEY` set in your local environment.
-
-## GCloud Setup
+- A Kaggle account that has
+ [accepted the Gemma 4 terms](https://www.kaggle.com/models/keras/gemma4).
+- `KAGGLE_USERNAME` and `KAGGLE_KEY` set in your local shell.
-Authenticate and configure your project:
+## Check the TPU availability in your zone
-```bash
-gcloud auth login
-gcloud auth application-default login
-export GOOGLE_CLOUD_PROJECT="your-project-id"
-export GOOGLE_CLOUD_ZONE="us-central1-a"
-```
-
-`us-central1-a` reliably has on-demand `v5litepod-8` availability. Before running, verify your zone has the hardware:
+Not every zone offers `v5litepod-8` on demand. `us-central1-a` does.
+Before you submit the job, make sure that the zone of your profile has
+the hardware. `kinetic config` shows the zone and the project of the
+active profile. Then run:
```bash
-gcloud compute tpus accelerator-types list --zone=your-zone --project=your-project-id
+gcloud compute tpus accelerator-types list --zone=us-central1-a --project=your-project-id
```
-Confirm `v5litepod-8` appears in the output before submitting a job.
+Look for `v5litepod-8` in the output. If the zone does not list it,
+create a profile for a zone that does. See
+[Profiles](../guides/profiles.md).
-## Forwarding Credentials
+## Forward the Kaggle credentials
-Kaggle credentials must be present in the remote pod to download the model weights. Use `capture_env_vars` to forward them automatically:
+The pod downloads the model weights from Kaggle, so the pod needs your
+Kaggle credentials. The `capture_env_vars` parameter copies the named
+variables from your shell into the pod:
```python
import kinetic
@@ -51,40 +68,87 @@ import kinetic
def fine_tune_gemma4(): ...
```
-This pattern is covered in depth in the [Environment Variables](../guides/env_vars.md) guide.
-
-`keras-hub` and its tokenizer backends are not installed in the Kinetic base container by default. Add a `requirements.txt` to your project so Kinetic picks them up automatically:
-
-```text
-keras==3.13.2
-keras-hub==0.27.1
-tokenizers==0.22.2
-sentencepiece==0.2.1
-```
-
-Kinetic detects changes to this file and rebuilds the container only when needed. See the [Managing Dependencies](../guides/dependencies.md) guide for details.
-
-## Fine-tuning with LoRA
-
-The full training function loads the model, enables LoRA, and fits on a small instruction-following dataset. Imports live inside the function so they run on the remote worker.
-
-Four things are worth understanding before reading the code:
-
-**Precision policy.** The 26B model stores ~52 GB of weights. Using `mixed_bfloat16` would keep float32 master copies (~13 GB/chip on 8 chips), which — combined with MoE activation tensors — exceeds the 15.75 GB/chip HBM limit. The `bfloat16` policy stores variables directly as bfloat16 (~6.5 GB/chip), which fits.
-
-**Sequence length.** MoE activation tensors scale with the compiled sequence length. The preset default (~1024 tokens) produces ~10 GB/chip of HLO temporaries. Setting `model.preprocessor.sequence_length = 128` before `compile()` keeps it under ~2 GB/chip.
+The pattern `KAGGLE_*` matches `KAGGLE_USERNAME` and `KAGGLE_KEY`. The
+pattern `GOOGLE_CLOUD_*` forwards a `GOOGLE_CLOUD_PROJECT` value if your
+shell has one. With an active profile, Kinetic does not read
+`GOOGLE_CLOUD_*` variables to find the project or the zone; the profile
+supplies both.
+
+:::{note}
+Kinetic logs a warning when it captures a variable whose name looks like
+a credential, such as `KAGGLE_KEY`. The value travels in plaintext inside
+the job payload in the jobs bucket. The warning is expected for this
+tutorial. See
+[Forward Environment Variables](../guides/env_vars.md#secrets).
+:::
-**Weight sharding.** The 26B model does not fit on a single 16 GB chip. `ModelParallel` with an explicit `LayoutMap` splits weights across all 8 chips at variable creation time. The `LayoutMap` must be set before calling `from_preset()` so that variables are created with the correct sharding specs from the start.
+## Add the dependency file
-**Custom weight loading.** The Kaggle preset stores weights across 6 sharded H5 files described by a `model.weights.json` manifest. Keras's built-in `load_weights()` on the full `CausalLM` prepends a `backbone/` prefix that mismatches every path in the manifest. Loading via `model.backbone.load_weights()` avoids that prefix, but Keras ≤ 3.14 has a bug in `ShardedH5IOStore`: after switching to a different shard file, the internal `current_shard_path` pointer is not updated. When a subsequent `keys()` call restores to the stale path, layers whose weights span multiple shards — every MoE expert bank and the token embedding — fail to load, producing a "received 0 variables" error. The solution is to bypass `ShardedH5IOStore` entirely and read the H5 files directly with h5py, pre-sharding each tensor with `jax.device_put` before assigning it to avoid a memory spike on device 0. The complete loader is implemented as `_load_sharded_weights()` in [`examples/gemma4_finetuning.py`](../../examples/gemma4_finetuning/gemma4_finetuning.py).
+The image that Kinetic builds does not contain `keras-hub` or the
+tokenizer backends. Put a `requirements.txt` next to your script:
-:::{admonition} TODO
-:class: note
+```{literalinclude} ../../examples/gemma4_finetuning/requirements.txt
+:language: text
+```
-Remove `_load_sharded_weights` once Keras exposes a public loading path that handles the `backbone/` prefix correctly and fixes the `ShardedH5IOStore` shard-switching bug.
+Kinetic finds the file, builds an image that contains these packages,
+and caches the image. A change to the file causes a new build. The
+shipped example keeps this file in `examples/gemma4_finetuning/`, next
+to the script. See [Dependencies](../guides/dependencies.md).
+
+## Fine-tune with LoRA
+
+The training function loads the model, enables LoRA, and fits the model
+on a small instruction-following dataset. The imports live inside the
+function, so they run on the pod.
+
+Four decisions in the code need an explanation:
+
+**Precision policy.** The 26B model stores about 52 GB of weights. The
+`mixed_bfloat16` policy keeps float32 master copies of every variable
+(about 13 GB per chip on 8 chips). Those copies, together with the MoE
+activation tensors, exceed the HBM of one chip. The `bfloat16` policy
+stores the variables directly as bfloat16 (about 6.5 GB per chip), which
+fits.
+
+**Sequence length.** The MoE activation tensors scale with the compiled
+sequence length. The preset default of about 1024 tokens produces about
+10 GB per chip of temporary HLO buffers. The line
+`model.preprocessor.sequence_length = 128`, before `compile()`, keeps
+those buffers under about 2 GB per chip.
+
+**Weight sharding.** The 26B model does not fit on one 16 GB chip.
+`ModelParallel` with an explicit `LayoutMap` splits the weights across
+all 8 chips when Keras creates the variables. Set the `LayoutMap` before
+you call `from_preset()`, so that every variable gets the correct
+sharding from the start.
+
+**Custom weight loading.** The Kaggle preset stores the weights in 6
+sharded H5 files that a `model.weights.json` manifest describes. The
+built-in `load_weights()` on the full `CausalLM` adds a `backbone/`
+prefix that matches no path in the manifest.
+`model.backbone.load_weights()` avoids that prefix, but the Keras
+`ShardedH5IOStore` has a bug. After the store switches to a different
+shard file, it does not update its internal `current_shard_path`
+pointer. A later `keys()` call restores the stale path. Every layer
+whose weights span more than one shard then fails to load with a
+"received 0 variables" error. Each MoE expert bank and the token
+embedding are such layers. The example therefore bypasses
+`ShardedH5IOStore` and reads the H5 files directly with `h5py`. The
+loader shards each tensor with `jax.device_put` before it assigns the
+tensor, which avoids a memory spike on device 0. The function
+`_load_sharded_weights()` in the
+[example script](https://github.com/keras-team/kinetic/blob/main/examples/gemma4_finetuning/gemma4_finetuning.py)
+contains the complete loader.
+
+:::{note}
+`_load_sharded_weights()` is a workaround. When Keras offers a public
+loading path that handles the `backbone/` prefix and that switches
+shards correctly, the loader is no longer necessary.
:::
-The code below assumes `_load_sharded_weights` and `_make_layout_map` are defined as in [`examples/gemma4_finetuning.py`](../../examples/gemma4_finetuning/gemma4_finetuning.py).
+The code below expects `_load_sharded_weights` and `_make_layout_map`
+as the example script defines them.
```python
import os
@@ -201,54 +265,93 @@ def fine_tune_gemma4():
if __name__ == "__main__":
os.environ["KERAS_BACKEND"] = "jax"
- os.environ["GOOGLE_CLOUD_PROJECT"] = "your-project-id"
- os.environ["GOOGLE_CLOUD_ZONE"] = "us-central1-a"
weights_path = fine_tune_gemma4()
print(f"Training complete. Weights at: {weights_path}")
```
-Only the LoRA adapter variables (a few MB) are saved — not the full 26B backbone. `KINETIC_OUTPUT_DIR` is automatically set to a unique GCS path (e.g. `gs://your-bucket/job-abc123/output/`) for every job. The full path is printed to your terminal so you can pass it to the inference job below.
+The function saves only the LoRA adapter variables, a few MB, and not
+the full 26B backbone. Kinetic sets `KINETIC_OUTPUT_DIR` in the pod to a
+Cloud Storage prefix that is unique to the job. The default is
+`gs://{project}-kn-{cluster}-jobs/outputs/{job_id}`, for example
+`gs://your-project-id-kn-kinetic-cluster-jobs/outputs/job-534ffeb6`.
+The function prints the full path of the weights file, and you pass that
+path to the inference job below.
+
+The shipped script also sets `GOOGLE_CLOUD_PROJECT`, `KINETIC_ZONE`, and
+`GOOGLE_CLOUD_ZONE` in its `__main__` block. With an active profile,
+Kinetic ignores the two `GOOGLE_CLOUD_*` variables; their only effect
+is that `capture_env_vars` forwards them to the pod. `KINETIC_ZONE`
+overrides the zone of the active profile. Set the three values to match
+your profile, or delete the three lines.
-## Monitoring the Job
+:::{warning}
+A lifecycle rule on the jobs bucket deletes every object 30 days after
+its creation. If you need the LoRA weights for longer, copy the file to
+a bucket that you manage, or pass `output_dir="gs://your-bucket/gemma4"`
+to `@kinetic.run()`. See
+[Outputs and Checkpoints](../guides/checkpointing.md#retention-and-cleanup).
+:::
+
+A blocking call to `fine_tune_gemma4()` blocks until the job ends. The
+first run builds the container image (5 to 10 minutes), and every run
+downloads about 52 GB of weights from Kaggle before training starts.
+If you do not want to keep a terminal open, call
+`fine_tune_gemma4.run_async()` and collect the result later. See
+[Detached Jobs](../guides/async_jobs.md).
-While the fine-tuning job is running you can inspect it from a separate terminal using the `kinetic jobs` CLI. All commands require `--project` (or the `KINETIC_PROJECT` env var) to locate your cluster.
+## Monitor the job
-List all live jobs:
+While the job runs, inspect it from a second terminal with the
+`kinetic jobs` commands. The active profile supplies the project and the
+cluster. Add `--project` or `--cluster` only when the job runs on a
+different cluster.
+
+List the live jobs:
```bash
-kinetic jobs list --project your-project-id
+kinetic jobs list
```
-Check the status of a specific job — the job ID is printed to your terminal when the job is submitted (e.g. `job-534ffeb6`):
+Show the status of one job. Kinetic prints the job ID, for example
+`job-534ffeb6`, in the log output when it submits the job:
```bash
-kinetic jobs status JOB_ID --project your-project-id
+kinetic jobs status JOB_ID
```
-Stream live logs until the job finishes:
+Stream the logs until the job ends:
```bash
-kinetic jobs logs --follow JOB_ID --project your-project-id
+kinetic jobs logs --follow JOB_ID
```
-If you need to stop the job early:
+Stop the job early:
```bash
-kinetic jobs cancel JOB_ID --project your-project-id
+kinetic jobs cancel JOB_ID
```
-If the job stays in `PENDING` for more than a few minutes, inspect the pod to diagnose scheduling failures:
+If the job stays in `PENDING` for more than 10 minutes, inspect the
+pod. A `tpu-v5litepod-8` job is a plain Kubernetes Job named
+`kinetic-JOB_ID`, and its pod carries the label `job-id=JOB_ID`:
```bash
-kubectl describe pod -l leaderworkerset.sigs.k8s.io/name=keras-pathways-JOB_ID -n default
+kubectl get pods -l job-id=JOB_ID -n default
+kubectl describe pod -l job-name=kinetic-JOB_ID -n default
```
-Check the **Events** section at the bottom — common causes are insufficient TPU quota, no matching node pool for the requested accelerator, or image pull errors.
+Replace `default` with the namespace of your profile if you changed it.
+Read the **Events** section at the end of the output. The common causes
+are insufficient TPU quota, no node pool that matches the accelerator,
+and an image pull error. See
+[Scheduling and quota issues](../troubleshooting.md#scheduling-and-quota-issues)
+in Troubleshooting.
-## Inference with Fine-tuned Weights
+## Run inference with the fine-tuned weights
-After the training job completes, copy the printed weights path and pass it to a separate inference job:
+After the training job ends, copy the printed weights path and pass it
+to a second job:
```python
import os
@@ -317,53 +420,79 @@ def run_inference(weights_path: str):
if __name__ == "__main__":
os.environ["KERAS_BACKEND"] = "jax"
- os.environ["GOOGLE_CLOUD_PROJECT"] = "your-project-id"
- os.environ["GOOGLE_CLOUD_ZONE"] = "us-central1-a"
- # Replace with the path printed at the end of the fine-tuning job.
- weights_path = "gs://your-bucket/job-abc123/output/gemma4_lora.weights.h5"
+ # Replace with the path that the fine-tuning job printed.
+ weights_path = "gs://your-project-id-kn-kinetic-cluster-jobs/outputs/job-534ffeb6/gemma4_lora.weights.h5"
response = run_inference(weights_path)
print(response)
```
-## Cleaning Up
+The inference job loads the base weights again from Kaggle, enables LoRA
+with the same rank, and assigns the saved adapter variables by path. It
+then generates one response and returns the text.
-:::{warning}
-TPU node pools accrue cost while they exist, even when no job is running. Remove resources when you are done to avoid unnecessary charges.
-:::
+## Clean up
+
+A node pool that you add with the default `--min-nodes 0` scales to zero
+after about 10 idle minutes. After that, the TPU nodes cost nothing. The
+cluster control plane and the default CPU node cost money while the
+cluster exists.
-Remove the v5litepod-8 pool while keeping the cluster intact for other workloads:
+To delete the `v5litepod-8` pool and keep the cluster for other jobs:
```bash
-# Find the exact pool name
-kinetic pool list --project your-project-id
+# Find the exact pool name.
+kinetic pool list
-# Remove it (use the name printed above, e.g. tpu-v5litepod-a1b2)
-kinetic pool remove POOL_NAME --project your-project-id
+# Delete the pool. Use the name from the list, for example tpu-v5litepod-a1b2.
+kinetic pool remove POOL_NAME
```
-To tear down the entire cluster, including all pools, the GKE cluster, and associated infrastructure:
+To delete the whole cluster, with every pool and the buckets:
```bash
-kinetic down --project your-project-id
+kinetic down
```
-## Next Steps
+:::{warning}
+`kinetic down` deletes the jobs bucket. The LoRA weights under the
+default `KINETIC_OUTPUT_DIR` are in that bucket. Copy the weights to a
+bucket that Kinetic does not manage before you run `kinetic down`.
+:::
+
+## Related pages
::::{grid} 1 1 2 2
:gutter: 3
-:::{grid-item-card} {octicon}`history;1em` Checkpointing during training
+:::{grid-item-card} {octicon}`book;1em` Fine-tuning LLMs
+:link: llm_finetuning
+:link-type: doc
+
+The general patterns for Keras Hub models: the dependency file, the
+Kaggle credentials, and the move to a multi-host slice.
+:::
+
+:::{grid-item-card} {octicon}`history;1em` Outputs and Checkpoints
:link: ../guides/checkpointing
:link-type: doc
-Use Orbax to save intermediate checkpoints so a long run can resume if interrupted.
+`KINETIC_OUTPUT_DIR`, retention, and how to write checkpoints that a
+long run can resume from.
+:::
+
+:::{grid-item-card} {octicon}`key;1em` Forward Environment Variables
+:link: ../guides/env_vars
+:link-type: doc
+
+How `capture_env_vars` works, which names a wildcard never matches, and
+how Kinetic handles secrets.
:::
-:::{grid-item-card} {octicon}`server;1em` Distributed training
+:::{grid-item-card} {octicon}`server;1em` Distributed Training
:link: ../guides/distributed_training
:link-type: doc
-Scale to larger TPU slices or multiple hosts.
+Move to a larger TPU slice that spans more than one host.
:::
::::
diff --git a/docs/examples/jax_training.md b/docs/examples/jax_training.md
index c4f8c051..5d75009c 100644
--- a/docs/examples/jax_training.md
+++ b/docs/examples/jax_training.md
@@ -1,18 +1,29 @@
# Native JAX Training
-:::{admonition} Who this is for
-:class: note
-
-Users who write training loops directly in JAX
-rather than going through Keras. Kinetic runs your JAX code on cloud
-TPUs and GPUs the same way it runs Keras code — wrap the function in
-`@kinetic.run()` and call it. JAX-specific details (multi-device
-parallelism, dependency filtering, multi-host coordination) are covered
-below.
-:::
+This page shows how to run a training loop that you write directly in
+JAX on a cloud TPU or GPU. Read this page if you use `jax.grad`,
+`jax.pmap`, or `jax.sharding` without Keras. Kinetic runs a JAX function
+the same way that it runs a Keras function: you decorate the function
+with `@kinetic.run()` and call it. The sections below cover the
+JAX-specific details: the JAX runtime in the image, single-host
+parallelism, multi-host slices, data, and outputs.
+
+## Before you start
+
+- Complete [Getting Started](../getting_started.md). You need an active
+ profile and a cluster.
+- Make sure that the cluster has a node pool for the accelerator that
+ you use. `kinetic pool list` shows the pools of the cluster.
+ `kinetic pool add --accelerator tpu-v5litepod-8` adds a pool for the
+ first example. See [Clusters and Node Pools](../guides/clusters.md).
+- Do not add `jax`, `jaxlib`, or `libtpu` to your dependency file.
+ Kinetic installs them for you. See
+ [The JAX runtime in the image](#the-jax-runtime-in-the-image).
## A first run
+Save this script and run it with `python`:
+
```python
import kinetic
@@ -32,10 +43,16 @@ def jax_computation():
print(jax_computation()) # 1000.0
```
-A standard JAX training loop with `jax.grad` runs without modification:
+Kinetic sends the function to the cluster, streams the log lines to your
+terminal, and returns the value to the script. The first run takes 5 to 10
+minutes, because Kinetic builds a container image with your
+dependencies. Later runs with the same dependency file start in less
+than 1 minute while a node still runs.
+
+A training loop with `jax.grad` runs without a change to the loop:
```python
-@kinetic.run(accelerator="tpu-v6e-8")
+@kinetic.run(accelerator="tpu-v5litepod-4")
def train():
import jax
import jax.numpy as jnp
@@ -61,31 +78,43 @@ def train():
return float(loss_fn(params, x, y))
```
-Imports for `jax`, `jaxlib`, and any other heavy library go **inside**
-the decorated function so the remote worker uses its accelerator-tuned
-install.
+Put `import jax` and the imports of other heavy libraries **inside** the
+decorated function. The import then runs on the pod, which has the JAX
+build for the accelerator. Your local machine does not need JAX.
+
+## The JAX runtime in the image
+
+JAX needs a `jaxlib` build and an accelerator runtime that match the
+hardware. The image that Kinetic builds contains JAX for your
+accelerator:
-## How to think about it
+| Accelerator | JAX package in the image |
+| ----------- | ------------------------ |
+| TPU (`tpu-...`) | `jax[tpu]`, which includes `libtpu` |
+| GPU (`gpu-...`) | `jax[cuda12]`, which includes the CUDA libraries |
+| CPU (`cpu`) | `jax` |
-JAX needs the right `jaxlib` and the right accelerator runtime
-(`libtpu`, CUDA) to be installed in the container. Kinetic handles this
-for you:
+Do not pin `jax`, `jaxlib`, or `libtpu` in your dependency file.
+Kinetic removes the entries for `jax`, `jaxlib`, `libtpu`, and
+`libtpu-nightly` from the dependency list before the install. A pin in
+your file therefore does not replace the installation in the image.
+Kinetic logs a warning for each removed entry. To keep a line in
+`requirements.txt`, append `# kn:keep` to the line. See
+[JAX and accelerator runtimes](../guides/dependencies.md#jax-and-accelerator-runtimes)
+for the filter rules and the override.
-- **Bundled and prebuilt images** ship with JAX matched to the
- accelerator type. You don't need to pin `jax`, `jaxlib`, or `libtpu`
- in `requirements.txt`.
-- **JAX packages in your `requirements.txt` are filtered out** before
- install so they don't shadow the accelerator-correct copy in the
- image. See [Dependencies](../guides/dependencies.md) for the filter behavior.
+Inside the function, `jax.devices()` returns the devices of the pod:
-Inside the function, `jax.devices()` returns whatever the pod sees: an
-8-chip TPU slice for `tpu-v6e-8`, an 8-device array for
-`tpu-v5litepod-8`, a single GPU for `l4`, etc.
+- `tpu-v5litepod-8`: 8 TPU devices on one host.
+- `tpu-v5litepod-4`: 4 TPU devices on one host.
+- `gpu-l4`: one GPU device.
+- `cpu`: one CPU device.
## Single-host parallelism
-Use `jax.pmap` (or `jax.sharding`) to spread computation across all
-devices on a single host:
+A single-host slice, for example `tpu-v5litepod-8`, holds all of its
+chips on one VM. `jax.pmap` or `jax.sharding` spreads the computation
+across those chips. Kinetic needs no extra setting for this case.
```python
@kinetic.run(accelerator="tpu-v5litepod-8")
@@ -105,39 +134,68 @@ def parallel_computation():
return float(result[0, 0, 0])
```
-## Scaling beyond a single host
+On `tpu-v5litepod-8`, `jax.local_device_count()` is 8, and each `pmap`
+replica runs on one chip.
-For multi-host slices (e.g., `tpu-v5litepod-2x4`) JAX needs a coordination
-runtime to set up cross-host collectives. Kinetic provides this through
-the Pathways backend:
+## Multi-host slices
+
+Some slices span more than one host. `tpu-v5litepod-16` and `tpu-v6e-16`
+each consist of four 4-chip VMs. Kinetic reads the host count from the
+accelerator name and selects the `pathways` backend for you. You do not
+set `backend=`. Kinetic runs one pod per host and sets the environment
+variables for multi-controller JAX on every pod. JAX then starts one
+process per host. `jax.process_count()` equals the host count,
+`jax.local_device_count()` is the chip count of one host, and
+`jax.device_count()` is the total. Collectives across hosts work, for
+example `jax.lax.psum` inside `pmap`.
```python
-@kinetic.run(accelerator="tpu-v5litepod-2x4", backend="pathways")
+@kinetic.run(accelerator="tpu-v6e-16")
def train_distributed():
import jax
- # jax.process_count() > 1 here; pmap/sharding work cross-host.
+ print(f"Host {jax.process_index()} of {jax.process_count()}")
+ print(f"Devices on this host: {jax.local_device_count()}")
+ print(f"Total devices: {jax.device_count()}")
+ # pmap and sharding work across hosts here.
...
```
-Without `backend="pathways"`, multi-host JAX collectives won't have a
-working coordinator. See [Distributed Training](../guides/distributed_training.md)
-for the full multi-host setup.
+On `tpu-v6e-16`, `jax.process_count()` is 4, `jax.local_device_count()`
+is 4, and `jax.device_count()` is 16.
+
+A multi-host job has two requirements:
+
+- The cluster needs a node pool for the multi-host accelerator, for
+ example `kinetic pool add --accelerator tpu-v6e-16`.
+- The cluster needs the LeaderWorkerSet controller. `kinetic up`
+ installs it.
+
+:::{note}
+The chip count alone does not tell you the host count. `tpu-v5litepod-8`
+is one 8-chip VM, but `tpu-v6e-8` is two 4-chip VMs and therefore a
+multi-host job. The **Hosts** column on the
+[Accelerators](../accelerators.md#tpus) page decides.
+:::
+
+See [Distributed Training](../guides/distributed_training.md) for the
+log that you see, the return value, and the failure behavior of a
+multi-host job.
## Data
To pass a dataset into a remote JAX function, construct a
-`kinetic.Data(...)` object **at the call site** in your local script and
-pass it as an argument. Kinetic uploads (or mounts) the source and
-delivers a plain filesystem path to the remote function. The decorated
-function only ever sees a `str` path:
+`kinetic.Data(...)` object **at the call site** in your local script.
+Pass the object as an argument. Kinetic uploads or mounts the source and
+gives the remote function a plain filesystem path. The decorated
+function only sees a `str` path:
```python
import kinetic
from kinetic import Data
-@kinetic.run(accelerator="tpu-v6e-8")
+@kinetic.run(accelerator="tpu-v5litepod-8")
def train(data_dir):
# `data_dir` is a local filesystem path on the remote pod.
import os
@@ -149,36 +207,44 @@ def train(data_dir):
# Local directory:
train(Data("./my_dataset/"))
-# Existing GCS bucket:
+# Existing Cloud Storage prefix (the trailing slash marks a directory):
train(Data("gs://my-bucket/dataset/"))
-# Large GCS dataset, streamed on demand via FUSE:
+# Large Cloud Storage dataset, read on demand through FUSE:
train(Data("gs://my-bucket/large/", fuse=True))
```
-`Data` accepts both local paths and `gs://` URIs. See [Data](../guides/data.md)
-for the decision matrix between downloaded, FUSE-mounted, and direct
-access patterns.
+`Data` accepts a local path, a `gs://` URI, or an `hf://` URI for a
+Hugging Face dataset. See [Working with Data](../guides/data.md) for the
+choice between a downloaded copy, a FUSE mount, and direct `gs://`
+access.
-## Next steps
+## Outputs and checkpoints
-::::{grid} 1 1 2 2
-:gutter: 3
+Kubernetes deletes the pod filesystem, including `/tmp`, when the pod
+ends. Write every file that you want to keep under `KINETIC_OUTPUT_DIR`.
+Kinetic sets that environment variable in the pod to a Cloud Storage
+prefix that stays after the pod ends. Orbax writes to a `gs://` path
+directly, so pass that prefix to the `CheckpointManager`:
-:::{grid-item-card} {octicon}`server;1em` Distributed Training
-:link: ../guides/distributed_training
-:link-type: doc
+```python
+@kinetic.run(accelerator="tpu-v5litepod-8")
+def train():
+ import os
-Multi-host JAX with Pathways.
-:::
+ import orbax.checkpoint as ocp
-:::{grid-item-card} {octicon}`history;1em` Checkpointing
-:link: ../guides/checkpointing
-:link-type: doc
+ output_dir = os.environ["KINETIC_OUTPUT_DIR"]
+ mngr = ocp.CheckpointManager(
+ f"{output_dir}/checkpoints", ocp.StandardCheckpointer()
+ )
+ ...
+```
-Orbax checkpoint patterns under `KINETIC_OUTPUT_DIR`.
-:::
-::::
+Add `orbax-checkpoint` to your dependency file for this example.
+
+See [Outputs and Checkpoints](../guides/checkpointing.md) for a full
+Orbax example, the default location, and how to resume a job.
## Related pages
@@ -189,20 +255,27 @@ Orbax checkpoint patterns under `KINETIC_OUTPUT_DIR`.
:link: ../guides/distributed_training
:link-type: doc
-Pathways and multi-host coordination.
+How a multi-host JAX job runs, what the log shows, and how it fails.
:::
:::{grid-item-card} {octicon}`package;1em` Dependencies
:link: ../guides/dependencies
:link-type: doc
-JAX filtering and what gets installed.
+The dependency file, the JAX filter, and the `# kn:keep` override.
:::
-:::{grid-item-card} {octicon}`history;1em` Checkpointing
+:::{grid-item-card} {octicon}`history;1em` Outputs and Checkpoints
:link: ../guides/checkpointing
:link-type: doc
-Orbax + `KINETIC_OUTPUT_DIR`.
+`KINETIC_OUTPUT_DIR`, retention, and Orbax checkpoints that resume.
+:::
+
+:::{grid-item-card} {octicon}`database;1em` Working with Data
+:link: ../guides/data
+:link-type: doc
+
+Ship local files and read Cloud Storage data from your function.
:::
::::
diff --git a/docs/examples/keras_training.md b/docs/examples/keras_training.md
index 085167eb..d715b4d4 100644
--- a/docs/examples/keras_training.md
+++ b/docs/examples/keras_training.md
@@ -1,14 +1,21 @@
# Training Keras Models
-:::{admonition} Who this is for
-:class: note
-
-Anyone with a working Keras training script who wants
-it to run on a cloud TPU or GPU without standing up infrastructure.
-Kinetic ships your existing `model.compile()` / `model.fit()` code to a
-remote accelerator with a single decorator change. You don't need to
-restructure your training loop.
-:::
+This page is for a reader who has a Keras training script and wants to
+run that script on a cloud TPU or GPU. Kinetic runs your
+`model.compile()` and `model.fit()` code on a remote accelerator after
+one decorator change. You do not restructure the training loop. This
+page shows the first run, explains what the function sees on the pod,
+and covers the Keras backend, multi-host slices, data, and outputs.
+
+## Before you start
+
+- Run `kinetic init` one time. The active profile supplies the project,
+ the zone, and the cluster to every job. See
+ [Getting Started](../getting_started.md).
+- Make sure that the cluster has a node pool for the accelerator in the
+ example. `kinetic pool list` shows the pools. Add one with
+ `kinetic pool add --accelerator tpu-v5litepod-4`, or change
+ `accelerator=` to an accelerator that the cluster has.
## A first run
@@ -16,7 +23,7 @@ restructure your training loop.
import kinetic
-@kinetic.run(accelerator="tpu-v6e-8")
+@kinetic.run(accelerator="tpu-v5litepod-4")
def train_model():
import keras
import numpy as np
@@ -40,132 +47,208 @@ final_loss = train_model()
print(f"Final loss: {final_loss}")
```
-A few things to note:
-
-- Imports for `keras`, `jax`, etc. live **inside** the function so the
- remote worker uses its hardware-tuned install.
-- The return value is serialized back to your local process. Keep it
- small — a final metric, a path under `KINETIC_OUTPUT_DIR`, a dict of
- numbers. Don't return the model object itself.
-- `accelerator="tpu-v6e-8"` picks an 8-chip TPU v6e slice. Use `cpu` while
- iterating; switch when you're ready for hardware. See
+Three points about this script:
+
+- Put the imports for `keras` and `numpy` inside the function. Your
+ machine then does not need those packages. The pod imports them from
+ the image, and the image contains JAX with the runtime for the
+ accelerator.
+- The return value goes back to your local process. Keep the value
+ small: a final metric, a dict of numbers, or a path under
+ `KINETIC_OUTPUT_DIR`. Do not return the model object.
+- `accelerator="tpu-v5litepod-4"` selects a 4-chip TPU v5e slice on one
+ host. Use `accelerator="cpu"` while you develop the code. Change the
+ accelerator when the code works. See
[Accelerators](../accelerators.md).
-For the canonical end-to-end example with a real dataset, see
-[`fashion_mnist.py`](../examples.md) (first entry under Quickstart).
+The first run takes 5 to 10 minutes, because Kinetic builds the
+container image with your dependencies. Later runs with the same
+dependencies start in less than 1 minute while a node still runs. See
+[How Kinetic Works](../concepts.md).
+
+For an end-to-end example with a real dataset, see
+[`fashion_mnist.py`](fashion_mnist.md).
-## How to think about it
+## What the function sees
-Your decorated function runs in a fresh process inside a container on a
-remote node. That has two practical consequences:
+Your decorated function runs in a new Python process, inside a
+container, on a node of the cluster. Kinetic serializes the function
+with `cloudpickle`, together with the objects that the function
+references. Two consequences follow:
-- **No local state crosses the boundary.** Anything the function needs
- must either be passed as an argument, captured by closure, or shipped
- via [`kinetic.Data`](../guides/data.md). Locally-loaded variables that you reference
- by global name will not be there on the remote.
-- **The Keras backend is whatever the remote has installed.** By default
- Kinetic's prebuilt and bundled images use JAX. Set `KERAS_BACKEND` if
- you need otherwise:
+- **The payload carries every referenced object.** A value from the
+ enclosing scope or a module-level global goes into the payload. A
+ small value, such as a config dict, is not a problem. A large object,
+ such as a dataset that you load at module level, makes the payload
+ large on every submit. Kinetic logs a warning when the payload is
+ larger than 50 MB. Load large data inside the function, or pass the
+ data as [`kinetic.Data`](../guides/data.md). See
+ [What Ships to the Pod](../guides/packaging.md).
+- **The image supplies the packages.** The image contains Keras, JAX,
+ and the packages from your `requirements.txt` or `pyproject.toml`. A
+ `pip install` in your local shell does not carry over. See
+ [Dependencies](../guides/dependencies.md).
- ```python
- @kinetic.run(accelerator="tpu-v6e-8", capture_env_vars=["KERAS_BACKEND"])
- def train(): ...
- ```
+### The Keras backend
-## Scaling beyond a single host
+The image that Kinetic builds sets `KERAS_BACKEND=jax`. Keras therefore
+uses the JAX backend on the pod, regardless of the backend on your
+machine. JAX is also the only accelerator runtime in the image:
+`jax[tpu]` for a TPU and `jax[cuda12]` for a GPU.
-For multi-host TPU slices like `tpu-v5litepod-2x4`, switch to the Pathways
-backend so Keras's distribution strategies have a working multi-host
-runtime to talk to:
+If your script needs another backend, do three things:
+
+1. Add the framework, for example `torch`, to your dependency file.
+2. Set `KERAS_BACKEND` in your shell, for example
+ `export KERAS_BACKEND=torch`.
+3. Forward the variable with `capture_env_vars`.
```python
-@kinetic.run(accelerator="tpu-v5litepod-2x4", backend="pathways")
-def train_distributed(): ...
+@kinetic.run(accelerator="gpu-l4", capture_env_vars=["KERAS_BACKEND"])
+def train(): ...
```
-See [Distributed Training](../guides/distributed_training.md) for the full
-multi-host setup, and [LLM Fine-tuning](llm_finetuning.md) for a
-concrete Gemma example.
+Name the variable exactly. A wildcard such as `"KERAS_*"` never
+captures `KERAS_BACKEND`, because that variable is on the wildcard
+blocklist. The pod applies the captured value before it calls your
+function, so the `import keras` inside the function sees the new
+backend. See
+[Forward Environment Variables](../guides/env_vars.md).
+
+## Scale to more than one host
+
+A single-host slice such as `tpu-v5litepod-8` has up to 8 chips on one
+VM. If the model or the batch does not fit on one host, select a
+multi-host slice, for example `tpu-v5litepod-16` or `tpu-v6e-16`. Each
+of those slices has four 4-chip VMs. Kinetic reads the host count from
+the accelerator name and selects the Pathways backend for you. You do
+not set `backend="pathways"`.
+
+The Keras distribution API sees every chip on every host. Set a
+`DataParallel` distribution before you build the model:
+
+```python
+@kinetic.run(accelerator="tpu-v5litepod-16")
+def train_distributed():
+ import keras
+
+ devices = keras.distribution.list_devices()
+ mesh = keras.distribution.DeviceMesh(
+ shape=(len(devices),), axis_names=["batch"], devices=devices
+ )
+ keras.distribution.set_distribution(
+ keras.distribution.DataParallel(device_mesh=mesh)
+ )
+
+ model = keras.Sequential([...])
+ model.compile(...)
+ model.fit(...)
+```
+
+The cluster needs a node pool for the multi-host accelerator. See
+[Distributed Training](../guides/distributed_training.md) for the node
+pool, the log that you see, and the return value of a multi-host job.
+See [Fine-tuning LLMs](llm_finetuning.md) for a Gemma example.
## Data
-Pulling NumPy arrays from inside the function works for tiny datasets,
-but breaks down quickly. For real data, construct a
-`kinetic.Data(...)` object **at the call site** in your local script
-and pass it as an argument. Kinetic uploads (or mounts) the source and
-delivers a plain filesystem path to the remote function. The decorated
-function only ever sees a `str` path:
+Random NumPy arrays inside the function are sufficient for a test. Real
+data must reach the pod. Construct a `kinetic.Data(...)` object **at the
+call site** in your local script. Pass the object as an argument.
+Kinetic uploads a local path one time and downloads the data to the
+pod, or mounts a Cloud Storage location. Your function receives a plain
+filesystem path (`str`):
```python
import kinetic
from kinetic import Data
-@kinetic.run(accelerator="tpu-v6e-8")
+@kinetic.run(accelerator="tpu-v5litepod-8")
def train(data_dir):
- # `data_dir` is a local filesystem path on the remote pod.
+ # `data_dir` is a local filesystem path on the pod.
import keras
...
-# Local directory:
+# A local directory:
train(Data("./my_dataset/"))
-# Existing GCS bucket:
+# A directory in Cloud Storage (the trailing slash marks a directory):
train(Data("gs://my-bucket/dataset/"))
-# Large GCS dataset, streamed on demand via FUSE:
+# A large Cloud Storage dataset, read on demand through a FUSE mount:
train(Data("gs://my-bucket/large/", fuse=True))
+
+# A Hugging Face dataset:
+train(Data("hf://imdb?split=train"))
```
-`Data` accepts both local paths and `gs://` URIs. See [Data](../guides/data.md)
-for the decision matrix between downloaded, FUSE-mounted, and direct
-access patterns.
+`Data` accepts a local path, a `gs://` URI, or an `hf://` Hugging Face
+dataset URI. An `hf://` URI needs the `datasets` package in your
+dependency file. See [Working with Data](../guides/data.md) for the
+choice between a download, a FUSE mount, and direct `gs://` access, and
+for the limits.
-## Next steps
+## Save the model
-::::{grid} 1 1 2 2
-:gutter: 3
+Kubernetes deletes the pod filesystem, including `/tmp`, when the pod
+ends. Write the files that you want to keep, for example the model
+weights and the checkpoints, under `KINETIC_OUTPUT_DIR`. Kinetic sets
+that variable on the pod to a per-job Cloud Storage location. Return the
+path, not the model:
-:::{grid-item-card} {octicon}`rocket;1em` fashion_mnist.py
-:link: ../examples
-:link-type: doc
+```python
+import os
-Full working example with a real dataset (first entry under Quickstart).
-:::
+import kinetic
-:::{grid-item-card} {octicon}`history;1em` Checkpointing
-:link: ../guides/checkpointing
-:link-type: doc
-Persist model weights and resume across runs.
-:::
-::::
+@kinetic.run(accelerator="tpu-v5litepod-4")
+def train():
+ import keras
+
+ output_dir = os.environ.get("KINETIC_OUTPUT_DIR", "/tmp/local_run")
+ # ... build, compile, and fit the model ...
+ # ... write the weights and the metrics under output_dir ...
+ return output_dir
+```
+
+See [Outputs and Checkpoints](../guides/checkpointing.md) for the output
+directory, retention, and a resumable Keras run with Orbax
+([`example_keras_checkpoint.py`](example_keras_checkpoint.md)).
## Related pages
::::{grid} 1 1 2 2
:gutter: 3
-:::{grid-item-card} {octicon}`database;1em` Data
+:::{grid-item-card} {octicon}`rocket;1em` fashion_mnist.py
+:link: fashion_mnist
+:link-type: doc
+
+A complete example with a real dataset on a TPU.
+:::
+
+:::{grid-item-card} {octicon}`database;1em` Working with Data
:link: ../guides/data
:link-type: doc
-Shipping local files and reading from GCS.
+Ship local files and read Cloud Storage data from your function.
:::
-:::{grid-item-card} {octicon}`history;1em` Checkpointing
+:::{grid-item-card} {octicon}`history;1em` Outputs and Checkpoints
:link: ../guides/checkpointing
:link-type: doc
-`KINETIC_OUTPUT_DIR` and resumable training.
+`KINETIC_OUTPUT_DIR`, retention, and resumable training.
:::
-:::{grid-item-card} {octicon}`cpu;1em` LLM Fine-tuning
+:::{grid-item-card} {octicon}`cpu;1em` Fine-tuning LLMs
:link: llm_finetuning
:link-type: doc
-KerasHub + Gemma walkthrough.
+Keras Hub, Kaggle credentials, and LoRA on Gemma.
:::
::::
diff --git a/docs/examples/llm_finetuning.md b/docs/examples/llm_finetuning.md
index b8d6e30c..9cdb572b 100644
--- a/docs/examples/llm_finetuning.md
+++ b/docs/examples/llm_finetuning.md
@@ -1,84 +1,242 @@
# Fine-tuning LLMs
-Kinetic integrates seamlessly with [Keras Hub](https://keras.io/keras_hub/) and the [Kaggle](https://www.kaggle.com/) ecosystem, making it easy to fine-tune large language models like Gemma on cloud TPUs.
+This page shows the Kinetic patterns for the fine-tuning of a large
+language model from [Keras Hub](https://keras.io/keras_hub/) on a TPU.
+It covers four things: the dependency file, the Kaggle credentials,
+LoRA on one TPU host, and the move to a multi-host slice. Read this page
+before your first fine-tuning job. For a complete tutorial with a saved
+model and an inference step, see
+[Fine-tuning Gemma 4 on TPU](gemma4_finetuning.md).
+
+## Before you start
+
+- A Kinetic cluster and an active profile. `kinetic init` creates both.
+ See [Getting Started](../getting_started.md).
+- A TPU node pool that matches the accelerator in your decorator.
+ `kinetic pool list` shows the pools of the cluster.
+ `kinetic pool add --accelerator tpu-v5litepod-8` adds one. See
+ [Clusters and Node Pools](../guides/clusters.md).
+- A [Kaggle](https://www.kaggle.com/) account. Accept the terms of the
+ model on its Kaggle page before the first download.
+- `KAGGLE_USERNAME` and `KAGGLE_KEY` set in your local shell.
+
+## Add `keras-hub` to your dependency file
+
+The image that Kinetic builds contains JAX, Keras, and Kinetic. The image
+does not contain `keras-hub`. Put a `requirements.txt` next to your
+script. List the packages that your function imports:
+
+```text
+# requirements.txt
+keras-hub
+```
+
+Kinetic finds the file, installs the packages into the image, and caches
+the image. Later runs with the same file reuse the image. If your
+function imports `keras_hub` and the file does not list `keras-hub`, the
+job fails on the pod with `ModuleNotFoundError`. Pin the versions to get
+the same packages each time Kinetic builds the image. The Gemma 4
+tutorial pins `keras==3.15.0` and `keras-hub==0.27.1`. See
+[Dependencies](../guides/dependencies.md).
+
+Two more points about the image:
-## Capturing Credentials
+- Do not list `jax`, `jaxlib`, or `libtpu`. Kinetic filters those lines
+ and installs the JAX version that matches the accelerator.
+- The pod already has `KERAS_BACKEND=jax`. You do not have to set the
+ Keras backend in your function.
-When fine-tuning models from Keras Hub or Kaggle, you often need to provide credentials (`KAGGLE_USERNAME`, `KAGGLE_KEY`). Use the `capture_env_vars` parameter to securely forward your local environment variables to the remote worker.
+## Forward the Kaggle credentials
+
+Keras Hub downloads model presets from Kaggle. The download reads
+`KAGGLE_USERNAME` and `KAGGLE_KEY` from the environment. The pod does not
+see the environment variables of your shell. List the names in
+`capture_env_vars`. Kinetic copies the values into the pod before your
+function runs.
```python
import kinetic
@kinetic.run(
- accelerator="tpu-v5litepod-1", capture_env_vars=["KAGGLE_*", "GOOGLE_CLOUD_*"]
+ accelerator="tpu-v5litepod-8",
+ capture_env_vars=["KAGGLE_USERNAME", "KAGGLE_KEY"],
)
def train_gemma():
import keras_hub
- # Credentials are automatically available in the remote environment
+ # The Kaggle credentials are set in the pod environment.
gemma_lm = keras_hub.models.Gemma3CausalLM.from_preset("gemma3_1b")
# ...
```
-## Low-Rank Adaptation (LoRA)
+A name that ends with `*` is a prefix pattern. `capture_env_vars=["KAGGLE_*"]`
+forwards every variable with the prefix `KAGGLE_`. See
+[Forward Environment Variables](../guides/env_vars.md) for the wildcard
+rules and the names that a wildcard never matches.
+
+:::{warning}
+Kinetic stores the captured values in plain text in the job payload in
+the jobs bucket. Every job pod in the cluster can read that bucket.
+Forward only the variables that the job needs. Use tokens with a short
+life if you can. Kinetic logs a warning if a captured name contains
+`KEY`, `TOKEN`, `SECRET`, `PASSWORD`, or `CREDENTIAL` in any letter
+case. The warning is information only. If you intend to forward the
+credential, take no action. See [Security](../security.md).
+:::
+
+## Fine-tune with LoRA
+
+A full fine-tune updates every weight of the model, and the optimizer
+keeps a state for every weight. Those states need a lot of memory, so a
+full fine-tune needs a large slice. LoRA (Low-Rank Adaptation) freezes
+the model weights and trains small additional matrices of a low rank.
+LoRA reduces the number of trainable parameters, so the fine-tune fits
+on a smaller slice.
-Fine-tuning large models often requires massive memory. LoRA significantly reduces the number of trainable parameters, enabling fine-tuning on smaller accelerator slices.
+Call `enable_lora(rank=...)` on the backbone before `fit`:
```python
-@kinetic.run(accelerator="tpu-v5litepod-8")
+@kinetic.run(
+ accelerator="tpu-v5litepod-8",
+ capture_env_vars=["KAGGLE_USERNAME", "KAGGLE_KEY"],
+)
def train_lora():
import keras_hub
gemma_lm = keras_hub.models.GemmaCausalLM.from_preset("gemma_2b_en")
- # Enable LoRA (rank=4)
- print("Enabling LoRA...")
+ # Freeze the backbone and add rank-4 LoRA matrices.
gemma_lm.backbone.enable_lora(rank=4)
- # Train as usual
+ # A short sequence length keeps the activation memory small.
+ gemma_lm.preprocessor.sequence_length = 128
+
+ train_data = [
+ "Question: What is the capital of India? Answer: New Delhi.",
+ "Question: What is the capital of South Africa? Answer: Pretoria.",
+ ]
gemma_lm.fit(train_data, epochs=3)
- return "Training complete!"
+ return "Training complete"
```
-## Distributed Fine-tuning
+The `rank` value sets the size of the LoRA matrices. A larger rank
+trains more parameters and needs more memory. The Gemma 4 tutorial uses
+`rank=4` for a 26B-parameter model on `tpu-v5litepod-8`.
-For larger models or datasets, use the Pathways backend to distribute training across multiple TPU hosts.
+## Save the adapted weights
+
+Kubernetes deletes the pod filesystem when the pod ends. Write the LoRA
+weights and every other file that you want to keep under
+`KINETIC_OUTPUT_DIR`. Kinetic sets that variable in every pod to a Cloud
+Storage location, by default `gs://{jobs bucket}/outputs/{job_id}`.
+Kinetic does not delete those files as part of the job cleanup. But a
+lifecycle rule on the jobs bucket deletes objects that are older than
+30 days. See
+[Outputs and Checkpoints](../guides/checkpointing.md). The
+[Gemma 4 tutorial](gemma4_finetuning.md) shows a complete save step and
+a second job that loads the weights for inference.
+
+## Distributed fine-tuning
+
+If the model or the batch does not fit on one host, select a multi-host
+slice, for example `tpu-v6e-16` or `tpu-v5litepod-16`. Each of these
+slices has 16 chips on four hosts. Kinetic reads the host count from the
+accelerator name and selects the `pathways` backend for you. You do not
+have to set `backend`. Kinetic runs one pod per host, and JAX sees the
+devices of every host.
+
+The Keras distribution API needs a device mesh over all devices. Build
+the mesh inside the function. Then load the model:
```python
-@kinetic.run(accelerator="tpu-v6e-8", backend="pathways")
+@kinetic.run(
+ accelerator="tpu-v6e-16",
+ capture_env_vars=["KAGGLE_USERNAME", "KAGGLE_KEY"],
+)
def train_distributed():
import keras
- import jax
- # Multi-host TPU environment is auto-initialized
- # ...
+ import keras_hub
+
+ # Every device on every host.
+ devices = keras.distribution.list_devices()
+ device_mesh = keras.distribution.DeviceMesh(
+ shape=(len(devices),),
+ axis_names=["batch"],
+ devices=devices,
+ )
+ keras.distribution.set_distribution(
+ keras.distribution.DataParallel(device_mesh=device_mesh)
+ )
+
+ gemma_lm = keras_hub.models.GemmaCausalLM.from_preset("gemma_2b_en")
+ gemma_lm.backbone.enable_lora(rank=4)
+ gemma_lm.preprocessor.sequence_length = 128
+
+ train_data = [...] # one string per example
+ gemma_lm.fit(train_data, batch_size=len(devices), epochs=3)
+ return "Training complete"
```
-See the [Distributed Training](../guides/distributed_training.md) guide for more details on scaling your workloads.
+Three things change when a job runs on more than one host:
+
+- The cluster needs a node pool for the slice, for example
+ `kinetic pool add --accelerator tpu-v6e-16`.
+- Your terminal shows the log of the leader host only.
+- Every host uploads a return value, and Kinetic returns the value of
+ the last host that wrote. Return the same value from every host.
+
+:::{warning}
+If your model and your batch fit on one host, stay on one host. A
+multi-host job starts more slowly, needs the LeaderWorkerSet controller,
+and fails as a whole if one host fails. See
+[Distributed Training](../guides/distributed_training.md) before you
+move to a multi-host slice.
+:::
+
+## Complete examples
+
+- [Fine-tuning Gemma 4 on TPU](gemma4_finetuning.md) — a full tutorial:
+ LoRA on `tpu-v5litepod-8`, weights saved under `KINETIC_OUTPUT_DIR`,
+ and inference in a second job.
+- [Single-TPU Gemma 3 fine-tune](gemma3_sft_demo.md) — the shortest
+ script: `gemma3_1b` on `tpu-v5litepod-1` with the Kaggle credentials
+ forwarded.
+- [Distributed Gemma 2B fine-tune](gemma_sft_pathways_distributed.md) —
+ Gemma 2B with LoRA and the Keras `DataParallel` distribution. The
+ script sets `backend="pathways"` on a single-host slice to test the
+ multi-host code path on one host.
## Related pages
::::{grid} 1 1 2 2
:gutter: 3
-:::{grid-item-card} {octicon}`server;1em` Distributed Training
-:link: ../guides/distributed_training
+:::{grid-item-card} {octicon}`rocket;1em` Fine-tuning Gemma 4 on TPU
+:link: gemma4_finetuning
:link-type: doc
-Multi-host Pathways setup that LLM fine-tuning typically needs.
+The complete tutorial: LoRA, saved weights, and inference.
:::
-:::{grid-item-card} {octicon}`history;1em` Checkpointing
-:link: ../guides/checkpointing
+:::{grid-item-card} {octicon}`key;1em` Forward Environment Variables
+:link: ../guides/env_vars
:link-type: doc
-Orbax + `KINETIC_OUTPUT_DIR` for resumable fine-tuning runs.
+`capture_env_vars`, the wildcard rules, and how to handle secrets.
:::
-:::{grid-item-card} {octicon}`book;1em` Examples
-:link: ../examples
+:::{grid-item-card} {octicon}`server;1em` Distributed Training
+:link: ../guides/distributed_training
+:link-type: doc
+
+Multi-host slices, the `pathways` backend, and the return value rules.
+:::
+
+:::{grid-item-card} {octicon}`history;1em` Outputs and Checkpoints
+:link: ../guides/checkpointing
:link-type: doc
-The Gemma SFT examples are full end-to-end LLM fine-tuning walkthroughs.
+`KINETIC_OUTPUT_DIR`, retention, and resumable fine-tuning runs.
:::
::::
diff --git a/docs/examples/pytorch_training.md b/docs/examples/pytorch_training.md
index 7dcc29be..a8133390 100644
--- a/docs/examples/pytorch_training.md
+++ b/docs/examples/pytorch_training.md
@@ -1,19 +1,74 @@
# PyTorch Training
-Kinetic can run PyTorch workloads on cloud GPUs. Since Kinetic executes arbitrary Python functions remotely, any PyTorch code that runs locally will run the same way on a provisioned GPU node.
+Kinetic runs a PyTorch function on a cloud GPU node in the same way that
+it runs a Keras or a JAX function. You add `torch` to your dependency
+file, you decorate the function with a GPU accelerator, and you call the
+function. This page shows a first GPU job and a job on several GPUs of
+one node. It also shows how to keep outputs and how to run on Spot
+capacity. Read [Getting Started](../getting_started.md) before this page.
-## Setup
+## Before you start
-Add `torch` to your project's `requirements.txt`:
+- You have run `kinetic init`, so an active profile supplies the project,
+ the zone, and the cluster. You do not pass those values in code.
+- Your cluster has a GPU node pool. List the pools with
+ `kinetic pool list`. If no GPU pool exists, add one:
+
+ ```bash
+ kinetic pool add --accelerator gpu-l4
+ ```
+
+ Use the same accelerator string for the pool and for the job. The pool
+ then has the GPU type and the GPU count that the job requests. When a
+ cluster has a GPU node pool, Kinetic installs the NVIDIA driver on the
+ GPU nodes. You do not install the driver.
+
+## Add `torch` to the dependency file
+
+Put a `requirements.txt` next to your script:
```text
torch
torchvision
```
-Kinetic will install these in the remote container automatically. See [Managing Dependencies](../guides/dependencies.md) for details on how dependency detection works.
+Kinetic finds the file and installs the packages into the image that it
+builds for the job. A `pyproject.toml` with a `[project.dependencies]`
+list works too. See [Dependencies](../guides/dependencies.md) for the
+rules that select the file.
+
+The image that Kinetic builds for a GPU job starts from a `python` image
+that matches your local Python minor version. Kinetic installs JAX with
+CUDA support, Keras, `cloudpickle`, `google-cloud-storage`, and the
+`keras-kinetic` package into that image. PyTorch is not in that set, so
+`torch` must be in your dependency file. The Linux `torch` wheels on
+PyPI include the CUDA libraries, and the GPU node supplies the NVIDIA
+driver. You do not add CUDA packages to the file.
+
+:::{note}
+The pod runs with `KERAS_BACKEND=jax`. If your function uses Keras 3
+with the PyTorch backend, set the variable inside the function before
+you import Keras:
+
+```python
+import os
+
+os.environ["KERAS_BACKEND"] = "torch"
+import keras
+```
+
+If `KERAS_BACKEND=torch` is set in your local shell, you can also pass
+`capture_env_vars=["KERAS_BACKEND"]` to the decorator. The exact name is
+necessary. A wildcard pattern never captures `KERAS_BACKEND`. See
+[Forward Environment Variables](../guides/env_vars.md).
+:::
-## Basic Usage
+## Run on one GPU
+
+`accelerator="gpu-l4"` requests one NVIDIA L4 GPU. The pod requests one
+GPU from Kubernetes, and `torch.cuda.is_available()` returns `True` on
+the pod. Keep the `torch` imports inside the function, so that your local
+process does not need PyTorch.
```python
import kinetic
@@ -56,9 +111,31 @@ def train():
final_loss = train()
```
-## Multi-GPU Training
+The call blocks until the job ends. Kinetic streams the `print` output to
+your terminal and returns the final loss as a Python float. The first run
+needs 5 to 10 minutes before the function starts, because Kinetic builds
+the image with `torch` and the autoscaler starts a GPU node. Later runs
+with the same dependency file reuse the image and start in less than
+1 minute while a node still runs.
+
+For a job that runs longer than a few minutes, call `train.run_async()`
+instead. That call returns a `JobHandle` as soon as Kinetic submits the
+job. See [Detached Jobs](../guides/async_jobs.md).
+
+## Use several GPUs on one node
+
+Append `xN` to the GPU name to request `N` GPUs on one node. Every GPU
+job is single-host: all GPUs of the job are on one VM, and the job runs
+in one pod. `gpu-a100x4` requests four A100 40GB GPUs on one
+`a2-highgpu-4g` VM. Add a node pool for that string before you run the
+job:
+
+```bash
+kinetic pool add --accelerator gpu-a100x4
+```
-For nodes with multiple GPUs, use `torch.nn.DataParallel` to split batches across devices.
+On the pod, `torch.cuda.device_count()` returns `4`. Use
+`torch.nn.DataParallel` to split each batch across the GPUs:
```python
import kinetic
@@ -95,17 +172,113 @@ def train_multi_gpu():
return loss.item()
```
-## GPU Selection
+Kinetic starts one Python process in the pod and does not run `torchrun`.
+If you use `DistributedDataParallel`, start the worker processes from
+inside the function, for example with `torch.multiprocessing.spawn`. Use
+`localhost` as the rendezvous address, because all workers are in one
+pod. Kinetic does not run a GPU job across more than one node.
+Multi-host jobs are available for TPU slices only. See
+[Distributed Training](../guides/distributed_training.md).
-See [Accelerator Support](../accelerators.md) for the full list of GPUs, multi-GPU counts, and TPU configurations.
+## Bring data and keep outputs
-Use `spot=True` to reduce costs for fault-tolerant workloads:
+Two rules from the other examples also apply to PyTorch:
+
+- **Inputs.** Wrap a local file, a local directory, or a `gs://` URI in
+ `kinetic.Data(...)`. Kinetic uploads local data one time and gives your
+ function a plain path on the pod. See
+ [Working with Data](../guides/data.md).
+- **Outputs.** Kubernetes deletes the pod and its filesystem when the
+ job ends. Write every file that you want to keep under
+ `KINETIC_OUTPUT_DIR`, a Cloud Storage prefix that Kinetic sets for
+ each job. See [Outputs and Checkpoints](../guides/checkpointing.md).
+
+`KINETIC_OUTPUT_DIR` is a `gs://` URI, and `torch.save` writes to a local
+file. Save the file to the local disk of the pod. Then upload the file
+with the `google-cloud-storage` client, which the image contains:
```python
-@kinetic.run(accelerator="gpu-a100", spot=True)
-def train(): ...
+import os
+
+import kinetic
+
+
+@kinetic.run(accelerator="gpu-l4")
+def train_and_save():
+ import torch
+ from google.cloud import storage
+
+ model = ... # build and train the model
+
+ local_path = "/tmp/model.pt"
+ torch.save(model.state_dict(), local_path)
+
+ # KINETIC_OUTPUT_DIR is gs:///outputs/ by default.
+ bucket_name, prefix = os.environ["KINETIC_OUTPUT_DIR"][5:].split("/", 1)
+ blob = storage.Client().bucket(bucket_name).blob(f"{prefix}/final/model.pt")
+ blob.upload_from_filename(local_path)
+ return f"gs://{bucket_name}/{prefix}/final/model.pt"
```
+The default output directory contains the job ID, so each call gets an
+empty directory. If a second call must find the checkpoints of the first
+call, pass the same `output_dir=` to the decorator on both calls.
+
+## Select a GPU
+
+The `accelerator` string names the GPU type and, with the `xN` suffix,
+the GPU count. Kinetic accepts these GPU names: `l4`, `t4`, `v100`,
+`a100`, `a100-80gb`, `h100`, `p4`, and `p100`. The `gpu-` prefix is
+optional, so `gpu-l4` and `l4` name the same hardware.
+[Accelerators](../accelerators.md) lists the counts and the machine
+types for each GPU. `kinetic accelerators` prints the same list in the
+shell, and `kinetic accelerators --live` marks each accelerator type
+that has a node pool on your cluster.
+
+A job runs only on a node pool with the same GPU type, and each node of
+the pool must have at least the requested GPU count. If no pool has the
+GPU type, a blocking call stops with the error `No GKE node pool
+exists`, and a detached job stays `PENDING`. If a pool has the GPU type
+but too few GPUs on each node, the job stays `PENDING`. Add a pool with
+`kinetic pool add`, or change the `accelerator` string to a pool that
+the cluster has.
+
+## Run on Spot capacity
+
+`spot=True` on the decorator asks for a Spot node. Spot VMs cost less,
+but Google Cloud can preempt a Spot VM at any time. Spot use in Kinetic
+has two parts, and both parts are necessary:
+
+1. Add a Spot node pool for the accelerator:
+
+ ```bash
+ kinetic pool add --accelerator gpu-a100 --spot
+ ```
+
+2. Set `spot=True` on the job. The accelerator string
+ `"gpu-a100:spot"` has the same effect:
+
+ ```python
+ @kinetic.run(accelerator="gpu-a100", spot=True)
+ def train(): ...
+ ```
+
+The job side adds the Spot node selector and the Spot toleration to the
+pod. Without `spot=True`, the pod cannot schedule on the Spot pool, and
+the job stays `PENDING`. Without a `--spot` pool, a job with `spot=True`
+has no node to run on: a blocking call stops with the error `No GKE node
+pool exists`, and a detached job stays `PENDING`.
+
+:::{warning}
+If Google Cloud preempts the node, the job ends as `FAILED`, and Kinetic
+does not submit the job again. Use Spot only for a job that can restart
+from a checkpoint. Write checkpoints under `KINETIC_OUTPUT_DIR`, and pass
+a fixed `output_dir=` so that the next submission finds them.
+:::
+
+See [Spot VMs](../guides/cost_optimization.md#spot-vms) for the full
+rules.
+
## Related pages
::::{grid} 1 1 2 2
@@ -115,20 +288,27 @@ def train(): ...
:link: ../guides/dependencies
:link-type: doc
-How `torch` gets installed in the remote container.
+How Kinetic finds your dependency file and installs `torch` into the image.
:::
:::{grid-item-card} {octicon}`cpu;1em` Accelerators
:link: ../accelerators
:link-type: doc
-Full list of GPUs and multi-GPU configurations.
+Every GPU name, GPU count, and machine type that Kinetic accepts.
+:::
+
+:::{grid-item-card} {octicon}`history;1em` Outputs and Checkpoints
+:link: ../guides/checkpointing
+:link-type: doc
+
+`KINETIC_OUTPUT_DIR`, retention, and how to resume a job from a checkpoint.
:::
:::{grid-item-card} {octicon}`zap;1em` Cost Optimization
:link: ../guides/cost_optimization
:link-type: doc
-Spot capacity for GPU workloads.
+Spot pools, scale to zero, and warm nodes for GPU jobs.
:::
::::
diff --git a/docs/getting_started.md b/docs/getting_started.md
index f0a33e0f..23cf2ca1 100644
--- a/docs/getting_started.md
+++ b/docs/getting_started.md
@@ -1,162 +1,228 @@
# Getting Started
-Install Kinetic, point it at a cluster, and run your first remote
-function. If your team has already provisioned a Kinetic cluster, skip
-ahead to [Run your first job](#run-your-first-job).
+This page takes you from an empty machine to your first remote job. You
+install Kinetic, run `kinetic init` one time, and run a Keras training
+script on a cloud TPU. If your team already has a Kinetic cluster,
+`kinetic init` connects you to that cluster.
-## Prerequisites
+## Before you start
-- Python 3.11+.
-- [uv](https://docs.astral.sh/uv/getting-started/installation/), used
- for the install command below.
-- Google Cloud SDK (`gcloud`): [install guide](https://cloud.google.com/sdk/docs/install).
+You need these tools and accounts:
+
+- Python 3.11 or later.
+- [uv](https://docs.astral.sh/uv/getting-started/installation/), for the
+ install command below. `pip` also works.
+- The Google Cloud SDK (`gcloud`): [install guide](https://cloud.google.com/sdk/docs/install).
- `kubectl`: [install guide](https://kubernetes.io/docs/tasks/tools/).
- Kinetic auto-installs the `gke-gcloud-auth-plugin` for you on first
- use, but `kubectl` itself must already be on your `PATH`.
+ Kinetic installs the `gke-gcloud-auth-plugin` for you on first use, but
+ `kubectl` itself must be on your `PATH`.
- A Google Cloud project with [billing enabled](https://docs.cloud.google.com/billing/docs/how-to/modify-project).
-Authenticate with Google Cloud once:
+Log in to Google Cloud one time:
```bash
gcloud auth login
gcloud auth application-default login
```
-## Install
+## Step 1: Install Kinetic
```bash
uv pip install keras-kinetic
```
-This installs the `@kinetic.run()` decorator and the `kinetic` CLI,
-which provisions and manages infrastructure (Pulumi and the GCP
-plugins).
+This command installs two things:
+
+- The `kinetic` Python package, with the `@kinetic.run()` decorator.
+- The `kinetic` command-line tool, which creates and manages your cloud
+ infrastructure.
:::{note}
-The [Pulumi](https://www.pulumi.com/) CLI (used for infrastructure
-provisioning) is bundled and managed automatically. It will be
-installed to `~/.kinetic/pulumi` on first use if not already present.
+Kinetic uses [Pulumi](https://www.pulumi.com/) to create cloud resources.
+Kinetic downloads the Pulumi CLI to `~/.kinetic/pulumi-cli` on first
+use. You do not install Pulumi yourself.
:::
-## Set up your environment
+## Step 2: Run `kinetic init`
```bash
kinetic init
```
-`kinetic init` checks your local tools, auth, and project, then routes
-you down one of two paths:
-
-- **Join** — if any Kinetic clusters already exist in this GCP project
- (provisioned by you or a teammate), `init` lists them, lets you pick
- one, and configures `kubectl` for it. Cluster discovery reads the
- project's shared state bucket
- (`gs://{project}-kinetic-state`), so collaborators with access to
- the bucket all see the same set.
-- **Create** — if no clusters exist yet, `init` calls `kinetic up` to
- enable APIs, provision a GKE cluster with an accelerator node pool,
- and wire up Docker / `kubectl` access.
-
-Either way, `init` ends by saving a **profile** and making it active.
-A profile is your saved infrastructure context like project, zone,
-cluster, and namespace - persisted at `~/.kinetic/profiles.json`. The
-active profile is what every `kinetic` command and every
-`@kinetic.run()` invocation targets, so you don't need to export env vars or pass `--project` / `--zone` / `--cluster` on
-the command line. Switch contexts with `kinetic profile use `, and
-see what's saved with `kinetic profile ls`.
-
-:::{admonition} Cleanup reminder
-:class: warning
-
-When you're done, run `kinetic down` to tear down all resources and
-stop incurring costs. See the [CLI Reference](cli) for the full set
-of commands.
+`kinetic init` checks your local tools and your login, and asks for your
+project ID. If the project does not exist, `init` offers to create it and
+to link a billing account. `init` then offers one of these paths:
+
+- **Join** — `init` lists the Kinetic clusters that already exist in
+ the project (yours or a teammate's), lets you select one, and
+ configures `kubectl` for it. Kinetic keeps the infrastructure state in
+ a bucket that the whole project shares, so every collaborator sees the
+ same list.
+- **Create** — `init` runs `kinetic up`. `kinetic up` enables the Google
+ Cloud APIs, creates a GKE cluster, adds one **node pool** (a group of
+ VMs of one accelerator type) of your choice, and configures `kubectl`.
+- **Troubleshoot** — `init` runs diagnostics and prints a fix command for
+ each failed check.
+
+`init` asks you which path to take. **Join** is available only when a
+cluster exists. When a prerequisite check fails, `init` offers
+**Troubleshoot** directly. Both the **Join** path and the **Create** path
+end with a saved **profile**, and that profile becomes active.
+
+:::{admonition} What is a profile?
+:class: tip
+
+A profile is a saved set of four values: the project, the zone, the
+cluster, and the Kubernetes namespace. Kinetic stores profiles in
+`~/.kinetic/profiles.json`. Every `kinetic` command and every
+`@kinetic.run()` call reads the active profile. You therefore do not
+export environment variables and you do not pass `--project` on the
+command line. `kinetic profile ls` lists your profiles.
+`kinetic profile use NAME` makes a different profile active. See
+[Profiles](guides/profiles.md).
:::
-:::{admonition} Sharing infrastructure with teammates?
-:class: seealso
+When `init` finishes, check the cluster:
-Kinetic stores Pulumi state in a per-project GCS bucket
-(`gs://{project}-kinetic-state`), so any teammate with
-`roles/storage.objectAdmin` on the bucket sees the same stack. The
-first `kinetic up` creates the bucket; the first admin needs
-`roles/storage.admin` on the project. See
-[Pulumi state](configuration.md#pulumi-state) for the full IAM story.
-:::
+```bash
+kinetic status
+```
+
+## Step 3: Make sure that a node pool matches your job
+
+A job runs only on a node pool with the same accelerator type and, for a
+TPU, the same slice shape. The **Create** path adds one node pool during
+`kinetic up`. The example script below uses
+`accelerator="tpu-v5litepod-1"`. List the pools of your cluster:
+
+```bash
+kinetic pool list
+```
-## Run your first job
+If the list has no `v5litepod` pool with the `1x1` topology, do one of
+these two things:
+
+- Add one: `kinetic pool add --accelerator tpu-v5litepod-1`.
+- Change `accelerator=` in the script to an accelerator that the cluster
+ has. For example, use `accelerator="tpu-v5litepod-4"` for a `2x2` v5e
+ pool, `accelerator="gpu-l4"` for an L4 GPU pool, or
+ `accelerator="cpu"` for a run without an accelerator.
+
+See [Accelerators](accelerators.md) for the accelerator names.
+
+## Step 4: Run your first job
+
+Save this script as `fashion_mnist.py`:
```{literalinclude} ../examples/fashion_mnist.py
:language: python
```
-Run it:
+Run the script:
```bash
python fashion_mnist.py
```
+Kinetic sends the function to the cluster, streams the log lines to your
+terminal, and prints the return value when the job ends.
+
:::{note}
-**Expected timing:**
-
-- **First run:** ~5 minutes. The slow part is the first container
- build via Cloud Build, which freezes your dependencies into an
- image tagged by their hash.
-- **Subsequent runs (same dependencies):** under a minute. The
- cached image is reused; only your code changes get re-uploaded.
-- **Subsequent runs (changed dependencies):** ~5 minutes again,
- since a new hash forces a fresh build.
- :::
-
-:::{tip}
-**Recommended defaults:**
-
-- Stay in **bundled mode** (the default — you don't need to pass
- `container_image=`). It's the only mode that works without
- publishing your own base image.
-- Use direct calls to **`@kinetic.run()`** decorated functions while you're
- iterating; switch to calling **`run_async()`** once your jobs run for
- more than a few minutes and you'd rather not block your local shell.
-- Write any artifacts you want to keep under `KINETIC_OUTPUT_DIR`,
- not under `/tmp`.
- :::
+**Expected time:**
+
+- **First run:** 5 to 10 minutes. Kinetic builds a container image with
+ your dependencies through Cloud Build, and the autoscaler starts a TPU
+ node. Kinetic caches the image. This script needs no dependency file,
+ because the image already contains Keras. Later, a `requirements.txt`
+ or a `pyproject.toml` next to your script decides the packages. See
+ [Dependencies](guides/dependencies.md).
+- **Later runs with the same dependencies:** less than 1 minute while a
+ node still runs. Kinetic reuses the cached image and uploads only your
+ code. After about 10 idle minutes the node scales down, and the next
+ run waits 2 to 5 minutes for a new node.
+- **Later runs after a change to the dependencies:** 5 to 10 minutes
+ again, because Kinetic builds a new image.
+:::
-## Next steps
+## What happened
+
+Kinetic did five things when you called `train_fashion_mnist()`:
+
+:::{container} kinetic-steps
+1. **Package.** Kinetic serialized the function and archived the
+ package root, the project directory that holds the script.
+2. **Build.** Kinetic built a container image with your dependencies, or
+ reused a cached image.
+3. **Schedule.** Kinetic created a Kubernetes Job on your cluster. The
+ cluster autoscaler started a TPU node for it.
+4. **Run.** The pod ran the function and streamed the logs to your
+ terminal.
+5. **Collect.** The pod uploaded the return value to Cloud Storage.
+ Kinetic downloaded the value, returned it, and deleted the job
+ resources.
+:::
+
+[How Kinetic Works](concepts.md) explains each phase, the vocabulary,
+and the choices that you make as your jobs get larger.
-After your first run works, the most useful follow-ups are:
+## Two habits for every job
+
+- **Use a blocking call while you iterate.** A blocking call, such as
+ `train_fashion_mnist()`, blocks until the job ends. When a job runs for
+ more than a few minutes, call `train_fashion_mnist.run_async()`
+ instead. That call returns a `JobHandle` as soon as Kinetic submitted
+ the job. See [Detached Jobs](guides/async_jobs.md).
+- **Write files that you want to keep under `KINETIC_OUTPUT_DIR`.** The
+ pod sets that environment variable to a Cloud Storage location that
+ stays after the pod ends. The pod discards the files under `/tmp` when
+ it ends. See [Outputs and Checkpoints](guides/checkpointing.md).
+
+## Clean up
+
+The cluster control plane costs money while the cluster exists, even
+when no job runs. When you no longer need the cluster, delete it:
+
+```bash
+kinetic down
+```
+
+`kinetic down` deletes the cluster, the node pools, the Artifact
+Registry repository, and the Cloud Storage buckets of the cluster,
+including the job outputs. See [Clusters and Node Pools](guides/clusters.md).
+
+## Next steps
::::{grid} 1 1 2 2
:gutter: 3
-:::{grid-item-card} {octicon}`code-square;1em` Examples
-:link: examples
+:::{grid-item-card} {octicon}`light-bulb;1em` How Kinetic Works
+:link: concepts
:link-type: doc
-A catalog of runnable scripts that cover async jobs, data,
-checkpoints, parallel sweeps, and LLM fine-tuning. The fastest way to
-see real patterns end to end.
+The vocabulary and the job lifecycle. Read this page before the guides.
:::
-:::{grid-item-card} {octicon}`package;1em` Execution Modes
-:link: guides/execution_modes
+:::{grid-item-card} {octicon}`code-square;1em` Examples
+:link: examples
:link-type: doc
-Bundled vs prebuilt vs custom image, and when to switch.
+Runnable scripts for detached jobs, data, checkpoints, parallel sweeps,
+and LLM fine-tuning.
:::
-:::{grid-item-card} {octicon}`clock;1em` Detached Jobs
-:link: guides/async_jobs
+:::{grid-item-card} {octicon}`database;1em` Working with Data
+:link: guides/data
:link-type: doc
-`run_async()`, reattach, and the job lifecycle for long-running work.
+`kinetic.Data(...)` for inputs. See also
+[Outputs and Checkpoints](guides/checkpointing.md) for
+`KINETIC_OUTPUT_DIR`.
:::
-:::{grid-item-card} {octicon}`database;1em` Data & Checkpointing
-:link: guides/data
+:::{grid-item-card} {octicon}`clock;1em` Detached Jobs
+:link: guides/async_jobs
:link-type: doc
-`kinetic.Data(...)` for inputs and `KINETIC_OUTPUT_DIR` for durable
-outputs and resumable checkpoints. See also
-[Checkpointing](guides/checkpointing.md).
+`run_async()`, the job lifecycle, and how to reattach to a job.
:::
::::
diff --git a/docs/guides/advanced.md b/docs/guides/advanced.md
deleted file mode 100644
index 88885d9d..00000000
--- a/docs/guides/advanced.md
+++ /dev/null
@@ -1,47 +0,0 @@
-# Advanced Workflows
-
-Guides for scaling Kinetic across multiple clusters, using capacity reservations, and managing long-running or batched jobs.
-
-::::{grid} 1 1 2 2
-:gutter: 3
-
-:::{grid-item-card} {octicon}`clock;1em` Detached Jobs
-:link: async_jobs
-:link-type: doc
-
-Submit with `run_async()`, then poll status, tail logs, and collect
-results — even from a different machine.
-:::
-
-:::{grid-item-card} {octicon}`stack;1em` Batched Jobs
-:link: batched_jobs
-:link-type: doc
-
-Run the same function over many inputs with `run_async_map()` and manage
-the whole fan-out through one `BatchHandle`.
-:::
-
-:::{grid-item-card} {octicon}`server;1em` Multiple Clusters
-:link: clusters
-:link-type: doc
-
-Run multiple independent clusters in the same GCP project for isolation,
-regions, and environments.
-:::
-
-:::{grid-item-card} {octicon}`cpu;1em` Capacity Reservations
-:link: reservations
-:link-type: doc
-
-Guarantee accelerator hardware is available when your node pool scales up.
-:::
-::::
-
-```{toctree}
-:hidden:
-
-async_jobs
-batched_jobs
-clusters
-reservations
-```
diff --git a/docs/guides/async_jobs.md b/docs/guides/async_jobs.md
index 6bb9e598..d9dbb69e 100644
--- a/docs/guides/async_jobs.md
+++ b/docs/guides/async_jobs.md
@@ -1,19 +1,25 @@
# Detached Jobs
-Most Kinetic users start with `@kinetic.run()`, which blocks the local
-process until the remote function returns. That's the right choice when the
-job is short, when you want the result inline in your script, or when
-you're iterating on code interactively.
-
-When the job is **long**, when you want to **walk away from your laptop**,
-or when you want to **fan out and check on multiple jobs in parallel**,
-switch to calling `func.run_async()`. It returns a `JobHandle` immediately
-and leaves the actual work running on the cluster. You can then poll status,
-tail logs, collect results, or reattach to the job from a different machine
-— all backed by metadata Kinetic persisted to GCS at submission time.
-
-This page covers the full submit → observe → collect → cleanup loop, both
-from Python and from the `kinetic jobs` CLI.
+A blocking call to a decorated function blocks your local process until the
+job ends. That is the right choice for a short job or for interactive
+work. A **detached job** is a job that you submit with `run_async()`. The
+call returns a `JobHandle` at once, and the job runs on the cluster
+without your local process. You then check the status, read the logs,
+collect the result, and delete the job. You do this from Python or from
+the `kinetic jobs` command group, on any machine.
+
+This page covers the loop from submit to cleanup. It shows a first
+example, the Python and CLI operations side by side, the job lifecycle,
+and how to reattach from another machine. It ends with timeouts, cleanup,
+and recommendations for long jobs.
+
+## Before you start
+
+- Complete [Getting Started](../getting_started.md). The active profile
+ supplies the project, the zone, the cluster, and the namespace for every
+ call and command on this page.
+- Read [How Kinetic Works](../concepts.md) for the vocabulary: job, job
+ ID, pod, jobs bucket, and output directory.
## A first detached job
@@ -21,59 +27,74 @@ from Python and from the `kinetic jobs` CLI.
import kinetic
-@kinetic.run(accelerator="tpu-v5e-1")
+@kinetic.run(accelerator="tpu-v5litepod-4")
def train_model():
# Long-running training code
return {"final_loss": 0.123}
job = train_model.run_async()
-print(f"Submitted: {job.job_id}")
+print(f"Submitted: {job.job_id}") # for example: job-3f9a1c2b
-# ... do something else, possibly close the script entirely ...
+# ... do other work, or exit this script ...
-final = job.result(timeout=3600) # blocks until done
+final = job.result(timeout=3600) # blocks until the job ends, or for 3600 s
print(final)
```
-The `@kinetic.run()` decorator accepts the same arguments regardless of
-whether you execute the function synchronously or asynchronously (accelerator,
-project, zone, cluster, container_image, env vars, data volumes, etc.). The
-only difference is how you invoke the function: calling it directly (e.g.,
-`train_model()`) runs it synchronously and blocks, while calling
-`train_model.run_async()` runs it asynchronously and returns a `JobHandle`
-immediately.
+`@kinetic.run()` takes the same arguments for a blocking call and for a
+detached job: `accelerator`, `volumes`, `capture_env_vars`, `output_dir`,
+and the others. Only the call changes. `train_model()` blocks and returns
+the value. `train_model.run_async()` returns a `JobHandle`.
+
+`run_async()` returns after Kinetic packages your code, builds or reuses
+the container image, uploads the artifacts, and creates the Kubernetes
+Job. If Kinetic must build a new image, `run_async()` blocks for the
+build, which takes about 5 to 10 minutes. After `run_async()` returns,
+the job does not need your local process. You can close the script.
+
+The job ID has the form `job-` plus 8 hexadecimal characters, for example
+`job-3f9a1c2b`. Read it from `job.job_id`, as the example does. Save the
+ID. With the active profile, the ID is all that you need to reattach.
## Python and CLI side by side
-Every operation is available both as a `JobHandle` method and as a
-`kinetic jobs` subcommand. Pick whichever fits your workflow.
+Each operation after submit exists as a `JobHandle` method and as a
+`kinetic jobs` subcommand. Use the one that fits your workflow.
Operation | Python | CLI
---------------- | --------------------------------- | ----------------------------------------------
-Submit | `job = train_model.run_async()` | (use the decorator from a script)
+Submit | `job = train_model.run_async()` | (no CLI command; call `run_async()` from a script)
Reattach | `job = kinetic.attach(job_id)` | (pass `` to any `kinetic jobs` subcommand)
List | `kinetic.list_jobs()` | `kinetic jobs list`
Check status | `job.status()` | `kinetic jobs status `
-Tail logs | `job.tail(n=100)` | `kinetic jobs logs --tail 100`
-Follow logs | `job.logs(follow=True)` | `kinetic jobs logs --follow`
+Read all logs | `job.logs()` | `kinetic jobs logs `
+Tail logs | `job.tail(n=100)` | `kinetic jobs logs --tail 100` (or `-n 100`)
+Follow logs | `job.logs(follow=True)` | `kinetic jobs logs --follow` (or `-f`)
Wait for result | `job.result(timeout=3600)` | `kinetic jobs result --timeout 3600`
Cancel | `job.cancel()` | `kinetic jobs cancel `
Clean up | `job.cleanup(k8s=True, gcs=True)` | `kinetic jobs cleanup `
+`--follow` and `--tail` are exclusive. `kinetic jobs logs` rejects a
+command that has both flags. Without a flag, the command prints the full
+log of the pod.
+
+Both the Python functions and the CLI read the active profile. See
+[Where `attach()` and the CLI find the cluster](#where-attach-and-the-cli-find-the-cluster).
+
## Job lifecycle
-A submitted job moves through five states (defined as `JobStatus` in
-`kinetic.job_status`):
+A job moves through five states. `JobStatus` in `kinetic.job_status`
+defines them.
```text
┌──────────┐
- run_async() ───▶ │ PENDING │ ── pod is waiting on a node
+ run_async() ───▶ │ PENDING │ ── the Job exists, no pod runs yet
└────┬─────┘
- │ pod scheduled
+ │ pod scheduled and started
▼
┌──────────┐
- │ RUNNING │ ── your function is executing
+ │ RUNNING │ ── the pod runs
└────┬─────┘
┌────────┴────────┐
▼ ▼
@@ -81,155 +102,286 @@ A submitted job moves through five states (defined as `JobStatus` in
│ SUCCEEDED │ │ FAILED │
└───────────┘ └──────────┘
- NOT_FOUND ── the k8s resource no longer exists (cleaned up,
- or never registered)
+ NOT_FOUND ── the Kubernetes resource no longer exists (deleted by
+ result(), cancel(), cleanup(), or the 10-minute timer)
```
-What each state means and what to do:
-
-- **PENDING** — Kubernetes has accepted the job but no pod is running yet.
- The cluster autoscaler may be provisioning a node; on a fresh accelerator
- pool this can take 2–5 minutes. *What to do:* wait. If it's stuck for
- much longer, run `kinetic init` and choose `troubleshoot`, and check
- your accelerator quota.
-- **RUNNING** — your function is executing inside the pod. Use
- `job.tail()` or `kinetic jobs logs --follow` to watch progress. *What to
- do:* nothing, unless you want to monitor.
-- **SUCCEEDED** — your function returned normally and Kinetic uploaded the
- result. *What to do:* call `job.result()` to get the return value. By
- default this also cleans up the k8s resource and GCS artifacts.
-- **FAILED** — the pod exited non-zero. The k8s resource is *not*
- auto-deleted so you can read logs. *What to do:* `job.tail()` or
- `kinetic jobs logs ` to see the error, then `job.cleanup()` when
- you're done debugging.
-- **NOT_FOUND** — the Kubernetes Job has already been deleted (typically
- by a successful `result()` call, or by an explicit `cleanup`). If the
- result was uploaded to GCS, `result()` can still return it; otherwise
- this state means the job is truly gone. *What to do:* if you need the
- return value, call `result()` once — it will read from GCS even after
- the pod is gone. If `result()` raises, the job is unrecoverable.
-
-The full submit-to-cleanup flow:
+What each state means, and what to do:
+
+- **PENDING** — Kubernetes accepted the Job, but no pod runs yet. The
+ cluster autoscaler starts a node if none is free. On a node pool that
+ is scaled to zero, the node start takes about 2 to 5 minutes. *What to
+ do:* wait. If the job stays `PENDING` for more than 10 minutes, run
+ `kinetic pool list` and make sure that a node pool for the accelerator
+ exists. Then run `kinetic init`, select `troubleshoot`, and check the
+ accelerator quota of the project. See
+ [Troubleshooting](../troubleshooting.md#a-job-stays-in-pending-for-more-than-10-minutes).
+- **RUNNING** — the pod runs. The pod first downloads the artifacts and
+ then calls your function. *What to do:* nothing. Use `job.tail()` or
+ `kinetic jobs logs -f` to watch the progress.
+- **SUCCEEDED** — your function returned, the pod uploaded the return
+ value, and the pod exited. *What to do:* call `job.result()` to get the
+ return value. With the default cleanup, `result()` also deletes the
+ Kubernetes Job and the Cloud Storage artifacts.
+- **FAILED** — your function raised an exception, or the pod exited with
+ a non-zero code. *What to do:* read the logs first, with `job.tail()`
+ or `kinetic jobs logs `. Then call `job.result()`. `result()`
+ raises the remote exception with the remote traceback. With the
+ default cleanup, `result()` deletes the Kubernetes Job and its pod for
+ a failed job too, so the logs are gone after that call. The GCS
+ artifacts of a failed job stay. See
+ [Cleanup and what remains](#cleanup-and-what-remains).
+- **NOT_FOUND** — the Kubernetes Job no longer exists. Four things cause
+ this state:
+ - a `result()` call with the default cleanup, on success and on
+ failure;
+ - a `cancel()` call;
+ - a `cleanup()` call;
+ - the Kubernetes timer that deletes a finished Job 10 minutes after
+ the job ends (2 hours for a job with `debug=True`); a multi-host
+ TPU job has no timer.
+
+ A job that you check one hour after it ended is therefore `NOT_FOUND`.
+ That state is normal. *What to do:* if you need the return value, call
+ `result()` one time. `result()` reads the result from Cloud Storage when the
+ artifacts still exist, and returns the value or raises the remote
+ exception. If `result()` raises `RuntimeError` with "no result payload
+ exists", the artifacts are gone and the job is not recoverable.
+
+The full flow from submit to cleanup:
:::{container} kinetic-steps
1. **Submit.**
- `run_async()` packages your code, builds (or reuses) a container image,
- uploads artifacts to GCS, creates a k8s Job, and returns a `JobHandle`.
- Status is `PENDING`.
+ `run_async()` packages your code and builds or reuses the container
+ image. It uploads the artifacts and a `handle.json` file to
+ `gs://{jobs bucket}/{job_id}/`. Then it creates a Kubernetes Job and
+ returns a `JobHandle`. The status is `PENDING`.
2. **Schedule.**
- The cluster autoscaler provisions a node if needed; the pod is
- scheduled. Status moves to `RUNNING`.
-3. **Execute.**
- Your function runs. The pod uploads its return value (or an exception
- payload) to GCS when it exits.
+ The cluster autoscaler starts a node if none is free. Kubernetes
+ schedules the pod. The status changes to `RUNNING`.
+3. **Run.**
+ The pod downloads the artifacts and calls your function. When the
+ function returns or raises, the pod uploads the return value or the
+ exception to `gs://{jobs bucket}/{job_id}/result.pkl` and exits.
4. **Finish.**
- Status moves to `SUCCEEDED` or `FAILED`.
+ The status changes to `SUCCEEDED` or `FAILED`. Kubernetes starts the
+ 10-minute timer.
5. **Collect and clean up.**
- Calling `job.result()` downloads the payload, returns it (or raises
- the user exception), and — by default — deletes both the k8s resource
- and the GCS artifacts. Status is now `NOT_FOUND` and the handle is
- spent.
+ `job.result()` downloads `result.pkl`, returns the value or raises the
+ exception, and deletes the Kubernetes Job. On success, `result()` also
+ deletes the Cloud Storage artifacts. The status is now `NOT_FOUND`. A second
+ `result()` call on a successful job fails, because the artifacts are
+ gone.
:::
-## Reattaching from another machine
+## Reattach from another machine
-The `JobHandle` is a small JSON-serializable dataclass that Kinetic
-persists to GCS at submit time. Anywhere you have Kinetic installed and
-GCP credentials for the same project, you can reconstruct it from the
-job ID:
+At submit time, Kinetic writes the `JobHandle` as a small JSON file to
+the jobs bucket. `kinetic.attach(job_id)` reads that file and rebuilds
+the handle. Reattach works from any machine that has Kinetic installed,
+Google Cloud credentials for the same project, and a profile for the same
+cluster:
```python
import kinetic
-job = kinetic.attach("v5e1-train-model-20260417-153012-abc1234")
+job = kinetic.attach("job-3f9a1c2b")
print(f"Status: {job.status().value}")
print(job.tail(n=20))
```
-If you don't remember the ID, list everything currently on the cluster:
+If you do not remember the ID, list the live jobs on the cluster:
```python
for j in kinetic.list_jobs():
- print(f"{j.job_id} {j.func_name} {j.status().value}")
+ print(f"{j.job_id} {j.func_name} {j.accelerator} {j.status().value}")
+```
+
+The CLI equivalent is `kinetic jobs list`. It prints the job ID, the
+function name, the accelerator, the backend, and the creation time.
+
+Two limits apply:
+
+- `list_jobs()` and `kinetic jobs list` show only jobs whose Kubernetes
+ resource still exists. A job that is `NOT_FOUND` does not appear in
+ the list. Keep the job ID if you need the job later.
+- A `result()` call with the default cleanup deletes the whole
+ `gs://{jobs bucket}/{job_id}/` prefix on success, including
+ `handle.json`. After that, `attach(job_id)` fails because the file no
+ longer exists.
+
+## Where `attach()` and the CLI find the cluster
+
+The jobs bucket is `gs://{project}-kn-{cluster}-jobs`, so `attach()` and
+`list_jobs()` need the project and the cluster name to find a job. Both
+functions accept `project=` and `cluster=` (`list_jobs()` also accepts
+`zone=` and `namespace=`). Kinetic resolves each value in this order, and
+the first value wins:
+
+1. The keyword argument, for example `kinetic.attach(job_id, cluster="research")`.
+2. The `KINETIC_*` environment variable, for example `KINETIC_CLUSTER`.
+3. The active profile.
+4. The built-in default.
+
+The `kinetic jobs` command group uses the same order. The flag
+(`--project`, `--zone`, `--cluster`, and `--namespace` for `list`) wins
+over the environment variable. The environment variable wins over the
+active profile. To run one command against a different profile, put
+`--profile NAME` before the subcommand:
+
+```bash
+kinetic --profile research jobs list
```
-The CLI equivalent is `kinetic jobs list`.
+With the active profile from `kinetic init`, you pass no flags. See
+[Profiles](profiles.md).
-## Timeouts and cleanup
+## Timeouts
-`result()` blocks indefinitely by default. Pass `timeout=` (in seconds) to
+`result()` blocks until the job ends. Pass `timeout=` (in seconds) to
bound the wait:
```python
try:
final = job.result(timeout=3600)
except TimeoutError:
- # Job is still running — handle is still valid; you can call .result()
- # again, .tail(), .cancel(), or just walk away.
+ # The job still runs, and the handle is still valid. You can call
+ # result() again, tail(), or cancel(). Or you can exit the script.
print(job.tail(n=50))
```
-By default `result()` cleans up after success: the k8s Job/pod and the
-GCS artifacts are deleted. Two ways to opt out:
+A `TimeoutError` does not stop the job. It only returns control to your
+script. `result()` polls the status every 5 seconds. Pass
+`stream_logs=True` to print the pod log to your terminal while `result()`
+waits.
+
+## Cleanup and what remains
+
+Three things belong to a job: the Kubernetes Job with its pod, the GCS
+artifacts under `gs://{jobs bucket}/{job_id}/`, and the files that your
+function wrote under `KINETIC_OUTPUT_DIR`. Cleanup touches the first two
+only.
+
+`result()` deletes the Kubernetes Job by default. It does this in every
+case: on success and on failure. `result()` deletes the Cloud Storage artifacts
+only when it collected a result, that is, on success. The artifacts of a
+failed job stay in the bucket, so `result()` can raise the same remote
+exception again from another machine.
+
+To keep the Kubernetes Job and the pod log after `result()`, pass
+`cleanup=False`. Delete the Job later with `cleanup()`:
```python
-final = job.result(cleanup=False) # keep everything
-job.cleanup(k8s=True, gcs=False) # later: delete pod, keep artifacts
+final = job.result(cleanup=False) # keep the Job and the artifacts
+job.cleanup(k8s=True, gcs=False) # later: delete the Job, keep the artifacts
```
-Failed jobs are not auto-cleaned, so logs survive until you delete them.
-Anything you wrote under `KINETIC_OUTPUT_DIR` is also kept regardless of
-cleanup — see [Checkpointing](../guides/checkpointing.md).
-
-## Recommendations for long-running jobs
-
-The following practices reduce the cost of failures on jobs that run for
-hours.
-
-- **Checkpoint regularly.** Anything written to `KINETIC_OUTPUT_DIR`
- survives a failed pod, but only the checkpoints already written can be
- used on resume. Pick a cadence that bounds how much progress a restart
- would lose. See [Checkpointing](../guides/checkpointing.md) for resume
- patterns.
-- **Persist the `job_id`.** Record it via stdout, a log file, or your
- workflow's tracking system. With the ID, you can reattach from any
- machine that has Kinetic installed and access to the same GCP project.
-- **Do not rely on the local Python process.** Once `run_async()` returns,
- the local script is no longer involved in the job's execution.
- Interrupting it (for example, with `Ctrl-C`) does not affect the
- remote job.
-- **Avoid `--follow` for jobs that run for hours.** Continuous log
- streaming is sensitive to transient network failures. Use
- `kinetic jobs logs --tail 200` from a fresh shell to check in
- periodically instead.
-- **Retain artifacts on multi-host or expensive jobs.** Pass
- `cleanup=False` to the first successful `result()` call so the
- Kubernetes resources and GCS artifacts remain available for
- inspection. Call `cleanup` explicitly once they are no longer needed.
+To read the logs of a failed job, do one of these:
+
+- Read the logs before you call `result()`, with `job.tail()`,
+ `job.logs()`, or `kinetic jobs logs `.
+- Call `job.result(cleanup=False)`, read the logs, and call
+ `job.cleanup()` when you are done.
+
+:::{note}
+If nobody calls `result()`, `cancel()`, or `cleanup()`, Kubernetes
+deletes a finished Job and its pod without a call from you: 10 minutes after the job
+ends, or 2 hours after the end for a job with `debug=True`. The pod log
+is deleted with the pod. The Cloud Storage artifacts stay until you delete them or until
+the 30-day rule of the jobs bucket deletes them. A multi-host TPU job
+uses a LeaderWorkerSet resource, and Kinetic sets no timer on that
+resource. Call `result()`, `cancel()`, or `cleanup()` to delete it. See
+[Distributed Training](distributed_training.md).
+:::
+
+The cleanup operations in detail:
+
+- `job.cancel()` deletes the Kubernetes Job and its pod, and keeps the
+ Cloud Storage artifacts. Use it to stop a running job. The status becomes
+ `NOT_FOUND`.
+- `job.cleanup(k8s=True, gcs=True)` deletes one or both parts. Set
+ `k8s=False` or `gcs=False` to keep a part.
+- A job with `debug=True` keeps its Kubernetes Job after `result()`
+ because `cleanup` defaults to `False` for debug jobs. See
+ [Interactive Debugging](debugging.md).
+
+The CLI has the same options:
+
+```bash
+kinetic jobs result --no-cleanup # collect, keep the Job and the artifacts
+kinetic jobs cleanup --no-gcs # delete the Job, keep the artifacts
+kinetic jobs cleanup --no-k8s # delete the artifacts, keep the Job
+kinetic jobs cancel # stop the job, keep the artifacts
+```
+
+`kinetic jobs result`, `cancel`, and `cleanup` also accept
+`--cleanup-timeout` (default 180 seconds) and `--cleanup-poll-interval`
+(default 2 seconds). Kinetic waits up to `--cleanup-timeout` for
+Kubernetes to confirm the deletion.
+
+Kinetic never deletes the files under `KINETIC_OUTPUT_DIR` as part of
+job cleanup. The default output directory is
+`gs://{jobs bucket}/outputs/{job_id}`, which is outside the
+`{job_id}/` prefix. The 30-day rule of the jobs bucket applies to those
+files. See [Outputs and Checkpoints](checkpointing.md).
+
+## Recommendations for long jobs
+
+These practices reduce the cost of a failure in a job that runs for hours.
+
+- **Write checkpoints at a regular interval.** The files under
+ `KINETIC_OUTPUT_DIR` survive a failed pod, but a restart can use only
+ the checkpoints that exist. Select an interval that bounds the work
+ that a restart loses. See [Outputs and Checkpoints](checkpointing.md)
+ for the resume pattern.
+- **Save the job ID.** Print it, write it to a log file, or record it in
+ your experiment tracker. With the ID, you reattach from any machine
+ that has Kinetic, credentials for the project, and a profile for the
+ cluster.
+- **Do not depend on the local Python process.** After `run_async()`
+ returns, the local script has no part in the job. If you stop the
+ script, for example with `Ctrl-C`, the remote job continues.
+- **Do not follow the logs of a job that runs for hours.** A log stream
+ breaks on a short network failure. Read the last lines from a new shell
+ with `kinetic jobs logs -n 200` instead, at the interval that you
+ choose.
+- **Keep the artifacts of a multi-host or high-cost job.** Pass
+ `cleanup=False` to the first `result()` call, so that the Kubernetes
+ resources and the Cloud Storage artifacts stay for inspection. Call `cleanup()`
+ when you no longer need them.
+- **Read the logs of a failed job before you collect it.** A `result()`
+ call with the default cleanup deletes the pod, and the pod log with it.
## Related pages
::::{grid} 1 1 2 2
:gutter: 3
-:::{grid-item-card} {octicon}`history;1em` Checkpointing
+:::{grid-item-card} {octicon}`history;1em` Outputs and Checkpoints
:link: checkpointing
:link-type: doc
-Make long jobs resumable.
+Write durable outputs and make a long job resumable.
+:::
+
+:::{grid-item-card} {octicon}`stack;1em` Batched Jobs
+:link: batched_jobs
+:link-type: doc
+
+Run one function over many inputs with `run_async_map()`.
:::
:::{grid-item-card} {octicon}`zap;1em` Cost Optimization
:link: cost_optimization
:link-type: doc
-Spot instances and scale-to-zero behavior for detached workloads.
+Spot capacity and scale-to-zero behavior for detached jobs.
:::
:::{grid-item-card} {octicon}`bug;1em` Troubleshooting
:link: ../troubleshooting
:link-type: doc
-What to do when a job is stuck in `PENDING` or repeatedly failing.
+What to do when a job stays in `PENDING` or fails repeatedly.
:::
::::
diff --git a/docs/guides/batched_jobs.md b/docs/guides/batched_jobs.md
index a0506c30..6bd1d505 100644
--- a/docs/guides/batched_jobs.md
+++ b/docs/guides/batched_jobs.md
@@ -1,78 +1,88 @@
# Batched Jobs
-`run_async()` is the tool for a single long-running job. When you
-need to run the **same function over many inputs**, such as a hyperparameter
-sweep, one job per dataset shard, an evaluation grid — wiring that up
-by hand means a loop that calls `run_async()`, your own bookkeeping for
-which handles are still live, your own error aggregation, your own
-cleanup. `run_async_map()` is that loop, done for you.
-
-You call `run_async_map()` on a `@kinetic.run()`-decorated function with a list of
-inputs. It returns a single `BatchHandle` that represents the whole
-collection: one place to observe progress, collect results in input
-order, handle failures, cancel siblings, and tear everything down. The
-underlying jobs are independent Kinetic jobs — each one gets a real
-`JobHandle`, runs on its own pod, and writes its own artifacts to GCS.
-
-This page builds on the single-job workflow covered in
-[Detached Jobs](async_jobs.md). Familiarity with `JobHandle` and the
-`PENDING`/`RUNNING`/`SUCCEEDED`/`FAILED`/`NOT_FOUND` lifecycle is
-assumed.
-
-## A first fan-out
-
-Pass a `@kinetic.run()`-decorated function and a list of inputs to
-`run_async_map()`. It returns a `BatchHandle` immediately while jobs are
-submitted in the background.
+This page explains `run_async_map()`, the call that runs one decorated
+function over many inputs as independent jobs. Use it for a
+hyperparameter sweep, for one job per dataset shard, or for an
+evaluation grid. The call returns one `BatchHandle` for the whole
+batch. With that handle you watch progress, collect the results in input
+order, handle failures, cancel jobs, and delete the resources.
+
+## Before you start
+
+- Read [Detached Jobs](async_jobs.md). Each job in a batch is a normal
+ detached job with its own `JobHandle`, its own pod, and its own
+ artifacts in the jobs bucket.
+- Know the job statuses: `PENDING`, `RUNNING`, `SUCCEEDED`, `FAILED`,
+ and `NOT_FOUND`. A batch reports the same statuses per job.
+
+## A first batch
+
+Call `run_async_map()` on a `@kinetic.run()`-decorated function with a
+list of inputs. Kinetic submits one job per input and returns a
+`BatchHandle`.
```python
import kinetic
-@kinetic.run(accelerator="tpu-v5e-1")
+@kinetic.run(accelerator="tpu-v5litepod-4")
def train(lr):
import keras
+ import numpy as np
+ x = np.random.rand(1000, 20).astype("float32")
+ y = x.sum(axis=1, keepdims=True)
model = keras.Sequential(
[keras.layers.Dense(64, activation="relu"), keras.layers.Dense(1)]
)
model.compile(optimizer=keras.optimizers.Adam(learning_rate=lr), loss="mse")
- history = model.fit(x_train, y_train, epochs=10, verbose=0)
+ history = model.fit(x, y, epochs=10, verbose=0)
return history.history["loss"][-1]
batch = train.run_async_map([0.001, 0.01, 0.1])
losses = batch.results()
-print(losses) # [0.32, 0.28, 0.41] — one result per input, in order
+print(losses) # [0.32, 0.28, 0.41] — one result per input, in input order
```
:::{note}
-You must use `run_async_map()` to fan out. Calling the decorated function
-directly will block until the job finishes and return the result directly,
-so it cannot be used for concurrent execution of multiple inputs.
+A blocking call to the decorated function blocks until that one job ends.
+To run many inputs at the same time, use `run_async_map()`.
:::
+`run_async_map()` accepts these keyword arguments:
+
+| Argument | Default | Meaning |
+| -------- | ------- | ------- |
+| `input_mode` | `"auto"` | How Kinetic passes each item to the function. See [Input modes](#input-modes). |
+| `max_concurrent` | `64` | The maximum number of jobs that are active at one time. `None` removes the limit. |
+| `retries` | `0` | The number of additional attempts for an input after a job failure. |
+| `fail_fast` | `False` | Stop the submission of new jobs after the first failure. |
+| `cancel_running_on_fail` | `False` | With `fail_fast=True`, also cancel the running jobs after the first failure. |
+| `name`, `tags` | `None` | A name and key-value metadata that Kinetic stores in the batch manifest. |
+| `project`, `cluster` | `None` | One-off overrides. Leave them unset; the active profile supplies them. |
+
## Input modes
-The `input_mode` parameter controls how each item in `inputs` is passed
-to the function.
+The `input_mode` argument controls how Kinetic passes each item in
+`inputs` to the function.
-| `input_mode` | Item type | How it's called | Example item |
+| `input_mode` | Item type | Call | Example item |
| ------------------ | --------------------------------- | --------------- | ----------------------------- |
| `"auto"` (default) | `dict` with valid identifier keys | `fn(**item)` | `{"lr": 0.01, "wd": 1e-4}` |
| `"auto"` (default) | `list` or `tuple` | `fn(*item)` | `[0.01, 32]` |
-| `"auto"` (default) | anything else | `fn(item)` | `0.01` |
-| `"single"` | any | `fn(item)` | always passed as a single arg |
+| `"auto"` (default) | any other type | `fn(item)` | `0.01` |
+| `"single"` | any | `fn(item)` | Always one positional argument |
| `"args"` | `list` or `tuple` (required) | `fn(*item)` | `[0.01, 32]` |
| `"kwargs"` | `dict` (required) | `fn(**item)` | `{"lr": 0.01}` |
-### Dict inputs (kwargs unpacking)
+### Dict inputs
-When using `"auto"` mode, dicts with valid Python identifier keys are
-unpacked as keyword arguments:
+In `"auto"` mode, Kinetic unpacks a dict with valid Python identifier
+keys as keyword arguments:
```python
-@kinetic.run(accelerator="tpu-v5e-1")
+@kinetic.run(accelerator="tpu-v5litepod-4")
def train(lr, batch_size): ...
@@ -83,10 +93,10 @@ configs = [
batch = train.run_async_map(configs)
```
-### Preventing unpacking
+### Prevent unpacking
-If your function takes a list or dict as a single argument, use
-`input_mode="single"` to prevent automatic unpacking:
+If your function takes a list or a dict as one argument, pass
+`input_mode="single"`:
```python
@kinetic.run(accelerator="cpu")
@@ -98,16 +108,16 @@ batch = process.run_async_map([[1, 2, 3], [4, 5, 6]], input_mode="single")
```
:::{note}
-In `"auto"` mode, dicts with non-identifier keys (like
-`{"not-an-id": 1}`) or Python keywords (like `{"class": 1}`) are passed
-as a single positional argument rather than unpacked. Use
-`input_mode="kwargs"` or `input_mode="single"` if you need explicit
-control.
+In `"auto"` mode, Kinetic does not unpack every dict. A dict with a key
+that is not a valid identifier, for example `{"not-an-id": 1}`, becomes
+one positional argument. The same applies to a dict with a key that is
+a Python keyword, for example `{"class": 1}`. Use `input_mode="kwargs"`
+or `input_mode="single"` if you need explicit control.
:::
-## Monitoring a batch
+## Monitor a batch
-You can inspect progress at any time through the `BatchHandle`.
+You can inspect the batch at any time through the `BatchHandle`.
```python
# Per-job status
@@ -117,61 +127,71 @@ for idx, status in batch.statuses():
# Aggregate counts
print(batch.status_counts())
# {'RUNNING': 2, 'SUCCEEDED': 1}
+
+# Block until every job is terminal (optional timeout in seconds)
+batch.wait(timeout=1800)
```
`statuses()` returns `(index, JobStatus)` pairs for each submitted job.
-Slots that haven't been submitted yet (when using bounded concurrency)
-are skipped. Job statuses follow the same lifecycle as single jobs —
-see [Detached Jobs](async_jobs.md) for details on `PENDING`, `RUNNING`,
-`SUCCEEDED`, `FAILED`, and `NOT_FOUND`.
+Kinetic skips a slot that is not submitted yet, for example under a
+concurrency limit. `wait()` blocks until the submission ends and every
+submitted job is terminal. `wait()` raises `TimeoutError` if the
+timeout expires.
-## Collecting results
+## Collect results
### `results()`
-The simplest way to collect all results. By default it blocks until
-every job finishes and returns results in input order.
+`results()` is the simplest way to collect every result. It blocks
+until every job ends and returns the results in input order.
```python
-# Input order (default)
losses = batch.results()
-# losses[0] corresponds to inputs[0], losses[1] to inputs[1], etc.
-```
-
-For faster access to early finishers, use `ordered=False` to collect in
-completion order:
-
-```python
-losses = batch.results(ordered=False)
-# Results appear in the order jobs finish, not input order
+# losses[0] belongs to inputs[0], losses[1] to inputs[1], and so on
```
**Parameters:**
-- **`timeout`** (`float | None`, default `None`): Maximum seconds to
- wait. Raises `TimeoutError` if exceeded.
-- **`ordered`** (`bool`, default `True`): `True` returns results aligned
- with `inputs`. `False` returns results in the order jobs complete.
-- **`cleanup`** (`bool`, default `True`): Delete each child's Kubernetes
- resources and GCS artifacts after downloading its result. The group
- manifest is preserved so `attach_batch()` still works.
+- **`timeout`** (`float | None`, default `None`): The maximum number of
+ seconds to wait. `results()` raises `TimeoutError` when the timeout
+ expires.
+- **`ordered`** (`bool`, default `True`): `True` returns the results
+ aligned with `inputs`. `False` returns the results in the order in
+ which the jobs ended.
+- **`cleanup`** (`bool`, default `True`): Delete the resources of each
+ child after Kinetic downloads its result. See the caution below.
- **`return_exceptions`** (`bool`, default `False`): When `True`, failed
positions contain the exception object instead of raising
`BatchError`. When `False`, any failure raises `BatchError`. A job
that fails and an input that fails at submission time both count as a
failure.
+:::{caution}
+With the default `cleanup=True`, `results()` deletes the Kubernetes Job
+of every child, and also the Cloud Storage artifacts of every child
+that succeeded. Those artifacts include the child's `handle.json`, so
+`attach_batch()` cannot load those children later and can block. If you
+want to reattach to the batch later, or to read the logs of a failed
+child, call `results(cleanup=False)`. Call `batch.cleanup()` when you
+no longer need the batch. See [Clean up](#clean-up).
+:::
+
:::{important}
-A `TimeoutError` does not cancel running jobs. They continue executing
-on the cluster. Call `batch.cancel()` explicitly if you want to stop
-them after a timeout.
+A `TimeoutError` does not cancel the jobs. The jobs continue to run on
+the cluster. Call `batch.cancel()` if you want to stop them after a
+timeout, and read [Manual cancellation](#manual-cancellation) first.
:::
+`ordered=False` does not give you earlier access to a result.
+`results(ordered=False)` also returns only after every job is terminal.
+It changes only the order of the list and the moment at which Kinetic
+cleans up each child. To process results as jobs end, use
+`as_completed()`.
+
### `as_completed()`
-For processing results incrementally as jobs finish, use the
-`as_completed()` iterator. It yields `JobHandle` objects in completion
-order.
+`as_completed()` yields each `JobHandle` as its job reaches a terminal
+state, in completion order.
```python
for job in batch.as_completed():
@@ -179,18 +199,21 @@ for job in batch.as_completed():
print(f"{job.job_id} finished: {result}")
```
-`as_completed()` streams results even while submission is still in
-progress. With bounded concurrency, you can start processing the first
-results before the last inputs have been submitted.
+`as_completed()` yields jobs while the submission of other inputs is
+still in progress. Under a concurrency limit, you can process the first
+results before Kinetic submits the last inputs. Each `job.result()` call
+in the loop cleans up that child by default; pass `cleanup=False` to
+keep its resources.
**Parameters:**
-- **`poll_interval`** (`float`, default `5.0`): Seconds between status
- polls.
-- **`timeout`** (`float | None`, default `None`): Maximum seconds to
- wait. Raises `TimeoutError` if exceeded.
+- **`poll_interval`** (`float`, default `5.0`): The number of seconds
+ between status polls.
+- **`timeout`** (`float | None`, default `None`): The maximum number of
+ seconds to wait. `as_completed()` raises `TimeoutError` when the
+ timeout expires.
-## Handling failures
+## Handle failures
When any job fails and `return_exceptions=False` (the default),
`results()` raises a `BatchError`. An input that fails at submission
@@ -198,39 +221,30 @@ time raises a `BatchError` too.
```python
try:
- results = batch.results()
+ results = batch.results(cleanup=False)
except kinetic.BatchError as e:
- print(e) # Batch grp-a1b2c3d4: 2 of 8 jobs failed
+ print(e) # "Batch grp-a1b2c3d4: 2 of 8 jobs failed"
for job in e.failures:
- print(f" job {job.job_id}: {job.status().value}")
- for index, exc in e.submission_failures.items():
- print(f" input {index} never started: {exc}")
- # e.partial_results has results at successful positions, None at failed ones
+ print(f"{job.job_id}: {job.status().value}")
+ print(job.tail(n=20))
+ for idx, exc in e.submission_failures.items():
+ print(f"Input {idx} failed at submission: {exc}")
+ # e.partial_results holds the result at each successful position
+ # and None at each failed position.
```
`BatchError` provides four attributes:
-- **`group_id`**: The batch identifier.
-- **`failures`**: A list of `JobHandle` objects for the jobs that
- started and then failed. The list holds only `JobHandle` objects, so
- `job.job_id` and `job.status()` are always safe to call.
-- **`submission_failures`**: A dict that maps an input index to the
- exception from the submission of that input. These inputs never
- became jobs. They have no `JobHandle`, and they never appear in
- `failures`.
-- **`partial_results`**: A list aligned with `inputs`. A successful
- position holds the result. A failed position holds `None`.
+* **`group_id`**: The batch identifier.
+* **`failures`**: A list of `JobHandle` objects for the jobs that started and then failed. The list holds only `JobHandle` objects, so `job.job_id` and `job.status()` are always safe to call.
+* **`submission_failures`**: A dict mapping input indices to the exceptions from submission. These inputs never became jobs, have no `JobHandle`, and never appear in `failures`.
+* **`partial_results`**: With `ordered=True`, a list aligned with `inputs`, where a successful position holds the result and a failed position holds `None`. With `ordered=False`, a shorter list in completion order holding only the successful results.
-:::{note}
-`partial_results` aligns with `inputs` only for the default
-`ordered=True`. With `ordered=False`, it holds the results that
-`results()` collected, in completion order.
-:::
+The example passes `cleanup=False`. With the default `cleanup=True`, `results()` deletes the Kubernetes Job of every child before it raises `BatchError`. After that deletion, `job.status()` returns `NOT_FOUND` and `job.tail()` raises `RuntimeError`, because the pod is gone. Even with `cleanup=False`, Kubernetes deletes a finished Job about 10 minutes after it ends, so read the logs soon after the failure. This retention window applies to single-host jobs on the GKE backend. A multi-host Pathways job has no retention window; its resources stay until a cleanup call deletes them.
-### Tolerating failures
+### Tolerate failures
-Use `return_exceptions=True` to collect results without raising. Failed
-positions contain the exception object.
+Pass `return_exceptions=True` to collect the results without raising an exception. A failed position holds the exception object.
```python
results = batch.results(return_exceptions=True)
@@ -241,25 +255,25 @@ for i, r in enumerate(results):
print(f"Job {i}: {r}")
```
-### Inspecting failures
+### Inspect failed jobs
-`failures()` returns handles for the jobs with status `FAILED`. It
-excludes `NOT_FOUND`, because that status is ambiguous. A job can be
-`NOT_FOUND` because Kinetic cleaned up its Kubernetes resources, and
-not because the job failed. Use `statuses()` for a more exact view.
+`failures()` returns the handles of the jobs with status `FAILED`. It excludes `NOT_FOUND`, because that status is ambiguous. A job can be `NOT_FOUND` because Kinetic cleaned up its Kubernetes resources, not because the job failed. Use `statuses()` for a finer inspection.
-After `results()` runs, `failures()` returns the failures from that
-collection pass, and not the live status of each job. This keeps the
-list correct after `cleanup=True` deletes the Kubernetes resources.
+After `results()` has run, `failures()` returns the failures from that collection pass, and not the live status of each job. This keeps the list correct after `cleanup=True` deletes the Kubernetes resources.
```python
+batch.wait()
for job in batch.failures():
print(f"{job.job_id}: {job.tail(n=20)}")
```
-`failures()` reports only the jobs that started. To see the inputs that
-failed before they became jobs, read `submission_failures`. It maps the
-input index to the exception from that submission.
+`job.tail()` reads the pod log, so call it while the pod exists. Call it after `wait()` and before `results()`, or after `results(cleanup=False)`, within the 10-minute Kubernetes retention window.
+
+### Submission failures
+
+The call that submits an input can raise, for example because of a packaging or validation error. Kinetic then records the exception and leaves `batch.jobs[idx]` as `None`. Kinetic does not retry a submission failure.
+
+`failures()` reports only the jobs that started. To inspect inputs that failed before they became jobs, read `batch.submission_failures`. The `batch.submission_failures` property returns a dict mapping the input index to the exception. `results()` reports these inputs as failures, but they have no `JobHandle` and do not appear in `failures()`. The position in the results holds the exception object only when `return_exceptions=True`. `wait()` logs a warning when a batch has submission failures.
```python
for index, exc in batch.submission_failures.items():
@@ -268,83 +282,65 @@ for index, exc in batch.submission_failures.items():
## Retries
-The `retries` parameter specifies how many additional attempts a job
-gets after failure. The total number of attempts per input is
-`1 + retries`.
+The `retries` argument sets the number of additional attempts that an input gets after a job failure. The total number of attempts per input is `1 + retries`.
```python
batch = train.run_async_map(configs, retries=2)
-# Each job gets up to 3 attempts (1 initial + 2 retries)
+# Each input gets up to 3 attempts (1 initial + 2 retries)
```
-- Retries are triggered when a job reaches `FAILED` or `NOT_FOUND`
- status.
-- Before each retry, Kinetic cleans up the previous attempt's
- Kubernetes resources (GCS artifacts are preserved for debugging).
-- The group manifest tracks the attempt count per job, so
- `attach_batch()` can distinguish retries from initial submissions.
-- Kinetic does not retry a submission error, which is an error that the
- call to the function raises. These errors are usually packaging errors
- or configuration errors, and they fail again.
-- Kinetic does not retry a cancelled job. `cancel()` marks its children,
- so the `NOT_FOUND` status that cancellation causes never starts a new
- attempt.
+* Kinetic starts a retry when a job reaches `FAILED` or `NOT_FOUND`.
+* Before each retry, Kinetic deletes the Kubernetes resources of the previous attempt and keeps its Cloud Storage (GCS) artifacts for debugging.
+* Each attempt is a new job with a new job ID. The group manifest tracks the attempt count per job, which allows `attach_batch()` to distinguish retries from initial submissions.
+* Kinetic does not retry a submission failure, such as an error raised during function packaging or validation, because repeated attempts would fail identically.
+* Kinetic does not retry a cancelled job. `cancel()` marks child jobs so that the resulting `NOT_FOUND` status never triggers a new attempt.
:::{note}
-When `retries > 0`, job submission runs in a background thread. This
-lets Kinetic poll for failures and submit the input again.
+When `retries > 0`, Kinetic runs the submission loop in a background thread, so that it can poll for failures and resubmit.
:::
## Concurrency control
-By default, `run_async_map()` limits the number of concurrently active
-jobs to 64. Use `max_concurrent` to tune this.
+By default, `run_async_map()` limits the number of active jobs to 64.
+Use `max_concurrent` to change the limit.
```python
-# At most 8 jobs running at once
+# At most 8 jobs run at one time
batch = train.run_async_map(configs, max_concurrent=8)
```
```python
-# Submit all jobs immediately (no concurrency limit)
+# Submit every job at once (no limit)
batch = train.run_async_map(configs, max_concurrent=None)
```
-- **Default:** `64`. Kinetic launches a new job each time a running job
- finishes.
-- **`None`:** Kinetic submits all inputs immediately, with no
- concurrency limit. The calling thread does this work when `retries=0`
- and when `fail_fast` and `cancel_running_on_fail` are not both `True`.
- See [Threading model](#threading-model).
-- Must be a positive integer when set. Passing `0` or a negative value
- raises `ValueError`.
+* **Default `64`:** Kinetic starts a new job each time an active job ends.
+* **`None`:** Kinetic submits all inputs immediately, with no concurrency limit. With `retries=0` (the default) and when `fail_fast` and `cancel_running_on_fail` are not both `True`, the submission runs in the calling thread before `run_async_map()` returns. See [Threading model](https://www.google.com/search?q=%23threading-model).
+* The value must be a positive integer when set. Passing `0` or a negative value raises `ValueError`.
In every case `run_async_map()` returns the `BatchHandle` as soon as the
submission work is handed off or complete. It never waits for the jobs
to finish. Use `wait()` or `results()` when you want to block.
:::{note}
-Kinetic logs a warning when submitting more than 100 jobs with
-`max_concurrent=None`, suggesting you set a limit to control resource
-usage.
+Kinetic logs a warning when you submit more than 100 inputs with
+`max_concurrent=None`. Set a limit to control the resource usage.
:::
-
## Cancellation and fail-fast
### Fail-fast behavior
-The `fail_fast` and `cancel_running_on_fail` parameters control what
-happens when a job fails.
+The `fail_fast` and `cancel_running_on_fail` arguments control what happens when a job fails.
-| `fail_fast` | `cancel_running_on_fail` | On first failure... |
-| ----------------- | ------------------------ | -------------------------------------------------------------------------------- |
-| `False` (default) | `False` (default) | All remaining jobs continue. Failures are collected at the end. |
-| `True` | `False` | No new jobs are launched. Already-running jobs continue to completion. |
-| `True` | `True` | No new jobs are launched. All running siblings are cancelled immediately. |
-| `False` | `True` | **No effect.** `cancel_running_on_fail` only takes effect when `fail_fast=True`. |
+| `fail_fast` | `cancel_running_on_fail` | On the first failure |
+| --- | --- | --- |
+| `False` (default) | `False` (default) | All remaining jobs continue. Kinetic reports the failures at the end. |
+| `True` | `False` | Kinetic starts no new jobs. Jobs that already run continue to the end. |
+| `True` | `True` | Kinetic starts no new jobs and cancels all running jobs at once. |
+| `False` | `True` | **No effect.** `cancel_running_on_fail` applies only when `fail_fast=True`. |
```python
-# Stop the batch as soon as any job fails, cancel all running siblings
+# Stop the batch as soon as any job fails, and cancel all running jobs
batch = train.run_async_map(
configs,
fail_fast=True,
@@ -352,99 +348,103 @@ batch = train.run_async_map(
)
```
-A failure is one of two events. The first is a submission error, when
-the call raises. The second is a runtime failure, when the remote job
-reaches `FAILED` or `NOT_FOUND` status after all of its attempts.
+A "failure" here is either a submission failure (when the call raises) or a runtime failure: a remote job reaching `FAILED` or `NOT_FOUND` status after all of its attempts.
+
+:::{note}
+`run_async_map(max_concurrent=None, retries=0, fail_fast=True)` does not return at once. In that configuration the submission loop runs in the calling thread, and `fail_fast` makes the loop poll until every job is terminal. If you want the call to return at once, set a concurrency limit or leave `fail_fast=False`.
+:::
### Manual cancellation
-`cancel()` stops the full collection at any time. It is independent of
-the `fail_fast` setting.
+`batch.cancel()` cancels every submitted job that is not terminal. It stops the full collection at any time and is independent of the `fail_fast` setting.
```python
batch.cancel()
```
-`cancel()` does three things:
+`batch.cancel()` performs three actions:
-- It deletes the Kubernetes resource of each job that is not terminal.
- The GCS artifacts of that job stay in place for debugging.
-- It drops the inputs that `max_concurrent` holds in the queue. Kinetic
- does not launch them.
-- It marks the children as cancelled. Kinetic does not submit them
- again, even when `retries` is above zero.
+* It deletes the Kubernetes resource of each job that is not terminal, while preserving Cloud Storage (GCS) artifacts for debugging.
+* It drops remaining inputs held in the `max_concurrent` queue so Kinetic does not launch them.
+* It marks child jobs as cancelled so Kinetic does not retry them, even when `retries > 0`.
-A cancelled job reports the status `NOT_FOUND`, because its Kubernetes
-resource is gone. `wait()` returns after each job that started is
-terminal, and the slot of an input that never launched stays `None`.
+A cancelled job reports the status `NOT_FOUND` because its Kubernetes resource has been deleted. `wait()` returns after every job that started becomes terminal, and the slot of an input that never launched remains `None`.
-A cancelled job has no result. `results()` therefore raises a
-`BatchError` that lists those jobs in `failures`. Use
-`results(return_exceptions=True)` to read the results of the children
-that finished before the cancellation.
+Because a cancelled job produces no result, `results()` raises a `BatchError` that lists those jobs in `failures`. Pass `return_exceptions=True` to `results()` to collect results from any jobs that finished before cancellation.
-## Reattaching to a batch
+:::{note}
+The subsequent trailing bullet points describing workaround configurations like `max_concurrent=None` or `fail_fast=True` are obsolete and should be removed, as `cancel()` now natively stops queued and retried jobs.
+:::
-If your local process exits or you want to check on a batch from a
+## Reattach to a batch
+
+If your local process exits, or if you want to check a batch from a
different machine, save the `group_id` and reattach later.
```python
# Original session
batch = train.run_async_map(configs)
print(f"Batch ID: {batch.group_id}") # e.g., "grp-a1b2c3d4"
+results = batch.results(cleanup=False) # keep the child handles
-# Later, from any machine with access to the same GCP project
-batch = kinetic.attach_batch("grp-a1b2c3d4")
-results = batch.results()
+# Later, from any machine with the same active profile
+batch = kinetic.attach_batch("grp-a1b2c3d4", poll_timeout=60)
+results = batch.results(cleanup=False)
+batch.cleanup() # when you are done
```
-`attach_batch()` downloads the group manifest from GCS and rebuilds a
-`JobHandle` for each child. It keeps the index alignment. If the
-original batch had 10 inputs, and a crash stopped it after 7, the
-`batch.jobs` list still has 10 entries. The 3 empty slots hold `None`.
-
-If the manifest names fewer children than the batch expects, the
-original `map()` is still at work. The handle then polls the manifest in
-a background thread until the rest of the children appear, or until
-`poll_timeout` ends the poll.
-
-:::{note}
-Kinetic writes a warning when the manifest of a reattached batch names
-fewer children than expected. This shows a partial submission.
+`attach_batch()` downloads the batch manifest from Cloud Storage and
+rebuilds a `JobHandle` for each child. Kinetic keeps the index
+alignment. If the original batch had 10 inputs and the process crashed
+after 7 submissions, `batch.jobs` has 10 entries. The 3 slots without a
+submission hold `None`.
+
+When the manifest has fewer children than expected, `attach_batch()`
+logs a warning and starts a background thread. That thread polls the
+manifest until all children appear or until `poll_timeout` expires.
+`wait()`, `results()`, and `as_completed()` block until that thread
+ends. After the timeout, the handle reports the submission as complete,
+and the empty slots stay `None`.
+
+:::{caution}
+If the original process crashed during submission, missing children
+never appear. With `poll_timeout=None`, `wait()`, `results()`, or
+`as_completed()` will block indefinitely. Always pass an explicit
+`poll_timeout` when you reattach. In the original session, use
+`results(cleanup=False)` if you plan to reattach later.
:::
**Parameters:**
-- **`group_id`** (`str`): The batch identifier (e.g., `"grp-a1b2c3d4"`).
-- **`project`** (`str | None`, default `None`): GCP project. Uses the
- default when `None`.
-- **`cluster`** (`str | None`, default `None`): GKE cluster name. Uses
- the default when `None`.
-- **`poll_interval`** (`float`, default `10.0`): Seconds between
- manifest polls while the batch is partially submitted.
-- **`poll_timeout`** (`float | None`, default `1800.0`): Maximum seconds
- to poll for the remaining children. After the timeout, the handle
- reports the submission as complete, and the empty slots stay `None`.
- Reattach again to pick up the children that started since then.
- `None` polls forever. Use `None` only when you are sure that the
- original process is alive, because a dead submitter then blocks
- `wait()` and `results()` forever.
+* **`group_id`** (`str`): The batch identifier (for example
+`"grp-a1b2c3d4"`).
+* **`project`** (`str | None`, default `None`): A one-off override. The
+active profile supplies the project when `None`.
+* **`cluster`** (`str | None`, default `None`): A one-off override. The
+active profile supplies the cluster when `None`.
+* **`poll_interval`** (`float`, default `10.0`): The number of seconds
+between manifest polls when children are missing.
+* **`poll_timeout`** (`float | None`, default `1800.0`): The maximum
+number of seconds to poll for missing children. After the timeout, the
+handle reports the submission as complete, and the empty slots stay
+`None`. Reattach again to pick up children that started since then.
+`None` polls forever. Use `None` only when you are sure that the
+original process is still running.
### Children that Kinetic cleaned up
-`results(cleanup=True)` deletes the GCS artifacts of each child that
-gives a result, and the `handle.json` file of the child is one of those
-artifacts. The group manifest stays in place, so `attach_batch()` still
-finds the batch. But it cannot rebuild a `JobHandle` for a child that it
-cleaned up.
+`results(cleanup=True)` deletes the Cloud Storage artifacts of each child
+that yields a result, including the child's `handle.json` file. The group
+manifest stays in place, so `attach_batch()` still finds the batch, but it
+cannot rebuild a `JobHandle` for a cleaned-up child.
-Kinetic treats such a child as terminal, and not as a child that is
-still on the way. The batch reports the submission as complete, and
-`wait()` and `results()` return immediately. The slot of that child
-stays `None`, and `results()` gives `None` at that position.
+Kinetic treats such a child as terminal rather than in-flight. The batch
+reports the submission as complete, and `wait()` and `results()` return
+immediately. The slot for that child stays `None`, and `results()` returns
+`None` at that position.
-`unavailable_children` shows which children are in this state. It maps
-the child index to the job ID from the manifest.
+The `unavailable_children` property shows which children are in this state.
+It maps the child index to the job ID from the manifest:
```python
batch = kinetic.attach_batch("grp-a1b2c3d4")
@@ -452,21 +452,22 @@ print(batch.unavailable_children)
# {0: 'job-1a2b3c4d', 1: 'job-5e6f7a8b'}
```
-A `None` slot that `unavailable_children` does not name is an input that
-the original `map()` never submitted.
+A `None` slot that does not appear in `unavailable_children` represents an
+input that the original `map()` never submitted.
-## Cleanup
+## Clean up
-There are two ways to clean up resources after a batch completes.
+There are two ways to delete the resources of a batch.
-### Automatic cleanup via `results()`
+### Automatic cleanup through `results()`
-By default, `results(cleanup=True)` deletes each child's Kubernetes
-resources and GCS artifacts after downloading its result. The group
-manifest is preserved, so `attach_batch()` still works.
+By default, `results()` cleans up each child after it downloads the
+result of that child. For every child, Kinetic deletes the Kubernetes
+Job. For a child that succeeded, Kinetic also deletes the Cloud Storage
+artifacts, including `handle.json`. A child that failed keeps its
+artifacts. Kinetic keeps the batch manifest.
```python
-# Each child is cleaned up as its result is downloaded
results = batch.results() # cleanup=True is the default
```
@@ -479,8 +480,8 @@ want to reattach later and read the results again. See
### Full teardown
-To delete everything — all children's resources and the group manifest
-itself — call `cleanup()` on the handle:
+To delete everything, including the batch manifest, call `cleanup()` on
+the handle:
```python
batch.cleanup(k8s=True, gcs=True)
@@ -488,53 +489,40 @@ batch.cleanup(k8s=True, gcs=True)
**Parameters:**
-- **`k8s`** (`bool`, default `True`): Delete Kubernetes resources
- (Jobs/pods) for each child.
-- **`gcs`** (`bool`, default `True`): Delete GCS artifacts for each
- child **and** the group manifest.
+- **`k8s`** (`bool`, default `True`): Delete the Kubernetes resources
+ (Jobs and pods) of each child.
+- **`gcs`** (`bool`, default `True`): Delete the Cloud Storage
+ artifacts of each child **and** the batch manifest.
:::{important}
-After calling `cleanup(gcs=True)`, the batch can no longer be reattached
-via `attach_batch()` because the manifest has been deleted.
+After `cleanup(gcs=True)`, `attach_batch()` cannot find the batch,
+because the manifest no longer exists.
:::
## How it works
### Threading model
-`run_async_map()` uses a non-daemon background thread when the
-submission loop must watch the jobs after it launches them. Three
-settings need this:
+`run_async_map()` uses a non-daemon background thread when the submission loop must monitor jobs after launching them. Three settings require this:
-- `max_concurrent` is set. The default is 64. The loop must wait for a
- free slot before it launches the next input.
-- `retries` is above zero. The loop must see a failure before it can
- submit that input again.
-- `fail_fast` and `cancel_running_on_fail` are both `True`. The loop
- must see the first failure before it can cancel the siblings.
+* **`max_concurrent` is set** (the default is 64): The loop must wait for a free slot before launching the next input.
+* **`retries > 0`**: The loop must detect a failure before resubmitting that input.
+* **`fail_fast` and `cancel_running_on_fail` are both `True**`: The loop must detect the first failure to cancel sibling jobs.
-In these cases the thread polls the active jobs, launches new jobs, and
-cancels jobs. `run_async_map()` returns the `BatchHandle` immediately.
+In these cases, the thread polls active jobs, starts new jobs, and cancels running jobs as needed. `run_async_map()` returns the `BatchHandle` immediately. Because the background thread is not a daemon thread, the Python process stays alive until the submission ends.
-In all other cases the calling thread submits every input, and then
-`run_async_map()` returns. Kinetic starts no background thread, and the
-loop does not poll the jobs. A terminal status cannot change what the
-loop does next, so the loop stops as soon as the last input is
-submitted.
+In all other cases, Kinetic submits every input in the calling thread and starts no background thread. The loop does not poll jobs, and `run_async_map()` returns as soon as the last input is submitted.
-`fail_fast` on its own is such a case. A submission error still stops
-the queue immediately, because the loop sees it inside the same
-submission pass. But after every input is launched, a runtime failure
-has nothing left for the loop to stop.
+`fail_fast` on its own falls into this category: a submission error stops the queue immediately because the loop catches it during submission, but once every input is launched, runtime failures have no remaining effect on the submission loop.
### Manifest
-A JSON manifest is written to GCS before the first job is submitted. It
-records the batch metadata (group ID, expected total, function name,
-tags) and is updated after each successful submission with the child's
-job ID and attempt count. This enables crash recovery: `attach_batch()`
-reads the manifest to determine which jobs were submitted and
-reconstructs the handle.
+Kinetic writes a JSON manifest to `gs://{jobs bucket}/_groups/{group_id}/manifest.json`
+before it submits the first job. The manifest records the batch metadata
+(group ID, expected total, function name, name, and tags). Kinetic
+updates the manifest after each successful submission with the child's
+index and job ID. `attach_batch()` reads the manifest to find the
+submitted jobs and rebuilds the handle from each child's `handle.json`.
### Group ID
@@ -565,29 +553,27 @@ jobs, so `BatchError.failures` does not list them.
:link: async_jobs
:link-type: doc
-The single-job `run_async()` workflow each child of a batch is built on.
+The `run_async()` workflow and the `JobHandle` that each child of a batch uses.
:::
:::{grid-item-card} {octicon}`zap;1em` Cost Optimization
:link: cost_optimization
:link-type: doc
-Fan-out amplifies both throughput and spend; concurrency limits and spot
-instances matter here.
+A batch multiplies the cost as well as the throughput; concurrency limits and Spot capacity help.
:::
-:::{grid-item-card} {octicon}`history;1em` Checkpointing
+:::{grid-item-card} {octicon}`history;1em` Outputs and Checkpoints
:link: checkpointing
:link-type: doc
-Each child writes to its own `KINETIC_OUTPUT_DIR`; useful for long
-per-job work inside a batch.
+Each child writes to its own `KINETIC_OUTPUT_DIR`; useful for long work inside a batch.
:::
:::{grid-item-card} {octicon}`bug;1em` Troubleshooting
:link: ../troubleshooting
:link-type: doc
-What to do when children stick in `PENDING` or repeatedly fail.
+What to do when children stay in `PENDING` or fail repeatedly.
:::
::::
diff --git a/docs/guides/checkpointing.md b/docs/guides/checkpointing.md
index b8852ef7..b7ae2a66 100644
--- a/docs/guides/checkpointing.md
+++ b/docs/guides/checkpointing.md
@@ -1,19 +1,18 @@
-# Checkpointing and Outputs
+# Outputs and Checkpoints
-Long jobs need somewhere durable to write to. Pods come and go — when
-your training script exits, anything that lived only inside the pod's
-filesystem is gone. Kinetic gives you `KINETIC_OUTPUT_DIR`: a per-job
-GCS prefix that survives the pod, so your checkpoints, logs, and final
-artifacts are still there when you come back.
+The filesystem of the pod is temporary. When your function returns, or
+when the pod stops for another reason, every file on the pod is lost.
+Kinetic gives each job an **output directory**: a Cloud Storage prefix
+that stays after the pod stops. The pod sees that prefix as the
+environment variable `KINETIC_OUTPUT_DIR`. This page explains what to
+write where and how to set the output directory. It also explains how to
+resume a training run from a checkpoint and how long the files stay.
-This page covers what to write where, how Orbax (or any other library)
-plugs into it, and how cleanup and TTLs work.
+## A first job that writes outputs
-## A first checkpointed job
-
-Inside the pod, `KINETIC_OUTPUT_DIR` is already set. Read it and write
-under it. Fall back to a local path when the variable is not present so
-that the same function works when you exercise it locally:
+Kinetic sets `KINETIC_OUTPUT_DIR` in the pod. Read the variable and write
+under that path. If the variable is not present, use a local path. Then
+the same function also works when you call it locally for a test:
```python
import os
@@ -23,163 +22,258 @@ import kinetic
@kinetic.run(accelerator="cpu")
def train():
- # Remote: KINETIC_OUTPUT_DIR resolves to gs://.../outputs/.
- # Local: fall back to a filesystem path under /tmp so the same code
- # works when you run the function directly for testing.
+ # Remote: KINETIC_OUTPUT_DIR is gs://.../outputs/.
+ # Local: fall back to a path under /tmp for a direct test call.
output_dir = os.environ.get("KINETIC_OUTPUT_DIR", "/tmp/local_checkpoints")
- # ... train and write checkpoints/artifacts under output_dir ...
+ # ... train, and write checkpoints and artifacts under output_dir ...
return f"saved to {output_dir}"
```
-:::{seealso}
-For full Orbax-managed auto-resume with JAX or Keras, the canonical
-runnable examples live in the repo:
+Every call of `train()` is a new job with a new job ID. By default, the
+output directory contains the job ID, so every call also gets a new,
+empty output directory. That default is correct for a job that runs one
+time. If a second call must find the files of the first call, both calls
+must use the same output directory. A resume from a checkpoint is one
+example. Pass `output_dir=` to the decorator:
-- [`examples/example_checkpoint.py`](https://github.com/keras-team/kinetic/blob/main/examples/example_checkpoint.py)
- — JAX + Orbax with auto-resume.
-- [`examples/example_keras_checkpoint.py`](https://github.com/keras-team/kinetic/blob/main/examples/example_keras_checkpoint.py)
- — same pattern using `model.get_weights()` / `set_weights()`.
-:::
+```python
+@kinetic.run(accelerator="cpu", output_dir="gs://my-bucket/runs/exp-01")
+def train(): ...
+```
+
+See [Resume a job from a checkpoint](#resume-a-job-from-a-checkpoint)
+below.
+
+## Three kinds of artifact
-## Outputs and checkpoints
+A Kinetic job produces three kinds of artifact. Each kind has its own
+location and its own lifecycle:
-A Kinetic job produces three distinct kinds of artifact, each with its
-own storage location and lifecycle:
+| Artifact | What it is | Where it goes |
+| -------- | ---------- | ------------- |
+| Job return value | The Python value that your function returns | `gs://{jobs bucket}/{job_id}/result.pkl`, then your local process |
+| Durable outputs | Files that you write during the run | `KINETIC_OUTPUT_DIR` in Cloud Storage |
+| Resumable checkpoints | Periodic snapshots of the training state | A fixed subdirectory under `KINETIC_OUTPUT_DIR` |
-Artifact | What it is | Where it lives
---------------------- | -------------------------------------- | -------------------------------------------------------------------------------
-Job return value | The Python value your function returns | Persisted to `gs://{bucket}/{job_id}/result.pkl`, then downloaded to your local process
-Durable outputs | Files you wrote during the run | `KINETIC_OUTPUT_DIR` (GCS)
-Resumable checkpoints | Periodic state snapshots for restart | `KINETIC_OUTPUT_DIR/` (GCS)
+Use the return value for **small** results: a final loss, a dict of
+metrics, a path string. Write large files under the output directory.
+Write checkpoints to a fixed subdirectory under the output directory, so
+that a later run finds them at a known path.
-The return value is the right channel for **small** results: a final
-loss, a metric dict, a path string. Large files belong on the output
-dir; checkpoints belong on a stable subpath under the output dir so
-restarts can find them.
+## The default output directory
-`KINETIC_OUTPUT_DIR` is set automatically when the job starts. By
-default it resolves to the jobs bucket for your cluster:
+Kinetic sets `KINETIC_OUTPUT_DIR` when the job starts. By default, the
+variable points to a per-job prefix in the jobs bucket of your cluster:
```text
gs://{project}-kn-{cluster}-jobs/outputs/{job_id}
```
-`{project}` is your GCP project (from `KINETIC_PROJECT`) and `{cluster}`
-is the Kinetic cluster name (from `KINETIC_CLUSTER`, defaulting to
-`kinetic-cluster`). The bucket is created by `kinetic up` and reused
-across all jobs submitted to that cluster.
+`{project}` and `{cluster}` come from the active profile, unless you
+override them for the job. `{job_id}` is the ID of the job, for example
+`job-3f9a1c2b`. `kinetic up` creates the jobs bucket one time, and every
+job on that cluster uses the same bucket.
+
+The pod runs as the node service account of the cluster,
+`kn-{cluster}-nodes@{project}.iam.gserviceaccount.com`. That account can
+read and write the jobs bucket. It has no access to other buckets unless
+you grant that access.
+
+## Set the output directory
+
+Pass `output_dir="gs://..."` to `@kinetic.run()` to replace the default
+location for a job.
+
+For a script that you cannot edit, export `KINETIC_OUTPUT_DIR` in your
+local shell before you submit the job. Kinetic reads the local variable
+at submit time and passes the value to the pod. The decorator argument
+wins over the local environment variable, and the local environment
+variable wins over the default. See [Configuration](../configuration.md)
+for the full precedence rules. `kinetic config` shows the value of
+`KINETIC_OUTPUT_DIR` if you set the variable in your shell.
+
+You cannot change the output directory from the `kinetic jobs` commands.
+The output directory is a property of the job that you set at submit
+time.
+
+:::{note}
+If `output_dir=` points to a bucket that `kinetic up` did not create, the
+pod cannot write there until you grant access. Give the node service
+account of the cluster the `roles/storage.objectAdmin` role and the
+`roles/storage.legacyBucketReader` role on that bucket:
+
+```bash
+gcloud storage buckets add-iam-policy-binding gs://my-bucket \
+ --member=serviceAccount:kn-kinetic-cluster-nodes@my-project.iam.gserviceaccount.com \
+ --role=roles/storage.objectAdmin
+gcloud storage buckets add-iam-policy-binding gs://my-bucket \
+ --member=serviceAccount:kn-kinetic-cluster-nodes@my-project.iam.gserviceaccount.com \
+ --role=roles/storage.legacyBucketReader
+```
+
+Replace `kinetic-cluster` and `my-project` with the cluster name and the
+project of your profile. Orbax and TensorStore need the second role to
+read the bucket metadata.
+:::
+
+## Resume a job from a checkpoint
+
+A job can stop before the work is done: Google Cloud preempts a Spot node, a node fails, or your code raises an error. Kinetic does not submit the job again for you. The checkpoints that the job wrote stay in Cloud Storage. To continue the work, submit the function again with the **same** output directory. Your code then finds the latest checkpoint under that directory and continues from that checkpoint.
+
+The output directory must be the same for each call. With the default output directory, each call gets a new job ID and therefore a new, empty prefix. The second call then starts from step 0. There are two ways to set a fixed directory:
+
+1. Pass `output_dir=` to the decorator:
+```python
+@kinetic.run(
+ accelerator="tpu-v5litepod-4",
+ output_dir="gs://my-bucket/runs/exp-01",
+)
+def train(): ...
+
+
+train() # writes checkpoints under gs://my-bucket/runs/exp-01
+train() # finds them and resumes
+```
+
+
+2. Set `KINETIC_OUTPUT_DIR` in your local environment before both submissions:
+```bash
+export KINETIC_OUTPUT_DIR=gs://my-bucket/runs/exp-01
+python train.py # first run
+python train.py # resumes
+
+```
-You can override the default for one job. Pass `output_dir=` to the
-decorator, or set `KINETIC_OUTPUT_DIR` in your local environment before
-you submit the job. Kinetic reads both sources at submit time. The
-output directory therefore does not change while the job runs. The
-precedence table in [Configuration](../configuration.md) shows how these
-sources combine.
+
+
+Kinetic reads both sources at submit time, so the output directory does not change while the job runs. The precedence table in [Configuration](https://www.google.com/search?q=../configuration.md) shows how these sources combine.
+
+A path in the jobs bucket also works, for example `gs://my-project-kn-kinetic-cluster-jobs/outputs/exp-01`, and needs no extra access grant. The 30-day rule of that bucket applies (see below).
## Recommended directory layout
-A simple convention that scales from one job to many:
+The layout below works for a single job and for many jobs:
```text
$KINETIC_OUTPUT_DIR/
├── checkpoints/ # Orbax / model.save_weights — periodic snapshots
-├── logs/ # extra logs your code writes (stdout already streams)
-├── metrics/ # tensorboard / json metric dumps
+├── logs/ # extra logs that your code writes (stdout already streams)
+├── metrics/ # TensorBoard / JSON metric dumps
└── final/ # post-training artifacts: exported model, eval results
```
-Use whichever subdirectories make sense for your workflow. The point is
-that the layout is yours to control — Kinetic only cares that you write
-under the prefix it gave you.
-
-## TTL and retention
-
-By default the GCS bucket Kinetic creates has a **30-day TTL** on its
-contents. Anything written to `KINETIC_OUTPUT_DIR` is auto-deleted
-after 30 days. That's the right default for ephemeral training, but if
-you want a checkpoint to outlive a month:
-
-- Copy it to a bucket with no lifecycle policy (`gsutil cp` or the GCS
- client library).
-- Or set `output_dir=` to a bucket you manage yourself, with whatever
- lifecycle rules you want.
-
-:::{note}
-`JobHandle.cleanup(gcs=True)` removes the per-job artifacts under the
-GCS prefix used for code and result payloads — it does **not** touch
-files you wrote under `KINETIC_OUTPUT_DIR`. Outputs survive cleanup.
+Use the subdirectories that fit your workflow. Kinetic does not read or
+interpret the layout. Kinetic only requires that you write under the
+prefix that it gives you.
+
+## Retention and cleanup
+
+**The 30-day rule.** The jobs bucket has a lifecycle rule that deletes
+every object 30 days after its creation. Cloud Storage therefore deletes
+the files under the default `KINETIC_OUTPUT_DIR` after 30 days. That
+default fits short experiments. If a checkpoint or a model must stay
+longer than 30 days, do one of these two things:
+
+- Copy the files to a bucket without a lifecycle rule, with
+ `gcloud storage cp` or the Cloud Storage client library.
+- Set `output_dir=` to a bucket that you manage, with the lifecycle
+ rules that you choose (see the access note above).
+
+**Job cleanup does not delete outputs.** A blocking call, a
+`JobHandle.result()` call, `JobHandle.cleanup(gcs=True)`, and
+`kinetic jobs cleanup JOB_ID` delete the job artifacts under
+`gs://{jobs bucket}/{job_id}/`. Those artifacts include the serialized
+function, the source archive, and the result. These calls never delete
+files under `KINETIC_OUTPUT_DIR`. `result()` deletes the artifacts only
+after it collects a usable result; a failed job keeps its artifacts. Pass
+`cleanup=False` to `result()`, or `--no-cleanup` to
+`kinetic jobs result`, to keep the artifacts of a job that succeeded.
+
+:::{warning}
+`kinetic down` deletes the cluster and the jobs bucket, with every output
+in that bucket. Before you run `kinetic down`, copy the files that you
+want to keep to a bucket that Kinetic does not manage.
:::
-## Copy-paste checklist
+## Checklist for a long job
-A short checklist for any long-running job that you don't want to redo
-from scratch:
+Follow these steps for a job that you do not want to run again from the
+start:
:::{container} kinetic-steps
-1. **Read `KINETIC_OUTPUT_DIR`** inside the function and write everything
- durable under it.
+1. **Read `KINETIC_OUTPUT_DIR`** inside the function, and write every
+ durable file under that path.
-2. **Write checkpoints to a stable subdirectory** (e.g.
- `$KINETIC_OUTPUT_DIR/checkpoints/`) so the resume path is
- predictable.
+2. **Write checkpoints to a fixed subdirectory**, for example
+ `$KINETIC_OUTPUT_DIR/checkpoints/`, so that you know the resume path.
-3. **Choose a checkpoint cadence** that bounds how much work a restart
- would lose (every N steps, or every M minutes).
+3. **Choose a checkpoint interval** that limits the work that a restart
+ loses: every N steps, or every M minutes.
-4. **Verify resume works locally** before the long run — submit the same
- function twice with the same `output_dir` and confirm the second
- call picks up where the first left off.
+4. **Set a fixed output directory** with `output_dir=` or with a local
+ `KINETIC_OUTPUT_DIR`. Without a fixed directory, the second submission
+ gets a new, empty prefix and cannot resume.
-5. **If the run is critical**, copy the final artifacts to a bucket
- without the 30-day TTL after success.
+5. **Test the resume before the long run.** Submit the same function two
+ times with the same output directory, and confirm that the second
+ call continues from the last checkpoint.
+
+6. **Copy the final artifacts** to a bucket without the 30-day rule if the
+ run is important.
:::
## JAX example
-```{literalinclude} ../../examples/example_checkpoint.py
-```
+The example below shows the write-and-read pattern with Orbax. The
+function points a `CheckpointManager` at `KINETIC_OUTPUT_DIR`. It
+restores the latest step if one exists, and it saves a checkpoint after
+each step.
-After the snippet:
+The example uses `@kinetic.run(accelerator="cpu")` with no `output_dir=`,
+so each call gets a new default output directory and starts from step 0.
+To see a real resume, add `output_dir="gs://..."` to the decorator, or
+export `KINETIC_OUTPUT_DIR` before you run the script two times.
-- The function reads `KINETIC_OUTPUT_DIR` and points Orbax's
- `CheckpointManager` at it.
-- Calling the function a second time picks up from the latest step
- rather than restarting from scratch.
+```{literalinclude} ../../examples/example_checkpoint.py
+```
## Keras example
+The Keras example uses the same pattern. `model.get_weights()` returns a
+list of NumPy arrays, and Orbax saves that list as a PyTree. After a
+restore, `model.set_weights()` loads the arrays back into the model. As
+with the JAX example, add `output_dir=` or export `KINETIC_OUTPUT_DIR`
+to make two calls share one output directory.
+
```{literalinclude} ../../examples/example_keras_checkpoint.py
```
-After the snippet:
-
-- `model.get_weights()` produces a PyTree of NumPy arrays that Orbax
- knows how to save.
-- `model.set_weights()` restores them on resume.
-
## Related pages
::::{grid} 1 1 2 2
:gutter: 3
-:::{grid-item-card} {octicon}`database;1em` Data
+:::{grid-item-card} {octicon}`database;1em` Working with Data
:link: data
:link-type: doc
-Input side of the I/O story.
+The input side: ship local files and read Cloud Storage data from your
+function.
:::
-:::{grid-item-card} {octicon}`clock;1em` Managing Async Jobs
+:::{grid-item-card} {octicon}`clock;1em` Detached Jobs
:link: async_jobs
:link-type: doc
-Long jobs are also the place where you most want detached submission.
+Submit long jobs with `run_async()`, and collect or clean up their
+results later.
:::
:::{grid-item-card} {octicon}`graph;1em` Cost Optimization
:link: cost_optimization
:link-type: doc
-Spot instances make checkpointing essential.
+Spot capacity lowers the cost, and a checkpoint makes a preempted run
+resumable.
:::
::::
diff --git a/docs/guides/clusters.md b/docs/guides/clusters.md
index 74db6dca..51184ad3 100644
--- a/docs/guides/clusters.md
+++ b/docs/guides/clusters.md
@@ -1,102 +1,238 @@
-# Multiple Clusters
+# Clusters and Node Pools
+
+A **cluster** is the GKE cluster, with its buckets and its image
+repository, that runs your jobs. A **node pool** is a group of VMs of one
+accelerator type inside a cluster. This page explains what `kinetic up`
+creates and how to add and remove node pools. It also explains how a
+team shares one cluster, when to run more than one cluster, and how to
+delete a cluster.
+
+## What `kinetic up` creates
+
+`kinetic init` runs `kinetic up` on the **Create** path. You can also run
+`kinetic up` yourself. One run creates one cluster and everything the
+cluster needs:
+
+| Resource | Name | Purpose |
+| -------- | ---- | ------- |
+| GKE cluster | `{cluster}` (default `kinetic-cluster`) | Runs the job pods. |
+| Node pool | `gpu-{type}-{4 hex}` or `tpu-{type}-{4 hex}`, for example `tpu-v5litepod-1a2b` | One accelerator type. `{type}` is the type name without the count, such as `l4` or `v5litepod`. `up` adds one pool of your choice; `kinetic pool add` adds more. |
+| Artifact Registry repository | `kn-{cluster}` | Holds the container images that Kinetic builds. |
+| Jobs bucket | `gs://{project}-kn-{cluster}-jobs` | Job artifacts, results, `KINETIC_OUTPUT_DIR`, and the data cache. Objects expire after 30 days. |
+| Builds bucket | `gs://{project}-kn-{cluster}-builds` | Cloud Build sources. Objects expire after 30 days. |
+| Service accounts | `kn-{cluster}-nodes`, and one for builds | Give the pods and Cloud Build access to the buckets and the repository. |
+| Cloud NAT gateway | `kn-{cluster}-nat` | Gives the private cluster nodes access to the internet. |
+| State bucket | `gs://{project}-kinetic-state` | One per project. Holds the infrastructure state for all clusters in the project. |
+
+`kinetic up` also enables the required Google Cloud APIs, installs the
+LeaderWorkerSet controller and the GPU driver installer on the cluster,
+configures `kubectl`, and saves a profile for the cluster. The profile
+becomes active.
+
+Useful flags:
+
+- `--cluster NAME` and `--zone ZONE` select the cluster name and zone.
+- `--accelerator SPEC` selects the first node pool without a prompt,
+ for example `--accelerator tpu-v5litepod-4`. Use `cpu` for a cluster
+ without an accelerator pool.
+- `--min-nodes N` keeps N nodes of the first pool warm. The default is
+ `0`.
+- `--preview` shows the changes without applying them.
+- `--yes` skips the confirmation prompt.
+
+A second `kinetic up` for the same cluster is safe. It keeps the existing
+node pools and ignores `--accelerator`. Use `kinetic pool add` and
+`kinetic pool remove` to change the pools.
+
+## Check the cluster
-Kinetic supports running multiple independent clusters within the same GCP project. Each cluster gets its own isolated set of cloud resources (GKE cluster, Artifact Registry, storage buckets), backed by a separate infrastructure stack.
+```bash
+kinetic status # cluster, buckets, repository, node pools
+kinetic pool list # node pools only
+kinetic accelerators # every accelerator name that Kinetic knows
+kinetic accelerators --live # marks the accelerators that have a pool
+```
+
+## Node pools
+
+A job runs only on a node pool with the same accelerator type and, for
+TPUs, the same topology. Add one pool for each accelerator that you use:
+
+```bash
+kinetic pool add --accelerator tpu-v5litepod-4
+kinetic pool add --accelerator gpu-l4
+kinetic pool add --accelerator gpu-a100x4 --spot
+kinetic pool add --accelerator tpu-v6e-16 --reservation my-v6e-reservation
+```
-## Why Use Multiple Clusters?
+`kinetic pool add` accepts the same accelerator names as
+`@kinetic.run(accelerator=...)`. See [Accelerators](../accelerators.md).
+The command prints the generated pool name.
-- **Isolation**: Separate your GPU and TPU workloads into different clusters.
-- **Regions**: Run jobs in different GCP regions or zones.
-- **Environment**: Maintain separate clusters for development, testing, and production.
+Options:
-## Creating a Cluster
+- `--min-nodes N` keeps N nodes running at all times. The default, `0`,
+ scales the pool to zero when no job runs. Nodes cost money while they
+ run, even without a job. See [Cost Optimization](cost_optimization.md).
+- `--spot` uses Spot VMs. Spot VMs cost less, but Google Cloud can
+ preempt them with 30 seconds of notice. A job runs on a Spot pool only
+ when the decorator also sets `spot=True`. See
+ [Cost Optimization](cost_optimization.md).
+- `--reservation NAME` binds a capacity reservation to the pool. You
+ cannot combine `--reservation` with `--spot`. See
+ [Capacity Reservations](reservations.md).
+- `--preview` shows the change without applying it.
-Use the `--cluster` flag with `kinetic up` to create a named cluster.
+Remove a pool by name:
```bash
-# Create a GPU cluster in us-east1-b
-kinetic up --cluster=gpu-cluster --zone=us-east1-b --accelerator=a100
+kinetic pool list
+kinetic pool remove tpu-v5litepod-1a2b
```
-If the `--cluster` flag is omitted, Kinetic uses the default name `kinetic-cluster`.
+Two limits apply to every pool that `kinetic pool add` creates:
-## Targeting a Cluster
+- A pool scales up to a fixed maximum: `--min-nodes` plus 10 nodes for a
+ GPU pool, or `--min-nodes` plus the hosts of one slice for a TPU pool.
+ One TPU pool therefore runs one slice at a time.
+- Each node has a maximum run duration of 24 hours, except a node in a
+ Spot TPU pool. GKE recycles the node after 24 hours, and a job on that
+ node fails. Write checkpoints and resume for a longer job.
-You can target a specific cluster from your code using the `cluster` parameter or an environment variable.
+You cannot add a CPU pool with `kinetic pool add`. Every cluster has a
+default pool with one `e2-standard-4` node that runs at all times. That
+node runs the cluster system pods and `accelerator="cpu"` jobs. The
+cluster autoscaler adds CPU nodes when a CPU job needs more.
-### Using the Decorator
+## Share a cluster with a team
-```python
-@kinetic.run(accelerator="a100", cluster="gpu-cluster")
-def train_on_gpu(): ...
-```
+The infrastructure state lives in the state bucket of the project, not
+on one laptop. A teammate who has access to the bucket runs
+`kinetic init`, sees the cluster on the **Join** path, and gets a
+profile for it. No one exports environment variables and no one copies
+files.
-### Using Environment Variables
+IAM requirements:
-Set `KINETIC_CLUSTER` to avoid repeating the cluster name in every decorator.
+- The first person who runs `kinetic up` in a project needs
+ `roles/storage.admin`, because that run creates the state bucket.
+- Every other person needs `roles/storage.objectAdmin` on the state
+ bucket, and the permissions to submit jobs: read and write access to
+ the jobs bucket, and Kubernetes access to the cluster.
+
+See [Security](../security.md) for the trust model.
+
+## Run more than one cluster
+
+Most people need one cluster. Create a second cluster when you have one
+of these reasons:
+
+- **Isolation** — GPU jobs and TPU jobs on separate clusters.
+- **Location** — jobs in two zones or two regions, for example to find
+ Spot capacity.
+- **Environments** — separate clusters for development and production.
+
+:::{warning}
+Each cluster has its own GKE control plane. The control plane costs about
+$0.10 per hour, and the Google Cloud free tier covers one cluster only.
+Each cluster also has its own repository and buckets. Do not create a
+second cluster without one of the reasons above.
+:::
+
+Create a named cluster:
```bash
-export KINETIC_CLUSTER="gpu-cluster"
+kinetic up --cluster gpu-cluster --zone us-east1-b --accelerator gpu-a100
```
-## Managing Clusters
+`kinetic up` saves a profile with the same name as the cluster and makes
+that profile active. Switch between clusters with the profile:
+
+```bash
+kinetic profile ls
+kinetic profile use gpu-cluster
+```
-All CLI commands accept the `--cluster` flag, allowing you to manage each cluster independently.
+For a one-off command against a different cluster, use the `--profile`
+flag, or pass both `--cluster` and `--zone`. Kinetic identifies a cluster by
+its project, its zone, and its name together:
```bash
-# Check status of a specific cluster
-kinetic status --cluster=gpu-cluster
+kinetic --profile gpu-cluster status
+kinetic status --cluster gpu-cluster --zone us-east1-b
+```
-# Add a node pool to a specific cluster
-kinetic pool add --cluster=gpu-cluster --accelerator=h100
+For a one-off job, pass `cluster=` to the decorator:
-# Tear down a specific cluster
-kinetic down --cluster=gpu-cluster
+```python
+@kinetic.run(accelerator="gpu-a100", cluster="gpu-cluster", zone="us-east1-b")
+def train_on_gpu(): ...
```
-## Resource Naming
+Each cluster has its own set of resources, named after the cluster:
+
+| Resource | Name for `gpu-cluster` |
+| -------- | ---------------------- |
+| GKE cluster | `gpu-cluster` |
+| Artifact Registry repository | `kn-gpu-cluster` |
+| Jobs bucket | `{project}-kn-gpu-cluster-jobs` |
+| Builds bucket | `{project}-kn-gpu-cluster-builds` |
+| Infrastructure stack | `{project}-gpu-cluster` |
+
+## Delete a cluster
-Kinetic uses the cluster name to scope its GCP resources. For a cluster named `gpu-cluster`, the resources will follow this pattern:
+```bash
+kinetic down
+```
-| Resource | Name |
-| --- | --- |
-| GKE Cluster | `gpu-cluster` |
-| Artifact Registry | `kn-gpu-cluster` |
-| Storage Bucket | `{project}-kn-gpu-cluster-jobs` |
+`kinetic down` deletes the cluster, the node pools, the Cloud NAT
+gateway, and the Artifact Registry repository with its images. It also
+deletes the jobs bucket and the builds bucket **with their contents**,
+including everything under `KINETIC_OUTPUT_DIR`. Copy the outputs that you want to keep to another
+bucket first. `kinetic down` leaves the Google Cloud APIs enabled, and it
+does not delete the state bucket or your profiles.
-:::{admonition} When not to use this
-:class: warning
+Pass `--cluster NAME --zone ZONE` to delete a cluster other than the
+one in the active profile. Kinetic identifies a cluster by its project,
+its zone, and its name together. Kinetic does not look up the zone from the
+name. The active profile still points at the deleted cluster
+afterwards. Run `kinetic profile use` or `kinetic profile rm` to update
+it.
-Most users only need one cluster.
-Each additional cluster has its own GKE control plane (~$0.10/hr,
-or ~$74/month) and its own Artifact Registry, so don't add a second
-cluster speculatively. Add one when you have a real reason: GPU vs
-TPU isolation, regional separation, or dev vs prod environments.
-:::
+By default, `kinetic down` empties the buckets before it deletes them.
+Run `kinetic up --no-force-destroy` to store the opposite choice in the
+stack. `kinetic down` then fails until you empty the buckets yourself.
## Related pages
::::{grid} 1 1 2 2
:gutter: 3
-:::{grid-item-card} {octicon}`graph;1em` Cost Optimization
-:link: cost_optimization
+:::{grid-item-card} {octicon}`stack;1em` Profiles
+:link: profiles
:link-type: doc
-Control plane costs and how the GKE free tier covers exactly one
-cluster.
+One profile per cluster; switch with one command.
:::
-:::{grid-item-card} {octicon}`checklist;1em` Capacity Reservations
-:link: reservations
+:::{grid-item-card} {octicon}`cpu;1em` Accelerators
+:link: ../accelerators
+:link-type: doc
+
+Every accelerator name and topology.
+:::
+
+:::{grid-item-card} {octicon}`graph;1em` Cost Optimization
+:link: cost_optimization
:link-type: doc
-When reservations make multi-cluster setups worth the overhead.
+Scale to zero, warm nodes, Spot VMs, and the control plane cost.
:::
-:::{grid-item-card} {octicon}`gear;1em` Configuration
-:link: ../configuration
+:::{grid-item-card} {octicon}`checklist;1em` Capacity Reservations
+:link: reservations
:link-type: doc
-`KINETIC_CLUSTER` and the precedence rules.
+Guaranteed hardware for a node pool.
:::
::::
diff --git a/docs/guides/containers.md b/docs/guides/containers.md
index 57635f3e..209e752c 100644
--- a/docs/guides/containers.md
+++ b/docs/guides/containers.md
@@ -1,173 +1,227 @@
# Container Images
-This page is the deep reference for the container image system: the
-three modes side by side, the prebuilt-base workflow, the custom-image
-contract, and the `kinetic build-image` command.
-
-For a higher-level overview and the recommendation matrix on which
-mode to pick, start with [Execution Modes](../guides/execution_modes.md).
-
-Kinetic supports three container image modes that control how your remote execution environment is built and deployed. Choose the mode that best fits your workflow by setting the `container_image` parameter in the `@kinetic.run()` decorator.
-
-:::{note}
-**Expected timing:**
-
-- **Bundled, cold (first run / dep change):** ~2–5 minutes for the
- Cloud Build step.
-- **Bundled, warm (cached image):** under a minute to schedule and
- start the pod.
-- **Prebuilt:** 30–60 seconds for the base image pull (cached after
- first use on a node), plus the time to `uv pip install` your
- `requirements.txt`.
-- **Custom:** a single image pull, then immediate execution. Cold
- pulls vary widely with image size and registry latency.
-:::
-
-:::{warning}
-**When not to use custom image mode:** if you only need different
-Python packages, bundled or prebuilt are simpler and cheaper. Reach
-for custom images when you have non-Python system libraries (CUDA
-builds, C++ deps), corporate compliance requirements, or you genuinely
-want to manage the image lifecycle yourself.
-:::
+Every job runs inside a container image. By default, Kinetic builds that
+image for you and caches it. Most people never change this default. This
+page is for three cases where the default does not fit. The build step
+is too slow for your workflow. You need system libraries that a Python
+install cannot provide. Or your organization requires a vetted base
+image.
-| Mode | `container_image=` | Build step | Dependencies installed |
-| --------------------- | --------------------- | -------------------------------- | ----------------------------------- |
-| **Bundled** (default) | `None` or `"bundled"` | Cloud Build (cached by dep hash) | Baked into the image |
-| **Prebuilt** | `"prebuilt"` | None | At pod startup via `uv pip install` |
-| **Custom** | `""` | None (you manage it) | Whatever is in your image |
+## The three modes
-## Bundled Mode (Default)
+The `container_image` argument of `@kinetic.run()` selects how Kinetic
+produces the image:
-Bundled mode builds a custom container image via Cloud Build with all your dependencies baked in. The image is tagged by a hash of your dependencies, so unchanged dependencies reuse the cached image.
+| Mode | `container_image=` | Build step | Where your dependencies install |
+| ---- | ------------------ | ---------- | ------------------------------- |
+| **Bundled** (default) | `None` or `"bundled"` | Cloud Build, cached by a hash of the inputs | In the image, at build time |
+| **Prebuilt** | `"prebuilt"` | None. Kinetic pulls a base image that you published. | In the pod, at start, with `uv pip install` |
+| **Custom** | An image URI | None. You build and push the image. | In your image. Kinetic installs nothing. |
```python
-import kinetic
+@kinetic.run(accelerator="tpu-v5litepod-4") # bundled (the default)
+def train_bundled(): ...
-# Bundled mode — these are equivalent:
-@kinetic.run(accelerator="tpu-v6e-8")
-def train(): ...
+@kinetic.run(accelerator="tpu-v5litepod-4", container_image="prebuilt")
+def train_prebuilt(): ...
-@kinetic.run(accelerator="tpu-v6e-8", container_image="bundled")
-def train(): ...
+@kinetic.run(
+ accelerator="tpu-v5litepod-4",
+ container_image="us-docker.pkg.dev/me/repo/img:v1", # custom
+)
+def train_custom(): ...
```
-### Tradeoffs
+In all three modes, Kinetic ships your function and your project source
+with the job. The image supplies the installed packages only. See
+[What Ships to the Pod](packaging.md).
+
+## Which mode to use
+
+| Situation | Mode | Reason |
+| --------- | ---- | ------ |
+| You start with Kinetic | Bundled | Works without setup. |
+| Your dependencies change less than once a day | Bundled | The cached image removes the build from later runs. |
+| Your dependency set is large | Bundled | You pay the install time one time, at build time. |
+| You need a reproducible environment | Bundled | The exact environment is frozen in a tagged image. |
+| You change the dependency file many times a day | Prebuilt | No build. The install runs at pod start. Needs a published base image. |
+| You need system libraries (custom CUDA builds, C++ libraries) | Custom, or prebuilt with your own Dockerfile | Bundled installs Python packages only. `kinetic build-image --dockerfile` lets a prebuilt base image contain system libraries. |
+| You need private packages that require credentials at install time | Custom | You control the build environment. |
+| Your organization requires a vetted base image | Custom | Use the image that your platform team approves. |
+
+## Bundled mode (default)
+
+Kinetic runs Cloud Build to produce an image and pushes the image to the
+Artifact Registry repository of the cluster. The image contains:
+
+- A `python:{X.Y}-slim` base image, where `X.Y` is the Python minor
+ version of your local interpreter. Pickled code is not portable across
+ minor versions, so this rule keeps the pod compatible with your client.
+- JAX with the runtime for the accelerator category: `jax[tpu]` with
+ `libtpu`, `jax[cuda12]`, or plain `jax` for CPU.
+- `keras`, `cloudpickle`, `google-cloud-storage`, and `keras-kinetic`
+ pinned to your client version.
+- The packages from your dependency file, without the JAX entries. See
+ [Dependencies](dependencies.md).
+- The Kinetic runner script.
+
+Kinetic tags the image with a hash of the base image, the accelerator
+category, the Kinetic version, the filtered dependency file, the runner
+script, and the Dockerfile template. Two jobs with the same inputs share
+one image.
+
+**Timing:**
+
+- **Cold** (first run, or after a change to the dependency file): about
+ 5 to 10 minutes for the build.
+- **Warm** (cached image): no build. The pod starts in less than 1
+ minute while a node still runs, or after the 2 to 5 minutes that a new
+ node needs.
+
+To inspect a build, list the Cloud Build history of the project:
-- **Reproducible**: The exact environment is frozen in the image.
-- **First-run cost**: The initial build takes ~2-5 minutes. Subsequent runs with unchanged dependencies use the cached image and start within a few seconds.
-- **Good for**: Production workloads, large dependency sets where you want to avoid per-run install overhead, or when you need a fully reproducible environment.
+```bash
+gcloud builds list --limit=5
+gcloud builds log
+```
-## Prebuilt Mode
+## Prebuilt mode
-Prebuilt mode uses a pre-published base image that already contains the accelerator runtime (JAX, CUDA/TPU libraries) and core dependencies. Your project's `requirements.txt` or `pyproject.toml` dependencies are installed at pod startup via `uv pip install`, so there is no Cloud Build step.
+In prebuilt mode, Kinetic does not build. Kinetic pulls a **base image**
+that already contains the accelerator runtime and the core packages, and
+the pod installs your dependency file at start with `uv pip install`.
-```python
-@kinetic.run(accelerator="tpu-v6e-8", container_image="prebuilt")
-def train(): ...
-```
+:::{warning}
+**Kinetic does not publish base images.** The default repository name is
+`kinetic` on Docker Hub, but no images exist there. Before you use
+prebuilt mode, publish base images to your own repository with
+`kinetic build-image`, and point Kinetic at that repository.
+:::
-### How it works
+### Set up prebuilt mode
-1. Kinetic resolves the base image from the image repository (see [Custom prebuilt images](#custom-prebuilt-images) below) using the accelerator category (`cpu`, `gpu`, or `tpu`) and the kinetic package version.
-2. Your project dependencies are filtered (JAX packages are removed to avoid conflicts) and uploaded to GCS alongside your code.
-3. At pod startup, the runner installs your dependencies with `uv pip install` before executing your function.
+1. Build and push the base images. One image per accelerator category:
-### Tradeoffs
+ ```bash
+ kinetic build-image --repo us-docker.pkg.dev/my-project/kinetic-base
+ ```
-- **Fast iteration**: No build step means jobs start quickly.
-- **Startup cost**: `uv pip install` runs on every job. For large dependency sets this adds time to each run.
-- **Good for**: Most workflows, especially during development and experimentation.
+2. Tell Kinetic where the images are. Set the environment variable, or
+ pass the repository in the decorator:
-### Custom prebuilt images
+ ::::{tab-set}
-By default, Kinetic pulls official base images from Docker Hub (`kinetic/base-{category}:{version}`). To use your own prebuilt images — for example, with additional system libraries or private packages — build and push them with the `kinetic build-image` command, then point Kinetic at your repository.
+ :::{tab-item} Environment variable
-Build and push images:
+ ```bash
+ export KINETIC_BASE_IMAGE_REPO=us-docker.pkg.dev/my-project/kinetic-base
+ ```
+ :::
-```bash
-kinetic build-image --repo us-docker.pkg.dev/my-project/kinetic-base
-```
+ :::{tab-item} Decorator argument
-Then set the repository so Kinetic uses your images:
+ ```python
+ @kinetic.run(
+ accelerator="gpu-l4",
+ container_image="prebuilt",
+ base_image_repo="us-docker.pkg.dev/my-project/kinetic-base",
+ )
+ def train(): ...
+ ```
+ :::
-::::{tab-set}
+ ::::
-:::{tab-item} Environment variable
+3. Select the mode with `container_image="prebuilt"`.
-```bash
-export KINETIC_BASE_IMAGE_REPO=us-docker.pkg.dev/my-project/kinetic-base
-```
-:::
+### How a prebuilt job starts
-:::{tab-item} Decorator argument
+1. Kinetic resolves the image name
+ `{repo}/base-{cpu|gpu|tpu}:{kinetic version}`. The version is the
+ version of your installed `keras-kinetic` package. If you upgrade the
+ client, run `kinetic build-image` again to publish images with the new
+ tag.
+2. Kinetic filters the JAX entries out of your dependency file and
+ uploads the result next to the job artifacts. Kinetic refuses a line
+ that points to a local path (`-r other.txt`, `-e .`, `./wheel.whl`),
+ because that path does not exist on the pod.
+3. The pod pulls the base image, runs `uv pip install` on the uploaded
+ file, and then runs your function.
-```python
-@kinetic.run(
- accelerator="l4", base_image_repo="us-docker.pkg.dev/my-project/kinetic-base"
-)
-def train(): ...
-```
-:::
+**Timing:** the image pull takes 30 to 60 seconds the first time on a
+node, and almost no time once the node has the image. The install time
+depends on your dependency file. A small file installs in less than 1
+minute.
-::::
-
-See [`kinetic build-image`](#kinetic-build-image) for the full command reference.
+The base image must have the same Python minor version as the client
+that submits the job. `kinetic build-image` uses the Python version of
+the machine that runs it.
-## Custom Image Mode
+## Custom image mode
-Provide a full container image URI to use your own image. Kinetic skips all build and dependency steps.
+Pass an image URI, and Kinetic uses the image without changes. Kinetic does not
+build and does not install. Your image is responsible for every package
+that your function imports.
```python
@kinetic.run(
- accelerator="tpu-v6e-8",
+ accelerator="tpu-v5litepod-4",
container_image="us-docker.pkg.dev/my-project/kinetic/my-image:v1.0",
)
def train(): ...
```
-### Requirements for custom images
-
-Your custom image must:
-
-1. Include `cloudpickle`, `google-cloud-storage`, `absl-py`, and a compatible Python environment. The runner imports these three packages at startup.
-2. Include the necessary dependencies for your function.
-3. Be accessible from the GKE nodes (e.g., Artifact Registry in the same GCP project, or a public registry).
-
-### When to use custom images
-
-- **Complex dependencies**: Non-Python system libraries (CUDA builds, C++ libs) that aren't in the default template.
-- **Corporate compliance**: Base images vetted by your security or platform team.
-- **Full control**: When you want to manage the entire image lifecycle yourself.
-
+The image must satisfy these requirements:
+
+* The Kinetic runner script is at `/app/remote_runner.py`. Kinetic
+starts the container with `python3 -u /app/remote_runner.py`, which
+replaces the `ENTRYPOINT` and `CMD` of the image. Copy the script from
+the `kinetic/runner/` directory of the installed package.
+* `python3` is on `PATH`, with the same minor version as your client.
+* The packages `cloudpickle`, `google-cloud-storage`, and `absl-py` are
+installed. The runner imports them.
+* Every other package that your function imports is installed.
+* We recommend that `keras-kinetic` is installed. The runner does not
+need it, but user code often imports it.
+* The GKE nodes can pull the image: Artifact Registry in the same
+project, or a public registry.
+
+Kinetic still ships your project source with the job and extracts it on
+the pod. Do not copy your source into the image.
+
+**Timing:** one image pull, then the function runs. The pull time depends
+on the image size and the registry.
## `kinetic build-image`
-Build and push prebuilt base images to a Docker Hub or Artifact Registry repository. One image is built per accelerator category (`cpu`, `gpu`, `tpu`) using Cloud Build.
+`kinetic build-image` builds base images with Cloud Build and pushes them
+to Docker Hub or to Artifact Registry. The command builds one image per
+accelerator category. The command needs an existing cluster: it uploads
+the build context to the builds bucket of the cluster and runs Cloud
+Build as the build service account of the cluster. Run `kinetic up`
+first.
```bash
-# Interactive mode — guides you through registry selection and setup
+# Interactive: the command asks for the registry and the settings.
kinetic build-image
-# Non-interactive with Artifact Registry
+# Artifact Registry, without prompts.
kinetic build-image \
--repo us-docker.pkg.dev/my-project/kinetic-base \
--project my-project \
--yes
-# Build only GPU and TPU images
+# GPU and TPU images only.
kinetic build-image --repo myuser/kinetic --category gpu --category tpu
-# Use a custom Dockerfile
+# A custom Dockerfile.
kinetic build-image --repo myuser/kinetic --dockerfile ./Dockerfile.custom
-# Specific version tag (default: kinetic package version)
+# A specific tag. The default is the kinetic package version.
kinetic build-image --repo myuser/kinetic --tag v2.0.0
```
-### Options
-
| Option | Description |
| ---------------------- | ------------------------------------------------------------------------------------------------------------------------------ |
| `--repo` | Image repository (Docker Hub or Artifact Registry). Omit to select interactively. |
@@ -177,36 +231,57 @@ kinetic build-image --repo myuser/kinetic --tag v2.0.0
| `--update-credentials` | Re-enter Docker Hub credentials even if they already exist in Secret Manager. |
| `--yes`, `-y` | Skip confirmation prompt. |
| `--project` | GCP project ID (default: `KINETIC_PROJECT`). |
-| `--cluster` | GKE cluster name (default: `kinetic-cluster`). |
+| `--cluster` | GKE cluster name (default: `kinetic-cluster`).
-### Registry support
+Registry notes:
-- **Docker Hub**: Credentials are stored in GCP Secret Manager and used by Cloud Build during the push. The command prompts for your Docker Hub username and access token on first use.
-- **Artifact Registry**: No additional credentials needed — the build service account authenticates automatically. The command prints the required `gcloud` setup commands for creating the repository and granting permissions.
+- **Docker Hub** — the command asks for your Docker Hub username and an
+ access token on first use, stores them in Secret Manager, and Cloud
+ Build uses them for the push.
+- **Artifact Registry** — no extra credentials. The build service
+ account pushes directly. The command prints the `gcloud` commands that
+ create the repository and grant the permissions.
+
+## How Kinetic decides
+
+At submit time, Kinetic reads `container_image`:
+
+1. `"prebuilt"` — resolve the base image for the accelerator category,
+ filter and upload the dependency file.
+2. `None` or `"bundled"` — hash the inputs, then reuse the cached image
+ or run Cloud Build.
+3. Any other string — use it as an image URI. No build, no install.
## Related pages
::::{grid} 1 1 2 2
:gutter: 3
-:::{grid-item-card} {octicon}`zap;1em` Execution Modes
-:link: execution_modes
+:::{grid-item-card} {octicon}`package;1em` Dependencies
+:link: dependencies
:link-type: doc
-Start here for the high-level mode-selection guidance.
+How Kinetic finds the dependency file, and what it filters out.
:::
-:::{grid-item-card} {octicon}`package;1em` Dependencies
-:link: dependencies
+:::{grid-item-card} {octicon}`file-directory;1em` What Ships to the Pod
+:link: packaging
+:link-type: doc
+
+The function, the source archive, and the Python version rule.
+:::
+
+:::{grid-item-card} {octicon}`graph;1em` Cost Optimization
+:link: cost_optimization
:link-type: doc
-What gets discovered and what gets installed in each mode.
+Cloud Build charges and how the image cache limits them.
:::
-:::{grid-item-card} {octicon}`gear;1em` Configuration
-:link: ../configuration
+:::{grid-item-card} {octicon}`bug;1em` Troubleshooting
+:link: ../troubleshooting
:link-type: doc
-`KINETIC_BASE_IMAGE_REPO` and other relevant env vars.
+Build failures and version-skew errors.
:::
::::
diff --git a/docs/guides/cost_optimization.md b/docs/guides/cost_optimization.md
index 9c9c5618..84385c61 100644
--- a/docs/guides/cost_optimization.md
+++ b/docs/guides/cost_optimization.md
@@ -1,116 +1,229 @@
-# Cost Optimization Guide
-
-Kinetic is designed to provide seamless cloud execution on Google Cloud Platform (GCP) while leveraging modern container orchestration (Google Kubernetes Engine, GKE) to keep operating costs extremely efficient. By default, Kinetic utilizes ephemeral functions and autoscaling node pools to ensure you only pay for compute while your workload is actively running.
-
-This guide covers the primary configurations and workflows for optimizing your cloud bill when using Kinetic.
-
-:::{tip}
-**Recommended defaults:**
-
-- Leave node pools at `--min-nodes 0` (the default) so accelerators
- scale to zero between jobs.
-- Pick a container mode that minimizes Cloud Build runs. Bundled mode
- (the default) reuses cached images across runs as long as your
- dependencies are stable, so you only pay for the first build per
- dependency set. Prebuilt mode skips Cloud Build entirely, but
- requires you to publish your own base image first with
- `kinetic build-image`.
-- Use `--spot` only on fault-tolerant single-host workloads with
- frequent checkpoints — spot is risky on multi-host slices because
- any host preemption fails the whole slice.
-- Run `kinetic pool list` periodically and remove pools you no
- longer use.
-:::
+# Cost Optimization
----
+This page explains what a Kinetic cluster costs while it is idle and
+while a job runs, and how to keep the bill low. Read this page after
+your [first run](../getting_started.md), when you decide how long the
+cluster stays up and which node pools it holds. Every command below
+reads the project, the zone, and the cluster from the active profile.
+
+## What you pay for
+
+A Kinetic cluster has two kinds of cost: the cost of a running job and
+the cost of an idle cluster.
+
+**While a job runs**, you pay for the accelerator node that runs the pod.
+The cluster autoscaler starts the node when the job needs it and deletes
+the node after an idle period. On a Spot node pool, the node costs less
+(see [Spot VMs](#spot-vms)).
-## 1. Understanding Scale-to-Zero Architecture
+**While the cluster exists**, you pay for these resources even when no
+job runs:
-By default, Kinetic provisions accelerator node pools with **scale-to-zero** capability.
+- The GKE control plane, about $0.10 per hour per cluster. Google Cloud
+ gives each billing account a
+ [monthly credit](https://cloud.google.com/kubernetes-engine/pricing)
+ that covers the control plane of one zonal cluster.
+- One `e2-standard-4` node in the default node pool. `kinetic up` creates
+ that node, and the node runs at all times. The node runs the cluster
+ system pods and `accelerator="cpu"` jobs. The free credit does not
+ cover the node.
+- One Cloud NAT gateway. The cluster nodes have private IP addresses,
+ and the gateway gives them internet access. Google Cloud bills the
+ gateway separately.
+- Storage in the jobs bucket and the builds bucket. Both buckets delete
+ objects that are older than 30 days.
-### How it works:
+Cloud Build bills each image build. Kinetic builds an image only when
+your dependency set changes (see
+[Container image builds](#container-image-builds)).
+
+`kinetic down` deletes all of these resources and stops all of these
+costs. See [Delete the cluster](#delete-the-cluster).
+
+## Scale to zero
+
+By default, every accelerator node pool scales to zero. A job goes
+through these steps:
:::{container} kinetic-steps
-1. **Job Submission:** When you invoke a remote function (via the
- `@kinetic.run()` decorator), Kinetic creates a Kubernetes `Job` on
- your GKE cluster.
-2. **Provisioning:** The GKE Cluster Autoscaler detects pending pods and
- spins up the required GPU/TPU virtual machines to support the job.
-3. **Execution & Termination:** Once your workload terminates (either
- successfully or due to an exception), the container exits and the pod
- goes away.
-4. **Scaling Down:** After the pod terminates, if no new jobs are
- submitted within the GKE idle window (typically around 10 minutes by
- default in GCP), the Cluster Autoscaler automatically shuts down the
- underlying VM instances, returning your accelerator compute
- consumption to zero.
+1. **Submit.** You call a function that has the `@kinetic.run()`
+ decorator. Kinetic creates a Kubernetes Job on the cluster.
+2. **Scale up.** The cluster autoscaler sees the pending pod and starts a
+ node in the matching node pool. This step takes about 2 to 5 minutes.
+3. **Run and exit.** The pod runs your function and exits, with a result
+ or with an exception.
+4. **Scale down.** If no new job arrives within the idle window, the
+ autoscaler deletes the node. The idle window is at most about 10
+ minutes. The accelerator cost returns to zero.
:::
-:::{important}
-Because the cluster control plane runs continuously, you will incur a baseline management fee of ~$0.10/hour per cluster. However, Google Cloud provides a [monthly free tier credit](https://cloud.google.com/kubernetes-engine/pricing) ($74.40/month) per billing account, which fully covers the management costs of exactly **one** active cluster.
-:::
+Kinetic creates the cluster with the `OPTIMIZE_UTILIZATION` GKE
+autoscaling mode. In this mode, GKE removes idle nodes faster than in
+the default balanced mode.
----
+## Warm nodes with `--min-nodes`
-## 2. Controlling Startup Latency vs. Idle Costs (`--min-nodes`)
+When you add a node pool, you choose between start latency and idle
+cost.
-When adding a node pool, you must trade off between **cold-start latency** and **idle compute costs**.
+**Default: scale to zero**
-### The Default Setup (Zero Idle Cost)
```bash
-# Creates a pool that scales from 0 to max required nodes
-kinetic pool add --accelerator l4
+kinetic pool add --accelerator gpu-l4
```
-- **Benefit:** Absolutely zero idle accelerator costs when you aren't running jobs.
-- **Trade-off:** Every time you submit a job after an idle period, you encounter a cold start (~2 to 5 minutes for VM provisioning, attach operations, and container image fetching).
-### Warm Persistent Nodes (Optimizing Developer Latency)
-If you are actively debugging or running highly iterative workloads, waiting for VMs to provision can interrupt your flow. You can keep a persistent warm node running using the `--min-nodes` parameter:
+- Benefit: no accelerator cost when no job runs.
+- Trade-off: after an idle period, each job waits about 2 to 5 minutes
+ for a node. The wait covers the VM start and the image pull.
+
+**Warm nodes**
+
+If you iterate on a script and submit a job every few minutes, the wait
+for a node interrupts your work. Keep one node warm with `--min-nodes`:
```bash
-# Keeps at least 1 node alive even when there are no jobs running
-kinetic pool add --accelerator l4 --min-nodes 1
+kinetic pool add --accelerator gpu-l4 --min-nodes 1
```
-- **Benefit:** Instant job scheduling and near-zero launch overhead because the container image is already cached locally on the node and the OS/driver layer is initialized.
-- **Trade-off:** You pay the continuous on-demand hourly rate for that node until you run `kinetic pool remove` or scale the `--min-nodes` count back down.
----
+- Benefit: a job starts as soon as Kinetic schedules the pod. The node
+ is up, and after the first job the node holds the image in its local
+ cache.
+- Trade-off: you pay for that node at all times, also at night and on
+ weekends.
-## 3. Optimizing Container Build Costs (`container_image`)
+`kinetic up --min-nodes N` applies the same setting to the first node
+pool of a new cluster.
-When you submit a job, Kinetic builds or retrieves a container image to execute your workload.
+:::{note}
+Kinetic has no command to change `--min-nodes` on an existing pool. To
+return a pool to scale-to-zero, remove the pool and add it again:
-By default, Kinetic uses **custom bundled images** (`container_image="bundled"`). This means it packages your local workspace and runs Cloud Build to freeze dependencies into an image.
+```bash
+kinetic pool list
+kinetic pool remove gpu-l4-1a2b
+kinetic pool add --accelerator gpu-l4
+```
-To save build time and Cloud Build execution charges, you can explicitly request prebuilt base configurations via `container_image="prebuilt"` in `@kinetic.run()`:
+The new pool gets a new generated name.
+:::
-```python
-@kinetic.run(accelerator="l4", container_image="prebuilt")
-def train(): ...
+For a multi-host TPU pool, `--min-nodes` counts VMs, not slices, and the
+value must be a multiple of the number of hosts in the slice. For
+example, a `tpu-v5litepod-16` slice has 4 hosts (see the **Hosts**
+column in [Accelerators](../accelerators.md#tpus)), so `--min-nodes 4`
+keeps one warm slice.
+
+## Container image builds
+
+Kinetic runs the pod in a container image that it builds with Cloud
+Build. Cloud Build bills the build minutes, and a cold build takes about
+5 to 10 minutes. You do not need to change anything to avoid repeated
+builds:
+
+- Kinetic tags each image with a hash of the dependency set and caches
+ the image in the Artifact Registry repository of the cluster.
+- A change to your code does not cause a build. Kinetic ships your code
+ in a separate archive for every job.
+- Kinetic builds a new image only when one of these inputs changes: the
+ dependency file (`requirements.txt` or `pyproject.toml`), your Python
+ minor version, the Kinetic version, or the accelerator category (CPU,
+ GPU, or TPU).
+
+Two habits keep the number of builds low:
+
+- Keep the dependency file stable. Add a package one time, not once per
+ experiment.
+- Use one Python minor version on your machine for all jobs on a
+ cluster.
+
+The other image modes do not lower the build cost. See
+[Container Images](containers.md).
+
+## Spot VMs
+
+Spot VMs use spare Google Cloud capacity at a large discount, up to 91%
+below the on-demand price (see the
+[Spot VM documentation](https://cloud.google.com/compute/docs/instances/spot)).
+Google Cloud can preempt a Spot VM at any time, with 30 seconds of
+notice.
+
+Spot use in Kinetic has two parts. Both parts are necessary.
+
+1. **The pool.** Add a node pool with `--spot`:
+
+ ```bash
+ kinetic pool add --accelerator gpu-a100 --spot
+ ```
+
+2. **The job.** Set `spot=True` on the decorator, or add the `:spot`
+ suffix to the accelerator string:
+
+ ```python
+ @kinetic.run(accelerator="gpu-a100", spot=True)
+ def train(): ...
+ ```
+
+ ```python
+ @kinetic.run(accelerator="gpu-a100:spot")
+ def train(): ...
+ ```
+
+The job side adds the Spot node selector and the Spot toleration to the
+pod. Without them, the pod cannot schedule on the Spot pool, and the
+job stays `PENDING`. A job with `spot=True` needs a `--spot` pool, and a
+job without `spot=True` needs an on-demand pool. If you run both kinds
+of job on one accelerator type, add two pools.
+
+Follow these rules when you use Spot:
+
+- **Run only fault-tolerant jobs on Spot.** Kinetic does not submit a
+ preempted single-host job again. The job ends as `FAILED`, and you
+ submit it again yourself. Do not use Spot for a job with a deadline or
+ for a job that must run without interruption.
+- **Write checkpoints under `KINETIC_OUTPUT_DIR`**, for example with
+ Orbax. Pass a fixed `output_dir=` to the decorator so that the second
+ submission finds the checkpoints of the first one. See
+ [Resume a job from a checkpoint](checkpointing.md#resume-a-job-from-a-checkpoint).
+- **Prefer single-host accelerators.** A multi-host TPU slice loses all
+ of its work when Google Cloud preempts any one of its hosts. Multi-host
+ slices include `tpu-v6e-8` and `tpu-v6e-16`, every `tpu-v5p` size,
+ `tpu-v5litepod-16` and larger, `tpu-v4-8` and larger, and `tpu-v3-16`
+ and larger. Use Spot for single-host jobs, for example on
+ `tpu-v5litepod-4`, `tpu-v5litepod-8`, or `gpu-l4`.
+- **Do not combine `debug=True` with `spot=True`.** A preemption ends the
+ debug session. Kinetic prints a warning for this combination.
+
+## Capacity reservations
+
+A capacity reservation does not lower the price. It guarantees that a
+node is available when the pool scales up. A reservation is useful for
+accelerators in short supply, for example H100 or TPU v6e. Bind a
+reservation to a pool with `--reservation`:
+
+```bash
+kinetic pool add --accelerator gpu-h100 --reservation my-h100-reservation
```
-### Why `"prebuilt"` reduces GCP expenses:
-- **Zero Cloud Build Fees:** Reuses existing Artifact Registry images immediately.
-- **Fast Execution Start:** Minor script changes are uploaded compactly as a context zip payload (`context.zip`) directly to GCS without re-compiling environment layers.
+You cannot combine `--reservation` with `--spot`. Google Cloud bills a
+reservation while it exists, also when no VM uses it. See
+[Capacity Reservations](reservations.md).
----
+## Node pool hygiene
-## 4. Slashing Compute Costs with Spot Instances (`--spot`)
+Each node pool with `--min-nodes 0` costs nothing while it is idle, but
+an unused pool with warm nodes costs money every hour. Review the pools
+of the cluster regularly:
-Spot instances run on unused Google Cloud capacity at significantly discounted rates (**up to 91% lower** than standard on-demand pricing according to the official [GCP Spot VM documentation](https://cloud.google.com/compute/docs/instances/spot)).
+```bash
+kinetic pool list
+```
-You can provision a fully automated Spot pool using the `--spot` flag:
+Remove each pool that you no longer use:
```bash
-# Add an A100 node pool backed by highly-discounted Spot capacity
-kinetic pool add --accelerator a100 --spot
+kinetic pool remove tpu-v5litepod-1a2b
```
-### Best Practices for Spot Nodes:
-1. **Fault-Tolerant Workloads Only:** Spot instances can be preempted by GCP with only 30 seconds of notice when capacity constraints occur. Do not use Spot pools for stateful production serving or time-critical jobs that cannot afford restarts.
-2. **Use Checkpointing:** Use Kinetic's integration with **Orbax** (see the [Checkpointing Guide](checkpointing.md)) to continuously flush state to Cloud Storage (`gs://`). If a Spot preemption kills your training run midway, you can resume from the last saved checkpoint instead of restarting from scratch.
-3. **Multi-Host TPUs:** While you can provision Spot pools for multi-node TPUs (such as `tpu-v3`, `tpu-v4`, `tpu-v5p`, or `tpu-v6e`), if any single host in the TPU slice is preempted, the entire slice job will fail. Spot pricing is therefore highly effective for single-host jobs (like `tpu-v5litepod-4`, `tpu-v5litepod-8`, or `l4` workloads) where you minimize the probability of aggregate preemption.
-
### Verify the Provisioning Model
Kinetic records the Spot setting in the cluster state. Later `kinetic pool add`, `kinetic pool remove`, and `kinetic up` commands keep the setting on the pool.
@@ -123,52 +236,74 @@ Kinetic records the Spot setting when it creates the pool. A pool from an earlie
---
-## 5. Managing Capacity Reservations
+## Delete the cluster
-If your Google Cloud project leverages centralized enterprise pricing via **On-Demand Capacity Reservations**, you can instruct Kinetic to consume that specific reserved capacity pool instead of competing for standard compute stock:
+If you do not use Kinetic for days or weeks, delete the cluster:
```bash
-kinetic pool add --accelerator h100 --reservation my-h100-reservation
+kinetic down
```
-:::{note}
-You cannot mix `--spot` pricing with `--reservation`. To utilize a reservation, your node pools must use standard on-demand billing tiers. For more details, see [Capacity Reservations](reservations.md).
+`kinetic down` deletes the cluster, the node pools, the default node,
+the NAT gateway, the Artifact Registry repository, and the jobs bucket
+and the builds bucket. After that, the cluster costs nothing. The state
+bucket of the project (`gs://{project}-kinetic-state`) stays, but it
+holds only small state files.
+
+:::{warning}
+`kinetic down` deletes the buckets with their contents, including
+everything under `KINETIC_OUTPUT_DIR`. Copy the outputs that you want to
+keep to another bucket first.
:::
----
+`kinetic up` creates a new cluster in about 5 to 10 minutes. The first
+job on the new cluster builds a new image, because the image cache is
+part of the cluster. See [Clusters and Node Pools](clusters.md).
-## 6. Summary Checklist for Cost Optimization
+## Checklist
-- [ ] **Utilize Prebuilt Images:** Apply `container_image="prebuilt"` to bypass unneeded cloud dependency compilation.
-- [ ] **Rely on defaults:** Let Kinetic's default `--min-nodes 0` configuration automatically scale your infrastructure down when you step away from your desk.
-- [ ] **Utilize Spot capacity:** Use `--spot` for long-running pretraining jobs, ensuring you regularly save model weights to Cloud Storage.
-- [ ] **Prune inactive pools:** Actively review your existing infrastructure by running `kinetic pool list` and drop idle pools using `kinetic pool remove `.
-- [ ] **Tear down unused clusters:** If you are not using Kinetic for days or weeks at a time, remove the entire underlying cluster via `kinetic down` to avoid baseline control plane charges. **Warning:** This will also delete your Cloud Storage buckets and any saved job data.
+1. Leave `--min-nodes` at `0` unless you submit a job every few minutes.
+2. Keep the dependency file stable so that Kinetic reuses the cached
+ image.
+3. For a fault-tolerant single-host job, add a `--spot` pool and set
+ `spot=True` on the job. Write checkpoints under `KINETIC_OUTPUT_DIR`.
+4. Run `kinetic pool list` and remove the pools that you no longer use.
+5. Run `kinetic down` when you do not need the cluster for days or
+ weeks. Copy your outputs first.
## Related pages
::::{grid} 1 1 2 2
:gutter: 3
-:::{grid-item-card} {octicon}`package;1em` Execution Modes
-:link: execution_modes
+:::{grid-item-card} {octicon}`server;1em` Clusters and Node Pools
+:link: clusters
+:link-type: doc
+
+What `kinetic up` creates, node pool options, and when a second
+cluster is justified.
+:::
+
+:::{grid-item-card} {octicon}`history;1em` Outputs and Checkpoints
+:link: checkpointing
:link-type: doc
-Bundled vs prebuilt tradeoffs when build cost matters.
+Write checkpoints under `KINETIC_OUTPUT_DIR` and resume a job after a
+preemption.
:::
:::{grid-item-card} {octicon}`checklist;1em` Capacity Reservations
:link: reservations
:link-type: doc
-Guaranteed capacity for newer accelerators.
+Guaranteed capacity for accelerators in short supply.
:::
-:::{grid-item-card} {octicon}`server;1em` Multiple Clusters
-:link: clusters
+:::{grid-item-card} {octicon}`package;1em` Container Images
+:link: containers
:link-type: doc
-When separate control planes are worth the extra ~$0.10/hr.
+How the image cache works, and when prebuilt or custom images make
+sense.
:::
::::
-
diff --git a/docs/guides/data.md b/docs/guides/data.md
index ca1cf853..89d9fa91 100644
--- a/docs/guides/data.md
+++ b/docs/guides/data.md
@@ -1,42 +1,53 @@
# Working with Data
-`kinetic.Data(...)` is the API for getting bytes into your remote function.
-It accepts a local file or directory path, or a `gs://` URI, and resolves
-to a plain filesystem path inside the pod. Your function code only sees
-paths — never URIs, never `Data` objects.
-
-That uniformity is the whole point: you write the same training code
-whether the data started on your laptop, in a GCS bucket, or as a
-FUSE-mounted dataset too large to fit on disk.
+`kinetic.Data(...)` is the API that makes input data available to your
+remote function. It accepts a local file, a local directory, a `gs://`
+Cloud Storage URI, or an `hf://` Hugging Face dataset URI. On the pod,
+Kinetic replaces each `Data` object with a plain filesystem path. Your
+function sees only paths, never URIs and never `Data` objects. This page
+explains the three ways to read data, when to use each way, and the
+limits. The appendix at the end is for contributors.
## A first example
```python
+import os
+
import kinetic
from kinetic import Data
@kinetic.run(accelerator="cpu")
def process_data(data_path):
- import os
-
print(f"Reading from: {data_path}")
return sorted(os.listdir(data_path))
-# Local directory
+# A local directory. Kinetic uploads it one time and downloads it to the pod.
process_data(Data("./my_dataset/"))
-# GCS directory — trailing slash signals it's a directory
+# A Cloud Storage directory. The trailing slash marks a directory.
process_data(Data("gs://my-bucket/training-set/"))
```
-`Data` works as a function argument, as a value inside a list/dict, and as
-a value in the `volumes={...}` decorator argument:
+The function code is the same for a local directory and for a Cloud
+Storage directory. In both cases `data_path` is a directory on the pod.
+
+## Two ways to pass `Data`
+
+You can pass a `Data` object in two places:
+
+- **As a function argument.** The function receives the path as the
+ argument value. A `Data` object can also sit inside a list, a tuple, or
+ a dict that you pass as an argument. Use this way when the function
+ takes the path explicitly.
+- **In the `volumes={...}` decorator argument.** Kinetic places the data
+ at the mount path that you give. Use this way when your training script
+ reads from a fixed absolute path.
```python
@kinetic.run(
- accelerator="tpu-v5e-4",
+ accelerator="tpu-v5litepod-4",
volumes={"/data": Data("./dataset/")},
)
def train():
@@ -46,74 +57,92 @@ def train():
return len(df)
```
-Use `volumes={...}` when your training script has hardcoded absolute
-paths it expects to read from. Pass `Data(...)` as a function argument
-when you'd rather receive the path explicitly.
-
-## Choosing a data access pattern
-
-Three patterns cover almost everything:
-
-1. **Downloaded `Data`** (default) — `Data("...")`. Kinetic copies the
- bytes onto the pod's local disk before your function runs. Reads are
- fast (local disk), but the pod has to wait for the download to finish.
-2. **FUSE-mounted `Data`** — `Data("gs://...", fuse=True)`. The bucket
- is mounted lazily; only files you actually `open()` are fetched from
- GCS. Pod startup is near-instant; per-file reads pay GCS latency.
-3. **Raw `gs://` streaming** — your code uses `tf.io.gfile`,
- `gcsfs`, or a similar library to talk to GCS directly without
- `Data(...)`. This bypasses the `Data` abstraction entirely; reach for
- it only when you have a specific reason to.
-
-Decision table:
-
-| Dataset size | Access pattern | Use |
-| ------------------ | ------------------------- | -------------------------------------------- |
-| Small (<10 GB) | Read most/all files | `Data(...)` (downloaded) |
-| Small (<10 GB) | Random access | `Data(...)` (downloaded) |
-| Medium (10–100 GB) | Streaming once-through | `Data(..., fuse=True)` |
-| Medium (10–100 GB) | Random access many epochs | `Data(...)` (downloaded) |
-| Large (>100 GB) | Streaming, sparse subset | `Data(..., fuse=True)` |
-| Large (>100 GB) | Need indexed shards | `Data(..., fuse=True)` + `tf.data` / `grain` |
-| Already in GCS | Any size | `Data("gs://...")` (with or without `fuse`) |
+The mount path must be an absolute path that starts with `/`. The mount
+path of a volume is a directory. If the `Data` object is a single file,
+the pod places that file inside the mount directory.
+
+## Choose an access pattern
+
+Three patterns cover almost every job:
+
+1. **Downloaded `Data`** (the default) — `Data("...")`. Kinetic copies
+ the data to the local disk of the pod before your function starts.
+ Reads are fast, but the pod waits for the download to finish.
+2. **FUSE-mounted `Data`** — `Data("gs://...", fuse=True)`. Kinetic
+ mounts the Cloud Storage prefix with the GCS FUSE CSI driver. The pod
+ does not wait for a download. Each read fetches the bytes from Cloud
+ Storage on demand.
+3. **Direct `gs://` access** — your code reads Cloud Storage with
+ `tf.io.gfile`, `gcsfs`, `tf.data`, `grain`, or a similar library.
+ You pass the URI as a plain string, not as a `Data` object. Kinetic
+ passes that string through unchanged. Use this pattern only when your
+ framework already has a Cloud Storage reader that you want to keep.
+
+Use this table to select a pattern:
+
+| Dataset size | Access | Use |
+| ------------------ | -------------------------- | -------------------------------------------- |
+| Small (<10 GB) | Read most or all files | `Data(...)` (downloaded) |
+| Small (<10 GB) | Random access | `Data(...)` (downloaded) |
+| Medium (10–100 GB) | Stream one time | `Data(..., fuse=True)` |
+| Medium (10–100 GB) | Random access, many epochs | `Data(...)` (downloaded) |
+| Large (>100 GB) | Stream a sparse subset | `Data(..., fuse=True)` |
+| Large (>100 GB) | Indexed shards | `Data(..., fuse=True)` + `tf.data` / `grain` |
+| Already in GCS | Any size | `Data("gs://...")` (with or without `fuse`) |
:::{tip}
**Recommended defaults:**
-- For small or medium datasets you read every epoch, use plain
- `Data(...)`. The download cost is paid once at pod startup; subsequent
- reads are local-disk fast.
-- For datasets that are too large to fit on the pod's disk, or where you
- only touch a fraction of the files, use `Data("gs://...", fuse=True)`.
-- Wrap GCS data in `Data(...)` even when it is already in GCS so your
- function uses the same path-based API regardless of source. Note that
- Kinetic's content-hash-based upload caching applies only to local
- data; GCS-hosted `Data` is passed through by URI without rehashing or
- re-uploading.
+- For a small or medium dataset that you read every epoch, use plain
+ `Data(...)`. The pod downloads the data one time at start. All later
+ reads come from the local disk.
+- For a dataset that does not fit on the disk of the pod, use
+ `Data("gs://...", fuse=True)`. Also use `fuse=True` for a dataset where
+ you read only a fraction of the files.
+- Wrap Cloud Storage data in `Data(...)` even when the data is already in
+ a bucket. Your function then uses the same path-based API for every
+ source. Kinetic passes a `gs://` `Data` object through by URI. Kinetic
+ does not hash it and does not upload it. The upload cache applies only
+ to local data.
+:::
+
+:::{note}
+The pod reads Cloud Storage as the node service account of the cluster,
+`kn-{cluster}-nodes@{project}.iam.gserviceaccount.com`. That account can
+read the jobs bucket. For a bucket that `kinetic up` did not create,
+grant that account the `roles/storage.objectViewer` role on the bucket
+before you submit the job. This rule applies to all three patterns.
:::
## FUSE mounting
`fuse=True` mounts the data through the
[GCS FUSE CSI driver](https://cloud.google.com/kubernetes-engine/docs/how-to/persistent-volumes/cloud-storage-fuse-csi-driver)
-instead of downloading it. Your function still receives a filesystem
-path; reads stream on demand from GCS.
+instead of a download. Your function still receives a filesystem path.
+Reads stream from Cloud Storage on demand.
```python
@kinetic.run(
- accelerator="tpu-v5e-4",
+ accelerator="tpu-v5litepod-4",
volumes={"/data": Data("gs://my-bucket/imagenet/", fuse=True)},
)
def train():
- # Only files you open() are fetched from GCS
+ # The pod fetches only the files that the function opens.
...
```
-FUSE works with both `volumes={...}` and function arguments, with both
-local paths and GCS URIs. Single files work transparently — the pod sees
-a file path, not a directory:
+FUSE works with `volumes={...}` and with function arguments. FUSE also
+works with a local path: Kinetic uploads the local data one time and
+then mounts the uploaded copy on the pod. A local single file resolves to
+a file path on the pod, not to a directory:
```python
+import json
+
+import kinetic
+from kinetic import Data
+
+
@kinetic.run(accelerator="cpu")
def read_config(config_path):
with open(config_path) as f:
@@ -123,11 +152,11 @@ def read_config(config_path):
read_config(Data("./config.json", fuse=True))
```
-You can mix FUSE-mounted and downloaded data in the same job:
+You can mix FUSE-mounted data and downloaded data in one job:
```python
@kinetic.run(
- accelerator="tpu-v5e-4",
+ accelerator="tpu-v5litepod-4",
volumes={
"/data": Data("gs://my-bucket/large-dataset/", fuse=True),
"/config": Data("./small-config/"),
@@ -136,58 +165,172 @@ You can mix FUSE-mounted and downloaded data in the same job:
def train(extra_data): ...
-train(Data("./labels.csv")) # downloaded function-argument data
+train(Data("./labels.csv")) # a downloaded function argument
```
+Two more rules apply to FUSE:
+
+- **A FUSE mount is read-only.** Your function cannot write under a FUSE
+ mount path. Write outputs under `KINETIC_OUTPUT_DIR` instead. See
+ [Outputs and Checkpoints](checkpointing.md).
+- **Kinetic reserves the prefix `/_kinetic/fuse-data/`.** Kinetic mounts
+ FUSE function arguments below that prefix. A `volumes` key below that
+ prefix raises a `ValueError` at submit time.
+
:::{admonition} Prerequisites
:class: important
-FUSE mounting needs the GCS FUSE CSI driver addon on
-the GKE cluster. `kinetic up` enables it by default.
+FUSE mounting needs the GCS FUSE CSI driver addon on the GKE cluster.
+`kinetic up` enables the addon by default.
:::
:::{seealso}
-For a runnable end-to-end walkthrough covering volume mounts, single
-files, multiple FUSE volumes, and mixed FUSE/downloaded data in the same
-job, see
+For a runnable script that covers volume mounts, single files, many FUSE
+volumes, and mixed FUSE and downloaded data in one job, see
[`examples/example_fuse.py`](https://github.com/keras-team/kinetic/blob/main/examples/example_fuse.py).
:::
-## How it caches
+## Hugging Face datasets
+
+`Data` also accepts an `hf://` URI. The pod downloads the dataset with
+the `datasets` library and saves it to a local directory. Your function
+receives that directory path and loads it with `datasets.load_from_disk`.
+
+```python
+import kinetic
+from kinetic import Data
+
-Local data is content-addressed: identical bytes upload only once,
-regardless of how many jobs reference them. SHA-256 of the contents
-becomes the cache key, and re-runs with unchanged data skip the upload
-entirely.
+@kinetic.run(accelerator="tpu-v5litepod-4")
+def train(dataset_path):
+ from datasets import load_from_disk
-This also means files inside your project root that you wrap in
-`Data(...)` are automatically excluded from the per-job `context.zip`
-payload — no redundant upload of the same bytes.
+ ds = load_from_disk(dataset_path)
+ return len(ds)
+
+
+train(Data("hf://imdb?split=train"))
+```
+
+Three rules apply to `hf://` URIs:
+
+- The dependency file of your project must list `datasets`. Kinetic does
+ not install it for you. See [Dependencies](dependencies.md).
+- The query string accepts `split`, `config_name`, and `revision`, for
+ example `hf://user/repo?config_name=reviews&split=train`.
+- `Data` rejects `fuse=True` with an `hf://` URI. `Data` raises a
+ `ValueError`.
+
+If the dataset repository runs its own loading code, pass
+`Data("hf://user/repo", hf_trust_remote_code=True)`. The pod then runs
+code from that repository, so use this option only for a repository that
+you trust. For a gated or private dataset, forward your token to the pod
+with `capture_env_vars=["HF_TOKEN"]`. The `datasets` library reads
+`HF_TOKEN` from the environment of the pod. Kinetic itself does not read
+the token. See [Forward Environment Variables](env_vars.md).
+
+:::{seealso}
+For a runnable script that loads a public Hugging Face dataset with
+`config_name` and `split`, see
+[`examples/hf_dataset_demo.py`](https://github.com/keras-team/kinetic/blob/main/examples/hf_dataset_demo.py).
+:::
+
+## Limits and pitfalls
+
+:::{warning}
+**Kinetic does not support a single Cloud Storage object today.**
+`Data("gs://my-bucket/dir/file.h5")`, without a trailing slash, does not
+give your function that file. The download path finds nothing, and the
+function receives an empty directory. The FUSE path mounts the parent
+prefix and can return a different file from the same prefix. Do one of
+these two things instead:
+
+- Point at the directory with a trailing slash:
+ `Data("gs://my-bucket/dir/")`. Then open `file.h5` inside the path
+ that your function receives.
+- Upload the file from your machine: `Data("./file.h5")`. A local single
+ file resolves to a file path on the pod.
+:::
+
+Other limits:
+
+- **A `gs://` directory needs a trailing slash.** Kinetic reads
+ `Data("gs://my-bucket/dataset/")` as a directory and
+ `Data("gs://my-bucket/dataset")` as a single object. Kinetic logs a
+ warning when a `gs://` path has no trailing slash and the last segment
+ has no file extension.
+- **A `Data` instance is a snapshot.** Kinetic hashes the local files one
+ time per `Data` instance and caches the hash. If you edit the files and
+ submit the same instance again, the job uses the old upload. Create a
+ new `Data` object to upload the changed files.
+- **A `Data` object cannot be a set member or a dict key.** Before the
+ upload, Kinetic replaces each `Data` object with a dict, and a dict is
+ not hashable. Kinetic raises a `ValueError` at submit time for a `Data`
+ object inside a set or a frozenset. Kinetic raises the same error for a
+ `Data` object that is a dict key. Pass the `Data` object as its own
+ argument, or inside a list, a tuple, or a dict value.
+- **Kinetic finds `Data` objects only in containers.** Kinetic walks
+ lists, tuples, dicts, and their subclasses. Kinetic does not find a
+ `Data` object that you store as an attribute of your own class.
+- **Large local data logs a warning.** If a local `Data` object is larger
+ than 10 GB, Kinetic logs a warning before the upload. The warning
+ recommends a `gs://` URI with framework-native I/O.
+
+## How Kinetic caches local data
+
+Kinetic uploads local data one time and reuses the upload for every later
+job that references the same data. The cache key is a SHA-256 hash of
+the relative path and the contents of every file. The hash also includes
+a marker that identifies a single file or a directory. Two consequences
+follow:
+
+- A second run with the same directory skips the upload. Kinetic logs
+ `Data cache hit` and passes the existing Cloud Storage location.
+- A rename or a move of a file inside the directory changes the hash.
+ Kinetic uploads the directory again.
+
+Kinetic stores the upload in the jobs bucket at
+`gs://{jobs bucket}/default/data-cache/{hash}/`. The `default` segment is
+a literal string, not the Kubernetes namespace of your profile. The jobs
+bucket deletes objects that are older than 30 days, so the cache is valid
+for 30 days after the upload.
+
+Kinetic also excludes a local `Data` path from the source archive of the
+job when that path sits inside the package root. Kinetic uploads those
+files one time, through the data cache. See
+[What Ships to the Pod](packaging.md).
## Related pages
::::{grid} 1 1 2 2
:gutter: 3
-:::{grid-item-card} {octicon}`history;1em` Checkpointing
+:::{grid-item-card} {octicon}`history;1em` Outputs and Checkpoints
:link: checkpointing
:link-type: doc
-Durable outputs and `KINETIC_OUTPUT_DIR`.
+Write durable outputs and checkpoints under `KINETIC_OUTPUT_DIR`.
:::
-:::{grid-item-card} {octicon}`beaker;1em` Examples
-:link: ../examples
+:::{grid-item-card} {octicon}`package;1em` What Ships to the Pod
+:link: packaging
+:link-type: doc
+
+The package root, the source archive, and how `Data` paths are excluded.
+:::
+
+:::{grid-item-card} {octicon}`key;1em` Forward Environment Variables
+:link: env_vars
:link-type: doc
-Walks through the Data API end-to-end.
+Copy tokens such as `HF_TOKEN` from your shell to the pod.
:::
-:::{grid-item-card} {octicon}`graph;1em` Cost Optimization
-:link: cost_optimization
+:::{grid-item-card} {octicon}`beaker;1em` Examples
+:link: ../examples
:link-type: doc
-FUSE vs download tradeoffs for repeated jobs.
+Runnable scripts for the `Data` API and for FUSE mounts.
:::
::::
@@ -195,19 +338,19 @@ FUSE vs download tradeoffs for repeated jobs.
## Appendix: implementation internals
-The rest of this page is for contributors and people debugging
-data-related issues. End users do not need to read it.
+This appendix is for contributors and for people who debug data issues.
+Users do not need to read it.
### `Data` reference serialization
-`Data` objects can't be sent directly to the remote pod. During
-`_prepare_artifacts()`, each `Data` is uploaded to GCS and replaced with
-a serializable `__data_ref__` dict:
+A `Data` object does not travel to the pod. During `_prepare_artifacts()`,
+Kinetic uploads each local `Data` object and replaces every `Data` object
+with a serializable `__data_ref__` dict:
```python
{
"__data_ref__": True,
- "uri": "gs://bucket/namespace/data-cache/abc123",
+ "uri": "gs://bucket/default/data-cache/abc123",
"is_dir": True,
"mount_path": "/data", # None unless the Data is a volume or uses FUSE
"fuse": False, # True when fuse=True was passed
@@ -229,7 +372,7 @@ Kinetic sets `mount_path` for two kinds of `Data`:
A plain function argument gets `mount_path: None`.
-On the remote pod, `resolve_data_refs()` in `remote_runner.py` walks the
+On the pod, `resolve_data_refs()` in `remote_runner.py` walks the
deserialized args and kwargs and replaces these dicts with local
filesystem paths. The walk uses an identity memo. One `Data` object that
you pass two times thus resolves one time. The aliasing between your
@@ -243,47 +386,71 @@ hash the replacement dict:
- A `Data` object inside a set or a frozenset.
- A `Data` object used as a dictionary key.
+For a downloaded reference, `resolve_data_refs()` returns a directory
+path. If the reference is not a directory and the download produced
+exactly one file, it returns the path of that file. For an `hf://`
+reference, `_download_hf_data()` calls `datasets.load_dataset()` with the
+`split`, `config_name`, and `revision` query parameters and then calls
+`save_to_disk()` on the target directory.
+
### Upload and caching pipeline
-Local data is uploaded to `gs://{bucket}/{namespace}/data-cache/{hash}/`,
-where `{hash}` is a SHA-256 computed over sorted file contents. The flow:
+`upload_data()` in `kinetic/utils/storage.py` uploads local data to
+`gs://{jobs bucket}/default/data-cache/{hash}/`. The function has a
+`namespace_prefix` parameter with the default value `"default"`. The
+callers in `execution.py` do not pass that parameter, so the prefix is
+always the literal `default`. The flow:
:::{container} kinetic-steps
-1. **Compute content hash.**
- Deterministic: sorted DFS order, per-file SHA-256, then combined.
+1. **Compute the content hash.** `Data.content_hash()` hashes each file
+ as SHA-256 of `relpath + "\0" + contents`, in sorted DFS order. It then
+ combines the per-file digests under a `dir:` or `file:` prefix. The
+ instance caches the result.
-2. **Check for a sentinel blob** at `{namespace}/data-markers/{hash}` — if
- present, skip upload.
+2. **Check for a sentinel blob** at `default/data-markers/{hash}`. If
+ the blob exists, skip the upload.
-3. **Upload files** preserving directory structure under the hash prefix.
+3. **Upload the files** under the hash prefix. The directory structure
+ is preserved.
-4. **Write the sentinel blob** last to signal upload-complete.
+4. **Write the sentinel blob** last. The blob signals that the upload is
+ complete.
:::
-For single files, the blob is stored at `{hash}/{filename}`. For
-directories, the full tree is preserved under `{hash}/`. The returned
-GCS URI always points to the hash prefix directory, not individual files.
+For a single file, the blob is stored at `{hash}/{filename}`. For a
+directory, the full tree is preserved under `{hash}/`. The returned Cloud
+Storage URI always points to the hash prefix directory, not to a file.
+For a `gs://` or `hf://` `Data` object, `upload_data()` returns the
+original URI without an upload.
### FUSE mount implementation
-GCS FUSE can only mount directories, not individual files. The system
-handles this through several layers:
+GCS FUSE mounts directories, not single files. Three layers handle
+single files:
-**Volume spec construction** (`execution.py`): for `fuse=True` Data, a
-FUSE volume spec is built with `gcs_uri`, `mount_path`, `is_dir`, and
-`read_only`. Specs live on `ctx.fuse_volume_specs` and pass through to
-the backend.
+**Volume spec construction** (`execution.py`): for a `fuse=True` `Data`
+object, Kinetic builds a FUSE volume spec. The spec has the keys
+`gcs_uri`, `mount_path`, `is_dir`, and `read_only`. `read_only` is always
+`True`. The specs live on `ctx.fuse_volume_specs` and pass through to the
+backend.
**URI adjustment for uploaded single files:** `upload_data()` returns a
-directory-level URI (`gs://bucket/ns/data-cache/{hash}`) since the hash
-prefix is a directory. For FUSE single-file mounts, `_fuse_gcs_uri()`
-appends the original filename (`gs://bucket/ns/data-cache/{hash}/config.json`)
-so the `only-dir` mount option scopes to the hash directory rather than
-the entire `data-cache/` tree. The data ref retains the directory-level
-URI for download compatibility.
-
-**K8s volume generation:** each spec becomes an inline ephemeral CSI
-volume. The `only-dir` mount option scopes the mount to a specific GCS
-prefix. For single files (`is_dir=False`), the parent directory is
-mounted. The pod receives a `gke-gcsfuse/volumes: "true"` annotation to
-trigger the GCS FUSE sidecar injection.
+directory-level URI (`gs://bucket/default/data-cache/{hash}`), because
+the hash prefix is a directory. For a FUSE single-file mount of a local
+file, `_fuse_gcs_uri()` appends the original filename
+(`gs://bucket/default/data-cache/{hash}/config.json`). The `only-dir`
+mount option then scopes the mount to the hash directory and not to the
+whole `data-cache/` tree. The data ref keeps the directory-level URI. A
+`gs://` URI passes through `_fuse_gcs_uri()` unchanged.
+
+**Kubernetes volume generation** (`k8s_utils.py`): each spec becomes an
+inline ephemeral CSI volume with the `implicit-dirs` mount option. The
+`only-dir` mount option scopes the mount to one Cloud Storage prefix. For
+a single file (`is_dir=False`), Kinetic mounts the parent directory. The
+pod receives a `gke-gcsfuse/volumes: "true"` annotation, and GKE then
+injects the GCS FUSE sidecar. On the pod, `_resolve_fuse_single_file()`
+returns the first entry of the mount directory. That result is correct
+for an uploaded local file, which is alone under its hash prefix. For a
+`gs://` single object, the parent prefix can hold many objects, and the
+first entry can be a different file. That is the reason for the
+limitation in the [Limits and pitfalls](#limits-and-pitfalls) section.
diff --git a/docs/guides/debugging.md b/docs/guides/debugging.md
index 263a89d3..abaad3e4 100644
--- a/docs/guides/debugging.md
+++ b/docs/guides/debugging.md
@@ -1,29 +1,40 @@
# Interactive Debugging
-Pass `debug=True` to `@kinetic.run()` to attach
-a VS Code debugger to the remote pod. Set breakpoints, step through
-your function, inspect variables, and evaluate expressions against the
-accelerator your code is running on.
-
-Kinetic prints a ready-to-paste `launch.json` entry — standard
-`debugpy` attach config — so every VS Code-derived editor picks it up
-as-is: **VS Code**, **Cursor**, **Windsurf**, **Antigravity**,
-**VSCodium**, and anything else that ships the Python / debugpy
-extension.
+This page explains how to attach a debugger to a job that runs on the
+cluster. Pass `debug=True` to `@kinetic.run()`. The pod then starts a
+`debugpy` server and waits for your editor before it calls your
+function. You set breakpoints, step through the function, and inspect
+variables on the accelerator that runs the code. Read this page when a
+job fails in a way that the logs do not explain, or when you want to
+explore data on the pod.
+
+Kinetic prints a `launch.json` entry that uses the standard `debugpy`
+attach request. VS Code, and other editors that use the VS Code Python
+debugger, for example Cursor and VSCodium, can use that entry without changes.
+
+## Before you start
+
+- An active profile. `kinetic init` creates the profile. Kinetic reads
+ the project and the cluster from the profile, so the commands on this
+ page do not need `--project` or `--cluster`.
+- `kubectl` on your `PATH`. Kinetic runs `kubectl port-forward` to
+ connect your editor to the pod.
+- VS Code, or another editor, with the Python and Python Debugger
+ extensions.
## A first debug session
-Add `debug=True` to `@kinetic.run()`:
+Add `debug=True` to `@kinetic.run()`. Then call the function directly:
```python
import kinetic
-@kinetic.run(accelerator="tpu-v5e-2x2", debug=True)
+@kinetic.run(accelerator="tpu-v5litepod-4", debug=True)
def train():
import jax
- breakpoint() # debugger will pause here
+ breakpoint() # the debugger pauses here
x = jax.numpy.arange(16)
return x.sum()
@@ -31,38 +42,72 @@ def train():
train()
```
-When you call `train()`, Kinetic:
+When you call `train()`, Kinetic does these things:
:::{container} kinetic-steps
-1. **Schedules the pod** with debugging enabled and an extended **2-hour**
- TTL (vs 10 minutes for normal jobs) so the session has time to
- breathe.
-2. **Pauses execution** just before your function runs and waits for a
- debugger to attach.
-3. **Prints a VS Code `launch.json` snippet** to your terminal — paste it
- into `.vscode/launch.json`.
-4. **Press F5** (Run → Start Debugging) in your editor. The debugger
- attaches and pauses inside Kinetic's runner. Press **F11** to step
- into your function, or **F10** to run straight through to your own
- `breakpoint()`.
+1. **Submits the job.** Kinetic packages the function and creates the
+ job on the cluster with debug mode enabled.
+2. **Waits for the pod.** The pod installs `debugpy`, starts a `debugpy`
+ server on port 5678, and writes a ready marker to Cloud Storage. Kinetic
+ waits for that marker and prints the job status while it waits.
+3. **Opens a tunnel.** Kinetic starts `kubectl port-forward` from
+ `localhost:5678` to the pod.
+4. **Prints a `launch.json` entry.** Paste the entry into
+ `.vscode/launch.json` in your workspace.
:::
-When your function returns, the debugger connection is torn down
-automatically and the pod cleans up.
+Then, in your editor, press **F5** (Run > Start Debugging). The debugger
+attaches and pauses inside the Kinetic runner, one line before the call
+to your function. Press **F11** to step into your function, or press
+**F10** to step over the call. In the second case, the debugger pauses
+at the first `breakpoint()` inside your function.
+
+The blocking call returns the return value of your function when the
+function ends. Kinetic then stops the port-forward process.
:::{tip}
-You don't need to call `breakpoint()` explicitly. Set breakpoints in
-your editor's UI like you would locally — Kinetic pauses before your
-function runs, and UI breakpoints work from there.
+You do not need to call `breakpoint()`. Set breakpoints in the editor as
+you do for a local run. The debugger attaches before your function
+starts, so those breakpoints are active from the first line. If the
+editor marks a breakpoint as unverified, see
+[Path mappings and source files](#path-mappings-and-source-files).
+:::
+
+:::{note}
+A blocking debug call does not stream the pod log to your terminal. To
+read the log during the session, run `kinetic jobs logs --follow`
+in a second terminal.
:::
-## Attaching to a submitted job
+## What `debug=True` changes
+
+`debug=True` changes the pod and the local call in these ways:
+
+- The pod runs `uv pip install --system debugpy` before it starts the
+ `debugpy` server. The image that Kinetic builds contains `uv`. If you
+ use your own image, `uv` must be on the `PATH`. See
+ [Container Images](containers.md).
+- The pod sets `PYTHONBREAKPOINT=debugpy.breakpoint`, so a `breakpoint()`
+ call in your code pauses in the attached debugger.
+- The pod waits up to 10 minutes for a debugger to attach. If no debugger
+ connects in that window, the pod runs your function without a debugger.
+ A `breakpoint()` call does not pause the function when no debugger is
+ attached.
+- Kinetic does not stream the pod log during a blocking debug call.
+- Kinetic does not delete the job resources when the function ends. See
+ [Clean up after a debug session](#clean-up-after-a-debug-session).
-For longer sessions, use `@kinetic.run(debug=True)` and attach later
-from the CLI:
+## Attach later from the command line
+
+For a long session, or when you want to attach from a different machine,
+submit the job as a detached job with `run_async()`. A blocking call blocks
+and returns the value of the function, not a `JobHandle`:
```python
-@kinetic.run(accelerator="tpu-v5e-2x2", debug=True)
+import kinetic
+
+
+@kinetic.run(accelerator="tpu-v5litepod-4", debug=True)
def train():
import jax
@@ -70,22 +115,24 @@ def train():
...
-job = train()
-print(job.job_id)
+job = train.run_async()
+print(job.job_id) # for example: job-a1b2c3d4
```
-Then from a terminal (same machine or a different one with access to
-the same GCP project):
+`run_async()` returns at once. The pod starts, and then waits up to 10
+minutes for a debugger. Attach from a terminal on the same machine, or on
+a different machine with the same active profile:
```bash
kinetic jobs debug
```
-`kinetic jobs debug` blocks until the job finishes or you hit Ctrl+C,
-then tears down the connection. The command fails fast if the job
-wasn't submitted with `debug=True`.
+`kinetic jobs debug` waits for the pod, opens the tunnel, and prints the
+`launch.json` entry. The command blocks until the job ends or until you
+press Ctrl+C. The command then stops the port-forward process. The
+command fails at once if you did not submit the job with `debug=True`.
-You can also drive it from Python:
+You can do the same from Python:
```python
import kinetic
@@ -94,17 +141,23 @@ from kinetic.debug import cleanup_port_forward
job = kinetic.attach("")
pf = job.debug_attach(local_port=5678)
try:
- job.result() # or job.status() in a loop
+ value = job.result()
finally:
cleanup_port_forward(pf)
+job.cleanup()
```
+`job.result()` waits for the job to end and returns the value of the
+function. For a debug job, `result()` does not delete the job resources.
+Call `job.cleanup()` when you no longer need the job resources.
+
## Port conflicts
-The default port is `5678` — debugpy's default, which VS Code's Python
-extension auto-fills in `launch.json`. If something else is already
-bound to `5678` locally, Kinetic raises a `RuntimeError` pointing you
-at a different port:
+The default local port is `5678`. That port is the `debugpy` default, and
+the VS Code Python extension fills it in for an attach configuration. If
+another process listens on `5678` on your machine, `kubectl port-forward`
+exits at once and Kinetic raises a `RuntimeError`. Attach again on a
+different port:
```bash
kinetic jobs debug --port 5679
@@ -116,84 +169,163 @@ Or from Python:
pf = job.debug_attach(local_port=5679)
```
-Remember to update the `port` field in your `launch.json` snippet to
-match.
-
-## Path mappings and source files
-
-Kinetic fills in `pathMappings` in the printed `launch.json` so
-breakpoints set in your local files hit the matching remote files —
-no "unverified breakpoint" warnings, no file mismatch.
+Change the `port` field in the printed `launch.json` entry to the same
+value.
-If you attach from a directory that isn't your project root, pass
-`working_dir=` to `debug_attach()` (or replace `${workspaceFolder}` in
-the printed snippet) so the mapping points at the sources you actually
-have open.
+:::{note}
+A blocking call always uses port `5678`. If the port is in use, the call
+raises `RuntimeError` after Kinetic submits the job. The job stays on the
+cluster and waits for a debugger. Find the job ID in the log output or
+with `kinetic jobs list`. Then attach to the job with
+`kinetic jobs debug --port `.
+:::
-## Timeouts and the attach window
+## Path mappings and source files
-The pod waits up to 10 minutes for a debugger client to attach. If no
-one connects in that window, it proceeds with your function running
-normally — the job does not hang indefinitely. To extend or shorten
-that window, set `KINETIC_DEBUG_WAIT_TIMEOUT` (seconds) in your local
-environment before submitting:
+The debugger must match the source files in your editor with the source
+files on the pod. Kinetic does two things:
+
+- The runner creates a symbolic link on the pod at the absolute path of
+ your **entry directory**, the directory of the file that defines your
+ function. The link points at the root of the extracted source. The
+ runner creates the link only if that path does not exist on the pod.
+ The source files in your entry directory therefore exist on the pod
+ under the same absolute path as on your machine. When the entry
+ directory sits below the package root, the link still points at the
+ root of the extracted source. See
+ [What Ships to the Pod](packaging.md#the-package-root).
+- The printed `launch.json` entry contains a `pathMappings` entry. The
+ `localRoot` value is the entry directory for a blocking call, and
+ `${workspaceFolder}` for `kinetic jobs debug`. The `working_dir=`
+ argument of `debug_attach()` sets that value. The `remoteRoot` value is
+ `/tmp/workspace`. The runner does not use that directory. The runner
+ extracts the source into a temporary directory with the prefix
+ `kinetic-run-`.
+
+If the editor marks a breakpoint as unverified, do one of these things:
+
+- Delete the `pathMappings` entry from the configuration. Then start the
+ debugger again. The symbolic link makes your local paths valid on the
+ pod, so the debugger does not need a mapping.
+- Put a `breakpoint()` call in the function. That call always pauses
+ when a debugger is attached.
+
+## The attach window
+
+Two waits have a fixed limit of 10 minutes each. Kinetic has no option
+to change these limits.
+
+- **On the pod.** After the `debugpy` server is ready, the pod waits up
+ to 10 minutes for a debugger. If no debugger connects, the pod runs the
+ function without a debugger. The job does not wait forever.
+- **On your machine.** `debug_attach()`, and therefore a blocking call
+ and `kinetic jobs debug`, wait up to 10 minutes for the ready marker.
+ This wait includes the time to schedule the pod and to install
+ `debugpy`. On a node pool that is scaled to zero, the wait can end with
+ a `TimeoutError` while the job continues. In that case, run
+ `kinetic jobs debug `. The command waits again for the pod.
+
+The 2-hour value that applies to debug jobs is not the attach window. That
+value is the retention time of the finished Kubernetes Job. See the next
+section.
+
+## Clean up after a debug session
+
+A blocking debug call, and `result()` on a debug job, use
+`cleanup=False`. Kinetic keeps the job resources so that you can inspect
+them after the session:
+
+- **The Kubernetes resource** stays on the cluster after the pod ends. On
+ the GKE backend, Kubernetes deletes the finished Job 2 hours after it
+ ends (10 minutes for a normal job). On the Pathways backend, Kinetic
+ sets no automatic deletion for the LeaderWorkerSet.
+- **The artifacts in the jobs bucket** (`context.zip`, `payload.pkl`,
+ `result.pkl`, `handle.json`) stay until you delete them, or until the
+ 30-day rule of the bucket deletes them.
+
+Delete both when you are done:
```bash
-export KINETIC_DEBUG_WAIT_TIMEOUT=1800 # 30 minutes
+kinetic jobs cleanup
```
+Or from Python, call `job.cleanup()`. To collect the result and delete
+the resources in one step, call `job.result(cleanup=True)`. The command
+`kinetic jobs result ` does the same by default.
+
## Multi-host debugging
-On multi-host TPU slices (Pathways backend), you attach once to the
-leader pod; Kinetic sequences the non-leader workers so the
-distributed runtime doesn't start until you're ready. `jax.process_index()`
-semantics stay predictable, and you don't need to attach to each host
-separately.
+On a multi-host TPU slice, for example `tpu-v5litepod-16`, Kinetic uses
+the Pathways backend and starts one pod per host. Only the leader pod
+runs the `debugpy` server. You attach one time, to the leader. The worker
+pods wait for a marker that the leader writes to Cloud Storage after the
+debugger attaches, or after the attach window ends. The workers therefore
+do not start the distributed runtime while the leader waits for a
+debugger.
+
+The workers wait up to 11 minutes for the marker (the 10-minute attach
+window plus a 1-minute margin). If the marker does not appear in that
+time, for example because the leader pod failed, the workers fail with a
+`RuntimeError` that names the missing marker.
+
+See [Distributed Training](distributed_training.md) for the multi-host
+slice names.
:::{warning}
-**Avoid `spot=True` with `debug=True`.** Preemption mid-session
-terminates the pod, dropping your debug connection. Kinetic warns at
-decoration time if both are set. Use on-demand capacity for interactive
-work.
+Do not use `spot=True` with `debug=True`. If Google Cloud preempts the
+node during the session, the pod stops and the debug connection drops.
+Kinetic warns at decoration time when both are set. Use on-demand
+capacity for interactive work.
:::
## Automated environments
-`@kinetic.run(debug=True)` requires an interactive terminal — if
-`stdin` isn't a TTY (CI, `nohup`, piped input), the local client
-raises `RuntimeError` before submission so your job doesn't silently
-hang waiting for someone to attach.
+A blocking call with `debug=True` needs an interactive terminal. If
+`stdin` is not a TTY (CI, `nohup`, piped input), the call raises
+`RuntimeError`. Kinetic runs this check after it submits the job. The
+submitted job therefore stays on the cluster: the pod waits 10 minutes
+for a debugger and then runs the function.
+
+For automation, use `run_async()`. The detached call does not check the
+terminal. Attach later with `kinetic jobs debug ` from an
+interactive shell, or let the job run without a debugger.
-For async submission there's no TTY requirement —
-`@kinetic.run(debug=True)` works fine in any environment, and
-`kinetic jobs debug` from an interactive shell attaches whenever
-you're ready.
+To permit a blocking debug call without a TTY, set
+`KINETIC_NO_TTY_DEBUG=1` in the environment of the process. See
+[Configuration](../configuration.md).
## Related pages
::::{grid} 1 1 2 2
:gutter: 3
-:::{grid-item-card} {octicon}`history;1em` Detached Jobs
+:::{grid-item-card} {octicon}`clock;1em` Detached Jobs
:link: async_jobs
:link-type: doc
-Pairs with `kinetic jobs debug ` for long debug sessions.
+`run_async()`, `JobHandle`, and the `kinetic jobs` commands that pair with
+`kinetic jobs debug`.
+:::
+
+:::{grid-item-card} {octicon}`server;1em` Distributed Training
+:link: distributed_training
+:link-type: doc
+
+Multi-host TPU slices and the Pathways backend.
:::
:::{grid-item-card} {octicon}`gear;1em` Configuration
:link: ../configuration
:link-type: doc
-`KINETIC_DEBUG_WAIT_TIMEOUT` and the other user-facing environment
-variables.
+`KINETIC_NO_TTY_DEBUG` and the other settings that Kinetic reads.
:::
:::{grid-item-card} {octicon}`bug;1em` Troubleshooting
:link: ../troubleshooting
:link-type: doc
-What to check when a pod doesn't reach `RUNNING` or the debugger fails
-to attach.
+What to check when a pod stays in `PENDING`, or when a job fails before
+it calls your function.
:::
::::
diff --git a/docs/guides/dependencies.md b/docs/guides/dependencies.md
index aa9a5464..786e0aed 100644
--- a/docs/guides/dependencies.md
+++ b/docs/guides/dependencies.md
@@ -1,26 +1,18 @@
-# Managing Dependencies
+# Dependencies
-There are three independent things going on when Kinetic runs your job:
-
-1. **Dependency discovery** — Kinetic figures out which packages your
- project needs by reading `requirements.txt` or `pyproject.toml` from
- your [entry directory](packaging.md#the-package-root) or from a
- directory above it.
-2. **Container mode choice** — those dependencies either get baked into
- a custom image (bundled mode), installed at pod startup (prebuilt
- mode), or ignored entirely (custom image mode). See
- [Execution Modes](execution_modes.md).
-3. **JAX filtering** — accelerator runtime packages (`jax`, `jaxlib`,
- `libtpu`) are filtered out before install so they don't shadow the
- hardware-correct versions in the container.
-
-This page focuses on (1) and (3). (2) lives on its own page:
-[Execution Modes](execution_modes.md).
+Kinetic reads one dependency file for each job: a `requirements.txt` or a
+`pyproject.toml`. Kinetic installs the packages in that file, and your
+function can import them on the pod. This page explains how Kinetic finds
+the file and which lines Kinetic filters. It also explains how to install
+packages from a private index, and lists the pitfalls that cause a missing
+import on the pod. The container image mode decides where the packages
+install: in the image at build time, or in the pod at start. See
+[Container Images](containers.md) for that choice.
## A first run
-Drop a `requirements.txt` next to your script and Kinetic picks it up
-automatically:
+Put a `requirements.txt` next to your script. Kinetic finds the file
+without configuration:
```text
# requirements.txt
@@ -30,44 +22,49 @@ pandas
```
```python
-@kinetic.run(accelerator="tpu-v6e-8")
+@kinetic.run(accelerator="tpu-v5litepod-4")
def train():
- import pandas as pd # installed automatically on the remote
+ import pandas as pd # installed on the pod
...
```
-`pyproject.toml` works equally well — Kinetic reads
-`[project.dependencies]`. If both files exist, `requirements.txt` wins.
+A `pyproject.toml` works too. Kinetic reads the `[project.dependencies]`
+list from that file. If both files are in the same directory, Kinetic
+uses `requirements.txt`.
:::{tip}
**Recommended defaults:**
-- Pin only the libraries you actually depend on. The fewer packages, the
- faster your image builds (or your prebuilt-mode pod start).
-- Don't pin `jax`, `jaxlib`, `libtpu`, or any other accelerator runtime
- — Kinetic filters them out and uses the version in the container.
-- Use a `pyproject.toml` if you already have one for local development
- rather than maintaining a separate `requirements.txt`.
+- List only the packages that your function imports. A short list makes
+ the image build faster.
+- Do not pin `jax`, `jaxlib`, `libtpu`, or `libtpu-nightly`. Kinetic
+ filters those lines and installs the JAX version that matches the
+ accelerator.
+- If you already have a `pyproject.toml` for local development, use that
+ file. You do not need a separate `requirements.txt`.
:::
## How discovery works
-When you call a decorated function, Kinetic starts at the entry
-directory. Kinetic then walks **up** one directory at a time. At each
-directory, Kinetic does these steps:
+When you call a decorated function, Kinetic starts at the
+[entry directory](packaging.md#the-package-root): the directory of the
+module file that defines the function. Kinetic then walks **up** one
+directory at a time. At each directory, Kinetic does these steps:
-1. If a file with the name `requirements.txt` is there, use that file.
-2. If not, and a file with the name `pyproject.toml` is there, read
- `[project.dependencies]` from that file.
-3. If not, walk up one directory and do the steps again.
+1. If a file with the name `requirements.txt` is in the directory,
+ Kinetic uses that file.
+2. If not, and a file with the name `pyproject.toml` is in the directory,
+ Kinetic uses that file and reads `[project.dependencies]` from it.
+3. If neither file is in the directory, Kinetic moves up one directory
+ and does the steps again.
Kinetic examines files only. If a directory has the name
`requirements.txt` or `pyproject.toml`, Kinetic ignores that directory
and continues the walk.
-The walk has bounds, so that it cannot leave your project and take a
-foreign file:
+The walk has bounds, so that Kinetic does not use a file from outside
+your project:
- Kinetic stops after it examines the first directory that holds a `.git`
entry. The entry can be a directory or a file, because a git worktree
@@ -76,90 +73,110 @@ foreign file:
- Kinetic stops at your home directory and at the root of the file
system. Kinetic examines these directories too.
-If Kinetic finds no dependency file, the job gets only the packages of the
-base image.
+If Kinetic finds no dependency file, the pod gets only the packages that
+Kinetic installs in every image: JAX, Keras, `cloudpickle`,
+`google-cloud-storage`, and Kinetic.
-Kinetic writes the name of the file that it selected into the log
-(`Using dependency file: ...`) on each submit. If the installed packages
-are not the packages that you expected, read this log line first. If both
-files are in the selected directory, the log also names the file that
-Kinetic selected. Kinetic also logs a warning when it finds no `.git`
-marker and the selected file is not in your entry directory. Kinetic thus
-never takes a foreign file without a message.
+Kinetic writes the name of the selected file into the log on each submit:
+`Using dependency file: ...`. If the installed packages are not the
+packages that you expected, read this log line first. If both files are
+in the selected directory, the log also names the file that Kinetic
+selected. Kinetic logs a warning when it finds no `.git` entry and the
+selected file is not in your entry directory. The log therefore always
+shows when Kinetic uses a file from outside your entry directory.
:::{warning}
-A `pyproject.toml` file with no `[project.dependencies]` table still stops
-the walk. If your dependencies are in a `requirements.txt` file above it,
-Kinetic does not find them. Kinetic logs a warning that the selected file
-declares no dependencies, and that the job will get only the base-image
-packages. Move or copy your dependency file into your entry directory.
-
-Kinetic reads `[project.dependencies]` only. Kinetic logs a warning when
-it finds dependencies in `[tool.poetry.dependencies]`, in
-`[project.optional-dependencies]`, in `[dependency-groups]`, or under
-`dynamic = ["dependencies"]`.
+A `pyproject.toml` file with no `[project.dependencies]` list still stops
+the walk. If your dependencies are in a `requirements.txt` file above
+that directory, Kinetic does not find them. Kinetic logs a warning that
+the selected file declares no dependencies, and that the pod gets only
+the packages that Kinetic installs in every image. Move or copy your
+dependency file into your entry directory.
+
+Kinetic reads `[project.dependencies]` only. If that list is empty or
+absent, Kinetic looks for dependencies in `[tool.poetry.dependencies]`,
+`[project.optional-dependencies]`, `[dependency-groups]`, and
+`dynamic = ["dependencies"]`. If Kinetic finds any of those, Kinetic logs
+a warning that names the tables. If `[project.dependencies]` is not
+empty, Kinetic installs that list and ignores the other tables without a
+warning.
:::
Kinetic ships only the [package root](packaging.md) in `context.zip`.
-Kinetic can thus select a dependency file above the package root. Kinetic
-installs the packages from that file, but the directory of that file is
-not on the pod.
-
-The image build and the pod see the generated requirements file alone.
-A line that points to a path on your machine (`-r base.txt`, `-e .`, or
-`./local-wheel`) thus cannot resolve. In prebuilt mode, Kinetic refuses
-such a line at submit time and names the line that it refuses. In bundled
-mode, Kinetic does not make this check, and the install fails inside Cloud
-Build instead.
-
-In bundled mode, the discovered file is hashed and used as part of the
-image cache key — change the file, and the next run rebuilds. In
-prebuilt mode, the same file is uploaded and installed at pod startup.
-In custom image mode, the file is ignored entirely.
+Kinetic can therefore select a dependency file above the package root.
+Kinetic installs the packages from that file, but the directory of that
+file is not on the pod.
+
+## What Kinetic does with the file
+
+Kinetic logs the selected file. Kinetic then filters the content (see
+[JAX and accelerator runtimes](#jax-and-accelerator-runtimes)), hashes
+the filtered content into the image tag, and writes the filtered content
+to a generated `requirements.txt` for Cloud Build. A change to the
+filtered content causes a new image build on the next run. The build
+sees the generated file alone. A line that points to a path on your
+machine (`-r base.txt`, `-e .`, or `./local-wheel`) therefore cannot
+resolve, and the install fails inside Cloud Build.
+
+The other [container image modes](containers.md) change where the
+install happens. The prebuilt mode installs a generated file on the pod
+at start, and rejects a local path line at submit time. A custom image
+ignores the dependency file.
## JAX and accelerator runtimes
-Kinetic's bundled and prebuilt images already have `jax`, `jaxlib`, and
-the right accelerator backend (`libtpu` on TPU, CUDA libs on GPU)
-installed and pinned to versions that match the container. To prevent
-your `requirements.txt` from clobbering that, Kinetic strips these
-entries before install:
+The image that Kinetic builds already contains `jax`, `jaxlib`, and the
+runtime for the accelerator category: `libtpu` on TPU, or the CUDA
+libraries on GPU. To prevent your dependency file from replacing that
+installation, Kinetic removes these entries before the install:
- `jax`
- `jaxlib`
- `libtpu`
- `libtpu-nightly`
-If you have a specific reason to override the in-container JAX —
-testing a new release, reproducing a bug — append `# kn:keep` to the
-line:
+Kinetic logs a warning for each removed line. The warning names the
+package and tells you how to keep the line.
+
+If you must override the JAX version, for example to test a new release,
+append `# kn:keep` to the line:
```text
jax==0.4.25 # kn:keep
jaxlib==0.4.25 # kn:keep
```
-This works in `requirements.txt`. Use it sparingly; getting JAX +
-`jaxlib` + accelerator runtime versions to line up by hand is a known
-source of obscure crashes.
+The marker works in `requirements.txt`. Use the marker with care. A
+mismatch between `jax`, `jaxlib`, and the accelerator runtime is a common
+cause of crashes that are hard to diagnose.
+
+:::{note}
+Kinetic filters physical lines for the image build. If a filtered `jax`
+entry continues onto more lines with a backslash, the continuation lines
+stay in the file and can break the install. The output of
+`pip-compile --generate-hashes` has such lines. Keep each entry on one
+line.
+:::
## Private packages
-Bundled-mode builds install your dependencies inside Cloud Build. Cloud
-Build does not inherit your local `pip.conf`, environment variables, or
-shell credentials, so anything the installer needs in order to find or
-authenticate to a private index has to be present in the project source
-that gets uploaded to the build.
+Cloud Build installs your dependencies into the image. Cloud Build
+never receives your project source. The build context holds exactly three
+things: the generated Dockerfile, the Kinetic runner script, and the
+generated `requirements.txt`. A `pip.conf`, an environment variable, or a
+credential on your machine or in your project tree does not reach the
+build. The only way to tell the installer about a private index is a line
+in the dependency file itself.
-You have two practical options:
+You have two options:
::::{tab-set}
-:::{tab-item} Bundled mode with an index URL
+:::{tab-item} An index URL in requirements.txt
-Add `--index-url` or `--extra-index-url` as a line in
-`requirements.txt`. The installer reads these directives and uses them
-when resolving every package in the file:
+Add `--index-url` or `--extra-index-url` as a line in `requirements.txt`.
+The installer reads these directives and uses them for every package in
+the file:
```text
--extra-index-url https://my-org-private-index.example.com/simple
@@ -167,59 +184,66 @@ my-private-package==1.2.3
some-public-dep==2.0.0
```
-This works without extra setup if the index is publicly reachable
-(no auth required), or if it sits behind network ACLs that the Cloud
-Build pool already satisfies (for example, a GCP-internal Artifact
-Registry repo that the build service account has read access to).
+This option does not work in a `pyproject.toml`, because `[project.dependencies]` holds
+package specifiers only. This option needs an index that requires no
+credentials. Examples: a public index, or an index that the Cloud Build
+worker and the pod can reach through network rules alone.
:::
-:::{tab-item} Custom image mode
+:::{tab-item} A custom image
If your private packages need credentials at install time, system
-libraries, or unusual build flags, prebuild a container image with
-them installed and pass it as `container_image=""`.
-This gives you full control over the build environment, including
-`pip.conf`, secret mounts, and `gcloud` authentication. See
-[Container Images](containers.md).
+libraries, or special build flags, build a container image that contains
+those packages. Pass the image as `container_image=""`.
+You control the build environment: `pip.conf`, secret mounts, and
+`gcloud` authentication. See
+[Container Images](containers.md#custom-image-mode).
:::
::::
:::{warning}
-Avoid embedding secrets in `requirements.txt`
-(`https://user:token@host/...`); the file is uploaded to GCS and used
-as part of the build context, so any credentials it contains will end
-up in build logs and cached artifacts.
+Do not put a secret into `requirements.txt`, for example
+`https://user:token@host/...`. Kinetic uploads the generated
+requirements file to the builds bucket as part of the Cloud Build source
+(or to the jobs bucket in the prebuilt mode). Anyone with read access to
+those buckets or to the build can read the token.
:::
## Common dependency pitfalls
-- **Pinning `jax` without `# kn:keep`** — the pin is silently dropped
- and you get the in-container version anyway. If you actually want a
- pin, use `# kn:keep`. If you don't, drop the line.
-- **Listing TensorFlow alongside JAX** — both ship their own copy of
- the accelerator runtime. They can co-exist, but on TPU you typically
- want only one. If `tf.data` is the only thing you need from
- TensorFlow, `tensorflow-cpu` is enough and won't fight with `libtpu`.
-- **Forgetting to add a new package locally** — Kinetic only sees what's
- in `requirements.txt` or `pyproject.toml`. A `pip install` in your
- shell that isn't reflected in those files won't carry over.
-- **Massive dependency sets** — every `requirements.txt` change forces
- a bundled rebuild. If your deps churn daily, consider prebuilt mode
- (after publishing a base image with `kinetic build-image`).
-- **Editable installs (`pip install -e`)** — an editable install does not
- show in `requirements.txt`, and Kinetic cannot carry it over. Keep the
- source inside the [package root](packaging.md), which Kinetic packages
- for you. As an alternative, publish the package and pin a released
- version. `pip` cannot install a line such as `-e .` on the pod. `pip`
- installs a line such as `-e git+https://example.com/pkg` correctly,
+- **A `jax` pin without `# kn:keep`.** Kinetic drops the line, logs a
+ warning, and installs the JAX version of the image. If you want the
+ pin, add `# kn:keep`. If you do not want the pin, delete the line.
+- **TensorFlow next to JAX.** The `tensorflow` package can also try to
+ use the TPU, and JAX then cannot open the TPU. If you need TensorFlow
+ for `tf.data` only, install `tensorflow-cpu`. That package does not use
+ the TPU.
+- **A package that you installed locally but did not list.** Kinetic
+ reads `requirements.txt` or `pyproject.toml` only. A package that you
+ installed with `pip install` in your shell, but did not list in one of
+ those files, is not on the pod.
+- **Extras in `pyproject.toml`.** Kinetic reads `[project.dependencies]`
+ only. Packages in `[project.optional-dependencies]` or in
+ `[dependency-groups]` do not install, and Kinetic logs no warning when
+ `[project.dependencies]` is not empty. Move those packages into
+ `[project.dependencies]`, or use a `requirements.txt`.
+- **A large dependency set that changes often.** A change to the
+ dependency file causes a new image build. If your dependencies change
+ many times a day, see [Container Images](containers.md) for a mode that
+ installs at pod start.
+- **An editable install (`pip install -e`).** An editable install does not
+ appear in `requirements.txt`, and Kinetic cannot install it on the pod.
+ Keep the source inside the [package root](packaging.md), which Kinetic
+ ships for you. As an alternative, publish the package and pin a released
+ version. A line such as `-e .` cannot install in the image or on the
+ pod. Kinetic accepts a line such as `-e git+https://example.com/pkg`,
because that line names a remote source.
-- **Local path references** — the lines `-r other.txt`,
+- **A local path reference.** Examples: `-r other.txt`,
`-c constraints.txt`, `./wheels/foo.whl`, `file://...`, and
- `mypkg @ ./vendor` point to paths that do not exist in the image or on
- the pod. In prebuilt mode, Kinetic refuses these lines at submit time
- and names the line. In bundled mode, the install fails inside Cloud
- Build.
+ `mypkg @ ./vendor`. These lines point to paths that do not exist in the
+ image or on the pod. The install fails inside Cloud Build. The prebuilt
+ mode rejects these lines at submit time and names the line.
## Related pages
@@ -234,24 +258,18 @@ The package root, and why the dependency file can sit outside the
archive.
:::
-:::{grid-item-card} {octicon}`zap;1em` Execution Modes
-:link: execution_modes
-:link-type: doc
-
-Where the discovered deps go.
-:::
-
:::{grid-item-card} {octicon}`stack;1em` Container Images
:link: containers
:link-type: doc
-Custom image and base-image workflows.
+Where the packages install: in the image at build time, in the pod at
+start, or in an image that you build.
:::
:::{grid-item-card} {octicon}`bug;1em` Troubleshooting
:link: ../troubleshooting
:link-type: doc
-What to check when an import fails on the remote.
+What to check when an import fails on the pod.
:::
::::
diff --git a/docs/guides/distributed_training.md b/docs/guides/distributed_training.md
index 48277ea4..134c990c 100644
--- a/docs/guides/distributed_training.md
+++ b/docs/guides/distributed_training.md
@@ -1,21 +1,54 @@
# Distributed Training
-:::{admonition} Who this is for
-:class: note
-
-Users whose model or batch size has outgrown a
-single TPU host. Kinetic + the [Pathways](https://cloud.google.com/ai-hypercomputer/docs/workloads/pathways-on-cloud/pathways-intro)
-backend lets you treat a multi-host TPU slice as one logical machine,
-without writing your own multi-process JAX coordination.
-:::
-
-For single-host slices (everything that fits on one TPU node like
-`tpu-v5litepod-8`), you don't need this page — your existing JAX or
-Keras code already uses every chip on the node.
+This page covers jobs on TPU slices that span more than one host. Read
+this page if your model or your batch no longer fits on one TPU host.
+A multi-host job runs one pod per host, and Kinetic runs it on the
+`pathways` backend. This page explains what changes: how Kinetic
+starts the job, which log you see, what the return value is, and how a
+multi-host job fails.
+
+If your slice fits on one host, for example `tpu-v5litepod-8`, you do
+not need this page. Your JAX or Keras code already uses every chip on
+that host.
+
+## Single-host or multi-host?
+
+The **Hosts** column of the TPU table on the [Accelerators](../accelerators.md#tpus)
+page decides. A slice with one host is single-host. A slice with more
+than one host is multi-host. Some examples:
+
+| Accelerator | Topology | Hosts | Backend |
+| ----------- | -------- | ----- | ------- |
+| `tpu-v5litepod-8` | 2x4, one 8-chip VM | 1 | `gke` |
+| `tpu-v6e-8` | 2x4, two 4-chip VMs | 2 | `pathways` |
+| `tpu-v5litepod-16` | 4x4, four 4-chip VMs | 4 | `pathways` |
+| `tpu-v6e-16` | 4x4, four 4-chip VMs | 4 | `pathways` |
+| `tpu-v5p-8` | 2x2x2, two 4-chip VMs | 2 | `pathways` |
+
+The chip count alone does not tell you the host count. `tpu-v5litepod-8`
+and `tpu-v6e-8` both have 8 chips, but the v5e slice is one 8-chip VM
+and the v6e slice is two 4-chip VMs. Every v6e slice and every v5p
+slice is multi-host. Run `kinetic accelerators` to print the names, chip
+counts, and topologies that your installed version accepts. That command
+does not print the host count, so use the table on the Accelerators page.
+
+Kinetic selects the backend from the host count. If the slice has more
+than one host, Kinetic selects `backend="pathways"`. You do not set
+`backend` yourself.
+
+## Before you start
+
+- The cluster needs a node pool for the multi-host accelerator, for
+ example `kinetic pool add --accelerator tpu-v6e-16`. See
+ [Clusters and Node Pools](clusters.md).
+- The cluster needs the LeaderWorkerSet controller. `kinetic up`
+ installs it. If a submit fails with `LeaderWorkerSet CRD not found`,
+ run `kinetic up` on that cluster.
## A first multi-host run
-Pick a multi-host accelerator:
+This example selects a multi-host accelerator and prints the process
+layout:
```python
import kinetic
@@ -28,25 +61,123 @@ def train_distributed():
print(f"Total devices across all hosts: {jax.device_count()}")
print(f"This host: {jax.process_index()} of {jax.process_count()}")
# ... your training code ...
+ return jax.device_count()
+
+
+print(train_distributed())
```
-Whether a slice is multi-host depends on the topology and the per-VM
-chip count, not on the accelerator string alone. For example,
-`tpu-v5litepod-2x2` (4 chips on one VM) and `tpu-v5litepod-2x4` (8 chips on one
-VM via `ct5lp-hightpu-8t`) are both single-host, while `tpu-v5litepod-16`
-(4×4 across four 4-chip VMs) and `tpu-v6e-16` (4×4 across four 4-chip VMs)
-are multi-host. See [Accelerators](../accelerators.md) for the full
-topology table.
-
-Pathways is **auto-selected** for multi-host slices — Kinetic resolves
-`backend="pathways"` whenever the accelerator's topology spans more than
-one node, so the example above doesn't need to set it explicitly. You
-only need to pass `backend="pathways"` yourself if you want to develop
-against the Pathways code path on a single-host slice — handy for
-shortening the iteration loop before you scale up.
+On `tpu-v6e-16`, `jax.process_count()` is 4, `jax.local_device_count()`
+is 4, and `jax.device_count()` is 16. Your terminal shows the log of
+the leader host only, so you see one `This host: 0 of 4` line.
+
+## How a multi-host job runs
+
+Kinetic creates one Kubernetes LeaderWorkerSet for the job. The set has
+one leader pod and one worker pod for each other host. Every pod runs
+the same command, downloads the same artifacts, and calls your function.
+Kinetic sets these environment variables on every pod:
+`MEGASCALE_COORDINATOR_ADDRESS`, `MEGASCALE_NUM_SLICES`,
+`TPU_WORKER_ID`, `JAX_PLATFORMS`, and `KERAS_BACKEND=jax`. Each pod
+runs one JAX process, and `jax.process_count()` equals the host count.
+
+Your code does the cross-host communication. JAX collectives
+(`jax.lax.psum`, sharding, `pmap`) and the Keras distribution API
+handle it. Kinetic and GKE do these things:
+
+- The LeaderWorkerSet creates one pod per host, and the GKE autoscaler
+ starts every VM of the slice for the job.
+- Kinetic streams the log of the **leader pod** to your terminal.
+ Kinetic does not stream or interleave the logs of the other pods. See
+ [Debugging distributed jobs](#debugging-distributed-jobs) for how to
+ read them.
+- Kinetic treats a failure on any host as a failure of the job.
+
+:::{note}
+The name `pathways` refers to the Kinetic backend for multi-host TPU
+jobs. Kinetic does not deploy Pathways-on-Cloud proxy servers or a
+resource manager. Each host runs its own JAX process, and
+`jax.process_count()` is greater than 1.
+:::
+
+:::{warning}
+If your model and your batch fit on one host, stay on one host. A
+multi-host job starts more slowly, needs the LeaderWorkerSet
+controller, and fails as a whole if one host fails. Move to
+multi-host only if one host is not enough.
+:::
+
+## The return value
+
+Every host uploads its return value to the same `result.pkl` object in
+the jobs bucket. Kinetic does not select the value of process 0. The
+value that Kinetic returns to you is the value that the last host wrote.
+
+Write your function so that this rule does not matter:
+
+- Return the same small value from every host, for example a metric that
+ every host computes.
+- Or return a value from process 0 only and return `None` from the other
+ hosts. Then check the returned value on the client, because the client
+ can receive the `None`.
+
+```python
+@kinetic.run(accelerator="tpu-v6e-16")
+def train():
+ import jax
+
+ final_loss = ... # every host holds the same value after the collective
+ return float(final_loss) # same value on every host
+```
+
+Write large outputs, such as model weights and checkpoints, under
+`KINETIC_OUTPUT_DIR` from process 0. See
+[Outputs and Checkpoints](checkpointing.md).
+
+## Failures and the exception that you see
+
+A failure on any host fails the job. What you see on the client depends
+on which pod wrote `result.pkl` last:
+
+- If the failed pod wrote last, Kinetic raises the remote exception on
+ the client with the traceback of that host attached.
+- If the leader finished and uploaded a success payload, but a worker
+ pod failed, Kinetic does not return the leader value. Kinetic raises a
+ `RuntimeError`. The message names the failed pods, their exit codes,
+ and the last 30 log lines of each failed pod. Kinetic keeps the job
+ artifacts in the jobs bucket for inspection.
+
+In both cases, a blocking call and `JobHandle.result()` with the default
+`cleanup` then delete the LeaderWorkerSet and its pods. Read the pod
+logs from Cloud Logging after that point.
+
+For a detached job, `JobHandle.status()` reports `FAILED` if any
+worker pod failed, also if the leader pod finished. See
+[Detached Jobs](async_jobs.md).
+
+## Test the multi-host code path on one host
+
+Set `backend="pathways"` on a single-host accelerator to run the same
+LeaderWorkerSet code path with one pod. Use this run to test a script
+before you request a large slice.
+
+```python
+@kinetic.run(accelerator="tpu-v5litepod-8", backend="pathways")
+def smoke_test():
+ import jax
+
+ return jax.process_count() # 1 on a single-host slice
+```
+
+On one host, `jax.process_count()` is 1, so this run does not test the
+cross-host collectives. The run tests the job flow, the image, and your
+code up to the first collective.
## Data parallelism with Keras
+The Keras distribution API sees every device on every host. Build a
+device mesh over all devices and set a `DataParallel` distribution:
+
```python
@kinetic.run(accelerator="tpu-v6e-16")
def train_data_parallel():
@@ -163,19 +294,21 @@ ones, with what to actually do:
:::{admonition} Recommended checkpoint frequency
:class: tip
-For any multi-host run, write a
-checkpoint at least every 10 minutes of wall time. The base rate of
-preemption, quota issues, and slice-wide failures is high enough that
-unbounded loss windows are not worth the throughput. See
-[Checkpointing](checkpointing.md) for the API.
+Write checkpoints at short intervals in a multi-host run, for example
+every 10 minutes of wall time. Preemption, quota problems, and
+slice-wide failures are frequent enough that a long gap between
+checkpoints costs more than the checkpoint itself. The default output
+directory is per job, so pass an explicit `output_dir=` if you resume
+a run from an earlier job. See
+[Outputs and Checkpoints](checkpointing.md).
:::
## Debugging distributed jobs
-`kinetic jobs logs ` (and `--follow` while the job is running)
-returns the **leader pod's** stdout, which is what `print()` calls on
-process index 0 produce. To gate output to that one process, guard
-print statements with `jax.process_index()`:
+`kinetic jobs logs ` returns the log of the **leader pod**. Add
+`--follow` while the job runs, or `--tail N` for the last `N` lines.
+Every host runs your `print()` calls. Guard the calls that you want to
+see one time with `jax.process_index()`:
```python
import jax
@@ -184,22 +317,41 @@ if jax.process_index() == 0:
print(f"epoch {epoch}: loss={loss}")
```
-For non-leader hosts, fetch logs directly from the per-host pods.
-`kubectl get pods -n | grep ` lists every pod in
-the slice; `kubectl logs ` then returns that host's stdout.
-Cloud Logging in the GCP Console offers the same view through a UI
-filter on the job name.
+Kinetic does not stream the logs of the other hosts. Read them with
+`kubectl`. Every pod of the job carries the label `job-id=`:
+
+```bash
+kubectl get pods -n -l job-id=
+kubectl logs -n
+
+```
+
+`` is the namespace of the active profile, `default` unless
+you changed it. The leader pod is `keras-pathways--0`. The
+worker pods have the same prefix and a further index suffix.
+
+A blocking call and `JobHandle.result()` with the default `cleanup`
+delete the pods after the job ends. Read the pod logs while the job
+runs, or use a detached job and read the logs before you call
+`result()`. Cloud Logging in the Cloud Console keeps the same logs after
+Kinetic deletes the pods. Filter on the pod name.
If a job fails on any host, Kinetic catches the exception and raises it
-locally. The local error gives the stack trace and the process index of
-that host. Usually you do not need the logs of the other pods.
+locally with the stack trace and the process index of that host. In addition,
+the failure message contains the last 30 log lines of each failed pod. Read
+those lines first.
+
+Usually you do not need the full logs of the other pods unless:
-Read the pod logs in two cases:
+* The local error is a collective timeout. All hosts report this error,
+so it does not tell you which host is at fault.
+* The local error indicates that Kubernetes stopped a pod before the pod
+could report a failure.
+
+`debug=True` attaches a debugger to the leader pod. Kinetic holds the
+worker pods until the leader is ready. See
+[Interactive Debugging](https://www.google.com/search?q=debugging.md%23multi-host-debugging).
-- The local error is a collective timeout. All hosts report this error,
- so it does not tell you which host is at fault.
-- The local error tells you that Kubernetes stopped a pod before the pod
- could report a failure.
## Related pages
@@ -210,20 +362,27 @@ Read the pod logs in two cases:
:link: ../accelerators
:link-type: doc
-Slice topologies and naming.
+The TPU table with the Hosts column that decides the backend.
:::
-:::{grid-item-card} {octicon}`history;1em` Checkpointing
+:::{grid-item-card} {octicon}`history;1em` Outputs and Checkpoints
:link: checkpointing
:link-type: doc
-Frequent checkpoints are essential here.
+`KINETIC_OUTPUT_DIR` and resumable training for long runs.
:::
-:::{grid-item-card} {octicon}`server;1em` Multiple Clusters
+:::{grid-item-card} {octicon}`server;1em` Clusters and Node Pools
:link: clusters
:link-type: doc
-When to isolate multi-host TPUs from the rest of your workloads.
+Add a node pool for a multi-host slice, or isolate it in its own cluster.
+:::
+
+:::{grid-item-card} {octicon}`clock;1em` Detached Jobs
+:link: async_jobs
+:link-type: doc
+
+Submit a long multi-host job and collect the result later.
:::
::::
diff --git a/docs/guides/env_vars.md b/docs/guides/env_vars.md
index 597089e1..c6afbdfc 100644
--- a/docs/guides/env_vars.md
+++ b/docs/guides/env_vars.md
@@ -1,94 +1,93 @@
# Forward Environment Variables
-Kinetic allows you to propagate local environment variables to the remote worker environment. This is useful for passing API keys, configuration, or credentials without hardcoding them in your script.
+The pod does not see the environment variables of your shell. When your
+function needs a value from your local environment — an API key, a
+Kaggle credential, a configuration flag — list the variable name in
+`capture_env_vars`. Kinetic copies the value into the pod at submit
+time. This page covers that parameter, its wildcard rules, and how to
+handle secrets.
-## Forwarding Variables
+This page is not about the `KINETIC_*` variables that select the
+project, the zone, the cluster, or the namespace. See
+[Configuration](../configuration.md) for those.
-Use the `capture_env_vars` parameter in the `@kinetic.run()` decorator. It accepts a list of environment variable names or wildcard patterns.
+## Forward a variable
+
+Pass a list of names or patterns to `capture_env_vars`:
```python
import kinetic
@kinetic.run(
- accelerator="tpu-v5litepod-1",
+ accelerator="tpu-v5litepod-4",
capture_env_vars=["KAGGLE_USERNAME", "KAGGLE_KEY", "WANDB_*"],
)
def train_model():
import os
- # These are available in the remote process
- user = os.environ.get("KAGGLE_USERNAME")
- # ...
+ user = os.environ.get("KAGGLE_USERNAME") # the value from your shell
+ ...
```
-## Wildcard Support
+Kinetic reads the values when you call the function, stores them in the
+job payload, and sets them in the pod before your function runs. A
+captured value replaces a variable with the same name in the image.
+
+## Wildcards
+
+A name that ends with `*` is a prefix pattern:
-You can use the `*` suffix to capture all environment variables that start with a specific prefix.
+- `capture_env_vars=["GOOGLE_CLOUD_*"]` captures `GOOGLE_CLOUD_PROJECT`,
+ `GOOGLE_CLOUD_REGION`, and every other name with that prefix.
+- `capture_env_vars=["*"]` captures your full local environment, except
+ the blocklist below. We do not recommend this pattern, because it also
+ captures every secret in your shell.
-- `capture_env_vars=["GOOGLE_CLOUD_*"]`: Captures `GOOGLE_CLOUD_PROJECT`, `GOOGLE_CLOUD_REGION`, etc.
-- `capture_env_vars=["*"]`: (Not recommended) Captures your full local environment, but not the names in the blocklist below.
+A `*` in any other position is not a wildcard.
-### Blocklist for wildcard matches
+### Names that a wildcard never matches
-Some variables describe your local machine, not your job. The pod applies
-the captured values over its own environment. A local `PATH` or
-`LD_LIBRARY_PATH` on the pod points to directories that do not exist
-there, and the job then fails before your code starts. A wildcard pattern
-thus never matches these names:
+Some variables describe your machine, not your job. A local `PATH` or
+`LD_LIBRARY_PATH` points to directories that do not exist in the pod,
+and the job fails before your code starts. A wildcard pattern therefore
+never matches these names:
`PATH`, `HOME`, `PYTHONPATH`, `LD_LIBRARY_PATH`, `LD_PRELOAD`,
`VIRTUAL_ENV`, `CONDA_PREFIX`, `CONDA_DEFAULT_ENV`, `SHELL`, `TMPDIR`,
`TEMP`, `TMP`, `HOSTNAME`, `USER`, `LOGNAME`, `SSH_AUTH_SOCK`,
`KUBERNETES_SERVICE_HOST`, `KERAS_BACKEND`
-Kinetic filters only the wildcard expansion. To forward one of these
-names, list the name exactly in `capture_env_vars`. For example,
-`capture_env_vars=["KERAS_BACKEND"]` replaces the default backend of the
-image. Kinetic writes the names that a wildcard did not capture into the
-log.
-
-## Secure Handling
-
-Kinetic serializes the values of the requested environment variables and sends them to the remote worker as part of the job payload. Make sure that you forward only the variables that the job needs.
-
-Kinetic writes the **names** that it captured into the log on each
-submit. Kinetic never writes the values.
+The filter applies to wildcard matches only. To forward one of these
+names, list the name exactly. For example,
+`capture_env_vars=["KERAS_BACKEND"]` replaces the default Keras backend
+of the image. Kinetic logs the names that a wildcard skipped.
-Kinetic also logs a warning when a captured name contains `TOKEN`,
-`SECRET`, `KEY`, `PASSWORD`, or `CREDENTIAL`. The comparison ignores the
-letter case. This warning is informational, and Kinetic shows it also for
-a name that you listed exactly. If you intend to forward the credential,
-you need no action.
+## Secrets
-The warning tells you where the value goes. Kinetic writes the value into
-`payload.pkl` in the job bucket, which each job pod in the cluster can
-read. Kinetic deletes the artifacts of a job when it collects a usable
-result from that job. If Kinetic does not collect a result, the artifacts
-stay until the lifecycle rule of the bucket deletes them. That rule
-deletes objects after 30 days on a bucket that `kinetic up` created.
+Kinetic stores the captured values in `payload.pkl` in the jobs bucket.
+Every job pod in the cluster can read that bucket. Kinetic deletes the
+payload when it collects a usable result with the default cleanup. In
+three cases the payload stays until you call `cleanup()` or until the
+30-day lifecycle rule of the bucket deletes it: Kinetic collects no
+result, you pass `cleanup=False`, or you use `debug=True`. Because of this:
-Use short-lived tokens for the values that you forward. For more
-information, see [Security](../security.md).
+- Forward only the variables that the job needs.
+- Use short-lived tokens where you can.
-## Precedence
+Kinetic logs the **names** that it captured on each submit. Kinetic never
+logs the values. Kinetic also logs a warning when a captured name contains
+`TOKEN`, `SECRET`, `KEY`, `PASSWORD`, or `CREDENTIAL`, in any letter case.
+The warning is informational and appears also for a name that you listed
+exactly. If you intend to forward the credential, no action is needed.
+See [Security](../security.md).
-Environment variables set via `capture_env_vars` will override any existing variables with the same name in the remote container's base environment.
+## Variables that Kinetic sets in the pod
-## Canonical Environment Variables
-
-Kinetic automatically sets some environment variables in the remote worker environment:
-
-- `KINETIC_OUTPUT_DIR`: The path to the directory where outputs should be saved. By default, this is a GCS path pointing to `gs://{bucket_name}/outputs/{job_id}`. This is useful for passing to checkpointing libraries like Orbax.
-
-:::{important}
-By default, Kinetic imposes a 30-day TTL (Time to Live) on the
-GCS buckets it creates. This means anything written to the default
-`KINETIC_OUTPUT_DIR` will be automatically deleted after 30 days. If you need
-to preserve outputs longer, you should copy them to a bucket without a
-lifecycle rule or specify a custom `output_dir` pointing to a different
-location.
-:::
+Kinetic sets `KINETIC_OUTPUT_DIR` in every pod. The value is a Cloud
+Storage location, by default `gs://{jobs bucket}/outputs/{job_id}`. Write
+every file that you want to keep under that location. See
+[Outputs and Checkpoints](checkpointing.md).
## Related pages
@@ -99,22 +98,20 @@ location.
:link: ../configuration
:link-type: doc
-Full list of `KINETIC_*` variables and precedence rules.
+The `KINETIC_*` variables that configure Kinetic itself.
:::
-:::{grid-item-card} {octicon}`history;1em` Checkpointing
-:link: checkpointing
+:::{grid-item-card} {octicon}`shield;1em` Security
+:link: ../security
:link-type: doc
-How `KINETIC_OUTPUT_DIR` fits into the durable-output story.
+The trust model, and where the payload lives.
:::
:::{grid-item-card} {octicon}`key;1em` LLM Fine-tuning
:link: ../examples/llm_finetuning
:link-type: doc
-`capture_env_vars` is the canonical way to forward Kaggle and other
-model-hub credentials.
+`capture_env_vars` for Kaggle and other model-hub credentials.
:::
::::
-
diff --git a/docs/guides/execution_modes.md b/docs/guides/execution_modes.md
deleted file mode 100644
index fc85a3ee..00000000
--- a/docs/guides/execution_modes.md
+++ /dev/null
@@ -1,236 +0,0 @@
-# Execution Modes
-
-Kinetic ships three ways of producing the container that runs your job. The
-mode you pick controls how long the first run takes, how much you can change
-between runs without paying a build cost, and how much of the image you own.
-
-The three modes:
-
-- **Bundled mode** — Kinetic builds a custom image with your dependencies
- baked in, via Cloud Build. This is the default.
-- **Prebuilt mode** — Kinetic pulls a published base image and installs your
- dependencies at pod startup with `uv pip install`.
-- **Custom image mode** — You provide a full image URI; Kinetic skips both
- the build and the install steps.
-
-You select the mode with the `container_image` argument on `@kinetic.run()`:
-
-```python
-@kinetic.run(accelerator="tpu-v6e-8") # bundled (default)
-@kinetic.run(accelerator="tpu-v6e-8", container_image="bundled") # bundled (explicit)
-@kinetic.run(accelerator="tpu-v6e-8", container_image="prebuilt") # prebuilt
-@kinetic.run(accelerator="tpu-v6e-8", container_image="us-docker.pkg.dev/me/repo/img:v1") # custom
-```
-
-:::{tip}
-**Recommended default:** **bundled mode**. It's the only mode that works
-out of the box, and it's the right choice for any workflow where your
-dependencies are reasonably stable. Cached images make warm runs fast;
-the build step only re-runs when your deps change.
-
-Reach for **prebuilt mode** only if you're iterating on `requirements.txt`
-several times a day and the per-iteration build cost is hurting you — and
-note that prebuilt currently requires you to publish your own base image
-with `kinetic build-image`, since no blessed base images ship with Kinetic
-today.
-:::
-
-## Recommendation matrix
-
-You are… | Use | Why
--------------------------------------------------- | ------------------------------- | ----------------------------------------------------------------
-A first-time user | **bundled** | The only mode that works without publishing your own base image.
-Iterating quickly on the same code | **bundled** | The dep-hashed image is cached; warm runs start in seconds.
-Changing dependencies multiple times a day | **prebuilt**\* | Skip the rebuild — install runs at pod startup instead.
-Running with a large dependency set | **bundled** | Pay the install cost once at build time, not on every run.
-Producing a reproducible production run | **bundled** | The exact environment is frozen into a tagged image.
-Needing custom system libs (CUDA builds, C++ deps) | **custom image** | Bundled and prebuilt can't add system packages.
-Pulling private packages | **bundled** or **custom image** | Bundled rebuilds on dep changes; custom gives full control.
-On a corporate base image | **custom image** | Use whatever your platform team blesses.
-
-\* Prebuilt mode requires a base image at the configured repo. Kinetic does
-not currently ship blessed base images, so you'll need to run
-`kinetic build-image` once and set `KINETIC_BASE_IMAGE_REPO` before this is a
-practical option.
-
-## Bundled mode
-
-Bundled mode runs Cloud Build to produce a tagged image with your project's
-dependencies installed. The image tag is a hash of those dependencies, so two
-jobs with the same `requirements.txt` reuse the same cached image.
-
-```python
-@kinetic.run(accelerator="tpu-v6e-8")
-def train():
- import keras
-
- ...
-```
-
-**Startup expectations:**
-
-- **Cold (first run, or after a dep change):** ~2–5 minutes for the build.
-- **Warm (cached image):** under a minute to schedule and start the pod.
-
-**Use it when:** any time you don't have a strong reason to do something else.
-This is the recommended default.
-
-**Avoid it when:** you change `requirements.txt` several times a day and the
-2–5 minute rebuilds are dominating your cycle time — at that point prebuilt
-mode is worth the setup cost.
-
-## Prebuilt mode
-
-Prebuilt mode pulls a published base image
-(`{repo}/base-{cpu|gpu|tpu}:{kinetic-version}`) that already contains the
-accelerator runtime, then runs `uv pip install` against your project's
-dependencies at pod startup.
-
-```python
-@kinetic.run(accelerator="tpu-v6e-8", container_image="prebuilt")
-def train(): ...
-```
-
-:::{warning}
-**You need to publish a base image first.** Kinetic does not currently ship
-blessed prebuilt base images. Before you can use prebuilt mode, run
-`kinetic build-image --repo ` once and set
-`KINETIC_BASE_IMAGE_REPO=` (or pass `base_image_repo=` to the
-decorator). Kinetic builds the image tag from the version of your
-installed Kinetic client. If you upgrade the client, publish a base image
-with the new tag first. See [Container Images](containers.md) for the full
-workflow.
-:::
-
-**Startup expectations:**
-
-- **Image pull:** typically 30–60 seconds the first time on a node, near zero
- once cached.
-- **Dependency install:** scales with the size of your `requirements.txt` —
- small projects start in under a minute, large ones a few minutes.
-
-**Use it when:** you've published a base image and you're churning
-`requirements.txt` often enough that bundled rebuilds are slowing you down.
-
-**Avoid it when:** every job has a long install step — bundled amortizes that
-cost into a single build, and once a bundled image is cached, warm runs are
-faster than prebuilt.
-
-## Custom image mode
-
-Pass a full image URI and Kinetic uses it as-is. No build, no install — your
-image is responsible for every dependency your function needs.
-
-```python
-@kinetic.run(
- accelerator="tpu-v6e-8",
- container_image="us-docker.pkg.dev/my-project/kinetic/my-image:v1.0",
-)
-def train(): ...
-```
-
-**Requirements:** the image must
-
-- Include the Kinetic runner script at `/app/remote_runner.py`. Kinetic
- invokes the container with `python3 -u /app/remote_runner.py`, which
- overrides whatever `ENTRYPOINT` or `CMD` the image declares — the only
- hard requirement is that the file is present at that path.
-- Have `python3` on `PATH`, with a version compatible with the one you
- used to pickle the function locally.
-- Install `cloudpickle`, `google-cloud-storage`, and `absl-py` — the
- runner imports them directly.
-- Install whatever other libraries your function imports.
-- It's also recommended to install the `kinetic` package itself if your
- function (or anything it imports) references it. The runner doesn't
- need it, but user code often does.
-- Be pullable from your GKE nodes (Artifact Registry in the same GCP
- project, or a public registry).
-
-**Startup expectations:** a single image pull, then immediate execution.
-Cold pulls vary widely with image size and registry latency.
-
-**Use it when:** you have system libraries that bundled or prebuilt can't add,
-you need a corporate-vetted base image, or you want full control over the
-image lifecycle.
-
-## How Kinetic decides what to build or install
-
-The dispatch happens at job submit time inside the backend execution path
-(`kinetic/backend/execution.py`):
-
-1. If `container_image == "prebuilt"`, Kinetic resolves the prebuilt base
- image for your accelerator category and uploads your filtered
- `requirements.txt` to GCS for runtime install.
-2. Else if `container_image` is `None` or `"bundled"`, Kinetic packages your
- working directory, computes a dependency hash, and either reuses a cached
- image or runs Cloud Build to produce a new one.
-3. Otherwise, Kinetic treats `container_image` as a literal image URI and
- uses it directly — no packaging of dependencies, no install.
-
-In all three modes, two more things happen regardless of which mode you
-picked:
-
-- **Your function and its captured closures.** Kinetic pickles them with
- `cloudpickle` and uploads them to Cloud Storage. Kinetic serializes the
- modules of your own project by value, so the pod does not import them.
- The runner in the pod downloads the payload, unpickles your function,
- and calls it.
-- **Your project source.** Kinetic zips it into a `context.zip` file and
- uploads that file to Cloud Storage. The archive starts at the *package
- root*, not at the directory of the script that you ran. Kinetic first
- walks up out of every directory that holds an `__init__.py` file.
- Kinetic then walks up to the nearest directory that holds a
- `pyproject.toml`, `requirements.txt`, `setup.py`, `setup.cfg`, or `.git`
- entry. Kinetic excludes the paths in your `Data(...)` objects and the
- default exclusion list (`.venv`, `node_modules`, and the cache
- directories).
-
-The runner extracts the archive into the workspace of the pod. The runner
-then rebuilds `sys.path` and changes to the workspace directory that
-matches your client working directory. Your imports and your relative-path
-reads thus operate as they do on your machine. This is most important in
-**custom image mode**: the image supplies the installed packages, but
-Kinetic still ships your project source with the job. You do not put the
-source into the image.
-
-[What Ships to the Pod](packaging.md) gives the full contract:
-
-- Root detection.
-- Exclusions and `.kineticignore`.
-- The rules for the working directory and for `sys.path`.
-- The guarantees for argument types.
-- Version matching between your client and the pod.
-
-## Related pages
-
-::::{grid} 1 1 2 2
-:gutter: 3
-
-:::{grid-item-card} {octicon}`file-directory;1em` What Ships to the Pod
-:link: packaging
-:link-type: doc
-
-The full packaging contract.
-:::
-
-:::{grid-item-card} {octicon}`package;1em` Dependencies
-:link: dependencies
-:link-type: doc
-
-How Kinetic discovers what to install.
-:::
-
-:::{grid-item-card} {octicon}`stack;1em` Container Images
-:link: containers
-:link-type: doc
-
-Base-image workflow and `kinetic build-image`.
-:::
-
-:::{grid-item-card} {octicon}`rocket;1em` Getting Started
-:link: ../getting_started
-:link-type: doc
-
-Your first run end-to-end.
-:::
-::::
diff --git a/docs/guides/faq.md b/docs/guides/faq.md
index edac43d2..62e9e2e0 100644
--- a/docs/guides/faq.md
+++ b/docs/guides/faq.md
@@ -1,81 +1,58 @@
# FAQ
-## When should I use `run()` vs `run_async()`?
-
-Use `@kinetic.run()` when you want your local script to wait for the
-result. Use `run_async()` when the job is long enough that you'd
-rather get a `JobHandle` back, walk away, and reattach later. `submit()` is
-the right call for anything multi-hour, anything you might want to monitor
-from a different machine, or anything you want to fan out and check on in
-parallel. See [Managing Async Jobs](async_jobs.md).
-
-## Why is the first run slower?
-
-The first run with a given set of dependencies builds a container image via
-Cloud Build (~2–5 minutes). The image is tagged by a hash of your
-dependencies, so any subsequent run with the same `requirements.txt` reuses
-the cached image and starts in under a minute. If your dependencies change,
-the build re-runs. When the build cost becomes a bottleneck (for example,
-when you change `requirements.txt` several times a day), switch to
-**prebuilt mode**, which installs deps at pod startup instead of baking
-them into a fresh image. See [Execution Modes](execution_modes.md) and
-[Dependencies](dependencies.md).
-
-## Should I use prebuilt or bundled mode?
-
-Default to **bundled**. It is the only mode that works without first
-publishing a base image. Reach for **prebuilt** when you change
-`requirements.txt` several times a day and the per-iteration build cost is
-hurting you. Prebuilt mode itself works with any base image at the
-configured repo, but the kinetic project does not currently publish public
-base images, so you will need to run `kinetic build-image` once to push your
-own before this becomes a usable option. See [Execution Modes](execution_modes.md).
-
-## When should I use `Data(...)` vs direct `gs://...` URIs?
-
-Always prefer `kinetic.Data(...)`. It accepts both local paths and
-`gs://` URIs and resolves to a plain filesystem path on the remote, so
-your function only sees paths regardless of where the bytes started.
-That is the whole point: one consistent API whether you are shipping a
-local directory, pointing at an existing GCS bucket, or asking for a
-FUSE mount via `Data(..., fuse=True)`. Reach for raw `gs://` URIs in
-your code only if you specifically want to bypass the `Data` abstraction.
-See [Data](data.md) for the decision matrix.
-
-## How do I save checkpoints and outputs?
-
-Write everything you want to keep under `KINETIC_OUTPUT_DIR`. Kinetic sets
-this env var inside the job pod to a per-job GCS prefix. Anything you write
-under it is durable: it outlives the pod and is reachable from your local
-machine. The job's Python return value is for small results; outputs and
-checkpoints belong on the output dir. See [Checkpointing](checkpointing.md).
-
-## How do I reattach to a job?
-
-Use `kinetic.attach(job_id)`. It reconstructs a `JobHandle` from the
-metadata Kinetic persisted to GCS at submission time, so you can call
-`.status()`, `.result()`, `.tail()`, or `.cleanup()` from any machine that
-has Kinetic and your GCP credentials. The `job_id` is available on the `JobHandle`
-returned by `run_async()`. If you have lost it, `kinetic.list_jobs()` enumerates
-jobs on the cluster. See [Managing Async Jobs](async_jobs.md).
-
-Note one condition: `.result()` unpickles the return value on your local
-machine. If the job returned an instance of one of your own classes, the
-machine that reattaches must import that class. Reattach from the project
-directory, or return plain data. If the unpickle fails, Kinetic raises an
-error, names the Cloud Storage artifact, and keeps that artifact.
+Short answers to the questions that new users ask most. Each answer
+links to the page with the details. The [glossary](#glossary) at the end
+defines the terms that the documentation uses.
+
+## How do I tell Kinetic which project and cluster to use?
+
+Run `kinetic init` one time. It saves a **profile** with your project,
+zone, cluster, and namespace, and makes that profile active. Every
+`kinetic` command and every `@kinetic.run()` call reads the active
+profile. To work with a second cluster, create a second profile and
+switch with `kinetic profile use NAME`. For a one-off override, pass
+`project=` or `cluster=` to the decorator, or `--project` and
+`--cluster` to a command. See [Profiles](profiles.md) and
+[Configuration](../configuration.md).
+
+## When do I call the function directly, and when do I use `run_async()`?
+
+Call the function directly (a **blocking call**) when you want your
+script to wait for the result. The call blocks and streams the logs. Use
+`func.run_async()` when the job runs for more than a few minutes. Use it
+also when you want to leave your machine, or when you want to watch
+several jobs at the same time. `run_async()` returns a `JobHandle` as
+soon as Kinetic submitted the job. See [Detached Jobs](async_jobs.md).
+
+## Why is the first run slow?
+
+The first run with a given dependency file builds a container image with
+Cloud Build. The build takes about 5 to 10 minutes. Kinetic tags the
+image with a hash of the dependency file, the accelerator category, the
+Python version, and the Kinetic version. Every later run with the same
+inputs reuses the image and starts in less than 1 minute. A change to the
+dependency file causes a new build. If you change the dependency file
+many times a day, see [Container Images](containers.md) for the prebuilt
+mode. See also [Dependencies](dependencies.md).
+
+## Which packages does the pod have?
+
+The image contains JAX for your accelerator, Keras, Kinetic, and the
+packages in the `requirements.txt` or `pyproject.toml` that Kinetic finds
+next to your script or above it. A `pip install` in your local shell
+does not carry over. Kinetic logs the file that it selected on each
+submit. See [Dependencies](dependencies.md).
## Which directory does Kinetic upload with my job?
Kinetic uploads the *package root*. Kinetic starts at the file that
defines your decorated function. Kinetic walks up out of each directory
that holds an `__init__.py` file. Kinetic then walks up to the nearest
-directory that holds a `pyproject.toml`, `requirements.txt`, `setup.py`,
-`setup.cfg`, or `.git` entry. Your home directory and the root of the
-file system are bounds for this search.
+directory that holds a `pyproject.toml`, a `requirements.txt`, a
+`setup.py`, a `setup.cfg`, or a `.git` entry. The walk stops at your
+home directory and at the root of the file system.
-Kinetic zips that directory. Kinetic excludes these paths from the
-archive:
+Kinetic zips that directory, without these paths:
- The paths in your `Data(...)` objects.
- The default exclusion list: `.venv`, `node_modules`, and the cache
@@ -83,138 +60,151 @@ archive:
- Each path that a `.kineticignore` pattern matches.
On the pod, the runner rebuilds `sys.path` and the working directory to
-match your client. Your imports and your relative paths thus operate as
-they do locally.
-
-To select the directory yourself, set `KINETIC_PACKAGE_ROOT`. The
-directory must exist, and it must contain the code that Kinetic packages.
-See [What Ships to the Pod](packaging.md).
-
-A notebook and a REPL have no source file. Kinetic packages your current
-working directory in this case.
-
-## What gets cleaned up automatically?
-
-A blocking `run()` and a `JobHandle.result()` call delete the Kubernetes
-Job and its pod. Kinetic does this for a job that succeeded and for a job
-that failed, so that dead resources do not stay in the cluster. A job with
-`debug=True` keeps its Kubernetes resources. To keep the resources for
-any other job, call `handle.result(cleanup=False)`. Read the logs, then
-call `handle.cleanup()`.
+match your client. Your imports and your relative paths therefore work
+as they do locally. To select the directory yourself, set
+`KINETIC_PACKAGE_ROOT`. A notebook and a REPL have no source file, so
+Kinetic uses your current working directory. See
+[What Ships to the Pod](packaging.md).
+
+## When do I use `Data(...)`, and when do I use a `gs://` URI directly?
+
+Use `kinetic.Data(...)`. It accepts a local path and a `gs://` URI, and
+it resolves to a plain filesystem path on the pod. Your function sees a
+path in every case: a local directory that Kinetic uploads, a bucket
+that already exists, or a FUSE mount with `Data(..., fuse=True)`. Read a
+`gs://` URI directly in your code only when you have a specific reason to
+bypass `Data`. See [Working with Data](data.md).
+
+## How do I keep checkpoints and other files?
+
+Write them under `KINETIC_OUTPUT_DIR`. Kinetic sets that environment
+variable in the pod to a per-job Cloud Storage location. Files under
+that location outlive the pod, and you can read them from your machine.
+The return value of the function is for small results. Files and
+checkpoints belong in the output directory. See
+[Outputs and Checkpoints](checkpointing.md).
-Kinetic deletes the Cloud Storage artifacts of the job (the uploaded code,
-the requirements, and the metadata) only when it collects a usable result.
-Kinetic keeps the artifacts in these conditions:
-
-- The job failed.
-- The job wrote no result.
-- The pod could not serialize the return value.
-- Your client could not deserialize the return value.
-
-For a job that you never collect, call `JobHandle.cleanup(gcs=True)` to
-delete the artifacts.
-
-Kinetic never deletes the outputs that you wrote under
-`KINETIC_OUTPUT_DIR` as part of job cleanup. The default output directory
-is in the job bucket, and the lifecycle rule of that bucket deletes
-objects after 30 days. Use one of these two methods to keep the outputs
-longer:
-
-- Copy the outputs to another bucket.
-- Set `output_dir` to a location outside that bucket.
+## How do I reattach to a job?
-## How do spot instances affect training?
+Call `kinetic.attach(job_id)`. It rebuilds a `JobHandle` from the
+metadata that Kinetic stored in Cloud Storage at submit time. You can
+then call `.status()`, `.result()`, `.tail()`, or `.cleanup()` from any
+machine that has Kinetic and your Google Cloud credentials. The job ID
+is on the `JobHandle` that `run_async()` returned. If you lost the ID,
+`kinetic.list_jobs()` or `kinetic jobs list` lists the jobs on the
+cluster. See [Detached Jobs](async_jobs.md).
+
+One condition: `.result()` unpickles the return value on your machine.
+If the job returned an instance of one of your own classes, the machine
+that reattaches must be able to import that class. Reattach from the
+project directory, or return plain data.
+
+## What does Kinetic clean up automatically?
+
+A blocking call and a `JobHandle.result()` call delete the Kubernetes
+Job and its pod, for a job that succeeded and for a job that failed. A
+job with `debug=True` keeps its Kubernetes resources until the 2-hour
+Kubernetes limit. To keep the resources for any other job, call
+`handle.result(cleanup=False)`, read the logs, and then call
+`handle.cleanup()`.
+
+Kinetic deletes the Cloud Storage artifacts of a job (the uploaded code,
+the requirements, and the metadata) only when it collects a usable
+result. Kinetic keeps the artifacts when the job failed, when the job
+wrote no result, when the pod could not serialize the return value, or
+when your client could not deserialize it. For a job that you never
+collect, call `JobHandle.cleanup(gcs=True)`.
+
+Kinetic never deletes what you wrote under `KINETIC_OUTPUT_DIR` as part
+of job cleanup. The default output directory is in the jobs bucket, and
+that bucket deletes objects after 30 days. To keep outputs longer, copy
+them to another bucket, or set `output_dir` to a location outside the
+jobs bucket. `kinetic down` deletes the jobs bucket with everything in it.
+
+## How do Spot VMs affect training?
+
+Spot capacity costs much less than on-demand capacity, but Google Cloud
+can preempt a Spot VM with 30 seconds of notice. A single-host job with
+frequent checkpoints recovers well. A multi-host TPU slice does not,
+because the loss of one host fails the whole slice. Use `--spot` on a
+node pool for fault-tolerant single-host work, and write checkpoints
+often enough to absorb a restart. See
+[Cost Optimization](cost_optimization.md).
+
+## When do I need more than one cluster?
+
+Most people do not. Create a second cluster to separate GPU and TPU
+work, to run in a second region, or to separate development from
+production. Each cluster has its own control plane cost. See
+[Clusters and Node Pools](clusters.md).
+
+## What is Pathways, and when does Kinetic use it?
+
+[Pathways](https://docs.cloud.google.com/ai-hypercomputer/docs/workloads/pathways-on-cloud/pathways-intro)
+is a JAX runtime that coordinates execution across many TPU hosts.
+Kinetic selects the Pathways backend automatically when the accelerator
+spans more than one host, for example `tpu-v5litepod-16` or
+`tpu-v6e-16`. JAX collectives and the Keras distribution API then work
+across hosts, and you write no coordination code. You pass
+`backend="pathways"` yourself only to test the Pathways code path on a
+single-host slice. See [Distributed Training](distributed_training.md)
+and the **Hosts** column in [Accelerators](../accelerators.md).
-Spot capacity costs significantly less than on-demand, but pods can be
-preempted with very little warning. Single-host jobs with frequent
-checkpoints recover well. Multi-host TPU slices do not, because losing
-any one host fails the whole slice. Use `--spot` for fault-tolerant
-single-host workloads, and write checkpoints often enough to absorb a
-restart. See [Cost Optimization](cost_optimization.md).
+## Glossary
-## When do I need multiple clusters?
+**Accelerator** — The string that you pass to `accelerator=`, for example `tpu-v5litepod-8`, `gpu-l4`, `gpu-a100x4`, or `cpu`. It names the hardware and, for a TPU, the slice size.
-Most users don't. Spin up a second cluster when you want to isolate GPU
-and TPU workloads, run jobs in different regions, or separate dev from
-prod environments. Each cluster has its own GKE control plane management
-fee, so don't add them speculatively. See [Multiple Clusters](clusters.md).
+**Topology** — The arrangement of TPU chips in a slice, for example `2x4`. Kinetic derives it from the accelerator name. A topology that spans more than one host makes the job multi-host.
-## What does Pathways mean in practice?
+**Pathways** — The JAX runtime that Kinetic uses for multi-host TPU slices. Kinetic selects it automatically.
-[Pathways](https://docs.cloud.google.com/ai-hypercomputer/docs/workloads/pathways-on-cloud/pathways-intro) is a JAX
-runtime that coordinates execution across many TPU hosts. Concretely,
-when you set `backend="pathways"` on a multi-host accelerator (e.g.,
-`tpu-v5litepod-2x4`), Kinetic launches your job against a
-Pathways-enabled cluster and JAX's collective communication (`jax.pmap`,
-sharding, etc.) Just Works across hosts. Without Pathways, you would have
-to manage multi-host JAX coordination yourself. See [Distributed Training](distributed_training.md).
+**Profile** — A saved set of project, zone, cluster, and namespace. One profile is active at a time. `kinetic init` creates the first one.
-## Glossary
+**Cluster** — A GKE cluster with its own image repository and buckets. Default name `kinetic-cluster`. Created by `kinetic up`, deleted by `kinetic down`.
-**Accelerator**: A TPU or GPU type identifier (e.g., `tpu-v6e-8`, `l4`,
-`a100`) passed to `accelerator=` on the decorator. Picks both the hardware
-and the topology.
+**Node pool** — A group of VMs of one accelerator type inside a cluster. Created by `kinetic pool add`. Scales between `--min-nodes` (default 0) and a fixed maximum: 10 more nodes for a GPU pool, or the hosts of one slice for a TPU pool.
-**Topology**: How many chips are arranged into the slice. For TPUs,
-encoded in the accelerator name (`tpu-v6e-8` is 8 chips; `tpu-v5litepod-2x4`
-is a 2×4 slice across hosts).
+**Job** — One execution of a decorated function on the cluster. A job has an ID such as `job-3f9a1c2b`.
-**Pathways**: JAX runtime for multi-host TPU coordination. Selected via
-`backend="pathways"` and required for cross-host collectives without
-hand-rolled setup.
+**Handle** — A `JobHandle` for one detached job, or a `BatchHandle` for a batch. Wraps `status()`, `result()`, `tail()`, `cancel()`, and `cleanup()`.
-**Node pool**: A GKE-managed group of VMs of one accelerator type.
-Created with `kinetic pool add`. Scales between `--min-nodes` and the max
-you need for the job.
+**Entry directory** — The directory of the file that defines the decorated function. Kinetic starts its search for the package root and for the dependency file there.
-**Cluster**: A GKE cluster with its own control plane and Artifact
-Registry. Default name `kinetic-cluster`. Managed with `kinetic up`,
-`kinetic down`, and `kinetic status`.
+**Package root** — The directory that Kinetic archives and ships with the job, in `context.zip`.
-**Bundled image**: A container image Kinetic builds for you via Cloud
-Build, with your dependencies baked in. The default execution mode. Tagged
-by a hash of your `requirements.txt`.
+**Workspace** — The directory on the pod into which the runner extracts the package root.
-**Prebuilt image**: A published base image that already has the
-accelerator runtime installed. Your project deps are installed at pod
-startup. Selected with `container_image="prebuilt"`. Requires you to
-publish base images with `kinetic build-image` first.
+**Dependency file** — The `requirements.txt` or `pyproject.toml` that decides which packages the image contains.
-**FUSE**: Filesystem-in-userspace mount. With `kinetic.Data(..., fuse=True)`,
-a GCS bucket is mounted lazily into the pod's filesystem so reads stream
-on demand instead of downloading up front.
+**Image** — The container image that the pod runs. Kinetic builds it by default. See [Container Images](containers.md) for the other modes.
-**Handle**: A `JobHandle` returned by `run_async()` (or
-`kinetic.attach()`). Wraps `status()`, `result()`, `tail()`, and
-`cleanup()` for one job.
+**Output directory** — The Cloud Storage location in `KINETIC_OUTPUT_DIR`. The place for checkpoints and every file that you want to keep.
-**Output dir**: The GCS prefix at `KINETIC_OUTPUT_DIR` inside the job
-pod. The canonical place to write checkpoints and any files you want to
-keep after the pod exits.
+**FUSE** — A mount of a Cloud Storage location into the pod filesystem. With `Data(..., fuse=True)`, the pod reads files on demand instead of a download at start.
## Related pages
::::{grid} 1 1 2 2
:gutter: 3
-:::{grid-item-card} {octicon}`gear;1em` Execution Modes
-:link: execution_modes
+:::{grid-item-card} {octicon}`light-bulb;1em` How Kinetic Works
+:link: ../concepts
:link-type: doc
-Bundled vs prebuilt vs custom.
+The model behind every answer on this page.
:::
:::{grid-item-card} {octicon}`bug;1em` Troubleshooting
:link: ../troubleshooting
:link-type: doc
-Symptom-first debugging.
+Symptom-first diagnosis.
:::
:::{grid-item-card} {octicon}`rocket;1em` Getting Started
:link: ../getting_started
:link-type: doc
-Your first run, end-to-end.
+Your first run, end to end.
:::
::::
diff --git a/docs/guides/packaging.md b/docs/guides/packaging.md
index 0094df49..147f6aa5 100644
--- a/docs/guides/packaging.md
+++ b/docs/guides/packaging.md
@@ -1,8 +1,12 @@
# What Ships to the Pod
-Kinetic uploads two artifacts to Cloud Storage for every job. The pod
-downloads both artifacts and rebuilds your project in a workspace
-directory.
+This page is the exact contract for the artifacts that Kinetic sends to
+the pod. It describes what goes into them, where the pod stores them, and
+how the pod rebuilds your project. Read this page when you debug a
+`ModuleNotFoundError`, a `FileNotFoundError`, or a job that uploads
+gigabytes. Then read [Troubleshooting](../troubleshooting.md).
+
+The pod downloads two artifacts from Cloud Storage for every job:
1. **`payload.pkl`** — your function, the objects that it closes over,
its arguments, the environment variables that you capture, and a
@@ -12,21 +16,14 @@ directory.
*package root*. Kinetic writes the packaging plan into the same
archive, at the reserved path `.kinetic/plan.json`.
-Kinetic reads one dependency file too: a `requirements.txt` or a
-`pyproject.toml`. In bundled mode Kinetic puts that content into the
-image build. In prebuilt mode Kinetic uploads the content beside the two
-artifacts, and the pod installs it at startup. See
-[Dependencies](dependencies.md).
-
-This page is the exact contract for the artifacts and their destinations.
-Read it first when you debug a `ModuleNotFoundError`, a
-`FileNotFoundError`, or a job that uploads gigabytes. Then read
-[Troubleshooting](../troubleshooting.md).
+Kinetic also reads one dependency file: a `requirements.txt` or a
+`pyproject.toml`. That file decides which Python packages the pod
+provides. See [Dependencies](dependencies.md).
## The package root
The package root is the directory that Kinetic archives. Kinetic resolves
-it at submit time, in three steps.
+the package root at submit time, in three steps.
**Step 1 — find the entry directory.** Kinetic finds the module that
defines the decorated function, and takes the directory of that module
@@ -36,7 +33,7 @@ those cases Kinetic uses your current working directory instead.
**Step 2 — escape the package.** Kinetic walks up from the entry
directory for as long as each directory holds an `__init__.py` file.
This step makes `from trainer.model import ...` resolvable on the pod,
-because the root ends above `trainer/`, and not inside it. This walk also
+because the root ends above `trainer/`, and not inside it. This walk
stops at your home directory and at the root of the file system.
**Step 3 — walk up to a project marker.** Kinetic then walks up to the
@@ -50,9 +47,10 @@ nearest directory that holds one of these markers:
The walk stops at your home directory and at the root of the file system.
Kinetic never adopts either one as the package root, unless step 2 already
-ended there. Kinetic keeps the directory from step 2 when it finds no marker.
+ended there. If Kinetic finds no marker, Kinetic keeps the directory from
+step 2.
-**Override.** Set `KINETIC_PACKAGE_ROOT` to pin the root:
+**Override.** Set `KINETIC_PACKAGE_ROOT` to select the root yourself:
```bash
export KINETIC_PACKAGE_ROOT=/home/me/monorepo/services/trainer
@@ -66,12 +64,12 @@ submit time in two cases:
- The value is neither the entry directory nor a parent of it.
:::{tip}
-Two habits keep the root predictable:
+Two rules keep the root predictable:
- Keep a `pyproject.toml`, a `requirements.txt`, or a git repository at
the top of the tree that you want to ship.
- Keep large data out of that tree. As an alternative, wrap the data in
- `kinetic.Data(...)`, which Kinetic excludes from `context.zip`
+ `kinetic.Data(...)`. Kinetic excludes such a path from `context.zip`
automatically.
:::
@@ -81,7 +79,7 @@ Two habits keep the root predictable:
| ------ | --------------------- | ------------ | --- |
| `proj/train.py` (+ `proj/requirements.txt`) | `train.py` | `proj/` | The entry directory already holds a marker. |
| `proj/trainer/model.py` (+ `proj/pyproject.toml`, `trainer/__init__.py`) | `model.py` | `proj/` | Step 2 escapes the package. Step 3 finds the marker. |
-| `proj/src/pkg/train.py` (+ `pkg/__init__.py`, `proj/pyproject.toml`) | `train.py` | `proj/` | `src/` ships too, so the plan can point `sys.path` at it. |
+| `proj/src/pkg/train.py` (+ `pkg/__init__.py`, `proj/pyproject.toml`) | `train.py` | `proj/` | Kinetic archives `src/` too, so the plan can add it to `sys.path`. |
| `~/scratch/one_off.py`, no markers anywhere | `one_off.py` | `~/scratch/` | No marker exists. Kinetic keeps the entry directory, and never `$HOME`. |
| A Jupyter notebook in `~/nb/`, no markers anywhere | the notebook cell | `~/nb/` | The cell has no `__file__`. Kinetic uses the current directory. |
@@ -91,16 +89,19 @@ Kinetic always excludes these two directory names, at any depth:
`.git`, `__pycache__`
+This rule applies to directory names only. In a git worktree, `.git` is
+a small file, and Kinetic archives that file.
+
Kinetic also excludes these names at any depth, unless you turn the
default exclusions off:
`.venv`, `venv`, `node_modules`, `.tox`, `.mypy_cache`, `.ruff_cache`,
`.pytest_cache`, `.ipynb_checkpoints`, `.DS_Store`
-Kinetic excludes a local path that you wrap in `kinetic.Data(...)` as
-well, when that path lies inside the package root. Those bytes travel
-through the content-addressed cache for data instead, so Kinetic does not
-upload them two times. See [Data](data.md).
+Kinetic also excludes a local path that you wrap in `kinetic.Data(...)`,
+if that path lies inside the package root. Kinetic uploads that data
+through the content-addressed data cache instead, and does not upload the
+data a second time. See [Working with Data](data.md).
To turn the default exclusions off, set this variable:
@@ -116,7 +117,8 @@ Put a `.kineticignore` file at the package root to exclude more paths.
Kinetic reads this file at the package root only. Write one pattern per
line, in `fnmatch` syntax. A line that starts with `#` is a comment.
Kinetic reads a `#` in any other position as part of the pattern. A
-trailing `/` restricts the pattern to directories.
+trailing `/` restricts the pattern to directories. Kinetic removes a
+leading `/`, so `/build/` and `build/` behave the same.
```text
# .kineticignore
@@ -153,16 +155,16 @@ of the files, do these steps:
### Size warnings
-Kinetic logs the archive size on every submit. Kinetic logs a warning
+Kinetic logs the archive size on every submission. Kinetic logs a warning
above 100 MB, and lists the five largest files in the archive. Change the
threshold with `KINETIC_CONTEXT_SIZE_WARN_MB`.
The pickled payload has a separate threshold of 50 MB
-(`KINETIC_PAYLOAD_SIZE_WARN_MB`). A large payload almost always means one
+(`KINETIC_PAYLOAD_SIZE_WARN_MB`). A large payload usually means one
module-level global that Kinetic captured **by value**. `cloudpickle`
serializes the objects that your function references. A module-level
-`DF = pd.read_parquet(...)` that your function reads therefore lands
-inside `payload.pkl` on every submit. Pass large data as
+`DF = pd.read_parquet(...)` that your function reads therefore goes
+into `payload.pkl` on every submission. Pass large data as
`kinetic.Data(...)`, or load it inside the function.
Set either threshold to `0` to turn that warning off. Kinetic ignores a
@@ -170,13 +172,13 @@ value that is not a number, and logs a warning about the bad value.
### Fidelity of the archive
-`context.zip` is a faithful snapshot, with these explicit rules:
+`context.zip` is an exact snapshot, with these rules:
- Kinetic follows a symlinked directory, and guards against a cycle.
Kinetic archives each real directory one time only, and names every
skipped directory in a warning.
- Kinetic skips a broken symlink and an unreadable file, and logs a
- warning for each one. One bad file never stops a submit.
+ warning for each one. One bad file never stops a submission.
- Kinetic stores the POSIX mode of each file. The runner restores the
read, write, and execute bits on the pod, for an archive that a Unix
client built. The runner never restores a setuid, setgid, or sticky
@@ -188,47 +190,54 @@ value that is not a number, and logs a warning about the bad value.
## What happens on the pod
The pod downloads both artifacts and verifies their SHA-256 hashes. The
-runner then does the following:
+runner then does these steps:
:::{container} kinetic-steps
1. The runner extracts `context.zip` into the workspace directory.
2. The runner rebuilds `sys.path` from the packaging plan. The workspace
root goes first. Each client `sys.path` entry that lived under the
- package root follows it, at the matching position inside the
+ package root follows it, at the same relative path inside the
workspace. This step makes a `src/` layout work without a
`PYTHONPATH` on the pod.
3. The runner changes the working directory to the workspace directory
that matches your client working directory. The runner uses the
workspace root when your client working directory was outside the
package root.
-4. The runner unpickles `payload.pkl` and calls your function.
+4. The runner unpickles `payload.pkl`.
+5. The runner applies the captured environment variables and resolves
+ the `Data` references.
+6. The runner calls your function.
:::
The runner does not replicate a client `sys.path` entry that points at
-site-packages, or at any directory outside the package root. The image
-provides those packages.
+site-packages, or at any directory outside the package root. The pod
+environment provides those packages.
Step 3 has one practical consequence: **a relative path behaves the same
on the pod as it does on your machine**. `open("configs/train.yaml")`
-works remotely if it worked locally from the same directory. The file
+works remotely if it works locally from the same directory. The file
must be inside the package root, and it must not be excluded.
An absolute client path is not a supported way to read a shipped file.
-The runner does create one symbolic link at the path of your client
-working directory. The link points at the workspace root. The link exists
-so the debugger can map source files. Use relative paths for portable
-access to files.
+The runner does create one symbolic link, at the path of your client
+**entry directory**. The entry directory is the directory of the module
+file that defines the decorated function. The link points at the
+workspace root, and the runner creates it only when that path does not
+exist on the pod. The link exists so that the debugger can map source
+files. When the entry directory sits below the package root, the link
+still points at the workspace root, and not at the matching
+subdirectory. Use relative paths for portable access to files.
The pod environment also holds `KINETIC_OUTPUT_DIR`. The value is a Cloud
-Storage URI, such as `gs://{bucket}/outputs/{job_id}`. The pod destroys
-the workspace at the end of the job. Write everything that you want to
-keep to that Cloud Storage location. See
-[Checkpointing](checkpointing.md).
+Storage URI, such as `gs://{project}-kn-{cluster}-jobs/outputs/{job_id}`.
+The pod filesystem, including the workspace, is gone when the pod ends.
+Write everything that you want to keep to that Cloud Storage location.
+See [Outputs and Checkpoints](checkpointing.md).
## How imports resolve remotely
-Kinetic uses two mechanisms. The mechanism that applies explains almost
-every remote import failure.
+Kinetic uses three mechanisms. The mechanism that applies to a module
+explains almost every remote import failure.
**Kinetic ships your first-party code by value.** At submit time Kinetic
inspects `sys.modules`. Kinetic registers each module whose file lives
@@ -236,8 +245,8 @@ under the package root with `cloudpickle.register_pickle_by_value`.
`cloudpickle` then serializes the functions and classes of those modules
into `payload.pkl` in full. The pod therefore does not import your
modules to unpickle your job. A helper in `trainer/utils.py` travels
-inside the payload. A module-level `import kinetic` in one of your modules
-does not have to succeed at unpickle time.
+inside the payload. The pod does not run the module-level statements of
+those modules, because the pod does not import them.
Kinetic never registers these two modules by value:
@@ -248,21 +257,28 @@ Kinetic never registers these two modules by value:
**The pod resolves third-party imports.** `cloudpickle` pickles anything
outside the package root by reference: `numpy`, `keras`, and your private
-packages. The image must provide those packages, and that is the purpose
-of the dependency file.
+packages. The pod environment must provide those packages, and that is
+the purpose of the dependency file.
**The workspace resolves runtime imports.** An `import` statement inside
your function runs at call time, against the pod `sys.path`. That path
starts with the workspace. The pod finds a first-party module that you
-import lazily in the extracted context. The image must provide a
-third-party package that you import lazily.
+import lazily in the extracted context. The pod environment must provide
+a third-party package that you import lazily.
:::{note}
**Identity limit of shipping by value.** A class that travels by value
is a different type object from the same class that the pod imports. An
-`isinstance` check across the two paths can return `False`. A class
-defined in `__main__` always behaves this way. Compare by attribute, or by
-name, when you must cross that boundary.
+`isinstance` check across the two paths can return `False`. Compare by
+attribute, or by name, when you must cross that boundary.
+:::
+
+:::{note}
+**A decorated function inside the payload.** The payload can reference a
+second `@kinetic.run()` function, for example as an attribute of an
+object that you pass. That function arrives on the pod as the plain,
+undecorated function. A call to it on the pod runs the body in the pod
+process. Kinetic does not support a nested job submission.
:::
## Arguments: which types Kinetic preserves
@@ -280,8 +296,9 @@ The walk keeps types and object identity:
- Kinetic rebuilds a `tuple` and a `NamedTuple` as the same type, with
the fields intact. Both `typing.NamedTuple` and
`collections.namedtuple` work.
-- Kinetic passes a `set` and a `frozenset` through unchanged, because
- neither one can hold a `Data` object.
+- Kinetic passes a `set` and a `frozenset` through unchanged when no
+ `Data` object is reachable inside it. If a `Data` object is reachable,
+ Kinetic rejects the argument (see below).
- Kinetic preserves a `dict` subclass. `OrderedDict` keeps its order,
`defaultdict` keeps its `default_factory`, and `Counter` keeps its
type.
@@ -296,8 +313,9 @@ built-in type instead, and logs a warning. Kinetic does not fail the job.
**Kinetic rejects three argument shapes at submit time.** Each message
names the position of the argument at fault:
-- A `Data` object inside a `set` or a `frozenset`. The replacement
- reference is a dict, and a dict is not hashable.
+- A `Data` object inside a `set` or a `frozenset`, at any depth (for
+ example, inside a tuple that sits in a set). The replacement reference
+ is a dict, and a dict is not hashable.
- A `Data` object used as a `dict` key.
- A self-referential structure whose cycle runs through a tuple, a set,
or a frozenset. A cycle through a list or a dict arrives on the pod
@@ -340,10 +358,10 @@ covers it does not.
Kinetic logs the names that it captured, and never the values. Kinetic
logs one more warning when a captured name contains `TOKEN`, `SECRET`,
`KEY`, `PASSWORD`, or `CREDENTIAL`, in any letter case. Kinetic stores
-those values inside `payload.pkl` in the job bucket, and every job pod in
-the cluster can read that bucket. This warning is informational, and it
-appears even for a name that you listed exactly. If you intend to send the
-credential, you need no action. See
+those values inside `payload.pkl` in the jobs bucket, and every job pod
+in the cluster can read that bucket. This warning is informational, and
+it appears even for a name that you listed exactly. If you intend to send
+the credential, you need no action. See
[Forward Environment Variables](env_vars.md) and
[Security](../security.md).
@@ -357,40 +375,49 @@ your current directory can move the root further up the tree. Change into
your project directory before you submit.
A notebook function has no importable module, so `cloudpickle` always
-pickles it by value. That behavior is the one that you want. A helper
+pickles it by value. That behavior is correct for a notebook. A helper
from another cell travels with the function. A helper in a `.py` file
beside the notebook travels in `context.zip`.
## Matching your local environment to the pod
Pickled code objects are **not** portable across Python minor versions. A
-function pickled on 3.12 does not unpickle on 3.11.
-
-- **Bundled mode** (the default) builds an image from
- `python:{your minor version}-slim`. The build installs `keras`,
- `cloudpickle`, `google-cloud-storage`, JAX for your accelerator
- category, and `keras-kinetic` pinned to your client version. The pod
- Python therefore always matches your client. Use this mode if you do
- not want to think about the question.
-- **Prebuilt mode** pulls a base image that you publish with
- `kinetic build-image`. Kinetic requests the image tag that matches your
- client Kinetic version. Build that image with the Python minor version
- of the client that submits the job.
-- **Custom image mode** makes you responsible for the Python version. See
- the image requirements in [Execution Modes](execution_modes.md).
+function that you pickle on Python 3.12 can fail to unpickle on Python
+3.11, or can crash the interpreter later. The pod Python must therefore
+match your client at the minor version (`X.Y`).
+
+- **The image that Kinetic builds** (the default) starts from
+ `python:{X.Y}-slim`, where `X.Y` is the minor version of your local
+ interpreter. The build installs `keras`, `cloudpickle`,
+ `google-cloud-storage`, JAX for your accelerator category, and
+ `keras-kinetic` pinned to your client version. The pod Python
+ therefore always matches your client, and you take no action.
+- **A base image that you publish** with `kinetic build-image` must have
+ the Python minor version of the client that submits the job. Kinetic
+ requests the image tag that matches your client Kinetic version.
+- **A custom image** makes you responsible for the Python version.
+
+See [Container Images](containers.md) for the three ways to produce the
+image and for the requirements of a custom image.
Every payload carries a fingerprint of the client: the Python version,
the `cloudpickle` version, and the Kinetic version. The runner compares
-the Python version and the `cloudpickle` version against the pod, and logs
-a warning about a difference. The error that `result()` raises names both
-sides, such as `client Python 3.12.2 / pod Python 3.11.9`. The pod log
-holds a skew warning only when the payload unpickled correctly.
+the fingerprint against the pod, and logs a warning for a mismatch. The
+runner compares the Python version at the minor version only (`X.Y`), so
+`3.12.2` and `3.12.7` produce no warning. The runner compares the
+`cloudpickle` version exactly. When the payload does not unpickle, the
+error that `result()` raises names both sides, such as
+`client Python 3.12.2 / pod Python 3.11.9`. The pod log holds a skew
+warning only when the payload unpickled correctly.
+
+## Settings reference
-## Environment knobs
+These settings are environment variables on the machine that submits the
+job. A profile does not store them.
| Variable | Default | Effect |
| -------- | ------- | ------ |
-| `KINETIC_PACKAGE_ROOT` | _(unset)_ | Pin the package root. The value must name an existing directory, and must be the entry directory or a parent of it. |
+| `KINETIC_PACKAGE_ROOT` | _(unset)_ | Set the package root. The value must name an existing directory, and must be the entry directory or a parent of it. |
| `KINETIC_NO_DEFAULT_EXCLUDES` | _(unset)_ | Set it to `1` to ship `.venv/`, `node_modules/`, and the rest of the default exclusion list. `.git` and `__pycache__` stay excluded. |
| `KINETIC_CONTEXT_SIZE_WARN_MB` | `100` | The `context.zip` size above which Kinetic logs a warning and lists the largest files. Set it to `0` to turn the warning off. |
| `KINETIC_PAYLOAD_SIZE_WARN_MB` | `50` | The `payload.pkl` size above which Kinetic logs a warning about capture by value. Set it to `0` to turn the warning off. |
@@ -404,21 +431,21 @@ holds a skew warning only when the payload unpickled correctly.
:link: dependencies
:link-type: doc
-How Kinetic finds the requirements file.
+How Kinetic finds the dependency file, and what the image contains.
:::
-:::{grid-item-card} {octicon}`database;1em` Data
+:::{grid-item-card} {octicon}`database;1em` Working with Data
:link: data
:link-type: doc
How to move large inputs without the archive.
:::
-:::{grid-item-card} {octicon}`zap;1em` Execution Modes
-:link: execution_modes
+:::{grid-item-card} {octicon}`container;1em` Container Images
+:link: containers
:link-type: doc
-How Kinetic produces the image.
+How Kinetic produces the image, and the three modes.
:::
:::{grid-item-card} {octicon}`bug;1em` Troubleshooting
diff --git a/docs/guides/profiles.md b/docs/guides/profiles.md
index a15e088b..cde7ff7a 100644
--- a/docs/guides/profiles.md
+++ b/docs/guides/profiles.md
@@ -1,34 +1,42 @@
# Profiles
-A **profile** bundles the four settings that determine where a Kinetic
-job runs — project, zone, cluster, and namespace (see
-[Configuration](../configuration.md) for what each one controls) —
-under a single name. Instead of re-exporting `KINETIC_*` environment
-variables for each cluster you target, you save a profile once and
-switch between them with a single command.
+A **profile** is a saved set of the four values that decide where a job
+runs: the Google Cloud project, the zone, the cluster, and the Kubernetes
+namespace. One profile is **active** at a time. Every `kinetic` command
+and every `@kinetic.run()` call reads the active profile. You therefore
+set the target one time, and you do not repeat it in code, in commands,
+or in environment variables.
-Profiles are optional and additive. If you have never run
-`kinetic profile create`, the existing env-var and prompt flow is
-unchanged.
+`kinetic init` creates your first profile and makes it active. This page
+explains how to create more profiles, how to switch between them, and
+how a profile combines with the other configuration sources.
-## When to use profiles
+## When you need more than one profile
-- You work with more than one cluster (for example, a dev cluster and a
- team-shared GPU cluster).
-- You want to chase spot capacity across zones or regions — keep a
- profile per (zone, cluster) pair and switch when one runs dry.
-- You move between projects or namespaces often enough that exporting
- env vars manually is slowing you down.
-- Multiple people share a machine or dotfiles and need clean switching
- between their own configurations.
+One profile per cluster is the normal pattern. Create a second profile
+when:
-For single-cluster setups, plain `KINETIC_*` env vars are still the
-simplest path — see [Configuration](../configuration.md) for the full
-list and what each one does.
+- You work with two clusters, for example a personal TPU cluster and a
+ shared GPU cluster.
+- You keep clusters in two zones or two regions, and you switch when one
+ zone has no Spot capacity.
+- You work in two projects, or in two namespaces of one cluster.
+- Two people share one user account on one machine and need separate
+ configurations.
-## Quick start
+## Create a profile
-Create two profiles — the first becomes active automatically:
+`kinetic init` and `kinetic up` save a profile automatically and make it
+active. The profile name defaults to the cluster name. Pass
+`--profile-name NAME` to either command to select a different name. If a
+profile with that name exists, `kinetic up` overwrites it, because the
+cluster that `up` just created is authoritative.
+
+A profile name starts with a letter or a digit, and contains only
+letters, digits, `-`, and `_`, up to 64 characters.
+
+To save a profile for a cluster that already exists, run
+`kinetic profile create`:
```bash
kinetic profile create dev-tpu \
@@ -39,101 +47,107 @@ kinetic profile create team-gpu \
--namespace alice
```
-Either pass the fields as flags (as above) or omit them to be prompted
-interactively. Any field you don't pass is read from `KINETIC_*` env
-vars, then prompted for if still unset.
+You can omit any flag. Kinetic then reads the value from the matching
+`KINETIC_*` environment variable, and prompts you if the variable is not
+set. The first profile that you create becomes active. After that,
+`kinetic profile create` does not change the active profile.
+`kinetic init`, `kinetic up`, `kinetic profile use`, `kinetic profile
+unset`, and `kinetic profile rm` do.
-List what's saved and switch between them:
+## Switch and inspect
```bash
-kinetic profile ls # active profile marked with *
+kinetic profile ls # lists profiles; * marks the active one
kinetic profile use team-gpu
-kinetic profile show # settings of the active profile
+kinetic profile show # prints the active profile
+kinetic config # prints each resolved value and its source
```
-The first profile you create is automatically promoted to active; after
-that, the only way to change the active profile is
-`kinetic profile use NAME`.
-
## Commands
| Command | What it does |
-| --- | --- |
-| `kinetic profile create [NAME]` | Save a new profile. Fills any field you don't pass as a flag from `KINETIC_*` env vars, then prompts for whatever is still missing. First profile becomes active. |
-| `kinetic profile ls` | List all profiles, with `*` on the active one. |
-| `kinetic profile use NAME` | Mark `NAME` as the persistent active profile. |
-| `kinetic profile show [NAME]` | Show a profile's settings. Defaults to the active profile. |
-| `kinetic profile rm NAME` | Delete a profile. Prompts unless `--yes`. If the removed profile was active, no fallback is chosen — run `profile use` to pick a new one. |
+| ------- | ------------ |
+| `kinetic profile create [NAME]` | Saves a new profile. Reads a missing field from the `KINETIC_*` environment variable, then prompts. The first profile becomes active. `--force` overwrites a profile with the same name. |
+| `kinetic profile ls` | Lists all profiles. `*` marks the active one. |
+| `kinetic profile use NAME` | Makes `NAME` the active profile. |
+| `kinetic profile unset` | Clears the active profile. Commands then read only environment variables, flags, and defaults. |
+| `kinetic profile show [NAME]` | Prints the fields of a profile. Defaults to the active profile. |
+| `kinetic profile rm NAME` | Deletes a profile. Prompts unless you pass `--yes`. If you delete the active profile, no other profile becomes active. Run `kinetic profile use` to select one. |
-## Selecting a profile for one invocation
+## Use a different profile for one command
-Use `--profile NAME` (or `KINETIC_PROFILE=NAME`) to target a non-active
-profile without changing the stored active one. The flag belongs to the
-root `kinetic` command, so it must come **before** the subcommand:
+Pass `--profile NAME` to the root `kinetic` command, or set
+`KINETIC_PROFILE=NAME`. The stored active profile does not change. The
+flag belongs to the root command, so it must come **before** the
+subcommand:
```bash
-# Works:
+# Correct: the flag is on the root command.
kinetic --profile team-gpu jobs list
-# Also works:
+# Also correct.
KINETIC_PROFILE=team-gpu kinetic jobs list
-# Does NOT work — --profile here would be a subcommand flag, not the global override.
+# Wrong: `--profile` after the subcommand is not the global flag.
kinetic jobs list --profile team-gpu
```
-When the selector differs from the stored active profile,
-`kinetic profile ls` marks the effective profile with `*` and prints
-an annotation like `Active profile: team-gpu (override; stored: dev-tpu)`
-underneath the table — so it's always obvious which profile is actually
-in effect for this invocation.
+When the selected profile is not the stored active profile,
+`kinetic profile ls` marks the selected profile with `*` and prints a
+line such as `Active profile: team-gpu (override; stored: dev-tpu)`.
+
+Python code reads `KINETIC_PROFILE` too. To run one script against a
+different profile, set the variable for that process:
-## Where profiles sit in the precedence chain
+```bash
+KINETIC_PROFILE=team-gpu python train.py
+```
-Profiles slot in below env vars and CLI flags:
+## How a profile combines with the other sources
+
+For each of the four values, Kinetic reads these sources in order and
+uses the first value that it finds:
```text
-CLI flag > KINETIC_* env var > active profile > built-in default
+decorator argument or CLI flag > KINETIC_* env var > active profile > built-in default
```
-In practice that means:
+Concretely:
-- Profiles give you good defaults without silencing explicit overrides.
-- `KINETIC_PROJECT=other-proj kinetic up` still targets `other-proj`
- even if the active profile names something else.
-- `kinetic --cluster adhoc status` still targets `adhoc` regardless of
- profile.
+- `@kinetic.run(cluster="adhoc")` runs on `adhoc`, whatever the profile
+ says.
+- `KINETIC_PROJECT=other-proj kinetic up` targets `other-proj`.
+- `kinetic status --cluster adhoc` inspects `adhoc`.
-`kinetic config show` prints the currently-active profile and, for each
-setting, where its resolved value came from (`profile`, a `KINETIC_*`
-env var, or `default`) — run it any time you want to see which layer
-won.
+Use the profile for the target that you use every day. Use a flag, an
+argument, or an environment variable for a one-off override. Run
+`kinetic config` to see which source supplied each value. See
+[Configuration](../configuration.md) for the full precedence table.
-## Storage
+## Where profiles are stored
-Profiles are local and per-user — they live on your machine only and
-are not synced across teammates or across machines. If a workflow
-needs to be reproducible by others, document it with `KINETIC_*` env
-vars (or CLI flags) rather than relying on a shared profile name.
+Profiles are local to your machine and to your user account. Kinetic does
+not sync them to teammates or to other machines. To make a workflow
+reproducible for others, write down the four values, or the
+`kinetic profile create` command that sets them.
-Profiles live in a single JSON file at `~/.kinetic/profiles.json`:
+Kinetic stores all profiles in one JSON file at
+`~/.kinetic/profiles.json`:
```json
{
"current": "team-gpu",
"profiles": {
- "dev-tpu": { "project": "my-ml-dev", "zone": "us-central2-b", "cluster": "dev-tpu", "namespace": "default" },
+ "dev-tpu": { "project": "my-ml-dev", "zone": "us-central2-b", "cluster": "dev-tpu", "namespace": "default" },
"team-gpu": { "project": "my-ml-prod", "zone": "us-east1-b", "cluster": "team-gpu", "namespace": "alice" }
}
}
```
-Editing the file by hand works, but the CLI is the supported path.
-Writes are atomic (tempfile + rename) so a crash mid-write will not
-corrupt the store.
-
-If you need to relocate the file (for example, in CI) set
-`KINETIC_PROFILES_FILE` to an alternate path.
+You can edit the file by hand, but the CLI is the supported path. Kinetic
+writes the file atomically, so an interrupted write does not corrupt the
+file. To store the file elsewhere, for example in CI, set
+`KINETIC_PROFILES_FILE` to a different path.
## Related pages
@@ -144,21 +158,20 @@ If you need to relocate the file (for example, in CI) set
:link: ../configuration
:link-type: doc
-The full precedence table and the list of `KINETIC_*` env vars.
+The full precedence table and every `KINETIC_*` variable.
:::
-:::{grid-item-card} {octicon}`server;1em` Multiple Clusters
+:::{grid-item-card} {octicon}`server;1em` Clusters and Node Pools
:link: clusters
:link-type: doc
-How Kinetic treats clusters as first-class targets; profiles are the
-ergonomic layer on top.
+Create, share, and delete the clusters that profiles point at.
:::
:::{grid-item-card} {octicon}`terminal;1em` CLI Reference
:link: ../cli
:link-type: doc
-Generated reference for every flag.
+Generated reference for every command and flag.
:::
::::
diff --git a/docs/guides/profiling.md b/docs/guides/profiling.md
index bcd3ae1b..8dcc2016 100644
--- a/docs/guides/profiling.md
+++ b/docs/guides/profiling.md
@@ -1,97 +1,154 @@
-# Performance Profiling (XProf)
+# Performance Profiling
-A job that *runs* on a TPU or GPU isn't necessarily using it well.
-Profiling shows where the accelerator time actually goes. This guide
-covers **XProf** — the profiler for XLA workloads — and how to capture a
-trace from a Kinetic job and view it on your own machine.
+A job that runs on a TPU or a GPU does not always use the accelerator
+well. A performance profile (a trace) shows where the accelerator time
+goes. This page covers
+**XProf**, the profiler for XLA workloads. It shows how to capture a
+trace inside a Kinetic job and how to keep the trace after the pod ends.
+It also shows how to view the trace on your own machine.
-Because a Kinetic pod scales to zero the moment your job finishes, the
-workflow is: **capture inside the function, write the trace to
-`KINETIC_OUTPUT_DIR` (durable GCS), then view it locally.**
+The pod filesystem is gone when your function returns. The workflow
+therefore has three parts:
+
+1. Capture the trace inside the decorated function.
+2. Write the trace under `KINETIC_OUTPUT_DIR`, which is a Cloud Storage
+ location that stays after the pod ends.
+3. View the trace on your machine.
## What XProf is
XProf is the open-source accelerator profiler from
-[OpenXLA](https://openxla.org/xprof) (formerly the TensorBoard "profile"
-plugin). It runs standalone or as a TensorBoard tab, reads hardware
-counters rather than wall-clock timers, and adds little overhead during
-the capture window.
+[OpenXLA](https://openxla.org/xprof). It was formerly the TensorBoard
+"profile" plugin. XProf runs as a standalone tool or as a TensorBoard
+tab. It reads accelerator hardware counters, so the capture window adds
+little overhead to the job.
-It is an **XLA** profiler, and that decides which tool you capture with:
+XProf is an **XLA** profiler. The framework of your job decides which
+tool captures the trace and which tool shows it:
-| Backend / workload | Capture with | View in |
+| Backend or workload | Capture with | View in |
| --- | --- | --- |
-| Keras-on-JAX (Kinetic default), native JAX | `jax.profiler` | XProf |
-| Keras-on-TensorFlow | `TensorBoard(profile_batch=…)` | XProf |
+| Keras on JAX (the Kinetic default), native JAX | `jax.profiler` | XProf |
+| Keras on TensorFlow | `TensorBoard(profile_batch=…)` | XProf |
| PyTorch/XLA (`torch_xla`) | `torch_xla.debug.profiler` | XProf |
| Native PyTorch (eager CUDA) | `torch.profiler` | Perfetto |
-Native eager PyTorch doesn't compile through XLA, so it uses
-`torch.profiler` and views in Perfetto — not XProf.
+Native eager PyTorch does not compile through XLA. For that workload,
+capture with `torch.profiler` and view the trace in Perfetto, not in
+XProf.
## Capture a profile
-Capture a handful of steps **after a warm-up**, inside the decorated
-function, and write to `$KINETIC_OUTPUT_DIR/profile` so the trace
-survives the pod (see [Checkpointing](checkpointing.md)). The job needs
-**no extra packages** — `jax.profiler` ships with JAX, and `xprof` is a
-*local viewer*, not a job dependency. If a profiling job does need an
-extra package, add it to a `requirements.txt` in that script's own
-directory (see [Dependencies](dependencies.md)).
+Follow this pattern inside the decorated function:
+
+:::{container} kinetic-steps
+1. Run one or more warm-up steps first. The warm-up keeps the XLA
+ compilation out of the trace.
+2. Start the trace with `jax.profiler.trace(trace_dir)`, where
+ `trace_dir` is `$KINETIC_OUTPUT_DIR/profile`.
+3. Run a small number of steps inside the trace region.
+4. Call `block_until_ready()` on the result before the trace region
+ closes.
+5. Print the trace path, so that you can find the trace later.
+:::
+
+The job needs **no extra packages**. `jax.profiler` is part of JAX, and
+JAX is in the image that Kinetic builds. `xprof` is a viewer that you
+install on your own machine, not a job dependency. If a profiling job
+does need an extra package, add the package to a `requirements.txt` in
+the directory of that script. See [Dependencies](dependencies.md).
-This is a deliberately minimal demo — a tiny JAX training loop — but the
-same capture pattern drops into any job (real training, KerasHub
-fine-tuning, vLLM serving, multi-host runs):
+The example below is a small JAX training loop. The same capture pattern
+works in a full training job, in a KerasHub fine-tuning job, and in a
+multi-host run.
```{literalinclude} ../../examples/jax_profiling_demo.py
:language: python
:caption: examples/jax_profiling_demo.py
```
-For **Keras-on-JAX**, it's the same idea: wrap a short `model.fit(...)`
-in `with jax.profiler.trace(trace_dir):` after a warm-up epoch.
+The example uses `accelerator="tpu-v5litepod-1x1"`, a single TPU v5e
+chip. `1x1` is the topology spelling of `tpu-v5litepod-1`; see
+[Accelerators](../accelerators.md). Your cluster needs a node pool for that accelerator. If the cluster
+has no such node pool, change the `accelerator=` argument to an
+accelerator that the cluster has. See
+[Clusters and Node Pools](clusters.md).
+
+For **Keras on JAX**, use the same pattern. Run one warm-up epoch. Then
+wrap a short `model.fit(...)` call in
+`with jax.profiler.trace(trace_dir):`.
:::{note}
-JAX dispatches asynchronously, so always `block_until_ready()` inside the
-trace region — otherwise the profiler can close before the device work
-lands. Keep the window to a few steps; traces grow fast.
+JAX dispatches work asynchronously. If you do not call
+`block_until_ready()` inside the trace region, the trace can close before
+the device work completes. The trace is then incomplete. Also keep the
+window to a few steps, because a trace grows fast.
:::
-Other backends: **native PyTorch** uses `torch.profiler` with
-`tensorboard_trace_handler(trace_dir)` (view in Perfetto);
-**PyTorch/XLA** uses `torch_xla.debug.profiler`, which produces
-XProf-readable traces.
+For **native PyTorch**, use `torch.profiler` with
+`tensorboard_trace_handler(trace_dir)`, and view the trace in Perfetto.
+For **PyTorch/XLA**, use `torch_xla.debug.profiler`, which writes traces
+that XProf reads.
+
+## Where the trace goes
+
+Kinetic sets `KINETIC_OUTPUT_DIR` in the pod to a per-job prefix in the
+jobs bucket of the cluster:
+
+```text
+gs://{project}-kn-{cluster}-jobs/outputs/{job_id}
+```
+
+The example therefore writes the trace to
+`gs://{project}-kn-{cluster}-jobs/outputs/{job_id}/profile`, and prints
+that path in the job log. The job ID has the form `job-a1b2c3d4`. For a
+detached job, `job.job_id` gives the job ID. The command `kinetic jobs
+list` lists the jobs of the cluster.
+
+Retention has two rules:
+
+- The trace survives result collection. After a blocking call or a
+ `result()` call, Kinetic deletes only the job artifacts under
+ `gs://{jobs bucket}/{job_id}/`. Kinetic does not delete anything under
+ `outputs/`.
+- The jobs bucket deletes objects that are older than 30 days. To keep a
+ trace for longer, pass `output_dir="gs://your-bucket/path"` to the
+ decorator. See [Outputs and Checkpoints](checkpointing.md).
## View the trace
-Install the viewer and point it at the trace path your job printed:
+Install the viewer on your machine. Then point the viewer at the trace
+path that the job printed:
```bash
-pip install xprof gcsfs # gcsfs lets XProf read gs:// directly
-xprof --logdir gs://-kn--jobs/outputs//profile --port 6006
-# --logdir is the directory that contains plugins/ (the path your job printed)
+pip install xprof gcsfs # gcsfs lets XProf read gs:// paths directly
+xprof --logdir gs://{project}-kn-{cluster}-jobs/outputs/{job_id}/profile --port 6006
+# --logdir is the directory that contains plugins/ (the path that the job printed)
# then open http://localhost:6006
```
-Or copy it down first — `gcloud storage cp -r ./trace` and
-`--logdir ./trace`. In the UI, use the tool dropdown: **Trace Viewer**
-for the step timeline, **Overview Page** for the summary. (The
-**Capture Profile** button does live, on-demand capture against a
-running profiler server, so it isn't used here — the trace was already
-captured inside the job.)
+You can also copy the trace to your machine first. Run
+`gcloud storage cp -r {trace-path} ./trace`. Then start XProf with
+`--logdir ./trace`.
+
+In the XProf user interface, select a tool from the tool list. Start
+with **Overview Page** for the summary. Then open **Trace Viewer** for
+the step timeline. Do not use the **Capture Profile** button. That
+button captures a live profile from a running profiler server. In this
+workflow, the job already captured the trace.
## What the tools show
-- **Overview Page** — top-level summary; whether you're host- or
- device-bound. Start here.
-- **Trace Viewer** — per-event timeline across host / TPU / GPU; where
- you spot gaps and stalls.
-- **Roofline** — memory-bound vs. compute-bound, which decides your
- optimization strategy.
-- **Framework / HLO Op Stats** — cost by framework op and by compiled
- HLO op.
-- **Memory Viewer / Profile** — usage over time and at peak; first stop
- after an OOM.
+- **Overview Page** — the top-level summary. It shows whether the job is
+ host-bound or device-bound. Start here.
+- **Trace Viewer** — the per-event timeline across the host, the TPU, and
+ the GPU. Use it to find gaps and stalls.
+- **Roofline** — memory-bound versus compute-bound. The result decides
+ the optimization strategy.
+- **Framework Op Stats** and **HLO Op Stats** — the cost by framework
+ operation and by compiled HLO operation.
+- **Memory Viewer** and **Memory Profile** — memory usage over time and
+ at the peak. Open these tools first after an out-of-memory error.
- **Megascale Stats** — cross-slice (DCN) communication on multi-host
[Pathways](distributed_training.md) runs.
@@ -104,28 +161,29 @@ captured inside the job.)
:link: cost_optimization
:link-type: doc
-A profile shows *where* the accelerator-hours go; that guide shows how
-to cut the bill.
+A profile shows where the accelerator hours go. That page shows how to
+reduce the cost.
:::
-:::{grid-item-card} {octicon}`database;1em` Checkpointing and Outputs
+:::{grid-item-card} {octicon}`database;1em` Outputs and Checkpoints
:link: checkpointing
:link-type: doc
-How `KINETIC_OUTPUT_DIR` keeps your trace after the pod exits.
+How `KINETIC_OUTPUT_DIR` keeps your trace after the pod ends, and how to
+send outputs to your own bucket.
:::
:::{grid-item-card} {octicon}`package;1em` Dependencies
:link: dependencies
:link-type: doc
-Per-job `requirements.txt`.
+How to add a package to one job with a `requirements.txt`.
:::
:::{grid-item-card} {octicon}`cpu;1em` Distributed Training
:link: distributed_training
:link-type: doc
-Multi-host runs where Megascale Stats applies.
+Multi-host runs, where Megascale Stats applies.
:::
::::
diff --git a/docs/guides/reservations.md b/docs/guides/reservations.md
index 46b637dd..3bc52db3 100644
--- a/docs/guides/reservations.md
+++ b/docs/guides/reservations.md
@@ -1,10 +1,58 @@
# Capacity Reservations
-GCP on-demand capacity for newer accelerators (TPU v6e, H100) is not guaranteed and may fail with `FailedScaleUp` errors. A GCP capacity reservation guarantees hardware is available when your node pool scales up.
-
-## GPU Reservations
-
-GPU reservations can be created directly through the GCP CLI. For example, for a single H100:
+Google Cloud does not guarantee on-demand capacity for newer
+accelerators, for example TPU v6e and H100. When a node pool scales up
+and the zone has no free VM, GKE reports a `FailedScaleUp` event and
+the job stays `PENDING`. A capacity reservation guarantees that the
+hardware is available when the pool scales up. This page shows how to
+get a reservation and how to bind it to a Kinetic node pool. Read this
+page if your jobs stay `PENDING` because of capacity, or if you need
+hardware at a fixed time.
+
+## Before you start
+
+- You need a cluster and an active profile. See
+ [Getting Started](../getting_started.md). The profile supplies the
+ project and the zone to `kinetic pool add`.
+- A reservation applies to one zone. Create the reservation in the zone
+ of your cluster. `kinetic config` prints the zone of the active
+ profile.
+- A reservation applies to one machine type. The machine type must be
+ the type that Kinetic uses for your accelerator. The
+ [Accelerators](../accelerators.md) page lists the machine type for each
+ accelerator name.
+
+## How a reservation works with Kinetic
+
+`kinetic pool add --reservation NAME` sets a `SPECIFIC_RESERVATION`
+affinity on the node pool, with the key
+`compute.googleapis.com/reservation-name` and the value `NAME`. The
+cluster autoscaler then creates the nodes of that pool from your
+reservation only. The pool does not compete for on-demand capacity.
+
+Kinetic does not validate the reservation. Kinetic passes the name to
+GKE unchanged. If the name is wrong, the scale-up fails and the job
+stays `PENDING`. The same happens if the zone or the machine type of the
+reservation does not match the pool.
+
+A reservation has a fixed number of VMs. Make sure that the reservation
+covers the number of VMs that the pool uses at one time:
+
+- A GPU job uses one VM. A GPU pool scales up to `--min-nodes` plus 10
+ VMs.
+- A single-host TPU job uses one VM.
+- A multi-host TPU job uses one VM per host. `tpu-v6e-16` uses 4 VMs of
+ type `ct6e-standard-4t`. A TPU pool scales up to `--min-nodes` plus
+ one slice.
+
+With the default `--min-nodes 0`, the pool creates no VM until a job
+runs. No VM consumes the reservation while the pool is idle.
+
+## Step 1: Create a GPU reservation
+
+Create a GPU reservation with the `gcloud` command. This example
+reserves one VM with one H100, which is the machine type
+`a3-highgpu-1g` for `accelerator="gpu-h100"`:
```bash
gcloud compute reservations create my-h100-reservation \
@@ -14,25 +62,69 @@ gcloud compute reservations create my-h100-reservation \
--project=your-project-id
```
-See the [GCP reservations documentation](https://cloud.google.com/compute/docs/instances/reservations-overview) for the full list of supported machine types and options.
+`gcloud` does not read the Kinetic profile. The `gcloud` commands on
+this page therefore name the project and the zone. Use the zone of your
+cluster.
+
+Two options are useful:
+
+- `--vm-count=N` reserves N VMs. Use the number of VMs from the section
+ above.
+- `--require-specific-reservation` makes sure that only a VM that names
+ the reservation can consume it. Without this option, any VM of the
+ same machine type in the project can consume the reservation.
-## TPU Reservations
+For a machine type in the `n1` family, for example `n1-standard-4` for
+one T4, add the accelerator to the reservation:
+`--accelerator=type=nvidia-tesla-t4,count=1`.
-TPU machine families do not support self-service reservations via `gcloud compute reservations create`. TPU capacity reservations must be requested through [Google Cloud support](https://cloud.google.com/support). Once approved, the reservation name will be available in your project for the specified zone and TPU type.
+See the [Compute Engine reservations
+documentation](https://cloud.google.com/compute/docs/instances/reservations-overview)
+for the full list of machine types and options.
-## Using a Reservation with Kinetic
+## Step 2: Request a TPU reservation
-Once you have a reservation name (from either path above), pass it to `kinetic pool add`:
+Google Cloud provides TPU capacity as a *future reservation*, not
+through `gcloud compute reservations create`. Two paths exist:
+
+- For a term of up to 90 days, request a future reservation in calendar
+ mode yourself, with the Google Cloud CLI or the console. Google Cloud
+ supports TPU v5p, TPU v6e, and TPU7x in calendar mode.
+- For a longer term, or for a TPU type that calendar mode does not
+ support, contact Cloud Sales or your account team.
+
+See [About Cloud TPU
+reservations](https://docs.cloud.google.com/tpu/docs/about-tpu-reservations)
+for the current list of TPU types and terms. In the request, give the
+zone, the TPU version, and the number of chips. For example,
+`tpu-v6e-16` needs 16 TPU v6e chips, which is 4 VMs of type
+`ct6e-standard-4t`, in one zone.
+
+When the reservation period starts, Google Cloud creates the reservation
+in your project. Find the name and the zone with:
+
+```bash
+gcloud compute reservations list --project=your-project-id
+```
+
+## Step 3: Bind the reservation to a node pool
+
+Pass the reservation name to `kinetic pool add`. The active profile
+supplies the project and the zone:
```bash
kinetic pool add \
- --accelerator tpu-v6e-8 \
- --reservation my-reservation-name \
- --project your-project-id
+ --accelerator gpu-h100 \
+ --reservation my-h100-reservation
```
-Kinetic sets `SPECIFIC_RESERVATION` affinity on the node pool so the autoscaler consumes nodes from your reservation instead of competing for on-demand capacity.
+For a TPU reservation, use the matching TPU accelerator name:
+```bash
+kinetic pool add \
+ --accelerator tpu-v6e-16 \
+ --reservation my-v6e-reservation
+```
## Verify the Reservation
Kinetic records the reservation in the cluster state. Later `kinetic pool add`, `kinetic pool remove`, and `kinetic up` commands keep the reservation on the pool.
@@ -45,13 +137,59 @@ kinetic pool list --project your-project-id
The `Reservation` row shows the reservation name. Kinetic does not show this row for a pool that uses on-demand or Spot capacity.
+## Cleaning Up
+
+The `KINETIC_RESERVATION` environment variable sets the same value for
+one command. See [Configuration](../configuration.md).
+
:::{note}
-Kinetic records the reservation when it creates the pool. A pool from an earlier version of Kinetic has no such record, and the `Reservation` row is absent. The next `kinetic pool add` or `kinetic pool remove` command removes the reservation affinity from that pool. To keep the reservation, remove the pool and add it again with `--reservation`.
+You cannot combine `--reservation` with `--spot`. A Spot VM cannot
+consume a reservation, and `kinetic pool add` rejects the two flags
+together. Create a separate pool for Spot jobs.
:::
-## Cleaning Up
+Every job that lands on the pool uses the reservation. You do not
+change the decorator. A job with `@kinetic.run(accelerator="gpu-h100")`
+runs on any pool of the cluster with that accelerator. If the cluster has
+a reserved H100 pool and a second H100 pool without a reservation, a job
+can land on either pool. Keep one pool per accelerator, unless the pools
+differ in Spot or in reservation and you accept that a job can land on
+either pool. See [Clusters and Node Pools](clusters.md).
+
+:::{warning}
+Kinetic does not save the reservation name, or the Spot setting, of a
+pool in its infrastructure state.
+A later `kinetic pool add`, `kinetic pool remove`, or `kinetic up`
+re-applies every existing pool from that state, without the
+reservation. The reserved pool then loses its reservation affinity.
+After you run one of these commands, remove the reserved pool and add
+it again with `--reservation`.
+:::
+
+## Check that a job uses the reservation
+
+Run a job with the accelerator of the pool. If the job stays `PENDING`
+for more than 10 minutes, do these checks:
+
+1. Describe the reservation:
-Remove the reservation when you are done to avoid ongoing charges:
+ ```bash
+ gcloud compute reservations describe my-h100-reservation \
+ --zone=us-central1-a \
+ --project=your-project-id
+ ```
+
+ Make sure that the zone and the machine type match the pool, and
+ that the reservation has a free VM.
+2. Run `kinetic pool list`. Make sure that a pool for the accelerator
+ exists.
+3. See [Troubleshooting](../troubleshooting.md) for the quota checks.
+
+## Clean up
+
+Google Cloud bills a reservation for the reserved machine type while
+the reservation exists, also when no VM uses it. Delete the reservation
+when you no longer need the hardware:
```bash
gcloud compute reservations delete my-h100-reservation \
@@ -59,8 +197,19 @@ gcloud compute reservations delete my-h100-reservation \
--project=your-project-id
```
+Also remove the node pool that names the reservation. After you delete
+the reservation, that pool cannot create nodes:
+
+```bash
+kinetic pool list
+kinetic pool remove gpu-h100-1a2b
+```
+
:::{note}
-Reservations accrue charges based on the reserved machine type regardless of whether VMs are running.
+The reservation charge starts when Google Cloud creates the reservation.
+The charge stops when you delete the reservation, or when a future
+reservation reaches its end time. A pool with `--min-nodes 0` does not
+stop the charge, because Google Cloud holds the reserved VMs for you.
:::
## Related pages
@@ -72,20 +221,20 @@ Reservations accrue charges based on the reserved machine type regardless of whe
:link: cost_optimization
:link-type: doc
-When to choose reservations over spot or on-demand.
+When to select a reservation, Spot VMs, or on-demand capacity.
:::
-:::{grid-item-card} {octicon}`server;1em` Multiple Clusters
+:::{grid-item-card} {octicon}`server;1em` Clusters and Node Pools
:link: clusters
:link-type: doc
-Isolating reserved capacity in its own cluster.
+Add, list, and remove node pools, and run more than one cluster.
:::
:::{grid-item-card} {octicon}`cpu;1em` Accelerators
:link: ../accelerators
:link-type: doc
-Which accelerators benefit most from reservations.
+The machine type and the host count behind each accelerator name.
:::
::::
diff --git a/docs/guides/vllm_tpu.md b/docs/guides/vllm_tpu.md
index 25ba9727..b07c1122 100644
--- a/docs/guides/vllm_tpu.md
+++ b/docs/guides/vllm_tpu.md
@@ -1,55 +1,195 @@
-# Running vLLM on TPU with Kinetic
+# Running vLLM on TPU
-This guide explains how to run vLLM (versatile Large Language Model) serving and inference on Cloud TPUs using the Kinetic framework.
+This page shows how to run [vLLM](https://docs.vllm.ai/) inference on a
+Cloud TPU with Kinetic. You add one dependency file, forward one
+environment variable, and run the example script. Read this page if you
+want to generate text with a large language model, such as Llama 3.1, on
+a TPU slice.
-## Overview
+## Before you start
-Kinetic allows you to easily offload heavy vLLM workloads to Cloud TPUs. This is particularly useful for serving large models like Llama 3.1 that require significant compute and memory.
+You need these things:
-## Prerequisites
+- A Kinetic cluster and an active profile. `kinetic init` creates both.
+ See [Getting Started](../getting_started.md).
+- A node pool that matches the accelerator of the example. The example
+ uses `accelerator="tpu-v5litepod"`, a 4-chip TPU v5e slice on one host.
+ Run `kinetic pool list`. If the list has no matching pool, add one:
-1. **Kinetic Cluster**: You need a provisioned Kinetic cluster with TPU nodes (e.g., `v5litepod`).
-2. **Hugging Face Token**: If you are using gated models like Llama 3.1, you need a Hugging Face token (`HF_TOKEN`) with access to the model.
+ ```bash
+ kinetic pool add --accelerator tpu-v5litepod-4
+ ```
-## Configuration
+- A Hugging Face token if the model is gated. The example uses
+ `meta-llama/Llama-3.1-8B`, which is a gated model. Request access on
+ the model page. Then create a token in your Hugging Face account
+ settings.
-To run vLLM successfully on TPU via Kinetic, you need to handle dependencies and environment variables properly.
+## Step 1: Add `vllm-tpu` to a dependency file
-### 1. Dependencies
-
-You need to ensure `vllm-tpu` is installed in the remote container. You can do this by creating a `requirements.txt` file in the directory of your script containing:
+Save a file with the name `requirements.txt` in the same directory as
+your script:
```text
vllm-tpu
```
-Kinetic will detect this file and build a container with vLLM installed.
-
-### 2. Environment Variables
-
-You must pass the following environment variables to ensure correct execution:
-
-| Variable | Purpose |
-|----------|---------|
-| `VLLM_TARGET_DEVICE="tpu"` | Tells vLLM to target TPU. |
-| `VLLM_USE_V1="0"` | Forces vLLM to use the stable v0 engine (recommended for TPU currently). |
-| `JAX_PLATFORMS="tpu,cpu"` | Allows JAX to see both TPU and CPU backends, avoiding initialization crashes. |
+When you call the decorated function, Kinetic looks for a dependency
+file. Kinetic starts the search in the directory of the script and
+continues in the parent directories. A `requirements.txt` next to the
+script is therefore the first file that Kinetic finds. See
+[Dependencies](dependencies.md) for the search rules.
+
+Kinetic then builds a container image that contains `vllm-tpu` and
+caches the image. The first run waits for that build. Later runs with
+an unchanged file reuse the image.
+
+:::{note}
+Do not add `jax`, `jaxlib`, or `libtpu` to the file. The image that
+Kinetic builds already contains JAX with the TPU runtime. If the file
+has such a line, Kinetic removes the line and logs a warning. See
+[JAX and accelerator runtimes](dependencies.md#jax-and-accelerator-runtimes).
+:::
-:::{tip}
-You can use `capture_env_vars` in the `@kinetic.run` decorator to pass these from your local environment.
+## Step 2: Review the environment variables
+
+The example uses two kinds of environment variable. The vLLM variables
+are constants, and the script sets them inside the function. The
+Hugging Face token is a secret, and the script forwards it from your
+shell. You do not set anything in this step. Step 4 sets the token.
+
+### The vLLM variables
+
+The function sets three variables with `os.environ` before the
+`from vllm import LLM, SamplingParams` line. The values therefore exist
+in the pod process when vLLM loads.
+
+| Variable | Value in the example | Purpose |
+| -------- | -------------------- | ------- |
+| `VLLM_TARGET_DEVICE` | `tpu` | Selects the TPU backend of vLLM. |
+| `JAX_PLATFORMS` | `tpu,cpu` | Lists the JAX backends. Kinetic sets `JAX_PLATFORMS=tpu` in every TPU pod. The example replaces that value with `tpu,cpu`. |
+| `VLLM_USE_V1` | `0` | Selects the vLLM engine version. |
+
+You do not set these variables in your shell, and you do not forward
+them with `capture_env_vars`. The code carries them.
+
+### The Hugging Face token
+
+One value comes from your shell: `HF_TOKEN`. The decorator lists the
+name in `capture_env_vars=["HF_TOKEN"]`. When you call the function,
+Kinetic reads `HF_TOKEN` from your environment and copies the value into
+the job payload. The pod applies the value before the pod calls your
+function, so vLLM can download a gated model. If your shell has no
+`HF_TOKEN`, Kinetic captures nothing and logs no error. The pod then
+has no token, and the download of a gated model fails. See
+[Forward Environment Variables](env_vars.md).
+
+:::{warning}
+The name `HF_TOKEN` contains `TOKEN`, so Kinetic logs a warning when it
+captures the value. Kinetic stores the value in plaintext inside the job
+payload in the jobs bucket, and every job pod in the cluster can read
+that bucket. Kinetic deletes the payload when a blocking call or a
+`result()` call collects a successful result with the default cleanup.
+Otherwise the payload stays until the 30-day lifecycle rule of the
+bucket deletes it. If you intend to forward the token, the warning needs
+no action. See [Secrets](env_vars.md#secrets) and
+[Security](../security.md).
:::
-## Example
+## Step 3: Read the example
+
+Save this script as `vllm_demo.py`, next to the `requirements.txt` from
+step 1:
```{literalinclude} ../../examples/vllm_demo.py
:language: python
```
-## Running the Example
+Four points in the script matter:
+
+- **The accelerator.** `accelerator="tpu-v5litepod"` resolves to the
+ default v5e slice: 4 chips on one host, topology `2x2`. The string
+ `tpu-v5litepod-4` names the same slice. See
+ [Accelerators](../accelerators.md).
+- **The parallelism.** `tensor_parallel_size=4` matches the 4 chips of
+ the slice.
+- **The import.** The `from vllm import ...` line is inside the
+ function. The pod runs that line, and the image on the pod contains
+ vLLM. Your machine does not need vLLM.
+- **The result.** The function prints the completions in the pod, and
+ Kinetic streams the pod log to your terminal. If you want the
+ completions in your local process, return them from the function.
-To run the script, set the required environment variables locally so they get captured:
+## Step 4: Run the script
+
+Set `HF_TOKEN` for the command and run the script:
```bash
-HF_TOKEN=your-hf-token \
-python3 your_script.py
+HF_TOKEN=your-hf-token python vllm_demo.py
```
+
+The active profile supplies the project, the zone, and the cluster. You
+do not pass them. Kinetic then does these things:
+
+:::{container} kinetic-steps
+1. **Package.** Kinetic captures `HF_TOKEN`, serializes the function, and
+ archives the directory of the script.
+2. **Build.** Kinetic builds an image with `vllm-tpu` on the first run,
+ or reuses the cached image.
+3. **Schedule.** The cluster autoscaler starts a TPU v5e node in the
+ matching node pool if no free node exists.
+4. **Run.** The pod downloads the model weights from Hugging Face, loads
+ the model on the TPU, and prints the completions. Kinetic streams the
+ log lines to your terminal.
+5. **Collect.** The call returns when the function ends. Kinetic deletes
+ the job resources.
+:::
+
+:::{note}
+The first run waits for the image build. Every run waits for a node
+start if the node pool has scaled to zero. Every run also downloads the
+model weights again, because the pod filesystem does not persist between
+jobs. If you expect a long run, call `run_vllm_inference.run_async()`
+instead of the blocking call, and collect the result later. See
+[Detached Jobs](async_jobs.md).
+:::
+
+## Change the model or the slice
+
+- **The model.** Change `model_id` to another Hugging Face model. If the
+ model is gated, your token must have access to the model.
+- **The slice.** Change `accelerator` to another single-host slice, for
+ example `tpu-v5litepod-8` (8 chips on one host). Set
+ `tensor_parallel_size` to the number of chips in the slice. Add a node
+ pool that matches the new accelerator with `kinetic pool add`.
+- **Larger slices.** A v5e slice with 16 or more chips spans more than
+ one host. Kinetic runs such a job on the Pathways backend, one pod per
+ host. See [Distributed Training](distributed_training.md).
+
+## Related pages
+
+::::{grid} 1 1 2 2
+:gutter: 3
+
+:::{grid-item-card} {octicon}`key;1em` Forward Environment Variables
+:link: env_vars
+:link-type: doc
+
+`capture_env_vars`, wildcards, and how Kinetic handles a secret such as
+`HF_TOKEN`.
+:::
+
+:::{grid-item-card} {octicon}`package;1em` Dependencies
+:link: dependencies
+:link-type: doc
+
+How Kinetic finds `requirements.txt`, and which lines Kinetic filters.
+:::
+
+:::{grid-item-card} {octicon}`cpu;1em` Accelerators
+:link: ../accelerators
+:link-type: doc
+
+Every accelerator name, and the topology behind each TPU name.
+:::
+::::
diff --git a/docs/index.md b/docs/index.md
index 1861a021..01138bba 100644
--- a/docs/index.md
+++ b/docs/index.md
@@ -5,28 +5,42 @@
:hidden:
getting_started
-guides/execution_modes
+concepts
troubleshooting
guides/faq
```
```{toctree}
-:caption: Core Workflows
+:caption: Run Jobs
:hidden:
guides/data
guides/checkpointing
-guides/packaging
guides/dependencies
guides/env_vars
-guides/profiles
+guides/async_jobs
+guides/batched_jobs
guides/debugging
+guides/distributed_training
guides/profiling
+```
+
+```{toctree}
+:caption: Manage Infrastructure
+:hidden:
+
+guides/profiles
+guides/clusters
guides/cost_optimization
-guides/distributed_training
-guides/vllm_tpu
+guides/reservations
+```
+
+```{toctree}
+:caption: Advanced
+:hidden:
+
guides/containers
-guides/advanced
+guides/packaging
```
```{toctree}
@@ -38,6 +52,7 @@ examples/jax_training
examples/pytorch_training
examples/gemma4_finetuning
examples/llm_finetuning
+guides/vllm_tpu
examples
```
@@ -47,8 +62,8 @@ examples
api
cli
-accelerators
configuration
+accelerators
security
```
@@ -62,9 +77,9 @@ code-of-conduct
```
:::{container} kinetic-hero
-Run any Python function on a cloud TPU or GPU with one decorator. No
-infrastructure to wire up, no images to build by hand, no multi-host
-boilerplate.
+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.
:::
::::{container} kinetic-hero-buttons
@@ -74,10 +89,10 @@ boilerplate.
Get started
:::
-:::{button-ref} examples
+:::{button-ref} concepts
:color: secondary
-Browse examples
+How it works
:::
::::
@@ -85,7 +100,7 @@ Browse examples
import kinetic
-@kinetic.run(accelerator="tpu-v6e-8")
+@kinetic.run(accelerator="tpu-v5litepod-4")
def train_model():
import keras
@@ -94,155 +109,114 @@ def train_model():
return model.history.history["loss"][-1]
-final_loss = train_model() # runs on a TPU v6e-8 slice
+final_loss = train_model() # runs on a 4-chip TPU v5e slice
```
## Start here
-Three entry points cover what most new users need first.
+Read these three pages in order. They take about 30 minutes. When
+something does not work, see [Troubleshooting](troubleshooting.md), and
+for short answers see the [FAQ](guides/faq.md).
::::{grid} 1 1 3 3
:gutter: 3
-:::{grid-item-card} {octicon}`rocket;1em` Your first run
+:::{grid-item-card} {octicon}`rocket;1em` 1. Getting Started
:link: getting_started
:link-type: doc
-Install, point at a cluster, and run a real Keras job in minutes.
+Install Kinetic, run `kinetic init`, and run a Keras job on a TPU.
:::
-:::{grid-item-card} {octicon}`clock;1em` Long-running jobs
-:link: guides/async_jobs
+:::{grid-item-card} {octicon}`light-bulb;1em` 2. How Kinetic Works
+:link: concepts
:link-type: doc
-Switch from blocking `run()` to detached `run_async()` for jobs that
-take hours.
+The vocabulary, the job lifecycle, and where your code, data, and
+results go.
:::
-:::{grid-item-card} {octicon}`database;1em` Data and checkpoints
-:link: guides/data
+:::{grid-item-card} {octicon}`code-square;1em` 3. Examples
+:link: examples
:link-type: doc
-Ship local files in, write durable artifacts back out via
-`KINETIC_OUTPUT_DIR`.
+Runnable scripts, from a first run to multi-host LLM fine-tuning.
:::
::::
-## How Kinetic works
-
-Five short phases on every job.
+## What happens on every job
:::{container} kinetic-steps
-1. **Discover.**
- Kinetic captures your function, its package root, and the
- `Data(...)` arguments. Kinetic reads `requirements.txt` or
- `pyproject.toml`. See
- [What Ships to the Pod](guides/packaging.md).
+1. **Package.**
+ Kinetic serializes your function and archives your project source.
+ `Data(...)` arguments upload one time, keyed by content.
-2. **Build or fetch.**
- A container image is produced — built with your dependencies
- (bundled mode) or pulled from a published base (prebuilt mode). See
- [Execution Modes](guides/execution_modes.md).
+2. **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.
3. **Schedule.**
- A Kubernetes resource (a `Job` for single-host workloads, a
- `LeaderWorkerSet` for multi-host TPU jobs on the Pathways backend)
- is submitted to your GKE cluster. The autoscaler provisions
- accelerator nodes if needed.
+ 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.
4. **Run.**
- Your function executes inside the pod with `KINETIC_OUTPUT_DIR`
- set; logs stream back to your terminal.
+ The pod runs your function with `KINETIC_OUTPUT_DIR` set. The logs
+ stream to your terminal.
5. **Collect.**
- The return value is serialized to GCS and pulled back to your local
- process. `@kinetic.run()` cleans up the pod and GCS artifacts as
- soon as the result is collected. `run_async()` leaves the pod
- running until you call `.result()` or `.cleanup()` on the returned
- `JobHandle` — important to remember on expensive accelerators.
+ 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.
:::
-## Choose your execution mode
-
-Three modes control how dependencies get into the container. See
-[Execution Modes](guides/execution_modes.md) for the full
-recommendation matrix and per-mode startup expectations.
+## Explore the guides
::::{grid} 1 1 3 3
:gutter: 3
-:::{grid-item-card} Bundled
-:link: guides/execution_modes
-:link-type: doc
-
-Kinetic builds a custom image with your deps baked in. Best for stable
-workflows and reproducible runs.
-
-+++
-{bdg-success}`default`
-:::
-
-:::{grid-item-card} Prebuilt
-:link: guides/execution_modes
-:link-type: doc
-
-Pulls a published base image, installs your deps at pod startup. Best
-for fast iteration when deps change often.
-:::
-
-:::{grid-item-card} Custom image
-:link: guides/execution_modes
-:link-type: doc
-
-Bring your own image URI. Best when you need custom system libraries
-or a corporate-vetted base.
-:::
-::::
-
-## Explore the docs
-
-::::{grid} 1 1 3 3
-:gutter: 3
-
-:::{grid-item-card} {octicon}`file-directory;1em` Working with data
+:::{grid-item-card} {octicon}`file-directory;1em` Working with Data
:link: guides/data
:link-type: doc
-Get inputs into the job and durable outputs back out.
+Get inputs into the job with `kinetic.Data(...)`.
:::
-:::{grid-item-card} {octicon}`history;1em` Checkpointing
+:::{grid-item-card} {octicon}`history;1em` Outputs and Checkpoints
:link: guides/checkpointing
:link-type: doc
-Make long jobs resumable with `KINETIC_OUTPUT_DIR`.
+Keep files and make long jobs resumable with `KINETIC_OUTPUT_DIR`.
:::
-:::{grid-item-card} {octicon}`server;1em` Distributed training
-:link: guides/distributed_training
+:::{grid-item-card} {octicon}`clock;1em` Detached Jobs
+:link: guides/async_jobs
:link-type: doc
-Scale beyond one host with the Pathways backend.
+`run_async()` for jobs that run more than a few minutes. Reattach from
+any machine.
:::
-:::{grid-item-card} {octicon}`graph;1em` Cost optimization
-:link: guides/cost_optimization
+:::{grid-item-card} {octicon}`stack;1em` Profiles
+:link: guides/profiles
:link-type: doc
-Spot capacity, autoscaling, and cleanup habits that save money.
+One saved project, zone, cluster, and namespace per cluster. Switch
+with one command.
:::
-:::{grid-item-card} {octicon}`bug;1em` Debugging
-:link: guides/debugging
+:::{grid-item-card} {octicon}`server;1em` Clusters and Node Pools
+:link: guides/clusters
:link-type: doc
-Interactive debugging and log streaming for remote jobs.
+Add accelerator pools, share a cluster with a team, and clean up.
:::
-:::{grid-item-card} {octicon}`code-square;1em` Examples
-:link: examples
+:::{grid-item-card} {octicon}`cpu;1em` Distributed Training
+:link: guides/distributed_training
:link-type: doc
-Runnable scripts from first run to multi-host LLM fine-tuning.
+Scale to a multi-host TPU slice with the Pathways backend.
:::
::::
diff --git a/docs/security.md b/docs/security.md
index 7e3b5318..cf9435be 100644
--- a/docs/security.md
+++ b/docs/security.md
@@ -1,45 +1,61 @@
-# Kinetic Security & Threat Model
+# Security and Threat Model
-Kinetic is a tool designed to seamlessly execute user-defined Python code on
-remote infrastructure (GCP/Kubernetes). Because its primary function is
-arbitrary code execution, understanding the security boundaries is critical.
+Kinetic runs user-defined Python code on remote infrastructure in your
+Google Cloud project. Its purpose is arbitrary code execution. This page
+states the security boundary, what Kinetic protects against, and what it
+does not.
-## The Security Boundary
+## The security boundary
:::{important}
-**Kinetic assumes that any user authorized to submit jobs is a trusted
-entity.** If a user is granted the IAM permissions to upload to the Kinetic GCS
-buckets and the Kubernetes RBAC permissions to create Jobs/Pods, they have the
-ability to execute arbitrary code on the cluster.
+**Kinetic treats every user who can submit jobs as a trusted user.** A
+user who has the IAM permissions to write to the jobs bucket and the
+Kubernetes permissions to create Jobs and Pods on the cluster can run
+arbitrary code on the cluster.
:::
-Kinetic **does not** attempt to sandbox, restrict, or monitor the Python code
-written by authorized users. Securing the cluster against authorized users must
-be done at the infrastructure level (e.g., network policies, minimal IAM roles
-for the GKE nodes, namespace isolation).
-
-## In-Scope Threats (What Kinetic Protects Against)
-
-1. **Payload Tampering (Man-in-the-Middle):** Kinetic relies on `cloudpickle`
- to serialize functions. Deserializing untrusted pickle data is dangerous
-(CWE-502). Kinetic protects against scenarios where an attacker gains write
-access to the GCS bucket but lacks K8s access.
- * **Mitigation:** The client computes a SHA-256 hash of the payload at
- submission time and embeds it immutably into the Kubernetes Pod Spec. The
-remote runner verifies the payload matches this hash before deserialization,
-ensuring the payload wasn't tampered with in transit or at rest in GCS.
-
-2. **Data Exfiltration via GCS:** Kinetic provisions cluster-scoped GCS
-buckets.
- * **Mitigation:** Buckets should be locked down using IAM so that only
- authorized developers and the cluster's service account can read/write.
-
-## Out-of-Scope Threats (What Kinetic Does Not Protect Against)
-
-1. **Malicious Insiders:** Kinetic does not prevent an authorized user from
-writing malicious code inside their `@kinetic.run()` function.
-2. **Container Escapes:** If a user exploits a container runtime vulnerability
-to escape the Pod, this is an infrastructure/GKE concern, not a Kinetic
-concern.
-3. **Compromised K8s Credentials:** If an attacker compromises a user's
-`kubeconfig` or GCP credentials, they inherit that user's ability to run code.
+Kinetic does not sandbox, restrict, or monitor the Python code of an
+authorized user. Protection of the cluster against its authorized users
+belongs at the infrastructure level: network policies, minimal IAM roles
+for the GKE nodes, and namespace isolation.
+
+## What Kinetic protects against
+
+1. **Payload tampering in transit or at rest.** Kinetic serializes your
+ function with `cloudpickle`. Deserialization of untrusted pickle data
+ is dangerous (CWE-502). An attacker who has write access to the jobs
+ bucket but no Kubernetes access could replace the payload.
+ *Mitigation:* the client computes a SHA-256 hash of `payload.pkl` and
+ of `context.zip` at submit time and writes both hashes into the pod
+ specification. The runner verifies both files against those hashes
+ before it deserializes anything. *Limit:* the runner does not verify
+ the other objects that it downloads from the bucket: the generated
+ `requirements.txt` in prebuilt mode, and the `Data` content in the
+ data cache. Restrict write access to the jobs bucket to the people
+ who can submit jobs.
+2. **Data exfiltration through the buckets.** Kinetic creates
+ cluster-scoped Cloud Storage buckets. *Mitigation:* the buckets use
+ uniform bucket-level access, and only the cluster service accounts get
+ IAM bindings. Grant access to developers with the same care.
+
+## What Kinetic does not protect against
+
+1. **Malicious insiders.** Kinetic does not prevent an authorized user
+ from writing malicious code inside a `@kinetic.run()` function.
+2. **Container escapes.** A container runtime vulnerability is a GKE
+ concern, not a Kinetic concern.
+3. **Compromised credentials.** An attacker who has a user's `kubeconfig`
+ or Google Cloud credentials inherits that user's ability to run code.
+
+## Secrets in jobs
+
+Values that you forward with `capture_env_vars` travel inside
+`payload.pkl` in the jobs bucket. Every job pod in the cluster can read
+that bucket. Forward only what the job needs, and prefer short-lived
+tokens. Kinetic warns when a captured name looks like a credential. See
+[Forward Environment Variables](guides/env_vars.md).
+
+Kinetic does not exclude secret files (`.env*`, `*.pem`, `id_rsa*`) from
+the source archive by default. Kinetic warns when the archive contains
+one. Add the file to `.kineticignore` if you did not intend to ship it.
+See [What Ships to the Pod](guides/packaging.md).
diff --git a/docs/troubleshooting.md b/docs/troubleshooting.md
index b42b56cb..e35e611e 100644
--- a/docs/troubleshooting.md
+++ b/docs/troubleshooting.md
@@ -1,170 +1,235 @@
# Troubleshooting
-This page is organized by symptom rather than by error message. Locate
-the section that best matches what you are observing and follow the
-guidance there.
+This page lists problems by symptom, not by error message. Find the
+section that matches what you see, and follow the steps in that section.
+The last section lists the checks that the troubleshoot path of
+`kinetic init` runs.
:::{admonition} Quick diagnostic
:class: tip
-For a quick diagnostic of common environment problems, run:
+For a fast check of the common environment problems, run:
```bash
kinetic init
```
-and choose `troubleshoot` at the prompt. It checks for missing tools,
-misconfigured credentials, and unhealthy infrastructure, and prints a
-concrete fix command for each failed check. `kinetic init` will also
-offer this path automatically when it detects that prerequisites are
-missing. The full list of categories the troubleshoot path covers is
-described at the end of this page.
+Select `troubleshoot` at the prompt. The troubleshoot path checks your
+local tools, your credentials, your Google Cloud project, and the
+cluster infrastructure. It prints a fix hint for each failed check.
+`kinetic init` offers this path directly when a prerequisite is
+missing. The troubleshoot path diagnoses only a cluster from one of your
+profiles. If you have no profile for the project, the troubleshoot path
+runs the environment checks only. See
+[What the troubleshoot path checks](#what-the-troubleshoot-path-checks).
:::
## Startup and build issues
### "Project must be specified"
-`KINETIC_PROJECT` (or `GOOGLE_CLOUD_PROJECT`) is not set. Set it once
-in your shell profile:
+Kinetic found no project. No profile is active, and no override supplies
+a project. Do one of these two things:
-```bash
-export KINETIC_PROJECT="your-project-id"
-```
+- Run `kinetic init`. Both the **Join** path and the **Create** path save
+ a profile and make it active.
+- If you already have a profile, run `kinetic profile use NAME`.
+ `kinetic profile ls` lists your profiles.
-Or pass `project=` to the decorator. See [Configuration](configuration.md).
+The environment variable `KINETIC_PROJECT` and the decorator argument
+`project=` are one-off overrides. Use them for one command or one job,
+not as the normal path. `kinetic config` shows the resolved project and
+its source. See [Profiles](guides/profiles.md) and
+[Configuration](configuration.md).
### "404 Requested entity was not found"
-A required GCP resource — usually an Artifact Registry repository or a
-GKE cluster — doesn't exist yet. Run the setup once:
+A required Google Cloud resource does not exist. Common examples are the
+Artifact Registry repository, a Cloud Storage bucket, or the GKE cluster
+itself. Run the setup:
```bash
kinetic up
```
-Or, if `up` already ran, enable the missing APIs and create the
-registry manually (uncommon):
+`kinetic up` creates the missing resources and enables the required
+APIs. A second `kinetic up` for the same cluster is safe. If the error
+remains, run the troubleshoot path of `kinetic init`. The **GCP
+Resources** group names the resource that is missing.
+
+If you enable the APIs by hand, enable all seven:
```bash
gcloud services enable compute.googleapis.com cloudbuild.googleapis.com \
artifactregistry.googleapis.com storage.googleapis.com \
- container.googleapis.com --project=$KINETIC_PROJECT
-
-gcloud artifacts repositories create "kn-${KINETIC_CLUSTER:-kinetic-cluster}" \
- --repository-format=docker --location=us \
- --project=$KINETIC_PROJECT
+ container.googleapis.com secretmanager.googleapis.com \
+ iam.googleapis.com --project
```
-### Container build is slow on first run
+### The first run takes a long time to start
+
+The first run with a given dependency file builds a container image with
+Cloud Build. The build takes about 5 to 10 minutes. Later runs with the
+same dependency file reuse the cached image and start in less than 1
+minute while a node still runs. Kinetic builds a new image when the
+dependency file, your Python minor version, the Kinetic version, or the
+accelerator category (CPU, GPU, or TPU) changes. See
+[Dependencies](guides/dependencies.md).
-The first run with a given `requirements.txt` builds a new container
-image via Cloud Build (~2–5 minutes). Subsequent runs reuse the cached
-image and start in under a minute. If you're churning dependencies
-multiple times a day and this is hurting you, see
-[Execution Modes](guides/execution_modes.md) for prebuilt mode.
+If your dependency file changes many times a day, see
+[Container Images](guides/containers.md) for a mode that installs the
+dependencies at pod start.
### Every submit uploads for a long time
-Kinetic logs the `context.zip` size on every submit. Kinetic logs a
-warning above 100 MB, and lists the five largest files. Read that warning
-first. The usual cause is a data directory, a checkpoint directory, or a
-virtualenv inside the [package root](guides/packaging.md). Wrap large
-inputs in `kinetic.Data(...)`, which Kinetic uploads one time and caches
-by content hash. As an alternative, add the paths to a `.kineticignore`
-file.
+Kinetic logs the size of `context.zip` on every submit. Above 100 MB,
+Kinetic logs a warning and lists the five largest files. Read that
+warning first. The usual cause is a data directory, a checkpoint
+directory, or a virtual environment inside the
+[package root](guides/packaging.md). Wrap large inputs in
+`kinetic.Data(...)`. Kinetic uploads such an input one time and caches
+it by content hash. As an alternative, add the paths to a
+`.kineticignore` file.
Kinetic logs a separate warning above 50 MB of *payload*. The payload
holds your arguments, the module-level globals that your function reads,
-and your first-party module code. The payload holds no files from the
-archive. Load a large object inside the function, or pass it as `Data`.
+and the code of your first-party modules. The payload holds no files
+from the archive. Load a large object inside the function, or pass it as
+`Data`.
-### Container build failures
+You can change both thresholds. Set `KINETIC_CONTEXT_SIZE_WARN_MB` or
+`KINETIC_PAYLOAD_SIZE_WARN_MB` to a number of megabytes.
-Check Cloud Build logs:
+### The container build fails
+
+Kinetic logs the build ID and a `View build:` URL when it submits the
+build. Open that URL to read the Cloud Build log. To list the recent
+builds of the project, run:
```bash
-gcloud builds list --project=$KINETIC_PROJECT --limit=5
-gcloud builds log --project=$KINETIC_PROJECT
+gcloud builds list --project --limit=5
```
-Common causes: a package in `requirements.txt` that doesn't exist,
-network issues during install, or a base image that's been updated
-since you last built. See [Dependencies](guides/dependencies.md).
+Common causes are a package name that does not exist on the index, a
+version pin that does not resolve, and a network error during the
+install. See [Dependencies](guides/dependencies.md).
## Auth and config issues
-### "Permission denied" on GCP operations
+### "Permission denied" on Google Cloud operations
-Your user (or service account) is missing IAM roles. The minimum set
-for Kinetic is `roles/storage.admin`, `roles/artifactregistry.admin`,
-`roles/container.admin`, and `roles/cloudbuild.builds.editor` on the
-project:
+Your user account, or your service account, lacks a permission on the
+project. `kinetic up` needs broad permissions, because one run creates
+many resources:
-```bash
-gcloud projects add-iam-policy-binding $KINETIC_PROJECT \
- --member="user:your-email@example.com" \
- --role="roles/storage.admin"
-```
+- A GKE cluster and its node pools.
+- Two service accounts and project-level IAM bindings for them.
+- A VPC network, a Cloud Router, and a Cloud NAT gateway.
+- An Artifact Registry repository and Cloud Storage buckets.
+- The required Google Cloud APIs, which `up` enables.
-Repeat for the other roles. `kinetic init`'s troubleshoot path flags
-missing roles by checking the actual operations that fail.
+The troubleshoot path reports a `Permission denied` on project access
+and recommends `roles/editor` or `roles/owner` on the project. Ask a
+project owner for one of these roles before you run `kinetic up`. For
+the roles that a teammate needs on the shared state bucket, see
+[Configuration](configuration.md#iam).
-### Application Default Credentials missing or expired
+### Application Default Credentials are missing or expired
+
+Log in again:
```bash
gcloud auth login
gcloud auth application-default login
```
-If you've previously set `GOOGLE_APPLICATION_CREDENTIALS` to a service
-account key, that takes precedence over user ADC.
+If you set `GOOGLE_APPLICATION_CREDENTIALS` to a service account key
+file, that key takes precedence over your user credentials.
+
+### A setting does not take effect
-### Settings aren't taking effect
+Kinetic resolves each of project, zone, cluster, and namespace in this
+order. The first value wins.
-Run `kinetic config` — it prints every config value and where it came
-from (decorator arg, CLI flag, env var, or default). The precedence
-rules are documented in [Configuration](configuration.md).
+1. The decorator argument or the CLI flag.
+2. The `KINETIC_*` environment variable.
+3. The active profile.
+4. The built-in default.
+
+Run `kinetic config`. It shows the resolved value of each setting and its
+source: an environment variable, the active profile, or the default.
+`kinetic config` cannot see a CLI flag or a decorator argument. If a
+value surprises you, look for a `KINETIC_*` variable in your shell,
+because an environment variable overrides the profile. See
+[Configuration](configuration.md).
## Scheduling and quota issues
-### Job stuck in `PENDING` for more than 10 minutes
+### A job stays in `PENDING` for more than 10 minutes
+
+The cluster cannot start a node for the pod. Check these causes in
+order:
+
+1. **No node pool matches the accelerator.** A job runs only on a node
+ pool with the same accelerator type and, for a TPU, the same
+ topology. The cluster does not create accelerator pools automatically.
+ Run `kinetic pool list`. If no pool matches, add one:
+ `kinetic pool add --accelerator tpu-v5litepod-4` (use your
+ accelerator). Kinetic logs an info line at submit when no node
+ matches, and then continues.
+2. **The Spot setting does not match.** A job with `spot=True` needs a
+ pool that you added with `--spot`. A job without `spot=True` needs an
+ on-demand pool. See the next section.
+3. **No quota for the accelerator** in your zone. Open
+ Cloud Console → IAM & Admin → Quotas and filter by your accelerator
+ type. Request more quota, or use a different zone.
+4. **No Spot capacity.** If the pool uses Spot VMs, Google Cloud can have
+ no Spot capacity at that moment. Use an on-demand pool, or try later.
+
+`kubectl describe pod ` shows the scheduling events for the
+pod. Find the pod with `kubectl get pods -l job-id=`. The
+troubleshoot path also lists pending pods and reports exhausted
+accelerator quotas in your region.
+
+### A `--spot` pool receives no jobs
+
+A pool that you added with `--spot` accepts only jobs that request Spot
+capacity. Set `spot=True` on the decorator, or add the `:spot` suffix to
+the accelerator string:
+
+```python
+@kinetic.run(accelerator="gpu-l4", spot=True)
+def train(): ...
+```
-The cluster autoscaler is trying to provision a node but can't. Two
-common reasons:
+A job without `spot=True` never runs on a Spot pool. See
+[Cost Optimization](guides/cost_optimization.md#spot-vms).
-- **No quota for the requested accelerator** in your zone. Check
- Cloud Console → IAM & Admin → Quotas, filter by your accelerator
- type. If quota is exhausted, request more or try a different zone.
-- **Spot capacity is unavailable.** If your node pool was created with
- `--spot`, GCP may have no spot capacity to allocate right now.
- Switch to on-demand or try later.
+### A multi-host TPU job fails at submit or soon after it starts
-`kinetic init`'s troubleshoot path includes a quota check that surfaces
-exhausted accelerator quotas in your region. If it doesn't flag
-anything, inspect the Cloud Console quota page directly for
-finer-grained breakdowns.
+Three causes are common:
-### Multi-host TPU job fails right after submit
+- A missing LeaderWorkerSet controller. The submit fails with
+ `LeaderWorkerSet CRD not found`.
+- A topology mismatch. Your code expects a device count that the slice
+ does not have.
+- One host that fails before the other hosts reach the first collective.
-Likely causes: topology mismatch (your code expected a different number
-of devices than the slice has), a stale Pathways context from a prior
-crashed job, or one host failing before the others can join the
-collective. See [Distributed Training](guides/distributed_training.md)
-for the full list of multi-host failure modes.
+See [Distributed Training](guides/distributed_training.md) for the full
+list of multi-host failure modes and their fixes.
## Runtime failures
-### `ModuleNotFoundError` / `ImportError` on the remote pod
+### `ModuleNotFoundError` / `ImportError` on the pod
Three different causes produce the same exception. Identify your cause
from the module name in the traceback:
| The missing module is | The cause | The fix |
| --------------------- | --------- | ------- |
-| A third-party package (`pandas`, `transformers`, and more) | The package is not in your `requirements.txt` or `pyproject.toml`. A local `pip install` does not carry over. | Add the package to the dependency file. Read the `Using dependency file: ...` log line. Kinetic walks up from the entry directory of your function, so Kinetic can select a different file than you expect. See [Dependencies](guides/dependencies.md). |
-| One of your own modules or packages (`trainer`, `mylib.utils`, and more) | The module lives outside the detected [package root](guides/packaging.md), so it never entered `context.zip`. An exclusion rule can also remove it. | Put a `pyproject.toml`, a `requirements.txt`, or a `.git` at the top of the tree that you want to ship. As an alternative, set `KINETIC_PACKAGE_ROOT`. Check your `.kineticignore` and the default exclusion list. |
-| `kinetic` itself | An old prebuilt image or custom image does not have Kinetic installed, and one of your own modules runs `import kinetic` at module scope. | Kinetic now ships your first-party modules by value, so the pod does not import them to unpickle the job. Rebuild your image: a bundled image installs `keras-kinetic` automatically. For a custom image, run `pip install keras-kinetic` in the image. |
+| A third-party package (`pandas`, `transformers`, and more) | The package is not in your `requirements.txt` or `pyproject.toml`. A local `pip install` has no effect on the pod. | Add the package to the dependency file. Read the `Using dependency file: ...` log line. Kinetic walks up from the entry directory of your function, so Kinetic can select a different file than you expect. See [Dependencies](guides/dependencies.md). |
+| One of your own modules or packages (`trainer`, `mylib.utils`, and more) | The module is outside the detected [package root](guides/packaging.md), so it never entered `context.zip`. An exclusion rule can also remove it. | Put a `pyproject.toml`, a `requirements.txt`, or a `.git` at the top of the tree that you want to ship. As an alternative, set `KINETIC_PACKAGE_ROOT`. Check your `.kineticignore` and the default exclusion list. |
+| `kinetic` itself | The image does not have Kinetic installed. Your function, or one of your own modules in the payload, references the `kinetic` module, so the pod must import `kinetic` to unpickle the job. An old base image or a custom image can lack it. | Rebuild your image. The image that Kinetic builds installs `keras-kinetic` for you. For a custom image, install `keras-kinetic` in the image. |
The traceback names the frame that failed. A failure inside
`cloudpickle.load` means that the pod needed the module to *deserialize*
@@ -179,21 +244,23 @@ Check these three points, in order:
1. **Is the file inside the [package root](guides/packaging.md)?**
Kinetic archives the package root only. A file one directory above
the root never ships.
-2. **Did an exclusion rule remove it?** Kinetic skips `.venv`,
+2. **Did an exclusion rule remove the file?** Kinetic skips `.venv`,
`node_modules`, the cache directories, and every path that matches
your `.kineticignore`. Kinetic also excludes a local path that you
wrap in `Data(...)`. Read such a path through the value that the
- `Data` argument resolves to, and not through its original location.
-3. **Is the path relative to the directory that you launched from?** The
+ `Data` argument resolves to, not through its original location.
+3. **Is the path relative to the directory from which you ran the script?** The
pod changes to the workspace directory that matches your client
working directory, so a relative path behaves as it does locally. If
your client working directory was *outside* the package root, the pod
uses the workspace root instead.
-An absolute client path is not a supported way to read a shipped file.
-The runner does create one symbolic link at the path of your client
-working directory. That link points at the workspace root. The link exists
-so the debugger can map source files. Use a relative path.
+Do not use an absolute client path to read a shipped file. The runner
+does create one symbolic link on the pod. The link is at the path of
+your client **entry directory**, which is the directory of the file that
+defines the decorated function. The link points at the workspace root,
+and it exists so that the debugger can map source files. Use a relative
+path.
### Pickle / cloudpickle errors at submit time
@@ -208,12 +275,12 @@ Your function, or one of its closures, references an object that
Move the initialization of these objects inside the decorated function.
-Kinetic bisects the payload after a failure, and names the component at
+Kinetic bisects the payload after a failure and names the component at
fault. One example message is
`kinetic could not serialize argument 2 (type socket): ...`. The message
-identifies the function, one positional argument, or one keyword
-argument, instead of a location inside `cloudpickle`. Kinetic counts
-positional arguments from `0`.
+names the function, one positional argument, or one keyword argument,
+instead of a location inside `cloudpickle`. Kinetic counts positional
+arguments from `0`.
### Kinetic rejects one of your arguments at submit time
@@ -235,27 +302,45 @@ guarantees about argument types.
Pickled code objects are not portable across Python minor versions. The
runner compares the client fingerprint in the payload against the pod.
The error that `result()` raises names both sides
-(`client Python 3.12.2 / pod Python 3.11.9`). The runner reports a
-difference in the `cloudpickle` version in the same way. The pod log holds
-a skew warning only when the payload unpickled correctly.
-
-Bundled mode always matches the Python of your client. For a prebuilt
-image or a custom image, you must match it yourself. See
+(`client Python 3.12.2 / pod Python 3.11.9`). The runner compares the
+Python minor version (`X.Y`) only, and the `cloudpickle` version exactly.
+The pod log holds a skew warning only when the payload unpickled
+correctly.
+
+The image that Kinetic builds always matches the Python of your client.
+For a base image that you publish, or for a custom image, you must match
+the version yourself. See
[Matching your local environment to the pod](guides/packaging.md#matching-your-local-environment-to-the-pod).
### JAX version mismatch errors
-You probably pinned `jax` or `jaxlib` in `requirements.txt`. Kinetic
-filters those out by default; if you need a specific version, use
-`# kn:keep` (see [Dependencies](guides/dependencies.md)), but expect
-to debug runtime/library alignment yourself.
+You pinned `jax`, `jaxlib`, or `libtpu` in your dependency file. Kinetic
+filters those lines out by default and logs a warning for each one. If
+you need a specific version, add `# kn:keep` to the line. You must then
+make sure that the JAX version matches the accelerator runtime. See
+[Dependencies](guides/dependencies.md).
+
+### The job fails but the logs look normal
+
+The pod exited with a non-zero code and wrote no result payload. The
+usual causes are an out-of-memory event (`OOMKilled`) and a signal that
+stopped the process. Read the pod events:
-### Job FAILS but logs look fine
+```bash
+kubectl get pods -l job-id=
+kubectl describe pod
+```
+
+Look for `OOMKilled` or an eviction event. If your profile uses a
+namespace other than `default`, add `-n ` to both commands.
-The pod exited non-zero without writing a result payload — usually
-caused by an OOM kill or the kernel reaping the process. Check pod
-events with `kubectl describe pod ` (find the pod name from
-`kinetic jobs status `).
+A blocking call, and `result()` on a `JobHandle`, delete the Kubernetes
+Job and its pod when they collect the failure. Kubernetes also deletes a
+finished Job 10 minutes after it ends (2 hours for a job with
+`debug=True`). To inspect the pod, run the job with `run_async()` and
+read the events before you call `result()`. As an alternative, call
+`result(cleanup=False)`, or run `kinetic jobs result
+--no-cleanup`, to keep the pod until the 10-minute limit.
On a multi-host job, `kinetic jobs logs` gives you the logs of the
leader pod only. A different host can be the host that failed. Read the
@@ -266,31 +351,40 @@ the host that failed.
Your function left a non-daemon thread alive: a data-loader worker, a
metrics uploader, or a `ThreadPoolExecutor` that your code did not shut
-down. CPython waits for non-daemon threads at interpreter exit. The pod
-therefore stays alive after the work ends. You continue to pay for the
-accelerator.
+down. CPython waits for non-daemon threads at interpreter exit. Without
+a forced exit, the pod stays alive after the work ends, and you pay for
+the accelerator.
Kinetic uploads the result first. Kinetic then logs a warning that names
-every thread still alive, and forces the process to exit. The job
-completes, but that warning reports a real leak. Call `.shutdown()` or
-`.join()` on your executors and pools. As an alternative, create your
+every thread that is still alive, and forces the process to exit. The
+job completes, but the warning reports a real leak. Call `.shutdown()`
+or `.join()` on your executors and pools. As an alternative, create your
threads with `daemon=True`.
### The job fails with "function called sys.exit(N)"
-Kinetic treats a `sys.exit()` inside your function as the process-level
-exit that it is. `sys.exit()` and `sys.exit(0)` report a **success**. The
-result is `None`. Any other exit code fails the job with that message.
-Return a value instead of an exit if you want Kinetic to collect a result.
+Kinetic catches a `sys.exit()` call inside your function. `sys.exit()`
+and `sys.exit(0)` report a **success**, and the result is `None`. Any
+other exit code fails the job with that message. Return a value instead
+of an exit if you want Kinetic to collect a result.
+
+### A job stops after about 24 hours
+
+The node pools that `kinetic up` and `kinetic pool add` create set a
+maximum run duration of 24 hours on each VM. Only a Spot TPU pool has no
+such limit. Google Cloud stops the VM at that time, the pod ends, and
+the job fails. For a job that runs longer than 24 hours, write
+checkpoints under `KINETIC_OUTPUT_DIR` and resume the work in a new job.
+See [Outputs and Checkpoints](guides/checkpointing.md).
## Missing outputs and results
### "Job failed but no result payload was found"
-The pod died before it wrote anything to
-`gs://{bucket}/{job_id}/result.pkl`. Kinetic uploads a failure payload
-for each of its own startup phases. That payload carries a `phase` field,
-which names the phase at fault:
+The pod ended before it wrote anything to
+`gs://{jobs bucket}/{job_id}/result.pkl`. Kinetic uploads a failure payload
+for each of its own startup phases. That payload carries a `phase`
+field, which names the phase at fault:
- `artifact download`
- `artifact verification`
@@ -301,25 +395,29 @@ which names the phase at fault:
- `data resolve`
- `debugger setup`
-If you see this message, the failure happened outside those phases. These
-causes remain:
+If you see this message, the failure happened outside those phases.
+These causes remain:
-- **Something killed the container** — an out-of-memory kill, a node
- preemption on Spot, or a reap by the kernel. Run
+- **Something stopped the container** — an out-of-memory event
+ (`OOMKilled`), a node preemption on Spot, or a signal that stopped the
+ process. Run
`kubectl describe pod ` and look for `OOMKilled` or an
- eviction event. Get the pod name from `kinetic jobs status `.
+ eviction event. Find the pod with
+ `kubectl get pods -l job-id=`.
- **The image cannot start the runner** — a custom image without
`/app/remote_runner.py`, `python3`, `cloudpickle`,
- `google-cloud-storage`, or `absl-py`. The pod logs stop before any
+ `google-cloud-storage`, or `absl-py`. The pod log stops before any
Kinetic output.
- **Cloud Storage was unreachable** — the runner could not upload the
- failure payload either. The pod logs hold the original error.
-- **Kinetic already deleted the artifacts** — a blocking `run()` deletes
- them after it collects a result, and `run_async()` deletes them on
- `.cleanup()`. A reattach after that point finds nothing.
+ failure payload either. The pod log holds the original error.
+- **Kinetic already deleted the artifacts** — a blocking call deletes
+ the artifacts after it collects a result. For a detached job,
+ `handle.result()` deletes the artifacts on success, and `.cleanup()`
+ deletes them at any time. A reattach after that point finds nothing.
For a long job, write your artifacts under `KINETIC_OUTPUT_DIR` instead
-of a return value. See [Checkpointing](guides/checkpointing.md).
+of a return value. See
+[Outputs and Checkpoints](guides/checkpointing.md).
On a multi-host job, only the leader (process 0) writes the result. You
therefore see this message only if no host of the job wrote a record. If
@@ -351,82 +449,99 @@ eviction event.
### Kinetic cannot serialize the return value of your job
-Your function returned an object that `cloudpickle` cannot serialize. The
-runner then uploads a payload with the flag `serialization_failed` and a
-`repr()` of the value. Kinetic reports the job as **failed**, because the
-return value is not retrievable, and a success report would hide the
-problem. The local error holds the truncated `repr`, so you can see what
-your function returned.
+Your function returned an object that `cloudpickle` cannot serialize.
+The runner then uploads a payload with the flag `serialization_failed`
+and a `repr()` of the value. Kinetic reports the job as **failed**,
+because the return value is not retrievable, and a success report would
+hide the problem. The local error holds the truncated `repr`, so you can
+see what your function returned.
Kinetic keeps the Cloud Storage artifacts in this case, so you can
inspect the run afterwards. Kinetic still deletes the Kubernetes
-resource, so collect the pod logs first if you need them. Return
-something serializable instead, such as a path or a dict of metrics.
-Write the heavy object, or the object that holds a handle, to
+resource, so collect the pod log first if you need it. Return a
+serializable value instead, such as a path or a dict of metrics. Write
+the heavy object, or the object that holds a handle, to
`KINETIC_OUTPUT_DIR`.
### `result()` cannot unpickle the result after a reattach
The result references classes from your project, or from a package that
-this client cannot import. Unpickling needs those types locally. A
-reattach from a different directory, a different virtualenv, or a
-different machine can therefore fail, although the client that submitted
-the job succeeds.
+this client cannot import. Unpickling needs those types locally. The
+client that submitted the job can import them. A reattach from a
+different directory, a different virtual environment, or a different
+machine can therefore fail.
Kinetic raises a `RuntimeError` that names the job and the artifact URI
-(`gs://{bucket}/{job_id}/result.pkl`), and chains the original error.
+(`gs://{jobs bucket}/{job_id}/result.pkl`), and chains the original error.
Kinetic does **not** delete the Cloud Storage artifacts, so you lose
nothing. Reattach from the project directory, with the same packages
installed. As an alternative, download the object and inspect it there.
-Return plain data to avoid the problem completely: a dict, a list, a
-number, a string, or an array.
+Return plain data to avoid the problem: a dict, a list, a number, a
+string, or an array.
-### Files I wrote inside the job are gone
+### Files that you wrote inside the job are gone
-Two possibilities:
+Two causes are possible:
-- You wrote them under `/tmp` or another pod-local path. The pod is
- destroyed when the job ends; pod-local files don't survive. Always
- write to `KINETIC_OUTPUT_DIR`.
-- You wrote them under `KINETIC_OUTPUT_DIR` but more than 30 days have
- passed. The default GCS bucket has a 30-day TTL. Copy critical
- artifacts to a bucket without lifecycle rules. See
- [Checkpointing](guides/checkpointing.md) for the TTL and retention
+- You wrote the files under `/tmp` or another pod-local path. Kinetic
+ deletes the pod when the job ends, and pod-local files do not survive.
+ Write the files under `KINETIC_OUTPUT_DIR`.
+- You wrote the files under `KINETIC_OUTPUT_DIR`, and more than 30 days
+ have passed. The jobs bucket has a 30-day lifecycle rule. Copy the
+ artifacts that you must keep to a bucket without a lifecycle rule. See
+ [Outputs and Checkpoints](guides/checkpointing.md) for the retention
details.
-### Logs aren't streaming back
-
-Network blip during a `--follow` stream is the most common cause. The
-pod is unaffected — log retrieval is read-only. Use
-`kinetic jobs logs ` (without `--follow`) or `--tail N` to fetch
-fresh logs from any machine.
-
-## What the troubleshoot path actually checks
-
-`kinetic init`'s troubleshoot path runs eight groups of checks and
-prints concrete fix commands when any fail. The groups (matching the
-source at `kinetic/cli/commands/doctor.py`):
-
-1. **Local Tools** — `gcloud`, `kubectl`, and
- `gke-gcloud-auth-plugin` are installed and on your PATH.
-2. **Authentication** — Application Default Credentials are present,
- refreshable, and not expired.
-3. **Configuration** — `KINETIC_PROJECT`, `KINETIC_ZONE`, and
- `KINETIC_CLUSTER` resolve to non-empty values.
-4. **GCP Project** — the project exists and has billing enabled.
-5. **GCP APIs** — Compute Engine, Cloud Build, Artifact Registry,
- Storage, and Container APIs are enabled.
-6. **GCP Resources** — the Kinetic service accounts, Artifact Registry
- repository, GCS buckets, VPC network, and Cloud NAT all exist.
-7. **Infrastructure** — Pulumi state is present and the GKE cluster is
- in the `RUNNING` state.
-8. **Kubernetes** — your `kubeconfig` points at the cluster, the API
- server responds, node pools are healthy, GPU drivers are installed
- where needed, and accelerator quotas are not exhausted.
-
-Each failing check prints a one-line fix suggestion. For multi-step
-fixes, the troubleshoot path prints a copy-paste command block.
+### The log does not stream
+
+A network interruption during a `--follow` stream is the most common
+cause. A stream interruption does not affect the pod, because log retrieval is read-only. Run
+`kinetic jobs logs ` without `--follow` for the full log. Run
+`kinetic jobs logs --tail N` for the last N lines. Both
+commands work from any machine.
+
+## What the troubleshoot path checks
+
+The troubleshoot path of `kinetic init` runs eight groups of checks and
+prints a fix hint for each check that fails. The groups match the
+source in `kinetic/cli/commands/doctor.py`:
+
+1. **Local Tools** — `gcloud`, `kubectl`, and `gke-gcloud-auth-plugin`
+ are on your `PATH`.
+2. **Authentication** — Application Default Credentials exist and
+ refresh, and `gcloud` has an active account.
+3. **Configuration** — the troubleshoot path has a project ID, a zone,
+ and a cluster name.
+4. **GCP Project** — the project exists, you can access it, and billing
+ is enabled.
+5. **GCP APIs** — the Compute Engine, Cloud Build, Artifact Registry,
+ Cloud Storage, Kubernetes Engine, Secret Manager, and IAM APIs are
+ enabled.
+6. **GCP Resources** — the resources of the cluster exist. These are
+ the node and build service accounts, the Artifact Registry
+ repository, the jobs and builds buckets, the VPC network, and the
+ Cloud NAT gateway.
+7. **Infrastructure** — the infrastructure state is present in the state
+ bucket, and the GKE cluster is in the `RUNNING` state.
+8. **Kubernetes** — the cluster is reachable and healthy. This group
+ checks these points:
+ - Your `kubeconfig` points at the cluster.
+ - The API server responds.
+ - The node pools are healthy.
+ - The LeaderWorkerSet CRD is installed.
+ - The Kinetic Kubernetes service account exists.
+ - The NVIDIA drivers are installed when a GPU pool exists.
+ - The nodes are healthy.
+ - No pod is stuck in `Pending`.
+ - The cluster has no warning events such as `FailedScheduling` or
+ `OOMKilling`.
+ - No accelerator quota in the region is exhausted.
+
+The troubleshoot path diagnoses only a cluster from one of your
+profiles. If you have no profile for the project, the troubleshoot path
+skips groups 6 to 8. Each failing check prints a fix hint. A hint can
+hold more than one command.
## Related pages
@@ -437,28 +552,28 @@ fixes, the troubleshoot path prints a copy-paste command block.
:link: getting_started
:link-type: doc
-First-run setup that shouldn't have to fail twice.
+Install Kinetic, run `kinetic init`, and run your first job.
:::
:::{grid-item-card} {octicon}`file-directory;1em` What Ships to the Pod
:link: guides/packaging
:link-type: doc
-The packaging contract behind most import failures and most failures
-of a file path.
+The packaging contract behind most import failures and most file path
+failures.
:::
:::{grid-item-card} {octicon}`book;1em` FAQ
:link: guides/faq
:link-type: doc
-Quick answers to common conceptual confusions.
+Short answers to common questions about how Kinetic works.
:::
:::{grid-item-card} {octicon}`gear;1em` Configuration
:link: configuration
:link-type: doc
-Env vars and precedence.
+Profiles, environment variables, and the precedence rules.
:::
::::
diff --git a/examples/example_gke.py b/examples/example_gke.py
index a0da5622..95b29c85 100644
--- a/examples/example_gke.py
+++ b/examples/example_gke.py
@@ -1,37 +1,21 @@
-"""
-Example: Using kinetic with GKE
-
-This demonstrates running remote functions on a GKE cluster with kinetic.
+"""Run remote functions on a CPU pool, a TPU pool, and a GPU pool in turn.
Prerequisites:
-1. A GKE cluster (CPU or with GPU node pools)
-2. kubectl configured to access the cluster
-3. KINETIC_PROJECT environment variable set
-
-Setup (CPU cluster - works out of the box):
- ./setup.sh # Answer 'yes' when prompted for GKE setup
-
-Setup (GPU cluster - for GPU examples):
- # Add a GPU node pool to existing cluster
- gcloud container node-pools create gpu-pool \\
- --cluster kinetic-cluster \\
- --zone us-central1-a \\
- --machine-type n1-standard-4 \\
- --accelerator type=nvidia-tesla-t4,count=1 \\
- --num-nodes 1 \\
- --scopes gke-default,storage-full
-
- # Install NVIDIA GPU drivers
- kubectl apply -f https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/master/nvidia-driver-installer/cos/daemonset-preloaded.yaml
-
-Supported accelerators:
- - cpu: CPU only (no GPU required)
- - gpu-t4: NVIDIA T4
- - gpu-l4: NVIDIA L4
- - gpu-v100: NVIDIA V100
- - gpu-a100: NVIDIA A100 (40GB)
- - gpu-a100-80gb: NVIDIA A100 (80GB)
- - gpu-h100: NVIDIA H100
+1. A Kinetic cluster and an active profile. Run `kinetic init` one time.
+ The profile supplies the project, the zone, and the cluster.
+2. A node pool for each accelerator that this script uses. Add a pool
+ with `kinetic pool add`, for example:
+
+ kinetic pool add --accelerator tpu-v6e-2x4 --spot
+ kinetic pool add --accelerator gpu-t4
+
+ `kinetic pool list` shows the pools of the cluster. Change the
+ `accelerator=` values below to match your pools, or remove the calls
+ for the pools that you do not have. `kinetic accelerators` lists every
+ accelerator name.
+
+Note: `tpu-v6e-2x4` is a two-host slice, so Kinetic runs it on the
+Pathways backend. The `spot=True` job needs a Spot node pool.
"""
import os
diff --git a/examples/gemma4_finetuning/gemma4_finetuning.py b/examples/gemma4_finetuning/gemma4_finetuning.py
index 0fa3d86f..aca278d1 100644
--- a/examples/gemma4_finetuning/gemma4_finetuning.py
+++ b/examples/gemma4_finetuning/gemma4_finetuning.py
@@ -5,12 +5,13 @@
memory (~52 GB in bfloat16), so a v5litepod-8 (8 chips × 16 GB = 128 GB HBM)
is the minimum supported configuration.
-Set the placeholder values in __main__ before running.
-See docs/guides/gemma4_finetuning.md for a full walkthrough.
+Run `kinetic init` one time so that an active profile supplies the project,
+the zone, and the cluster. Make sure that the cluster has a v5litepod-8
+node pool (`kinetic pool add --accelerator tpu-v5litepod-8`). See
+docs/examples/gemma4_finetuning.md for the full walkthrough.
-Dependencies are declared in requirements.txt in this directory. Run this
-script from examples/gemma4_finetuning/ so Kinetic picks it up from the
-current working directory.
+Kinetic reads the requirements.txt in this directory: it starts the search
+at the directory of this file, not at your current working directory.
"""
import os
@@ -326,9 +327,9 @@ def run_inference(weights_path: str):
if __name__ == "__main__":
os.environ["KERAS_BACKEND"] = "jax"
- os.environ["GOOGLE_CLOUD_PROJECT"] = "your-project-id"
- os.environ["KINETIC_ZONE"] = "us-central1-a"
- os.environ["GOOGLE_CLOUD_ZONE"] = "us-central1-a"
+ # The active profile (from `kinetic init`) supplies the project, the zone,
+ # and the cluster. Set KINETIC_PROJECT / KINETIC_ZONE here only for a
+ # one-off override.
weights_path = fine_tune_gemma4()
print(f"Training complete. Weights at: {weights_path}")
diff --git a/examples/hf_dataset_demo.py b/examples/hf_dataset_demo.py
index 3d5a6339..304adf11 100644
--- a/examples/hf_dataset_demo.py
+++ b/examples/hf_dataset_demo.py
@@ -3,8 +3,10 @@
This script demonstrates how to use Hugging Face datasets in a Kinetic job.
It downloads the dataset on the remote pod and makes it available to your function.
-Note: To run this demo with a gated or private dataset, you must set the
-`HF_TOKEN` environment variable in your environment before running this script.
+Note: For a gated or private dataset, set `HF_TOKEN` in your shell and add
+`capture_env_vars=["HF_TOKEN"]` to the decorator. Kinetic then copies the
+token into the pod, where the `datasets` library reads it. Kinetic itself
+does not read `HF_TOKEN`.
"""
import time