Skip to content
Merged
Show file tree
Hide file tree
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 7 additions & 4 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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
Expand All @@ -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
Expand Down
135 changes: 101 additions & 34 deletions docs/accelerators.md
Original file line number Diff line number Diff line change
@@ -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.
:::
::::
131 changes: 113 additions & 18 deletions docs/architecture.md
Original file line number Diff line number Diff line change
@@ -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.
Loading
Loading