Introduction

In this book, we will learn about running GPU workloads on Kubernetes and dive deep into Kubernetes machinery with hands-on examples, both in Rust and Go. We will first build intuition by reasoning and quickly prototype the data structures behind Kubernetes machinery and concepts. Then we will build the same machinery with the libraries that production controllers use, controller-runtime in Go and kube-rs in Rust, and run it on a real cluster.

GPU workloads range from distributed training jobs that run across many machines to inference services that answer requests all day. Running them by hand means choosing a machine for each job, tracking which GPUs are free, and restarting work when something fails, and doing that manually is error-prone and cumbersome. Kubernetes automates it: we declare what we want to run, and its scheduler and controllers place the work and keep it running. It is also extensible, so we can teach it about GPUs and about new kinds of work, such as training jobs.

When a model is too large to train on one machine, we split the training across many machines, and every training step needs every GPU in the job. This makes scheduling harder. All of the job’s workers have to start together, because a job that gets only some of its workers holds those GPUs and makes no progress. The workers should also run close to each other, because GPUs in the same machine exchange data much faster than GPUs in different machines. And when several teams share one cluster, we need queues and quotas so that every team gets its share of the GPUs. In this book, we will learn how Kubernetes solves each of these problems, starting with how a controller learns about a new job.

The examples run on any Kubernetes cluster with GPU nodes: a managed cluster from a cloud provider, a cluster in our own data center, or a small cluster on a single machine.

By default, Kubernetes knows nothing about GPUs. It knows how many CPU cores and how much memory each machine has, and it places containers on machines that have enough of both. Two pieces of software, supplied by the GPU vendor, teach it about GPUs.

The first is the device plugin, a small program that runs on every machine with GPUs. It counts the GPUs, reports the count to the kubelet (the Kubernetes agent on each machine), and from then on the machine advertises the GPUs as a resource next to its CPU and memory; for NVIDIA GPUs the resource is called nvidia.com/gpu. A pod, the object in which Kubernetes runs one or more containers, can then ask for nvidia.com/gpu: 1 in the same way that it asks for two CPU cores, and the scheduler, the Kubernetes component that chooses a machine for each new pod, places it only on a machine with a free GPU. By default, one unit is one whole physical GPU; chapter 2 looks at what that unit contains and at the ways a GPU can be shared or divided.

The second is a GPU-aware container runtime. A container normally sees none of the machine’s hardware devices, and its image does not contain the machine’s GPU driver. When a pod has been given a GPU, the runtime adds the GPU’s device files and the driver’s libraries to the container just before it starts, so that a program inside it can use the GPU.

Figure I.1 follows one request for a GPU past both pieces.

We develop and test every example on a single-node cluster created with kind, a tool that runs Kubernetes nodes as Docker containers. Appendix A shows how we build that cluster and install both pieces on it.

We assume you can read Go and Rust and have used containers and kubectl. We explain each Kubernetes concept when we first need it.

Running the code

The book and all of its code live in the gpu-accelerated-kubernetes directory of the cracking-the-systems-programming-interview repository on GitHub:

$ git clone https://github.com/arpanpathak/cracking-the-systems-programming-interview.git
$ cd cracking-the-systems-programming-interview/gpu-accelerated-kubernetes

The listings in each chapter are included from the files under code/, so the code printed in the book is the code that runs, and the source: link above each listing opens that file on GitHub.

• code/go is one Go module. informer is the in-memory model from chapter 1, gpujob is the controller from section 1.1, built on controller-runtime, and gpuusage is the GPU accounting tool from chapter 2:

  $ cd code/go
  $ go run ./informer
  $ go run ./gpujob/cmd/controller
  

• code/rust is one Cargo workspace with the same three programs as packages. gpujob and gpuusage are built on kube-rs:

  $ cd code/rust
  $ cargo run -p informer
  $ cargo run -p gpujob
  

• code/manifests holds the example objects that both controllers act on.

• code/gpu-sharing holds the CUDA programs and manifests for chapter 2.

• code/jetson-gpu holds the script, manifests, and test program for the kind cluster in Appendix A.

Every program has the same name in both trees. From chapter 3 on, each is one directory under code/go, run with go run ./<name>, and one package under code/rust, run with cargo run -p <name>:

The programs of chapters 4 to 10 are models: they need no cluster, finish in well under a second, and print the same output in Go and in Rust. The output shown in each chapter is included from code/output, where a run of the program wrote it.

Each chapter introduces its programs in the same order. It starts from the problem and a small example worked by hand, draws the design, and states the types and the rules they keep in a box headed “Low-level design”. The code then appears in short parts, each next to the text that explains it, and the complete program is printed at the end of the chapter.

The Go module declares Go 1.26, and the Rust programs need Rust 1.89 or later, the minimum version of the kube crate. The controllers need a Kubernetes cluster in the current kubeconfig context.

The diagrams are drawn by tools/draw_figures.py and the animations by tools/animations.py; make figures and make animations run them, and their output in src/figures is part of the repository. The book’s build removes the ANCHOR comments that mark the parts of each program, with tools/strip_anchors.py, so Python 3 has to be installed.

To read the book in a browser, install mdBook and run:

$ mdbook serve --open

1. Keeping a local copy of the cluster

Suppose we want to write a program that runs GPU training jobs on Kubernetes. A user creates a job and says “this job needs one GPU”. Our program has to notice the new job, start its container on a machine with a free GPU, and write down that the job is running. When the user changes the job to ask for two GPUs, our program has to notice that too and replace the container. When the user deletes the job, our program has to stop the container and free the GPU.

In Kubernetes, a program like this is called a controller. In this chapter, we will learn how a controller finds out what is happening in the cluster, which turns out to be the most interesting part of the whole design. We will first build intuition by reasoning about the problems a controller runs into, and quickly prototype the data structures that solve them: a small API server, a cache, a loop that keeps the cache up to date, a work queue, and the controller itself. Then we will run our prototype through a burst of edits, a dropped connection, and a network partition, and finally we will map each piece onto client-go, the library that real Go controllers use.

We will build the model twice, in Go and in Rust, with the same six parts and the same behaviour. The Go files are in code/go/informer, and we run them with go run ./informer from code/go. The Rust files are in code/rust/informer, and we run them with cargo run -p informer from code/rust. Neither version needs anything beyond its language’s standard library.

Objects, specs, and statuses

Kubernetes describes a cluster as a collection of objects. An object is a record with a name and two main parts. The spec holds what a user asked for: “run this image with one GPU”. The status holds what has actually happened so far: “running on one GPU”. A training job in our cluster will be one such object.

All objects live behind a single HTTP service, the API server. Every read and every write of every object goes through it, and it stores the objects in etcd, a replicated key-value database. Users, command-line tools, and controllers are all just clients of the API server.

A controller’s job is to compare each object’s spec with the real world and act on any difference, then record what it did in the status. To do that well, it needs two things from the API server: an accurate view of every object of its kind, and news of every change within a fraction of a second. Let us see how to get both.

Numbering every write

The most obvious approach is to ask the API server for every object every few seconds. It works for a handful of objects, and it falls apart as the cluster grows. A cluster with 50,000 pods stores a few kilobytes of JSON per pod, so a single request for all of them transfers around 200 megabytes. If a dozen controllers each did that every five seconds, the API server would be sending almost half a gigabyte per second, and between two of those requests only a handful of pods would have changed.

What we really want is for the API server to tell us about each change as it happens. That idea brings a problem with it. Connections break: networks fail, servers restart, and proxies close connections they consider idle. When our connection breaks and we open a new one, we need to say “send me everything I missed”, and that requires a way to name the point where we left off.

The API server solves this by numbering its writes. It keeps a single counter for the whole store. Every time any object is created, updated, or deleted, it increments the counter and stamps the new value onto the object that changed. Kubernetes calls this number the object’s resource version. The numbers come straight from etcd, which gives every write a revision number from one counter for the whole database.

The same number helps with a second problem. Imagine two clients read a job at version 7. The first changes the GPU count and writes the job back. The second changes the image and writes the job back a moment later, still based on what it read at version 7. Without any check, the second write would replace the whole object and silently undo the first client’s change. With resource versions, the second client sends “I read this at version 7” along with its write. The object is now at version 8, so the server refuses the write with the HTTP status 409 Conflict, and the second client knows it has to read the job again and redo its change on top of the new version.

Listing, then watching

With numbered writes, a client can build an exact view of the cluster in two steps.

First it asks for every object. This request is called a list, and the response carries the objects together with the resource version at which the server took the snapshot. Then the client opens a watch: a request that stays open and asks the server to send every change made after a given version, one event per change. The client passes the version it got from the list.

Passing that version is the important detail. The list and the watch are two separate requests, and other clients keep writing in between. Suppose our list is taken at version 10, then a user edits a job at version 11, and only then does our watch open, when the store is at version 12:

version:   10          11             12
           list        job-a edited   watch opened
           snapshot                   from version 10
                                      -> the server replays event 11,
                                         then sends new events as they occur

Because the watch asks for everything after version 10, the server first sends us event 11 from its record of recent changes and then carries on with new events. Our view is the snapshot at version 10 plus every change after it: nothing missing, nothing counted twice. If we had opened the watch at the current version, 12, we would have skipped the edit at version 11 and kept an out-of-date copy of job-a until the job happened to change again.

When the history runs out

The server can only replay changes it still has. etcd keeps every revision of every key until it is told to throw old ones away, and that clean-up is called compaction. Without it the database would grow forever. The Kubernetes API server asks etcd to compact every five minutes by default. After a compaction, a watch that asks to resume from an older version cannot be served, because the events it needs have been thrown away. The API server answers with the HTTP status 410 Gone, and the only way for us to get back to an exact view is to list again and watch from the new list’s version.

Watches also end for ordinary reasons. The API server closes every watch after a timeout, chosen at random between 30 and 60 minutes by default, so that thousands of clients do not all reconnect at the same moment. A load balancer between us and the server may close a connection that has been quiet for too long. And the network may fail for a while, so that even new connections fail.

Each of these calls for a different response, and our client has to tell them apart. If the watch simply closed, we reopen it from the last version we received. If we cannot connect at all, we wait and try again. If the server says 410 Gone, we list again.

Building the server

Now we have enough understanding to build a model. We start with the server. It keeps the objects in memory, numbers every write, records recent events, and answers lists and watches. We also give it three extra methods, DropWatches, Partition, and Heal, so that our demonstration at the end of the chapter can cause each of the failures we just discussed.

The server’s data comes first: the object, the event that reports a change to it, and the two errors a client has to tell apart.

Listing 1.1 Objects, events, and the two errors, in Go. code/go/informer/server.go

// Spec is the part of an object that a user writes: the desired state.
type Spec struct {
	// GPUs is the number of GPUs the job asks for.
	GPUs int
}

// Status is the part of an object that the controller writes: the observed state.
type Status struct {
	// GPUs is the number of GPUs the controller has provided to the job.
	GPUs int
}

// Object is one stored record, such as a training job.
type Object struct {
	// Name identifies the object on the server.
	Name string
	// Spec holds what the user asked for.
	Spec Spec
	// Status holds what the controller has done about it.
	Status Status
	// ResourceVersion is the server's write counter at the object's latest change.
	ResourceVersion int
}

// String formats an object for the demonstration's log lines.
func (o Object) String() string {
	return fmt.Sprintf("%s{spec.GPUs: %d, status.GPUs: %d, v%d}",
		o.Name, o.Spec.GPUs, o.Status.GPUs, o.ResourceVersion)
}

// EventType is the kind of change that an Event records.
type EventType int

// The kinds of change that a watch reports.
const (
	Added    EventType = iota // the object was created
	Modified                  // the object was updated
	Deleted                   // the object was removed
)

// String returns the name that the Kubernetes watch protocol uses for t.
func (t EventType) String() string {
	return [...]string{"ADDED", "MODIFIED", "DELETED"}[t]
}

// Event records one change to one object.
type Event struct {
	// Type is the kind of change.
	Type EventType
	// Object is the object after the change, or its final state if it was deleted.
	Object Object
}

// Errors that the server returns. The first two carry the HTTP status codes
// that the Kubernetes API server uses for the same conditions.
var (
	// ErrGone reports that the history a watch needs has been discarded.
	ErrGone = errors.New("410 Gone: resource version too old")
	// ErrConflict reports a conditional write whose resource version is out of date.
	ErrConflict = errors.New("409 Conflict: object has been modified")
	// ErrUnreachable reports that the client cannot reach the server.
	ErrUnreachable = errors.New("server unreachable")
)

Listing 1.2 Objects, events, and the two errors, in Rust. code/rust/informer/src/server.rs

/// The part of an object that a user writes: the desired state.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct Spec {
    /// Number of GPUs the job asks for.
    pub gpus: u32,
}

/// The part of an object that the controller writes: the observed state.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct Status {
    /// Number of GPUs the controller has provided to the job.
    pub gpus: u32,
}

/// One stored record, such as a training job.
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct Object {
    /// Identifies the object on the server.
    pub name: String,
    /// What the user asked for.
    pub spec: Spec,
    /// What the controller has done about it.
    pub status: Status,
    /// The server's write counter at the object's latest change.
    pub resource_version: u64,
}

impl Object {
    /// Returns a new object with the given name and GPU request, an empty
    /// status, and no resource version.
    pub fn new(name: &str, gpus: u32) -> Self {
        Self {
            name: name.to_string(),
            spec: Spec { gpus },
            ..Self::default()
        }
    }
}

impl fmt::Display for Object {
    /// Formats an object for the demonstration's log lines.
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "{}{{spec.gpus: {}, status.gpus: {}, v{}}}",
            self.name, self.spec.gpus, self.status.gpus, self.resource_version
        )
    }
}

/// The kind of change that an [`Event`] records.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum EventType {
    /// The object was created.
    Added,
    /// The object was updated.
    Modified,
    /// The object was removed.
    Deleted,
}

/// One change to one object.
#[derive(Clone, Debug)]
pub struct Event {
    /// The kind of change.
    pub kind: EventType,
    /// The object after the change, or its final state if it was deleted.
    pub object: Object,
}

/// Errors that the server returns. The first two carry the HTTP status codes
/// that the Kubernetes API server uses for the same conditions.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ApiError {
    /// The history that a watch needs has been discarded.
    Gone,
    /// A conditional write carried an out-of-date resource version.
    Conflict,
    /// The client cannot reach the server.
    Unreachable,
}

impl fmt::Display for ApiError {
    /// Formats the error the way the Kubernetes API server reports it.
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(match self {
            Self::Gone => "410 Gone: resource version too old",
            Self::Conflict => "409 Conflict: object has been modified",
            Self::Unreachable => "server unreachable",
        })
    }
}

The model reduces a job to a GPU count in its spec and in its status.

The server itself is a lock around four things: the objects, the counter, the events it still has, and the watches that are open.

Listing 1.3 The server's state, in Go. code/go/informer/server.go

// APIServer stores objects and streams their changes to watchers. Its methods
// are safe for concurrent use.
type APIServer struct {
	// mu guards every field below.
	mu sync.Mutex
	// version is the last resource version the server handed out.
	version int
	// objects holds the current state of every object, keyed by name.
	objects map[string]Object
	// history holds the events newer than compacted, oldest first.
	history []Event
	// compacted is the oldest version from which a watch may resume.
	compacted int
	// watches holds one channel for each open watch.
	watches []chan Event
	// unreachable makes new watches fail, to simulate a network partition.
	unreachable bool
}

// NewAPIServer returns an empty server.
func NewAPIServer() *APIServer {
	return &APIServer{objects: map[string]Object{}}
}

Listing 1.4 The server's state, in Rust. code/rust/informer/src/server.rs

/// Everything the server stores, kept behind one lock.
#[derive(Default)]
struct State {
    /// The last resource version the server handed out.
    version: u64,
    /// The current state of every object, keyed by name.
    objects: HashMap<String, Object>,
    /// The events newer than `compacted`, oldest first.
    history: Vec<Event>,
    /// The oldest version from which a watch may resume.
    compacted: u64,
    /// One sender for each open watch.
    watches: Vec<Sender<Event>>,
    /// Makes new watches fail, to simulate a network partition.
    unreachable: bool,
}

impl State {
    /// Appends `event` to the history and sends it to every open watch.
    /// Watches whose receiver has gone away are forgotten.
    fn publish(&mut self, event: Event) {
        self.watches.retain(|w| w.send(event.clone()).is_ok());
        self.history.push(event);
    }
}

/// Stores objects and streams their changes to watchers. Its methods are safe
/// to call from several threads.
#[derive(Default)]
pub struct ApiServer {
    /// Guards all of the server's data.
    state: Mutex<State>,
}

impl ApiServer {
    /// Locks the server's data.
    fn state(&self) -> std::sync::MutexGuard<'_, State> {
        self.state.lock().expect("server lock poisoned")
    }
}

Two methods change objects, and both end by publishing an event.

Listing 1.5 Writing: update, delete, and publish, in Go. code/go/informer/server.go

// Update creates or replaces an object and returns ErrConflict if the write's
// condition fails. A non-zero o.ResourceVersion is a condition: the caller read
// the object at that version and wants the write applied only if the object is
// still at it. A zero ResourceVersion writes unconditionally.
func (s *APIServer) Update(o Object) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	current, exists := s.objects[o.Name]

	if o.ResourceVersion != 0 && o.ResourceVersion != current.ResourceVersion {
		return ErrConflict
	}

	eventType := Modified

	if !exists {
		eventType = Added
	}

	s.version++
	o.ResourceVersion = s.version
	s.objects[o.Name] = o
	s.publish(Event{eventType, o})

	return nil
}

// Delete removes the named object, if it exists.
func (s *APIServer) Delete(name string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	o, exists := s.objects[name]

	if !exists {
		return
	}

	delete(s.objects, name)
	s.version++
	o.ResourceVersion = s.version
	s.publish(Event{Deleted, o})
}

// publish appends e to the history and sends it to every open watch.
// The caller must hold s.mu.
func (s *APIServer) publish(e Event) {
	s.history = append(s.history, e)

	for _, w := range s.watches {
		w <- e
	}
}

Listing 1.6 Writing: update, delete, and publish, in Rust. code/rust/informer/src/server.rs

impl ApiServer {
    /// Creates or replaces an object. A non-zero `resource_version` is a
    /// condition: the caller read the object at that version and wants the
    /// write applied only if the object is still at it. Zero writes
    /// unconditionally.
    pub fn update(&self, mut object: Object) -> Result<(), ApiError> {
        let mut state = self.state();
        let current = state.objects.get(&object.name);
        let current_version = current.map_or(0, |o| o.resource_version);

        if object.resource_version != 0 && object.resource_version != current_version {
            return Err(ApiError::Conflict);
        }

        let kind = if current.is_some() {
            EventType::Modified
        } else {
            EventType::Added
        };
        state.version += 1;
        object.resource_version = state.version;
        state.objects.insert(object.name.clone(), object.clone());
        state.publish(Event { kind, object });

        Ok(())
    }

    /// Removes the named object, if it exists.
    pub fn delete(&self, name: &str) {
        let mut state = self.state();
        let Some(mut object) = state.objects.remove(name) else {
            return;
        };
        state.version += 1;
        object.resource_version = state.version;
        state.publish(Event {
            kind: EventType::Deleted,
            object,
        });
    }
}

Update carries both uses of the resource version. A zero version skips the check, which the demonstration uses to play a user editing a job. Any other version must equal the stored one, and a missing object counts as version 0. Each successful write increments the counter and publishes one event. Delete publishes the object’s final state under a new version.

Two methods read: one for the snapshot and one for the stream of changes after it.

Listing 1.7 Reading: list and watch, in Go. code/go/informer/server.go

// List returns every object, sorted by name, and the version at which the
// snapshot was taken.
func (s *APIServer) List() ([]Object, int) {
	s.mu.Lock()
	defer s.mu.Unlock()
	byName := func(a, b Object) int { return strings.Compare(a.Name, b.Name) }

	return slices.SortedFunc(maps.Values(s.objects), byName), s.version
}

// Watch returns a channel that first delivers the stored events newer than
// version and then every new event, until the server closes it.
func (s *APIServer) Watch(version int) (<-chan Event, error) {
	s.mu.Lock()
	defer s.mu.Unlock()

	switch {
	case s.unreachable:
		return nil, ErrUnreachable
	case version < s.compacted:
		return nil, ErrGone
	}

	w := make(chan Event, 1024)

	for _, e := range s.history {
		if e.Object.ResourceVersion > version {
			w <- e
		}
	}

	s.watches = append(s.watches, w)

	return w, nil
}

Listing 1.8 Reading: list and watch, in Rust. code/rust/informer/src/server.rs

impl ApiServer {
    /// Returns every object, sorted by name, and the version at which the
    /// snapshot was taken.
    pub fn list(&self) -> (Vec<Object>, u64) {
        let state = self.state();
        let mut objects: Vec<Object> = state.objects.values().cloned().collect();
        objects.sort_by(|a, b| a.name.cmp(&b.name));

        (objects, state.version)
    }

    /// Returns a receiver that first delivers the stored events newer than
    /// `version` and then every new event, until the server drops the watch.
    pub fn watch(&self, version: u64) -> Result<Receiver<Event>, ApiError> {
        let mut state = self.state();

        if state.unreachable {
            return Err(ApiError::Unreachable);
        }

        if version < state.compacted {
            return Err(ApiError::Gone);
        }

        let (sender, receiver) = channel();

        for event in &state.history {
            if event.object.resource_version > version {
                sender.send(event.clone()).expect("receiver is held here");
            }
        }

        state.watches.push(sender);

        Ok(receiver)
    }
}

publish runs while the server lock is held, and Watch holds the same lock while it replays the history and registers the new watch. Every event is therefore either in the replayed history or delivered after registration: no watch misses an event or receives one twice.

A Go watch buffers 1,024 events. A client that stops reading would eventually block publish while it holds the lock and stall every writer; the real API server closes such a watch. The Rust channel is unbounded. In Go the server ends a watch by closing its channel; in Rust it drops the sender, which ends the receiver’s loop.

The last four methods cause the failures of the previous section on request.

Listing 1.9 Compaction, dropped watches, and partitions, in Go. code/go/informer/server.go

// Compact discards the event history, as etcd discards old versions.
func (s *APIServer) Compact() {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.compacted = s.version
	s.history = nil
}

// DropWatches closes every open watch, as a load balancer's idle timeout would.
func (s *APIServer) DropWatches() {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.dropWatches()
}

// Partition closes every open watch and makes new watches fail until Heal.
func (s *APIServer) Partition() {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.unreachable = true
	s.dropWatches()
}

// Heal ends a partition, so new watches succeed again.
func (s *APIServer) Heal() {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.unreachable = false
}

// dropWatches closes and forgets every watch channel. The caller must hold s.mu.
func (s *APIServer) dropWatches() {
	for _, w := range s.watches {
		close(w)
	}

	s.watches = nil
}

Listing 1.10 Compaction, dropped watches, and partitions, in Rust. code/rust/informer/src/server.rs

impl ApiServer {
    /// Discards the event history, as etcd discards old versions.
    pub fn compact(&self) {
        let mut state = self.state();
        state.compacted = state.version;
        state.history.clear();
    }

    /// Closes every open watch, as a load balancer's idle timeout would.
    pub fn drop_watches(&self) {
        self.state().watches.clear();
    }

    /// Closes every open watch and makes new watches fail until [`Self::heal`].
    pub fn partition(&self) {
        let mut state = self.state();
        state.unreachable = true;
        state.watches.clear();
    }

    /// Ends a partition, so new watches succeed again.
    pub fn heal(&self) {
        self.state().unreachable = false;
    }
}

Watch fails with the gone error when the requested version is older than compacted. Compact, DropWatches, Partition, and Heal exist only to reproduce the failures described above.

A local copy of the objects

Once we have listed and are watching, we can keep our own copy of every object in memory. The copy is exact up to the last event we received. Reading an object becomes a map lookup instead of a network request, and the server only has to send each change to us once. We call this copy the cache.

The cache is a map behind a lock, with a method for one event.

Listing 1.11 The cache, a lookup, and applying one event, in Go. code/go/informer/cache.go

// Cache is the client's copy of the server's objects, kept current by a
// Reflector. Its methods are safe for concurrent use.
type Cache struct {
	// mu guards objects.
	mu sync.RWMutex
	// objects holds the latest known state of every object, keyed by name.
	objects map[string]Object
}

// NewCache returns an empty cache.
func NewCache() *Cache {
	return &Cache{objects: map[string]Object{}}
}

// Get returns the cached copy of the named object and whether it exists.
func (c *Cache) Get(name string) (Object, bool) {
	c.mu.RLock()
	defer c.mu.RUnlock()
	o, ok := c.objects[name]

	return o, ok
}

// Apply records one event from a watch.
func (c *Cache) Apply(e Event) {
	c.mu.Lock()
	defer c.mu.Unlock()

	if e.Type == Deleted {
		delete(c.objects, e.Object.Name)

		return
	}

	c.objects[e.Object.Name] = e.Object
}

Listing 1.12 The cache, a lookup, and applying one event, in Rust. code/rust/informer/src/cache.rs

/// The client's copy of the server's objects, kept current by a
/// [`Reflector`](crate::reflector::Reflector). Its methods are safe to call
/// from several threads.
#[derive(Default)]
pub struct Cache {
    /// The latest known state of every object, keyed by name.
    objects: RwLock<HashMap<String, Object>>,
}

impl Cache {
    /// Returns the cached copy of the named object, if it exists.
    pub fn get(&self, name: &str) -> Option<Object> {
        let objects = self.objects.read().expect("cache lock poisoned");

        objects.get(name).cloned()
    }

    /// Records one event from a watch.
    pub fn apply(&self, event: Event) {
        let mut objects = self.objects.write().expect("cache lock poisoned");

        match event.kind {
            EventType::Deleted => objects.remove(&event.object.name),
            EventType::Added | EventType::Modified => {
                objects.insert(event.object.name.clone(), event.object)
            }
        };
    }
}

Both caches use a read-write lock: every reconcile reads the cache, and only incoming events write to it.

Apply records a single event. It stores the object for an addition or a modification and removes it for a deletion.

A second method installs a whole list.

Listing 1.13 Replacing the cache with a fresh list, in Go. code/go/informer/cache.go

// Replace installs the result of a fresh list and returns the names of the
// objects that changed, including objects that disappeared while the client
// was not watching.
func (c *Cache) Replace(list []Object) []string {
	c.mu.Lock()
	defer c.mu.Unlock()
	old := c.objects
	c.objects = make(map[string]Object, len(list))
	var changed []string

	for _, o := range list {
		c.objects[o.Name] = o
		// A name missing from old reads as version 0, so it counts as changed.
		if old[o.Name].ResourceVersion != o.ResourceVersion {
			changed = append(changed, o.Name)
		}
	}

	for name := range old {
		if _, ok := c.objects[name]; !ok {
			changed = append(changed, name)
		}
	}

	return changed
}

Listing 1.14 Replacing the cache with a fresh list, in Rust. code/rust/informer/src/cache.rs

impl Cache {
    /// Installs the result of a fresh list and returns the names of the
    /// objects that changed, including objects that disappeared while the
    /// client was not watching.
    pub fn replace(&self, list: &[Object]) -> Vec<String> {
        let mut objects = self.objects.write().expect("cache lock poisoned");
        let old = std::mem::take(&mut *objects);
        let mut changed = Vec::new();

        for object in list {
            let old_version = old.get(&object.name).map(|o| o.resource_version);

            if old_version != Some(object.resource_version) {
                changed.push(object.name.clone());
            }

            objects.insert(object.name.clone(), object.clone());
        }

        for name in old.keys() {
            if !objects.contains_key(name) {
                changed.push(name.clone());
            }
        }

        changed
    }
}

Replace installs the result of a fresh list, and it has one more job that events alone cannot do. Think about what happens when we relist after 410 Gone. Some objects may have been deleted while we were disconnected, and their deletion events were thrown away by the compaction. Nobody will ever send us those events. The only evidence that the objects are gone is that they are in our old cache and missing from the new list. So Replace builds the new contents from the list, reports every object whose version differs from the one we had, and then reports every name that was in the old cache but is missing from the list. Objects that are new to us have no cached version at all; Go reads the missing entry as version 0 and Rust compares against None, so new objects get reported as changed too.

The reflector

Next we need a loop that keeps the cache in step with the server. It lists, installs the result with Replace, and queues the names of the objects that Replace reported, so that the controller will look at them. Then it watches from the list’s version and applies each event as it arrives. When the history it needs is gone, it lists again. We call this loop the reflector.

The outer loop lists, installs the list, and hands over to the watch.

Listing 1.15 The reflector and its list step, in Go. code/go/informer/reflector.go

// Reflector copies the server's objects into a cache and queues the name of
// every object that changes.
type Reflector struct {
	// server is the source of objects and events.
	server *APIServer
	// cache receives every object the reflector sees.
	cache *Cache
	// queue receives the name of every object that changes.
	queue *WorkQueue
}

// Run lists every object, watches from the list's version, and lists again
// whenever the server can no longer resume the watch. It does not return.
func (r *Reflector) Run() {
	for {
		list, version := r.server.List()
		r.queue.Add(r.cache.Replace(list)...)
		logf("reflector: list at version %d: %v", version, list)
		r.watchFrom(version)
	}
}

Listing 1.16 The reflector and its list step, in Rust. code/rust/informer/src/reflector.rs

/// Copies the server's objects into a cache and queues the name of every
/// object that changes.
pub struct Reflector {
    /// The source of objects and events.
    pub server: Arc<ApiServer>,
    /// Receives every object the reflector sees.
    pub cache: Arc<Cache>,
    /// Receives the name of every object that changes.
    pub queue: Arc<WorkQueue>,
}

impl Reflector {
    /// Lists every object, watches from the list's version, and lists again
    /// whenever the server can no longer resume the watch. It does not return.
    pub fn run(&self) {
        loop {
            let (list, version) = self.server.list();
            self.queue.add(self.cache.replace(&list));
            let objects: Vec<String> = list.iter().map(ToString::to_string).collect();
            log!(
                "reflector: list at version {version}: [{}]",
                objects.join(" ")
            );
            self.watch_from(version);
        }
    }
}

The inner loop watches, and decides what each kind of ending calls for.

Listing 1.17 Watching from a version, in Go. code/go/informer/reflector.go

// watchFrom applies events from version onward, reopens the watch after each
// disconnect, and returns when the server reports that the history is gone.
func (r *Reflector) watchFrom(version int) {
	for {
		w, err := r.server.Watch(version)

		if errors.Is(err, ErrGone) {
			logf("reflector: watch from version %d: %v", version, err)

			return
		}

		if err != nil {
			time.Sleep(10 * time.Millisecond) // client-go backs off exponentially
			continue
		}

		for e := range w {
			// Update the cache before queueing the name, so that a worker
			// that takes the name reads the object as of this event.
			r.cache.Apply(e)
			r.queue.Add(e.Object.Name)
			version = e.Object.ResourceVersion
		}

		logf("reflector: watch closed, resuming from version %d", version)
	}
}

Listing 1.18 Watching from a version, in Rust. code/rust/informer/src/reflector.rs

impl Reflector {
    /// Applies events from `version` onward, reopens the watch after each
    /// disconnect, and returns when the server reports that the history is gone.
    fn watch_from(&self, mut version: u64) {
        loop {
            let events = match self.server.watch(version) {
                Ok(events) => events,
                Err(ApiError::Gone) => {
                    log!(
                        "reflector: watch from version {version}: {}",
                        ApiError::Gone
                    );

                    return;
                }
                Err(_) => {
                    thread::sleep(Duration::from_millis(10)); // client-go backs off exponentially
                    continue;
                }
            };

            for event in events {
                version = event.object.resource_version;
                let name = event.object.name.clone();
                // Update the cache before queueing the name, so that a worker
                // that takes the name reads the object as of this event.
                self.cache.apply(event);
                self.queue.add([name]);
            }

            log!("reflector: watch closed, resuming from version {version}");
        }
    }
}

The heart of the reflector is watchFrom (watch_from in Rust). It keeps a variable, version, that always holds the version of the last event it applied to the cache. That means our cache is complete up to version, so a new watch opened from version picks up exactly where the old one stopped. With that in mind, here is what the function does for each of the three ways a watch can end:

• If the watch channel closes because the server dropped the connection, we log the version and open a new watch from it. The server replays anything that happened while we were reconnecting.
• If Watch fails because the server is unreachable, we wait 10 milliseconds and try again with the same version. client-go waits a little longer after each failure in a row, so that thousands of clients do not hammer a server that is trying to recover.
• If Watch fails with the gone error, our version is older than the compaction point, and no watch can continue from it. watchFrom returns, and Run starts over with a new list.

Look closely at the order of the two calls inside the event loop: we update the cache first, and only then queue the object’s name. Suppose we did it the other way round. A worker could take the name straight away, read the object from the cache before our update landed, see the old state, decide there was nothing to do, and finish. The name would not be queued again, and the change would sit unhandled until the object happened to change once more. Updating the cache first rules that out.

client-go uses the same name for this loop, reflector, because it reflects the server’s state into local memory. A reflector together with its cache, plus the plumbing that tells other code about changes, is what client-go calls an informer.

A queue of names

The reflector knows when an object changes, and the controller needs to hear about it. The obvious design is for the reflector to call the controller once for every event. Let us see why that design breaks down.

The first problem is bursts. A user who edits a job five times in one second produces five events. A controller called five times would start work for four specs that were already out of date before the work finished. The second problem is lost events. After a relist, Replace tells us which names changed, but the individual events that led there are gone for good. A controller that relies on seeing every event would have nothing to react to, and it would miss changes every time the history was compacted.

So we give the controller only the names of objects that changed, and the controller reads the object’s current state from the cache when it gets to the name. Five edits now put the same name in the queue five times, and if the queue simply ignores a name it already holds, they turn into one reconcile that acts on the latest spec. The controller takes names from the queue in one or more workers, each processing one name at a time. Our model runs a single worker.

The queue’s state is a list and two sets behind one lock.

Listing 1.19 The queue's state, in Go. code/go/informer/queue.go

// WorkQueue stores the names of objects that need reconciling. A name waits
// in the queue at most once, and a name that a worker holds is handed to no
// other worker until the first calls Done.
type WorkQueue struct {
	// mu guards every field below.
	mu sync.Mutex
	// cond wakes workers blocked in Get.
	cond *sync.Cond
	// queue holds the names ready for a worker, in arrival order.
	queue []string
	// dirty holds every name that needs processing.
	dirty map[string]bool
	// processing holds every name that a worker currently holds.
	processing map[string]bool
}

// NewWorkQueue returns an empty queue.
func NewWorkQueue() *WorkQueue {
	q := &WorkQueue{dirty: map[string]bool{}, processing: map[string]bool{}}
	q.cond = sync.NewCond(&q.mu)

	return q
}

Listing 1.20 The queue's state, in Rust. code/rust/informer/src/queue.rs

/// The queue's data, kept behind one lock.
#[derive(Default)]
struct State {
    /// The names ready for a worker, in arrival order.
    ready: VecDeque<String>,
    /// Every name that needs processing.
    dirty: HashSet<String>,
    /// Every name that a worker currently holds.
    processing: HashSet<String>,
}

/// Stores the names of objects that need reconciling. A name waits in the
/// queue at most once, and a name that a worker holds is handed to no other
/// worker until the first calls [`WorkQueue::done`].
#[derive(Default)]
pub struct WorkQueue {
    /// Guards the queue's data.
    state: Mutex<State>,
    /// Wakes workers blocked in [`WorkQueue::get`].
    ready: Condvar,
}

impl WorkQueue {
    /// Locks the queue's data.
    fn state(&self) -> MutexGuard<'_, State> {
        self.state.lock().expect("queue lock poisoned")
    }
}

Besides the list of names that are ready for a worker, the queue keeps two sets. dirty holds every name that needs processing. processing holds every name that a worker has taken and not yet handed back.

Three operations work on that state.

Listing 1.21 Add, Get, and Done, in Go. code/go/informer/queue.go

// Add marks each name as needing processing. A name that is already pending
// is left where it is. A name that a worker holds is queued again by Done.
func (q *WorkQueue) Add(names ...string) {
	q.mu.Lock()
	defer q.mu.Unlock()

	for _, name := range names {
		if q.dirty[name] {
			continue
		}

		q.dirty[name] = true

		if !q.processing[name] {
			q.queue = append(q.queue, name)
			q.cond.Signal()
		}
	}
}

// Get blocks until a name is ready and hands it to the caller.
func (q *WorkQueue) Get() string {
	q.mu.Lock()
	defer q.mu.Unlock()

	for len(q.queue) == 0 {
		q.cond.Wait()
	}

	name := q.queue[0]
	q.queue = q.queue[1:]
	delete(q.dirty, name)
	q.processing[name] = true

	return name
}

// Done releases a name that Get handed out, and queues it again if it was
// added while the worker held it.
func (q *WorkQueue) Done(name string) {
	q.mu.Lock()
	defer q.mu.Unlock()
	delete(q.processing, name)

	if q.dirty[name] {
		q.queue = append(q.queue, name)
		q.cond.Signal()
	}
}

Listing 1.22 Add, Get, and Done, in Rust. code/rust/informer/src/queue.rs

impl WorkQueue {
    /// Marks each name as needing processing. A name that is already pending
    /// is left where it is. A name that a worker holds is queued again by `done`.
    pub fn add(&self, names: impl IntoIterator<Item = String>) {
        let mut state = self.state();

        for name in names {
            if !state.dirty.insert(name.clone()) {
                continue;
            }

            if !state.processing.contains(&name) {
                state.ready.push_back(name);
                self.ready.notify_one();
            }
        }
    }

    /// Blocks until a name is ready and hands it to the caller.
    pub fn get(&self) -> String {
        let mut state = self.state();

        loop {
            if let Some(name) = state.ready.pop_front() {
                state.dirty.remove(&name);
                state.processing.insert(name.clone());

                return name;
            }

            state = self.ready.wait(state).expect("queue lock poisoned");
        }
    }

    /// Releases a name that `get` handed out, and queues it again if it was
    /// added while the worker held it.
    pub fn done(&self, name: &str) {
        let mut state = self.state();
        state.processing.remove(name);

        if state.dirty.contains(name) {
            state.ready.push_back(name.to_string());
            self.ready.notify_one();
        }
    }
}

Why the second set? Because an object can change while a worker is in the middle of reconciling it. Without processing, the new event would put the name back in the ready list, a second worker would take it, and two workers would be acting on the same object at once. If both write a status with a version check, one of them fails with a conflict and its work is wasted. If both do something outside Kubernetes, like starting a container, both succeed, and our job ends up with two containers. So the queue holds the name back instead. When Add sees a name that a worker is holding, it marks the name dirty but leaves it out of the ready list. When the worker calls Done, the queue notices the mark and makes the name ready again, and the next reconcile sees the object as it is after the change. The table follows one name through that sequence:

call             ready      dirty      processing   effect
---------------  ---------  ---------  -----------  --------------------------------
Add(job-a)       [job-a]    {job-a}    {}           ready for a worker
Get() -> job-a   []         {}         {job-a}      a worker takes it
Add(job-a)       []         {job-a}    {job-a}      edited during the run: marked
Add(job-a)       []         {job-a}    {job-a}      already marked: no change
Done(job-a)      [job-a]    {job-a}    {}           marked, so ready again

Animation 1.1 plays that sequence with the burst of five edits from the run at the end of this chapter. Its last part repeats the calls on a queue that has no processing set and gives the name to a second worker.

One more detail: a worker that calls Get on an empty queue has to go to sleep until a name arrives, and it must not hold the lock while it sleeps, or Add could never get in to add one. Both versions solve this with a condition variable, sync.Cond in Go and std::sync::Condvar in Rust. Waiting on a condition variable releases the lock and puts the worker to sleep in one step. Add and Done signal it whenever they make a name ready. The wait sits inside a loop that checks the list again after waking up, because with several workers another one may have taken the name first.

Reconciling

Now we can write the controller itself. It takes a name from the queue, reads the object from the cache, and makes the world match it. Kubernetes calls this reconciling: bringing what actually exists into agreement with what the user asked for.

The part that does the work is one function.

Listing 1.23 Reconciling one name, in Go. code/go/informer/controller.go

// reconcile reads the named object from the cache and, if its status differs
// from its spec, writes a new status conditioned on the version it read.
func (c *Controller) reconcile(name string) error {
	o, ok := c.cache.Get(name)

	if !ok {
		logf("reconcile %s: deleted, releasing its GPUs", name)

		return nil
	}

	if o.Status.GPUs == o.Spec.GPUs {
		return nil
	}

	time.Sleep(20 * time.Millisecond) // slow work, such as starting containers
	o.Status.GPUs = o.Spec.GPUs

	if err := c.server.Update(o); err != nil {
		logf("reconcile %s: write at version %d: %v", name, o.ResourceVersion, err)

		return err
	}

	logf("reconcile %s: status.GPUs = %d", name, o.Status.GPUs)

	return nil
}

Listing 1.24 Reconciling one name, in Rust. code/rust/informer/src/controller.rs

impl Controller {
    /// Reads the named object from the cache and, if its status differs from
    /// its spec, writes a new status conditioned on the version it read.
    fn reconcile(&self, name: &str) -> Result<(), ApiError> {
        let Some(mut object) = self.cache.get(name) else {
            log!("reconcile {name}: deleted, releasing its GPUs");

            return Ok(());
        };

        if object.status.gpus == object.spec.gpus {
            return Ok(());
        }

        thread::sleep(Duration::from_millis(20)); // slow work, such as starting containers
        object.status.gpus = object.spec.gpus;
        let version = object.resource_version;

        match self.server.update(object.clone()) {
            Ok(()) => {
                log!("reconcile {name}: status.gpus = {}", object.status.gpus);

                Ok(())
            }
            Err(error) => {
                log!("reconcile {name}: write at version {version}: {error}");

                Err(error)
            }
        }
    }
}

reconcile first looks the object up in the cache. If the object is missing, it was deleted, and a real controller would now release whatever it had set up for it; our model logs a message. If the status already matches the spec, there is nothing to do. Otherwise the controller does its work, which our model stands in for with a 20-millisecond sleep, sets the status, and writes the object back.

Notice that reconcile looks only at the object’s current state. It does not know, or care, why it was called. A name queued once and a name queued ten times lead to the same action, and so do a change reported by an event and a change discovered by a relist. Engineers call this style level-triggered, a term borrowed from digital electronics, where a level-triggered circuit responds to the present level of a signal and an edge-triggered circuit responds to each change of the signal. A controller that reacted to individual events would be edge-triggered, and it would do wasted work during bursts of edits and miss changes whose events a compaction had discarded.

Around it sits the loop that feeds it names.

Listing 1.25 The controller and its worker loop, in Go. code/go/informer/controller.go

// Controller makes each object's status agree with its spec.
type Controller struct {
	// server receives the controller's status writes.
	server *APIServer
	// cache supplies the current state of each object.
	cache *Cache
	// queue supplies the names of objects to reconcile.
	queue *WorkQueue
}

// Run is a single worker: it takes names from the queue one at a time and
// reconciles them. It does not return.
func (c *Controller) Run() {
	for {
		name := c.queue.Get()
		err := c.reconcile(name)
		c.queue.Done(name)

		if err != nil {
			c.queue.Add(name) // client-go requeues with a backoff
		}
	}
}

Listing 1.26 The controller and its worker loop, in Rust. code/rust/informer/src/controller.rs

/// Makes each object's status agree with its spec.
pub struct Controller {
    /// Receives the controller's status writes.
    pub server: Arc<ApiServer>,
    /// Supplies the current state of each object.
    pub cache: Arc<Cache>,
    /// Supplies the names of objects to reconcile.
    pub queue: Arc<WorkQueue>,
}

impl Controller {
    /// A single worker: takes names from the queue one at a time and
    /// reconciles them. It does not return.
    pub fn run(&self) {
        loop {
            let name = self.queue.get();
            let result = self.reconcile(&name);
            self.queue.done(&name);

            if result.is_err() {
                self.queue.add([name]); // client-go requeues with a backoff
            }
        }
    }
}

Run is the worker loop. It takes a name, reconciles it, and hands it back with Done. If the reconcile failed, it adds the name again so that the next attempt starts from whatever the cache holds by then.

Writing back to the server

Our cache always trails the server by the time it takes an event to arrive. During the 20 milliseconds that reconcile spends working, a user might change the spec. If the controller then wrote its copy of the object back unconditionally, it would store the old spec together with a status computed for the old spec, and the user’s edit would simply disappear.

This is exactly the situation the resource version check was made for. reconcile writes the object with the version it read. If the object changed in the meantime, the server refuses the write with a conflict, Run puts the name back in the queue, and the next reconcile reads the newer object from the cache. This approach is called optimistic concurrency: we take no lock on the object, we assume nobody else will change it, and we rely on the server to catch the cases where that assumption was wrong at the moment we write.

Running the model

Figure 1.1 shows the five parts together and numbers the two steps the reflector takes for each event.

Time to see it all work together. main creates two jobs, starts the reflector and the controller, and then puts the machinery through each situation from this chapter: a burst of edits that arrives while a reconcile is running, a dropped watch, and a partition during which a job is deleted and the history is compacted. The Go version runs the reflector and the controller as goroutines. The Rust version runs them on threads, each holding an Arc handle to the shared server, cache, and queue. Both write their log to standard error, with the number of milliseconds since the program started; Rust does this with a small log! macro, and Go with a logf function.

Listing 1.27 The demonstration, in Go. code/go/informer/main.go

func main() {
	log.SetFlags(0)
	server := NewAPIServer()
	server.Update(Object{Name: "job-a", Spec: Spec{GPUs: 1}})
	server.Update(Object{Name: "job-b", Spec: Spec{GPUs: 8}})

	cache := NewCache()
	queue := NewWorkQueue()
	go (&Reflector{server, cache, queue}).Run()
	go (&Controller{server, cache, queue}).Run()
	time.Sleep(100 * time.Millisecond)

	logf("--- job-a is edited five times in a burst")

	for gpus := 2; gpus <= 6; gpus++ {
		server.Update(Object{Name: "job-a", Spec: Spec{GPUs: gpus}})

		if gpus == 2 {
			time.Sleep(5 * time.Millisecond) // the controller starts on the first edit
		}
	}

	time.Sleep(150 * time.Millisecond)

	logf("--- the watch connection drops")
	server.DropWatches()
	time.Sleep(50 * time.Millisecond)

	logf("--- partition: job-b is deleted and the history compacted meanwhile")
	server.Partition()
	server.Delete("job-b")
	server.Compact()
	time.Sleep(30 * time.Millisecond)
	server.Heal()
	time.Sleep(150 * time.Millisecond)

	list, version := server.List()
	logf("final state at version %d: %v", version, list)
}

Listing 1.28 The demonstration, in Rust. code/rust/informer/src/main.rs

fn main() {
    START.get_or_init(Instant::now);
    let server = Arc::new(ApiServer::default());
    server
        .update(Object::new("job-a", 1))
        .expect("unconditional write");
    server
        .update(Object::new("job-b", 8))
        .expect("unconditional write");

    let cache = Arc::new(Cache::default());
    let queue = Arc::new(WorkQueue::default());
    let reflector = Reflector {
        server: Arc::clone(&server),
        cache: Arc::clone(&cache),
        queue: Arc::clone(&queue),
    };
    let controller = Controller {
        server: Arc::clone(&server),
        cache,
        queue,
    };
    thread::spawn(move || reflector.run());
    thread::spawn(move || controller.run());
    sleep_ms(100);

    log!("--- job-a is edited five times in a burst");

    for gpus in 2..=6 {
        server
            .update(Object::new("job-a", gpus))
            .expect("unconditional write");
        if gpus == 2 {
            sleep_ms(5); // the controller starts on the first edit
        }
    }

    sleep_ms(150);

    log!("--- the watch connection drops");
    server.drop_watches();
    sleep_ms(50);

    log!("--- partition: job-b is deleted and the history compacted meanwhile");
    server.partition();
    server.delete("job-b");
    server.compact();
    sleep_ms(30);
    server.heal();
    sleep_ms(150);

    let (list, version) = server.list();
    let objects: Vec<String> = list.iter().map(ToString::to_string).collect();
    log!("final state at version {version}: [{}]", objects.join(" "));
}

The Go version prints:

$ go run ./informer
   0ms  reflector: list at version 2: [job-a{spec.GPUs: 1, status.GPUs: 0, v1} job-b{spec.GPUs: 8, status.GPUs: 0, v2}]
  20ms  reconcile job-a: status.GPUs = 1
  40ms  reconcile job-b: status.GPUs = 8
 101ms  --- job-a is edited five times in a burst
 121ms  reconcile job-a: write at version 5: 409 Conflict: object has been modified
 141ms  reconcile job-a: status.GPUs = 6
 257ms  --- the watch connection drops
 257ms  reflector: watch closed, resuming from version 10
 307ms  --- partition: job-b is deleted and the history compacted meanwhile
 307ms  reflector: watch closed, resuming from version 10
 338ms  reflector: watch from version 10: 410 Gone: resource version too old
 338ms  reflector: list at version 11: [job-a{spec.GPUs: 6, status.GPUs: 6, v10}]
 338ms  reconcile job-b: deleted, releasing its GPUs
 489ms  final state at version 11: [job-a{spec.GPUs: 6, status.GPUs: 6, v10}]

The Rust version prints the same sequence of events. Its timings differ by a millisecond or two, and its field names follow Rust’s naming conventions:

$ cargo run -p informer
   0ms  reflector: list at version 2: [job-a{spec.gpus: 1, status.gpus: 0, v1} job-b{spec.gpus: 8, status.gpus: 0, v2}]
  20ms  reconcile job-a: status.gpus = 1
  40ms  reconcile job-b: status.gpus = 8
 100ms  --- job-a is edited five times in a burst
 120ms  reconcile job-a: write at version 5: 409 Conflict: object has been modified
 141ms  reconcile job-a: status.gpus = 6
 255ms  --- the watch connection drops
 256ms  reflector: watch closed, resuming from version 10
 306ms  --- partition: job-b is deleted and the history compacted meanwhile
 306ms  reflector: watch closed, resuming from version 10
 336ms  reflector: watch from version 10: 410 Gone: resource version too old
 337ms  reflector: list at version 11: [job-a{spec.gpus: 6, status.gpus: 6, v10}]
 337ms  reconcile job-b: deleted, releasing its GPUs
 486ms  final state at version 11: [job-a{spec.gpus: 6, status.gpus: 6, v10}]

Let us read the log from the top. The first three lines are the initial list and one reconcile for each job. Each status write is itself a change, so the reflector receives an event for it and queues the name again. The reconcile that follows finds the status equal to the spec and returns quietly, so each job appears only once.

The burst of five edits led to just two reconciles of job-a. The controller took the name right after the first edit, read a spec of 2 GPUs at version 5, and started its 20 milliseconds of work. The other four edits arrived while it was busy, moving the object to version 9, and because the controller was holding the name, the queue marked it dirty. When the controller wrote its status with version 5, the server refused it with a conflict. Done made the name ready again, the next reconcile read a spec of 6 GPUs at version 9, and it wrote a status of 6 at version 10. The specs of 3, 4, and 5 GPUs had been replaced before the controller ever read them, so it did no work for them at all.

When the watch was dropped, the reflector reopened it from version 10, the version of the last event it had applied, and there was no need to list. During the partition, job-b was deleted at version 11 and the history was compacted, which moved the compaction point to 11. When the partition ended, the reflector asked to resume from version 10, the server answered 410 Gone, and the reflector listed again. The new list held only job-a, Replace reported job-b as missing, and the controller found no job-b in the cache and released its GPUs. Our client learned about the deletion by comparing the list with its cache, because the compaction had thrown the deletion event away.

Animation 1.2 replays the log from the list at version 2 to the list at version 11, with the server’s objects on the left and the cache on the right.

The same design in client-go

Real Go controllers use client-go, either directly or through controller-runtime, a higher-level library built on top of it. Every part of our model has a counterpart there:

model                     client-go
------------------------  ---------------------------------------------------------
APIServer.List, Watch     ListerWatcher, backed by HTTP list and watch requests
Reflector                 Reflector
Cache                     Indexer: a thread-safe store with secondary indexes,
                          for example all pods on one node
Reflector and Cache       SharedIndexInformer
WorkQueue                 workqueue, with the same dirty and processing sets
Controller.reconcile      the Reconcile method of a controller

client-go also has a few mechanisms that our model leaves out, and each one fixes a problem we can now recognize.

Between the reflector and the cache, client-go puts a queue called the DeltaFIFO. For each object it collects the changes in the order they arrived, and the code that updates the cache and notifies controllers takes one object at a time along with all of its pending changes. When a relist shows that an object has disappeared, the DeltaFIFO emits a deletion whose object is wrapped in a DeletedFinalStateUnknown marker. The marker warns event handlers that the object they receive is the last state the cache had seen, which may be older than the state the object was in when it was deleted.

The API server can send bookmark events, which carry a resource version and no object. They solve a problem with quiet watches. A watch on a kind of object that rarely changes can go a long time without any events, so the version it would resume from grows old. After the next compaction, reconnecting from that version would fail with 410 Gone and force a full relist. Bookmarks move the watch’s position forward even when nothing changes, so a reconnect usually succeeds.

An informer can be given a resync period. Every so often it queues the name of every object in its cache again, straight from memory and without asking the server for anything. This is for controllers that manage something outside Kubernetes, such as a cloud load balancer. If someone changes that load balancer by hand, no Kubernetes event tells the controller, and a periodic resync gives it a regular chance to notice and fix the drift.

The work queue waits before retrying a name whose reconcile failed. By default the wait starts at 5 milliseconds and doubles after each failure, up to 1,000 seconds, and the queue also limits how fast it hands out names overall. Our model retries at once. A real controller that did that while its writes kept failing, say because of a bug in how it builds the object, would flood the API server with requests.

Finally, informers are shared within a process. If ten controllers in one program all care about pods, they share a single pod watch and a single cache through a SharedInformerFactory, and each registers its own event handlers and its own work queue. The API server then serves one watch per process instead of one per controller.

Sharing a cache safely

Our model’s cache hands out copies, so a caller can change what it gets without affecting anyone else. client-go’s cache hands out pointers into the shared store instead, to avoid copying every object on every read. That has a sharp edge. If a client-go controller changes a field of an object it read from the cache, it is changing the very object every other controller in the process reads, while the server still holds the original. The other controllers now see a state that exists nowhere else. The rule is to call the object’s generated DeepCopy method, which copies every nested slice and map, and to change the copy.

Reading your own writes

A controller’s own writes reach its cache the same way everyone else’s do: the server stores the write, emits an event, and the reflector applies it. Until that event arrives, the cache still shows the state from before the write. Picture a reconcile that creates a pod for a job. If a second event for the same job arrives before the pod’s creation event, the next reconcile finds no pod in the cache and creates a second one.

Kubernetes controllers guard against this in two ways. The ReplicaSet controller, which keeps a requested number of identical pods running, records expectations: after creating three pods, it notes that it expects to see three creation events, and it leaves that ReplicaSet alone until the cache has shown them or a timeout has passed. The second technique uses two fields that every object carries. The API server increments metadata.generation every time the spec changes, and a controller copies the generation it acted on into status.observedGeneration. When the two numbers are equal, the status describes the current spec. When they differ, both the controller and any user reading the object can see that the latest change to the spec has not been handled yet.

The complete program

The model is six source files in each language, printed here in the order the chapter built them.

Listing 1.29 The API server, complete, in Go. code/go/informer/server.go

// server.go models the Kubernetes API server. It stores objects, numbers every
// write, keeps a history of recent changes, and serves lists and watches.

package main

import (
	"errors"
	"fmt"
	"maps"
	"slices"
	"strings"
	"sync"
)

// Spec is the part of an object that a user writes: the desired state.
type Spec struct {
	// GPUs is the number of GPUs the job asks for.
	GPUs int
}

// Status is the part of an object that the controller writes: the observed state.
type Status struct {
	// GPUs is the number of GPUs the controller has provided to the job.
	GPUs int
}

// Object is one stored record, such as a training job.
type Object struct {
	// Name identifies the object on the server.
	Name string
	// Spec holds what the user asked for.
	Spec Spec
	// Status holds what the controller has done about it.
	Status Status
	// ResourceVersion is the server's write counter at the object's latest change.
	ResourceVersion int
}

// String formats an object for the demonstration's log lines.
func (o Object) String() string {
	return fmt.Sprintf("%s{spec.GPUs: %d, status.GPUs: %d, v%d}",
		o.Name, o.Spec.GPUs, o.Status.GPUs, o.ResourceVersion)
}

// EventType is the kind of change that an Event records.
type EventType int

// The kinds of change that a watch reports.
const (
	Added    EventType = iota // the object was created
	Modified                  // the object was updated
	Deleted                   // the object was removed
)

// String returns the name that the Kubernetes watch protocol uses for t.
func (t EventType) String() string {
	return [...]string{"ADDED", "MODIFIED", "DELETED"}[t]
}

// Event records one change to one object.
type Event struct {
	// Type is the kind of change.
	Type EventType
	// Object is the object after the change, or its final state if it was deleted.
	Object Object
}

// Errors that the server returns. The first two carry the HTTP status codes
// that the Kubernetes API server uses for the same conditions.
var (
	// ErrGone reports that the history a watch needs has been discarded.
	ErrGone = errors.New("410 Gone: resource version too old")
	// ErrConflict reports a conditional write whose resource version is out of date.
	ErrConflict = errors.New("409 Conflict: object has been modified")
	// ErrUnreachable reports that the client cannot reach the server.
	ErrUnreachable = errors.New("server unreachable")
)


// APIServer stores objects and streams their changes to watchers. Its methods
// are safe for concurrent use.
type APIServer struct {
	// mu guards every field below.
	mu sync.Mutex
	// version is the last resource version the server handed out.
	version int
	// objects holds the current state of every object, keyed by name.
	objects map[string]Object
	// history holds the events newer than compacted, oldest first.
	history []Event
	// compacted is the oldest version from which a watch may resume.
	compacted int
	// watches holds one channel for each open watch.
	watches []chan Event
	// unreachable makes new watches fail, to simulate a network partition.
	unreachable bool
}

// NewAPIServer returns an empty server.
func NewAPIServer() *APIServer {
	return &APIServer{objects: map[string]Object{}}
}


// Update creates or replaces an object and returns ErrConflict if the write's
// condition fails. A non-zero o.ResourceVersion is a condition: the caller read
// the object at that version and wants the write applied only if the object is
// still at it. A zero ResourceVersion writes unconditionally.
func (s *APIServer) Update(o Object) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	current, exists := s.objects[o.Name]

	if o.ResourceVersion != 0 && o.ResourceVersion != current.ResourceVersion {
		return ErrConflict
	}

	eventType := Modified

	if !exists {
		eventType = Added
	}

	s.version++
	o.ResourceVersion = s.version
	s.objects[o.Name] = o
	s.publish(Event{eventType, o})

	return nil
}

// Delete removes the named object, if it exists.
func (s *APIServer) Delete(name string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	o, exists := s.objects[name]

	if !exists {
		return
	}

	delete(s.objects, name)
	s.version++
	o.ResourceVersion = s.version
	s.publish(Event{Deleted, o})
}

// publish appends e to the history and sends it to every open watch.
// The caller must hold s.mu.
func (s *APIServer) publish(e Event) {
	s.history = append(s.history, e)

	for _, w := range s.watches {
		w <- e
	}
}


// List returns every object, sorted by name, and the version at which the
// snapshot was taken.
func (s *APIServer) List() ([]Object, int) {
	s.mu.Lock()
	defer s.mu.Unlock()
	byName := func(a, b Object) int { return strings.Compare(a.Name, b.Name) }

	return slices.SortedFunc(maps.Values(s.objects), byName), s.version
}

// Watch returns a channel that first delivers the stored events newer than
// version and then every new event, until the server closes it.
func (s *APIServer) Watch(version int) (<-chan Event, error) {
	s.mu.Lock()
	defer s.mu.Unlock()

	switch {
	case s.unreachable:
		return nil, ErrUnreachable
	case version < s.compacted:
		return nil, ErrGone
	}

	w := make(chan Event, 1024)

	for _, e := range s.history {
		if e.Object.ResourceVersion > version {
			w <- e
		}
	}

	s.watches = append(s.watches, w)

	return w, nil
}


// Compact discards the event history, as etcd discards old versions.
func (s *APIServer) Compact() {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.compacted = s.version
	s.history = nil
}

// DropWatches closes every open watch, as a load balancer's idle timeout would.
func (s *APIServer) DropWatches() {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.dropWatches()
}

// Partition closes every open watch and makes new watches fail until Heal.
func (s *APIServer) Partition() {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.unreachable = true
	s.dropWatches()
}

// Heal ends a partition, so new watches succeed again.
func (s *APIServer) Heal() {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.unreachable = false
}

// dropWatches closes and forgets every watch channel. The caller must hold s.mu.
func (s *APIServer) dropWatches() {
	for _, w := range s.watches {
		close(w)
	}

	s.watches = nil
}

Listing 1.30 The API server, complete, in Rust. code/rust/informer/src/server.rs

//! A model of the Kubernetes API server. It stores objects, numbers every
//! write, keeps a history of recent changes, and serves lists and watches.

use std::collections::HashMap;
use std::fmt;
use std::sync::Mutex;
use std::sync::mpsc::{Receiver, Sender, channel};

/// The part of an object that a user writes: the desired state.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct Spec {
    /// Number of GPUs the job asks for.
    pub gpus: u32,
}

/// The part of an object that the controller writes: the observed state.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct Status {
    /// Number of GPUs the controller has provided to the job.
    pub gpus: u32,
}

/// One stored record, such as a training job.
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct Object {
    /// Identifies the object on the server.
    pub name: String,
    /// What the user asked for.
    pub spec: Spec,
    /// What the controller has done about it.
    pub status: Status,
    /// The server's write counter at the object's latest change.
    pub resource_version: u64,
}

impl Object {
    /// Returns a new object with the given name and GPU request, an empty
    /// status, and no resource version.
    pub fn new(name: &str, gpus: u32) -> Self {
        Self {
            name: name.to_string(),
            spec: Spec { gpus },
            ..Self::default()
        }
    }
}

impl fmt::Display for Object {
    /// Formats an object for the demonstration's log lines.
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "{}{{spec.gpus: {}, status.gpus: {}, v{}}}",
            self.name, self.spec.gpus, self.status.gpus, self.resource_version
        )
    }
}

/// The kind of change that an [`Event`] records.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum EventType {
    /// The object was created.
    Added,
    /// The object was updated.
    Modified,
    /// The object was removed.
    Deleted,
}

/// One change to one object.
#[derive(Clone, Debug)]
pub struct Event {
    /// The kind of change.
    pub kind: EventType,
    /// The object after the change, or its final state if it was deleted.
    pub object: Object,
}

/// Errors that the server returns. The first two carry the HTTP status codes
/// that the Kubernetes API server uses for the same conditions.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ApiError {
    /// The history that a watch needs has been discarded.
    Gone,
    /// A conditional write carried an out-of-date resource version.
    Conflict,
    /// The client cannot reach the server.
    Unreachable,
}

impl fmt::Display for ApiError {
    /// Formats the error the way the Kubernetes API server reports it.
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(match self {
            Self::Gone => "410 Gone: resource version too old",
            Self::Conflict => "409 Conflict: object has been modified",
            Self::Unreachable => "server unreachable",
        })
    }
}

/// Everything the server stores, kept behind one lock.
#[derive(Default)]
struct State {
    /// The last resource version the server handed out.
    version: u64,
    /// The current state of every object, keyed by name.
    objects: HashMap<String, Object>,
    /// The events newer than `compacted`, oldest first.
    history: Vec<Event>,
    /// The oldest version from which a watch may resume.
    compacted: u64,
    /// One sender for each open watch.
    watches: Vec<Sender<Event>>,
    /// Makes new watches fail, to simulate a network partition.
    unreachable: bool,
}

impl State {
    /// Appends `event` to the history and sends it to every open watch.
    /// Watches whose receiver has gone away are forgotten.
    fn publish(&mut self, event: Event) {
        self.watches.retain(|w| w.send(event.clone()).is_ok());
        self.history.push(event);
    }
}

/// Stores objects and streams their changes to watchers. Its methods are safe
/// to call from several threads.
#[derive(Default)]
pub struct ApiServer {
    /// Guards all of the server's data.
    state: Mutex<State>,
}

impl ApiServer {
    /// Locks the server's data.
    fn state(&self) -> std::sync::MutexGuard<'_, State> {
        self.state.lock().expect("server lock poisoned")
    }
}

impl ApiServer {
    /// Creates or replaces an object. A non-zero `resource_version` is a
    /// condition: the caller read the object at that version and wants the
    /// write applied only if the object is still at it. Zero writes
    /// unconditionally.
    pub fn update(&self, mut object: Object) -> Result<(), ApiError> {
        let mut state = self.state();
        let current = state.objects.get(&object.name);
        let current_version = current.map_or(0, |o| o.resource_version);

        if object.resource_version != 0 && object.resource_version != current_version {
            return Err(ApiError::Conflict);
        }

        let kind = if current.is_some() {
            EventType::Modified
        } else {
            EventType::Added
        };
        state.version += 1;
        object.resource_version = state.version;
        state.objects.insert(object.name.clone(), object.clone());
        state.publish(Event { kind, object });

        Ok(())
    }

    /// Removes the named object, if it exists.
    pub fn delete(&self, name: &str) {
        let mut state = self.state();
        let Some(mut object) = state.objects.remove(name) else {
            return;
        };
        state.version += 1;
        object.resource_version = state.version;
        state.publish(Event {
            kind: EventType::Deleted,
            object,
        });
    }
}

impl ApiServer {
    /// Returns every object, sorted by name, and the version at which the
    /// snapshot was taken.
    pub fn list(&self) -> (Vec<Object>, u64) {
        let state = self.state();
        let mut objects: Vec<Object> = state.objects.values().cloned().collect();
        objects.sort_by(|a, b| a.name.cmp(&b.name));

        (objects, state.version)
    }

    /// Returns a receiver that first delivers the stored events newer than
    /// `version` and then every new event, until the server drops the watch.
    pub fn watch(&self, version: u64) -> Result<Receiver<Event>, ApiError> {
        let mut state = self.state();

        if state.unreachable {
            return Err(ApiError::Unreachable);
        }

        if version < state.compacted {
            return Err(ApiError::Gone);
        }

        let (sender, receiver) = channel();

        for event in &state.history {
            if event.object.resource_version > version {
                sender.send(event.clone()).expect("receiver is held here");
            }
        }

        state.watches.push(sender);

        Ok(receiver)
    }
}

impl ApiServer {
    /// Discards the event history, as etcd discards old versions.
    pub fn compact(&self) {
        let mut state = self.state();
        state.compacted = state.version;
        state.history.clear();
    }

    /// Closes every open watch, as a load balancer's idle timeout would.
    pub fn drop_watches(&self) {
        self.state().watches.clear();
    }

    /// Closes every open watch and makes new watches fail until [`Self::heal`].
    pub fn partition(&self) {
        let mut state = self.state();
        state.unreachable = true;
        state.watches.clear();
    }

    /// Ends a partition, so new watches succeed again.
    pub fn heal(&self) {
        self.state().unreachable = false;
    }
}

Listing 1.31 The cache, complete, in Go. code/go/informer/cache.go

// cache.go holds the client's local copy of the server's objects.

package main

import "sync"

// Cache is the client's copy of the server's objects, kept current by a
// Reflector. Its methods are safe for concurrent use.
type Cache struct {
	// mu guards objects.
	mu sync.RWMutex
	// objects holds the latest known state of every object, keyed by name.
	objects map[string]Object
}

// NewCache returns an empty cache.
func NewCache() *Cache {
	return &Cache{objects: map[string]Object{}}
}

// Get returns the cached copy of the named object and whether it exists.
func (c *Cache) Get(name string) (Object, bool) {
	c.mu.RLock()
	defer c.mu.RUnlock()
	o, ok := c.objects[name]

	return o, ok
}

// Apply records one event from a watch.
func (c *Cache) Apply(e Event) {
	c.mu.Lock()
	defer c.mu.Unlock()

	if e.Type == Deleted {
		delete(c.objects, e.Object.Name)

		return
	}

	c.objects[e.Object.Name] = e.Object
}


// Replace installs the result of a fresh list and returns the names of the
// objects that changed, including objects that disappeared while the client
// was not watching.
func (c *Cache) Replace(list []Object) []string {
	c.mu.Lock()
	defer c.mu.Unlock()
	old := c.objects
	c.objects = make(map[string]Object, len(list))
	var changed []string

	for _, o := range list {
		c.objects[o.Name] = o
		// A name missing from old reads as version 0, so it counts as changed.
		if old[o.Name].ResourceVersion != o.ResourceVersion {
			changed = append(changed, o.Name)
		}
	}

	for name := range old {
		if _, ok := c.objects[name]; !ok {
			changed = append(changed, name)
		}
	}

	return changed
}

Listing 1.32 The cache, complete, in Rust. code/rust/informer/src/cache.rs

//! The client's local copy of the server's objects.

use crate::server::{Event, EventType, Object};
use std::collections::HashMap;
use std::sync::RwLock;

/// The client's copy of the server's objects, kept current by a
/// [`Reflector`](crate::reflector::Reflector). Its methods are safe to call
/// from several threads.
#[derive(Default)]
pub struct Cache {
    /// The latest known state of every object, keyed by name.
    objects: RwLock<HashMap<String, Object>>,
}

impl Cache {
    /// Returns the cached copy of the named object, if it exists.
    pub fn get(&self, name: &str) -> Option<Object> {
        let objects = self.objects.read().expect("cache lock poisoned");

        objects.get(name).cloned()
    }

    /// Records one event from a watch.
    pub fn apply(&self, event: Event) {
        let mut objects = self.objects.write().expect("cache lock poisoned");

        match event.kind {
            EventType::Deleted => objects.remove(&event.object.name),
            EventType::Added | EventType::Modified => {
                objects.insert(event.object.name.clone(), event.object)
            }
        };
    }
}

impl Cache {
    /// Installs the result of a fresh list and returns the names of the
    /// objects that changed, including objects that disappeared while the
    /// client was not watching.
    pub fn replace(&self, list: &[Object]) -> Vec<String> {
        let mut objects = self.objects.write().expect("cache lock poisoned");
        let old = std::mem::take(&mut *objects);
        let mut changed = Vec::new();

        for object in list {
            let old_version = old.get(&object.name).map(|o| o.resource_version);

            if old_version != Some(object.resource_version) {
                changed.push(object.name.clone());
            }

            objects.insert(object.name.clone(), object.clone());
        }

        for name in old.keys() {
            if !objects.contains_key(name) {
                changed.push(name.clone());
            }
        }

        changed
    }
}

Listing 1.33 The reflector, complete, in Go. code/go/informer/reflector.go

// reflector.go keeps a Cache current by listing and watching the server.

package main

import (
	"errors"
	"time"
)

// Reflector copies the server's objects into a cache and queues the name of
// every object that changes.
type Reflector struct {
	// server is the source of objects and events.
	server *APIServer
	// cache receives every object the reflector sees.
	cache *Cache
	// queue receives the name of every object that changes.
	queue *WorkQueue
}

// Run lists every object, watches from the list's version, and lists again
// whenever the server can no longer resume the watch. It does not return.
func (r *Reflector) Run() {
	for {
		list, version := r.server.List()
		r.queue.Add(r.cache.Replace(list)...)
		logf("reflector: list at version %d: %v", version, list)
		r.watchFrom(version)
	}
}


// watchFrom applies events from version onward, reopens the watch after each
// disconnect, and returns when the server reports that the history is gone.
func (r *Reflector) watchFrom(version int) {
	for {
		w, err := r.server.Watch(version)

		if errors.Is(err, ErrGone) {
			logf("reflector: watch from version %d: %v", version, err)

			return
		}

		if err != nil {
			time.Sleep(10 * time.Millisecond) // client-go backs off exponentially
			continue
		}

		for e := range w {
			// Update the cache before queueing the name, so that a worker
			// that takes the name reads the object as of this event.
			r.cache.Apply(e)
			r.queue.Add(e.Object.Name)
			version = e.Object.ResourceVersion
		}

		logf("reflector: watch closed, resuming from version %d", version)
	}
}

Listing 1.34 The reflector, complete, in Rust. code/rust/informer/src/reflector.rs

//! Keeps a [`Cache`] current by listing and watching the server.

use crate::cache::Cache;
use crate::log;
use crate::queue::WorkQueue;
use crate::server::{ApiError, ApiServer};
use std::sync::Arc;
use std::thread;
use std::time::Duration;

/// Copies the server's objects into a cache and queues the name of every
/// object that changes.
pub struct Reflector {
    /// The source of objects and events.
    pub server: Arc<ApiServer>,
    /// Receives every object the reflector sees.
    pub cache: Arc<Cache>,
    /// Receives the name of every object that changes.
    pub queue: Arc<WorkQueue>,
}

impl Reflector {
    /// Lists every object, watches from the list's version, and lists again
    /// whenever the server can no longer resume the watch. It does not return.
    pub fn run(&self) {
        loop {
            let (list, version) = self.server.list();
            self.queue.add(self.cache.replace(&list));
            let objects: Vec<String> = list.iter().map(ToString::to_string).collect();
            log!(
                "reflector: list at version {version}: [{}]",
                objects.join(" ")
            );
            self.watch_from(version);
        }
    }
}

impl Reflector {
    /// Applies events from `version` onward, reopens the watch after each
    /// disconnect, and returns when the server reports that the history is gone.
    fn watch_from(&self, mut version: u64) {
        loop {
            let events = match self.server.watch(version) {
                Ok(events) => events,
                Err(ApiError::Gone) => {
                    log!(
                        "reflector: watch from version {version}: {}",
                        ApiError::Gone
                    );

                    return;
                }
                Err(_) => {
                    thread::sleep(Duration::from_millis(10)); // client-go backs off exponentially
                    continue;
                }
            };

            for event in events {
                version = event.object.resource_version;
                let name = event.object.name.clone();
                // Update the cache before queueing the name, so that a worker
                // that takes the name reads the object as of this event.
                self.cache.apply(event);
                self.queue.add([name]);
            }

            log!("reflector: watch closed, resuming from version {version}");
        }
    }
}

Listing 1.35 The work queue, complete, in Go. code/go/informer/queue.go

// queue.go implements the work queue between the reflector and the controller.

package main

import "sync"

// WorkQueue stores the names of objects that need reconciling. A name waits
// in the queue at most once, and a name that a worker holds is handed to no
// other worker until the first calls Done.
type WorkQueue struct {
	// mu guards every field below.
	mu sync.Mutex
	// cond wakes workers blocked in Get.
	cond *sync.Cond
	// queue holds the names ready for a worker, in arrival order.
	queue []string
	// dirty holds every name that needs processing.
	dirty map[string]bool
	// processing holds every name that a worker currently holds.
	processing map[string]bool
}

// NewWorkQueue returns an empty queue.
func NewWorkQueue() *WorkQueue {
	q := &WorkQueue{dirty: map[string]bool{}, processing: map[string]bool{}}
	q.cond = sync.NewCond(&q.mu)

	return q
}


// Add marks each name as needing processing. A name that is already pending
// is left where it is. A name that a worker holds is queued again by Done.
func (q *WorkQueue) Add(names ...string) {
	q.mu.Lock()
	defer q.mu.Unlock()

	for _, name := range names {
		if q.dirty[name] {
			continue
		}

		q.dirty[name] = true

		if !q.processing[name] {
			q.queue = append(q.queue, name)
			q.cond.Signal()
		}
	}
}

// Get blocks until a name is ready and hands it to the caller.
func (q *WorkQueue) Get() string {
	q.mu.Lock()
	defer q.mu.Unlock()

	for len(q.queue) == 0 {
		q.cond.Wait()
	}

	name := q.queue[0]
	q.queue = q.queue[1:]
	delete(q.dirty, name)
	q.processing[name] = true

	return name
}

// Done releases a name that Get handed out, and queues it again if it was
// added while the worker held it.
func (q *WorkQueue) Done(name string) {
	q.mu.Lock()
	defer q.mu.Unlock()
	delete(q.processing, name)

	if q.dirty[name] {
		q.queue = append(q.queue, name)
		q.cond.Signal()
	}
}

Listing 1.36 The work queue, complete, in Rust. code/rust/informer/src/queue.rs

//! The work queue between the reflector and the controller.

use std::collections::{HashSet, VecDeque};
use std::sync::{Condvar, Mutex, MutexGuard};

/// The queue's data, kept behind one lock.
#[derive(Default)]
struct State {
    /// The names ready for a worker, in arrival order.
    ready: VecDeque<String>,
    /// Every name that needs processing.
    dirty: HashSet<String>,
    /// Every name that a worker currently holds.
    processing: HashSet<String>,
}

/// Stores the names of objects that need reconciling. A name waits in the
/// queue at most once, and a name that a worker holds is handed to no other
/// worker until the first calls [`WorkQueue::done`].
#[derive(Default)]
pub struct WorkQueue {
    /// Guards the queue's data.
    state: Mutex<State>,
    /// Wakes workers blocked in [`WorkQueue::get`].
    ready: Condvar,
}

impl WorkQueue {
    /// Locks the queue's data.
    fn state(&self) -> MutexGuard<'_, State> {
        self.state.lock().expect("queue lock poisoned")
    }
}

impl WorkQueue {
    /// Marks each name as needing processing. A name that is already pending
    /// is left where it is. A name that a worker holds is queued again by `done`.
    pub fn add(&self, names: impl IntoIterator<Item = String>) {
        let mut state = self.state();

        for name in names {
            if !state.dirty.insert(name.clone()) {
                continue;
            }

            if !state.processing.contains(&name) {
                state.ready.push_back(name);
                self.ready.notify_one();
            }
        }
    }

    /// Blocks until a name is ready and hands it to the caller.
    pub fn get(&self) -> String {
        let mut state = self.state();

        loop {
            if let Some(name) = state.ready.pop_front() {
                state.dirty.remove(&name);
                state.processing.insert(name.clone());

                return name;
            }

            state = self.ready.wait(state).expect("queue lock poisoned");
        }
    }

    /// Releases a name that `get` handed out, and queues it again if it was
    /// added while the worker held it.
    pub fn done(&self, name: &str) {
        let mut state = self.state();
        state.processing.remove(name);

        if state.dirty.contains(name) {
            state.ready.push_back(name.to_string());
            self.ready.notify_one();
        }
    }
}

Listing 1.37 The controller, complete, in Go. code/go/informer/controller.go

// controller.go reconciles objects whose names arrive on the work queue.

package main

import "time"

// Controller makes each object's status agree with its spec.
type Controller struct {
	// server receives the controller's status writes.
	server *APIServer
	// cache supplies the current state of each object.
	cache *Cache
	// queue supplies the names of objects to reconcile.
	queue *WorkQueue
}

// Run is a single worker: it takes names from the queue one at a time and
// reconciles them. It does not return.
func (c *Controller) Run() {
	for {
		name := c.queue.Get()
		err := c.reconcile(name)
		c.queue.Done(name)

		if err != nil {
			c.queue.Add(name) // client-go requeues with a backoff
		}
	}
}


// reconcile reads the named object from the cache and, if its status differs
// from its spec, writes a new status conditioned on the version it read.
func (c *Controller) reconcile(name string) error {
	o, ok := c.cache.Get(name)

	if !ok {
		logf("reconcile %s: deleted, releasing its GPUs", name)

		return nil
	}

	if o.Status.GPUs == o.Spec.GPUs {
		return nil
	}

	time.Sleep(20 * time.Millisecond) // slow work, such as starting containers
	o.Status.GPUs = o.Spec.GPUs

	if err := c.server.Update(o); err != nil {
		logf("reconcile %s: write at version %d: %v", name, o.ResourceVersion, err)

		return err
	}

	logf("reconcile %s: status.GPUs = %d", name, o.Status.GPUs)

	return nil
}

Listing 1.38 The controller, complete, in Rust. code/rust/informer/src/controller.rs

//! Reconciles objects whose names arrive on the work queue.

use crate::cache::Cache;
use crate::log;
use crate::queue::WorkQueue;
use crate::server::{ApiError, ApiServer};
use std::sync::Arc;
use std::thread;
use std::time::Duration;

/// Makes each object's status agree with its spec.
pub struct Controller {
    /// Receives the controller's status writes.
    pub server: Arc<ApiServer>,
    /// Supplies the current state of each object.
    pub cache: Arc<Cache>,
    /// Supplies the names of objects to reconcile.
    pub queue: Arc<WorkQueue>,
}

impl Controller {
    /// A single worker: takes names from the queue one at a time and
    /// reconciles them. It does not return.
    pub fn run(&self) {
        loop {
            let name = self.queue.get();
            let result = self.reconcile(&name);
            self.queue.done(&name);

            if result.is_err() {
                self.queue.add([name]); // client-go requeues with a backoff
            }
        }
    }
}

impl Controller {
    /// Reads the named object from the cache and, if its status differs from
    /// its spec, writes a new status conditioned on the version it read.
    fn reconcile(&self, name: &str) -> Result<(), ApiError> {
        let Some(mut object) = self.cache.get(name) else {
            log!("reconcile {name}: deleted, releasing its GPUs");

            return Ok(());
        };

        if object.status.gpus == object.spec.gpus {
            return Ok(());
        }

        thread::sleep(Duration::from_millis(20)); // slow work, such as starting containers
        object.status.gpus = object.spec.gpus;
        let version = object.resource_version;

        match self.server.update(object.clone()) {
            Ok(()) => {
                log!("reconcile {name}: status.gpus = {}", object.status.gpus);

                Ok(())
            }
            Err(error) => {
                log!("reconcile {name}: write at version {version}: {error}");

                Err(error)
            }
        }
    }
}

Listing 1.39 The demonstration, complete, in Go. code/go/informer/main.go

// Command informer models how a Kubernetes controller learns about objects:
// a server that numbers every write, a reflector that lists and watches into
// a local cache, a work queue of names, and a controller that reconciles them.
// It then causes a burst of edits, a dropped connection, and a partition, and
// logs how each part responds.
//
// Run it from code/go with:
//
//	go run ./informer
package main

import (
	"fmt"
	"log"
	"time"
)

// start is the time the program began, for the relative timestamps in the log.
var start = time.Now()

// logf prints one log line prefixed with the milliseconds since start.
func logf(format string, args ...any) {
	log.Printf("%4dms  %s", time.Since(start).Milliseconds(), fmt.Sprintf(format, args...))
}

func main() {
	log.SetFlags(0)
	server := NewAPIServer()
	server.Update(Object{Name: "job-a", Spec: Spec{GPUs: 1}})
	server.Update(Object{Name: "job-b", Spec: Spec{GPUs: 8}})

	cache := NewCache()
	queue := NewWorkQueue()
	go (&Reflector{server, cache, queue}).Run()
	go (&Controller{server, cache, queue}).Run()
	time.Sleep(100 * time.Millisecond)

	logf("--- job-a is edited five times in a burst")

	for gpus := 2; gpus <= 6; gpus++ {
		server.Update(Object{Name: "job-a", Spec: Spec{GPUs: gpus}})

		if gpus == 2 {
			time.Sleep(5 * time.Millisecond) // the controller starts on the first edit
		}
	}

	time.Sleep(150 * time.Millisecond)

	logf("--- the watch connection drops")
	server.DropWatches()
	time.Sleep(50 * time.Millisecond)

	logf("--- partition: job-b is deleted and the history compacted meanwhile")
	server.Partition()
	server.Delete("job-b")
	server.Compact()
	time.Sleep(30 * time.Millisecond)
	server.Heal()
	time.Sleep(150 * time.Millisecond)

	list, version := server.List()
	logf("final state at version %d: %v", version, list)
}

Listing 1.40 The demonstration, complete, in Rust. code/rust/informer/src/main.rs

//! A model of how a Kubernetes controller learns about objects: a server that
//! numbers every write, a reflector that lists and watches into a local cache,
//! a work queue of names, and a controller that reconciles them. It then
//! causes a burst of edits, a dropped connection, and a partition, and logs
//! how each part responds.
//!
//! Run it from `code/rust` with `cargo run -p informer`.

mod cache;
mod controller;
mod queue;
mod reflector;
mod server;

use cache::Cache;
use controller::Controller;
use queue::WorkQueue;
use reflector::Reflector;
use server::{ApiServer, Object};
use std::sync::{Arc, OnceLock};
use std::thread;
use std::time::{Duration, Instant};

/// The time the program began, for the relative timestamps in the log.
static START: OnceLock<Instant> = OnceLock::new();

/// Prints one log line to standard error, prefixed with the milliseconds
/// since the program began.
#[macro_export]
macro_rules! log {
    ($($arg:tt)*) => {
        eprintln!("{:4}ms  {}", $crate::elapsed_ms(), format!($($arg)*))
    };
}

/// Returns the milliseconds since the program began.
pub fn elapsed_ms() -> u128 {
    START.get_or_init(Instant::now).elapsed().as_millis()
}

/// Sleeps for the given number of milliseconds.
fn sleep_ms(ms: u64) {
    thread::sleep(Duration::from_millis(ms));
}

fn main() {
    START.get_or_init(Instant::now);
    let server = Arc::new(ApiServer::default());
    server
        .update(Object::new("job-a", 1))
        .expect("unconditional write");
    server
        .update(Object::new("job-b", 8))
        .expect("unconditional write");

    let cache = Arc::new(Cache::default());
    let queue = Arc::new(WorkQueue::default());
    let reflector = Reflector {
        server: Arc::clone(&server),
        cache: Arc::clone(&cache),
        queue: Arc::clone(&queue),
    };
    let controller = Controller {
        server: Arc::clone(&server),
        cache,
        queue,
    };
    thread::spawn(move || reflector.run());
    thread::spawn(move || controller.run());
    sleep_ms(100);

    log!("--- job-a is edited five times in a burst");

    for gpus in 2..=6 {
        server
            .update(Object::new("job-a", gpus))
            .expect("unconditional write");
        if gpus == 2 {
            sleep_ms(5); // the controller starts on the first edit
        }
    }

    sleep_ms(150);

    log!("--- the watch connection drops");
    server.drop_watches();
    sleep_ms(50);

    log!("--- partition: job-b is deleted and the history compacted meanwhile");
    server.partition();
    server.delete("job-b");
    server.compact();
    sleep_ms(30);
    server.heal();
    sleep_ms(150);

    let (list, version) = server.list();
    let objects: Vec<String> = list.iter().map(ToString::to_string).collect();
    log!("final state at version {version}: [{}]", objects.join(" "));
}

1.1 A GPU job controller on a real cluster

In chapter 1, we built a server, a cache, a reflector, a work queue, and a controller, and ran all of them inside one process. In this chapter, we will take that machinery to a real Kubernetes cluster and write a controller that runs GPU jobs on a real GPU. We will first teach the API server a new kind of object, the GPU job, and then watch those objects arrive the way our reflector did. Then we will write the controller, run it, and see a CUDA program run on a GPU inside a pod that our controller created. Finally we will make two clients collide on the same object and watch the API server’s version check stop the second one.

This time we will not write the machinery ourselves. We will use each language’s standard Kubernetes library, which provides the same parts we built by hand:

• In Go, we will use controller-runtime. It is built on client-go, the library whose parts we listed at the end of chapter 1, and it is what Kubebuilder, the usual tool for starting a Go controller project, generates code for.
• In Rust, we will use kube-rs, published as the kube crate.

We will introduce each part of these libraries at the moment we need it, and relate it to the piece we built in chapter 1.

Our cluster is created with kind, which runs each Kubernetes node as a Docker container. We talk to it with kubectl, the Kubernetes command-line client:

$ kind get clusters
gpu-lab
$ kubectl get nodes
NAME                    STATUS   ROLES           AGE   VERSION
gpu-lab-control-plane   Ready    control-plane   22s   v1.37.0

Pods in this cluster can use a GPU. Appendix A shows how we set that up.

The projects

Our Go code lives in the module in code/go, next to the model from chapter 1:

Listing 1.41 The Go module. code/go/go.mod

module gpuk8s

go 1.26.0

require (
	k8s.io/api v0.37.1
	k8s.io/apimachinery v0.37.1
	k8s.io/client-go v0.37.0
	sigs.k8s.io/controller-runtime v0.25.1
)

We depend on four modules. k8s.io/api holds the Go types of the built-in Kubernetes objects, such as Pod. k8s.io/apimachinery holds what all objects share: metadata, resource quantities such as “2 CPUs”, and the errors the API server returns. k8s.io/client-go is the client library, and sigs.k8s.io/controller-runtime is the controller framework built on it. We pick versions that match our cluster’s Kubernetes 1.37.

Our Rust code is the gpujob package in the Cargo workspace in code/rust, next to the model’s informer package:

Listing 1.42 The Cargo workspace. code/rust/Cargo.toml

[workspace]
members = ["informer", "gpujob", "gpuusage", "heartbeats", "phi", "leaseagent", "leasewatch", "fleethealth", "fencing", "provisioner", "migration", "allreduce", "gangsched", "counter", "cas", "gpuctl"]
resolver = "3"

Listing 1.43 The `gpujob` package. code/rust/gpujob/Cargo.toml

[package]
name = "gpujob"
version = "0.1.0"
edition = "2024"
publish = false
default-run = "controller"

[dependencies]
futures = "0.3"
k8s-openapi = { version = "0.28", features = ["v1_36"] }
kube = { version = "4.2", features = ["runtime", "derive"] }
schemars = "1"
serde = { version = "1", features = ["derive"] }
serde_json = "1"
thiserror = "2"
tokio = { version = "1", features = ["macros", "rt-multi-thread"] }
tracing = "0.1"
tracing-subscriber = { version = "0.3", features = ["env-filter"] }

kube is the client. We turn on its runtime feature for the controller machinery and its derive feature for generating our own object types. k8s-openapi provides the Rust types of the built-in objects; its v1_36 feature selects the newest Kubernetes version it supports. The other dependencies do supporting work: schemars generates JSON schemas, serde and serde_json convert objects to and from JSON, tokio runs our asynchronous code, futures gives us operations on streams, thiserror derives our error type, and tracing writes the controller’s log. Finally, default-run lets us start the controller with a plain cargo run -p gpujob.

Teaching the API server a new kind of object

In chapter 1 our jobs were records with a GPU count. A real cluster has no such kind of object: it knows about pods, services, deployments, and the other built-in kinds, and that is all. Fortunately, the API server can be taught new kinds. We describe the new kind in a CustomResourceDefinition (CRD): its API group and version, its name, and the schema of its fields. Once we install the CRD, the API server stores, validates, lists, and watches objects of our new kind exactly as it does its built-in kinds. That means everything from chapter 1 applies to our jobs too: they get resource versions, clients can list and watch them, and writes can be conditional.

We will call our kind GpuJob, in the API group gpucloud.dev, version v1. Its spec holds a container image, an optional command, and a GPU count. Its status holds a phase, a message, and the generation it describes. In both languages, we write the kind as ordinary types and let a tool generate the CRD from them.

In Go

A Go API package usually has two files. The first one names the group and version, and it registers our types in a scheme. A scheme is a table that maps Go types to the API kinds they represent. A client needs it to know that JSON with "kind": "GpuJob" should become a GpuJob value, and which URL to use when it sends one to the server.

Listing 1.44 The group, the version, and the scheme. code/go/gpujob/api/v1/groupversion.go

// Package v1 defines version v1 of the GpuJob custom resource in the API group
// gpucloud.dev: a request to run one container with a number of GPUs, and the
// status that the controller reports back.
//
// The comments on the types below become the field descriptions in the
// generated CustomResourceDefinition, which kubectl explain prints. The
// +kubebuilder comments are markers that controller-gen reads.
//
// +kubebuilder:object:generate=true
// +groupName=gpucloud.dev
package v1

import (
	"k8s.io/apimachinery/pkg/runtime/schema"
	"sigs.k8s.io/controller-runtime/pkg/scheme"
)

//go:generate controller-gen object crd paths=. output:crd:dir=../../config/crd

// GroupVersion is the API group and version of the types in this package.
var GroupVersion = schema.GroupVersion{Group: "gpucloud.dev", Version: "v1"}

// SchemeBuilder registers the types in this package with a runtime.Scheme.
var SchemeBuilder = &scheme.Builder{GroupVersion: GroupVersion}

// AddToScheme adds the types in this package to a runtime.Scheme, so that a
// client can encode and decode them.
var AddToScheme = SchemeBuilder.AddToScheme

The second file defines the types themselves, in three parts: the spec, the status, and the object that holds both.

Listing 1.45 The spec and the phases a job passes through. code/go/gpujob/api/v1/types.go

// GpuJobSpec is a request to run one container with a number of GPUs.
type GpuJobSpec struct {
	// Container image to run.
	Image string `json:"image"`
	// Command to run in the container. When empty, the image's own command runs.
	// +optional
	Command []string `json:"command,omitempty"`
	// Number of GPUs the container requests.
	// +kubebuilder:validation:Minimum=0
	GPUs int32 `json:"gpus"`
}

// Phase is the progress of a job, following the phase of its pod.
// +kubebuilder:validation:Enum=Pending;Running;Succeeded;Failed
type Phase string

// The phases a job moves through.
const (
	// PhasePending means the pod is waiting to be scheduled or for its containers to start.
	PhasePending Phase = "Pending"
	// PhaseRunning means the pod's container is running.
	PhaseRunning Phase = "Running"
	// PhaseSucceeded means the container exited with status 0.
	PhaseSucceeded Phase = "Succeeded"
	// PhaseFailed means the container exited with a non-zero status.
	PhaseFailed Phase = "Failed"
)

Listing 1.46 The status. code/go/gpujob/api/v1/types.go

// GpuJobStatus is the state of a job as the controller last observed it.
type GpuJobStatus struct {
	// Progress of the job's pod.
	Phase Phase `json:"phase"`
	// The scheduler's reason, when the pod cannot be placed on a node.
	// +optional
	Message string `json:"message,omitempty"`
	// The metadata.generation of the spec that this status describes.
	// +optional
	ObservedGeneration int64 `json:"observedGeneration,omitempty"`
}

Listing 1.47 The object, the list, and their registration. code/go/gpujob/api/v1/types.go

// GpuJob runs one container with a number of GPUs.
//
// +kubebuilder:object:root=true
// +kubebuilder:subresource:status
// +kubebuilder:resource:shortName=gj
// +kubebuilder:printcolumn:name="GPUs",type=integer,JSONPath=".spec.gpus"
// +kubebuilder:printcolumn:name="Phase",type=string,JSONPath=".status.phase"
type GpuJob struct {
	metav1.TypeMeta `json:",inline"`
	// Standard object metadata: name, namespace, labels, resource version.
	metav1.ObjectMeta `json:"metadata,omitempty"`
	// What the user asks for.
	Spec GpuJobSpec `json:"spec"`
	// What the controller has observed.
	// +optional
	Status GpuJobStatus `json:"status,omitempty"`
}

// GpuJobList is a list of GpuJob objects, as the API server returns it.
//
// +kubebuilder:object:root=true
type GpuJobList struct {
	metav1.TypeMeta `json:",inline"`
	// Standard list metadata, including the list's resource version.
	metav1.ListMeta `json:"metadata,omitempty"`
	// The jobs in the list.
	Items []GpuJob `json:"items"`
}

func init() {
	SchemeBuilder.Register(&GpuJob{}, &GpuJobList{})
}

The comments that begin with +kubebuilder: are markers. They are instructions for controller-gen, a code generator from the Kubebuilder project, which reads our source code and writes the CRD for us. Let us go through the ones we use:

• object:root=true marks GpuJob and GpuJobList as complete API objects, as opposed to structs that only appear inside other objects.
• subresource:status turns on the status subresource, which we explain below.
• resource:shortName=gj lets us type kubectl get gj instead of kubectl get gpujobs.
• The two printcolumn markers add the GPU count and the phase to the table that kubectl get prints.
• The validation markers become rules in the schema. The API server enforces them, so a job with a negative GPU count or an unknown phase is rejected before our controller ever sees it.

The ordinary comments on our types and fields matter as well: controller-gen copies them into the CRD as field descriptions, and kubectl explain gpujob.spec shows them to users.

Listing 1.48 Working out a status from a pod. code/go/gpujob/api/v1/types.go

// StatusFromPod derives a job's status from its pod. generation is the job's
// metadata.generation, recorded as the generation that the status describes.
func StatusFromPod(pod *corev1.Pod, generation int64) GpuJobStatus {
	status := GpuJobStatus{Phase: PhasePending, ObservedGeneration: generation}

	switch pod.Status.Phase {
	case corev1.PodRunning:
		status.Phase = PhaseRunning
	case corev1.PodSucceeded:
		status.Phase = PhaseSucceeded
	case corev1.PodFailed:
		status.Phase = PhaseFailed
	}

	// A pod that cannot be placed explains why in its PodScheduled condition.
	for _, c := range pod.Status.Conditions {
		if c.Type == corev1.PodScheduled && c.Status == corev1.ConditionFalse {
			status.Message = c.Message
		}
	}

	return status
}

StatusFromPod works out a job’s status from its pod. The pod’s phase maps directly onto our Phase. There is one extra piece of information worth passing on. When the scheduler cannot find a machine for a pod, it sets the pod’s PodScheduled condition to False and writes the reason in the condition’s message. StatusFromPod copies that message into the job’s status, so that a user who runs kubectl get gpujob can see why the job is still waiting.

The go:generate line in groupversion.go tells Go how to run controller-gen. It writes two files for us. The first is zz_generated.deepcopy.go, with a DeepCopy method for each type. Every Kubernetes API type needs these, because the shared caches described at the end of chapter 1 hand out pointers, and a controller must copy an object before changing it. The second file is the CRD itself, in gpujob/config/crd. Both generated files are checked into the repository, so we only need to run the generator again when we change the types:

$ go install sigs.k8s.io/controller-tools/cmd/controller-gen@v0.22.0
$ go generate ./gpujob/api/...
$ kubectl apply -f gpujob/config/crd/
customresourcedefinition.apiextensions.k8s.io/gpujobs.gpucloud.dev configured

In Rust

In Rust, there is no separate generator to run. kube-rs builds the CRD while our program compiles, from derive macros on our types:

Listing 1.49 The spec, with the attributes that describe the kind. code/rust/gpujob/src/lib.rs

/// A request to run one container with a number of GPUs.
///
/// The `CustomResource` derive generates the `GpuJob` type, which holds this
/// spec together with the object's metadata and an optional `GpuJobStatus`.
#[derive(CustomResource, Deserialize, Serialize, Clone, Debug, JsonSchema)]
#[kube(
    group = "gpucloud.dev",
    version = "v1",
    kind = "GpuJob",
    namespaced,
    status = "GpuJobStatus",
    shortname = "gj",
    printcolumn = r#"{"name":"GPUs","type":"integer","jsonPath":".spec.gpus"}"#,
    printcolumn = r#"{"name":"Phase","type":"string","jsonPath":".status.phase"}"#
)]
pub struct GpuJobSpec {
    /// Container image to run.
    pub image: String,
    /// Command to run in the container. When empty, the image's own command runs.
    #[serde(default)]
    pub command: Vec<String>,
    /// Number of GPUs the container requests.
    // Kubernetes schemas support int32 and int64, so the u32 is published as a
    // non-negative int32.
    #[schemars(with = "i32", range(min = 0))]
    pub gpus: u32,
}

Listing 1.50 The status and the phases. code/rust/gpujob/src/lib.rs

/// The state of a job as the controller last observed it.
#[derive(Deserialize, Serialize, Clone, Debug, Default, PartialEq, JsonSchema)]
#[serde(rename_all = "camelCase")]
pub struct GpuJobStatus {
    /// Progress of the job's pod.
    pub phase: Phase,
    /// The scheduler's reason, when the pod cannot be placed on a node.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub message: Option<String>,
    /// The `metadata.generation` of the spec that this status describes.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub observed_generation: Option<i64>,
}

/// Progress of a job, following the phase of its pod.
#[derive(Deserialize, Serialize, Clone, Copy, Debug, Default, PartialEq, Eq, JsonSchema)]
pub enum Phase {
    /// The pod is waiting to be scheduled or for its containers to start.
    #[default]
    Pending,
    /// The pod's container is running.
    Running,
    /// The container exited with status 0.
    Succeeded,
    /// The container exited with a non-zero status.
    Failed,
}

Listing 1.51 Working out a status from a pod. code/rust/gpujob/src/lib.rs

impl GpuJobStatus {
    /// Derives a job's status from its pod. `generation` is the job's
    /// `metadata.generation`, recorded as the generation this status describes.
    #[must_use]
    pub fn from_pod(pod: &Pod, generation: Option<i64>) -> Self {
        let status = pod.status.as_ref();
        let phase = match status.and_then(|s| s.phase.as_deref()) {
            Some("Running") => Phase::Running,
            Some("Succeeded") => Phase::Succeeded,
            Some("Failed") => Phase::Failed,
            _ => Phase::Pending,
        };
        // A pod that cannot be placed explains why in its PodScheduled condition.
        let message = status
            .and_then(|s| s.conditions.as_ref())
            .into_iter()
            .flatten()
            .find(|c| c.type_ == "PodScheduled" && c.status == "False")
            .and_then(|c| c.message.clone());
        Self {
            phase,
            message,
            observed_generation: generation,
        }
    }
}

The CustomResource derive generates a GpuJob type for us. It holds our spec together with the object’s metadata and an optional status, the same shape as the Go struct we wrote by hand. The #[kube(...)] attribute does the job of the Go markers. group and version place the kind at gpucloud.dev/v1. namespaced makes each job belong to a namespace, a named partition of the cluster’s objects, so that two teams can each have a job called train. status turns on the status subresource, and shortname and printcolumn do the same as their Go counterparts. The /// comments become the field descriptions, and our Phase enum becomes a string field whose allowed values are the variant names.

One attribute deserves a closer look. The schema that kube-rs generates for a u32 has the format uint32, which the API server does not recognise, and it complains with a warning every time the CRD is applied. The attribute #[schemars(with = "i32", range(min = 0))] publishes the field as a non-negative 32-bit integer instead, which is exactly the schema the Go field of type int32 with its Minimum=0 marker produces, while the Rust field itself stays unsigned. from_pod is the Rust version of StatusFromPod.

A tiny program prints the CRD, and we pipe it straight into kubectl:

Listing 1.52 Printing the CRD. code/rust/gpujob/src/bin/crdgen.rs

//! Prints the `CustomResourceDefinition` for `GpuJob` as JSON, ready for kubectl:
//!
//! ```console
//! cargo run -p gpujob --bin crdgen | kubectl apply -f -
//! ```

use gpujob::GpuJob;
use kube::CustomResourceExt;

/// Generates the definition from the `GpuJob` type and prints it.
fn main() -> serde_json::Result<()> {
    println!("{}", serde_json::to_string_pretty(&GpuJob::crd())?);

    Ok(())
}

$ cargo run -p gpujob --bin crdgen | kubectl apply -f -
customresourcedefinition.apiextensions.k8s.io/gpujobs.gpucloud.dev created

Our Go and Rust definitions describe the same kind with the same fields. We can install either one, and either controller can then serve the jobs.

The status subresource

Why did we turn on the status subresource? Think about two writers touching the same job: a user who edits the spec, and our controller, which writes the status. Without the subresource, both write the whole object, and whoever writes second can overwrite the other’s part with an older copy. With the subresource, the status gets its own endpoint, /status. The API server takes only status fields from writes to /status, and only spec fields from writes to the main endpoint, so the user and the controller can no longer step on each other.

The subresource has a second effect. The API server now increments metadata.generation only when the spec changes. That gives our status field observedGeneration its meaning: our controller copies the generation it acted on into it, and anyone comparing the two numbers can tell whether the latest spec change has been handled yet.

Watching the jobs

In chapter 1, our reflector listed the objects, watched from the list’s version, and listed again when the history ran out. Both libraries give us that loop ready-made. To see exactly what it delivers, we will write a small program in each language that prints every event it receives.

In Go, the loop comes packaged inside an informer, just like in client-go: controller-runtime’s cache package creates one informer per kind of object. To hear about changes, we register an event handler: a set of three functions that the informer calls whenever it adds, updates, or deletes an object in its cache. Our handlers simply print what they receive:

Listing 1.53 Describing an object for the log. code/go/gpujob/cmd/watch-jobs/main.go

// namespace is the namespace whose jobs the program watches.
const namespace = "default"

// describe formats a job as its name, GPU count, and resource version.
func describe(obj any) string {
	// A deletion that a relist discovered arrives wrapped in a tombstone.
	if tombstone, ok := obj.(toolscache.DeletedFinalStateUnknown); ok {
		obj = tombstone.Obj
	}

	job := obj.(*gpujobv1.GpuJob)

	return fmt.Sprintf("%s gpus=%d v%s", job.Name, job.Spec.GPUs, job.ResourceVersion)
}

Listing 1.54 Starting the informer and registering the handlers. code/go/gpujob/cmd/watch-jobs/main.go

// run does the work of main and returns the first error it meets.
func run() error {
	scheme := runtime.NewScheme()

	if err := gpujobv1.AddToScheme(scheme); err != nil {
		return err
	}

	informers, err := cache.New(ctrl.GetConfigOrDie(), cache.Options{
		Scheme:            scheme,
		DefaultNamespaces: map[string]cache.Config{namespace: {}},
	})

	if err != nil {
		return err
	}

	ctx := ctrl.SetupSignalHandler()
	informer, err := informers.GetInformer(ctx, &gpujobv1.GpuJob{})

	if err != nil {
		return err
	}

	registration, err := informer.AddEventHandler(toolscache.ResourceEventHandlerDetailedFuncs{
		AddFunc: func(obj any, isInInitialList bool) {
			if isInInitialList {
				fmt.Println("  listed  ", describe(obj))
			} else {
				fmt.Println("  changed ", describe(obj))
			}
		},
		UpdateFunc: func(_, obj any) { fmt.Println("  changed ", describe(obj)) },
		DeleteFunc: func(obj any) { fmt.Println("  deleted ", describe(obj)) },
	})

	if err != nil {
		return err
	}

	fmt.Println("list started")
	go informers.Start(ctx)

	if !toolscache.WaitForCacheSync(ctx.Done(), registration.HasSynced) {
		return fmt.Errorf("informer did not sync")
	}

	fmt.Println("list complete, watching")
	<-ctx.Done()

	return nil
}

The add handler gets an extra flag, isInInitialList, which is true for objects that came from the first list, so we can tell the objects that already existed from the ones created later. The handler’s registration has a HasSynced method that becomes true once our handler has seen the whole initial list, and WaitForCacheSync waits for that moment before we print “list complete”. The delete handler may receive a DeletedFinalStateUnknown value instead of a job. That is the marker from chapter 1 for a deletion that the informer only discovered by relisting, and describe unwraps it.

In Rust, kube-rs hands us the loop directly as a stream: a sequence of values that arrive over time, which we read one at a time with await. It is the asynchronous counterpart of an iterator. The function watcher returns a stream of Event values for one kind of object:

Listing 1.55 Printing every event, in Rust. code/rust/gpujob/src/bin/watch_jobs.rs

//! Prints the event stream that a watcher produces for `GpuJob` objects in the
//! current namespace: the initial list, one object at a time, and then one
//! event per change. Stop it with Ctrl-C.
//!
//! Run it from `code/rust` with `cargo run -p gpujob --bin watch_jobs`.

use futures::TryStreamExt;
use gpujob::GpuJob;
use kube::runtime::watcher::{self, Event, watcher};
use kube::{Api, Client, ResourceExt};

/// Formats a job as its name, GPU count, and resource version.
fn describe(job: &GpuJob) -> String {
    let version = job.resource_version().unwrap_or_default();

    format!("{} gpus={} v{version}", job.name_any(), job.spec.gpus)
}

/// Watches `GpuJob` objects and prints one line per event.
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let jobs: Api<GpuJob> = Api::default_namespaced(Client::try_default().await?);
    let mut events = std::pin::pin!(watcher(jobs, watcher::Config::default()));

    while let Some(event) = events.try_next().await? {
        match event {
            Event::Init => println!("list started"),
            Event::InitApply(job) => println!("  listed   {}", describe(&job)),
            Event::InitDone => println!("list complete, watching"),
            Event::Apply(job) => println!("  changed  {}", describe(&job)),
            Event::Delete(job) => println!("  deleted  {}", describe(&job)),
        }
    }

    Ok(())
}

The Event enum spells out the list-then-watch sequence we reasoned about in chapter 1. Init, then one InitApply for each object that existed at the time of the list, then InitDone, together make up a list. After that, each event is an Apply for an object that was created or changed, or a Delete. If the watch has to start over, for example after 410 Gone, the stream begins a new Init sequence by itself. A consumer that rebuilds its state at every Init therefore gets the same result as the Replace method of our chapter 1 cache. kube-rs provides exactly such a consumer and calls it a reflector: it applies the events to a store that the rest of the program can read.

We started each program before any jobs existed and left it running while we did everything else in this chapter. The Go program printed:

$ go run ./gpujob/cmd/watch-jobs
list started
list complete, watching
  changed  smoke-test gpus=0 v41900
  changed  train gpus=1 v41901
  changed  smoke-test gpus=0 v41904
  changed  train gpus=1 v41909
  changed  smoke-test gpus=0 v41920
  changed  train gpus=1 v41923
  changed  train gpus=1 v41931
  changed  smoke-test gpus=0 v41960
  changed  train gpus=2 v41972
  changed  train gpus=2 v41977
  changed  train gpus=2 v41979
  deleted  smoke-test gpus=0 v41992

And the Rust program printed:

$ cargo run -p gpujob --bin watch_jobs
list started
list complete, watching
  changed  smoke-test gpus=0 v656
  changed  train gpus=1 v657
  changed  train gpus=1 v661
  changed  smoke-test gpus=0 v662
  changed  train gpus=1 v676
  changed  smoke-test gpus=0 v684
  changed  train gpus=1 v688
  changed  smoke-test gpus=0 v722
  changed  train gpus=2 v731
  changed  train gpus=2 v737
  changed  train gpus=2 v739
  deleted  smoke-test gpus=0 v752

When we started them again later, with train still in the cluster, each one listed the job before it started watching. In Go:

$ go run ./gpujob/cmd/watch-jobs
list started
  listed   train gpus=2 v41979
list complete, watching

In Rust:

$ cargo run -p gpujob --bin watch_jobs
list started
  listed   train gpus=2 v739
list complete, watching

The resource versions are etcd’s revision numbers. Since etcd counts every write in the cluster, including those made by Kubernetes’ own components, the numbers of one job jump by more than one between its changes. The Go run happened later in the cluster’s life, so its numbers are larger. Most of the changed lines are our controller writing a new status: the object changes, even though its spec stays the same.

The controller

Now we can write the controller. It will run each GpuJob as a pod and report the pod’s progress in the job’s status. Both versions follow the same plan, so we walk through them together, one part at a time, and print the complete programs at the end of the chapter.

The pod for a job

Listing 1.56 The pod a job should have, in Go. code/go/gpujob/cmd/controller/main.go

// manager is the field manager name under which the controller writes with
// server-side apply.
const manager = "gpujob-controller"

// gpu is the extended resource through which a container requests GPUs.
const gpu corev1.ResourceName = "nvidia.com/gpu"

// desiredPod builds the pod that should exist for job, as an apply
// configuration: a partial object that lists only the fields this controller
// sets.
//
// The owner reference makes Kubernetes delete the pod together with the job,
// and it lets the controller map each pod event to the job that owns the pod.
func desiredPod(job *gpujobv1.GpuJob) *corev1ac.PodApplyConfiguration {
	gpus := corev1.ResourceList{gpu: *resource.NewQuantity(int64(job.Spec.GPUs), resource.DecimalSI)}
	container := corev1ac.Container().
		WithName("job").
		WithImage(job.Spec.Image).
		WithResources(corev1ac.ResourceRequirements().WithLimits(gpus))
	if len(job.Spec.Command) > 0 {
		container.WithCommand(job.Spec.Command...)
	}

	owner := metav1ac.OwnerReference().
		WithAPIVersion(gpujobv1.GroupVersion.String()).
		WithKind("GpuJob").
		WithName(job.Name).
		WithUID(job.UID).
		WithController(true).
		WithBlockOwnerDeletion(true)
	return corev1ac.Pod(job.Name, job.Namespace).
		WithOwnerReferences(owner).
		WithSpec(corev1ac.PodSpec().
			WithRestartPolicy(corev1.RestartPolicyNever).
			WithContainers(container))
}

// podGPUs returns the number of GPUs that an existing pod's container
// requests, read from its limits, and whether the pod has a GPU limit.
func podGPUs(pod *corev1.Pod) (int32, bool) {
	if len(pod.Spec.Containers) == 0 {
		return 0, false
	}

	quantity, ok := pod.Spec.Containers[0].Resources.Limits[gpu]

	return int32(quantity.Value()), ok
}

Listing 1.57 The pod a job should have, in Rust. code/rust/gpujob/src/bin/controller.rs

/// The field manager name under which the controller writes with server-side apply.
const MANAGER: &str = "gpujob-controller";

/// The extended resource through which a container requests GPUs.
const GPU: &str = "nvidia.com/gpu";

/// Everything a reconcile can fail with.
#[derive(Debug, thiserror::Error)]
enum Error {
    /// A request to the API server failed.
    #[error("kubernetes API: {0}")]
    Kube(#[from] kube::Error),
    /// The status could not be serialized into a patch.
    #[error("serializing status: {0}")]
    Json(#[from] serde_json::Error),
}

/// State shared by every reconcile.
struct Context {
    /// Client for the cluster's API server.
    client: Client,
}

/// Builds the pod that should exist for `job`.
///
/// The owner reference makes Kubernetes delete the pod together with the job,
/// and it lets the controller map each pod event to the job that owns the pod.
fn desired_pod(job: &GpuJob) -> Pod {
    let gpus = BTreeMap::from([(GPU.to_string(), Quantity(job.spec.gpus.to_string()))]);
    Pod {
        metadata: ObjectMeta {
            name: Some(job.name_any()),
            owner_references: job.controller_owner_ref(&()).map(|owner| vec![owner]),
            ..Default::default()
        },
        spec: Some(PodSpec {
            restart_policy: Some("Never".into()),
            containers: vec![Container {
                name: "job".into(),
                image: Some(job.spec.image.clone()),
                command: Some(job.spec.command.clone()).filter(|c| !c.is_empty()),
                resources: Some(ResourceRequirements {
                    limits: Some(gpus),
                    ..Default::default()
                }),
                ..Default::default()
            }],
            ..Default::default()
        }),
        ..Default::default()
    }
}

/// Returns the number of GPUs that an existing pod's container requests, read
/// from its limits, or `None` if the pod has no GPU limit.
fn pod_gpus(pod: &Pod) -> Option<u32> {
    let container = pod.spec.as_ref()?.containers.first()?;
    let limits = container.resources.as_ref()?.limits.as_ref()?;

    limits.get(GPU)?.0.parse().ok()
}

desiredPod (desired_pod in Rust) builds the pod that a job should have. The interesting part is how the pod asks for a GPU. It puts nvidia.com/gpu in its resource limits, just as it would put CPU or memory there. Kubernetes calls a resource like this an extended resource: a name that Kubernetes itself does not understand, but whose quantity a machine can advertise as a whole number, and which the scheduler counts like any other resource. The name nvidia.com/gpu is advertised by the NVIDIA device plugin, a small program on each GPU machine that counts the GPUs and reports them to Kubernetes; Appendix A shows it at work. Extended resources can only be requested in whole units, and they are requested through limits. One unit of nvidia.com/gpu is, by default, one whole physical GPU; chapter 2 examines what that includes and how a GPU can be divided.

The pod also carries an owner reference, a note in its metadata that says “this pod belongs to that job”. It has two uses. When the job is deleted, Kubernetes’ garbage collector, a controller built into every cluster, sees the note and deletes the pod as well, so we never have to clean up after deleted jobs ourselves. And because the note has its controller flag set, both libraries can take any event about the pod and work out which job it belongs to, which is how our controller hears about its pods.

The two languages build the pod differently. The Go version builds an apply configuration, a type from client-go in which every field is optional and is filled in through With methods. The value ends up containing exactly the fields our controller cares about and nothing else, which matters for the kind of write we are about to do. The Rust version builds an ordinary Pod value and leaves everything else at its default; serde leaves the empty fields out when it turns the pod into JSON, which has the same effect.

podGPUs (pod_gpus in Rust) reads the GPU count back out of an existing pod. In Go it returns the count and a flag saying whether the pod had a GPU limit at all. In Rust it returns an Option<u32>, and each ? bails out with None as soon as any field along the way is missing.

Reconciling a job

Listing 1.58 Reconciling one job, in Go. code/go/gpujob/cmd/controller/main.go

// GpuJobReconciler makes the cluster match each GpuJob.
type GpuJobReconciler struct {
	// Client reads objects from the manager's cache and writes to the API server.
	client.Client
}

// Reconcile makes the cluster match one job: it ensures the job's pod exists
// with the requested GPUs, and writes the pod's progress into the job's status.
func (r *GpuJobReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) {
	log := ctrl.Log.WithName("controller").WithValues("name", req.Name)

	var job gpujobv1.GpuJob

	if err := r.Get(ctx, req.NamespacedName, &job); err != nil {
		// A job deleted since it was queued needs nothing: its pod goes with it.
		return ctrl.Result{}, client.IgnoreNotFound(err)
	}

	// The resources of a pod are fixed when it is created. When the job asks
	// for a different number of GPUs, delete the pod and come back after it is gone.
	var pod corev1.Pod
	err := r.Get(ctx, req.NamespacedName, &pod)

	switch {
	case err == nil:
		if gpus, ok := podGPUs(&pod); !ok || gpus != job.Spec.GPUs {
			log.Info("GPU count changed, replacing pod")

			if err := r.Delete(ctx, &pod); client.IgnoreNotFound(err) != nil {
				return ctrl.Result{}, err
			}

			return ctrl.Result{RequeueAfter: 2 * time.Second}, nil
		}
	case !apierrors.IsNotFound(err):
		return ctrl.Result{}, err
	}

	if err := r.Apply(ctx, desiredPod(&job), client.FieldOwner(manager), client.ForceOwnership); err != nil {
		return ctrl.Result{}, err
	}

	// A pod that was just created is not in the cache yet; its zero value
	// yields the Pending phase, and the pod's creation event triggers the
	// reconcile that reports its real progress.
	status := gpujobv1.StatusFromPod(&pod, job.Generation)

	if job.Status != status {
		log.Info("status changed", "phase", status.Phase, "message", status.Message)
		patch := client.MergeFrom(job.DeepCopy())
		job.Status = status

		if err := r.Status().Patch(ctx, &job, patch); err != nil {
			return ctrl.Result{}, err
		}
	}

	// Changes to the pod trigger the next reconcile through Owns.
	return ctrl.Result{}, nil
}

Listing 1.59 Reconciling one job, in Rust. code/rust/gpujob/src/bin/controller.rs

/// Makes the cluster match one job: ensures its pod exists with the requested
/// GPUs, and writes the pod's progress into the job's status.
async fn reconcile(job: Arc<GpuJob>, ctx: Arc<Context>) -> Result<Action, Error> {
    let namespace = job.namespace().unwrap_or_default();
    let name = job.name_any();
    let pods: Api<Pod> = Api::namespaced(ctx.client.clone(), &namespace);
    let jobs: Api<GpuJob> = Api::namespaced(ctx.client.clone(), &namespace);

    // The resources of a pod are fixed when it is created. When the job asks
    // for a different number of GPUs, delete the pod and come back after it is gone.
    if let Some(pod) = pods.get_opt(&name).await?
        && pod_gpus(&pod) != Some(job.spec.gpus)
    {
        tracing::info!(%name, "GPU count changed, replacing pod");
        pods.delete(&name, &DeleteParams::default()).await?;

        return Ok(Action::requeue(Duration::from_secs(2)));
    }

    let apply = PatchParams::apply(MANAGER).force();
    let pod = pods
        .patch(&name, &apply, &Patch::Apply(desired_pod(&job)))
        .await?;

    let status = GpuJobStatus::from_pod(&pod, job.metadata.generation);

    if job.status.as_ref() != Some(&status) {
        let message = status.message.as_deref().unwrap_or("");
        tracing::info!(%name, phase = ?status.phase, message, "status changed");
        let patch = serde_json::json!({
            "apiVersion": GpuJob::api_version(&()),
            "kind": GpuJob::kind(&()),
            "status": status,
        });
        jobs.patch_status(&name, &apply, &Patch::Apply(patch))
            .await?;
    }

    // Changes to the pod trigger the next reconcile through `owns`.
    Ok(Action::await_change())
}

/// Decides what happens after a failed reconcile: log it and retry in five seconds.
#[allow(clippy::needless_pass_by_value)] // the signature is fixed by `Controller::run`
fn error_policy(job: Arc<GpuJob>, error: &Error, _ctx: Arc<Context>) -> Action {
    tracing::warn!(name = %job.name_any(), %error, "reconcile failed");

    Action::requeue(Duration::from_secs(5))
}

Our reconcile function receives a job and brings the cluster into line with it in three steps.

The first step covers a change that cannot be made in place. Kubernetes fixes a pod’s resources when the pod is created, so if the user changes the job’s GPU count, we cannot simply edit the pod. Instead we delete the pod and ask to be called again in two seconds, by which time it should be gone. The Go version asks for that with ctrl.Result{RequeueAfter: 2 * time.Second}, and the Rust version with Action::requeue. The Rust condition uses a let chain, if let ... && ..., available since the 2024 edition, to check for the pod and compare its GPU count in one expression.

The second step creates or updates the pod with a single call, using server-side apply. We send the pod we want, and the API server merges it into whatever it has stored. The server also records which fields each writer set; a writer in this sense is called a field manager, and our controller calls itself gpujob-controller. Applying the same object twice changes nothing, which makes server-side apply a natural fit for a level-triggered controller: we can apply on every single reconcile, without first checking whether the pod exists. We also force ownership, with client.ForceOwnership in Go and force() in Rust. If some other writer has set a field that we also set, we take that field over, which is what we want from a controller that is the only authority over its pods.

The third step works out the job’s status and writes it through the status subresource, but only if it differs from the status already stored. Skipping unchanged writes is more than a saving. Every write produces a watch event, every event causes a reconcile, and a controller that wrote on every reconcile would keep waking itself up forever. The Go version writes the status with a merge patch, which sends only the fields that differ from a copy of the job we saved before changing it. The Rust version applies the status with server-side apply, just as it does the pod.

The two versions also read the pod differently, and the difference is worth understanding. controller-runtime’s client reads from the informer caches that the manager keeps, so the Go Get of the pod is a lookup in memory, exactly like the cache reads in chapter 1. A pod we have only just created is not in the cache yet; in that case the Go version computes the status from an empty pod, which comes out as Pending, and the pod’s creation event soon triggers another reconcile that reports the real state. kube-rs’s Api, on the other hand, talks to the API server directly, so the Rust get_opt is a real GET request. That is fine for our small controller. A Rust controller that looks after thousands of pods would keep a reflector store of pods in its Context and read them from memory, just as the Go version does.

When a reconcile fails, the Go version returns the error, and controller-runtime puts the job back in the queue with the growing delay we described at the end of chapter 1. The Rust version returns the error to kube-rs, which asks our error_policy function what to do; ours logs the error and tries again in five seconds.

Starting the controller

Listing 1.60 Starting the controller, in Go. code/go/gpujob/cmd/controller/main.go

// main builds a manager, registers the reconciler, and runs until Ctrl-C.
func main() {
	ctrl.SetLogger(zap.New(zap.UseDevMode(true))) // readable console output
	setup := ctrl.Log.WithName("setup")

	scheme := runtime.NewScheme()

	if err := clientgoscheme.AddToScheme(scheme); err != nil {
		setup.Error(err, "registering built-in types")
		os.Exit(1)
	}

	if err := gpujobv1.AddToScheme(scheme); err != nil {
		setup.Error(err, "registering GpuJob")
		os.Exit(1)
	}

	mgr, err := ctrl.NewManager(ctrl.GetConfigOrDie(), ctrl.Options{
		Scheme:  scheme,
		Metrics: metricsserver.Options{BindAddress: "0"}, // no metrics endpoint
	})

	if err != nil {
		setup.Error(err, "creating manager")
		os.Exit(1)
	}

	err = ctrl.NewControllerManagedBy(mgr).
		For(&gpujobv1.GpuJob{}).
		Owns(&corev1.Pod{}).
		Complete(&GpuJobReconciler{Client: mgr.GetClient()})
	if err != nil {
		setup.Error(err, "creating controller")
		os.Exit(1)
	}

	if err := mgr.Start(ctrl.SetupSignalHandler()); err != nil {
		setup.Error(err, "running manager")
		os.Exit(1)
	}
}

Listing 1.61 Starting the controller, in Rust. code/rust/gpujob/src/bin/controller.rs

/// Starts the controller and runs it until Ctrl-C.
#[tokio::main]
async fn main() -> Result<(), kube::Error> {
    tracing_subscriber::fmt()
        .with_env_filter("controller=info,kube=warn")
        .with_target(false)
        .with_ansi(std::io::IsTerminal::is_terminal(&std::io::stdout()))
        .init();
    let client = Client::try_default().await?;
    let jobs: Api<GpuJob> = Api::all(client.clone());
    let pods: Api<Pod> = Api::all(client.clone());

    Controller::new(jobs, watcher::Config::default())
        .owns(pods, watcher::Config::default())
        .shutdown_on_signal()
        .run(reconcile, error_policy, Arc::new(Context { client }))
        .for_each(|result| async move {
            if let Err(error) = result {
                tracing::debug!(%error, "reconcile error");
            }
        })
        .await;
    Ok(())
}

In Go, everything a controller process shares lives in a manager: the connection to the cluster, the scheme, the informer caches, and the client that reads from them. ctrl.NewControllerManagedBy(mgr) builds our controller inside the manager, and each call in the chain wires up one piece. For(&GpuJob{}) watches jobs and queues the name of each job that changes. Owns(&corev1.Pod{}) watches pods and queues the name of each pod’s owning job, using the owner reference we set. Complete registers our reconciler. Then mgr.Start starts the caches, waits for their first lists, and runs the workers until the process receives a signal such as Ctrl-C. We switch off the metrics endpoint that a manager starts by default, since we do not need it here.

In Rust, Controller::new starts a watcher and a store for GpuJob objects, and owns adds a second watcher, for pods, that queues each pod’s owning job in the same way. run starts the work queue and the workers and gives us back a stream with one item per finished reconcile, and for_each keeps reading that stream until the program stops. shutdown_on_signal ends the stream when we press Ctrl-C.

In both libraries, the work queue behaves exactly like the one we built in chapter 1: a job waits in it at most once, and a job that one worker is reconciling is held back from all the others until that worker is done.

Running it

Figure 1.2 shows what one reconcile writes for a job named train, and what Kubernetes does with the pod.

We will test the controller with two jobs. smoke-test asks for no GPUs and runs a short shell command, which tells us the basic machinery works. train asks for one GPU and runs gpu-check, a small CUDA program from Appendix A that adds two vectors on the GPU and prints which device it used.

Listing 1.62 The two jobs. code/manifests/gpujobs.yaml

apiVersion: gpucloud.dev/v1
kind: GpuJob
metadata:
  name: smoke-test
spec:
  image: busybox:1.37
  command: ["sh", "-c", "echo hello from the job; sleep 20"]
  gpus: 0
---
apiVersion: gpucloud.dev/v1
kind: GpuJob
metadata:
  name: train
spec:
  image: gpu-check:1
  gpus: 1

We start the Go controller from code/go, create the jobs, and check on them:

$ go run ./gpujob/cmd/controller
$ kubectl apply -f ../manifests/gpujobs.yaml
$ kubectl get gpujobs
NAME         GPUS   PHASE
smoke-test   0      Succeeded
train        1      Succeeded
$ kubectl logs train
device: Orin, compute 8.7, 8 SMs
kernel: no error, out[0] = 3.0, out[n-1] = 3.0

The log of train shows our CUDA kernel running on the cluster’s GPU, inside the pod our controller created. After its startup messages, the Go controller logged the following lines, shown here with the date and time zone trimmed from the timestamps:

15:12:37  INFO  controller  status changed  {"name": "smoke-test", "phase": "Pending", "message": ""}
15:12:37  INFO  controller  status changed  {"name": "train", "phase": "Pending", "message": ""}
15:12:38  INFO  controller  status changed  {"name": "smoke-test", "phase": "Running", "message": ""}
15:12:38  INFO  controller  status changed  {"name": "train", "phase": "Running", "message": ""}
15:12:40  INFO  controller  status changed  {"name": "train", "phase": "Succeeded", "message": ""}
15:13:00  INFO  controller  status changed  {"name": "smoke-test", "phase": "Succeeded", "message": ""}

The Rust controller runs from code/rust, and the same steps give the same result:

$ cargo run -p gpujob
$ kubectl apply -f ../manifests/gpujobs.yaml
$ kubectl get gpujobs
NAME         GPUS   PHASE
smoke-test   0      Succeeded
train        1      Succeeded

14:01:15.949  INFO status changed name=train phase=Pending
14:01:15.954  INFO status changed name=smoke-test phase=Pending
14:01:15.985  INFO status changed name=smoke-test phase=Pending
14:01:17.124  INFO status changed name=train phase=Running
14:01:18.504  INFO status changed name=smoke-test phase=Running
14:01:18.959  INFO status changed name=train phase=Succeeded
14:01:39.957  INFO status changed name=smoke-test phase=Succeeded

In both runs, each job moved from Pending to Running to Succeeded along with its pod, and every step reached our controller through the pod watch that Owns or owns set up. train finished about two seconds after its pod was created, and smoke-test finished after its twenty-second sleep.

Look at the Rust log once more: smoke-test reported Pending twice, 31 milliseconds apart. Here is what happened. The first reconcile wrote the status. The pod’s creation then triggered a second reconcile before our own status write had come back through the watch, so the job in the cache still had no status, the comparison saw a difference, and the controller wrote the same status again. This is the read-your-own-writes gap from the end of chapter 1. Here it costs one redundant write, which server-side apply turns into a no-op.

Two writers, one object

Our controllers never send a resource version when they write, because server-side apply does not need one. Many clients do, though: a client that reads an object, changes it, and sends the whole object back with an update includes the version it read, and the API server applies the same check as our Update in chapter 1. To see that check on a real server, we will write a program in which two clients read the train job and then both try to write it back.

In Go:

Listing 1.63 Two clients and one refused write, in Go. code/go/gpujob/cmd/conflict/main.go

// run does the work of main and returns the first unexpected error.
func run(ctx context.Context) error {
	scheme := runtime.NewScheme()

	if err := gpujobv1.AddToScheme(scheme); err != nil {
		return err
	}

	c, err := client.New(ctrl.GetConfigOrDie(), client.Options{Scheme: scheme})

	if err != nil {
		return err
	}

	var first gpujobv1.GpuJob

	if err := c.Get(ctx, job, &first); err != nil {
		return err
	}

	second := first.DeepCopy()
	fmt.Printf("both clients read %s at v%s\n", job.Name, first.ResourceVersion)

	first.Spec.GPUs++

	if err := c.Update(ctx, &first); err != nil {
		return err
	}

	fmt.Printf("first write stored, now v%s\n", first.ResourceVersion)

	second.Spec.GPUs += 2
	err = c.Update(ctx, second)
	var status *apierrors.StatusError

	if errors.As(err, &status) && status.ErrStatus.Code == http.StatusConflict {
		fmt.Printf("second write refused: %d %s\n", status.ErrStatus.Code, status.ErrStatus.Message)

		return nil
	}

	return fmt.Errorf("unexpected result: %v", err)
}

In Rust:

Listing 1.64 Two clients and one refused write, in Rust. code/rust/gpujob/src/bin/conflict.rs

//! Shows the API server refusing a write with an outdated resource version.
//! Two clients read the `train` job, and both write it back with `replace`.
//! The first write succeeds and changes the version, so the second fails with
//! `409 Conflict`.
//!
//! Run it from `code/rust`, while a `GpuJob` named `train` exists, with
//! `cargo run -p gpujob --bin conflict`.

use gpujob::GpuJob;
use kube::api::PostParams;
use kube::{Api, Client, ResourceExt};

/// The job that both clients read and write.
const JOB: &str = "train";

/// Performs the two writes and reports how the server answered each.
#[tokio::main]
async fn main() -> Result<(), kube::Error> {
    let jobs: Api<GpuJob> = Api::default_namespaced(Client::try_default().await?);
    let mut first = jobs.get(JOB).await?;
    let mut second = first.clone();
    println!(
        "both clients read {JOB} at v{}",
        first.resource_version().unwrap_or_default()
    );

    first.spec.gpus += 1;
    let stored = jobs.replace(JOB, &PostParams::default(), &first).await?;
    println!(
        "first write stored, now v{}",
        stored.resource_version().unwrap_or_default()
    );

    second.spec.gpus += 2;

    match jobs.replace(JOB, &PostParams::default(), &second).await {
        Err(kube::Error::Api(status)) if status.code == 409 => {
            println!("second write refused: {} {}", status.code, status.message);
        }
        other => println!("unexpected: {other:?}"),
    }

    Ok(())
}

The Go version recognises the conflict with errors.As, which digs the API server’s status out of the returned error, and compares its code with http.StatusConflict. The Rust version matches the error against a pattern with a guard, if status.code == 409, and every other outcome falls through to the second arm.

$ go run ./gpujob/cmd/conflict
both clients read train at v41931
first write stored, now v41972
second write refused: 409 Operation cannot be fulfilled on gpujobs.gpucloud.dev "train":
the object has been modified; please apply your changes to the latest version and try again

$ cargo run -p gpujob --bin conflict
both clients read train at v688
first write stored, now v731
second write refused: 409 Operation cannot be fulfilled on gpujobs.gpucloud.dev "train":
the object has been modified; please apply your changes to the latest version and try again

The first write raised train to two GPUs, and our controllers noticed. The Go controller logged:

15:13:07  INFO  controller  GPU count changed, replacing pod  {"name": "train"}
15:13:07  INFO  controller  status changed  {"name": "train", "phase": "Pending", "message": ""}
15:13:07  INFO  controller  status changed  {"name": "train", "phase": "Pending", "message":
          "0/1 nodes are available: 1 Insufficient nvidia.com/gpu. preemption: 0/1 nodes are
          available: 1 Preemption is not helpful for scheduling."}

And the Rust controller:

14:01:46.101  INFO GPU count changed, replacing pod name=train
14:01:46.137  INFO status changed name=train phase=Pending
14:01:46.170  INFO status changed name=train phase=Pending 0/1 nodes are available:
                   1 Insufficient nvidia.com/gpu. preemption: 0/1 nodes are available:
                   1 Preemption is not helpful for scheduling.

Each controller deleted the old pod, because a pod’s resources cannot change, and the deletion event triggered the reconcile that created the replacement. The replacement asks for two GPUs on a machine that has only one, so the scheduler cannot place it anywhere. Thanks to the message we copy from the pod’s PodScheduled condition, the job’s status tells the user exactly why, and kubectl get gpujobs shows the job as Pending until its request is lowered or a machine with two free GPUs joins the cluster.

The complete programs

Listing 1.65 The API types, complete, in Go. code/go/gpujob/api/v1/types.go

package v1

import (
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)


// GpuJobSpec is a request to run one container with a number of GPUs.
type GpuJobSpec struct {
	// Container image to run.
	Image string `json:"image"`
	// Command to run in the container. When empty, the image's own command runs.
	// +optional
	Command []string `json:"command,omitempty"`
	// Number of GPUs the container requests.
	// +kubebuilder:validation:Minimum=0
	GPUs int32 `json:"gpus"`
}

// Phase is the progress of a job, following the phase of its pod.
// +kubebuilder:validation:Enum=Pending;Running;Succeeded;Failed
type Phase string

// The phases a job moves through.
const (
	// PhasePending means the pod is waiting to be scheduled or for its containers to start.
	PhasePending Phase = "Pending"
	// PhaseRunning means the pod's container is running.
	PhaseRunning Phase = "Running"
	// PhaseSucceeded means the container exited with status 0.
	PhaseSucceeded Phase = "Succeeded"
	// PhaseFailed means the container exited with a non-zero status.
	PhaseFailed Phase = "Failed"
)



// GpuJobStatus is the state of a job as the controller last observed it.
type GpuJobStatus struct {
	// Progress of the job's pod.
	Phase Phase `json:"phase"`
	// The scheduler's reason, when the pod cannot be placed on a node.
	// +optional
	Message string `json:"message,omitempty"`
	// The metadata.generation of the spec that this status describes.
	// +optional
	ObservedGeneration int64 `json:"observedGeneration,omitempty"`
}



// GpuJob runs one container with a number of GPUs.
//
// +kubebuilder:object:root=true
// +kubebuilder:subresource:status
// +kubebuilder:resource:shortName=gj
// +kubebuilder:printcolumn:name="GPUs",type=integer,JSONPath=".spec.gpus"
// +kubebuilder:printcolumn:name="Phase",type=string,JSONPath=".status.phase"
type GpuJob struct {
	metav1.TypeMeta `json:",inline"`
	// Standard object metadata: name, namespace, labels, resource version.
	metav1.ObjectMeta `json:"metadata,omitempty"`
	// What the user asks for.
	Spec GpuJobSpec `json:"spec"`
	// What the controller has observed.
	// +optional
	Status GpuJobStatus `json:"status,omitempty"`
}

// GpuJobList is a list of GpuJob objects, as the API server returns it.
//
// +kubebuilder:object:root=true
type GpuJobList struct {
	metav1.TypeMeta `json:",inline"`
	// Standard list metadata, including the list's resource version.
	metav1.ListMeta `json:"metadata,omitempty"`
	// The jobs in the list.
	Items []GpuJob `json:"items"`
}

func init() {
	SchemeBuilder.Register(&GpuJob{}, &GpuJobList{})
}



// StatusFromPod derives a job's status from its pod. generation is the job's
// metadata.generation, recorded as the generation that the status describes.
func StatusFromPod(pod *corev1.Pod, generation int64) GpuJobStatus {
	status := GpuJobStatus{Phase: PhasePending, ObservedGeneration: generation}

	switch pod.Status.Phase {
	case corev1.PodRunning:
		status.Phase = PhaseRunning
	case corev1.PodSucceeded:
		status.Phase = PhaseSucceeded
	case corev1.PodFailed:
		status.Phase = PhaseFailed
	}

	// A pod that cannot be placed explains why in its PodScheduled condition.
	for _, c := range pod.Status.Conditions {
		if c.Type == corev1.PodScheduled && c.Status == corev1.ConditionFalse {
			status.Message = c.Message
		}
	}

	return status
}

Listing 1.66 The API types, complete, in Rust. code/rust/gpujob/src/lib.rs

//! The `GpuJob` custom resource: a request to run one container with a number
//! of GPUs, and the status that the controller reports back.
//!
//! The `///` comments on the spec and status types become the field
//! descriptions in the generated `CustomResourceDefinition`, which
//! `kubectl explain gpujob.spec` prints.

use k8s_openapi::api::core::v1::Pod;
use kube::CustomResource;
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};

/// A request to run one container with a number of GPUs.
///
/// The `CustomResource` derive generates the `GpuJob` type, which holds this
/// spec together with the object's metadata and an optional `GpuJobStatus`.
#[derive(CustomResource, Deserialize, Serialize, Clone, Debug, JsonSchema)]
#[kube(
    group = "gpucloud.dev",
    version = "v1",
    kind = "GpuJob",
    namespaced,
    status = "GpuJobStatus",
    shortname = "gj",
    printcolumn = r#"{"name":"GPUs","type":"integer","jsonPath":".spec.gpus"}"#,
    printcolumn = r#"{"name":"Phase","type":"string","jsonPath":".status.phase"}"#
)]
pub struct GpuJobSpec {
    /// Container image to run.
    pub image: String,
    /// Command to run in the container. When empty, the image's own command runs.
    #[serde(default)]
    pub command: Vec<String>,
    /// Number of GPUs the container requests.
    // Kubernetes schemas support int32 and int64, so the u32 is published as a
    // non-negative int32.
    #[schemars(with = "i32", range(min = 0))]
    pub gpus: u32,
}

/// The state of a job as the controller last observed it.
#[derive(Deserialize, Serialize, Clone, Debug, Default, PartialEq, JsonSchema)]
#[serde(rename_all = "camelCase")]
pub struct GpuJobStatus {
    /// Progress of the job's pod.
    pub phase: Phase,
    /// The scheduler's reason, when the pod cannot be placed on a node.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub message: Option<String>,
    /// The `metadata.generation` of the spec that this status describes.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub observed_generation: Option<i64>,
}

/// Progress of a job, following the phase of its pod.
#[derive(Deserialize, Serialize, Clone, Copy, Debug, Default, PartialEq, Eq, JsonSchema)]
pub enum Phase {
    /// The pod is waiting to be scheduled or for its containers to start.
    #[default]
    Pending,
    /// The pod's container is running.
    Running,
    /// The container exited with status 0.
    Succeeded,
    /// The container exited with a non-zero status.
    Failed,
}

impl GpuJobStatus {
    /// Derives a job's status from its pod. `generation` is the job's
    /// `metadata.generation`, recorded as the generation this status describes.
    #[must_use]
    pub fn from_pod(pod: &Pod, generation: Option<i64>) -> Self {
        let status = pod.status.as_ref();
        let phase = match status.and_then(|s| s.phase.as_deref()) {
            Some("Running") => Phase::Running,
            Some("Succeeded") => Phase::Succeeded,
            Some("Failed") => Phase::Failed,
            _ => Phase::Pending,
        };
        // A pod that cannot be placed explains why in its PodScheduled condition.
        let message = status
            .and_then(|s| s.conditions.as_ref())
            .into_iter()
            .flatten()
            .find(|c| c.type_ == "PodScheduled" && c.status == "False")
            .and_then(|c| c.message.clone());
        Self {
            phase,
            message,
            observed_generation: generation,
        }
    }
}

Listing 1.67 The event printer, complete, in Go. code/go/gpujob/cmd/watch-jobs/main.go

// Command watch-jobs prints the events that an informer delivers for GpuJob
// objects in the default namespace: the initial list, one object at a time,
// and then one event per change. Stop it with Ctrl-C.
//
// Run it from code/go with:
//
//	go run ./gpujob/cmd/watch-jobs
package main

import (
	"fmt"
	"os"

	"k8s.io/apimachinery/pkg/runtime"
	toolscache "k8s.io/client-go/tools/cache"
	ctrl "sigs.k8s.io/controller-runtime"
	"sigs.k8s.io/controller-runtime/pkg/cache"

	gpujobv1 "gpuk8s/gpujob/api/v1"
)

// namespace is the namespace whose jobs the program watches.
const namespace = "default"

// describe formats a job as its name, GPU count, and resource version.
func describe(obj any) string {
	// A deletion that a relist discovered arrives wrapped in a tombstone.
	if tombstone, ok := obj.(toolscache.DeletedFinalStateUnknown); ok {
		obj = tombstone.Obj
	}

	job := obj.(*gpujobv1.GpuJob)

	return fmt.Sprintf("%s gpus=%d v%s", job.Name, job.Spec.GPUs, job.ResourceVersion)
}


// main starts an informer for GpuJob objects and prints one line per event.
func main() {
	if err := run(); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}

// run does the work of main and returns the first error it meets.
func run() error {
	scheme := runtime.NewScheme()

	if err := gpujobv1.AddToScheme(scheme); err != nil {
		return err
	}

	informers, err := cache.New(ctrl.GetConfigOrDie(), cache.Options{
		Scheme:            scheme,
		DefaultNamespaces: map[string]cache.Config{namespace: {}},
	})

	if err != nil {
		return err
	}

	ctx := ctrl.SetupSignalHandler()
	informer, err := informers.GetInformer(ctx, &gpujobv1.GpuJob{})

	if err != nil {
		return err
	}

	registration, err := informer.AddEventHandler(toolscache.ResourceEventHandlerDetailedFuncs{
		AddFunc: func(obj any, isInInitialList bool) {
			if isInInitialList {
				fmt.Println("  listed  ", describe(obj))
			} else {
				fmt.Println("  changed ", describe(obj))
			}
		},
		UpdateFunc: func(_, obj any) { fmt.Println("  changed ", describe(obj)) },
		DeleteFunc: func(obj any) { fmt.Println("  deleted ", describe(obj)) },
	})

	if err != nil {
		return err
	}

	fmt.Println("list started")
	go informers.Start(ctx)

	if !toolscache.WaitForCacheSync(ctx.Done(), registration.HasSynced) {
		return fmt.Errorf("informer did not sync")
	}

	fmt.Println("list complete, watching")
	<-ctx.Done()

	return nil
}

Listing 1.68 The controller, complete, in Go. code/go/gpujob/cmd/controller/main.go

// Command controller is a controller for GpuJob objects. It runs each job as a
// pod that requests the job's GPUs, and reports the pod's progress in the
// job's status.
//
// Run it from code/go, against the cluster in the current kubeconfig context,
// with:
//
//	go run ./gpujob/cmd/controller
package main

import (
	"context"
	"os"
	"time"

	corev1 "k8s.io/api/core/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/apimachinery/pkg/api/resource"
	"k8s.io/apimachinery/pkg/runtime"
	corev1ac "k8s.io/client-go/applyconfigurations/core/v1"
	metav1ac "k8s.io/client-go/applyconfigurations/meta/v1"
	clientgoscheme "k8s.io/client-go/kubernetes/scheme"
	ctrl "sigs.k8s.io/controller-runtime"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/log/zap"
	metricsserver "sigs.k8s.io/controller-runtime/pkg/metrics/server"

	gpujobv1 "gpuk8s/gpujob/api/v1"
)

// manager is the field manager name under which the controller writes with
// server-side apply.
const manager = "gpujob-controller"

// gpu is the extended resource through which a container requests GPUs.
const gpu corev1.ResourceName = "nvidia.com/gpu"

// desiredPod builds the pod that should exist for job, as an apply
// configuration: a partial object that lists only the fields this controller
// sets.
//
// The owner reference makes Kubernetes delete the pod together with the job,
// and it lets the controller map each pod event to the job that owns the pod.
func desiredPod(job *gpujobv1.GpuJob) *corev1ac.PodApplyConfiguration {
	gpus := corev1.ResourceList{gpu: *resource.NewQuantity(int64(job.Spec.GPUs), resource.DecimalSI)}
	container := corev1ac.Container().
		WithName("job").
		WithImage(job.Spec.Image).
		WithResources(corev1ac.ResourceRequirements().WithLimits(gpus))
	if len(job.Spec.Command) > 0 {
		container.WithCommand(job.Spec.Command...)
	}

	owner := metav1ac.OwnerReference().
		WithAPIVersion(gpujobv1.GroupVersion.String()).
		WithKind("GpuJob").
		WithName(job.Name).
		WithUID(job.UID).
		WithController(true).
		WithBlockOwnerDeletion(true)
	return corev1ac.Pod(job.Name, job.Namespace).
		WithOwnerReferences(owner).
		WithSpec(corev1ac.PodSpec().
			WithRestartPolicy(corev1.RestartPolicyNever).
			WithContainers(container))
}

// podGPUs returns the number of GPUs that an existing pod's container
// requests, read from its limits, and whether the pod has a GPU limit.
func podGPUs(pod *corev1.Pod) (int32, bool) {
	if len(pod.Spec.Containers) == 0 {
		return 0, false
	}

	quantity, ok := pod.Spec.Containers[0].Resources.Limits[gpu]

	return int32(quantity.Value()), ok
}


// GpuJobReconciler makes the cluster match each GpuJob.
type GpuJobReconciler struct {
	// Client reads objects from the manager's cache and writes to the API server.
	client.Client
}

// Reconcile makes the cluster match one job: it ensures the job's pod exists
// with the requested GPUs, and writes the pod's progress into the job's status.
func (r *GpuJobReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) {
	log := ctrl.Log.WithName("controller").WithValues("name", req.Name)

	var job gpujobv1.GpuJob

	if err := r.Get(ctx, req.NamespacedName, &job); err != nil {
		// A job deleted since it was queued needs nothing: its pod goes with it.
		return ctrl.Result{}, client.IgnoreNotFound(err)
	}

	// The resources of a pod are fixed when it is created. When the job asks
	// for a different number of GPUs, delete the pod and come back after it is gone.
	var pod corev1.Pod
	err := r.Get(ctx, req.NamespacedName, &pod)

	switch {
	case err == nil:
		if gpus, ok := podGPUs(&pod); !ok || gpus != job.Spec.GPUs {
			log.Info("GPU count changed, replacing pod")

			if err := r.Delete(ctx, &pod); client.IgnoreNotFound(err) != nil {
				return ctrl.Result{}, err
			}

			return ctrl.Result{RequeueAfter: 2 * time.Second}, nil
		}
	case !apierrors.IsNotFound(err):
		return ctrl.Result{}, err
	}

	if err := r.Apply(ctx, desiredPod(&job), client.FieldOwner(manager), client.ForceOwnership); err != nil {
		return ctrl.Result{}, err
	}

	// A pod that was just created is not in the cache yet; its zero value
	// yields the Pending phase, and the pod's creation event triggers the
	// reconcile that reports its real progress.
	status := gpujobv1.StatusFromPod(&pod, job.Generation)

	if job.Status != status {
		log.Info("status changed", "phase", status.Phase, "message", status.Message)
		patch := client.MergeFrom(job.DeepCopy())
		job.Status = status

		if err := r.Status().Patch(ctx, &job, patch); err != nil {
			return ctrl.Result{}, err
		}
	}

	// Changes to the pod trigger the next reconcile through Owns.
	return ctrl.Result{}, nil
}


// main builds a manager, registers the reconciler, and runs until Ctrl-C.
func main() {
	ctrl.SetLogger(zap.New(zap.UseDevMode(true))) // readable console output
	setup := ctrl.Log.WithName("setup")

	scheme := runtime.NewScheme()

	if err := clientgoscheme.AddToScheme(scheme); err != nil {
		setup.Error(err, "registering built-in types")
		os.Exit(1)
	}

	if err := gpujobv1.AddToScheme(scheme); err != nil {
		setup.Error(err, "registering GpuJob")
		os.Exit(1)
	}

	mgr, err := ctrl.NewManager(ctrl.GetConfigOrDie(), ctrl.Options{
		Scheme:  scheme,
		Metrics: metricsserver.Options{BindAddress: "0"}, // no metrics endpoint
	})

	if err != nil {
		setup.Error(err, "creating manager")
		os.Exit(1)
	}

	err = ctrl.NewControllerManagedBy(mgr).
		For(&gpujobv1.GpuJob{}).
		Owns(&corev1.Pod{}).
		Complete(&GpuJobReconciler{Client: mgr.GetClient()})
	if err != nil {
		setup.Error(err, "creating controller")
		os.Exit(1)
	}

	if err := mgr.Start(ctrl.SetupSignalHandler()); err != nil {
		setup.Error(err, "running manager")
		os.Exit(1)
	}
}

Listing 1.69 The controller, complete, in Rust. code/rust/gpujob/src/bin/controller.rs

//! A controller for `GpuJob` objects. It runs each job as a pod that requests
//! the job's GPUs, and reports the pod's progress in the job's status.
//!
//! Run it against the cluster in the current kubeconfig context with
//! `cargo run -p gpujob` from `code/rust` (the package's default binary).

use futures::StreamExt;
use gpujob::{GpuJob, GpuJobStatus};
use k8s_openapi::api::core::v1::{Container, Pod, PodSpec, ResourceRequirements};
use k8s_openapi::apimachinery::pkg::api::resource::Quantity;
use kube::api::{DeleteParams, ObjectMeta, Patch, PatchParams};
use kube::runtime::controller::{Action, Controller};
use kube::runtime::watcher;
use kube::{Api, Client, Resource, ResourceExt};
use std::collections::BTreeMap;
use std::sync::Arc;
use std::time::Duration;

/// The field manager name under which the controller writes with server-side apply.
const MANAGER: &str = "gpujob-controller";

/// The extended resource through which a container requests GPUs.
const GPU: &str = "nvidia.com/gpu";

/// Everything a reconcile can fail with.
#[derive(Debug, thiserror::Error)]
enum Error {
    /// A request to the API server failed.
    #[error("kubernetes API: {0}")]
    Kube(#[from] kube::Error),
    /// The status could not be serialized into a patch.
    #[error("serializing status: {0}")]
    Json(#[from] serde_json::Error),
}

/// State shared by every reconcile.
struct Context {
    /// Client for the cluster's API server.
    client: Client,
}

/// Builds the pod that should exist for `job`.
///
/// The owner reference makes Kubernetes delete the pod together with the job,
/// and it lets the controller map each pod event to the job that owns the pod.
fn desired_pod(job: &GpuJob) -> Pod {
    let gpus = BTreeMap::from([(GPU.to_string(), Quantity(job.spec.gpus.to_string()))]);
    Pod {
        metadata: ObjectMeta {
            name: Some(job.name_any()),
            owner_references: job.controller_owner_ref(&()).map(|owner| vec![owner]),
            ..Default::default()
        },
        spec: Some(PodSpec {
            restart_policy: Some("Never".into()),
            containers: vec![Container {
                name: "job".into(),
                image: Some(job.spec.image.clone()),
                command: Some(job.spec.command.clone()).filter(|c| !c.is_empty()),
                resources: Some(ResourceRequirements {
                    limits: Some(gpus),
                    ..Default::default()
                }),
                ..Default::default()
            }],
            ..Default::default()
        }),
        ..Default::default()
    }
}

/// Returns the number of GPUs that an existing pod's container requests, read
/// from its limits, or `None` if the pod has no GPU limit.
fn pod_gpus(pod: &Pod) -> Option<u32> {
    let container = pod.spec.as_ref()?.containers.first()?;
    let limits = container.resources.as_ref()?.limits.as_ref()?;

    limits.get(GPU)?.0.parse().ok()
}

/// Makes the cluster match one job: ensures its pod exists with the requested
/// GPUs, and writes the pod's progress into the job's status.
async fn reconcile(job: Arc<GpuJob>, ctx: Arc<Context>) -> Result<Action, Error> {
    let namespace = job.namespace().unwrap_or_default();
    let name = job.name_any();
    let pods: Api<Pod> = Api::namespaced(ctx.client.clone(), &namespace);
    let jobs: Api<GpuJob> = Api::namespaced(ctx.client.clone(), &namespace);

    // The resources of a pod are fixed when it is created. When the job asks
    // for a different number of GPUs, delete the pod and come back after it is gone.
    if let Some(pod) = pods.get_opt(&name).await?
        && pod_gpus(&pod) != Some(job.spec.gpus)
    {
        tracing::info!(%name, "GPU count changed, replacing pod");
        pods.delete(&name, &DeleteParams::default()).await?;

        return Ok(Action::requeue(Duration::from_secs(2)));
    }

    let apply = PatchParams::apply(MANAGER).force();
    let pod = pods
        .patch(&name, &apply, &Patch::Apply(desired_pod(&job)))
        .await?;

    let status = GpuJobStatus::from_pod(&pod, job.metadata.generation);

    if job.status.as_ref() != Some(&status) {
        let message = status.message.as_deref().unwrap_or("");
        tracing::info!(%name, phase = ?status.phase, message, "status changed");
        let patch = serde_json::json!({
            "apiVersion": GpuJob::api_version(&()),
            "kind": GpuJob::kind(&()),
            "status": status,
        });
        jobs.patch_status(&name, &apply, &Patch::Apply(patch))
            .await?;
    }

    // Changes to the pod trigger the next reconcile through `owns`.
    Ok(Action::await_change())
}

/// Decides what happens after a failed reconcile: log it and retry in five seconds.
#[allow(clippy::needless_pass_by_value)] // the signature is fixed by `Controller::run`
fn error_policy(job: Arc<GpuJob>, error: &Error, _ctx: Arc<Context>) -> Action {
    tracing::warn!(name = %job.name_any(), %error, "reconcile failed");

    Action::requeue(Duration::from_secs(5))
}

/// Starts the controller and runs it until Ctrl-C.
#[tokio::main]
async fn main() -> Result<(), kube::Error> {
    tracing_subscriber::fmt()
        .with_env_filter("controller=info,kube=warn")
        .with_target(false)
        .with_ansi(std::io::IsTerminal::is_terminal(&std::io::stdout()))
        .init();
    let client = Client::try_default().await?;
    let jobs: Api<GpuJob> = Api::all(client.clone());
    let pods: Api<Pod> = Api::all(client.clone());

    Controller::new(jobs, watcher::Config::default())
        .owns(pods, watcher::Config::default())
        .shutdown_on_signal()
        .run(reconcile, error_policy, Arc::new(Context { client }))
        .for_each(|result| async move {
            if let Err(error) = result {
                tracing::debug!(%error, "reconcile error");
            }
        })
        .await;
    Ok(())
}

Listing 1.70 The two-client program, complete, in Go. code/go/gpujob/cmd/conflict/main.go

// Command conflict shows the API server refusing a write with an out-of-date
// resource version. Two clients read the train job, and both write it back
// with Update. The first write succeeds and changes the version, so the second
// fails with 409 Conflict.
//
// Run it from code/go, while a GpuJob named train exists, with:
//
//	go run ./gpujob/cmd/conflict
package main

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"os"

	apierrors "k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/apimachinery/pkg/runtime"
	ctrl "sigs.k8s.io/controller-runtime"
	"sigs.k8s.io/controller-runtime/pkg/client"

	gpujobv1 "gpuk8s/gpujob/api/v1"
)

// job identifies the job that both clients read and write.
var job = client.ObjectKey{Namespace: "default", Name: "train"}

// main performs the two writes and reports how the server answered each.
func main() {
	if err := run(context.Background()); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}

// run does the work of main and returns the first unexpected error.
func run(ctx context.Context) error {
	scheme := runtime.NewScheme()

	if err := gpujobv1.AddToScheme(scheme); err != nil {
		return err
	}

	c, err := client.New(ctrl.GetConfigOrDie(), client.Options{Scheme: scheme})

	if err != nil {
		return err
	}

	var first gpujobv1.GpuJob

	if err := c.Get(ctx, job, &first); err != nil {
		return err
	}

	second := first.DeepCopy()
	fmt.Printf("both clients read %s at v%s\n", job.Name, first.ResourceVersion)

	first.Spec.GPUs++

	if err := c.Update(ctx, &first); err != nil {
		return err
	}

	fmt.Printf("first write stored, now v%s\n", first.ResourceVersion)

	second.Spec.GPUs += 2
	err = c.Update(ctx, second)
	var status *apierrors.StatusError

	if errors.As(err, &status) && status.ErrStatus.Code == http.StatusConflict {
		fmt.Printf("second write refused: %d %s\n", status.ErrStatus.Code, status.ErrStatus.Message)

		return nil
	}

	return fmt.Errorf("unexpected result: %v", err)
}

2. What one GPU means

So far our pods have asked for nvidia.com/gpu: 1 as if a GPU were a single, indivisible thing. In this chapter, we will open that unit up. We will first look at what a GPU is made of and measure the one in our cluster, then see exactly what Kubernetes gives a pod when it grants one unit. Then we will go through the ways a GPU can be shared or divided: time-slicing, the Multi-Process Service, Multi-Instance GPU, virtual GPUs, and Dynamic Resource Allocation. We will try time-slicing on our own cluster and measure what it does.

The programs in this chapter

We use four programs, two that run on the GPU and two that look at the cluster.

gpu_info and gpu_load are CUDA programs in code/gpu-sharing. We compile them on the host with nvcc, the CUDA compiler, for the compute capability of our GPU, and link the CUDA runtime library into each binary. Each binary then needs only the driver library, libcuda, at run time, and that library is one of the files the NVIDIA container runtime mounts into a GPU container (Appendix A). We copy both binaries into one image, gpu-sharing:1, and run them in pods that request nvidia.com/gpu. When such a pod starts, the device plugin tells the kubelet to set NVIDIA_VISIBLE_DEVICES in the container, the NVIDIA runtime adds the GPU’s device files and libcuda, and the CUDA calls in our programs reach the GPU through them.

• gpu_info asks the CUDA runtime for the properties of the GPU it sees and prints them, one per line. It shows what one nvidia.com/gpu unit contains.
• gpu_load launches a compute-heavy kernel ten times, waits for all ten to finish, and prints the elapsed time. Several copies running at once show how much GPU time each one receives.

gpuusage is written in Go (code/go/gpuusage) and in Rust (code/rust/gpuusage). It is an ordinary Kubernetes client that runs on our workstation. It lists every node and every pod through the API server, adds up the nvidia.com/gpu units that running pods hold on each node, and prints them next to the node’s capacity. It shows the scheduler’s view of the GPU: how many units exist and who holds them. We run it in a second terminal, from code/go with go run ./gpuusage or from code/rust with cargo run -p gpuusage, while the pods from code/gpu-sharing run.

What a GPU is made of

A GPU is a large number of simple processors grouped into streaming multiprocessors (SMs). Each SM contains arithmetic units (the CUDA cores for ordinary floating-point and integer work, and Tensor Cores for matrix multiplication), a large register file, schedulers, and a block of fast on-chip memory that is used as shared memory and L1 cache. A CUDA program launches a kernel, a function that runs on the GPU, as a grid of thread blocks. The GPU assigns each block to an SM, and the SM runs its threads in groups of 32, called warps. An SM keeps many warps resident at once and switches between them whenever one waits for memory, which is how a GPU hides memory latency: by having other work ready, where a CPU would rely on large caches.

Memory is arranged in levels, from fastest and smallest to slowest and largest:

• registers, private to each thread;
• shared memory and L1 cache, one block per SM;
• the L2 cache, shared by all SMs;
• device memory: HBM on data center GPUs, GDDR on graphics cards, and, on integrated GPUs such as the Jetson’s, the same LPDDR memory that the CPU uses.

Next to the SMs sit fixed-function engines. Copy engines move data between host memory and device memory while kernels run, so transfers and computation can overlap. Video decoders and encoders (NVDEC and NVENC) and JPEG decoders handle media without using the SMs.

A discrete GPU is a separate card with its own memory, attached over PCIe or a dedicated interconnect; a program copies data into that memory before a kernel can use it. An integrated GPU sits on the same chip as the CPU and shares the system’s memory with it.

Measuring our GPU

The CUDA runtime reports most of these properties through cudaGetDeviceProperties. This program prints them, along with the current free and total device memory:

Listing 2.1 Reading the GPU's properties. code/gpu-sharing/gpu_info.cu

// gpu_info prints what one Kubernetes nvidia.com/gpu unit contains on the
// machine it runs on: the GPU's compute units, memory, caches, and engines.
//
// How it runs, end to end:
//   1. nvcc compiles this file on the host for the Orin's GPU (-arch=sm_87) and
//      links the CUDA runtime library statically (-cudart static), so the
//      binary needs only the CUDA driver library, libcuda, at run time.
//   2. The binary is copied into the gpu-sharing:1 image and run in a pod that
//      requests nvidia.com/gpu: 1.
//   3. The device plugin's Allocate answer sets NVIDIA_VISIBLE_DEVICES in the
//      container, and the NVIDIA container runtime mounts the GPU's device
//      files and libcuda into it.
//   4. The CUDA runtime calls below load libcuda, which opens the device files
//      and asks the kernel driver for the device's properties.
//
// Build: nvcc -O2 -arch=sm_87 -cudart static -o gpu_info gpu_info.cu
#include <cstdio>
#include <cuda_runtime.h>

// kGiB is the number of bytes in one gibibyte, for printing memory sizes.
constexpr double kGiB = 1024.0 * 1024.0 * 1024.0;

// main queries device 0 and prints one property per line. It exits with
// status 1 when no CUDA device is visible, for example in a container that
// was not given a GPU.
int main() {
    cudaDeviceProp p;
    if (cudaGetDeviceProperties(&p, 0) != cudaSuccess) {
        std::printf("no CUDA device visible\n");
        return 1;
    }

    // cudaMemGetInfo reports the memory that the driver can still allocate and
    // the total. On an integrated GPU both come from the system's shared RAM.
    size_t free_bytes = 0, total_bytes = 0;
    cudaMemGetInfo(&free_bytes, &total_bytes);

    std::printf("device                      %s\n", p.name);
    std::printf("compute capability          %d.%d\n", p.major, p.minor);
    std::printf("integrated with the CPU     %s\n", p.integrated ? "yes" : "no");
    std::printf("streaming multiprocessors   %d\n", p.multiProcessorCount);
    std::printf("max resident threads/SM     %d\n", p.maxThreadsPerMultiProcessor);
    std::printf("registers per SM            %d\n", p.regsPerMultiprocessor);
    std::printf("shared memory per SM        %zu KiB\n", p.sharedMemPerMultiprocessor / 1024);
    std::printf("L2 cache                    %d KiB\n", p.l2CacheSize / 1024);
    std::printf("device memory (total)       %.1f GiB\n", total_bytes / kGiB);
    std::printf("device memory (free now)    %.1f GiB\n", free_bytes / kGiB);
    std::printf("copy engines                %d\n", p.asyncEngineCount);
    std::printf("concurrent kernels          %s\n", p.concurrentKernels ? "yes" : "no");
    std::printf("managed memory              %s\n", p.managedMemory ? "yes" : "no");
    return 0;
}

The image for this chapter holds gpu_info and a second program, gpu_load, which we use later:

Listing 2.2 The image that holds both CUDA programs. code/gpu-sharing/Dockerfile

FROM ubuntu:22.04
COPY gpu_info gpu_load /usr/local/bin/
CMD ["gpu_info"]

$ cd code/gpu-sharing
$ /usr/local/cuda/bin/nvcc -O2 -arch=sm_87 -cudart static -o gpu_info gpu_info.cu
$ /usr/local/cuda/bin/nvcc -O2 -arch=sm_87 -cudart static -o gpu_load gpu_load.cu
$ docker build -t gpu-sharing:1 .
$ kind load docker-image gpu-sharing:1 --name gpu-lab

We run it in a pod that asks for one GPU:

Listing 2.3 A pod that runs `gpu_info` with one GPU. code/gpu-sharing/gpu-info-pod.yaml

apiVersion: v1
kind: Pod
metadata:
  name: gpu-info
spec:
  restartPolicy: Never
  containers:
    - name: info
      image: gpu-sharing:1
      imagePullPolicy: Never
      command: ["gpu_info"]
      resources:
        limits:
          nvidia.com/gpu: 1

$ kubectl apply -f gpu-info-pod.yaml
$ kubectl logs gpu-info
device                      Orin
compute capability          8.7
integrated with the CPU     yes
streaming multiprocessors   8
max resident threads/SM     1536
registers per SM            65536
shared memory per SM        164 KiB
L2 cache                    2048 KiB
device memory (total)       15.3 GiB
device memory (free now)    1.8 GiB
copy engines                2
concurrent kernels          yes
managed memory              yes

The GPU in our cluster has 8 SMs. An Ampere SM of this generation has 128 CUDA cores, which gives the 1,024 cores in NVIDIA’s specification for the Orin NX, and each SM can hold 1,536 threads at once, 12,288 for the whole GPU. All SMs share a 2 MiB L2 cache, and two copy engines move data. For comparison, a data center H100 SXM5 has 132 SMs and 80 GB of HBM3 memory.

Figure 2.1 draws those numbers as one picture.

The memory lines show what “integrated” means in practice. The GPU’s device memory is the board’s 15.3 GiB of LPDDR5, the same memory the operating system and every CPU program use. Only 1.8 GiB was free when the pod ran, because other programs on the board were using the rest. On a discrete GPU, device memory belongs to the GPU alone; on an integrated GPU, the amount available to a kernel depends on everything else running on the machine.

The same measurement on a Thor

We also have a Jetson AGX Thor developer kit on the network. It runs Jetson Linux R39.2.1 with CUDA 13.2 and has no Kubernetes on it, so we compile gpu_info on the board and run it directly:

$ /usr/local/cuda/bin/nvcc -O2 -cudart static -o gpu_info gpu_info.cu
$ ./gpu_info
device                      NVIDIA Thor
compute capability          11.0
integrated with the CPU     yes
streaming multiprocessors   20
max resident threads/SM     1536
registers per SM            65536
shared memory per SM        228 KiB
L2 cache                    32768 KiB
device memory (total)       122.8 GiB
device memory (free now)    44.8 GiB
copy engines                2
concurrent kernels          yes
managed memory              yes

The program is the same source file, and the two boards answer differently in four lines:

Twenty SMs with 128 CUDA cores each give the 2,560 cores in NVIDIA’s specification for the Thor. The L2 cache is sixteen times as large, and the device memory is eight times as large. It is again the board’s own memory: 44.8 GiB were free when we ran the program, because two other programs on the Thor, a model server among them, held 24 GiB and 39 GiB at the time.

Both boards would advertise the same thing to a Kubernetes scheduler: nvidia.com/gpu: 1. A job that needs 60 GiB of device memory fits on one of them, and the count alone does not say which. Chapter 5 returns to this with node labels.

What one unit gives a pod

By default, the device plugin advertises one nvidia.com/gpu unit for each physical GPU, and a pod that is granted a unit receives the whole device: every SM, all of its memory, and all of its engines. Kubernetes counts extended resources as integers. A pod cannot ask for half a GPU, and the node never hands out more units than it advertises.

gpuusage shows the same accounting that the scheduler does. It sums the GPU limits of every pod that is bound to a node and has not finished, and compares the sum with the node’s allocatable units. The Go version uses client-go’s typed clientset, and the Rust version uses kube-rs.

The first part decides what one pod holds.

Listing 2.4 What one pod holds, in Go. code/go/gpuusage/main.go

// gpu is the extended resource that the NVIDIA device plugin advertises.
const gpu corev1.ResourceName = "nvidia.com/gpu"

// holder is a pod that holds GPU units on a node.
type holder struct {
	// pod is the pod's namespace and name, as "namespace/name".
	pod string
	// units is the number of nvidia.com/gpu units the pod holds.
	units int64
}

// podGPUs returns the number of GPU units a pod requests: the sum of the
// nvidia.com/gpu limits of its containers.
func podGPUs(pod *corev1.Pod) int64 {
	var units int64

	for _, c := range pod.Spec.Containers {
		if q, ok := c.Resources.Limits[gpu]; ok {
			units += q.Value()
		}
	}

	return units
}

// holdsResources reports whether a pod currently holds its resources on a
// node. A pod holds them from the moment it is bound to a node until it
// finishes; a pod that has succeeded or failed gives them back.
func holdsResources(pod *corev1.Pod) bool {
	finished := pod.Status.Phase == corev1.PodSucceeded || pod.Status.Phase == corev1.PodFailed

	return pod.Spec.NodeName != "" && !finished
}

Listing 2.5 What one pod holds, in Rust. code/rust/gpuusage/src/main.rs

/// The extended resource that the NVIDIA device plugin advertises.
const GPU: &str = "nvidia.com/gpu";

/// A pod that holds GPU units on a node.
struct Holder {
    /// The pod's namespace and name, as `namespace/name`.
    pod: String,
    /// The number of `nvidia.com/gpu` units the pod holds.
    units: i64,
}

/// Returns the number of GPU units a pod requests: the sum of the
/// `nvidia.com/gpu` limits of its containers.
fn pod_gpus(pod: &Pod) -> i64 {
    let Some(spec) = &pod.spec else { return 0 };
    spec.containers
        .iter()
        .filter_map(|c| c.resources.as_ref()?.limits.as_ref()?.get(GPU))
        .filter_map(|q| q.0.parse::<i64>().ok())
        .sum()
}

/// Returns the node a pod currently holds resources on. A pod holds them from
/// the moment it is bound to a node until it finishes; a pod that has
/// succeeded or failed gives them back.
fn holding_node(pod: &Pod) -> Option<&str> {
    let phase = pod.status.as_ref().and_then(|s| s.phase.as_deref());

    if matches!(phase, Some("Succeeded" | "Failed")) {
        return None;
    }

    pod.spec.as_ref()?.node_name.as_deref()
}

/// Reads the number of GPU units from a node's capacity or allocatable map.
fn gpu_units(resources: Option<&BTreeMap<String, Quantity>>) -> i64 {
    resources
        .and_then(|r| r.get(GPU))
        .and_then(|q| q.0.parse().ok())
        .unwrap_or(0)
}

The second part asks the API server and prints the table.

Listing 2.6 Listing, grouping, and printing, in Go. code/go/gpuusage/main.go

// run lists every node and every pod, groups the GPU-holding pods by node,
// and prints one table row per node followed by the pods that hold its units.
func run(ctx context.Context) error {
	clientset, err := kubernetes.NewForConfig(ctrl.GetConfigOrDie())

	if err != nil {
		return err
	}

	nodes, err := clientset.CoreV1().Nodes().List(ctx, metav1.ListOptions{})

	if err != nil {
		return err
	}

	pods, err := clientset.CoreV1().Pods(metav1.NamespaceAll).List(ctx, metav1.ListOptions{})

	if err != nil {
		return err
	}

	holders := map[string][]holder{}

	for i := range pods.Items {
		pod := &pods.Items[i]

		if units := podGPUs(pod); units > 0 && holdsResources(pod) {
			name := pod.Namespace + "/" + pod.Name
			holders[pod.Spec.NodeName] = append(holders[pod.Spec.NodeName], holder{name, units})
		}
	}

	table := tabwriter.NewWriter(os.Stdout, 0, 0, 2, ' ', 0)
	fmt.Fprintln(table, "NODE\tCAPACITY\tALLOCATABLE\tIN USE\tFREE")

	for _, node := range nodes.Items {
		capacity := node.Status.Capacity[gpu]
		allocatable := node.Status.Allocatable[gpu]
		var inUse int64

		for _, h := range holders[node.Name] {
			inUse += h.units
		}

		fmt.Fprintf(table, "%s\t%d\t%d\t%d\t%d\n",
			node.Name, capacity.Value(), allocatable.Value(), inUse, allocatable.Value()-inUse)
	}

	if err := table.Flush(); err != nil {
		return err
	}

	for _, node := range nodes.Items {
		list := holders[node.Name]
		slices.SortFunc(list, func(a, b holder) int { return cmp.Compare(a.pod, b.pod) })

		for _, h := range list {
			fmt.Printf("  %s holds %d on %s\n", h.pod, h.units, node.Name)
		}
	}

	return nil
}

Listing 2.7 Listing, grouping, and printing, in Rust. code/rust/gpuusage/src/main.rs

/// Lists every node and every pod, groups the GPU-holding pods by node, and
/// prints one table row per node followed by the pods that hold its units.
#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    let client = Client::try_default().await?;
    let nodes = Api::<Node>::all(client.clone())
        .list(&ListParams::default())
        .await?;
    let pods = Api::<Pod>::all(client).list(&ListParams::default()).await?;

    let mut holders: BTreeMap<String, Vec<Holder>> = BTreeMap::new();

    for pod in &pods {
        let units = pod_gpus(pod);
        if let Some(node) = holding_node(pod)
            && units > 0
        {
            let name = format!("{}/{}", pod.namespace().unwrap_or_default(), pod.name_any());
            holders
                .entry(node.to_string())
                .or_default()
                .push(Holder { pod: name, units });
        }
    }

    println!(
        "{:<23}{:<10}{:<13}{:<8}FREE",
        "NODE", "CAPACITY", "ALLOCATABLE", "IN USE"
    );

    for node in &nodes {
        let name = node.name_any();
        let in_use: i64 = holders
            .get(&name)
            .map_or(0, |h| h.iter().map(|h| h.units).sum());
        let status = node.status.as_ref();
        let capacity = gpu_units(status.and_then(|s| s.capacity.as_ref()));
        let allocatable = gpu_units(status.and_then(|s| s.allocatable.as_ref()));
        println!(
            "{:<23}{:<10}{:<13}{:<8}{}",
            name,
            capacity,
            allocatable,
            in_use,
            allocatable - in_use
        );
    }

    for (node, list) in &mut holders {
        list.sort_by(|a, b| a.pod.cmp(&b.pod));

        for h in list.iter() {
            println!("  {} holds {} on {node}", h.pod, h.units);
        }
    }

    Ok(())
}

Both programs list the nodes and the pods with one request each, keep the pods that hold GPU units, and group them by node. A pod counts from the moment the scheduler binds it to a node until it succeeds or fails; a finished pod has released its units even though the object still exists. With no GPU pods running, both print the same report:

$ go run ./gpuusage
NODE                   CAPACITY  ALLOCATABLE  IN USE  FREE
gpu-lab-control-plane  1         1            0       1
$ cargo run -p gpuusage
NODE                   CAPACITY  ALLOCATABLE  IN USE  FREE
gpu-lab-control-plane  1         1            0       1

This is simple and gives each pod exclusive use of its GPU, but it is coarse. An inference service that keeps 5 percent of the SMs busy still holds the entire GPU, and the other 95 percent sits idle while other pods wait. The rest of this chapter covers the ways to share or divide a GPU, each with a different trade-off between isolation and flexibility.

Sharing by taking turns: time-slicing

Without further configuration, a GPU runs work from one process’s CUDA context at a time. When several processes use the same GPU, the driver switches between their contexts, giving each a slice of time on all the SMs. Time-slicing in Kubernetes uses exactly this: the device plugin advertises each GPU as several units, called replicas, and every pod that receives a replica of the same GPU gets the same device.

NVIDIA’s GPU Operator documentation is direct about the cost: “Unlike Multi-Instance GPU (MIG), there is no memory or fault-isolation between replicas.” Every pod sees all of the GPU’s memory, one pod can allocate so much that the others fail, and, in the words of the device plugin’s README, workloads run “in the same fault-domain as of all the others (meaning if one workload crashes, they all do)”.

We can measure the effect. gpu_load launches a kernel ten times, each launch keeping 1,024 blocks of 256 threads busy with 200,000 multiply-adds per thread, and prints how long the work took:

Listing 2.8 The sizes, and the kernel that keeps one thread busy. code/gpu-sharing/gpu_load.cu

// kBlocks and kThreads give 262,144 threads per launch, enough to occupy every
// SM of the GPU many times over.
constexpr int kBlocks = 1024;
constexpr int kThreads = 256;

// kIterations is the number of multiply-adds each thread performs per launch.
constexpr int kIterations = 200000;

// kRounds is the number of kernel launches in one measurement.
constexpr int kRounds = 10;

// spin performs `iterations` dependent multiply-adds on a per-thread value and
// stores the result, so the compiler cannot remove the loop.
__global__ void spin(float *out, int iterations) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    float x = i * 1e-6f;
    for (int k = 0; k < iterations; k++) {
        x = x * 1.000001f + 0.5f;
    }
    out[i] = x;
}

Listing 2.9 Launching the kernel ten times and timing it. code/gpu-sharing/gpu_load.cu

// main runs the measurement and prints "<name>: <result>, <milliseconds> ms",
// where name is the first argument. It exits with status 1 if the GPU reported
// an error.
int main(int argc, char **argv) {
    const char *name = argc > 1 ? argv[1] : "load";
    float *out;
    cudaMalloc(&out, kBlocks * kThreads * sizeof(float));

    auto start = std::chrono::steady_clock::now();
    for (int r = 0; r < kRounds; r++) {
        spin<<<kBlocks, kThreads>>>(out, kIterations);
    }
    cudaError_t err = cudaDeviceSynchronize();
    double ms = std::chrono::duration<double, std::milli>(
                    std::chrono::steady_clock::now() - start).count();

    std::printf("%s: %s, %.0f ms\n", name, cudaGetErrorString(err), ms);
    cudaFree(out);
    return err == cudaSuccess ? 0 : 1;
}

A Job runs four copies of it at once, each asking for one nvidia.com/gpu:

Listing 2.10 A Job that runs four copies of `gpu_load`. code/gpu-sharing/gpu-load-job.yaml

# Four pods, each running the same fixed amount of GPU work and asking for one
# nvidia.com/gpu. How they run depends on how many units the node advertises.
apiVersion: batch/v1
kind: Job
metadata:
  name: gpu-load
spec:
  completions: 4
  parallelism: 4
  template:
    spec:
      restartPolicy: Never
      containers:
        - name: load
          image: gpu-sharing:1
          imagePullPolicy: Never
          command: ["sh", "-c", "date +%T.%N | cut -c1-12; gpu_load $HOSTNAME"]
          resources:
            limits:
              nvidia.com/gpu: 1

With the device plugin as configured in Appendix A, the node advertises one unit. Three of the four pods wait in Pending while the first runs:

$ kubectl apply -f gpu-load-job.yaml
$ kubectl get pods -l job-name=gpu-load
gpu-load-5wf4d Completed
gpu-load-mbw4m Pending
gpu-load-w46sf Pending
gpu-load-w9h7x ContainerCreating
$ go run ./gpuusage
NODE                   CAPACITY  ALLOCATABLE  IN USE  FREE
gpu-lab-control-plane  1         1            1       0
  default/gpu-load-w9h7x holds 1 on gpu-lab-control-plane

gpuusage shows the reason: the node’s one unit is held by the running pod, so the scheduler has nowhere to place the others. Each pod prints its start time and its result. They ran one after another, about three seconds apart, and each needed about 620 milliseconds of GPU time:

23:29:59.987 gpu-load-5wf4d: no error, 612 ms
23:30:03.466 gpu-load-w9h7x: no error, 622 ms
23:30:06.487 gpu-load-w46sf: no error, 616 ms
23:30:09.495 gpu-load-mbw4m: no error, 634 ms

Time-slicing is switched on through the device plugin’s configuration file. This ConfigMap advertises each GPU as four replicas:

Listing 2.11 The time-slicing configuration. code/gpu-sharing/time-slicing-config.yaml

# Device plugin configuration that advertises each GPU as four
# nvidia.com/gpu units. Pods that receive units of the same GPU take turns on it.
apiVersion: v1
kind: ConfigMap
metadata:
  name: nvidia-device-plugin-config
  namespace: kube-system
data:
  config.yaml: |
    version: v1
    sharing:
      timeSlicing:
        resources:
          - name: nvidia.com/gpu
            replicas: 4

The plugin reads the file named by its CONFIG_FILE setting. This manifest is the plugin from Appendix A with the ConfigMap mounted and CONFIG_FILE set:

Listing 2.12 The device plugin, reading that configuration. code/gpu-sharing/device-plugin-time-sliced.yaml

# The device plugin from jetson-gpu/device-plugin.yaml, with its configuration
# file taken from the ConfigMap in time-slicing-config.yaml. Apply
# jetson-gpu/device-plugin.yaml again to switch time-slicing off.
apiVersion: apps/v1
kind: DaemonSet
metadata:
  name: nvidia-device-plugin
  namespace: kube-system
spec:
  selector:
    matchLabels:
      app: nvidia-device-plugin
  template:
    metadata:
      labels:
        app: nvidia-device-plugin
    spec:
      priorityClassName: system-node-critical
      tolerations:
        - key: nvidia.com/gpu
          operator: Exists
          effect: NoSchedule
      containers:
        - name: plugin
          image: nvcr.io/nvidia/k8s-device-plugin:v0.20.1
          env:
            - name: DEVICE_DISCOVERY_STRATEGY
              value: tegra
            - name: NVIDIA_DRIVER_ROOT      # driver root as the node sees it
              value: /
            - name: DRIVER_ROOT_CTR_PATH    # the same root inside this container
              value: /driver-root
            - name: FAIL_ON_INIT_ERROR
              value: "true"
            - name: CONFIG_FILE
              value: /etc/nvidia-device-plugin/config.yaml
          volumeMounts:
            - name: device-plugins
              mountPath: /var/lib/kubelet/device-plugins
            - name: driver-root
              mountPath: /driver-root
              readOnly: true
            - name: config
              mountPath: /etc/nvidia-device-plugin
      volumes:
        - name: device-plugins
          hostPath:
            path: /var/lib/kubelet/device-plugins
        - name: driver-root
          hostPath:
            path: /
        - name: config
          configMap:
            name: nvidia-device-plugin-config

$ kubectl apply -f time-slicing-config.yaml -f device-plugin-time-sliced.yaml
$ kubectl get node -o jsonpath='{.items[0].status.allocatable.nvidia\.com/gpu}'; echo
4

The same GPU now counts as four units. Running the same Job again, all four pods start at once:

$ kubectl delete job gpu-load
$ kubectl apply -f gpu-load-job.yaml
$ kubectl get pods -l job-name=gpu-load
gpu-load-dk7h8 Running
gpu-load-grrkx Running
gpu-load-s4lb2 Running
gpu-load-t28sd Running
$ cargo run -p gpuusage
NODE                   CAPACITY  ALLOCATABLE  IN USE  FREE
gpu-lab-control-plane  4         4            4       0
  default/gpu-load-dk7h8 holds 1 on gpu-lab-control-plane
  default/gpu-load-grrkx holds 1 on gpu-lab-control-plane
  default/gpu-load-s4lb2 holds 1 on gpu-lab-control-plane
  default/gpu-load-t28sd holds 1 on gpu-lab-control-plane

23:30:25.915 gpu-load-dk7h8: no error, 2198 ms
23:30:25.996 gpu-load-s4lb2: no error, 2231 ms
23:30:26.048 gpu-load-grrkx: no error, 2249 ms
23:30:26.052 gpu-load-t28sd: no error, 2260 ms

Each pod now took about 2,230 milliseconds for the same work, roughly 3.6 times as long as it took alone. The four pods took turns on the same 8 SMs. Time-slicing added no compute: it changed how many pods can hold the GPU at once and how long each one waits. Together, the four runs occupied the GPU for about 2.3 seconds, close to the 4 × 620 = 2,480 milliseconds they needed one after another.

The Thor shows the same behaviour with different numbers. We ran gpu_load there directly on the host, first alone and then as three processes started together:

$ ./gpu_load thor-alone
thor-alone: no error, 160 ms
$ ./gpu_load copy-1 & ./gpu_load copy-2 & ./gpu_load copy-3 & wait
copy-3: no error, 467 ms
copy-2: no error, 477 ms
copy-1: no error, 479 ms

Alone, the work that takes the Orin NX about 620 milliseconds takes the Thor 160. Three runs gave 190, 160, and 159 milliseconds, the first including the start-up of the CUDA context. With three processes on the GPU at once, each took about 475 milliseconds, which is three times 160. No device plugin was involved here. The driver switched between the three processes’ CUDA contexts by itself. Time-slicing in Kubernetes relies on that switching and adds the extra units that the scheduler counts.

The scheduler now believes the node has four GPUs. A pod that asks for two replicas gets two units of the same GPU and no more of its time; the plugin’s failRequestsGreaterThanOne option rejects such requests to make that explicit. To switch time-slicing off, we apply the original plugin manifest again:

$ kubectl apply -f ../jetson-gpu/device-plugin.yaml
$ go run ./gpuusage
NODE                   CAPACITY  ALLOCATABLE  IN USE  FREE
gpu-lab-control-plane  1         1            0       1

Running kernels side by side: MPS

Time-slicing runs one process’s kernels at a time. NVIDIA’s MPS documentation states the underlying rule: without MPS, work “from work queues belonging to different CUDA contexts cannot execute concurrently”. A process whose kernels fill only a few SMs leaves the rest idle during its slice.

The Multi-Process Service (MPS) removes that restriction. An MPS server becomes “the clients’ shared connection to the GPU and owner of the GPU scheduling resources”, and kernels from different processes then run on the GPU at the same time. MPS also lets an administrator limit each client: CUDA_MPS_ACTIVE_THREAD_PERCENTAGE caps the share of the GPU’s threads a client may use, and CUDA_MPS_PINNED_DEVICE_MEM_LIMIT caps the memory it may allocate. Each client has its own GPU address space, but error containment is limited: a fatal fault is reported to all clients on the affected GPU, “without indicating which client generated the error”.

The device plugin can hand out MPS shares in the same way as time-slicing replicas, each with an equal fraction of the GPU’s memory and compute. Its README marks this support as experimental and limits it to whole GPUs without MIG. MPS runs on Tegra devices such as ours with restrictions, so we do not use it here.

Dividing the hardware: Multi-Instance GPU

Time-slicing and MPS share one GPU in software. Multi-Instance GPU (MIG) divides it in hardware. According to NVIDIA’s MIG user guide, a GPU can be “securely partitioned into up to seven separate GPU Instances”, and within each instance “the on-chip crossbar ports, L2 cache banks, memory controllers, and DRAM address busses are all assigned uniquely to an individual instance”. SMs, copy engines, and decoders are divided as well. An instance therefore keeps its cache and memory bandwidth “even if other tasks are thrashing their own caches or saturating their DRAM interfaces”, and a fault in one instance does not affect the others.

Instances come in fixed sizes, called profiles. A profile name gives the number of compute slices and the memory. The profiles of an H100 with 80 GB, from the MIG user guide:

profile      memory   SMs    copy engines   instances per GPU
-----------  -------  -----  -------------  -----------------
1g.10gb      1/8      1/7    1              7
1g.20gb      1/4      1/7    1              4
2g.20gb      2/8      2/7    2              3
3g.40gb      4/8      3/7    3              2
4g.40gb      4/8      4/7    4              1
7g.80gb      full     7/7    8              1

The device plugin exposes MIG instances according to its MIG_STRATEGY setting: with the single strategy every instance is advertised as nvidia.com/gpu, and with the mixed strategy each profile gets its own resource name, such as nvidia.com/mig-1g.10gb. A pod then asks for a slice of known size and gets guaranteed memory and bandwidth.

MIG requires a supporting GPU: the A100, H100, H200, and B200 families support seven instances, and smaller models support fewer. The Orin in our cluster does not support MIG.

Virtual GPUs for virtual machines: vGPU

Everything so far shares a GPU between containers on the same operating system. NVIDIA’s vGPU software shares a GPU between virtual machines: it “enables multiple virtual machines (VMs) to have simultaneous, direct access to a single physical GPU”. A vGPU is either time-sliced, in which case the GPU “performs the work for each VM serially” under a scheduler that can give each VM a fixed or equal share, or backed by a MIG instance, in which case VMs run in parallel with MIG’s isolation. Clusters that run virtual machines on Kubernetes can use vGPUs to give each VM its own GPU device.

Asking for GPUs by their properties: Dynamic Resource Allocation

The device plugin model reduces every GPU to an integer count of an opaque name. A pod cannot ask for “a GPU with at least 40 GB of memory”, and two containers cannot explicitly share one GPU. Dynamic Resource Allocation (DRA), generally available since Kubernetes 1.34, replaces the count with structured requests:

• a driver publishes the devices on each node, with their attributes such as model and memory size, as ResourceSlice objects;
• a DeviceClass groups devices of one kind, for example all NVIDIA GPUs;
• a pod refers to a ResourceClaim, which requests devices from a class and can filter them with expressions over their attributes; several pods can refer to the same claim to share its devices.

The scheduler then allocates concrete devices that satisfy each claim. Our cluster already serves these APIs:

$ kubectl api-resources --api-group=resource.k8s.io
NAME                     SHORTNAMES   APIVERSION           NAMESPACED   KIND
deviceclasses                         resource.k8s.io/v1   false        DeviceClass
devicetaintrules                      resource.k8s.io/v1   false        DeviceTaintRule
resourceclaims                        resource.k8s.io/v1   true         ResourceClaim
resourceclaimtemplates                resource.k8s.io/v1   true         ResourceClaimTemplate
resourceslices                        resource.k8s.io/v1   false        ResourceSlice

A driver has to publish the devices, and NVIDIA’s DRA driver describes its GPU support as not yet official: “While some GPU allocation features can be tried out, they are not yet officially supported.” Our cluster therefore stays with the device plugin.

Choosing a method

Figure 2.2 draws the four ways of handing out a GPU to containers as pictures of one GPU over time, and the table after it adds vGPU.

method          what a pod gets                 memory isolated   faults isolated
--------------  ------------------------------  ----------------  ---------------
whole GPU       the entire device               yes               yes
time-slicing    turns on the entire device      no                no
MPS             a share of threads and memory   by limits         partly
MIG             a hardware partition            yes               yes
vGPU            a device inside a VM            yes               yes

Whole GPUs suit training jobs that keep a GPU busy. Time-slicing suits development and light, trusted workloads that tolerate each other. MPS suits many small inference processes that each use part of the SMs. MIG suits multi-tenant clusters that need guaranteed memory and bandwidth for each tenant. vGPU suits workloads that must run in virtual machines.

The complete programs

Listing 2.13 `gpu_load`, complete. code/gpu-sharing/gpu_load.cu

// gpu_load keeps the GPU busy with a fixed amount of arithmetic and prints how
// long the work took. Running it alone and running several copies that share
// one GPU shows how much GPU time each copy receives.
//
// How it runs, end to end:
//   1. nvcc compiles this file on the host (-arch=sm_87, -cudart static), and
//      the binary is copied into the gpu-sharing:1 image.
//   2. Each pod of the gpu-load Job runs it with the pod's name as argument and
//      receives one nvidia.com/gpu unit, so the NVIDIA container runtime gives
//      it the GPU's device files and libcuda.
//   3. main launches the spin kernel kRounds times on the default stream. The
//      launches queue up on the GPU, and cudaDeviceSynchronize waits until all
//      of them have finished, so the measured time covers the GPU work.
//   4. When several pods share the GPU through time-slicing, the driver
//      switches between their CUDA contexts, and each copy's time grows by the
//      share of the GPU that the others take.
//
// Build: nvcc -O2 -arch=sm_87 -cudart static -o gpu_load gpu_load.cu
#include <chrono>
#include <cstdio>
#include <cuda_runtime.h>

// kBlocks and kThreads give 262,144 threads per launch, enough to occupy every
// SM of the GPU many times over.
constexpr int kBlocks = 1024;
constexpr int kThreads = 256;

// kIterations is the number of multiply-adds each thread performs per launch.
constexpr int kIterations = 200000;

// kRounds is the number of kernel launches in one measurement.
constexpr int kRounds = 10;

// spin performs `iterations` dependent multiply-adds on a per-thread value and
// stores the result, so the compiler cannot remove the loop.
__global__ void spin(float *out, int iterations) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    float x = i * 1e-6f;
    for (int k = 0; k < iterations; k++) {
        x = x * 1.000001f + 0.5f;
    }
    out[i] = x;
}

// main runs the measurement and prints "<name>: <result>, <milliseconds> ms",
// where name is the first argument. It exits with status 1 if the GPU reported
// an error.
int main(int argc, char **argv) {
    const char *name = argc > 1 ? argv[1] : "load";
    float *out;
    cudaMalloc(&out, kBlocks * kThreads * sizeof(float));

    auto start = std::chrono::steady_clock::now();
    for (int r = 0; r < kRounds; r++) {
        spin<<<kBlocks, kThreads>>>(out, kIterations);
    }
    cudaError_t err = cudaDeviceSynchronize();
    double ms = std::chrono::duration<double, std::milli>(
                    std::chrono::steady_clock::now() - start).count();

    std::printf("%s: %s, %.0f ms\n", name, cudaGetErrorString(err), ms);
    cudaFree(out);
    return err == cudaSuccess ? 0 : 1;
}

Listing 2.14 `gpuusage`, complete, in Go. code/go/gpuusage/main.go

// Command gpuusage reports, for every node in the cluster, how many
// nvidia.com/gpu units the node advertises, how many are held by pods, and
// which pods hold them. It is a read-only client: it lists nodes and pods
// through the API server and does the accounting itself, the same way the
// scheduler counts extended resources.
//
// Run it from code/go, against the cluster in the current kubeconfig context,
// with:
//
//	go run ./gpuusage
package main

import (
	"cmp"
	"context"
	"fmt"
	"os"
	"slices"
	"text/tabwriter"

	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/kubernetes"
	ctrl "sigs.k8s.io/controller-runtime"
)

// gpu is the extended resource that the NVIDIA device plugin advertises.
const gpu corev1.ResourceName = "nvidia.com/gpu"

// holder is a pod that holds GPU units on a node.
type holder struct {
	// pod is the pod's namespace and name, as "namespace/name".
	pod string
	// units is the number of nvidia.com/gpu units the pod holds.
	units int64
}

// podGPUs returns the number of GPU units a pod requests: the sum of the
// nvidia.com/gpu limits of its containers.
func podGPUs(pod *corev1.Pod) int64 {
	var units int64

	for _, c := range pod.Spec.Containers {
		if q, ok := c.Resources.Limits[gpu]; ok {
			units += q.Value()
		}
	}

	return units
}

// holdsResources reports whether a pod currently holds its resources on a
// node. A pod holds them from the moment it is bound to a node until it
// finishes; a pod that has succeeded or failed gives them back.
func holdsResources(pod *corev1.Pod) bool {
	finished := pod.Status.Phase == corev1.PodSucceeded || pod.Status.Phase == corev1.PodFailed

	return pod.Spec.NodeName != "" && !finished
}


// main runs the report and exits with status 1 on any error.
func main() {
	if err := run(context.Background()); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}

// run lists every node and every pod, groups the GPU-holding pods by node,
// and prints one table row per node followed by the pods that hold its units.
func run(ctx context.Context) error {
	clientset, err := kubernetes.NewForConfig(ctrl.GetConfigOrDie())

	if err != nil {
		return err
	}

	nodes, err := clientset.CoreV1().Nodes().List(ctx, metav1.ListOptions{})

	if err != nil {
		return err
	}

	pods, err := clientset.CoreV1().Pods(metav1.NamespaceAll).List(ctx, metav1.ListOptions{})

	if err != nil {
		return err
	}

	holders := map[string][]holder{}

	for i := range pods.Items {
		pod := &pods.Items[i]

		if units := podGPUs(pod); units > 0 && holdsResources(pod) {
			name := pod.Namespace + "/" + pod.Name
			holders[pod.Spec.NodeName] = append(holders[pod.Spec.NodeName], holder{name, units})
		}
	}

	table := tabwriter.NewWriter(os.Stdout, 0, 0, 2, ' ', 0)
	fmt.Fprintln(table, "NODE\tCAPACITY\tALLOCATABLE\tIN USE\tFREE")

	for _, node := range nodes.Items {
		capacity := node.Status.Capacity[gpu]
		allocatable := node.Status.Allocatable[gpu]
		var inUse int64

		for _, h := range holders[node.Name] {
			inUse += h.units
		}

		fmt.Fprintf(table, "%s\t%d\t%d\t%d\t%d\n",
			node.Name, capacity.Value(), allocatable.Value(), inUse, allocatable.Value()-inUse)
	}

	if err := table.Flush(); err != nil {
		return err
	}

	for _, node := range nodes.Items {
		list := holders[node.Name]
		slices.SortFunc(list, func(a, b holder) int { return cmp.Compare(a.pod, b.pod) })

		for _, h := range list {
			fmt.Printf("  %s holds %d on %s\n", h.pod, h.units, node.Name)
		}
	}

	return nil
}

Listing 2.15 `gpuusage`, complete, in Rust. code/rust/gpuusage/src/main.rs

//! Reports, for every node in the cluster, how many `nvidia.com/gpu` units the
//! node advertises, how many are held by pods, and which pods hold them. It is
//! a read-only client: it lists nodes and pods through the API server and does
//! the accounting itself, the same way the scheduler counts extended resources.
//!
//! Run it from `code/rust`, against the cluster in the current kubeconfig
//! context, with `cargo run -p gpuusage`.

use k8s_openapi::api::core::v1::{Node, Pod};
use k8s_openapi::apimachinery::pkg::api::resource::Quantity;
use kube::api::ListParams;
use kube::{Api, Client, ResourceExt};
use std::collections::BTreeMap;
use std::error::Error;

/// The extended resource that the NVIDIA device plugin advertises.
const GPU: &str = "nvidia.com/gpu";

/// A pod that holds GPU units on a node.
struct Holder {
    /// The pod's namespace and name, as `namespace/name`.
    pod: String,
    /// The number of `nvidia.com/gpu` units the pod holds.
    units: i64,
}

/// Returns the number of GPU units a pod requests: the sum of the
/// `nvidia.com/gpu` limits of its containers.
fn pod_gpus(pod: &Pod) -> i64 {
    let Some(spec) = &pod.spec else { return 0 };
    spec.containers
        .iter()
        .filter_map(|c| c.resources.as_ref()?.limits.as_ref()?.get(GPU))
        .filter_map(|q| q.0.parse::<i64>().ok())
        .sum()
}

/// Returns the node a pod currently holds resources on. A pod holds them from
/// the moment it is bound to a node until it finishes; a pod that has
/// succeeded or failed gives them back.
fn holding_node(pod: &Pod) -> Option<&str> {
    let phase = pod.status.as_ref().and_then(|s| s.phase.as_deref());

    if matches!(phase, Some("Succeeded" | "Failed")) {
        return None;
    }

    pod.spec.as_ref()?.node_name.as_deref()
}

/// Reads the number of GPU units from a node's capacity or allocatable map.
fn gpu_units(resources: Option<&BTreeMap<String, Quantity>>) -> i64 {
    resources
        .and_then(|r| r.get(GPU))
        .and_then(|q| q.0.parse().ok())
        .unwrap_or(0)
}

/// Lists every node and every pod, groups the GPU-holding pods by node, and
/// prints one table row per node followed by the pods that hold its units.
#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    let client = Client::try_default().await?;
    let nodes = Api::<Node>::all(client.clone())
        .list(&ListParams::default())
        .await?;
    let pods = Api::<Pod>::all(client).list(&ListParams::default()).await?;

    let mut holders: BTreeMap<String, Vec<Holder>> = BTreeMap::new();

    for pod in &pods {
        let units = pod_gpus(pod);
        if let Some(node) = holding_node(pod)
            && units > 0
        {
            let name = format!("{}/{}", pod.namespace().unwrap_or_default(), pod.name_any());
            holders
                .entry(node.to_string())
                .or_default()
                .push(Holder { pod: name, units });
        }
    }

    println!(
        "{:<23}{:<10}{:<13}{:<8}FREE",
        "NODE", "CAPACITY", "ALLOCATABLE", "IN USE"
    );

    for node in &nodes {
        let name = node.name_any();
        let in_use: i64 = holders
            .get(&name)
            .map_or(0, |h| h.iter().map(|h| h.units).sum());
        let status = node.status.as_ref();
        let capacity = gpu_units(status.and_then(|s| s.capacity.as_ref()));
        let allocatable = gpu_units(status.and_then(|s| s.allocatable.as_ref()));
        println!(
            "{:<23}{:<10}{:<13}{:<8}{}",
            name,
            capacity,
            allocatable,
            in_use,
            allocatable - in_use
        );
    }

    for (node, list) in &mut holders {
        list.sort_by(|a, b| a.pod.cmp(&b.pod));

        for h in list.iter() {
            println!("  {} holds {} on {node}", h.pod, h.units);
        }
    }

    Ok(())
}

References

• NVIDIA Multi-Instance GPU user guide: https://docs.nvidia.com/datacenter/tesla/mig-user-guide/index.html
• NVIDIA GPU Operator, time-slicing GPUs in Kubernetes: https://docs.nvidia.com/datacenter/cloud-native/gpu-operator/latest/gpu-sharing.html
• NVIDIA device plugin for Kubernetes, “Shared Access to GPUs”: https://github.com/NVIDIA/k8s-device-plugin
• NVIDIA Multi-Process Service documentation: https://docs.nvidia.com/deploy/mps/index.html
• NVIDIA vGPU software user guide: https://docs.nvidia.com/vgpu/latest/grid-vgpu-user-guide/index.html
• Kubernetes, Dynamic Resource Allocation: https://kubernetes.io/docs/concepts/resource-management/dynamic-resource-allocation/
• Kubernetes v1.34: DRA has graduated to GA: https://kubernetes.io/blog/2025/09/01/kubernetes-v1-34-dra-updates/
• DRA Driver for NVIDIA GPUs: https://github.com/NVIDIA/k8s-dra-driver-gpu
• NVIDIA Jetson Orin technical specifications: https://www.nvidia.com/en-us/autonomous-machines/embedded-systems/jetson-orin/

3. Knowing which machines are alive

A training job can run on a machine for hours or days. If that machine loses power, or its network cable is pulled, the job stops, and nothing on the machine can report the failure. The control plane has to notice on its own, so that it can restart the job somewhere else.

In this chapter, we will learn how a control plane decides which machines are alive. We will first build intuition by reasoning about heartbeats and timeouts, and quickly prototype a liveness registry and a failure detector from scratch. Then we will use the Lease objects that Kubernetes itself uses to track its nodes: we will write an agent that renews a Lease for a host and a watcher that notices when the renewals stop, and run both on our cluster.

As before, every program is written in Go and in Rust. The two models, heartbeats and phi, use only the standard library. The two Kubernetes programs, leaseagent and leasewatch, use client-go and kube-rs. The Go programs are in code/go, and we run them with go run ./<name>; the Rust programs are packages in the workspace in code/rust, and we run them with cargo run -p <name>.

Heartbeats and timeouts

A machine that has crashed cannot announce that it has crashed. The only evidence the control plane gets is silence. So each machine runs a small agent that sends a message, a heartbeat, at a fixed interval, and the control plane treats a long enough silence as a failure.

The difficulty is that silence has more than one cause. The machine may be dead, or it may be alive but overloaded, or the network between it and the control plane may be broken while the machine keeps running its job. From the control plane’s side, the three look the same. Every failure detector therefore makes a bet on how long a silence has to last before it is more likely to mean death than delay, and the bet has a cost in both directions. A short timeout declares healthy machines dead, and their jobs are restarted for nothing, sometimes while the original copy is still running. A long timeout leaves the jobs of truly dead machines stopped for longer.

Keeping heartbeats out of the database

The control plane has to remember the last heartbeat of every machine. The obvious place for that is its database, and the obvious design writes every heartbeat there. A fleet of 10,000 hosts sending one heartbeat a second would then make 10,000 writes a second. In a replicated store such as etcd, every write is sent to a majority of the replicas and flushed to disk before it is acknowledged, and a steady stream of 10,000 writes a second would leave little capacity for the writes that matter, such as creating jobs.

Most heartbeats carry no news: they say that a host that was alive is still alive. The registry can keep the time of each host’s last heartbeat in memory and write to the database only when a host changes state. Our model uses three states. A host is Up while heartbeats arrive. It becomes Suspect after 3 seconds of silence, which lets the control plane stop placing new work on it without yet giving up on the work it has. It becomes Dead after 10 seconds, at which point its jobs would be restarted elsewhere.

The model simulates 10,000 hosts for 60 seconds, without real time passing. Host 42 crashes at second 20, and host 7 misses four heartbeats, from second 30 to second 33.

The types come first. A host’s record is two fields, and a Transition is the only thing the registry writes to the database.

Listing 3.1 The states, the thresholds, and the registry's records, in Go. code/go/heartbeats/main.go

// State is what the registry believes about a host.
type State int

// The states a host moves through.
const (
	Up      State = iota // heartbeats are arriving
	Suspect              // heartbeats stopped recently; the host may be slow
	Dead                 // heartbeats stopped long enough to give up on the host
)

// String returns the state's name for the log.
func (s State) String() string {
	return [...]string{"Up", "Suspect", "Dead"}[s]
}

// Thresholds and the size of the simulation, in seconds and hosts.
const (
	suspectAfter = 3      // seconds without a heartbeat before a host is Suspect
	deadAfter    = 10     // seconds without a heartbeat before a host is Dead
	hostCount    = 10_000 // hosts in the fleet
	duration     = 60     // seconds simulated
)

// entry is the registry's in-memory record of one host.
type entry struct {
	// lastSeen is the second of the host's most recent heartbeat.
	lastSeen int
	// state is the host's current state.
	state State
}

// Transition is one durable write: a host entering a new state.
type Transition struct {
	// At is the second at which the sweep or heartbeat noticed the change.
	At int
	// Host is the host's number.
	Host int
	// State is the state the host entered.
	State State
}

// Registry holds the liveness of every host.
type Registry struct {
	// hosts is indexed by host number.
	hosts []entry
	// heartbeats counts the heartbeats handled in memory.
	heartbeats int
	// writes lists the durable writes, in order.
	writes []Transition
}

Listing 3.2 The states, the thresholds, and the registry's records, in Rust. code/rust/heartbeats/src/main.rs

/// Seconds without a heartbeat before a host is Suspect.
const SUSPECT_AFTER: u32 = 3;
/// Seconds without a heartbeat before a host is Dead.
const DEAD_AFTER: u32 = 10;
/// Hosts in the fleet.
const HOST_COUNT: usize = 10_000;
/// Seconds simulated.
const DURATION: u32 = 60;

/// What the registry believes about a host.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
enum State {
    /// Heartbeats are arriving.
    #[default]
    Up,
    /// Heartbeats stopped recently; the host may be slow.
    Suspect,
    /// Heartbeats stopped long enough to give up on the host.
    Dead,
}

/// The registry's in-memory record of one host.
#[derive(Clone, Copy, Default)]
struct Entry {
    /// The second of the host's most recent heartbeat.
    last_seen: u32,
    /// The host's current state.
    state: State,
}

/// One durable write: a host entering a new state.
struct Transition {
    /// The second at which the sweep or heartbeat noticed the change.
    at: u32,
    /// The host's number.
    host: usize,
    /// The state the host entered.
    state: State,
}

/// Holds the liveness of every host.
struct Registry {
    /// Indexed by host number.
    hosts: Vec<Entry>,
    /// Counts the heartbeats handled in memory.
    heartbeats: usize,
    /// The durable writes, in order.
    writes: Vec<Transition>,
}

The registry has one operation for an arriving heartbeat and one for the sweep, and a small function that turns a length of silence into a state.

Listing 3.3 A heartbeat, the sweep, and the rule from silence to state, in Go. code/go/heartbeats/main.go

// NewRegistry returns a registry in which every host is Up, as it would be
// after the hosts registered.
func NewRegistry(n int) *Registry {
	return &Registry{hosts: make([]entry, n)}
}

// Heartbeat records a heartbeat from host at second now. Only a host that was
// not Up costs a durable write.
func (r *Registry) Heartbeat(now, host int) {
	r.heartbeats++
	e := &r.hosts[host]
	e.lastSeen = now

	if e.state != Up {
		e.state = Up
		r.writes = append(r.writes, Transition{now, host, Up})
	}
}

// Sweep re-evaluates every host at second now and records each change of state.
func (r *Registry) Sweep(now int) {
	for host := range r.hosts {
		e := &r.hosts[host]

		if next := stateAfter(now - e.lastSeen); next != e.state {
			e.state = next
			r.writes = append(r.writes, Transition{now, host, next})
		}
	}
}

// stateAfter returns the state of a host that has been silent for the given
// number of seconds.
func stateAfter(silent int) State {
	switch {
	case silent >= deadAfter:
		return Dead
	case silent >= suspectAfter:
		return Suspect
	default:
		return Up
	}
}

Listing 3.4 A heartbeat, the sweep, and the rule from silence to state, in Rust. code/rust/heartbeats/src/main.rs

impl Registry {
    /// Returns a registry in which every host is Up, as it would be after the
    /// hosts registered.
    fn new(n: usize) -> Self {
        Self {
            hosts: vec![Entry::default(); n],
            heartbeats: 0,
            writes: Vec::new(),
        }
    }

    /// Records a heartbeat from `host` at second `now`. Only a host that was
    /// not Up costs a durable write.
    fn heartbeat(&mut self, now: u32, host: usize) {
        self.heartbeats += 1;
        let entry = &mut self.hosts[host];
        entry.last_seen = now;

        if entry.state != State::Up {
            entry.state = State::Up;
            self.writes.push(Transition {
                at: now,
                host,
                state: State::Up,
            });
        }
    }

    /// Re-evaluates every host at second `now` and records each change of state.
    fn sweep(&mut self, now: u32) {
        for (host, entry) in self.hosts.iter_mut().enumerate() {
            let next = state_after(now - entry.last_seen);

            if next != entry.state {
                entry.state = next;
                self.writes.push(Transition {
                    at: now,
                    host,
                    state: next,
                });
            }
        }
    }
}

/// Returns the state of a host that has been silent for `silent` seconds.
fn state_after(silent: u32) -> State {
    match silent {
        s if s >= DEAD_AFTER => State::Dead,
        s if s >= SUSPECT_AFTER => State::Suspect,
        _ => State::Up,
    }
}

The run feeds the registry one heartbeat from every host each second, except from the two hosts that the scenario silences.

Listing 3.5 The simulated fleet and the main loop, in Go. code/go/heartbeats/main.go

// sendsHeartbeat reports whether host sends a heartbeat at second now in the
// simulation.
func sendsHeartbeat(host, now int) bool {
	switch host {
	case 42:
		return now < 20 // crashes at second 20
	case 7:
		return now < 30 || now >= 34 // silent from second 30 to 33
	default:
		return true
	}
}

// main runs the simulation and prints the registry's work.
func main() {
	registry := NewRegistry(hostCount)

	for now := 0; now < duration; now++ {
		for host := 0; host < hostCount; host++ {
			if sendsHeartbeat(host, now) {
				registry.Heartbeat(now, host)
			}
		}

		registry.Sweep(now)
	}

	fmt.Printf("heartbeats handled in memory: %d\n", registry.heartbeats)
	fmt.Printf("durable writes:               %d\n", len(registry.writes))

	for _, w := range registry.writes {
		fmt.Printf("  second %2d  host %2d -> %s\n", w.At, w.Host, w.State)
	}
}

Listing 3.6 The simulated fleet and the main loop, in Rust. code/rust/heartbeats/src/main.rs

/// Reports whether `host` sends a heartbeat at second `now` in the simulation.
fn sends_heartbeat(host: usize, now: u32) -> bool {
    match host {
        42 => now < 20,                // crashes at second 20
        7 => !(30..34).contains(&now), // silent from second 30 to 33
        _ => true,
    }
}

/// Runs the simulation and prints the registry's work.
fn main() {
    let mut registry = Registry::new(HOST_COUNT);

    for now in 0..DURATION {
        for host in 0..HOST_COUNT {
            if sends_heartbeat(host, now) {
                registry.heartbeat(now, host);
            }
        }

        registry.sweep(now);
    }

    println!("heartbeats handled in memory: {}", registry.heartbeats);
    println!("durable writes:               {}", registry.writes.len());

    for w in &registry.writes {
        println!("  second {:2}  host {:2} -> {:?}", w.at, w.host, w.state);
    }
}

Heartbeat updates the host’s last-seen time in memory and writes to the database only when the host was not Up. Sweep runs once a second, computes each host’s state from the length of its silence, and writes only the hosts whose state changed. A heartbeat from a Suspect or Dead host brings it back at once; the sweep is what moves hosts the other way.

Both versions print the same result:

$ go run ./heartbeats
heartbeats handled in memory: 599956
durable writes:               4
  second 22  host 42 -> Suspect
  second 29  host 42 -> Dead
  second 32  host  7 -> Suspect
  second 34  host  7 -> Up

The registry handled almost 600,000 heartbeats and made four writes. Host 42 sent its last heartbeat at second 19; it became Suspect at second 22 and Dead at second 29. Host 7 became Suspect at second 32, three seconds after its last heartbeat at second 29, and came back Up with its heartbeat at second 34, before it reached the Dead threshold.

Figure 3.1 shows the two paths, and the three states with the silences that separate them.

Keeping the state in memory has one consequence. When the registry itself restarts, it has lost every host’s last-seen time. Our model starts by marking every host Up with a last-seen time of 0, which gives each host a full timeout to report in before anything is declared dead. A real registry has to do the same after a restart, or it would declare the whole fleet dead at once.

Choosing the timeout: the phi accrual detector

A fixed timeout treats every host alike. Hosts are not alike: one on a quiet network sends heartbeats exactly a second apart, and another on a congested link sends them anywhere from half a second to two and a half seconds apart. Two seconds of silence is alarming for the first host and ordinary for the second.

The phi accrual failure detector, described by Hayashibara, Défago, Yared, and Katayama in 2004 and used today in Apache Cassandra and Akka, makes the timeout depend on each host’s history. It keeps the most recent gaps between a host’s heartbeats, treats them as samples of a normal distribution, and computes the probability that a heartbeat would still arrive after the current silence. It reports that probability as phi, minus its base-10 logarithm, so that each unit of phi is a factor of ten: phi 1 means a 10 percent chance that the host is fine, phi 2 means 1 percent, phi 3 means 0.1 percent. The control plane picks a phi threshold instead of a timeout, and each host’s effective timeout follows from its own behaviour.

A detector belongs to one host. It remembers a window of gaps and the time of the latest heartbeat.

Listing 3.7 The detector and the recording of a heartbeat, in Go. code/go/phi/main.go

// timeout is the fixed timeout that the program compares against, in seconds.
const timeout = 2.0

// minStdDev is the smallest standard deviation the detector assumes, in
// seconds. Without it, a perfectly regular host would make any delay look
// infinitely suspicious.
const minStdDev = 0.2

// Detector is a phi accrual failure detector for one host.
type Detector struct {
	// intervals holds the most recent gaps between heartbeats, in seconds.
	intervals []float64
	// window is the number of gaps the detector remembers.
	window int
	// last is the time of the most recent heartbeat, in seconds.
	last float64
}

// NewDetector returns a detector that remembers window gaps and has seen a
// first heartbeat at time first.
func NewDetector(window int, first float64) *Detector {
	return &Detector{window: window, last: first}
}

// Heartbeat records a heartbeat that arrived at time now.
func (d *Detector) Heartbeat(now float64) {
	d.intervals = append(d.intervals, now-d.last)

	if len(d.intervals) > d.window {
		d.intervals = d.intervals[1:]
	}

	d.last = now
}

Listing 3.8 The detector and the recording of a heartbeat, in Rust. code/rust/phi/src/main.rs

/// The fixed timeout that the program compares against, in seconds.
const TIMEOUT: f64 = 2.0;

/// The smallest standard deviation the detector assumes, in seconds. Without
/// it, a perfectly regular host would make any delay look infinitely
/// suspicious.
const MIN_STD_DEV: f64 = 0.2;

/// A phi accrual failure detector for one host.
struct Detector {
    /// The most recent gaps between heartbeats, in seconds.
    intervals: VecDeque<f64>,
    /// The number of gaps the detector remembers.
    window: usize,
    /// The time of the most recent heartbeat, in seconds.
    last: f64,
}

impl Detector {
    /// Returns a detector that remembers `window` gaps and has seen a first
    /// heartbeat at time `first`.
    fn new(window: usize, first: f64) -> Self {
        Self {
            intervals: VecDeque::new(),
            window,
            last: first,
        }
    }

    /// Records a heartbeat that arrived at time `now`.
    fn heartbeat(&mut self, now: f64) {
        self.intervals.push_back(now - self.last);

        if self.intervals.len() > self.window {
            self.intervals.pop_front();
        }

        self.last = now;
    }
}

Phi turns the window into a number.

Listing 3.9 Computing phi from the recorded gaps, in Go. code/go/phi/main.go

// Phi returns the suspicion level at time now: minus the base-10 logarithm of
// the probability that a heartbeat would still arrive after this much silence,
// given the gaps seen so far. Phi 1 means a 10% chance the host is fine, phi 2
// means 1%, and phi 3 means 0.1%.
func (d *Detector) Phi(now float64) float64 {
	var mean float64

	for _, gap := range d.intervals {
		mean += gap
	}

	mean /= float64(len(d.intervals))
	var variance float64

	for _, gap := range d.intervals {
		variance += (gap - mean) * (gap - mean)
	}

	stdDev := math.Max(math.Sqrt(variance/float64(len(d.intervals))), minStdDev)

	// The probability that a gap is longer than the current silence, from a
	// logistic approximation of the normal distribution.
	silence := now - d.last
	y := (silence - mean) / stdDev
	e := math.Exp(-y * (1.5976 + 0.070566*y*y))
	pLater := 1 - 1/(1+e)

	if silence > mean {
		pLater = e / (1 + e)
	}

	return -math.Log10(pLater)
}

Listing 3.10 Computing phi from the recorded gaps, in Rust. code/rust/phi/src/main.rs

impl Detector {
    /// Returns the suspicion level at time `now`: minus the base-10 logarithm
    /// of the probability that a heartbeat would still arrive after this much
    /// silence, given the gaps seen so far. Phi 1 means a 10% chance the host
    /// is fine, phi 2 means 1%, and phi 3 means 0.1%.
    #[allow(clippy::cast_precision_loss)] // the window holds at most 100 gaps
    fn phi(&self, now: f64) -> f64 {
        let n = self.intervals.len() as f64;
        let mean = self.intervals.iter().sum::<f64>() / n;
        let variance = self
            .intervals
            .iter()
            .map(|gap| (gap - mean).powi(2))
            .sum::<f64>()
            / n;
        let std_dev = variance.sqrt().max(MIN_STD_DEV);

        // The probability that a gap is longer than the current silence, from
        // a logistic approximation of the normal distribution.
        let silence = now - self.last;
        let y = (silence - mean) / std_dev;
        let e = (-y * (1.5976 + 0.070_566 * y * y)).exp();
        let p_later = if silence > mean {
            e / (1.0 + e)
        } else {
            1.0 - 1.0 / (1.0 + e)
        };

        -p_later.log10()
    }
}

The run trains one detector on a steady host and one on a jittery host, then asks each for phi after silences of increasing length.

Listing 3.11 Two hosts and the table of results, in Go. code/go/phi/main.go

// hosts lists the simulated hosts and the gaps between their heartbeats.
var hosts = []struct {
	name string
	gaps []float64
}{
	{"steady", []float64{1.0, 1.02, 0.98, 1.01, 0.99, 1.0, 1.03, 0.97, 1.0, 1.0}},
	{"jittery", []float64{0.4, 1.9, 0.6, 1.7, 0.5, 1.8, 0.7, 1.6, 0.4, 2.4}},
}

// main feeds each host's heartbeats to a detector and prints phi after
// increasing lengths of silence.
func main() {
	for _, host := range hosts {
		detector := NewDetector(100, 0)
		var now float64

		for _, gap := range host.gaps {
			now += gap
			detector.Heartbeat(now)
		}

		fmt.Printf("%s host, silent for:\n", host.name)

		for _, silence := range []float64{1.0, 1.5, 2.0, 2.5, 3.0, 4.0} {
			verdict := "alive"

			if silence > timeout {
				verdict = "dead"
			}

			fmt.Printf("  %.1fs  phi = %6.2f   %.0fs timeout says %s\n",
				silence, detector.Phi(now+silence), timeout, verdict)
		}
	}
}

Listing 3.12 Two hosts and the table of results, in Rust. code/rust/phi/src/main.rs

/// The simulated hosts and the gaps between their heartbeats.
const HOSTS: [(&str, [f64; 10]); 2] = [
    (
        "steady",
        [1.0, 1.02, 0.98, 1.01, 0.99, 1.0, 1.03, 0.97, 1.0, 1.0],
    ),
    (
        "jittery",
        [0.4, 1.9, 0.6, 1.7, 0.5, 1.8, 0.7, 1.6, 0.4, 2.4],
    ),
];

/// Feeds each host's heartbeats to a detector and prints phi after increasing
/// lengths of silence.
fn main() {
    for (name, gaps) in HOSTS {
        let mut detector = Detector::new(100, 0.0);
        let mut now = 0.0;

        for gap in gaps {
            now += gap;
            detector.heartbeat(now);
        }

        println!("{name} host, silent for:");

        for silence in [1.0, 1.5, 2.0, 2.5, 3.0, 4.0] {
            let verdict = if silence > TIMEOUT { "dead" } else { "alive" };
            println!(
                "  {silence:.1}s  phi = {:6.2}   {TIMEOUT:.0}s timeout says {verdict}",
                detector.phi(now + silence)
            );
        }
    }
}

Phi computes the mean and standard deviation of the remembered gaps and measures how many standard deviations the current silence lies beyond the mean. The logistic expression is a standard approximation of the normal distribution’s tail, the same one Akka uses, and it avoids computing the error function. The standard deviation has a floor of 0.2 seconds. Without it, a host whose heartbeats were perfectly regular would have a standard deviation near zero, and the smallest delay would push phi towards infinity.

Both versions print the same table:

$ go run ./phi
steady host, silent for:
  1.0s  phi =   0.30   2s timeout says alive
  1.5s  phi =   2.22   2s timeout says alive
  2.0s  phi =   7.30   2s timeout says alive
  2.5s  phi =  18.13   2s timeout says dead
  3.0s  phi =  37.58   2s timeout says dead
  4.0s  phi = 113.84   2s timeout says dead
jittery host, silent for:
  1.0s  phi =   0.21   2s timeout says alive
  1.5s  phi =   0.47   2s timeout says alive
  2.0s  phi =   0.88   2s timeout says alive
  2.5s  phi =   1.47   2s timeout says dead
  3.0s  phi =   2.25   2s timeout says dead
  4.0s  phi =   4.58   2s timeout says dead

For the steady host, two seconds of silence already gives phi 7.3: a heartbeat that late is roughly a one-in-twenty-million event, and the host is almost certainly gone, while the 2-second timeout still calls it alive. For the jittery host, two and a half seconds gives phi 1.47, about a 3 percent chance that the host is fine, and the timeout has already declared it dead. With a threshold of 3, the detector would give up on the steady host after about 1.6 seconds and on the jittery host only after about 3.4 seconds.

Leases: how Kubernetes tracks its nodes

Kubernetes applies the same ideas to its own nodes. Every node has a Lease, a small object in the coordination.k8s.io API group, in the namespace kube-node-lease. The kubelet renews its node’s Lease every 10 seconds by writing the current time into the Lease’s renewTime field. Our cluster has one node and one node Lease:

$ kubectl get lease -n kube-node-lease
NAME                    HOLDER                  AGE
gpu-lab-control-plane   gpu-lab-control-plane   4d11h
$ kubectl get lease -n kube-node-lease gpu-lab-control-plane -o yaml
...
spec:
  holderIdentity: gpu-lab-control-plane
  leaseDurationSeconds: 40
  renewTime: "2026-10-01T01:31:36.294522Z"

A Lease has five main fields. holderIdentity names whoever holds it, here the node itself. leaseDurationSeconds says how long the Lease stays valid after its last renewal. renewTime is the time of the last renewal, acquireTime the time the holder first took the Lease, and leaseTransitions counts how often the Lease has changed hands. Kubernetes also uses Leases for leader election, so that only one copy of the controller manager or the scheduler is active at a time, and for each API server to publish its identity.

The kubelet used to send its heartbeats by updating the Node object’s status, which holds the node’s capacity, conditions, images, and addresses. The Kubernetes documentation explains the change: “a Lease is a lightweight resource. Using Leases for heartbeats reduces the performance impact of these updates”. The kubelet still updates the Node’s status, but only when something changes or every five minutes.

On the other side, the node lifecycle controller, part of the controller manager, checks every node every 5 seconds. If it has seen no renewal for longer than its grace period, 50 seconds by default, it sets the node’s Ready condition to Unknown. By default, the pods on that node are evicted 5 minutes later, so that their controllers can recreate them elsewhere. The long wait before eviction reflects the cost of a wrong guess: a node that only lost contact with the control plane may still be running its pods.

A lease agent

We will give each simulated GPU host a Lease of its own, in a namespace called fleet:

Listing 3.13 The namespace for the fleet's Leases. code/manifests/fleet-namespace.yaml

# The namespace that holds one Lease per simulated GPU host in chapter 3.
apiVersion: v1
kind: Namespace
metadata:
  name: fleet

The agent plays the part of the kubelet. It writes the Lease named after its host every two seconds, with a lease duration of six seconds, so a host has to miss three renewals in a row before its Lease expires.

The settings are constants.

Listing 3.14 The agent's settings, in Go. code/go/leaseagent/main.go

// Lease settings. The renewal interval is a third of the lease duration, so a
// host has to miss three renewals in a row before its Lease expires.
const (
	namespace     = "fleet"
	renewEvery    = 2 * time.Second
	leaseDuration = 6 // seconds
	fieldManager  = "leaseagent"
)

Listing 3.15 The agent's settings, in Rust. code/rust/leaseagent/src/main.rs

/// The namespace that holds the hosts' Leases.
const NAMESPACE: &str = "fleet";
/// How often the agent renews the Lease: a third of the lease duration, so a
/// host has to miss three renewals in a row before its Lease expires.
const RENEW_EVERY: Duration = Duration::from_secs(2);
/// How long a Lease stays valid after its last renewal, in seconds.
const LEASE_DURATION: i32 = 6;
/// The field manager name for server-side apply.
const FIELD_MANAGER: &str = "leaseagent";

The Go version builds the Lease inside its loop with client-go’s apply configuration types. The Rust version builds it in a function of its own.

Listing 3.16 The renewal loop, in Go. code/go/leaseagent/main.go

// run writes the host's Lease with server-side apply once per renewEvery.
// A renewal that fails is reported and retried at the next tick.
func run(host string) error {
	clientset, err := kubernetes.NewForConfig(ctrl.GetConfigOrDie())

	if err != nil {
		return err
	}

	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()

	leases := clientset.CoordinationV1().Leases(namespace)
	acquired := metav1.NewMicroTime(time.Now())
	ticker := time.NewTicker(renewEvery)
	defer ticker.Stop()

	for {
		now := metav1.NewMicroTime(time.Now())
		lease := coordinationv1ac.Lease(host, namespace).
			WithSpec(coordinationv1ac.LeaseSpec().
				WithHolderIdentity(host).
				WithLeaseDurationSeconds(leaseDuration).
				WithAcquireTime(acquired).
				WithRenewTime(now))
		_, err := leases.Apply(ctx, lease, metav1.ApplyOptions{FieldManager: fieldManager, Force: true})

		switch {
		case ctx.Err() != nil:
			return nil
		case err != nil:
			fmt.Printf("%s %s renewal failed: %v\n", now.Format(time.TimeOnly), host, err)
		default:
			fmt.Printf("%s %s renewed\n", now.Format(time.TimeOnly), host)
		}

		select {
		case <-ctx.Done():
			return nil
		case <-ticker.C:
		}
	}
}

Listing 3.17 Building the Lease, in Rust. code/rust/leaseagent/src/main.rs

/// Builds the Lease that `host` writes at time `now`, having first acquired
/// it at time `acquired`.
fn lease(host: &str, acquired: Timestamp, now: Timestamp) -> Lease {
    Lease {
        metadata: ObjectMeta {
            name: Some(host.to_string()),
            namespace: Some(NAMESPACE.to_string()),
            ..ObjectMeta::default()
        },
        spec: Some(LeaseSpec {
            holder_identity: Some(host.to_string()),
            lease_duration_seconds: Some(LEASE_DURATION),
            acquire_time: Some(MicroTime(acquired)),
            renew_time: Some(MicroTime(now)),
            ..LeaseSpec::default()
        }),
    }
}

Listing 3.18 The renewal loop, in Rust. code/rust/leaseagent/src/main.rs

/// Reads the host name from the command line and renews its Lease with
/// server-side apply until the process receives Ctrl-C. A renewal that fails
/// is reported and retried at the next tick.
#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    let Some(host) = std::env::args().nth(1) else {
        eprintln!("usage: leaseagent <host>");
        std::process::exit(2);
    };
    let leases: Api<Lease> = Api::namespaced(Client::try_default().await?, NAMESPACE);
    let params = PatchParams::apply(FIELD_MANAGER).force();
    let acquired = Timestamp::now();
    let mut ticker = tokio::time::interval(RENEW_EVERY);

    loop {
        tokio::select! {
            _ = tokio::signal::ctrl_c() => return Ok(()),
            _ = ticker.tick() => {}
        }
        let now = Timestamp::now();
        match leases
            .patch(&host, &params, &Patch::Apply(lease(&host, acquired, now)))
            .await
        {
            Ok(_) => println!("{} {host} renewed", time_of_day(now)),
            Err(error) => println!("{} {host} renewal failed: {error}", time_of_day(now)),
        }
    }
}

Both agents write with server-side apply, which we met in section 1.1: each renewal sends the whole Lease, and the API server creates it on the first renewal and updates it on the others. A failed renewal, for example while the API server restarts, is logged and retried at the next tick; the Lease simply ages until a renewal succeeds. The renewal time is written in UTC, as every Kubernetes timestamp is, and it comes from the host’s own clock.

A lease watcher

The watcher plays the part of the node lifecycle controller. It has to decide when a Lease has expired, and the obvious rule, “expired when renewTime plus the lease duration is in the past”, has a flaw: renewTime was written with the host’s clock and is compared with the watcher’s clock. If the two clocks disagree by a few seconds, healthy hosts look expired, or dead hosts look alive for longer than they should.

The node lifecycle controller avoids this by never comparing the two clocks. A comment in its source explains that it records its own time of observation “to avoid the problem with time skew across the cluster”: whenever it sees a Lease’s renewTime change, it notes its own current time, and it declares the node unhealthy when its own clock shows that the grace period has passed since that note. The value of renewTime matters only in that it changed. Our watcher works the same way.

An observation is three fields, and the second is the one that matters: the local time at which the watcher first saw the current renewTime.

Listing 3.19 What the watcher remembers about one Lease, in Go. code/go/leasewatch/main.go

// observation is what the program knows about one host's Lease.
type observation struct {
	// renewTime is the latest renewTime seen in the Lease.
	renewTime time.Time
	// observedAt is the local time at which that renewTime was first seen.
	observedAt time.Time
	// expired records whether the host has been reported as expired.
	expired bool
}

Listing 3.20 What the watcher remembers about one Lease, in Rust. code/rust/leasewatch/src/main.rs

/// What the program knows about one host's Lease.
struct Observation {
    /// The latest `renewTime` seen in the Lease.
    renew_time: Timestamp,
    /// The local time at which that `renewTime` was first seen.
    observed_at: Instant,
    /// Whether the host has been reported as expired.
    expired: bool,
}

check holds the rule from the node lifecycle controller.

Listing 3.21 Checking one Lease against the watcher's own clock, in Go. code/go/leasewatch/main.go

// check updates the observation of one Lease at local time now and prints a
// line when the host is first seen, expires, or renews again after expiring.
func check(seen map[string]*observation, lease *coordinationv1.Lease, now time.Time) {
	if lease.Spec.RenewTime == nil || lease.Spec.LeaseDurationSeconds == nil {
		return
	}

	renew := lease.Spec.RenewTime.Time
	duration := time.Duration(*lease.Spec.LeaseDurationSeconds) * time.Second
	stamp := now.Format(time.TimeOnly)

	obs, known := seen[lease.Name]

	switch {
	case !known:
		seen[lease.Name] = &observation{renewTime: renew, observedAt: now}
		fmt.Printf("%s %s seen, lease duration %s\n", stamp, lease.Name, duration)
	case !renew.Equal(obs.renewTime):
		obs.renewTime, obs.observedAt = renew, now

		if obs.expired {
			obs.expired = false
			fmt.Printf("%s %s renewing again\n", stamp, lease.Name)
		}
	case !obs.expired && now.Sub(obs.observedAt) > duration:
		obs.expired = true
		fmt.Printf("%s %s expired: no renewal seen for %s\n",
			stamp, lease.Name, now.Sub(obs.observedAt).Round(time.Second))
	}
}

Listing 3.22 Checking one Lease against the watcher's own clock, in Rust. code/rust/leasewatch/src/main.rs

/// Updates the observation of one Lease at local time `now` and prints a line
/// when the host is first seen, expires, or renews again after expiring.
fn check(seen: &mut HashMap<String, Observation>, lease: &Lease, now: Instant) {
    let Some(spec) = &lease.spec else { return };
    let (Some(renew), Some(seconds)) = (&spec.renew_time, spec.lease_duration_seconds) else {
        return;
    };
    let duration = Duration::from_secs(u64::from(seconds.unsigned_abs()));
    let stamp = Zoned::now().strftime("%H:%M:%S").to_string();
    let name = lease.name_any();

    match seen.get_mut(&name) {
        None => {
            seen.insert(
                name.clone(),
                Observation {
                    renew_time: renew.0,
                    observed_at: now,
                    expired: false,
                },
            );
            println!(
                "{stamp} {name} seen, lease duration {}s",
                duration.as_secs()
            );
        }
        Some(obs) if obs.renew_time != renew.0 => {
            obs.renew_time = renew.0;
            obs.observed_at = now;

            if obs.expired {
                obs.expired = false;
                println!("{stamp} {name} renewing again");
            }
        }
        Some(obs) if !obs.expired && now.duration_since(obs.observed_at) > duration => {
            obs.expired = true;
            let silent = now.duration_since(obs.observed_at).as_secs_f64().round();
            println!("{stamp} {name} expired: no renewal seen for {silent}s");
        }
        Some(_) => {}
    }
}

The rest of the program keeps a cache of Leases and calls check for each of them once a second.

Listing 3.23 The cache of Leases and the once-a-second loop, in Go. code/go/leasewatch/main.go

// run starts an informer for Leases in the namespace, waits for its first
// list, and then checks every Lease in its cache once per checkEvery.
func run() error {
	clientset, err := kubernetes.NewForConfig(ctrl.GetConfigOrDie())

	if err != nil {
		return err
	}

	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()

	factory := informers.NewSharedInformerFactoryWithOptions(clientset, 0, informers.WithNamespace(namespace))
	lister := factory.Coordination().V1().Leases().Lister().Leases(namespace)
	factory.Start(ctx.Done())
	factory.WaitForCacheSync(ctx.Done())

	seen := map[string]*observation{}
	ticker := time.NewTicker(checkEvery)
	defer ticker.Stop()

	for {
		leases, err := lister.List(labels.Everything())

		if err != nil {
			return err
		}

		now := time.Now()

		for _, lease := range leases {
			check(seen, lease, now)
		}

		select {
		case <-ctx.Done():
			return nil
		case <-ticker.C:
		}
	}
}

Listing 3.24 The cache of Leases and the once-a-second loop, in Rust. code/rust/leasewatch/src/main.rs

/// Starts a watcher that fills a reflector store with the namespace's Leases,
/// waits for the first list, and then checks every Lease in the store once
/// per `CHECK_EVERY` until the process receives Ctrl-C.
#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    let leases: Api<Lease> = Api::namespaced(Client::try_default().await?, NAMESPACE);
    let (store, writer) = reflector::store();
    let stream = watcher(leases, watcher::Config::default())
        .default_backoff()
        .reflect(writer)
        .applied_objects();
    tokio::spawn(stream.for_each(|_| std::future::ready(())));
    store.wait_until_ready().await?;

    let mut seen = HashMap::new();
    let mut ticker = tokio::time::interval(CHECK_EVERY);

    loop {
        tokio::select! {
            _ = tokio::signal::ctrl_c() => return Ok(()),
            _ = ticker.tick() => {}
        }
        let now = Instant::now();

        for lease in store.state() {
            check(&mut seen, &lease, now);
        }
    }
}

Neither version sends a request per check. The Go version starts an informer for Leases in the fleet namespace from client-go’s shared informer factory and reads its cache through a lister. The Rust version runs a watcher stream into a reflector store in a background task, with default_backoff so that the stream reconnects after errors, and reads the store’s contents. Both are the list-and-watch machinery from chapter 1, and both read the cache once a second.

check keeps one observation per Lease and moves through three cases. A Lease seen for the first time is recorded with the current local time. A Lease whose renewTime has changed gets a new local observation time, and if the host had been reported as expired, the watcher reports that it is renewing again. A Lease whose renewTime has not changed for longer than its duration, measured on the watcher’s clock, is reported as expired, once. The Rust version measures with Instant, a monotonic clock that is unaffected if the system’s wall clock is adjusted.

Running it

We create the namespace, start three agents, two in Go and one in Rust, and then start both watchers. Each command runs in its own terminal:

$ kubectl apply -f code/manifests/fleet-namespace.yaml
$ go run ./leaseagent gpu-host-a          # in code/go
$ cargo run -p leaseagent -- gpu-host-b   # in code/rust
$ go run ./leaseagent gpu-host-c          # in code/go
$ go run ./leasewatch                     # in code/go
$ cargo run -p leasewatch                 # in code/rust

Figure 3.2 shows the five programs and the Lease objects between them.

The Leases appear in the cluster like any other object:

$ kubectl get lease -n fleet
NAME         HOLDER       AGE
gpu-host-a   gpu-host-a   29s
gpu-host-b   gpu-host-b   29s
gpu-host-c   gpu-host-c   29s
$ kubectl get lease -n fleet gpu-host-a -o yaml
...
spec:
  acquireTime: "2026-10-01T01:31:16.509475Z"
  holderIdentity: gpu-host-a
  leaseDurationSeconds: 6
  renewTime: "2026-10-01T01:31:44.511013Z"

At 18:31:25 we kill the agent for gpu-host-b with kill -9, which stops it without giving it a chance to clean up, as a power failure would. Fourteen seconds later, at 18:31:39, we start it again. The agent’s own log shows the gap:

18:31:16 gpu-host-b renewed
18:31:18 gpu-host-b renewed
18:31:20 gpu-host-b renewed
18:31:22 gpu-host-b renewed
18:31:24 gpu-host-b renewed
18:31:39 gpu-host-b renewed
18:31:41 gpu-host-b renewed
18:31:43 gpu-host-b renewed
18:31:45 gpu-host-b renewed

The Go watcher printed:

18:31:19 gpu-host-c seen, lease duration 6s
18:31:19 gpu-host-a seen, lease duration 6s
18:31:19 gpu-host-b seen, lease duration 6s
18:31:31 gpu-host-b expired: no renewal seen for 7s
18:31:39 gpu-host-b renewing again

And the Rust watcher printed:

18:31:19 gpu-host-c seen, lease duration 6s
18:31:19 gpu-host-a seen, lease duration 6s
18:31:19 gpu-host-b seen, lease duration 6s
18:31:30 gpu-host-b expired: no renewal seen for 6s
18:31:40 gpu-host-b renewing again

The last renewal of gpu-host-b happened at 18:31:24. Each watcher sees a renewal only at its next once-a-second check, so its observation time is up to a second later than the renewal itself, and it reports the expiry at the first check that falls more than six seconds after that observation. The two watchers check at slightly different moments, so one reported the expiry at 18:31:30 and the other at 18:31:31. The same one-second granularity explains the second difference: the restarted agent renewed at once, at 18:31:39, and the two watchers noticed at 18:31:39 and 18:31:40.

Animation 3.1 replays the run for gpu-host-b as the Go watcher saw it. Its last part gives the host a clock that runs eight seconds behind and a watcher that subtracts renewTime from its own clock.

Throughout, the Go watchers read a Lease written by the Rust agent and the Rust watcher read Leases written by the Go agents. The Lease in the API server is the contract between them, and neither side needs to know what language the other is written in.

Declaring a host dead is where a control plane’s harder decisions begin. The host’s jobs should run again elsewhere, but a host that only lost its network connection may still be running them, and a GPU job that writes checkpoints or holds a reservation must not run twice at once. Kubernetes answers part of this with the 5-minute delay before eviction. A control plane for GPU jobs needs a firmer answer as well.

The complete programs

Listing 3.25 The complete program in Go. code/go/heartbeats/main.go

// Command heartbeats models how a control plane tracks which machines in a
// fleet are alive. Every host sends a heartbeat once a second. The registry
// keeps each host's last heartbeat in memory, runs a sweep once a second, and
// records a durable write only when a host changes state: Up, Suspect, or Dead.
//
// The program simulates 10,000 hosts for 60 seconds without real time passing.
// Host 42 crashes at second 20, and host 7 misses four heartbeats from second
// 30 to second 33. It prints how many heartbeats the registry handled, how many
// durable writes it made, and each write.
//
// Run it from code/go with:
//
//	go run ./heartbeats
package main

import "fmt"

// State is what the registry believes about a host.
type State int

// The states a host moves through.
const (
	Up      State = iota // heartbeats are arriving
	Suspect              // heartbeats stopped recently; the host may be slow
	Dead                 // heartbeats stopped long enough to give up on the host
)

// String returns the state's name for the log.
func (s State) String() string {
	return [...]string{"Up", "Suspect", "Dead"}[s]
}

// Thresholds and the size of the simulation, in seconds and hosts.
const (
	suspectAfter = 3      // seconds without a heartbeat before a host is Suspect
	deadAfter    = 10     // seconds without a heartbeat before a host is Dead
	hostCount    = 10_000 // hosts in the fleet
	duration     = 60     // seconds simulated
)

// entry is the registry's in-memory record of one host.
type entry struct {
	// lastSeen is the second of the host's most recent heartbeat.
	lastSeen int
	// state is the host's current state.
	state State
}

// Transition is one durable write: a host entering a new state.
type Transition struct {
	// At is the second at which the sweep or heartbeat noticed the change.
	At int
	// Host is the host's number.
	Host int
	// State is the state the host entered.
	State State
}

// Registry holds the liveness of every host.
type Registry struct {
	// hosts is indexed by host number.
	hosts []entry
	// heartbeats counts the heartbeats handled in memory.
	heartbeats int
	// writes lists the durable writes, in order.
	writes []Transition
}


// NewRegistry returns a registry in which every host is Up, as it would be
// after the hosts registered.
func NewRegistry(n int) *Registry {
	return &Registry{hosts: make([]entry, n)}
}

// Heartbeat records a heartbeat from host at second now. Only a host that was
// not Up costs a durable write.
func (r *Registry) Heartbeat(now, host int) {
	r.heartbeats++
	e := &r.hosts[host]
	e.lastSeen = now

	if e.state != Up {
		e.state = Up
		r.writes = append(r.writes, Transition{now, host, Up})
	}
}

// Sweep re-evaluates every host at second now and records each change of state.
func (r *Registry) Sweep(now int) {
	for host := range r.hosts {
		e := &r.hosts[host]

		if next := stateAfter(now - e.lastSeen); next != e.state {
			e.state = next
			r.writes = append(r.writes, Transition{now, host, next})
		}
	}
}

// stateAfter returns the state of a host that has been silent for the given
// number of seconds.
func stateAfter(silent int) State {
	switch {
	case silent >= deadAfter:
		return Dead
	case silent >= suspectAfter:
		return Suspect
	default:
		return Up
	}
}


// sendsHeartbeat reports whether host sends a heartbeat at second now in the
// simulation.
func sendsHeartbeat(host, now int) bool {
	switch host {
	case 42:
		return now < 20 // crashes at second 20
	case 7:
		return now < 30 || now >= 34 // silent from second 30 to 33
	default:
		return true
	}
}

// main runs the simulation and prints the registry's work.
func main() {
	registry := NewRegistry(hostCount)

	for now := 0; now < duration; now++ {
		for host := 0; host < hostCount; host++ {
			if sendsHeartbeat(host, now) {
				registry.Heartbeat(now, host)
			}
		}

		registry.Sweep(now)
	}

	fmt.Printf("heartbeats handled in memory: %d\n", registry.heartbeats)
	fmt.Printf("durable writes:               %d\n", len(registry.writes))

	for _, w := range registry.writes {
		fmt.Printf("  second %2d  host %2d -> %s\n", w.At, w.Host, w.State)
	}
}

Listing 3.26 The complete program in Rust. code/rust/heartbeats/src/main.rs

//! Models how a control plane tracks which machines in a fleet are alive.
//! Every host sends a heartbeat once a second. The registry keeps each host's
//! last heartbeat in memory, runs a sweep once a second, and records a durable
//! write only when a host changes state: Up, Suspect, or Dead.
//!
//! The program simulates 10,000 hosts for 60 seconds without real time
//! passing. Host 42 crashes at second 20, and host 7 misses four heartbeats
//! from second 30 to second 33. It prints how many heartbeats the registry
//! handled, how many durable writes it made, and each write.
//!
//! Run it from `code/rust` with `cargo run -p heartbeats`.

/// Seconds without a heartbeat before a host is Suspect.
const SUSPECT_AFTER: u32 = 3;
/// Seconds without a heartbeat before a host is Dead.
const DEAD_AFTER: u32 = 10;
/// Hosts in the fleet.
const HOST_COUNT: usize = 10_000;
/// Seconds simulated.
const DURATION: u32 = 60;

/// What the registry believes about a host.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
enum State {
    /// Heartbeats are arriving.
    #[default]
    Up,
    /// Heartbeats stopped recently; the host may be slow.
    Suspect,
    /// Heartbeats stopped long enough to give up on the host.
    Dead,
}

/// The registry's in-memory record of one host.
#[derive(Clone, Copy, Default)]
struct Entry {
    /// The second of the host's most recent heartbeat.
    last_seen: u32,
    /// The host's current state.
    state: State,
}

/// One durable write: a host entering a new state.
struct Transition {
    /// The second at which the sweep or heartbeat noticed the change.
    at: u32,
    /// The host's number.
    host: usize,
    /// The state the host entered.
    state: State,
}

/// Holds the liveness of every host.
struct Registry {
    /// Indexed by host number.
    hosts: Vec<Entry>,
    /// Counts the heartbeats handled in memory.
    heartbeats: usize,
    /// The durable writes, in order.
    writes: Vec<Transition>,
}

impl Registry {
    /// Returns a registry in which every host is Up, as it would be after the
    /// hosts registered.
    fn new(n: usize) -> Self {
        Self {
            hosts: vec![Entry::default(); n],
            heartbeats: 0,
            writes: Vec::new(),
        }
    }

    /// Records a heartbeat from `host` at second `now`. Only a host that was
    /// not Up costs a durable write.
    fn heartbeat(&mut self, now: u32, host: usize) {
        self.heartbeats += 1;
        let entry = &mut self.hosts[host];
        entry.last_seen = now;

        if entry.state != State::Up {
            entry.state = State::Up;
            self.writes.push(Transition {
                at: now,
                host,
                state: State::Up,
            });
        }
    }

    /// Re-evaluates every host at second `now` and records each change of state.
    fn sweep(&mut self, now: u32) {
        for (host, entry) in self.hosts.iter_mut().enumerate() {
            let next = state_after(now - entry.last_seen);

            if next != entry.state {
                entry.state = next;
                self.writes.push(Transition {
                    at: now,
                    host,
                    state: next,
                });
            }
        }
    }
}

/// Returns the state of a host that has been silent for `silent` seconds.
fn state_after(silent: u32) -> State {
    match silent {
        s if s >= DEAD_AFTER => State::Dead,
        s if s >= SUSPECT_AFTER => State::Suspect,
        _ => State::Up,
    }
}

/// Reports whether `host` sends a heartbeat at second `now` in the simulation.
fn sends_heartbeat(host: usize, now: u32) -> bool {
    match host {
        42 => now < 20,                // crashes at second 20
        7 => !(30..34).contains(&now), // silent from second 30 to 33
        _ => true,
    }
}

/// Runs the simulation and prints the registry's work.
fn main() {
    let mut registry = Registry::new(HOST_COUNT);

    for now in 0..DURATION {
        for host in 0..HOST_COUNT {
            if sends_heartbeat(host, now) {
                registry.heartbeat(now, host);
            }
        }

        registry.sweep(now);
    }

    println!("heartbeats handled in memory: {}", registry.heartbeats);
    println!("durable writes:               {}", registry.writes.len());

    for w in &registry.writes {
        println!("  second {:2}  host {:2} -> {:?}", w.at, w.host, w.state);
    }
}

Listing 3.27 The complete program in Go. code/go/phi/main.go

// Command phi compares two ways of deciding that a host has failed: a fixed
// timeout, and the phi accrual failure detector, which learns how regular each
// host's heartbeats are and reports how unlikely the current silence is.
//
// Two simulated hosts both send a heartbeat about once a second. One is
// steady; the other sits on a congested network and its heartbeats arrive
// between 0.4 and 2.4 seconds apart. For each host, the program prints the
// detector's suspicion level, phi, after different lengths of silence, next to
// the verdict of a 2-second timeout.
//
// Run it from code/go with:
//
//	go run ./phi
package main

import (
	"fmt"
	"math"
)

// timeout is the fixed timeout that the program compares against, in seconds.
const timeout = 2.0

// minStdDev is the smallest standard deviation the detector assumes, in
// seconds. Without it, a perfectly regular host would make any delay look
// infinitely suspicious.
const minStdDev = 0.2

// Detector is a phi accrual failure detector for one host.
type Detector struct {
	// intervals holds the most recent gaps between heartbeats, in seconds.
	intervals []float64
	// window is the number of gaps the detector remembers.
	window int
	// last is the time of the most recent heartbeat, in seconds.
	last float64
}

// NewDetector returns a detector that remembers window gaps and has seen a
// first heartbeat at time first.
func NewDetector(window int, first float64) *Detector {
	return &Detector{window: window, last: first}
}

// Heartbeat records a heartbeat that arrived at time now.
func (d *Detector) Heartbeat(now float64) {
	d.intervals = append(d.intervals, now-d.last)

	if len(d.intervals) > d.window {
		d.intervals = d.intervals[1:]
	}

	d.last = now
}


// Phi returns the suspicion level at time now: minus the base-10 logarithm of
// the probability that a heartbeat would still arrive after this much silence,
// given the gaps seen so far. Phi 1 means a 10% chance the host is fine, phi 2
// means 1%, and phi 3 means 0.1%.
func (d *Detector) Phi(now float64) float64 {
	var mean float64

	for _, gap := range d.intervals {
		mean += gap
	}

	mean /= float64(len(d.intervals))
	var variance float64

	for _, gap := range d.intervals {
		variance += (gap - mean) * (gap - mean)
	}

	stdDev := math.Max(math.Sqrt(variance/float64(len(d.intervals))), minStdDev)

	// The probability that a gap is longer than the current silence, from a
	// logistic approximation of the normal distribution.
	silence := now - d.last
	y := (silence - mean) / stdDev
	e := math.Exp(-y * (1.5976 + 0.070566*y*y))
	pLater := 1 - 1/(1+e)

	if silence > mean {
		pLater = e / (1 + e)
	}

	return -math.Log10(pLater)
}


// hosts lists the simulated hosts and the gaps between their heartbeats.
var hosts = []struct {
	name string
	gaps []float64
}{
	{"steady", []float64{1.0, 1.02, 0.98, 1.01, 0.99, 1.0, 1.03, 0.97, 1.0, 1.0}},
	{"jittery", []float64{0.4, 1.9, 0.6, 1.7, 0.5, 1.8, 0.7, 1.6, 0.4, 2.4}},
}

// main feeds each host's heartbeats to a detector and prints phi after
// increasing lengths of silence.
func main() {
	for _, host := range hosts {
		detector := NewDetector(100, 0)
		var now float64

		for _, gap := range host.gaps {
			now += gap
			detector.Heartbeat(now)
		}

		fmt.Printf("%s host, silent for:\n", host.name)

		for _, silence := range []float64{1.0, 1.5, 2.0, 2.5, 3.0, 4.0} {
			verdict := "alive"

			if silence > timeout {
				verdict = "dead"
			}

			fmt.Printf("  %.1fs  phi = %6.2f   %.0fs timeout says %s\n",
				silence, detector.Phi(now+silence), timeout, verdict)
		}
	}
}

Listing 3.28 The complete program in Rust. code/rust/phi/src/main.rs

//! Compares two ways of deciding that a host has failed: a fixed timeout, and
//! the phi accrual failure detector, which learns how regular each host's
//! heartbeats are and reports how unlikely the current silence is.
//!
//! Two simulated hosts both send a heartbeat about once a second. One is
//! steady; the other sits on a congested network and its heartbeats arrive
//! between 0.4 and 2.4 seconds apart. For each host, the program prints the
//! detector's suspicion level, phi, after different lengths of silence, next
//! to the verdict of a 2-second timeout.
//!
//! Run it from `code/rust` with `cargo run -p phi`.

use std::collections::VecDeque;

/// The fixed timeout that the program compares against, in seconds.
const TIMEOUT: f64 = 2.0;

/// The smallest standard deviation the detector assumes, in seconds. Without
/// it, a perfectly regular host would make any delay look infinitely
/// suspicious.
const MIN_STD_DEV: f64 = 0.2;

/// A phi accrual failure detector for one host.
struct Detector {
    /// The most recent gaps between heartbeats, in seconds.
    intervals: VecDeque<f64>,
    /// The number of gaps the detector remembers.
    window: usize,
    /// The time of the most recent heartbeat, in seconds.
    last: f64,
}

impl Detector {
    /// Returns a detector that remembers `window` gaps and has seen a first
    /// heartbeat at time `first`.
    fn new(window: usize, first: f64) -> Self {
        Self {
            intervals: VecDeque::new(),
            window,
            last: first,
        }
    }

    /// Records a heartbeat that arrived at time `now`.
    fn heartbeat(&mut self, now: f64) {
        self.intervals.push_back(now - self.last);

        if self.intervals.len() > self.window {
            self.intervals.pop_front();
        }

        self.last = now;
    }
}

impl Detector {
    /// Returns the suspicion level at time `now`: minus the base-10 logarithm
    /// of the probability that a heartbeat would still arrive after this much
    /// silence, given the gaps seen so far. Phi 1 means a 10% chance the host
    /// is fine, phi 2 means 1%, and phi 3 means 0.1%.
    #[allow(clippy::cast_precision_loss)] // the window holds at most 100 gaps
    fn phi(&self, now: f64) -> f64 {
        let n = self.intervals.len() as f64;
        let mean = self.intervals.iter().sum::<f64>() / n;
        let variance = self
            .intervals
            .iter()
            .map(|gap| (gap - mean).powi(2))
            .sum::<f64>()
            / n;
        let std_dev = variance.sqrt().max(MIN_STD_DEV);

        // The probability that a gap is longer than the current silence, from
        // a logistic approximation of the normal distribution.
        let silence = now - self.last;
        let y = (silence - mean) / std_dev;
        let e = (-y * (1.5976 + 0.070_566 * y * y)).exp();
        let p_later = if silence > mean {
            e / (1.0 + e)
        } else {
            1.0 - 1.0 / (1.0 + e)
        };

        -p_later.log10()
    }
}

/// The simulated hosts and the gaps between their heartbeats.
const HOSTS: [(&str, [f64; 10]); 2] = [
    (
        "steady",
        [1.0, 1.02, 0.98, 1.01, 0.99, 1.0, 1.03, 0.97, 1.0, 1.0],
    ),
    (
        "jittery",
        [0.4, 1.9, 0.6, 1.7, 0.5, 1.8, 0.7, 1.6, 0.4, 2.4],
    ),
];

/// Feeds each host's heartbeats to a detector and prints phi after increasing
/// lengths of silence.
fn main() {
    for (name, gaps) in HOSTS {
        let mut detector = Detector::new(100, 0.0);
        let mut now = 0.0;

        for gap in gaps {
            now += gap;
            detector.heartbeat(now);
        }

        println!("{name} host, silent for:");

        for silence in [1.0, 1.5, 2.0, 2.5, 3.0, 4.0] {
            let verdict = if silence > TIMEOUT { "dead" } else { "alive" };
            println!(
                "  {silence:.1}s  phi = {:6.2}   {TIMEOUT:.0}s timeout says {verdict}",
                detector.phi(now + silence)
            );
        }
    }
}

Listing 3.29 The complete program in Go. code/go/leaseagent/main.go

// Command leaseagent keeps a Kubernetes Lease alive for one host, the way the
// kubelet keeps its node's Lease in the kube-node-lease namespace alive. Every
// two seconds it writes the Lease named after the host in the fleet namespace,
// with the current time as renewTime and a lease duration of six seconds.
// Other programs judge the host alive for as long as these renewals continue.
//
// Run it from code/go, against the cluster in the current kubeconfig context,
// with:
//
//	go run ./leaseagent gpu-host-a
package main

import (
	"context"
	"fmt"
	"os"
	"os/signal"
	"syscall"
	"time"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	coordinationv1ac "k8s.io/client-go/applyconfigurations/coordination/v1"
	"k8s.io/client-go/kubernetes"
	ctrl "sigs.k8s.io/controller-runtime"
)

// Lease settings. The renewal interval is a third of the lease duration, so a
// host has to miss three renewals in a row before its Lease expires.
const (
	namespace     = "fleet"
	renewEvery    = 2 * time.Second
	leaseDuration = 6 // seconds
	fieldManager  = "leaseagent"
)


// main reads the host name from the command line and renews its Lease until
// the process receives Ctrl-C or SIGTERM.
func main() {
	if len(os.Args) != 2 {
		fmt.Fprintln(os.Stderr, "usage: leaseagent <host>")
		os.Exit(2)
	}

	if err := run(os.Args[1]); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}

// run writes the host's Lease with server-side apply once per renewEvery.
// A renewal that fails is reported and retried at the next tick.
func run(host string) error {
	clientset, err := kubernetes.NewForConfig(ctrl.GetConfigOrDie())

	if err != nil {
		return err
	}

	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()

	leases := clientset.CoordinationV1().Leases(namespace)
	acquired := metav1.NewMicroTime(time.Now())
	ticker := time.NewTicker(renewEvery)
	defer ticker.Stop()

	for {
		now := metav1.NewMicroTime(time.Now())
		lease := coordinationv1ac.Lease(host, namespace).
			WithSpec(coordinationv1ac.LeaseSpec().
				WithHolderIdentity(host).
				WithLeaseDurationSeconds(leaseDuration).
				WithAcquireTime(acquired).
				WithRenewTime(now))
		_, err := leases.Apply(ctx, lease, metav1.ApplyOptions{FieldManager: fieldManager, Force: true})

		switch {
		case ctx.Err() != nil:
			return nil
		case err != nil:
			fmt.Printf("%s %s renewal failed: %v\n", now.Format(time.TimeOnly), host, err)
		default:
			fmt.Printf("%s %s renewed\n", now.Format(time.TimeOnly), host)
		}

		select {
		case <-ctx.Done():
			return nil
		case <-ticker.C:
		}
	}
}

Listing 3.30 The complete program in Rust. code/rust/leaseagent/src/main.rs

//! Keeps a Kubernetes Lease alive for one host, the way the kubelet keeps its
//! node's Lease in the `kube-node-lease` namespace alive. Every two seconds it
//! writes the Lease named after the host in the `fleet` namespace, with the
//! current time as `renewTime` and a lease duration of six seconds. Other
//! programs judge the host alive for as long as these renewals continue.
//!
//! Run it from `code/rust`, against the cluster in the current kubeconfig
//! context, with `cargo run -p leaseagent -- gpu-host-a`.

use jiff::Timestamp;
use jiff::tz::TimeZone;
use k8s_openapi::api::coordination::v1::{Lease, LeaseSpec};
use k8s_openapi::apimachinery::pkg::apis::meta::v1::MicroTime;
use kube::api::{ObjectMeta, Patch, PatchParams};
use kube::{Api, Client};
use std::error::Error;
use std::time::Duration;

/// The namespace that holds the hosts' Leases.
const NAMESPACE: &str = "fleet";
/// How often the agent renews the Lease: a third of the lease duration, so a
/// host has to miss three renewals in a row before its Lease expires.
const RENEW_EVERY: Duration = Duration::from_secs(2);
/// How long a Lease stays valid after its last renewal, in seconds.
const LEASE_DURATION: i32 = 6;
/// The field manager name for server-side apply.
const FIELD_MANAGER: &str = "leaseagent";

/// Builds the Lease that `host` writes at time `now`, having first acquired
/// it at time `acquired`.
fn lease(host: &str, acquired: Timestamp, now: Timestamp) -> Lease {
    Lease {
        metadata: ObjectMeta {
            name: Some(host.to_string()),
            namespace: Some(NAMESPACE.to_string()),
            ..ObjectMeta::default()
        },
        spec: Some(LeaseSpec {
            holder_identity: Some(host.to_string()),
            lease_duration_seconds: Some(LEASE_DURATION),
            acquire_time: Some(MicroTime(acquired)),
            renew_time: Some(MicroTime(now)),
            ..LeaseSpec::default()
        }),
    }
}

/// Formats a timestamp as local time of day, for the log.
fn time_of_day(t: Timestamp) -> String {
    t.to_zoned(TimeZone::system())
        .strftime("%H:%M:%S")
        .to_string()
}

/// Reads the host name from the command line and renews its Lease with
/// server-side apply until the process receives Ctrl-C. A renewal that fails
/// is reported and retried at the next tick.
#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    let Some(host) = std::env::args().nth(1) else {
        eprintln!("usage: leaseagent <host>");
        std::process::exit(2);
    };
    let leases: Api<Lease> = Api::namespaced(Client::try_default().await?, NAMESPACE);
    let params = PatchParams::apply(FIELD_MANAGER).force();
    let acquired = Timestamp::now();
    let mut ticker = tokio::time::interval(RENEW_EVERY);

    loop {
        tokio::select! {
            _ = tokio::signal::ctrl_c() => return Ok(()),
            _ = ticker.tick() => {}
        }
        let now = Timestamp::now();
        match leases
            .patch(&host, &params, &Patch::Apply(lease(&host, acquired, now)))
            .await
        {
            Ok(_) => println!("{} {host} renewed", time_of_day(now)),
            Err(error) => println!("{} {host} renewal failed: {error}", time_of_day(now)),
        }
    }
}

Listing 3.31 The complete program in Go. code/go/leasewatch/main.go

// Command leasewatch reports when a host in the fleet namespace stops renewing
// its Lease, and when it starts again. It keeps an informer cache of the
// Leases and checks it once a second.
//
// It judges a Lease by its own clock, the way the Kubernetes node lifecycle
// controller judges node Leases: it notes the local time at which it first
// sees each new renewTime, and declares the host expired when no new
// renewTime has appeared for longer than the lease duration. The renewTime
// value itself was written with the host's clock, which may differ from ours.
//
// Run it from code/go, against the cluster in the current kubeconfig context,
// with:
//
//	go run ./leasewatch
package main

import (
	"context"
	"fmt"
	"os"
	"os/signal"
	"syscall"
	"time"

	coordinationv1 "k8s.io/api/coordination/v1"
	"k8s.io/apimachinery/pkg/labels"
	"k8s.io/client-go/informers"
	"k8s.io/client-go/kubernetes"
	ctrl "sigs.k8s.io/controller-runtime"
)

// namespace is the namespace whose Leases the program watches.
const namespace = "fleet"

// checkEvery is how often the program re-evaluates every Lease.
const checkEvery = time.Second

// observation is what the program knows about one host's Lease.
type observation struct {
	// renewTime is the latest renewTime seen in the Lease.
	renewTime time.Time
	// observedAt is the local time at which that renewTime was first seen.
	observedAt time.Time
	// expired records whether the host has been reported as expired.
	expired bool
}


// main runs the watcher until the process receives Ctrl-C or SIGTERM.
func main() {
	if err := run(); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}

// run starts an informer for Leases in the namespace, waits for its first
// list, and then checks every Lease in its cache once per checkEvery.
func run() error {
	clientset, err := kubernetes.NewForConfig(ctrl.GetConfigOrDie())

	if err != nil {
		return err
	}

	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()

	factory := informers.NewSharedInformerFactoryWithOptions(clientset, 0, informers.WithNamespace(namespace))
	lister := factory.Coordination().V1().Leases().Lister().Leases(namespace)
	factory.Start(ctx.Done())
	factory.WaitForCacheSync(ctx.Done())

	seen := map[string]*observation{}
	ticker := time.NewTicker(checkEvery)
	defer ticker.Stop()

	for {
		leases, err := lister.List(labels.Everything())

		if err != nil {
			return err
		}

		now := time.Now()

		for _, lease := range leases {
			check(seen, lease, now)
		}

		select {
		case <-ctx.Done():
			return nil
		case <-ticker.C:
		}
	}
}


// check updates the observation of one Lease at local time now and prints a
// line when the host is first seen, expires, or renews again after expiring.
func check(seen map[string]*observation, lease *coordinationv1.Lease, now time.Time) {
	if lease.Spec.RenewTime == nil || lease.Spec.LeaseDurationSeconds == nil {
		return
	}

	renew := lease.Spec.RenewTime.Time
	duration := time.Duration(*lease.Spec.LeaseDurationSeconds) * time.Second
	stamp := now.Format(time.TimeOnly)

	obs, known := seen[lease.Name]

	switch {
	case !known:
		seen[lease.Name] = &observation{renewTime: renew, observedAt: now}
		fmt.Printf("%s %s seen, lease duration %s\n", stamp, lease.Name, duration)
	case !renew.Equal(obs.renewTime):
		obs.renewTime, obs.observedAt = renew, now

		if obs.expired {
			obs.expired = false
			fmt.Printf("%s %s renewing again\n", stamp, lease.Name)
		}
	case !obs.expired && now.Sub(obs.observedAt) > duration:
		obs.expired = true
		fmt.Printf("%s %s expired: no renewal seen for %s\n",
			stamp, lease.Name, now.Sub(obs.observedAt).Round(time.Second))
	}
}

Listing 3.32 The complete program in Rust. code/rust/leasewatch/src/main.rs

//! Reports when a host in the `fleet` namespace stops renewing its Lease, and
//! when it starts again. It keeps a reflector store of the Leases and checks
//! it once a second.
//!
//! It judges a Lease by its own clock, the way the Kubernetes node lifecycle
//! controller judges node Leases: it notes the local time at which it first
//! sees each new `renewTime`, and declares the host expired when no new
//! `renewTime` has appeared for longer than the lease duration. The
//! `renewTime` value itself was written with the host's clock, which may
//! differ from ours.
//!
//! Run it from `code/rust`, against the cluster in the current kubeconfig
//! context, with `cargo run -p leasewatch`.

use futures::StreamExt;
use jiff::{Timestamp, Zoned};
use k8s_openapi::api::coordination::v1::Lease;
use kube::runtime::{WatchStreamExt, reflector, watcher};
use kube::{Api, Client, ResourceExt};
use std::collections::HashMap;
use std::error::Error;
use std::time::{Duration, Instant};

/// The namespace whose Leases the program watches.
const NAMESPACE: &str = "fleet";
/// How often the program re-evaluates every Lease.
const CHECK_EVERY: Duration = Duration::from_secs(1);

/// What the program knows about one host's Lease.
struct Observation {
    /// The latest `renewTime` seen in the Lease.
    renew_time: Timestamp,
    /// The local time at which that `renewTime` was first seen.
    observed_at: Instant,
    /// Whether the host has been reported as expired.
    expired: bool,
}

/// Updates the observation of one Lease at local time `now` and prints a line
/// when the host is first seen, expires, or renews again after expiring.
fn check(seen: &mut HashMap<String, Observation>, lease: &Lease, now: Instant) {
    let Some(spec) = &lease.spec else { return };
    let (Some(renew), Some(seconds)) = (&spec.renew_time, spec.lease_duration_seconds) else {
        return;
    };
    let duration = Duration::from_secs(u64::from(seconds.unsigned_abs()));
    let stamp = Zoned::now().strftime("%H:%M:%S").to_string();
    let name = lease.name_any();

    match seen.get_mut(&name) {
        None => {
            seen.insert(
                name.clone(),
                Observation {
                    renew_time: renew.0,
                    observed_at: now,
                    expired: false,
                },
            );
            println!(
                "{stamp} {name} seen, lease duration {}s",
                duration.as_secs()
            );
        }
        Some(obs) if obs.renew_time != renew.0 => {
            obs.renew_time = renew.0;
            obs.observed_at = now;

            if obs.expired {
                obs.expired = false;
                println!("{stamp} {name} renewing again");
            }
        }
        Some(obs) if !obs.expired && now.duration_since(obs.observed_at) > duration => {
            obs.expired = true;
            let silent = now.duration_since(obs.observed_at).as_secs_f64().round();
            println!("{stamp} {name} expired: no renewal seen for {silent}s");
        }
        Some(_) => {}
    }
}

/// Starts a watcher that fills a reflector store with the namespace's Leases,
/// waits for the first list, and then checks every Lease in the store once
/// per `CHECK_EVERY` until the process receives Ctrl-C.
#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    let leases: Api<Lease> = Api::namespaced(Client::try_default().await?, NAMESPACE);
    let (store, writer) = reflector::store();
    let stream = watcher(leases, watcher::Config::default())
        .default_backoff()
        .reflect(writer)
        .applied_objects();
    tokio::spawn(stream.for_each(|_| std::future::ready(())));
    store.wait_until_ready().await?;

    let mut seen = HashMap::new();
    let mut ticker = tokio::time::interval(CHECK_EVERY);

    loop {
        tokio::select! {
            _ = tokio::signal::ctrl_c() => return Ok(()),
            _ = ticker.tick() => {}
        }
        let now = Instant::now();

        for lease in store.state() {
            check(&mut seen, &lease, now);
        }
    }
}

References

• Kubernetes, Leases: https://kubernetes.io/docs/concepts/architecture/leases/
• Kubernetes, Node status and heartbeats: https://kubernetes.io/docs/reference/node/node-status/#heartbeats
• Kubernetes, Nodes and the node controller: https://kubernetes.io/docs/concepts/architecture/nodes/
• Kubernetes source, node lifecycle controller: https://github.com/kubernetes/kubernetes/blob/master/pkg/controller/nodelifecycle/node_lifecycle_controller.go
• N. Hayashibara, X. Défago, R. Yared, T. Katayama, “The φ Accrual Failure Detector”, 23rd IEEE International Symposium on Reliable Distributed Systems (SRDS), 2004.

4. A control plane for a fleet of hosts

In chapter 3, we learned how a control plane notices that a machine has gone silent. A silent machine is the simple case. More often a machine keeps sending heartbeats while something on it is going wrong: a GPU reports a memory error, a job crashes with a GPU fault, or a network link drops for a second and returns. In this chapter, we will build the part of a control plane that reads those reports and decides what to do with the machine.

We will work with a fleet of real machines, small enough to keep in our heads: five Jetson AGX Thor developer kits on one switch, two DGX Sparks joined by a cable, and the Jetson Orin NX that has run every example so far. We will first reason about what each kind of report tells us, then draw the design, then write it as two programs in Go and in Rust. fleethealth decides which reports blame a machine and walks a blamed machine through its repair. fencing deals with the question chapter 3 ended on: what stops a job from running twice when a machine that was declared dead turns out to be alive.

The fleet

Figure 4.1 shows the eight machines. Each has one GPU that shares memory with the CPU, as we saw for the Orin in chapter 2, so a machine’s memory is also the limit on what its GPU can hold.

Three facts about this fleet shape the chapter. The machines are few, so losing two of them removes a quarter of the capacity. They are connected in groups: the five Thors reach each other through one switch, and the two Sparks through one cable, so one failed switch looks like five failed machines. And every machine runs long jobs whose progress lives in the GPU’s memory until it is saved, so taking a machine away without warning destroys work.

A fleet of ten thousand machines has the same three properties with larger numbers. The design that follows is the one such fleets use, and our eight machines are enough to exercise every part of it.

What a failure report tells us

Each machine produces several kinds of report about its GPU and its network link.

The GPU driver writes a message to the kernel log when the GPU reports an error. On discrete and data center GPUs, the NVIDIA driver numbers these messages, and the number is called an Xid. NVIDIA publishes a catalog that gives the usual cause of each one. The table lists the Xids our model has a class for:

Of our two boards, the Thor runs this driver, version 595.78, and reports its errors as Xids. Its kernel log held none when we looked. The Orin NX runs the older nvgpu driver, which writes its errors to the kernel log without Xid numbers. The model reduces every report to a class, so it does not depend on the format.

ECC stands for error-correcting code: GPU memory stores a few extra bits with each word, enough to repair one flipped bit and to detect two. An uncorrectable error is one the extra bits detected and could not repair.

The network interface reports each time its link goes down and comes back. We call one such event a link flap.

The same report can have three different causes, and the right response differs for each:

• A job with a bug in its GPU code produces an application fault on every machine it runs on. Taking those machines out of service repairs nothing and removes capacity.
• A machine with a failing GPU produces faults for every job placed on it. The machine has to leave the pool before the scheduler places another job there.
• A switch that restarts makes every machine attached to it report a link flap within a few seconds. The machines are healthy, and draining them would turn a short network interruption into a loss of five machines’ worth of running jobs.

One report cannot distinguish these cases. Several reports can, by counting along two directions. A fault that follows one job across three machines blames the job. Faults from three different jobs on one machine blame the machine. Flaps from three machines on one link, close together in time, blame the link. A report of a memory error or of a GPU off the bus needs no counting, because no program can cause one.

Two paths for one report

A report has to reach two places, at two speeds. Figure 4.2 shows both.

The fast path stays on the machine. The device plugin, which we met in the introduction, reports the health of each GPU to the kubelet along with the count. When a GPU fails, the plugin marks it unhealthy, the kubelet lowers the number of GPUs the node offers, and the scheduler stops placing pods that ask for that GPU. This takes seconds and needs no decision from anyone. It stops new work only: a pod that already holds the GPU keeps running.

The slow path leaves the machine. Every report goes to a health service, which holds the counts from the previous section and turns them into a verdict for a machine: leave it alone, or repair it. A remediation controller acts on the verdict. It is a controller in the sense of chapter 1: it holds a desired state for each machine and moves the machine towards it one step at a time, through the Kubernetes API. In a production cluster the reports travel through a metrics pipeline, and the verdict is written to a custom resource that the controller watches. Our model keeps both halves in one process so that we can see every decision in one log.

The two halves stay separate in the design for a practical reason. The rules that produce verdicts change often, as new failure patterns appear. The controller that drains and reboots machines changes rarely and is dangerous when it is wrong. Keeping them apart lets us change a rule, or run a new rule in a mode that only logs its verdicts, without touching the code that can take machines away.

The repair lifecycle

A machine with a verdict against it goes through six states. Figure 4.3 shows them.

Cordoned means the scheduler places no new pods on the machine. In Kubernetes this is one field, spec.unschedulable, on the Node object, and kubectl cordon sets it. Cordoning is cheap and fully reversible, so the controller does it at once.

Draining removes the pods that are already there. A training job loses everything since its last checkpoint when its pod is deleted, so the controller first tells the job to write a checkpoint and gives it time to do so. Kubernetes has a request made for this purpose, the Eviction API. An eviction asks for a pod to be removed and can be refused: a PodDisruptionBudget is an object in which a workload’s owner states how many of its pods may be down at once, and the API server refuses an eviction that would exceed it. kubectl drain cordons a node and then evicts each of its pods in this way.

Remediating runs one repair action. The actions form a ladder from the cheapest to the most thorough: reset the GPU, reboot the machine, and write a fresh operating system image to it. The controller starts at the bottom and climbs one rung at a time.

Validating runs a burn-in test: a short GPU workload with a known answer. A machine that boots and joins the cluster has shown only that its operating system starts. The burn-in shows that its GPU computes correctly under load, which is the property that failed. A machine that passes returns to Healthy. A machine that fails goes back for the next rung, and a machine that fails after the last rung is Quarantined: it stays out of the pool, and a ticket asks a person to look at the hardware.

A limit on the repair system itself

A controller that can drain machines can also drain all of them. If a faulty rule condemned every machine in the fleet, a controller that obeyed would cause a larger outage than any hardware failure. So the controller carries a disruption budget: the largest number of machines that may be out of service at once. A machine whose verdict arrives when the budget is spent stays Cordoned, which costs nothing that is running, and waits for a slot. A budget that stays spent is itself a signal worth a page to a person, because it means failures are arriving faster than repairs finish.

Our fleet’s budget is two machines. A quarantined machine keeps its slot until someone repairs or removes it, since it is still missing from the fleet.

Building fleethealth

The program is a simulation with a clock that counts seconds. Each second it feeds the signals that arrive at that second to Observe, then calls Step. No real time passes, so a run takes a few milliseconds and prints the same log every time.

The fleet and its signals

A machine is a name, a model, and the name of the switch or cable it shares with its neighbours. The link name is what lets the program see that five reports came from one place.

Listing 4.1 The machines, in Go. code/go/fleethealth/main.go

// Node is one machine in the fleet.
type Node struct {
	// Name is the machine's host name.
	Name string
	// Model is the product the machine is.
	Model string
	// Link names the switch or cable that connects the machine to its peers.
	Link string
}

// fleet lists the machines in the order the program reports them.
var fleet = []Node{
	{"thor-1", "Jetson AGX Thor", "switch-25g"},
	{"thor-2", "Jetson AGX Thor", "switch-25g"},
	{"thor-3", "Jetson AGX Thor", "switch-25g"},
	{"thor-4", "Jetson AGX Thor", "switch-25g"},
	{"thor-5", "Jetson AGX Thor", "switch-25g"},
	{"spark-1", "DGX Spark", "cable-200g"},
	{"spark-2", "DGX Spark", "cable-200g"},
	{"orin-nx-1", "Jetson Orin NX", "bench"},
}

Listing 4.2 The machines, in Rust. code/rust/fleethealth/src/main.rs

/// One machine in the fleet.
struct Node {
    /// The machine's host name.
    name: &'static str,
    /// The product the machine is.
    model: &'static str,
    /// The switch or cable that connects the machine to its peers.
    link: &'static str,
}

/// The machines, in the order the program reports them.
const FLEET: [Node; 8] = [
    Node {
        name: "thor-1",
        model: "Jetson AGX Thor",
        link: "switch-25g",
    },
    Node {
        name: "thor-2",
        model: "Jetson AGX Thor",
        link: "switch-25g",
    },
    Node {
        name: "thor-3",
        model: "Jetson AGX Thor",
        link: "switch-25g",
    },
    Node {
        name: "thor-4",
        model: "Jetson AGX Thor",
        link: "switch-25g",
    },
    Node {
        name: "thor-5",
        model: "Jetson AGX Thor",
        link: "switch-25g",
    },
    Node {
        name: "spark-1",
        model: "DGX Spark",
        link: "cable-200g",
    },
    Node {
        name: "spark-2",
        model: "DGX Spark",
        link: "cable-200g",
    },
    Node {
        name: "orin-nx-1",
        model: "Jetson Orin NX",
        link: "bench",
    },
];

A signal carries the second it arrived, the machine, the job that was using the GPU, and its class. Hardware signals and link flaps have no job.

Listing 4.3 Signals and their classes, in Go. code/go/fleethealth/main.go

// Kind is the class of a failure signal.
type Kind int

// The classes of signal the model understands.
const (
	AppFault  Kind = iota // the GPU rejected work from one program
	MemoryECC             // an uncorrectable memory error on the GPU
	OffBus                // the GPU stopped answering the driver
	LinkFlap              // the machine's network link went down and up
)

// String returns the class as it appears in the log.
func (k Kind) String() string {
	return [...]string{"application fault", "uncorrectable ECC", "GPU off the bus", "link flap"}[k]
}

// Signal is one failure report from a machine.
type Signal struct {
	// At is the second at which the signal arrived.
	At int
	// Node is the machine that reported it.
	Node string
	// Job is the job that was using the GPU, or "" for a signal with no job.
	Job string
	// Kind is the signal's class.
	Kind Kind
}

Listing 4.4 Signals and their classes, in Rust. code/rust/fleethealth/src/main.rs

/// The class of a failure signal.
#[derive(Clone, Copy)]
enum Kind {
    /// The GPU rejected work from one program.
    AppFault,
    /// An uncorrectable memory error on the GPU.
    MemoryEcc,
    /// The GPU stopped answering the driver.
    OffBus,
    /// The machine's network link went down and up.
    LinkFlap,
}

impl fmt::Display for Kind {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(match self {
            Kind::AppFault => "application fault",
            Kind::MemoryEcc => "uncorrectable ECC",
            Kind::OffBus => "GPU off the bus",
            Kind::LinkFlap => "link flap",
        })
    }
}

/// One failure report from a machine.
struct Signal {
    /// The second at which the signal arrived.
    at: u32,
    /// The machine that reported it.
    node: &'static str,
    /// The job that was using the GPU, or "" for a signal with no job.
    job: &'static str,
    /// The signal's class.
    kind: Kind,
}

The states and the ladder

The states are declared in the order a failing machine passes through them, and both languages can compare them. That ordering does real work later: “every state from Draining on” becomes one comparison.

Listing 4.5 The lifecycle states and the repair ladder, in Go. code/go/fleethealth/main.go

// State is a machine's place in the repair lifecycle.
type State int

// The lifecycle, in the order a failing machine passes through it.
const (
	Healthy     State = iota // schedulable
	Cordoned                 // no new work; waiting for the disruption budget
	Draining                 // jobs are checkpointing and leaving
	Remediating              // a repair action is running
	Validating               // a burn-in test is running
	Quarantined              // every repair failed; a person has to look
)

// String returns the state's name for the log.
func (s State) String() string {
	return [...]string{"Healthy", "Cordoned", "Draining", "Remediating", "Validating", "Quarantined"}[s]
}

// ladder lists the repair actions from the cheapest to the most thorough.
var ladder = []string{"reset the GPU", "reboot", "reimage"}

Listing 4.6 The lifecycle states and the repair ladder, in Rust. code/rust/fleethealth/src/main.rs

/// A machine's place in the repair lifecycle, in the order a failing machine
/// passes through it.
#[derive(Clone, Copy, Debug, PartialEq, PartialOrd)]
enum State {
    /// Schedulable.
    Healthy,
    /// No new work; waiting for the disruption budget.
    Cordoned,
    /// Jobs are checkpointing and leaving.
    Draining,
    /// A repair action is running.
    Remediating,
    /// A burn-in test is running.
    Validating,
    /// Every repair failed; a person has to look.
    Quarantined,
}

impl fmt::Display for State {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        // `pad` honours the width given in a format string such as `{:<11}`.
        f.pad(&format!("{self:?}"))
    }
}

/// The repair actions, from the cheapest to the most thorough.
const LADDER: [&str; 3] = ["reset the GPU", "reboot", "reimage"];

Deciding what a signal blames

Observe holds the counting rules. A hardware signal condemns its machine at once. An application fault is recorded twice, under the machine and under the job, and then both counts are checked. Three machines for one job blames the job, and the program prints that verdict once and leaves every machine in service. Three jobs on one machine condemns the machine. A link flap is stamped on its machine, and the program counts the machines on the same link whose latest flap lies within five seconds.

Listing 4.7 Triage: which of a job, a machine, or a link a signal blames, in Go. code/go/fleethealth/main.go

// Observe takes one signal and decides what it blames. A hardware signal
// blames the machine. An application fault blames the job once the job has
// faulted on three machines, and blames the machine once three jobs have
// faulted on it. Link flaps blame the link when three machines on it flap
// within five seconds of each other.
func (p *Plane) Observe(s Signal) {
	r := p.find(s.Node)

	switch s.Kind {
	case MemoryECC, OffBus:
		p.condemn(s.At, r, s.Kind.String())
	case AppFault:
		r.faultJobs[s.Job] = true

		if p.jobNodes[s.Job] == nil {
			p.jobNodes[s.Job] = map[string]bool{}
		}

		p.jobNodes[s.Job][s.Node] = true

		switch {
		case len(p.jobNodes[s.Job]) >= blameAfter && !p.blamed[s.Job]:
			p.blamed[s.Job] = true
			fmt.Printf("%4ds  job %s faulted on %d machines: the job is at fault, no machine is drained\n",
				s.At, s.Job, len(p.jobNodes[s.Job]))
		case len(r.faultJobs) >= blameAfter:
			p.condemn(s.At, r, fmt.Sprintf("%d different jobs faulted here", len(r.faultJobs)))
		}
	case LinkFlap:
		r.lastFlap = s.At
		together := 0

		for _, other := range p.records {
			if other.node.Link == r.node.Link && other.lastFlap >= 0 && s.At-other.lastFlap <= fabricWindow {
				together++
			}
		}

		if together >= fabricAfter && !p.blamed[r.node.Link] {
			p.blamed[r.node.Link] = true
			fmt.Printf("%4ds  %d machines on %s flapped together: the link is at fault, no machine is drained\n",
				s.At, together, r.node.Link)
		}
	}
}

Listing 4.8 Triage: which of a job, a machine, or a link a signal blames, in Rust. code/rust/fleethealth/src/main.rs

impl Plane {
    /// Takes one signal and decides what it blames. A hardware signal blames
    /// the machine. An application fault blames the job once the job has
    /// faulted on three machines, and blames the machine once three jobs have
    /// faulted on it. Link flaps blame the link when three machines on it flap
    /// within five seconds of each other.
    fn observe(&mut self, s: &Signal) {
        let index = self.find(s.node);

        match s.kind {
            Kind::MemoryEcc | Kind::OffBus => self.condemn(s.at, index, &s.kind.to_string()),
            Kind::AppFault => {
                self.records[index].fault_jobs.insert(s.job);
                let nodes = self.job_nodes.entry(s.job).or_default();
                nodes.insert(s.node);
                let (machines, jobs) = (nodes.len(), self.records[index].fault_jobs.len());

                if machines >= BLAME_AFTER && self.blamed.insert(s.job) {
                    println!(
                        "{:4}s  job {} faulted on {machines} machines: the job is at fault, no machine is drained",
                        s.at, s.job
                    );
                } else if machines < BLAME_AFTER && jobs >= BLAME_AFTER {
                    self.condemn(s.at, index, &format!("{jobs} different jobs faulted here"));
                }
            }
            Kind::LinkFlap => {
                self.records[index].last_flap = Some(s.at);
                let link = self.records[index].node.link;
                let together = self
                    .records
                    .iter()
                    .filter(|r| r.node.link == link)
                    .filter(|r| r.last_flap.is_some_and(|at| s.at - at <= FABRIC_WINDOW))
                    .count();
                if together >= FABRIC_AFTER && self.blamed.insert(link) {
                    println!(
                        "{:4}s  {together} machines on {link} flapped together: the link is at fault, no machine is drained",
                        s.at
                    );
                }
            }
        }
    }
}

The Go version keeps each set as a map from a name to true. The Rust version uses BTreeSet, whose insert returns whether the value was new, so “blame this once” is the single expression self.blamed.insert(s.job).

The budget

The budget check is a count of the machines in the four states that take a machine away from the fleet. Cordoned is left out on purpose: a cordoned machine still runs its jobs.

Listing 4.9 Counting the machines that are out of service, in Go. code/go/fleethealth/main.go

// outOfService counts the machines that hold a slot of the disruption budget:
// every machine that is draining, under repair, in burn-in, or quarantined.
func (p *Plane) outOfService() int {
	n := 0

	for _, r := range p.records {
		if r.state >= Draining {
			n++
		}
	}

	return n
}

Listing 4.10 Counting the machines that are out of service, in Rust. code/rust/fleethealth/src/main.rs

impl Plane {
    /// Counts the machines that hold a slot of the disruption budget: every
    /// machine that is draining, under repair, in burn-in, or quarantined.
    fn out_of_service(&self) -> usize {
        self.records
            .iter()
            .filter(|r| r.state >= State::Draining)
            .count()
    }
}

One second of the lifecycle

Step visits every machine and asks how long it has been in its state. A Cordoned machine starts to drain as soon as the budget has room, and reports once that it is blocked when the budget has none. Each of the other states lasts a fixed number of seconds in the model: four for a job to write its checkpoint and leave, three for a repair action, three for the burn-in. The result of a burn-in comes from a table the run supplies, which says how many tests each machine fails before one passes.

Listing 4.11 Advancing every machine by one second, in Go. code/go/fleethealth/main.go

// Step advances every machine by one second. A Cordoned machine starts to
// drain when the budget has a free slot. The other states each last a fixed
// time and then move on; a failed burn-in sends the machine back for the next
// action on the ladder, and a machine with no actions left is quarantined.
func (p *Plane) Step(now int) {
	for _, r := range p.records {
		waited := now - r.since

		switch {
		case r.state == Cordoned && p.outOfService() < budget:
			p.move(now, r, Draining, "jobs told to checkpoint")
		case r.state == Cordoned && waited == 1:
			fmt.Printf("%4ds  %-10s blocked: %d machines are already out of service\n", now, r.node.Name, budget)
		case r.state == Draining && waited >= drainTime:
			p.move(now, r, Remediating, ladder[r.step])
		case r.state == Remediating && waited >= repairTime:
			p.move(now, r, Validating, "burn-in after "+ladder[r.step])
		case r.state == Validating && waited >= burnInTime:
			p.finishBurnIn(now, r)
		}
	}
}

// finishBurnIn acts on the result of a burn-in test. A pass returns the
// machine to service. A failure sends it back for the next action on the
// ladder, and a machine with no actions left is quarantined.
func (p *Plane) finishBurnIn(now int, r *record) {
	if p.burnIn[r.node.Name] == 0 {
		r.faultJobs = map[string]bool{}
		p.move(now, r, Healthy, "burn-in passed")

		return
	}

	p.burnIn[r.node.Name]--

	if r.step+1 == len(ladder) {
		p.move(now, r, Quarantined, "burn-in failed after every action; ticket opened")

		return
	}

	r.step++
	p.move(now, r, Remediating, "burn-in failed; "+ladder[r.step])
}

Listing 4.12 Advancing every machine by one second, in Rust. code/rust/fleethealth/src/main.rs

impl Plane {
    /// Advances every machine by one second. A Cordoned machine starts to
    /// drain when the budget has a free slot. The other states each last a
    /// fixed time and then move on.
    fn step(&mut self, now: u32) {
        for index in 0..self.records.len() {
            let (state, step, name) = {
                let r = &self.records[index];
                (r.state, r.step, r.node.name)
            };
            let waited = now - self.records[index].since;

            match state {
                State::Cordoned if self.out_of_service() < BUDGET => {
                    self.transition(now, index, State::Draining, "jobs told to checkpoint");
                }
                State::Cordoned if waited == 1 => {
                    println!(
                        "{now:4}s  {name:<10} blocked: {BUDGET} machines are already out of service"
                    );
                }
                State::Draining if waited >= DRAIN_TIME => {
                    self.transition(now, index, State::Remediating, LADDER[step]);
                }
                State::Remediating if waited >= REPAIR_TIME => {
                    let why = format!("burn-in after {}", LADDER[step]);
                    self.transition(now, index, State::Validating, &why);
                }
                State::Validating if waited >= BURN_IN_TIME => self.finish_burn_in(now, index),
                _ => {}
            }
        }
    }

    /// Acts on the result of a burn-in test. A pass returns the machine to
    /// service. A failure sends it back for the next action on the ladder, and
    /// a machine with no actions left is quarantined.
    fn finish_burn_in(&mut self, now: u32, index: usize) {
        let name = self.records[index].node.name;
        let failures = self.burn_in.entry(name).or_default();

        if *failures == 0 {
            self.records[index].fault_jobs.clear();
            self.transition(now, index, State::Healthy, "burn-in passed");

            return;
        }

        *failures -= 1;

        if self.records[index].step + 1 == LADDER.len() {
            let why = "burn-in failed after every action; ticket opened";
            self.transition(now, index, State::Quarantined, why);

            return;
        }

        self.records[index].step += 1;
        let why = format!("burn-in failed; {}", LADDER[self.records[index].step]);
        self.transition(now, index, State::Remediating, &why);
    }
}

Step reads the current state and acts on it, with no memory of which signal started the repair. This is the level-triggered style of chapter 1 applied to machines. If the process stopped and started again with the same records, the next call to Step would continue every repair from where it was.

The run

The script holds eleven signals that exercise each rule: one job faulting on three Thors, a memory error on a Spark, a Thor whose GPU stops answering, a restart of the 25 GbE switch, and two machines failing in the same second.

Listing 4.13 The scripted signals and the main loop, in Go. code/go/fleethealth/main.go

// script is the run: the signals in the order they arrive.
var script = []Signal{
	// One job faults on three different Thors.
	{5, "thor-1", "finetune-7b", AppFault},
	{9, "thor-2", "finetune-7b", AppFault},
	{14, "thor-3", "finetune-7b", AppFault},
	// A memory error on a Spark, which one GPU reset repairs.
	{20, "spark-1", "", MemoryECC},
	// A Thor whose GPU stops answering, and which no action repairs.
	{40, "thor-4", "", OffBus},
	// The 25 GbE switch restarts: four Thors lose their link together.
	{60, "thor-1", "", LinkFlap},
	{61, "thor-2", "", LinkFlap},
	{61, "thor-3", "", LinkFlap},
	{62, "thor-5", "", LinkFlap},
	// Two machines fail in the same second while thor-4 is quarantined.
	{80, "spark-2", "", OffBus},
	{80, "orin-nx-1", "", MemoryECC},
}

func main() {
	// thor-4 fails every burn-in; every other machine passes its first.
	p := NewPlane(map[string]int{"thor-4": len(ladder)})
	next := 0

	for now := 0; now <= duration; now++ {
		for ; next < len(script) && script[next].At == now; next++ {
			p.Observe(script[next])
		}

		p.Step(now)
	}

	fmt.Println("final states:")

	for _, r := range p.records {
		fmt.Printf("  %-10s %-16s %s\n", r.node.Name, r.node.Model, r.state)
	}
}

Listing 4.14 The scripted signals and the main loop, in Rust. code/rust/fleethealth/src/main.rs

/// The run: the signals in the order they arrive.
const SCRIPT: [Signal; 11] = [
    // One job faults on three different Thors.
    Signal {
        at: 5,
        node: "thor-1",
        job: "finetune-7b",
        kind: Kind::AppFault,
    },
    Signal {
        at: 9,
        node: "thor-2",
        job: "finetune-7b",
        kind: Kind::AppFault,
    },
    Signal {
        at: 14,
        node: "thor-3",
        job: "finetune-7b",
        kind: Kind::AppFault,
    },
    // A memory error on a Spark, which one GPU reset repairs.
    Signal {
        at: 20,
        node: "spark-1",
        job: "",
        kind: Kind::MemoryEcc,
    },
    // A Thor whose GPU stops answering, and which no action repairs.
    Signal {
        at: 40,
        node: "thor-4",
        job: "",
        kind: Kind::OffBus,
    },
    // The 25 GbE switch restarts: four Thors lose their link together.
    Signal {
        at: 60,
        node: "thor-1",
        job: "",
        kind: Kind::LinkFlap,
    },
    Signal {
        at: 61,
        node: "thor-2",
        job: "",
        kind: Kind::LinkFlap,
    },
    Signal {
        at: 61,
        node: "thor-3",
        job: "",
        kind: Kind::LinkFlap,
    },
    Signal {
        at: 62,
        node: "thor-5",
        job: "",
        kind: Kind::LinkFlap,
    },
    // Two machines fail in the same second while thor-4 is quarantined.
    Signal {
        at: 80,
        node: "spark-2",
        job: "",
        kind: Kind::OffBus,
    },
    Signal {
        at: 80,
        node: "orin-nx-1",
        job: "",
        kind: Kind::MemoryEcc,
    },
];

fn main() {
    // thor-4 fails every burn-in; every other machine passes its first.
    let mut plane = Plane::new(BTreeMap::from([("thor-4", LADDER.len())]));
    let mut signals = SCRIPT.iter().peekable();

    for now in 0..=DURATION {
        while let Some(signal) = signals.next_if(|s| s.at == now) {
            plane.observe(signal);
        }

        plane.step(now);
    }

    println!("final states:");

    for r in &plane.records {
        println!("  {:<10} {:<16} {}", r.node.name, r.node.model, r.state);
    }
}

Both versions print the same log:

$ go run ./fleethealth
  14s  job finetune-7b faulted on 3 machines: the job is at fault, no machine is drained
  20s  spark-1    Healthy     -> Cordoned    uncorrectable ECC
  20s  spark-1    Cordoned    -> Draining    jobs told to checkpoint
  24s  spark-1    Draining    -> Remediating reset the GPU
  27s  spark-1    Remediating -> Validating  burn-in after reset the GPU
  30s  spark-1    Validating  -> Healthy     burn-in passed
  40s  thor-4     Healthy     -> Cordoned    GPU off the bus
  40s  thor-4     Cordoned    -> Draining    jobs told to checkpoint
  44s  thor-4     Draining    -> Remediating reset the GPU
  47s  thor-4     Remediating -> Validating  burn-in after reset the GPU
  50s  thor-4     Validating  -> Remediating burn-in failed; reboot
  53s  thor-4     Remediating -> Validating  burn-in after reboot
  56s  thor-4     Validating  -> Remediating burn-in failed; reimage
  59s  thor-4     Remediating -> Validating  burn-in after reimage
  61s  3 machines on switch-25g flapped together: the link is at fault, no machine is drained
  62s  thor-4     Validating  -> Quarantined burn-in failed after every action; ticket opened
  80s  spark-2    Healthy     -> Cordoned    GPU off the bus
  80s  orin-nx-1  Healthy     -> Cordoned    uncorrectable ECC
  80s  spark-2    Cordoned    -> Draining    jobs told to checkpoint
  81s  orin-nx-1  blocked: 2 machines are already out of service
  84s  spark-2    Draining    -> Remediating reset the GPU
  87s  spark-2    Remediating -> Validating  burn-in after reset the GPU
  90s  spark-2    Validating  -> Healthy     burn-in passed
  90s  orin-nx-1  Cordoned    -> Draining    jobs told to checkpoint
  94s  orin-nx-1  Draining    -> Remediating reset the GPU
  97s  orin-nx-1  Remediating -> Validating  burn-in after reset the GPU
 100s  orin-nx-1  Validating  -> Healthy     burn-in passed
final states:
  thor-1     Jetson AGX Thor  Healthy
  thor-2     Jetson AGX Thor  Healthy
  thor-3     Jetson AGX Thor  Healthy
  thor-4     Jetson AGX Thor  Quarantined
  thor-5     Jetson AGX Thor  Healthy
  spark-1    DGX Spark        Healthy
  spark-2    DGX Spark        Healthy
  orin-nx-1  Jetson Orin NX   Healthy

The log has five episodes.

At second 14 the job finetune-7b has faulted on thor-1, thor-2, and thor-3. The verdict names the job, and all three Thors stay Healthy.

At second 20 spark-1 reports an uncorrectable memory error. It is cordoned and, with the whole budget free, starts to drain in the same second. Four seconds later its jobs have left, the GPU is reset, the burn-in passes, and at second 30 the Spark is back in service. The whole repair took ten seconds and no person was involved.

At second 40 the GPU of thor-4 falls off the bus. The reset, the reboot, and the reimage each end in a failed burn-in, and at second 62 the Thor is quarantined with a ticket. It holds one of the two budget slots from then on.

At seconds 60 to 62 four Thors report link flaps. The third report, at second 61, brings the count for switch-25g to three within the window, and the verdict names the switch. No Thor is cordoned.

At second 80 spark-2 and orin-nx-1 fail together. spark-2 comes first in the fleet and takes the last budget slot. orin-nx-1 is cordoned and reports that it is blocked. It waits ten seconds, until spark-2 passes its burn-in at second 90 and frees the slot, and then goes through the same repair. At no point are more than two of the eight machines out of service.

Animation 4.1 replays the log from second 20 as a board with one column for each state.

A job that runs twice

The lifecycle above covers a machine that still answers. Chapter 3 left us with the harder case: a machine whose Lease has expired. The control plane cannot drain it, because it cannot reach it, and it cannot tell a machine that has lost power from one that has lost only its network link. If it waits, the job on a truly dead machine stays stopped. If it starts the job elsewhere, the job on a machine that was only cut off is now running twice.

Two copies of a training job are harmless while they only compute. They become harmful at the moment they write. Both save checkpoints to the same place, and the copy on the cut-off machine started from an older state. When its link returns, its next checkpoint replaces a newer one from the second copy, and the job’s saved progress moves backwards.

The control plane cannot prevent the second copy from existing, so the protection has to sit at the place where the writes land. The mechanism has two parts:

• Every time the control plane assigns the job to a machine, it increments a counter and gives the new value to that machine. The value is called an epoch. A later assignment always has a higher epoch than an earlier one.
• Every write to the checkpoint store carries the writer’s epoch. The store remembers the highest epoch it has accepted and refuses a write with a lower one.

The second part is called fencing, and the epoch in this role is a fencing token. Once the new copy has written, the old copy is fenced off: each of its writes is refused, and the refusal tells it that it has been replaced, so it can stop.

The store is a single checkpoint and one comparison.

Listing 4.15 The checkpoint store and its epoch check, in Go. code/go/fencing/main.go

// Checkpoint is what a worker saves: the training step it reached and the
// epoch of the assignment under which it was running.
type Checkpoint struct {
	// Step is the training step the checkpoint holds.
	Step int
	// Epoch is the assignment number the writer was given.
	Epoch int
	// Writer is the host that wrote the checkpoint.
	Writer string
}

// ErrFenced is returned to a writer whose epoch is older than one the store
// has already accepted.
var ErrFenced = errors.New("fenced: a newer assignment has written")

// Store holds the latest checkpoint of one job.
type Store struct {
	// Fencing turns the epoch check on.
	Fencing bool
	// latest is the stored checkpoint.
	latest Checkpoint
}

// Write stores a checkpoint. With fencing on, a checkpoint whose epoch is
// lower than the stored one is refused.
func (s *Store) Write(c Checkpoint) error {
	if s.Fencing && c.Epoch < s.latest.Epoch {
		return ErrFenced
	}

	s.latest = c

	return nil
}

Listing 4.16 The checkpoint store and its epoch check, in Rust. code/rust/fencing/src/main.rs

/// What a worker saves: the training step it reached and the epoch of the
/// assignment under which it was running.
#[derive(Clone, Copy, Default)]
struct Checkpoint {
    /// The training step the checkpoint holds.
    step: u32,
    /// The assignment number the writer was given.
    epoch: u32,
    /// The host that wrote the checkpoint.
    writer: &'static str,
}

/// Returned to a writer whose epoch is older than one the store has already
/// accepted.
struct Fenced;

impl fmt::Display for Fenced {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str("fenced: a newer assignment has written")
    }
}

/// Holds the latest checkpoint of one job.
#[derive(Default)]
struct Store {
    /// Turns the epoch check on.
    fencing: bool,
    /// The stored checkpoint.
    latest: Checkpoint,
}

impl Store {
    /// Stores a checkpoint. With fencing on, a checkpoint whose epoch is lower
    /// than the stored one is refused.
    fn write(&mut self, c: Checkpoint) -> Result<(), Fenced> {
        if self.fencing && c.epoch < self.latest.epoch {
            return Err(Fenced);
        }

        self.latest = c;

        Ok(())
    }
}

The control plane is a counter.

Listing 4.17 Handing out epochs, in Go. code/go/fencing/main.go

// ControlPlane hands out assignments of one job. Each assignment gets the
// next epoch, so a later assignment always carries a higher number.
type ControlPlane struct {
	epoch int
}

// Assign gives the job to a host and returns the epoch of the assignment.
func (c *ControlPlane) Assign(host string) int {
	c.epoch++
	fmt.Printf("  control plane: job assigned to %s with epoch %d\n", host, c.epoch)

	return c.epoch
}

Listing 4.18 Handing out epochs, in Rust. code/rust/fencing/src/main.rs

/// Hands out assignments of one job. Each assignment gets the next epoch, so a
/// later assignment always carries a higher number.
#[derive(Default)]
struct ControlPlane {
    epoch: u32,
}

impl ControlPlane {
    /// Gives the job to a host and returns the epoch of the assignment.
    fn assign(&mut self, host: &str) -> u32 {
        self.epoch += 1;
        println!(
            "  control plane: job assigned to {host} with epoch {}",
            self.epoch
        );

        self.epoch
    }
}

The run plays the same sequence against both kinds of store.

Listing 4.19 The sequence of writes, in Go. code/go/fencing/main.go

// save writes one checkpoint and prints what the store answered.
func save(s *Store, host string, epoch, step int) {
	err := s.Write(Checkpoint{Step: step, Epoch: epoch, Writer: host})

	if err != nil {
		fmt.Printf("  %s writes step %d with epoch %d: refused, %v\n", host, step, epoch, err)

		return
	}

	fmt.Printf("  %s writes step %d with epoch %d: stored\n", host, step, epoch)
}

// run plays the sequence against one store.
func run(s *Store) {
	var plane ControlPlane
	first := plane.Assign("thor-1")
	save(s, "thor-1", first, 1000)
	fmt.Println("  thor-1 loses its link and keeps training; its Lease expires")
	second := plane.Assign("thor-2")
	fmt.Printf("  thor-2 resumes from step %d\n", s.latest.Step)
	save(s, "thor-2", second, 2000)
	fmt.Println("  thor-1's link returns; it still holds epoch 1")
	save(s, "thor-1", first, 1500)
	fmt.Printf("  a restart now resumes from step %d, written by %s\n", s.latest.Step, s.latest.Writer)
}

func main() {
	fmt.Println("a store that accepts every write:")
	run(&Store{})
	fmt.Println("a store that checks the epoch:")
	run(&Store{Fencing: true})
}

Listing 4.20 The sequence of writes, in Rust. code/rust/fencing/src/main.rs

/// Writes one checkpoint and prints what the store answered.
fn save(store: &mut Store, writer: &'static str, epoch: u32, step: u32) {
    match store.write(Checkpoint {
        step,
        epoch,
        writer,
    }) {
        Ok(()) => println!("  {writer} writes step {step} with epoch {epoch}: stored"),
        Err(e) => println!("  {writer} writes step {step} with epoch {epoch}: refused, {e}"),
    }
}

/// Plays the sequence against one store.
fn run(mut store: Store) {
    let mut plane = ControlPlane::default();
    let first = plane.assign("thor-1");
    save(&mut store, "thor-1", first, 1000);
    println!("  thor-1 loses its link and keeps training; its Lease expires");
    let second = plane.assign("thor-2");
    println!("  thor-2 resumes from step {}", store.latest.step);
    save(&mut store, "thor-2", second, 2000);
    println!("  thor-1's link returns; it still holds epoch 1");
    save(&mut store, "thor-1", first, 1500);
    println!(
        "  a restart now resumes from step {}, written by {}",
        store.latest.step, store.latest.writer
    );
}

fn main() {
    println!("a store that accepts every write:");
    run(Store::default());
    println!("a store that checks the epoch:");
    run(Store {
        fencing: true,
        ..Store::default()
    });
}

$ go run ./fencing
a store that accepts every write:
  control plane: job assigned to thor-1 with epoch 1
  thor-1 writes step 1000 with epoch 1: stored
  thor-1 loses its link and keeps training; its Lease expires
  control plane: job assigned to thor-2 with epoch 2
  thor-2 resumes from step 1000
  thor-2 writes step 2000 with epoch 2: stored
  thor-1's link returns; it still holds epoch 1
  thor-1 writes step 1500 with epoch 1: stored
  a restart now resumes from step 1500, written by thor-1
a store that checks the epoch:
  control plane: job assigned to thor-1 with epoch 1
  thor-1 writes step 1000 with epoch 1: stored
  thor-1 loses its link and keeps training; its Lease expires
  control plane: job assigned to thor-2 with epoch 2
  thor-2 resumes from step 1000
  thor-2 writes step 2000 with epoch 2: stored
  thor-1's link returns; it still holds epoch 1
  thor-1 writes step 1500 with epoch 1: refused, fenced: a newer assignment has written
  a restart now resumes from step 2000, written by thor-2

In the first run, thor-1’s write of step 1500 lands after thor-2’s write of step 2000, and a restart would resume from step 1500: five hundred steps of training are gone, and the stored checkpoint comes from the copy that should have stopped. In the second run the store refuses that write, and a restart resumes from step 2000.

Kubernetes gives us the pieces for both parts. A Lease’s leaseTransitions field counts how often the Lease has changed hands, and it rises by one with each new holder, so it can serve as the epoch. The check has to be made by whatever stores the checkpoints: an object store that supports conditional writes, a database row with a version column, or a file server that rejects clients holding an old token. The resource version check from chapter 1 is the same idea applied to Kubernetes objects.

The complete programs

Listing 4.21 The complete program in Go. code/go/fleethealth/main.go

// Command fleethealth models the part of a control plane that keeps a fleet of
// GPU machines healthy. It reads failure signals, decides whether each one
// blames a job, a machine, or the network between machines, and moves each
// machine through a lifecycle: cordon, drain, repair, validate, and return.
//
// The fleet is eight machines: five Jetson AGX Thor developer kits on one
// 25 GbE switch, two DGX Sparks joined by a 200 GbE cable, and one Jetson
// Orin NX. The program replays a scripted list of signals without real time
// passing and prints every decision it makes.
//
// Run it from code/go with:
//
//	go run ./fleethealth
package main

import "fmt"

// Node is one machine in the fleet.
type Node struct {
	// Name is the machine's host name.
	Name string
	// Model is the product the machine is.
	Model string
	// Link names the switch or cable that connects the machine to its peers.
	Link string
}

// fleet lists the machines in the order the program reports them.
var fleet = []Node{
	{"thor-1", "Jetson AGX Thor", "switch-25g"},
	{"thor-2", "Jetson AGX Thor", "switch-25g"},
	{"thor-3", "Jetson AGX Thor", "switch-25g"},
	{"thor-4", "Jetson AGX Thor", "switch-25g"},
	{"thor-5", "Jetson AGX Thor", "switch-25g"},
	{"spark-1", "DGX Spark", "cable-200g"},
	{"spark-2", "DGX Spark", "cable-200g"},
	{"orin-nx-1", "Jetson Orin NX", "bench"},
}


// Kind is the class of a failure signal.
type Kind int

// The classes of signal the model understands.
const (
	AppFault  Kind = iota // the GPU rejected work from one program
	MemoryECC             // an uncorrectable memory error on the GPU
	OffBus                // the GPU stopped answering the driver
	LinkFlap              // the machine's network link went down and up
)

// String returns the class as it appears in the log.
func (k Kind) String() string {
	return [...]string{"application fault", "uncorrectable ECC", "GPU off the bus", "link flap"}[k]
}

// Signal is one failure report from a machine.
type Signal struct {
	// At is the second at which the signal arrived.
	At int
	// Node is the machine that reported it.
	Node string
	// Job is the job that was using the GPU, or "" for a signal with no job.
	Job string
	// Kind is the signal's class.
	Kind Kind
}


// State is a machine's place in the repair lifecycle.
type State int

// The lifecycle, in the order a failing machine passes through it.
const (
	Healthy     State = iota // schedulable
	Cordoned                 // no new work; waiting for the disruption budget
	Draining                 // jobs are checkpointing and leaving
	Remediating              // a repair action is running
	Validating               // a burn-in test is running
	Quarantined              // every repair failed; a person has to look
)

// String returns the state's name for the log.
func (s State) String() string {
	return [...]string{"Healthy", "Cordoned", "Draining", "Remediating", "Validating", "Quarantined"}[s]
}

// ladder lists the repair actions from the cheapest to the most thorough.
var ladder = []string{"reset the GPU", "reboot", "reimage"}


// Thresholds and durations, in seconds unless stated otherwise.
const (
	blameAfter   = 3 // distinct jobs (or nodes) before an application fault blames a node (or a job)
	fabricAfter  = 3 // machines on one link that must flap together to blame the link
	fabricWindow = 5 // how close in time those flaps have to be
	budget       = 2 // machines that may be out of service at once
	drainTime    = 4 // time a job gets to write a checkpoint and leave
	repairTime   = 3 // time one repair action takes
	burnInTime   = 3 // time the burn-in test takes
	duration     = 110
)

// record is what the control plane remembers about one machine.
type record struct {
	node  Node
	state State
	// since is the second at which the machine entered its state.
	since int
	// step indexes the next action on the ladder.
	step int
	// faultJobs holds the distinct jobs that have faulted on this machine.
	faultJobs map[string]bool
	// lastFlap is the second of the machine's latest link flap, or -1.
	lastFlap int
}

// Plane is the control plane's state: one record per machine, the nodes each
// job has faulted on, and a scripted answer for each burn-in test.
type Plane struct {
	records []*record
	// jobNodes holds, for each job, the distinct nodes it has faulted on.
	jobNodes map[string]map[string]bool
	// blamed holds the jobs and links already reported, so each is reported once.
	blamed map[string]bool
	// burnIn says how many burn-in tests a machine fails before one passes.
	burnIn map[string]int
}

// NewPlane returns a plane in which every machine is Healthy.
func NewPlane(burnIn map[string]int) *Plane {
	p := &Plane{jobNodes: map[string]map[string]bool{}, blamed: map[string]bool{}, burnIn: burnIn}

	for _, n := range fleet {
		p.records = append(p.records, &record{node: n, faultJobs: map[string]bool{}, lastFlap: -1})
	}

	return p
}

// find returns the record of the named machine.
func (p *Plane) find(name string) *record {
	for _, r := range p.records {
		if r.node.Name == name {
			return r
		}
	}

	panic("unknown node " + name)
}

// move puts a machine in a new state and logs the transition with its reason.
func (p *Plane) move(now int, r *record, to State, why string) {
	fmt.Printf("%4ds  %-10s %-11s -> %-11s %s\n", now, r.node.Name, r.state, to, why)
	r.state, r.since = to, now
}

// Observe takes one signal and decides what it blames. A hardware signal
// blames the machine. An application fault blames the job once the job has
// faulted on three machines, and blames the machine once three jobs have
// faulted on it. Link flaps blame the link when three machines on it flap
// within five seconds of each other.
func (p *Plane) Observe(s Signal) {
	r := p.find(s.Node)

	switch s.Kind {
	case MemoryECC, OffBus:
		p.condemn(s.At, r, s.Kind.String())
	case AppFault:
		r.faultJobs[s.Job] = true

		if p.jobNodes[s.Job] == nil {
			p.jobNodes[s.Job] = map[string]bool{}
		}

		p.jobNodes[s.Job][s.Node] = true

		switch {
		case len(p.jobNodes[s.Job]) >= blameAfter && !p.blamed[s.Job]:
			p.blamed[s.Job] = true
			fmt.Printf("%4ds  job %s faulted on %d machines: the job is at fault, no machine is drained\n",
				s.At, s.Job, len(p.jobNodes[s.Job]))
		case len(r.faultJobs) >= blameAfter:
			p.condemn(s.At, r, fmt.Sprintf("%d different jobs faulted here", len(r.faultJobs)))
		}
	case LinkFlap:
		r.lastFlap = s.At
		together := 0

		for _, other := range p.records {
			if other.node.Link == r.node.Link && other.lastFlap >= 0 && s.At-other.lastFlap <= fabricWindow {
				together++
			}
		}

		if together >= fabricAfter && !p.blamed[r.node.Link] {
			p.blamed[r.node.Link] = true
			fmt.Printf("%4ds  %d machines on %s flapped together: the link is at fault, no machine is drained\n",
				s.At, together, r.node.Link)
		}
	}
}


// condemn cordons a Healthy machine. A machine already under repair stays
// where it is.
func (p *Plane) condemn(now int, r *record, why string) {
	if r.state == Healthy {
		r.step = 0
		p.move(now, r, Cordoned, why)
	}
}

// outOfService counts the machines that hold a slot of the disruption budget:
// every machine that is draining, under repair, in burn-in, or quarantined.
func (p *Plane) outOfService() int {
	n := 0

	for _, r := range p.records {
		if r.state >= Draining {
			n++
		}
	}

	return n
}


// Step advances every machine by one second. A Cordoned machine starts to
// drain when the budget has a free slot. The other states each last a fixed
// time and then move on; a failed burn-in sends the machine back for the next
// action on the ladder, and a machine with no actions left is quarantined.
func (p *Plane) Step(now int) {
	for _, r := range p.records {
		waited := now - r.since

		switch {
		case r.state == Cordoned && p.outOfService() < budget:
			p.move(now, r, Draining, "jobs told to checkpoint")
		case r.state == Cordoned && waited == 1:
			fmt.Printf("%4ds  %-10s blocked: %d machines are already out of service\n", now, r.node.Name, budget)
		case r.state == Draining && waited >= drainTime:
			p.move(now, r, Remediating, ladder[r.step])
		case r.state == Remediating && waited >= repairTime:
			p.move(now, r, Validating, "burn-in after "+ladder[r.step])
		case r.state == Validating && waited >= burnInTime:
			p.finishBurnIn(now, r)
		}
	}
}

// finishBurnIn acts on the result of a burn-in test. A pass returns the
// machine to service. A failure sends it back for the next action on the
// ladder, and a machine with no actions left is quarantined.
func (p *Plane) finishBurnIn(now int, r *record) {
	if p.burnIn[r.node.Name] == 0 {
		r.faultJobs = map[string]bool{}
		p.move(now, r, Healthy, "burn-in passed")

		return
	}

	p.burnIn[r.node.Name]--

	if r.step+1 == len(ladder) {
		p.move(now, r, Quarantined, "burn-in failed after every action; ticket opened")

		return
	}

	r.step++
	p.move(now, r, Remediating, "burn-in failed; "+ladder[r.step])
}


// script is the run: the signals in the order they arrive.
var script = []Signal{
	// One job faults on three different Thors.
	{5, "thor-1", "finetune-7b", AppFault},
	{9, "thor-2", "finetune-7b", AppFault},
	{14, "thor-3", "finetune-7b", AppFault},
	// A memory error on a Spark, which one GPU reset repairs.
	{20, "spark-1", "", MemoryECC},
	// A Thor whose GPU stops answering, and which no action repairs.
	{40, "thor-4", "", OffBus},
	// The 25 GbE switch restarts: four Thors lose their link together.
	{60, "thor-1", "", LinkFlap},
	{61, "thor-2", "", LinkFlap},
	{61, "thor-3", "", LinkFlap},
	{62, "thor-5", "", LinkFlap},
	// Two machines fail in the same second while thor-4 is quarantined.
	{80, "spark-2", "", OffBus},
	{80, "orin-nx-1", "", MemoryECC},
}

func main() {
	// thor-4 fails every burn-in; every other machine passes its first.
	p := NewPlane(map[string]int{"thor-4": len(ladder)})
	next := 0

	for now := 0; now <= duration; now++ {
		for ; next < len(script) && script[next].At == now; next++ {
			p.Observe(script[next])
		}

		p.Step(now)
	}

	fmt.Println("final states:")

	for _, r := range p.records {
		fmt.Printf("  %-10s %-16s %s\n", r.node.Name, r.node.Model, r.state)
	}
}

Listing 4.22 The complete program in Rust. code/rust/fleethealth/src/main.rs

//! Models the part of a control plane that keeps a fleet of GPU machines
//! healthy. It reads failure signals, decides whether each one blames a job, a
//! machine, or the network between machines, and moves each machine through a
//! lifecycle: cordon, drain, repair, validate, and return.
//!
//! The fleet is eight machines: five Jetson AGX Thor developer kits on one
//! 25 GbE switch, two DGX Sparks joined by a 200 GbE cable, and one Jetson
//! Orin NX. The program replays a scripted list of signals without real time
//! passing and prints every decision it makes.
//!
//! Run it from code/rust with:
//!
//! ```text
//! cargo run -p fleethealth
//! ```

use std::collections::{BTreeMap, BTreeSet};
use std::fmt;

/// One machine in the fleet.
struct Node {
    /// The machine's host name.
    name: &'static str,
    /// The product the machine is.
    model: &'static str,
    /// The switch or cable that connects the machine to its peers.
    link: &'static str,
}

/// The machines, in the order the program reports them.
const FLEET: [Node; 8] = [
    Node {
        name: "thor-1",
        model: "Jetson AGX Thor",
        link: "switch-25g",
    },
    Node {
        name: "thor-2",
        model: "Jetson AGX Thor",
        link: "switch-25g",
    },
    Node {
        name: "thor-3",
        model: "Jetson AGX Thor",
        link: "switch-25g",
    },
    Node {
        name: "thor-4",
        model: "Jetson AGX Thor",
        link: "switch-25g",
    },
    Node {
        name: "thor-5",
        model: "Jetson AGX Thor",
        link: "switch-25g",
    },
    Node {
        name: "spark-1",
        model: "DGX Spark",
        link: "cable-200g",
    },
    Node {
        name: "spark-2",
        model: "DGX Spark",
        link: "cable-200g",
    },
    Node {
        name: "orin-nx-1",
        model: "Jetson Orin NX",
        link: "bench",
    },
];

/// The class of a failure signal.
#[derive(Clone, Copy)]
enum Kind {
    /// The GPU rejected work from one program.
    AppFault,
    /// An uncorrectable memory error on the GPU.
    MemoryEcc,
    /// The GPU stopped answering the driver.
    OffBus,
    /// The machine's network link went down and up.
    LinkFlap,
}

impl fmt::Display for Kind {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(match self {
            Kind::AppFault => "application fault",
            Kind::MemoryEcc => "uncorrectable ECC",
            Kind::OffBus => "GPU off the bus",
            Kind::LinkFlap => "link flap",
        })
    }
}

/// One failure report from a machine.
struct Signal {
    /// The second at which the signal arrived.
    at: u32,
    /// The machine that reported it.
    node: &'static str,
    /// The job that was using the GPU, or "" for a signal with no job.
    job: &'static str,
    /// The signal's class.
    kind: Kind,
}

/// A machine's place in the repair lifecycle, in the order a failing machine
/// passes through it.
#[derive(Clone, Copy, Debug, PartialEq, PartialOrd)]
enum State {
    /// Schedulable.
    Healthy,
    /// No new work; waiting for the disruption budget.
    Cordoned,
    /// Jobs are checkpointing and leaving.
    Draining,
    /// A repair action is running.
    Remediating,
    /// A burn-in test is running.
    Validating,
    /// Every repair failed; a person has to look.
    Quarantined,
}

impl fmt::Display for State {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        // `pad` honours the width given in a format string such as `{:<11}`.
        f.pad(&format!("{self:?}"))
    }
}

/// The repair actions, from the cheapest to the most thorough.
const LADDER: [&str; 3] = ["reset the GPU", "reboot", "reimage"];

/// Distinct jobs (or nodes) before an application fault blames a node (or a job).
const BLAME_AFTER: usize = 3;
/// Machines on one link that must flap together to blame the link.
const FABRIC_AFTER: usize = 3;
/// How close in time those flaps have to be, in seconds.
const FABRIC_WINDOW: u32 = 5;
/// Machines that may be out of service at once.
const BUDGET: usize = 2;
/// Seconds a job gets to write a checkpoint and leave.
const DRAIN_TIME: u32 = 4;
/// Seconds one repair action takes.
const REPAIR_TIME: u32 = 3;
/// Seconds the burn-in test takes.
const BURN_IN_TIME: u32 = 3;
/// Seconds simulated.
const DURATION: u32 = 110;

/// What the control plane remembers about one machine.
struct Record {
    node: &'static Node,
    state: State,
    /// The second at which the machine entered its state.
    since: u32,
    /// The index of the next action on the ladder.
    step: usize,
    /// The distinct jobs that have faulted on this machine.
    fault_jobs: BTreeSet<&'static str>,
    /// The second of the machine's latest link flap.
    last_flap: Option<u32>,
}

/// The control plane's state: one record per machine, the nodes each job has
/// faulted on, and a scripted answer for each burn-in test.
struct Plane {
    records: Vec<Record>,
    /// For each job, the distinct nodes it has faulted on.
    job_nodes: BTreeMap<&'static str, BTreeSet<&'static str>>,
    /// The jobs and links already reported, so each is reported once.
    blamed: BTreeSet<&'static str>,
    /// How many burn-in tests a machine fails before one passes.
    burn_in: BTreeMap<&'static str, usize>,
}

impl Plane {
    /// Returns a plane in which every machine is Healthy.
    fn new(burn_in: BTreeMap<&'static str, usize>) -> Self {
        let records = FLEET
            .iter()
            .map(|node| Record {
                node,
                state: State::Healthy,
                since: 0,
                step: 0,
                fault_jobs: BTreeSet::new(),
                last_flap: None,
            })
            .collect();
        Plane {
            records,
            job_nodes: BTreeMap::new(),
            blamed: BTreeSet::new(),
            burn_in,
        }
    }

    /// Returns the index of the named machine's record.
    fn find(&self, name: &str) -> usize {
        self.records
            .iter()
            .position(|r| r.node.name == name)
            .expect("unknown node")
    }

    /// Puts a machine in a new state and logs the transition with its reason.
    fn transition(&mut self, now: u32, index: usize, to: State, why: &str) {
        let r = &mut self.records[index];
        println!(
            "{now:4}s  {:<10} {:<11} -> {:<11} {why}",
            r.node.name, r.state, to
        );
        r.state = to;
        r.since = now;
    }
}

impl Plane {
    /// Takes one signal and decides what it blames. A hardware signal blames
    /// the machine. An application fault blames the job once the job has
    /// faulted on three machines, and blames the machine once three jobs have
    /// faulted on it. Link flaps blame the link when three machines on it flap
    /// within five seconds of each other.
    fn observe(&mut self, s: &Signal) {
        let index = self.find(s.node);

        match s.kind {
            Kind::MemoryEcc | Kind::OffBus => self.condemn(s.at, index, &s.kind.to_string()),
            Kind::AppFault => {
                self.records[index].fault_jobs.insert(s.job);
                let nodes = self.job_nodes.entry(s.job).or_default();
                nodes.insert(s.node);
                let (machines, jobs) = (nodes.len(), self.records[index].fault_jobs.len());

                if machines >= BLAME_AFTER && self.blamed.insert(s.job) {
                    println!(
                        "{:4}s  job {} faulted on {machines} machines: the job is at fault, no machine is drained",
                        s.at, s.job
                    );
                } else if machines < BLAME_AFTER && jobs >= BLAME_AFTER {
                    self.condemn(s.at, index, &format!("{jobs} different jobs faulted here"));
                }
            }
            Kind::LinkFlap => {
                self.records[index].last_flap = Some(s.at);
                let link = self.records[index].node.link;
                let together = self
                    .records
                    .iter()
                    .filter(|r| r.node.link == link)
                    .filter(|r| r.last_flap.is_some_and(|at| s.at - at <= FABRIC_WINDOW))
                    .count();
                if together >= FABRIC_AFTER && self.blamed.insert(link) {
                    println!(
                        "{:4}s  {together} machines on {link} flapped together: the link is at fault, no machine is drained",
                        s.at
                    );
                }
            }
        }
    }
}

impl Plane {
    /// Cordons a Healthy machine. A machine already under repair stays where
    /// it is.
    fn condemn(&mut self, now: u32, index: usize, why: &str) {
        if self.records[index].state == State::Healthy {
            self.records[index].step = 0;
            self.transition(now, index, State::Cordoned, why);
        }
    }
}

impl Plane {
    /// Counts the machines that hold a slot of the disruption budget: every
    /// machine that is draining, under repair, in burn-in, or quarantined.
    fn out_of_service(&self) -> usize {
        self.records
            .iter()
            .filter(|r| r.state >= State::Draining)
            .count()
    }
}

impl Plane {
    /// Advances every machine by one second. A Cordoned machine starts to
    /// drain when the budget has a free slot. The other states each last a
    /// fixed time and then move on.
    fn step(&mut self, now: u32) {
        for index in 0..self.records.len() {
            let (state, step, name) = {
                let r = &self.records[index];
                (r.state, r.step, r.node.name)
            };
            let waited = now - self.records[index].since;

            match state {
                State::Cordoned if self.out_of_service() < BUDGET => {
                    self.transition(now, index, State::Draining, "jobs told to checkpoint");
                }
                State::Cordoned if waited == 1 => {
                    println!(
                        "{now:4}s  {name:<10} blocked: {BUDGET} machines are already out of service"
                    );
                }
                State::Draining if waited >= DRAIN_TIME => {
                    self.transition(now, index, State::Remediating, LADDER[step]);
                }
                State::Remediating if waited >= REPAIR_TIME => {
                    let why = format!("burn-in after {}", LADDER[step]);
                    self.transition(now, index, State::Validating, &why);
                }
                State::Validating if waited >= BURN_IN_TIME => self.finish_burn_in(now, index),
                _ => {}
            }
        }
    }

    /// Acts on the result of a burn-in test. A pass returns the machine to
    /// service. A failure sends it back for the next action on the ladder, and
    /// a machine with no actions left is quarantined.
    fn finish_burn_in(&mut self, now: u32, index: usize) {
        let name = self.records[index].node.name;
        let failures = self.burn_in.entry(name).or_default();

        if *failures == 0 {
            self.records[index].fault_jobs.clear();
            self.transition(now, index, State::Healthy, "burn-in passed");

            return;
        }

        *failures -= 1;

        if self.records[index].step + 1 == LADDER.len() {
            let why = "burn-in failed after every action; ticket opened";
            self.transition(now, index, State::Quarantined, why);

            return;
        }

        self.records[index].step += 1;
        let why = format!("burn-in failed; {}", LADDER[self.records[index].step]);
        self.transition(now, index, State::Remediating, &why);
    }
}

/// The run: the signals in the order they arrive.
const SCRIPT: [Signal; 11] = [
    // One job faults on three different Thors.
    Signal {
        at: 5,
        node: "thor-1",
        job: "finetune-7b",
        kind: Kind::AppFault,
    },
    Signal {
        at: 9,
        node: "thor-2",
        job: "finetune-7b",
        kind: Kind::AppFault,
    },
    Signal {
        at: 14,
        node: "thor-3",
        job: "finetune-7b",
        kind: Kind::AppFault,
    },
    // A memory error on a Spark, which one GPU reset repairs.
    Signal {
        at: 20,
        node: "spark-1",
        job: "",
        kind: Kind::MemoryEcc,
    },
    // A Thor whose GPU stops answering, and which no action repairs.
    Signal {
        at: 40,
        node: "thor-4",
        job: "",
        kind: Kind::OffBus,
    },
    // The 25 GbE switch restarts: four Thors lose their link together.
    Signal {
        at: 60,
        node: "thor-1",
        job: "",
        kind: Kind::LinkFlap,
    },
    Signal {
        at: 61,
        node: "thor-2",
        job: "",
        kind: Kind::LinkFlap,
    },
    Signal {
        at: 61,
        node: "thor-3",
        job: "",
        kind: Kind::LinkFlap,
    },
    Signal {
        at: 62,
        node: "thor-5",
        job: "",
        kind: Kind::LinkFlap,
    },
    // Two machines fail in the same second while thor-4 is quarantined.
    Signal {
        at: 80,
        node: "spark-2",
        job: "",
        kind: Kind::OffBus,
    },
    Signal {
        at: 80,
        node: "orin-nx-1",
        job: "",
        kind: Kind::MemoryEcc,
    },
];

fn main() {
    // thor-4 fails every burn-in; every other machine passes its first.
    let mut plane = Plane::new(BTreeMap::from([("thor-4", LADDER.len())]));
    let mut signals = SCRIPT.iter().peekable();

    for now in 0..=DURATION {
        while let Some(signal) = signals.next_if(|s| s.at == now) {
            plane.observe(signal);
        }

        plane.step(now);
    }

    println!("final states:");

    for r in &plane.records {
        println!("  {:<10} {:<16} {}", r.node.name, r.node.model, r.state);
    }
}

Listing 4.23 The complete program in Go. code/go/fencing/main.go

// Command fencing shows why a checkpoint store has to check who is writing.
//
// A training job runs on thor-1. The Thor loses its network link and keeps
// training. The control plane sees silence, declares thor-1 dead, and starts
// the job again on thor-2 from the last checkpoint. For a while two copies of
// the job run, and both write checkpoints to the same store.
//
// The program plays that sequence twice: against a store that accepts every
// write, and against a store that remembers the highest epoch it has seen
// and refuses writes that carry an older one.
//
// Run it from code/go with:
//
//	go run ./fencing
package main

import (
	"errors"
	"fmt"
)

// Checkpoint is what a worker saves: the training step it reached and the
// epoch of the assignment under which it was running.
type Checkpoint struct {
	// Step is the training step the checkpoint holds.
	Step int
	// Epoch is the assignment number the writer was given.
	Epoch int
	// Writer is the host that wrote the checkpoint.
	Writer string
}

// ErrFenced is returned to a writer whose epoch is older than one the store
// has already accepted.
var ErrFenced = errors.New("fenced: a newer assignment has written")

// Store holds the latest checkpoint of one job.
type Store struct {
	// Fencing turns the epoch check on.
	Fencing bool
	// latest is the stored checkpoint.
	latest Checkpoint
}

// Write stores a checkpoint. With fencing on, a checkpoint whose epoch is
// lower than the stored one is refused.
func (s *Store) Write(c Checkpoint) error {
	if s.Fencing && c.Epoch < s.latest.Epoch {
		return ErrFenced
	}

	s.latest = c

	return nil
}


// ControlPlane hands out assignments of one job. Each assignment gets the
// next epoch, so a later assignment always carries a higher number.
type ControlPlane struct {
	epoch int
}

// Assign gives the job to a host and returns the epoch of the assignment.
func (c *ControlPlane) Assign(host string) int {
	c.epoch++
	fmt.Printf("  control plane: job assigned to %s with epoch %d\n", host, c.epoch)

	return c.epoch
}


// save writes one checkpoint and prints what the store answered.
func save(s *Store, host string, epoch, step int) {
	err := s.Write(Checkpoint{Step: step, Epoch: epoch, Writer: host})

	if err != nil {
		fmt.Printf("  %s writes step %d with epoch %d: refused, %v\n", host, step, epoch, err)

		return
	}

	fmt.Printf("  %s writes step %d with epoch %d: stored\n", host, step, epoch)
}

// run plays the sequence against one store.
func run(s *Store) {
	var plane ControlPlane
	first := plane.Assign("thor-1")
	save(s, "thor-1", first, 1000)
	fmt.Println("  thor-1 loses its link and keeps training; its Lease expires")
	second := plane.Assign("thor-2")
	fmt.Printf("  thor-2 resumes from step %d\n", s.latest.Step)
	save(s, "thor-2", second, 2000)
	fmt.Println("  thor-1's link returns; it still holds epoch 1")
	save(s, "thor-1", first, 1500)
	fmt.Printf("  a restart now resumes from step %d, written by %s\n", s.latest.Step, s.latest.Writer)
}

func main() {
	fmt.Println("a store that accepts every write:")
	run(&Store{})
	fmt.Println("a store that checks the epoch:")
	run(&Store{Fencing: true})
}

Listing 4.24 The complete program in Rust. code/rust/fencing/src/main.rs

//! Shows why a checkpoint store has to check who is writing.
//!
//! A training job runs on thor-1. The Thor loses its network link and keeps
//! training. The control plane sees silence, declares thor-1 dead, and starts
//! the job again on thor-2 from the last checkpoint. For a while two copies of
//! the job run, and both write checkpoints to the same store.
//!
//! The program plays that sequence twice: against a store that accepts every
//! write, and against a store that remembers the highest epoch it has seen
//! and refuses writes that carry an older one.
//!
//! Run it from code/rust with:
//!
//! ```text
//! cargo run -p fencing
//! ```

use std::fmt;

/// What a worker saves: the training step it reached and the epoch of the
/// assignment under which it was running.
#[derive(Clone, Copy, Default)]
struct Checkpoint {
    /// The training step the checkpoint holds.
    step: u32,
    /// The assignment number the writer was given.
    epoch: u32,
    /// The host that wrote the checkpoint.
    writer: &'static str,
}

/// Returned to a writer whose epoch is older than one the store has already
/// accepted.
struct Fenced;

impl fmt::Display for Fenced {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str("fenced: a newer assignment has written")
    }
}

/// Holds the latest checkpoint of one job.
#[derive(Default)]
struct Store {
    /// Turns the epoch check on.
    fencing: bool,
    /// The stored checkpoint.
    latest: Checkpoint,
}

impl Store {
    /// Stores a checkpoint. With fencing on, a checkpoint whose epoch is lower
    /// than the stored one is refused.
    fn write(&mut self, c: Checkpoint) -> Result<(), Fenced> {
        if self.fencing && c.epoch < self.latest.epoch {
            return Err(Fenced);
        }

        self.latest = c;

        Ok(())
    }
}

/// Hands out assignments of one job. Each assignment gets the next epoch, so a
/// later assignment always carries a higher number.
#[derive(Default)]
struct ControlPlane {
    epoch: u32,
}

impl ControlPlane {
    /// Gives the job to a host and returns the epoch of the assignment.
    fn assign(&mut self, host: &str) -> u32 {
        self.epoch += 1;
        println!(
            "  control plane: job assigned to {host} with epoch {}",
            self.epoch
        );

        self.epoch
    }
}

/// Writes one checkpoint and prints what the store answered.
fn save(store: &mut Store, writer: &'static str, epoch: u32, step: u32) {
    match store.write(Checkpoint {
        step,
        epoch,
        writer,
    }) {
        Ok(()) => println!("  {writer} writes step {step} with epoch {epoch}: stored"),
        Err(e) => println!("  {writer} writes step {step} with epoch {epoch}: refused, {e}"),
    }
}

/// Plays the sequence against one store.
fn run(mut store: Store) {
    let mut plane = ControlPlane::default();
    let first = plane.assign("thor-1");
    save(&mut store, "thor-1", first, 1000);
    println!("  thor-1 loses its link and keeps training; its Lease expires");
    let second = plane.assign("thor-2");
    println!("  thor-2 resumes from step {}", store.latest.step);
    save(&mut store, "thor-2", second, 2000);
    println!("  thor-1's link returns; it still holds epoch 1");
    save(&mut store, "thor-1", first, 1500);
    println!(
        "  a restart now resumes from step {}, written by {}",
        store.latest.step, store.latest.writer
    );
}

fn main() {
    println!("a store that accepts every write:");
    run(Store::default());
    println!("a store that checks the epoch:");
    run(Store {
        fencing: true,
        ..Store::default()
    });
}

References

• NVIDIA, Xid errors: https://docs.nvidia.com/deploy/xid-errors/
• NVIDIA, DGX Spark hardware overview: https://docs.nvidia.com/dgx/dgx-spark/hardware.html
• NVIDIA, Jetson modules technical specifications: https://developer.nvidia.com/embedded/jetson-modules
• Kubernetes, Safely drain a node: https://kubernetes.io/docs/tasks/administer-cluster/safely-drain-node/
• Kubernetes, API-initiated eviction: https://kubernetes.io/docs/concepts/scheduling-eviction/api-eviction/
• Kubernetes, Device plugins: https://kubernetes.io/docs/concepts/extend-kubernetes/compute-storage-net/device-plugins/
• M. Kleppmann, “How to do distributed locking”, 2016: https://martin.kleppmann.com/2016/02/08/how-to-do-distributed-locking.html

5. Kubernetes in a new cloud account

So far our cluster has existed before each chapter began. In this chapter, we will look at where a cluster comes from. We start with an account at a provider of machines, which holds nothing, and end with a set of machines that the API server knows as nodes and the scheduler can place GPU pods on.

Two separate things have to happen. Something has to ask the provider for machines and keep asking until the right number exist, and each machine, once it boots, has to prove to the cluster that it is allowed to join. We will take the second first, because we can watch it on the cluster we already have. Then we will write provisioner, in Go and in Rust, a model of the controller that asks for machines, and use it to study the one failure that makes this job harder than it looks: a request that succeeds while its reply is lost.

What the account gives us

A provider of machines, whether a public cloud or the team that racks hardware in a company’s own data center, offers the same small interface. We can ask for a machine of a given model, list the machines we have, and give one back. The account also has a quota for each model: the largest number of machines of that model it may hold at once.

Our fleet from chapter 4 fits this picture exactly, and its quotas are physical. There are five Jetson AGX Thors on the shelf and two DGX Sparks on the desk. A request for a sixth Thor fails for the same reason a request for a sixth GPU instance fails in a cloud account whose quota is five: the provider has no more to give us.

A new machine is a computer with an operating system and a network address. It knows nothing about our cluster, and our cluster knows nothing about it.

How a machine becomes a node

A machine becomes a node when the kubelet on it registers with the API server. From chapter 1 we know that every request to the API server is an HTTP request, and the API server requires each one to identify its sender. A kubelet identifies itself with a client certificate: a file that states a name, here the node’s name, and carries a signature from the cluster’s certificate authority (CA), a key pair the cluster created for itself when it was first set up. The API server trusts any certificate that its CA signed.

A new machine has no such certificate, and it cannot be handed the CA’s private key, since any holder of that key can impersonate every node in the cluster. The machine is given something much weaker instead, and trades it for a certificate. Figure 5.1 shows the exchange.

The weaker thing is a bootstrap token: a short string that the API server accepts for a limited time and for one purpose. Our kind cluster was created with kubeadm, the standard tool for setting up a cluster, and kubeadm can print the command a new machine would run to join. We have cut the two secret values from the output:

$ docker exec gpu-lab-control-plane kubeadm token create --ttl 10m --print-join-command
kubeadm join gpu-lab-control-plane:6443 --token <token> --discovery-token-ca-cert-hash sha256:<hash>

The command carries three facts, and each one answers a question the new machine has.

• gpu-lab-control-plane:6443 is where the API server listens.
• The token is what the machine will show to prove it was invited. We asked for it to expire after ten minutes.
• The hash is a fingerprint of the cluster’s CA public key. The machine uses it in the other direction: before it sends the token anywhere, it checks that the server it reached presents a certificate signed by the CA with this fingerprint. A machine without the hash could be tricked into sending its token to an impostor.

With the token accepted, the kubelet generates a key pair of its own and submits a CertificateSigningRequest, a Kubernetes object that says “please sign a certificate that names me system:node:thor-3”. The token is permitted to do only this. A controller in the control plane checks that the request comes from a bootstrap token and asks for a node identity, approves it, and signs it with the CA. The kubelet downloads the certificate and from then on uses it for every request. The private key was generated on the machine and has stayed there.

With its certificate, the kubelet creates the two objects we already know: a Node object, and the Lease from chapter 3 that it will renew every ten seconds.

What a node reports about itself

The Node object is where the kubelet tells the scheduler what the machine has. Here is the capacity our Orin NX reports, with the huge-page lines removed:

$ kubectl get node gpu-lab-control-plane -o jsonpath='{.status.capacity}'
{"cpu":"8","ephemeral-storage":"1966327652352","memory":"16031232Ki","nvidia.com/gpu":"1","pods":"110"}

Eight CPU cores, 16,031,232 KiB of memory, which is the board’s 15.3 GiB, and one nvidia.com/gpu, reported by the device plugin from Appendix A. The scheduler subtracts what running pods have requested from these numbers to find out what is free.

The Node object also carries labels: pairs of a key and a value that describe the machine.

$ kubectl get node gpu-lab-control-plane -o jsonpath='{.metadata.labels}'
{"beta.kubernetes.io/arch":"arm64","beta.kubernetes.io/os":"linux","kubernetes.io/arch":"arm64",
"kubernetes.io/hostname":"gpu-lab-control-plane","kubernetes.io/os":"linux",
"node-role.kubernetes.io/control-plane":""}

The kubelet sets the architecture, the operating system, and the host name by itself. Nothing here says which GPU the machine has. A count of nvidia.com/gpu: 1 is the same on an Orin NX with 16 GB and on a Thor with 128 GB, and a job that needs 60 GB of GPU memory fits on only one of them. On clusters that the GPU Operator supports, a component called GPU Feature Discovery adds labels such as nvidia.com/gpu.product and nvidia.com/gpu.memory to each node. On our fleet we would set a label by hand as each machine joins:

$ kubectl label node thor-3 gpucloud.dev/model=jetson-agx-thor

A pod then restricts itself to machines with that label through the nodeSelector field of its spec, and the scheduler considers only nodes whose labels match. Chapter 8 uses labels of this kind to place jobs.

Asking for machines

Now for the first half: getting the machines to exist. We describe what we want as pools. A pool is a model and a count: five Thors, two Sparks. The controller’s job is to make the account hold exactly that.

This is reconciliation again, with the provider’s interface in place of the API server. The controller compares the wanted count with the machines it has and creates the difference. The difficulty lies in a property of the create call that the calls in chapter 1 did not have. Creating a machine changes the world outside the cluster, costs money while the machine runs, and counts against a quota.

Consider what the controller knows after a create call. There are three outcomes:

1. A reply arrives with the new machine’s identifier. The machine exists.
2. A reply arrives with an error such as “quota exceeded”. No machine was created.
3. No reply arrives before the controller’s timeout.

The third outcome is the awkward one, because it covers two different situations that look identical from the controller’s side. The request may have been lost on the way, in which case nothing was created. Or the provider may have created the machine and the reply was lost on the way back. Figure 5.2 follows the second situation.

If the controller assumes nothing was created and sends another create, it may end up with two machines where it wanted one. It owns the second and has no record of the first, which runs, costs money, and uses up quota until someone finds it. If the controller assumes the machine was created, it has no identifier to manage it by.

A token that names the request

The remedy is to let the controller name each request. The controller attaches a client token to the call: a string it chose, which identifies this particular machine-to-be. The provider stores the token with the machine it creates. When a second call arrives with a token it has already seen, the provider creates nothing and returns the existing machine.

A call with this property is called idempotent: sending it twice has the same effect as sending it once. With an idempotent create, the controller needs no way to tell the two timeout situations apart. It sends the same call again, and the answer is the same machine either way.

The token has to be the same on the retry, so it cannot be random. Our controller derives it from the pool’s name and the slot it is filling: the third Thor’s token is thor-3, on every attempt. Real providers work the same way. The call that creates a virtual machine in Amazon EC2 takes a ClientToken parameter for exactly this purpose.

Building provisioner

The provider

The provider is the account. Create does three things in order. It looks for a machine that already carries the caller’s token and, if it finds one, returns it. It checks the quota. Then it creates the machine, and only after that does the run get the chance to lose the reply. That ordering is what makes the lost reply dangerous: the machine is in the list before the caller learns anything.

Listing 5.1 The provider and its create call, in Go. code/go/provisioner/main.go

// Machine is one machine the provider has created.
type Machine struct {
	// ID is the provider's identifier for the machine.
	ID string
	// Model is the kind of machine.
	Model string
	// Token is the client token of the call that created it, or "".
	Token string
}

// Errors a create call can return.
var (
	// ErrTimeout means no reply arrived. The machine may or may not exist.
	ErrTimeout = errors.New("timed out")
	// ErrQuota means the account already holds as many machines of this model
	// as it is allowed.
	ErrQuota = errors.New("quota exceeded")
)

// Provider is the account that machines are created in.
type Provider struct {
	// quota is the number of machines of each model the account may hold.
	quota map[string]int
	// machines lists every machine that exists, in creation order.
	machines []Machine
	// calls counts create calls, so the run can lose the reply to one of them.
	calls int
	// loseReply is the number of the call whose reply is lost.
	loseReply int
}

// count returns how many machines of a model exist.
func (p *Provider) count(model string) int {
	n := 0

	for _, m := range p.machines {
		if m.Model == model {
			n++
		}
	}

	return n
}

// Create makes one machine. A call that carries a token already used by an
// existing machine returns that machine and creates nothing.
func (p *Provider) Create(model, token string) (Machine, error) {
	p.calls++

	if token != "" {
		for _, m := range p.machines {
			if m.Token == token {
				return m, nil
			}
		}
	}

	if p.count(model) >= p.quota[model] {
		return Machine{}, ErrQuota
	}

	m := Machine{ID: fmt.Sprintf("m-%02d", len(p.machines)+1), Model: model, Token: token}
	p.machines = append(p.machines, m)

	if p.calls == p.loseReply {
		return Machine{}, ErrTimeout
	}

	return m, nil
}

Listing 5.2 The provider and its create call, in Rust. code/rust/provisioner/src/main.rs

/// One machine the provider has created.
#[derive(Clone)]
struct Machine {
    /// The provider's identifier for the machine.
    id: String,
    /// The kind of machine.
    model: &'static str,
    /// The client token of the call that created it, if it carried one.
    token: Option<String>,
}

/// The ways a create call can fail.
enum CreateError {
    /// No reply arrived. The machine may or may not exist.
    Timeout,
    /// The account already holds as many machines of this model as it is
    /// allowed.
    Quota,
}

impl fmt::Display for CreateError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(match self {
            CreateError::Timeout => "timed out",
            CreateError::Quota => "quota exceeded",
        })
    }
}

/// The account that machines are created in.
struct Provider {
    /// The number of machines of each model the account may hold.
    quota: BTreeMap<&'static str, usize>,
    /// Every machine that exists, in creation order.
    machines: Vec<Machine>,
    /// Counts create calls, so the run can lose the reply to one of them.
    calls: usize,
    /// The number of the call whose reply is lost.
    lose_reply: usize,
}

impl Provider {
    /// Returns how many machines of a model exist.
    fn count(&self, model: &str) -> usize {
        self.machines.iter().filter(|m| m.model == model).count()
    }

    /// Makes one machine. A call that carries a token already used by an
    /// existing machine returns that machine and creates nothing.
    fn create(&mut self, model: &'static str, token: Option<&str>) -> Result<Machine, CreateError> {
        self.calls += 1;
        if let Some(token) = token
            && let Some(existing) = self
                .machines
                .iter()
                .find(|m| m.token.as_deref() == Some(token))
        {
            return Ok(existing.clone());
        }

        if self.count(model) >= self.quota[model] {
            return Err(CreateError::Quota);
        }

        let machine = Machine {
            id: format!("m-{:02}", self.machines.len() + 1),
            model,
            token: token.map(str::to_owned),
        };
        self.machines.push(machine.clone());

        if self.calls == self.lose_reply {
            return Err(CreateError::Timeout);
        }

        Ok(machine)
    }
}

The Rust version represents a missing token as None, and the lookup is an if let ... && let ... chain, the same construct we used in section 1.1: the first let unwraps the caller’s token, and the second searches for a machine that carries it.

The controller

Reconcile works on one pool. It names the next empty slot, builds the token from that name when tokens are switched on, and calls the provider. A failed call ends the pass for this pool, and the controller records nothing for it. The next pass starts from the same slot.

Listing 5.3 The pool and the controller's reconcile, in Go. code/go/provisioner/main.go

// Pool is a group of identical machines and the size its owner wants.
type Pool struct {
	// Name prefixes the name of each slot: thor-1, thor-2, and so on.
	Name string
	// Model is the kind of machine in the pool.
	Model string
	// Want is the number of machines the pool should have.
	Want int
}

// Controller brings each pool to its wanted size.
type Controller struct {
	provider *Provider
	// tokens says whether create calls carry a client token.
	tokens bool
	// owned maps a pool's name to the machines the controller knows it has.
	owned map[string][]Machine
}

// Reconcile fills the empty slots of one pool, in order, and stops at the
// first call that fails. It returns true when the pool is full.
func (c *Controller) Reconcile(pool Pool) bool {
	for len(c.owned[pool.Name]) < pool.Want {
		slot := fmt.Sprintf("%s-%d", pool.Name, len(c.owned[pool.Name])+1)
		token := ""

		if c.tokens {
			token = slot
		}

		m, err := c.provider.Create(pool.Model, token)

		if err != nil {
			fmt.Printf("    %-8s create failed: %v\n", slot, err)

			return false
		}

		fmt.Printf("    %-8s is machine %s\n", slot, m.ID)
		c.owned[pool.Name] = append(c.owned[pool.Name], m)
	}

	return true
}

Listing 5.4 The pool and the controller's reconcile, in Rust. code/rust/provisioner/src/main.rs

/// A group of identical machines and the size its owner wants.
struct Pool {
    /// Prefixes the name of each slot: thor-1, thor-2, and so on.
    name: &'static str,
    /// The kind of machine in the pool.
    model: &'static str,
    /// The number of machines the pool should have.
    want: usize,
}

/// Brings each pool to its wanted size.
struct Controller {
    provider: Provider,
    /// Whether create calls carry a client token.
    tokens: bool,
    /// Maps a pool's name to the machines the controller knows it has.
    owned: BTreeMap<&'static str, Vec<Machine>>,
}

impl Controller {
    /// Fills the empty slots of one pool, in order, and stops at the first
    /// call that fails. It returns true when the pool is full.
    fn reconcile(&mut self, pool: &Pool) -> bool {
        let owned = self.owned.entry(pool.name).or_default();

        while owned.len() < pool.want {
            let slot = format!("{}-{}", pool.name, owned.len() + 1);
            let token = self.tokens.then_some(slot.as_str());

            match self.provider.create(pool.model, token) {
                Ok(machine) => {
                    println!("    {slot:<8} is machine {}", machine.id);
                    owned.push(machine);
                }
                Err(e) => {
                    println!("    {slot:<8} create failed: {e}");

                    return false;
                }
            }
        }

        true
    }
}

The controller counts only what it has been told. owned grows by one machine for each successful reply, and Reconcile reads the pool’s size from it. This is where a lost reply does its damage in the first run, as the output will show.

The run

The run asks for five Thors and two Sparks, with quotas of exactly five and two, and loses the reply to the third create call. It allows three passes and then prints what the controller owns beside what the provider holds.

Listing 5.5 The pools, the three passes, and the final comparison, in Go. code/go/provisioner/main.go

// pools is the fleet we ask for: the Thors on the shelf and the two Sparks.
var pools = []Pool{
	{"thor", "Jetson AGX Thor", 5},
	{"spark", "DGX Spark", 2},
}

// run reconciles every pool for three passes and then compares what the
// controller owns with what the provider holds.
func run(tokens bool) {
	provider := &Provider{quota: map[string]int{"Jetson AGX Thor": 5, "DGX Spark": 2}, loseReply: 3}
	c := &Controller{provider: provider, tokens: tokens, owned: map[string][]Machine{}}

	for pass := 1; pass <= 3; pass++ {
		fmt.Printf("  pass %d\n", pass)
		full := true

		for _, pool := range pools {
			full = c.Reconcile(pool) && full
		}

		if full {
			fmt.Println("    every pool is full")
			break
		}
	}

	known := map[string]bool{}

	for _, pool := range pools {
		for _, m := range c.owned[pool.Name] {
			known[m.ID] = true
		}

		fmt.Printf("  %-5s wanted %d, controller owns %d, provider holds %d\n",
			pool.Name, pool.Want, len(c.owned[pool.Name]), provider.count(pool.Model))
	}

	for _, m := range provider.machines {
		if !known[m.ID] {
			fmt.Printf("  machine %s (%s) is running and belongs to no pool\n", m.ID, m.Model)
		}
	}
}

func main() {
	fmt.Println("create calls without a client token:")
	run(false)
	fmt.Println("create calls with the slot name as the client token:")
	run(true)
}

Listing 5.6 The pools, the three passes, and the final comparison, in Rust. code/rust/provisioner/src/main.rs

/// The fleet we ask for: the Thors on the shelf and the two Sparks.
const POOLS: [Pool; 2] = [
    Pool {
        name: "thor",
        model: "Jetson AGX Thor",
        want: 5,
    },
    Pool {
        name: "spark",
        model: "DGX Spark",
        want: 2,
    },
];

/// Reconciles every pool for three passes and then compares what the
/// controller owns with what the provider holds.
fn run(tokens: bool) {
    let provider = Provider {
        quota: BTreeMap::from([("Jetson AGX Thor", 5), ("DGX Spark", 2)]),
        machines: Vec::new(),
        calls: 0,
        lose_reply: 3,
    };
    let mut c = Controller {
        provider,
        tokens,
        owned: BTreeMap::new(),
    };

    for pass in 1..=3 {
        println!("  pass {pass}");
        let mut full = true;

        for pool in &POOLS {
            full = c.reconcile(pool) && full;
        }

        if full {
            println!("    every pool is full");
            break;
        }
    }

    let mut known = BTreeSet::new();

    for pool in &POOLS {
        let owned = &c.owned[pool.name];
        known.extend(owned.iter().map(|m| m.id.as_str()));
        println!(
            "  {:<5} wanted {}, controller owns {}, provider holds {}",
            pool.name,
            pool.want,
            owned.len(),
            c.provider.count(pool.model)
        );
    }

    for m in c
        .provider
        .machines
        .iter()
        .filter(|m| !known.contains(m.id.as_str()))
    {
        println!(
            "  machine {} ({}) is running and belongs to no pool",
            m.id, m.model
        );
    }
}

fn main() {
    println!("create calls without a client token:");
    run(false);
    println!("create calls with the slot name as the client token:");
    run(true);
}

$ go run ./provisioner
create calls without a client token:
  pass 1
    thor-1   is machine m-01
    thor-2   is machine m-02
    thor-3   create failed: timed out
    spark-1  is machine m-04
    spark-2  is machine m-05
  pass 2
    thor-3   is machine m-06
    thor-4   is machine m-07
    thor-5   create failed: quota exceeded
  pass 3
    thor-5   create failed: quota exceeded
  thor  wanted 5, controller owns 4, provider holds 5
  spark wanted 2, controller owns 2, provider holds 2
  machine m-03 (Jetson AGX Thor) is running and belongs to no pool
create calls with the slot name as the client token:
  pass 1
    thor-1   is machine m-01
    thor-2   is machine m-02
    thor-3   create failed: timed out
    spark-1  is machine m-04
    spark-2  is machine m-05
  pass 2
    thor-3   is machine m-03
    thor-4   is machine m-06
    thor-5   is machine m-07
    every pool is full
  thor  wanted 5, controller owns 5, provider holds 5
  spark wanted 2, controller owns 2, provider holds 2

Both runs begin the same way. The third call creates machine m-03 and its reply is lost, so the controller reports a timeout for thor-3 and moves on to the Sparks.

In the first run, pass 2 sends a fresh create for thor-3. The provider has no way to connect it with the earlier call and creates m-06. The controller fills thor-4 with m-07, and then the create for thor-5 fails: the account now holds five Thors, which is the whole quota. Pass 3 fails in the same way, and so would every later pass. The controller owns four Thors and has no way to get a fifth, while m-03 runs with no owner. The pool is one machine short and the account is full.

In the second run, pass 2 sends the create for thor-3 again with the token thor-3. The provider finds m-03 under that token and returns it. The controller then creates thor-4 and thor-5, and every pool is full after two passes, with nothing left over.

Animation 5.1 replays both runs for the pool of Thors.

The same design as Kubernetes objects

Cluster API is the Kubernetes project that manages clusters with Kubernetes objects. It runs in one cluster, called the management cluster, and creates others. Its objects map onto our model directly:

A Machine object is created first and the provider’s machine second, so the object’s own name can serve as the client token, and a controller that restarts finds every slot it was working on by listing Machine objects. The record of ownership lives in the API server, where the resource versions of chapter 1 protect it, and the provider’s account is reconciled towards it.

The complete program

Listing 5.7 The complete program in Go. code/go/provisioner/main.go

// Command provisioner models a controller that asks a provider for machines
// until a pool has the size its owner asked for.
//
// The provider's create call can fail in an awkward way: the machine is
// created and the reply is lost. The program runs the same controller twice.
// The first run retries a lost reply with a plain create call. The second
// sends a client token with every call, so a retry returns the machine that
// the first attempt created.
//
// Run it from code/go with:
//
//	go run ./provisioner
package main

import (
	"errors"
	"fmt"
)

// Machine is one machine the provider has created.
type Machine struct {
	// ID is the provider's identifier for the machine.
	ID string
	// Model is the kind of machine.
	Model string
	// Token is the client token of the call that created it, or "".
	Token string
}

// Errors a create call can return.
var (
	// ErrTimeout means no reply arrived. The machine may or may not exist.
	ErrTimeout = errors.New("timed out")
	// ErrQuota means the account already holds as many machines of this model
	// as it is allowed.
	ErrQuota = errors.New("quota exceeded")
)

// Provider is the account that machines are created in.
type Provider struct {
	// quota is the number of machines of each model the account may hold.
	quota map[string]int
	// machines lists every machine that exists, in creation order.
	machines []Machine
	// calls counts create calls, so the run can lose the reply to one of them.
	calls int
	// loseReply is the number of the call whose reply is lost.
	loseReply int
}

// count returns how many machines of a model exist.
func (p *Provider) count(model string) int {
	n := 0

	for _, m := range p.machines {
		if m.Model == model {
			n++
		}
	}

	return n
}

// Create makes one machine. A call that carries a token already used by an
// existing machine returns that machine and creates nothing.
func (p *Provider) Create(model, token string) (Machine, error) {
	p.calls++

	if token != "" {
		for _, m := range p.machines {
			if m.Token == token {
				return m, nil
			}
		}
	}

	if p.count(model) >= p.quota[model] {
		return Machine{}, ErrQuota
	}

	m := Machine{ID: fmt.Sprintf("m-%02d", len(p.machines)+1), Model: model, Token: token}
	p.machines = append(p.machines, m)

	if p.calls == p.loseReply {
		return Machine{}, ErrTimeout
	}

	return m, nil
}


// Pool is a group of identical machines and the size its owner wants.
type Pool struct {
	// Name prefixes the name of each slot: thor-1, thor-2, and so on.
	Name string
	// Model is the kind of machine in the pool.
	Model string
	// Want is the number of machines the pool should have.
	Want int
}

// Controller brings each pool to its wanted size.
type Controller struct {
	provider *Provider
	// tokens says whether create calls carry a client token.
	tokens bool
	// owned maps a pool's name to the machines the controller knows it has.
	owned map[string][]Machine
}

// Reconcile fills the empty slots of one pool, in order, and stops at the
// first call that fails. It returns true when the pool is full.
func (c *Controller) Reconcile(pool Pool) bool {
	for len(c.owned[pool.Name]) < pool.Want {
		slot := fmt.Sprintf("%s-%d", pool.Name, len(c.owned[pool.Name])+1)
		token := ""

		if c.tokens {
			token = slot
		}

		m, err := c.provider.Create(pool.Model, token)

		if err != nil {
			fmt.Printf("    %-8s create failed: %v\n", slot, err)

			return false
		}

		fmt.Printf("    %-8s is machine %s\n", slot, m.ID)
		c.owned[pool.Name] = append(c.owned[pool.Name], m)
	}

	return true
}


// pools is the fleet we ask for: the Thors on the shelf and the two Sparks.
var pools = []Pool{
	{"thor", "Jetson AGX Thor", 5},
	{"spark", "DGX Spark", 2},
}

// run reconciles every pool for three passes and then compares what the
// controller owns with what the provider holds.
func run(tokens bool) {
	provider := &Provider{quota: map[string]int{"Jetson AGX Thor": 5, "DGX Spark": 2}, loseReply: 3}
	c := &Controller{provider: provider, tokens: tokens, owned: map[string][]Machine{}}

	for pass := 1; pass <= 3; pass++ {
		fmt.Printf("  pass %d\n", pass)
		full := true

		for _, pool := range pools {
			full = c.Reconcile(pool) && full
		}

		if full {
			fmt.Println("    every pool is full")
			break
		}
	}

	known := map[string]bool{}

	for _, pool := range pools {
		for _, m := range c.owned[pool.Name] {
			known[m.ID] = true
		}

		fmt.Printf("  %-5s wanted %d, controller owns %d, provider holds %d\n",
			pool.Name, pool.Want, len(c.owned[pool.Name]), provider.count(pool.Model))
	}

	for _, m := range provider.machines {
		if !known[m.ID] {
			fmt.Printf("  machine %s (%s) is running and belongs to no pool\n", m.ID, m.Model)
		}
	}
}

func main() {
	fmt.Println("create calls without a client token:")
	run(false)
	fmt.Println("create calls with the slot name as the client token:")
	run(true)
}

Listing 5.8 The complete program in Rust. code/rust/provisioner/src/main.rs

//! Models a controller that asks a provider for machines until a pool has the
//! size its owner asked for.
//!
//! The provider's create call can fail in an awkward way: the machine is
//! created and the reply is lost. The program runs the same controller twice.
//! The first run retries a lost reply with a plain create call. The second
//! sends a client token with every call, so a retry returns the machine that
//! the first attempt created.
//!
//! Run it from code/rust with:
//!
//! ```text
//! cargo run -p provisioner
//! ```

use std::collections::{BTreeMap, BTreeSet};
use std::fmt;

/// One machine the provider has created.
#[derive(Clone)]
struct Machine {
    /// The provider's identifier for the machine.
    id: String,
    /// The kind of machine.
    model: &'static str,
    /// The client token of the call that created it, if it carried one.
    token: Option<String>,
}

/// The ways a create call can fail.
enum CreateError {
    /// No reply arrived. The machine may or may not exist.
    Timeout,
    /// The account already holds as many machines of this model as it is
    /// allowed.
    Quota,
}

impl fmt::Display for CreateError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(match self {
            CreateError::Timeout => "timed out",
            CreateError::Quota => "quota exceeded",
        })
    }
}

/// The account that machines are created in.
struct Provider {
    /// The number of machines of each model the account may hold.
    quota: BTreeMap<&'static str, usize>,
    /// Every machine that exists, in creation order.
    machines: Vec<Machine>,
    /// Counts create calls, so the run can lose the reply to one of them.
    calls: usize,
    /// The number of the call whose reply is lost.
    lose_reply: usize,
}

impl Provider {
    /// Returns how many machines of a model exist.
    fn count(&self, model: &str) -> usize {
        self.machines.iter().filter(|m| m.model == model).count()
    }

    /// Makes one machine. A call that carries a token already used by an
    /// existing machine returns that machine and creates nothing.
    fn create(&mut self, model: &'static str, token: Option<&str>) -> Result<Machine, CreateError> {
        self.calls += 1;
        if let Some(token) = token
            && let Some(existing) = self
                .machines
                .iter()
                .find(|m| m.token.as_deref() == Some(token))
        {
            return Ok(existing.clone());
        }

        if self.count(model) >= self.quota[model] {
            return Err(CreateError::Quota);
        }

        let machine = Machine {
            id: format!("m-{:02}", self.machines.len() + 1),
            model,
            token: token.map(str::to_owned),
        };
        self.machines.push(machine.clone());

        if self.calls == self.lose_reply {
            return Err(CreateError::Timeout);
        }

        Ok(machine)
    }
}

/// A group of identical machines and the size its owner wants.
struct Pool {
    /// Prefixes the name of each slot: thor-1, thor-2, and so on.
    name: &'static str,
    /// The kind of machine in the pool.
    model: &'static str,
    /// The number of machines the pool should have.
    want: usize,
}

/// Brings each pool to its wanted size.
struct Controller {
    provider: Provider,
    /// Whether create calls carry a client token.
    tokens: bool,
    /// Maps a pool's name to the machines the controller knows it has.
    owned: BTreeMap<&'static str, Vec<Machine>>,
}

impl Controller {
    /// Fills the empty slots of one pool, in order, and stops at the first
    /// call that fails. It returns true when the pool is full.
    fn reconcile(&mut self, pool: &Pool) -> bool {
        let owned = self.owned.entry(pool.name).or_default();

        while owned.len() < pool.want {
            let slot = format!("{}-{}", pool.name, owned.len() + 1);
            let token = self.tokens.then_some(slot.as_str());

            match self.provider.create(pool.model, token) {
                Ok(machine) => {
                    println!("    {slot:<8} is machine {}", machine.id);
                    owned.push(machine);
                }
                Err(e) => {
                    println!("    {slot:<8} create failed: {e}");

                    return false;
                }
            }
        }

        true
    }
}

/// The fleet we ask for: the Thors on the shelf and the two Sparks.
const POOLS: [Pool; 2] = [
    Pool {
        name: "thor",
        model: "Jetson AGX Thor",
        want: 5,
    },
    Pool {
        name: "spark",
        model: "DGX Spark",
        want: 2,
    },
];

/// Reconciles every pool for three passes and then compares what the
/// controller owns with what the provider holds.
fn run(tokens: bool) {
    let provider = Provider {
        quota: BTreeMap::from([("Jetson AGX Thor", 5), ("DGX Spark", 2)]),
        machines: Vec::new(),
        calls: 0,
        lose_reply: 3,
    };
    let mut c = Controller {
        provider,
        tokens,
        owned: BTreeMap::new(),
    };

    for pass in 1..=3 {
        println!("  pass {pass}");
        let mut full = true;

        for pool in &POOLS {
            full = c.reconcile(pool) && full;
        }

        if full {
            println!("    every pool is full");
            break;
        }
    }

    let mut known = BTreeSet::new();

    for pool in &POOLS {
        let owned = &c.owned[pool.name];
        known.extend(owned.iter().map(|m| m.id.as_str()));
        println!(
            "  {:<5} wanted {}, controller owns {}, provider holds {}",
            pool.name,
            pool.want,
            owned.len(),
            c.provider.count(pool.model)
        );
    }

    for m in c
        .provider
        .machines
        .iter()
        .filter(|m| !known.contains(m.id.as_str()))
    {
        println!(
            "  machine {} ({}) is running and belongs to no pool",
            m.id, m.model
        );
    }
}

fn main() {
    println!("create calls without a client token:");
    run(false);
    println!("create calls with the slot name as the client token:");
    run(true);
}

References

• Kubernetes, TLS bootstrapping: https://kubernetes.io/docs/reference/access-authn-authz/kubelet-tls-bootstrapping/
• Kubernetes, Authenticating with bootstrap tokens: https://kubernetes.io/docs/reference/access-authn-authz/bootstrap-tokens/
• Kubernetes, kubeadm join: https://kubernetes.io/docs/reference/setup-tools/kubeadm/kubeadm-join/
• Amazon EC2, Ensuring idempotency in API requests: https://docs.aws.amazon.com/ec2/latest/devguide/ec2-api-idempotency.html
• The Cluster API Book: https://cluster-api.sigs.k8s.io/
• NVIDIA, GPU Feature Discovery: https://github.com/NVIDIA/k8s-device-plugin/blob/main/docs/gpu-feature-discovery/README.md

6. Virtual machines on Kubernetes

Every workload so far has run in a container. A container shares the host’s kernel, and some users of a GPU cluster need a kernel of their own: a customer who must run a particular operating system, a tenant who is not trusted to share a kernel with others, or a team with an appliance that ships as a disk image. For them the unit of work is a virtual machine (VM), a complete computer simulated in software, with its own kernel, on a real one.

In this chapter, we will see how a VM runs on a Kubernetes cluster and what that means for the two things this book cares about: giving the VM a GPU, and moving the VM off a machine that has to be repaired. The second leads to live migration, moving a running VM to another host without stopping it. We will reason about when that can work, and then write migration, in Go and in Rust, a model of the algorithm that we can run with the memory sizes and link speeds of our own fleet.

A virtual machine is a process

On Linux, a VM is an ordinary process on the host. The program is usually QEMU. It allocates a large block of its own memory and presents it to the guest as the guest’s physical memory, and it creates one thread for each of the guest’s CPUs.

Those threads run the guest’s code directly on the host’s CPU cores, at close to full speed, through a kernel feature called KVM (Kernel-based Virtual Machine). KVM uses the CPU’s virtualization support to run guest code in a restricted mode, and returns control to QEMU whenever the guest does something that needs the host’s help, such as talking to a disk or a network card. A program asks for this by opening the device file /dev/kvm.

The Orin NX that runs our cluster has it:

$ ls -l /dev/kvm
crw-rw----+ 1 root kvm 10, 232 Nov 21  2023 /dev/kvm
$ zcat /proc/config.gz | grep -E '^CONFIG_(KVM|VFIO)='
CONFIG_KVM=y
CONFIG_VFIO=y

Since a VM is a process, it can run in a container, and so in a pod. That observation is the whole basis for running VMs on Kubernetes.

KubeVirt: a VM inside a pod

KubeVirt is the project that adds virtual machines to Kubernetes. It follows the pattern of section 1.1: it teaches the API server new kinds of object with CustomResourceDefinitions and runs controllers that reconcile them. Figure 6.1 shows the parts.

A user creates a VirtualMachine object, whose spec describes the VM: its CPUs, its memory, its disks, and its devices. A controller, virt-controller, creates a pod for it, named virt-launcher followed by the VM’s name. The pod’s resource requests match the VM’s size, so the scheduler places it exactly as it places any pod, and everything we have built so far applies: the pod is counted against the node’s capacity, it shows up in gpuusage from chapter 2 if it holds a GPU, and it is evicted when chapter 4’s controller drains the node.

On each node runs virt-handler, an agent that plays the kubelet’s part for VMs. When a launcher pod starts on its node, it has the launcher start a QEMU process inside the pod’s container. The guest’s memory is that process’s memory, so it counts against the pod’s memory limit.

The running VM is represented by a second object, the VirtualMachineInstance, whose status reports the node it is on and the phase it is in. The VirtualMachine object outlives it, in the way a GpuJob outlived its pod: stopping a VM deletes the instance and keeps the VirtualMachine.

Giving a VM a GPU

A container reaches a GPU through the host’s driver, as Appendix A showed. A VM has its own kernel and so needs its own driver, which has to talk to a GPU device. There are two ways to give it one.

Passthrough hands a whole physical GPU to one VM. The host detaches the GPU from its own driver and attaches it to a stub driver, VFIO, whose job is to let a process map the device’s registers and receive its interrupts. QEMU presents the device to the guest, and the guest’s NVIDIA driver then operates the real GPU directly. The guest gets the full performance of the card. The host can no longer use it, and can no longer see what it is doing.

A vGPU, which we met in chapter 2, divides one physical GPU among several VMs. The host keeps its own driver, which creates several virtual devices and schedules their work on the real GPU. Each VM is given one virtual device.

In KubeVirt, the VM’s spec names the device it wants in a gpus list, by the same kind of resource name a pod uses, and a device plugin on the node advertises the GPUs that have been set aside for VMs. The scheduler counts them as it counts nvidia.com/gpu.

The last row of the table decides what the rest of this chapter can do for a VM with a GPU, as we will see once we have built the model.

$ nvidia-smi | sed -n 9,10p        # on the Orin NX
|   0  Orin (nvgpu)                  N/A  | N/A              N/A |                  N/A |
| N/A   N/A  N/A               N/A /  N/A | Not Supported        |     N/A          N/A |

$ lspci -nnk -s 01:00.0            # on the Thor
0000:01:00.0 3D controller [0302]: NVIDIA Corporation Device [10de:2b00] (rev a1)
	Kernel driver in use: nvidia
	Kernel modules: nouveau, nvgpu, nvidia_drm, nvidia

Moving a running VM

Chapter 4’s controller drains a machine before repairing it. For a container, draining means stopping the process and starting a new one elsewhere. A VM often runs software that expects to stay up for months, and its users did not write it to be restarted. For a VM we would like something better: to move it to another host while it runs.

A VM’s state is its memory, the registers of its CPUs, and the state of its devices. The registers and most device state amount to a few kilobytes. The memory is the problem. A VM sized for a Thor may hold 100 GB, and at the 3.125 GB per second that one 25 GbE lane carries, copying it takes 32 seconds. If the VM were paused for that long, its network connections would time out and its users would see an outage.

So the copy happens while the VM runs, and that creates a second problem: the VM keeps writing to its memory during the copy, and the pages it writes after they were sent are stale on the other side. Memory that has been written since it was last sent is called dirty.

Pre-copy, round by round

The standard algorithm, pre-copy, turns this into a loop.

1. Mark all of the VM’s memory as dirty.
2. Send every dirty page to the target host, and clear the mark on each page as it is sent. The VM keeps running, and the hypervisor marks each page it writes.
3. When the round ends, some pages are dirty again. If they are few enough to send within the pause we are willing to accept, go to step 4. Otherwise go back to step 2.
4. Pause the VM, send the remaining dirty pages and the CPU and device state, and resume the VM on the target.

Each round sends less than the one before, provided the VM dirties memory more slowly than the link carries it. Let us work one case by hand. A 100 GB VM dirties 4 GB each second, and the link carries 12.5 GB each second, which is four bonded 25 GbE lanes.

• Round 1 sends 100 GB. That takes 100 / 12.5 = 8 seconds, during which the VM dirties 4 × 8 = 32 GB.
• Round 2 sends 32 GB in 2.56 seconds, during which the VM dirties 10.24 GB.
• Round 3 sends 10.24 GB in 0.82 seconds, during which the VM dirties 3.28 GB.
• 3.28 GB takes 0.26 seconds to send. If we allow a pause of 0.3 seconds, the VM is paused now and the last 3.28 GB are sent.

Each round is smaller than the last by the same factor: the dirty rate divided by the link rate, here 4 / 12.5 = 0.32. Figure 6.2 shows the rounds as bars.

That factor is the whole story of convergence:

• When the VM dirties memory more slowly than the link carries it, the factor is below one, the rounds shrink, and the migration finishes.
• When the VM dirties memory as fast as the link carries it, or faster, the factor is one or more. Every round ends with as much dirty memory as it started with, and the migration does not reach step 4.

A hypervisor has two remedies for the second case. It can slow the guest down, by giving its CPU threads less time to run, which lowers the rate at which they dirty memory. QEMU calls this auto-converge. Or it can switch to post-copy: pause the VM at once, move only the CPU and device state, resume it on the target, and fetch each page from the source at the moment the guest first touches it. Post-copy always finishes, at the price that a network failure during the migration leaves the VM’s memory split between two hosts with no complete copy on either.

Building migration

The VM and the link

The model needs three numbers: the VM’s memory, the rate at which the VM dirties it, and the rate of the link. Link speeds are quoted in gigabits per second and memory in gigabytes, so each link is entered in gigabytes per second: 25 GbE is 25 / 8 = 3.125.

Listing 6.1 The VM, the link, and the limits, in Go. code/go/migration/main.go

// VM describes the virtual machine being moved.
type VM struct {
	// Name labels the VM in the output.
	Name string
	// MemoryGB is the size of the VM's memory.
	MemoryGB float64
	// DirtyGBps is how much memory the VM's programs change each second.
	DirtyGBps float64
}

// Link is the network path between the two hosts.
type Link struct {
	// Name labels the link in the output.
	Name string
	// GBps is what the link carries, in gigabytes per second. A link of
	// 8 gigabits per second carries 1 gigabyte per second.
	GBps float64
}

// Limits of the algorithm.
const (
	// maxPauseSeconds is the longest the VM may be paused for the last round.
	maxPauseSeconds = 0.3
	// maxRounds ends a migration that keeps making progress too slowly.
	maxRounds = 30
	// throttleStep is how much more the guest is slowed each time a round
	// fails to shrink the remaining memory, as a fraction of full speed.
	throttleStep = 0.2
	// maxThrottle is the most the guest may be slowed.
	maxThrottle = 0.8
)

Listing 6.2 The VM, the link, and the limits, in Rust. code/rust/migration/src/main.rs

/// The virtual machine being moved.
struct Vm {
    /// Labels the VM in the output.
    name: &'static str,
    /// The size of the VM's memory.
    memory_gb: f64,
    /// How much memory the VM's programs change each second.
    dirty_gbps: f64,
}

/// The network path between the two hosts.
struct Link {
    /// Labels the link in the output.
    name: &'static str,
    /// What the link carries, in gigabytes per second. A link of 8 gigabits
    /// per second carries 1 gigabyte per second.
    gbps: f64,
}

/// The longest the VM may be paused for the last round, in seconds.
const MAX_PAUSE_SECONDS: f64 = 0.3;
/// Ends a migration that keeps making progress too slowly.
const MAX_ROUNDS: u32 = 30;
/// How much more the guest is slowed each time a round fails to shrink the
/// remaining memory, as a fraction of full speed.
const THROTTLE_STEP: f64 = 0.2;
/// The most the guest may be slowed.
const MAX_THROTTLE: f64 = 0.8;

One migration

Migrate keeps one quantity between rounds: remaining, the memory still to be sent. It starts as the whole memory. Each pass of the loop first asks whether remaining can be sent within the pause, and if so pauses the VM and finishes. Otherwise it computes how long the round takes and how much memory the guest dirties in that time. The dirtied amount is capped at the VM’s size, since a page that is written twice is still one dirty page.

Listing 6.3 The pre-copy loop, in Go. code/go/migration/main.go

// Migrate runs pre-copy rounds until the remaining memory can be sent within
// the allowed pause, and prints each round. With throttle set, it slows the
// guest whenever a round leaves as much memory to send as it started with.
func Migrate(vm VM, link Link, throttle bool) {
	fmt.Printf("%s over %s", vm.Name, link.Name)

	if throttle {
		fmt.Print(", with throttling")
	}

	fmt.Println(":")

	remaining := vm.MemoryGB // memory still to be sent
	slowed := 0.0            // fraction by which the guest is slowed
	elapsed := 0.0

	for round := 1; round <= maxRounds; round++ {
		if remaining <= link.GBps*maxPauseSeconds {
			pause := remaining / link.GBps
			fmt.Printf("  round %2d: VM paused, %.2f GB sent in %.2f s\n", round, remaining, pause)
			fmt.Printf("  done in %.1f s; the VM was paused for %.0f ms\n", elapsed+pause, pause*1000)

			return
		}

		seconds := remaining / link.GBps
		dirtied := min(vm.DirtyGBps*(1-slowed)*seconds, vm.MemoryGB)
		fmt.Printf("  round %2d: %6.2f GB sent in %6.1f s, %6.2f GB changed meanwhile\n",
			round, remaining, seconds, dirtied)
		elapsed += seconds

		if dirtied >= remaining {
			if !throttle || slowed >= maxThrottle {
				fmt.Println("  does not converge: memory changes as fast as the link carries it")

				return
			}

			slowed += throttleStep
			fmt.Printf("            no progress: guest slowed to %.0f%% of full speed\n", (1-slowed)*100)
		}

		remaining = dirtied
	}

	fmt.Printf("  gave up after %d rounds\n", maxRounds)
}

Listing 6.4 The pre-copy loop, in Rust. code/rust/migration/src/main.rs

/// Runs pre-copy rounds until the remaining memory can be sent within the
/// allowed pause, and prints each round. With `throttle` set, it slows the
/// guest whenever a round leaves as much memory to send as it started with.
fn migrate(vm: &Vm, link: &Link, throttle: bool) {
    let suffix = if throttle { ", with throttling" } else { "" };
    println!("{} over {}{suffix}:", vm.name, link.name);

    let mut remaining = vm.memory_gb; // memory still to be sent
    let mut slowed = 0.0; // fraction by which the guest is slowed
    let mut elapsed = 0.0;

    for round in 1..=MAX_ROUNDS {
        if remaining <= link.gbps * MAX_PAUSE_SECONDS {
            let pause = remaining / link.gbps;
            println!("  round {round:2}: VM paused, {remaining:.2} GB sent in {pause:.2} s");
            println!(
                "  done in {:.1} s; the VM was paused for {:.0} ms",
                elapsed + pause,
                pause * 1000.0
            );

            return;
        }

        let seconds = remaining / link.gbps;
        let dirtied = (vm.dirty_gbps * (1.0 - slowed) * seconds).min(vm.memory_gb);
        println!(
            "  round {round:2}: {remaining:6.2} GB sent in {seconds:6.1} s, {dirtied:6.2} GB changed meanwhile"
        );
        elapsed += seconds;

        if dirtied >= remaining {
            if !throttle || slowed >= MAX_THROTTLE {
                println!("  does not converge: memory changes as fast as the link carries it");

                return;
            }

            slowed += THROTTLE_STEP;
            println!(
                "            no progress: guest slowed to {:.0}% of full speed",
                (1.0 - slowed) * 100.0
            );
        }

        remaining = dirtied;
    }

    println!("  gave up after {MAX_ROUNDS} rounds");
}

A round that dirties at least as much as it sent has made no progress. Without throttling the function reports that the migration does not converge. With throttling it slows the guest by another 20 percent and tries again.

The run

The run uses two VMs and the four links our fleet has: the Orin NX’s 1 GbE, one of a Thor’s 25 GbE lanes, all four lanes bonded together, and the 200 GbE cable between the two Sparks.

Listing 6.5 The VMs, the links, and the six migrations, in Go. code/go/migration/main.go

func main() {
	// An inference VM sized for the Orin NX, and a training VM sized for a
	// Thor or a Spark.
	inference := VM{"16 GB inference VM (0.2 GB/s changed)", 16, 0.2}
	training := VM{"100 GB training VM (4 GB/s changed)", 100, 4}

	// The links in the fleet, in gigabytes per second.
	orin := Link{"1 GbE", 0.125}
	lane := Link{"one 25 GbE lane", 3.125}
	bonded := Link{"four bonded 25 GbE lanes", 12.5}
	spark := Link{"the Sparks' 200 GbE cable", 25}

	Migrate(inference, orin, false)
	Migrate(inference, lane, false)
	Migrate(training, lane, false)
	Migrate(training, lane, true)
	Migrate(training, bonded, false)
	Migrate(training, spark, false)
}

Listing 6.6 The VMs, the links, and the six migrations, in Rust. code/rust/migration/src/main.rs

fn main() {
    // An inference VM sized for the Orin NX, and a training VM sized for a
    // Thor or a Spark.
    let inference = Vm {
        name: "16 GB inference VM (0.2 GB/s changed)",
        memory_gb: 16.0,
        dirty_gbps: 0.2,
    };
    let training = Vm {
        name: "100 GB training VM (4 GB/s changed)",
        memory_gb: 100.0,
        dirty_gbps: 4.0,
    };

    // The links in the fleet, in gigabytes per second.
    let orin = Link {
        name: "1 GbE",
        gbps: 0.125,
    };
    let lane = Link {
        name: "one 25 GbE lane",
        gbps: 3.125,
    };
    let bonded = Link {
        name: "four bonded 25 GbE lanes",
        gbps: 12.5,
    };
    let spark = Link {
        name: "the Sparks' 200 GbE cable",
        gbps: 25.0,
    };

    migrate(&inference, &orin, false);
    migrate(&inference, &lane, false);
    migrate(&training, &lane, false);
    migrate(&training, &lane, true);
    migrate(&training, &bonded, false);
    migrate(&training, &spark, false);
}

$ go run ./migration
16 GB inference VM (0.2 GB/s changed) over 1 GbE:
  round  1:  16.00 GB sent in  128.0 s,  16.00 GB changed meanwhile
  does not converge: memory changes as fast as the link carries it
16 GB inference VM (0.2 GB/s changed) over one 25 GbE lane:
  round  1:  16.00 GB sent in    5.1 s,   1.02 GB changed meanwhile
  round  2:   1.02 GB sent in    0.3 s,   0.07 GB changed meanwhile
  round  3: VM paused, 0.07 GB sent in 0.02 s
  done in 5.5 s; the VM was paused for 21 ms
100 GB training VM (4 GB/s changed) over one 25 GbE lane:
  round  1: 100.00 GB sent in   32.0 s, 100.00 GB changed meanwhile
  does not converge: memory changes as fast as the link carries it
100 GB training VM (4 GB/s changed) over one 25 GbE lane, with throttling:
  round  1: 100.00 GB sent in   32.0 s, 100.00 GB changed meanwhile
            no progress: guest slowed to 80% of full speed
  round  2: 100.00 GB sent in   32.0 s, 100.00 GB changed meanwhile
            no progress: guest slowed to 60% of full speed
  round  3: 100.00 GB sent in   32.0 s,  76.80 GB changed meanwhile
  round  4:  76.80 GB sent in   24.6 s,  58.98 GB changed meanwhile
  round  5:  58.98 GB sent in   18.9 s,  45.30 GB changed meanwhile
  round  6:  45.30 GB sent in   14.5 s,  34.79 GB changed meanwhile
  round  7:  34.79 GB sent in   11.1 s,  26.72 GB changed meanwhile
  round  8:  26.72 GB sent in    8.5 s,  20.52 GB changed meanwhile
  round  9:  20.52 GB sent in    6.6 s,  15.76 GB changed meanwhile
  round 10:  15.76 GB sent in    5.0 s,  12.10 GB changed meanwhile
  round 11:  12.10 GB sent in    3.9 s,   9.30 GB changed meanwhile
  round 12:   9.30 GB sent in    3.0 s,   7.14 GB changed meanwhile
  round 13:   7.14 GB sent in    2.3 s,   5.48 GB changed meanwhile
  round 14:   5.48 GB sent in    1.8 s,   4.21 GB changed meanwhile
  round 15:   4.21 GB sent in    1.3 s,   3.23 GB changed meanwhile
  round 16:   3.23 GB sent in    1.0 s,   2.48 GB changed meanwhile
  round 17:   2.48 GB sent in    0.8 s,   1.91 GB changed meanwhile
  round 18:   1.91 GB sent in    0.6 s,   1.46 GB changed meanwhile
  round 19:   1.46 GB sent in    0.5 s,   1.12 GB changed meanwhile
  round 20:   1.12 GB sent in    0.4 s,   0.86 GB changed meanwhile
  round 21: VM paused, 0.86 GB sent in 0.28 s
  done in 201.0 s; the VM was paused for 276 ms
100 GB training VM (4 GB/s changed) over four bonded 25 GbE lanes:
  round  1: 100.00 GB sent in    8.0 s,  32.00 GB changed meanwhile
  round  2:  32.00 GB sent in    2.6 s,  10.24 GB changed meanwhile
  round  3:  10.24 GB sent in    0.8 s,   3.28 GB changed meanwhile
  round  4: VM paused, 3.28 GB sent in 0.26 s
  done in 11.6 s; the VM was paused for 262 ms
100 GB training VM (4 GB/s changed) over the Sparks' 200 GbE cable:
  round  1: 100.00 GB sent in    4.0 s,  16.00 GB changed meanwhile
  round  2:  16.00 GB sent in    0.6 s,   2.56 GB changed meanwhile
  round  3: VM paused, 2.56 GB sent in 0.10 s
  done in 4.7 s; the VM was paused for 102 ms

The inference VM dirties 0.2 GB per second. The Orin NX’s 1 GbE link carries 0.125 GB per second, which is less, so after 128 seconds of copying all 16 GB are dirty again. Over one 25 GbE lane the same VM moves in 5.5 seconds with a pause of 21 milliseconds. The VM is the same in both runs. The link decides the outcome.

The training VM dirties 4 GB per second, which is more than one 25 GbE lane carries, so that migration does not converge either. With throttling, the first two rounds make no progress and the guest is slowed twice, to 60 percent of full speed. At that speed it dirties 2.4 GB per second, the factor becomes 2.4 / 3.125 = 0.768, and the rounds shrink, slowly: the migration takes 21 rounds and 201 seconds, and the training job runs at reduced speed for most of that time.

Over the four bonded lanes the factor is 0.32 and the migration takes 11.6 seconds. Over the Sparks’ cable it is 0.16 and the migration takes 4.7 seconds, with a pause of 102 milliseconds against 262.

Animation 6.1 plays three of these migrations for the training VM, with the memory still to send beside the memory that changes during the round.

This is the reason clusters that migrate large VMs give migration a fast network of its own. The memory size sets how long the first round takes. The ratio of dirty rate to link rate sets whether there is a last round at all.

Migration in KubeVirt

KubeVirt runs this algorithm through objects. A VirtualMachineInstanceMigration names a running instance. A controller creates a second launcher pod on another node, the two QEMU processes connect to each other, and the source sends its memory in rounds. When the migration completes, the instance’s status names the new node and the old pod is removed.

A VM’s spec can set evictionStrategy: LiveMigrate. An eviction of the launcher pod, the request chapter 4’s drain sends for every pod on a node, then starts a migration and is refused until the VM has moved. This is how the two chapters meet: the remediation controller drains with the Eviction API and needs no knowledge of VMs, and each VM leaves in the way its owner chose.

The VM that cannot move

Step 4 of pre-copy sends “the CPU and device state”. For a virtual disk or a virtual network card, the hypervisor implements the device, so it has that state in its own memory and can send it.

A GPU given to the VM by passthrough is different. The device’s state, including everything in the GPU’s memory, is inside the physical card on the source host, and the host gave up its driver for that card when it attached the stub. The hypervisor has nothing to read and nothing to send, and the target host’s GPU is a different physical device. KubeVirt therefore marks a VM with a passed-through device as not migratable. Draining its node means shutting the VM down and starting it again elsewhere, and the checkpointing discipline of chapter 4 applies to the jobs inside it.

With a vGPU, the host’s driver holds the device’s state and can save and restore it, so migration is possible when the host driver and the hypervisor both support it. That support depends on the driver and hypervisor versions, and has to be checked for each combination.

The complete program

Listing 6.7 The complete program in Go. code/go/migration/main.go

// Command migration models the live migration of a virtual machine: moving a
// running VM's memory from one host to another while the VM keeps running.
//
// The model is the pre-copy algorithm. Each round sends the memory that
// changed during the round before it. The VM is paused only for the last
// round, when what remains is small enough to send within the allowed pause.
// The program runs the algorithm for VMs of two sizes over the links in our
// fleet and prints each round.
//
// Run it from code/go with:
//
//	go run ./migration
package main

import "fmt"

// VM describes the virtual machine being moved.
type VM struct {
	// Name labels the VM in the output.
	Name string
	// MemoryGB is the size of the VM's memory.
	MemoryGB float64
	// DirtyGBps is how much memory the VM's programs change each second.
	DirtyGBps float64
}

// Link is the network path between the two hosts.
type Link struct {
	// Name labels the link in the output.
	Name string
	// GBps is what the link carries, in gigabytes per second. A link of
	// 8 gigabits per second carries 1 gigabyte per second.
	GBps float64
}

// Limits of the algorithm.
const (
	// maxPauseSeconds is the longest the VM may be paused for the last round.
	maxPauseSeconds = 0.3
	// maxRounds ends a migration that keeps making progress too slowly.
	maxRounds = 30
	// throttleStep is how much more the guest is slowed each time a round
	// fails to shrink the remaining memory, as a fraction of full speed.
	throttleStep = 0.2
	// maxThrottle is the most the guest may be slowed.
	maxThrottle = 0.8
)


// Migrate runs pre-copy rounds until the remaining memory can be sent within
// the allowed pause, and prints each round. With throttle set, it slows the
// guest whenever a round leaves as much memory to send as it started with.
func Migrate(vm VM, link Link, throttle bool) {
	fmt.Printf("%s over %s", vm.Name, link.Name)

	if throttle {
		fmt.Print(", with throttling")
	}

	fmt.Println(":")

	remaining := vm.MemoryGB // memory still to be sent
	slowed := 0.0            // fraction by which the guest is slowed
	elapsed := 0.0

	for round := 1; round <= maxRounds; round++ {
		if remaining <= link.GBps*maxPauseSeconds {
			pause := remaining / link.GBps
			fmt.Printf("  round %2d: VM paused, %.2f GB sent in %.2f s\n", round, remaining, pause)
			fmt.Printf("  done in %.1f s; the VM was paused for %.0f ms\n", elapsed+pause, pause*1000)

			return
		}

		seconds := remaining / link.GBps
		dirtied := min(vm.DirtyGBps*(1-slowed)*seconds, vm.MemoryGB)
		fmt.Printf("  round %2d: %6.2f GB sent in %6.1f s, %6.2f GB changed meanwhile\n",
			round, remaining, seconds, dirtied)
		elapsed += seconds

		if dirtied >= remaining {
			if !throttle || slowed >= maxThrottle {
				fmt.Println("  does not converge: memory changes as fast as the link carries it")

				return
			}

			slowed += throttleStep
			fmt.Printf("            no progress: guest slowed to %.0f%% of full speed\n", (1-slowed)*100)
		}

		remaining = dirtied
	}

	fmt.Printf("  gave up after %d rounds\n", maxRounds)
}


func main() {
	// An inference VM sized for the Orin NX, and a training VM sized for a
	// Thor or a Spark.
	inference := VM{"16 GB inference VM (0.2 GB/s changed)", 16, 0.2}
	training := VM{"100 GB training VM (4 GB/s changed)", 100, 4}

	// The links in the fleet, in gigabytes per second.
	orin := Link{"1 GbE", 0.125}
	lane := Link{"one 25 GbE lane", 3.125}
	bonded := Link{"four bonded 25 GbE lanes", 12.5}
	spark := Link{"the Sparks' 200 GbE cable", 25}

	Migrate(inference, orin, false)
	Migrate(inference, lane, false)
	Migrate(training, lane, false)
	Migrate(training, lane, true)
	Migrate(training, bonded, false)
	Migrate(training, spark, false)
}

Listing 6.8 The complete program in Rust. code/rust/migration/src/main.rs

//! Models the live migration of a virtual machine: moving a running VM's
//! memory from one host to another while the VM keeps running.
//!
//! The model is the pre-copy algorithm. Each round sends the memory that
//! changed during the round before it. The VM is paused only for the last
//! round, when what remains is small enough to send within the allowed pause.
//! The program runs the algorithm for VMs of two sizes over the links in our
//! fleet and prints each round.
//!
//! Run it from code/rust with:
//!
//! ```text
//! cargo run -p migration
//! ```

/// The virtual machine being moved.
struct Vm {
    /// Labels the VM in the output.
    name: &'static str,
    /// The size of the VM's memory.
    memory_gb: f64,
    /// How much memory the VM's programs change each second.
    dirty_gbps: f64,
}

/// The network path between the two hosts.
struct Link {
    /// Labels the link in the output.
    name: &'static str,
    /// What the link carries, in gigabytes per second. A link of 8 gigabits
    /// per second carries 1 gigabyte per second.
    gbps: f64,
}

/// The longest the VM may be paused for the last round, in seconds.
const MAX_PAUSE_SECONDS: f64 = 0.3;
/// Ends a migration that keeps making progress too slowly.
const MAX_ROUNDS: u32 = 30;
/// How much more the guest is slowed each time a round fails to shrink the
/// remaining memory, as a fraction of full speed.
const THROTTLE_STEP: f64 = 0.2;
/// The most the guest may be slowed.
const MAX_THROTTLE: f64 = 0.8;

/// Runs pre-copy rounds until the remaining memory can be sent within the
/// allowed pause, and prints each round. With `throttle` set, it slows the
/// guest whenever a round leaves as much memory to send as it started with.
fn migrate(vm: &Vm, link: &Link, throttle: bool) {
    let suffix = if throttle { ", with throttling" } else { "" };
    println!("{} over {}{suffix}:", vm.name, link.name);

    let mut remaining = vm.memory_gb; // memory still to be sent
    let mut slowed = 0.0; // fraction by which the guest is slowed
    let mut elapsed = 0.0;

    for round in 1..=MAX_ROUNDS {
        if remaining <= link.gbps * MAX_PAUSE_SECONDS {
            let pause = remaining / link.gbps;
            println!("  round {round:2}: VM paused, {remaining:.2} GB sent in {pause:.2} s");
            println!(
                "  done in {:.1} s; the VM was paused for {:.0} ms",
                elapsed + pause,
                pause * 1000.0
            );

            return;
        }

        let seconds = remaining / link.gbps;
        let dirtied = (vm.dirty_gbps * (1.0 - slowed) * seconds).min(vm.memory_gb);
        println!(
            "  round {round:2}: {remaining:6.2} GB sent in {seconds:6.1} s, {dirtied:6.2} GB changed meanwhile"
        );
        elapsed += seconds;

        if dirtied >= remaining {
            if !throttle || slowed >= MAX_THROTTLE {
                println!("  does not converge: memory changes as fast as the link carries it");

                return;
            }

            slowed += THROTTLE_STEP;
            println!(
                "            no progress: guest slowed to {:.0}% of full speed",
                (1.0 - slowed) * 100.0
            );
        }

        remaining = dirtied;
    }

    println!("  gave up after {MAX_ROUNDS} rounds");
}

fn main() {
    // An inference VM sized for the Orin NX, and a training VM sized for a
    // Thor or a Spark.
    let inference = Vm {
        name: "16 GB inference VM (0.2 GB/s changed)",
        memory_gb: 16.0,
        dirty_gbps: 0.2,
    };
    let training = Vm {
        name: "100 GB training VM (4 GB/s changed)",
        memory_gb: 100.0,
        dirty_gbps: 4.0,
    };

    // The links in the fleet, in gigabytes per second.
    let orin = Link {
        name: "1 GbE",
        gbps: 0.125,
    };
    let lane = Link {
        name: "one 25 GbE lane",
        gbps: 3.125,
    };
    let bonded = Link {
        name: "four bonded 25 GbE lanes",
        gbps: 12.5,
    };
    let spark = Link {
        name: "the Sparks' 200 GbE cable",
        gbps: 25.0,
    };

    migrate(&inference, &orin, false);
    migrate(&inference, &lane, false);
    migrate(&training, &lane, false);
    migrate(&training, &lane, true);
    migrate(&training, &bonded, false);
    migrate(&training, &spark, false);
}

References

• KubeVirt user guide, Architecture: https://kubevirt.io/user-guide/architecture/
• KubeVirt user guide, Live migration: https://kubevirt.io/user-guide/compute/live_migration/
• KubeVirt user guide, Host devices assignment: https://kubevirt.io/user-guide/compute/host-devices/
• QEMU, Migration framework: https://www.qemu.org/docs/master/devel/migration/main.html
• The Linux kernel, VFIO: https://docs.kernel.org/driver-api/vfio.html
• C. Clark and others, “Live Migration of Virtual Machines”, 2nd USENIX Symposium on Networked Systems Design and Implementation (NSDI), 2005.

7. Collective communication and training failures

The jobs in the earlier chapters each ran on one machine. A model that is too large or too slow to train on one machine is trained on several at once, and that changes what a failure means. In this chapter, we will look at what the machines of a training job say to each other during every step, why the failure of one machine stops all of them, and how often a job should save its state so that a failure costs little.

We will first reason about the communication with pencil and paper, on the five Thors from chapter 4. Then we will write allreduce, in Go and in Rust: five workers running at the same time, connected in a ring, exchanging the pieces of a gradient. We will run it healthy, run it again with one worker failing part of the way through, and finish with the arithmetic of checkpoints.

One step of training on five machines

Training adjusts a model’s parameters, the numbers that make up the model, a little at a time. One step takes a batch of training examples, measures how wrong the model’s output is for them, and computes a gradient: a list with one number for each parameter, saying in which direction and how strongly that parameter should change. The step ends by changing every parameter by a small multiple of its gradient entry.

With five machines, each holds a full copy of the model and takes a different batch. Each computes its own gradient. Then all five have to apply the same change, or their copies of the model drift apart. The change they apply is the average of the five gradients, so before any machine can finish the step, every machine needs the sum of all five gradients.

That requirement has a name. An operation in which every member of a group contributes a value and every member receives the combined result is called all-reduce. It belongs to a family called collective communication, a term from the MPI standard for operations that a group of processes performs together, as opposed to one process sending a message to one other. On NVIDIA GPUs the library that implements these operations is NCCL, the NVIDIA Collective Communications Library, and training frameworks call it once per step.

The gradient is as large as the model. A model with 7 billion parameters stored as 16-bit numbers has a gradient of 14 GB. The step cannot finish until 14 GB from each machine have been combined and the result has reached every machine, so the network between the machines is part of the inner loop of training.

Two ways to add five gradients

The obvious design gives the job to one machine. Every worker sends its gradient to a coordinator, the coordinator adds them, and sends the sum back to every worker.

Count what the coordinator’s network link carries in one step: five gradients in and five out. Every other link carries one in each direction. The coordinator’s link is the bottleneck, and it gets worse in proportion to the number of workers. With 14 GB gradients and five workers, 70 GB pass through one link in each direction, every step.

The second design has no coordinator. Arrange the workers in a ring, so each has a left and a right neighbour, and cut every gradient into as many chunks as there are workers. Figure 7.1 compares the two shapes.

The ring works in two phases. We follow one chunk, chunk 0, through both. Our five workers start with the values 1, 2, 3, 4, and 5 in chunk 0, so its sum is 15.

In the first phase, each worker passes a partial sum to the right, and the receiver adds its own value:

• thor-1 sends its chunk 0, the value 1, to thor-2.
• thor-2 adds its own 2 and sends 3 to thor-3.
• thor-3 adds 3 and sends 6 to thor-4.
• thor-4 adds 4 and sends 10 to thor-5.
• thor-5 adds 5 and now holds 15, the finished sum of chunk 0.

That took four exchanges, one fewer than the number of workers. Meanwhile the other four chunks made the same journey from different starting points: chunk 1 started at thor-2 and finished at thor-1, and so on. Every worker sent one chunk in every exchange, each worker ends the phase holding the finished sum of exactly one chunk, and no two hold the same one.

In the second phase the finished chunks travel around the ring so that everyone gets a copy. thor-5 sends the 15 to thor-1, which stores it and passes it on, and after four more exchanges every worker holds it. The other finished chunks circulate at the same time.

Now count what a link carries. With N workers there are 2(N − 1) exchanges, and each sends one chunk, which is one Nth of a gradient, across every link. A worker therefore sends 2(N − 1)/N gradients’ worth of data per step. For five workers that is 1.6, and however many workers join, it stays below 2. Every link carries the same amount, and none is special.

Here is what that costs on our machines for the 14 GB gradient:

A step whose computation takes a second is dominated by seven seconds of communication on a single 25 GbE lane. This is the reason the fast links on these machines exist, and the reason chapter 8 places a job’s workers by the links between them.

One more property follows from the ring’s shape. Every chunk crosses every link. If one link is slower than the rest, every chunk waits for it, and the whole step runs at the speed of the slowest link.

When one worker fails

The same shape decides what happens when a worker fails in the middle of a step. A worker’s every exchange depends on its left-hand neighbour: it cannot add to a chunk it has not received. So when one worker stops, its right-hand neighbour is stuck at its next exchange. That neighbour had already sent its chunk for the exchange, so the worker after it completes one more exchange and is stuck at the following one. The stall travels around the ring, one worker per exchange, until it has reached everyone.

When the stall is complete, no worker holds the sum. Each holds some finished chunks, some partial sums, and some of its own original values, in a mixture that depends on where it was stopped. The step cannot be finished by the workers that remain, because the failed worker’s values for some chunks were not sent.

The model in every worker’s memory is intact, since the step had not yet changed it. But a training job is built as a fixed group: its workers were started together with a fixed ring, and the frameworks do not shrink the ring in mid-run. The practical response to one failed worker is to stop all of them, replace the failed one, and start the whole group again from the last saved state.

This has a consequence for the control plane that the single-machine jobs did not have. The controller that owns a training job treats its pods as one unit: when one pod fails, it deletes all of them and creates them all again. And it makes the two numbers from chapter 3 matter in a new way. The time to detect the failed machine is time during which the surviving machines hold their GPUs and do nothing.

Building allreduce

Each worker in the program stands for one Thor. One number stands for each chunk, which keeps the output readable. The algorithm does the same thing to a chunk of three gigabytes.

A worker and its two channels

A worker has a position in the ring, its gradient, and two channels: it receives on one and sends on the other. Both languages use a channel, a queue that one concurrently running piece of code puts values into and another takes them out of. Nothing else is shared between workers, which mirrors the real situation: machines in a training job share no memory, only a network.

Listing 7.1 The message, the worker, and how long it waits, in Go. code/go/allreduce/main.go

// message carries one chunk of a gradient from a worker to its right-hand
// neighbour.
type message struct {
	// chunk is the index of the chunk within the gradient.
	chunk int
	// value is the chunk's contents: one number stands for the whole chunk.
	value int
}

// Worker is one machine's part in the step.
type Worker struct {
	// id is the worker's position in the ring, from 0.
	id int
	// gradient holds one value per chunk. It starts as this worker's own
	// gradient and ends as the sum over all workers.
	gradient []int
	// in receives from the left-hand neighbour, and out sends to the right.
	in  <-chan message
	out chan<- message
	// failAfter is the number of messages the worker exchanges before it
	// fails, or -1 for a worker that stays healthy.
	failAfter int
}

// errFailed is the result of the worker that the run makes fail.
var errFailed = errors.New("failed")

// patience is how long a worker waits for its neighbour before giving up.
const patience = 100 * time.Millisecond

Listing 7.2 The message, the worker, and how long it waits, in Rust. code/rust/allreduce/src/main.rs

/// Carries one chunk of a gradient from a worker to its right-hand neighbour.
struct Message {
    /// The index of the chunk within the gradient.
    chunk: usize,
    /// The chunk's contents: one number stands for the whole chunk.
    value: i64,
}

/// One machine's part in the step.
struct Worker {
    /// The worker's position in the ring, from 0.
    id: usize,
    /// One value per chunk. It starts as this worker's own gradient and ends
    /// as the sum over all workers.
    gradient: Vec<i64>,
    /// Receives from the left-hand neighbour.
    inbox: Receiver<Message>,
    /// Sends to the right-hand neighbour.
    out: Sender<Message>,
    /// The number of messages the worker exchanges before it fails, if the
    /// run makes it fail.
    fail_after: Option<usize>,
}

/// How long a worker waits for its neighbour before giving up.
const PATIENCE: Duration = Duration::from_millis(100);

The exchanges

Run is the algorithm. Each pass of its loop is one exchange: send a chunk to the right, then wait for a chunk from the left. The chunk to send is found by counting backwards from the worker’s own position, one place per exchange. Working through the hand example shows why that is right: after the first exchange, the chunk a worker sends next is always the one it received last.

Listing 7.3 A worker's 2(n − 1) exchanges, in Go. code/go/allreduce/main.go

// Run performs the worker's 2(n-1) exchanges. In each one it sends a chunk to
// the right and receives a chunk from the left. During the first n-1
// exchanges it adds what it receives to its own chunk. During the last n-1
// the chunks it receives are finished sums, and it stores them. Run returns
// the number of exchanges completed and the reason it stopped early, if any.
func (w *Worker) Run(n int) (int, error) {
	for step := range 2 * (n - 1) {
		if step == w.failAfter {
			return step, errFailed
		}

		// The chunk to send moves one place back around the gradient each
		// step, so after the first step it is the chunk just received.
		chunk := mod(w.id-step, n)
		w.out <- message{chunk, w.gradient[chunk]}
		reducing := step < n-1

		select {
		case m := <-w.in:
			if reducing {
				w.gradient[m.chunk] += m.value
			} else {
				w.gradient[m.chunk] = m.value
			}
		case <-time.After(patience):
			return step, fmt.Errorf("no chunk from %s within %v", name(mod(w.id-1, n)), patience)
		}
	}

	return 2 * (n - 1), nil
}

Listing 7.4 A worker's 2(n − 1) exchanges, in Rust. code/rust/allreduce/src/main.rs

impl Worker {
    /// Performs the worker's 2(n-1) exchanges. In each one it sends a chunk to
    /// the right and receives a chunk from the left. During the first n-1
    /// exchanges it adds what it receives to its own chunk. During the last
    /// n-1 the chunks it receives are finished sums, and it stores them.
    /// Returns the worker with the number of exchanges completed and the
    /// reason it stopped early, if any.
    fn run(mut self, n: usize) -> (Self, usize, Option<String>) {
        for step in 0..2 * (n - 1) {
            if self.fail_after == Some(step) {
                return (self, step, Some("failed".to_owned()));
            }

            // The chunk to send moves one place back around the gradient each
            // step, so after the first step it is the chunk just received.
            // Adding 2n keeps the subtraction from going below zero.
            let chunk = (self.id + 2 * n - step) % n;
            // The right-hand neighbour may have stopped; its channel still accepts.
            let _ = self.out.send(Message {
                chunk,
                value: self.gradient[chunk],
            });
            let reducing = step < n - 1;

            match self.inbox.recv_timeout(PATIENCE) {
                Ok(m) if reducing => self.gradient[m.chunk] += m.value,
                Ok(m) => self.gradient[m.chunk] = m.value,
                Err(_) => {
                    let why = format!(
                        "no chunk from {} within {PATIENCE:?}",
                        name((self.id + n - 1) % n)
                    );

                    return (self, step, Some(why));
                }
            }
        }

        (self, 2 * (n - 1), None)
    }
}

The wait for the left-hand neighbour has a time limit, which the program calls patience. In Go the wait is a select with two cases, a message on the channel or the timer from time.After, and whichever is ready first runs. In Rust it is the channel’s own recv_timeout, which returns an error if nothing arrives in time. Without the limit a worker whose neighbour had failed would wait forever, and the program, like a training job with no timeouts, would hang with no report of what went wrong.

The Rust version of run takes self by value and returns it. The worker is moved into its thread when the thread starts and moved back out with the result when it ends, so the main thread can read each worker’s final gradient without any lock.

Building the ring

ring creates one channel per link and gives each worker the receiving end of its own link and the sending end of its right-hand neighbour’s. It then starts every worker at once and waits for all of them.

Listing 7.5 Connecting the workers and running them together, in Go. code/go/allreduce/main.go

// result is what one worker reports when it stops.
type result struct {
	steps int
	err   error
}

// ring connects n workers in a circle, runs them all at once, and returns the
// workers and their results. The worker at position failing stops after
// failAfter exchanges; pass -1 for a run with no failure.
func ring(n, failing, failAfter int) ([]*Worker, []result) {
	// Each link holds every message a worker can send in one step, so a send
	// goes through even when the receiver has stopped.
	links := make([]chan message, n)

	for i := range links {
		links[i] = make(chan message, 2*(n-1))
	}

	workers := make([]*Worker, n)

	for i := range workers {
		gradient := make([]int, n)

		for chunk := range gradient {
			gradient[chunk] = i + 1 + 10*chunk
		}

		workers[i] = &Worker{id: i, gradient: gradient, in: links[i], out: links[mod(i+1, n)], failAfter: -1}
	}

	if failing >= 0 {
		workers[failing].failAfter = failAfter
	}

	results := make([]result, n)
	var wg sync.WaitGroup

	for i, w := range workers {
		wg.Go(func() {
			steps, err := w.Run(n)
			results[i] = result{steps, err}
		})
	}

	wg.Wait()

	return workers, results
}

Listing 7.6 Connecting the workers and running them together, in Rust. code/rust/allreduce/src/main.rs

/// Connects n workers in a circle, runs them all at once, and returns each
/// worker with its result. The worker named in `failing` stops after that
/// many exchanges.
fn ring(n: usize, failing: Option<(usize, usize)>) -> Vec<(Worker, usize, Option<String>)> {
    // One channel per link. `senders[i]` leads to worker i.
    let (senders, receivers): (Vec<_>, Vec<_>) = (0..n).map(|_| mpsc::channel()).unzip();
    let workers: Vec<Worker> = receivers
        .into_iter()
        .enumerate()
        .map(|(id, inbox)| Worker {
            id,
            gradient: own_gradient(id, n),
            inbox,
            out: senders[(id + 1) % n].clone(),
            fail_after: failing
                .filter(|(who, _)| *who == id)
                .map(|(_, after)| after),
        })
        .collect();

    let handles: Vec<_> = workers
        .into_iter()
        .map(|w| thread::spawn(move || w.run(n)))
        .collect();

    handles
        .into_iter()
        .map(|h| h.join().expect("a worker panicked"))
        .collect()
}

In Go each worker runs in a goroutine started with wg.Go, a method of sync.WaitGroup that starts a function and counts it, and wg.Wait blocks until all have returned. Each goroutine writes its result into its own element of results, so no two goroutines write the same memory. In Rust each worker runs in a thread from thread::spawn, and join waits for the thread and hands back what run returned.

The Go channels have a buffer large enough for every message a worker can send in a step, and the Rust channels have no limit. A send therefore completes even when the receiving worker has stopped, as it does on a real network, where data sent to a dead machine leaves the sender and is acknowledged by nobody. The sender finds out only when it waits for something in return.

The time between checkpoints

The last piece is arithmetic. Saving a checkpoint, the model’s parameters and the position in the training data, takes time during which no training happens. Saving rarely wastes little time on checkpoints and loses much work at each failure. Saving often does the opposite. Somewhere between lies the interval that wastes the least.

Let C be the time one checkpoint takes, and M the job’s mean time between failures (MTBF), the average time from one failure to the next. With a checkpoint every T seconds:

• the fraction of time spent writing checkpoints is C / T;
• a failure arrives, on average, halfway between two checkpoints, so it destroys T / 2 of work, and with one failure every M seconds the fraction of time lost that way is T / (2M).

The total, C / T + T / (2M), is smallest when the two parts are equal, which happens at T = √(2CM). This is Young’s formula, published in 1974.

The M in the formula is the job’s, and a job fails when any one of its machines fails. A job on N machines therefore fails N times as often as one machine: its MTBF is one machine’s divided by N.

Listing 7.7 Young's formula, in Go. code/go/allreduce/main.go

// Costs and failure rates for the checkpoint arithmetic, in seconds.
const (
	// checkpointCost is the time one checkpoint takes to write.
	checkpointCost = 50.0
	// machineMTBF is the mean time between failures of one machine: 30 days.
	machineMTBF = 30 * 24 * 3600.0
)

// interval returns the time between checkpoints that loses the least work,
// by Young's formula: the square root of twice the checkpoint cost times the
// job's mean time between failures.
func interval(jobMTBF float64) float64 {
	return math.Sqrt(2 * checkpointCost * jobMTBF)
}

Listing 7.8 Young's formula, in Rust. code/rust/allreduce/src/main.rs

/// The time one checkpoint takes to write, in seconds.
const CHECKPOINT_COST: f64 = 50.0;
/// The mean time between failures of one machine, in seconds: 30 days.
const MACHINE_MTBF: f64 = 30.0 * 24.0 * 3600.0;

/// Returns the time between checkpoints that loses the least work, by Young's
/// formula: the square root of twice the checkpoint cost times the job's mean
/// time between failures.
fn interval(job_mtbf: f64) -> f64 {
    (2.0 * CHECKPOINT_COST * job_mtbf).sqrt()
}

The run

Listing 7.9 The healthy step, the failed step, and the checkpoint table, in Go. code/go/allreduce/main.go

func main() {
	const n = 5

	fmt.Printf("ring all-reduce over %d workers, %d chunks each:\n", n, n)
	workers, _ := ring(n, -1, 0)

	for i := range n {
		own := make([]int, n)

		for chunk := range own {
			own[chunk] = i + 1 + 10*chunk
		}

		fmt.Printf("  %s starts with %-17s and ends with %s\n", name(i), show(own), show(workers[i].gradient))
	}

	sent := 2 * (n - 1)
	fmt.Printf("  each worker sent %d chunks, %.1f times its gradient; %d messages in all\n",
		sent, float64(sent)/n, sent*n)

	fmt.Printf("the same step, with %s failing after 2 of its %d exchanges:\n", name(2), sent)
	_, results := ring(n, 2, 2)

	for i, r := range results {
		fmt.Printf("  %s stopped after %d of %d: %v\n", name(i), r.steps, sent, r.err)
	}

	fmt.Println("checkpoint interval, for a 50 s checkpoint and one failure per machine every 30 days:")
	fmt.Println("  machines  job fails every  checkpoint every  time lost")

	for _, machines := range []int{1, 5, 64, 512} {
		jobMTBF := machineMTBF / float64(machines)
		every := interval(jobMTBF)
		lost := checkpointCost/every + every/(2*jobMTBF)
		fmt.Printf("  %8d  %10.1f hours  %9.0f minutes  %8.1f%%\n", machines, jobMTBF/3600, every/60, lost*100)
	}
}

Listing 7.10 The healthy step, the failed step, and the checkpoint table, in Rust. code/rust/allreduce/src/main.rs

fn main() {
    const N: usize = 5;

    println!("ring all-reduce over {N} workers, {N} chunks each:");

    for (worker, _, _) in ring(N, None) {
        println!(
            "  {} starts with {:<17} and ends with {}",
            name(worker.id),
            show(&own_gradient(worker.id, N)),
            show(&worker.gradient)
        );
    }

    let sent = 2 * (N - 1);
    println!(
        "  each worker sent {sent} chunks, {:.1} times its gradient; {} messages in all",
        sent as f64 / N as f64,
        sent * N
    );

    println!(
        "the same step, with {} failing after 2 of its {sent} exchanges:",
        name(2)
    );

    for (worker, steps, why) in ring(N, Some((2, 2))) {
        let why = why.unwrap_or_else(|| "finished".to_owned());
        println!(
            "  {} stopped after {steps} of {sent}: {why}",
            name(worker.id)
        );
    }

    println!(
        "checkpoint interval, for a 50 s checkpoint and one failure per machine every 30 days:"
    );
    println!("  machines  job fails every  checkpoint every  time lost");

    for machines in [1, 5, 64, 512] {
        let job_mtbf = MACHINE_MTBF / f64::from(machines);
        let every = interval(job_mtbf);
        let lost = CHECKPOINT_COST / every + every / (2.0 * job_mtbf);
        println!(
            "  {machines:8}  {:10.1} hours  {:9.0} minutes  {:8.1}%",
            job_mtbf / 3600.0,
            every / 60.0,
            lost * 100.0
        );
    }
}

$ go run ./allreduce
ring all-reduce over 5 workers, 5 chunks each:
  thor-1 starts with 1 11 21 31 41     and ends with 15 65 115 165 215
  thor-2 starts with 2 12 22 32 42     and ends with 15 65 115 165 215
  thor-3 starts with 3 13 23 33 43     and ends with 15 65 115 165 215
  thor-4 starts with 4 14 24 34 44     and ends with 15 65 115 165 215
  thor-5 starts with 5 15 25 35 45     and ends with 15 65 115 165 215
  each worker sent 8 chunks, 1.6 times its gradient; 40 messages in all
the same step, with thor-3 failing after 2 of its 8 exchanges:
  thor-1 stopped after 4 of 8: no chunk from thor-5 within 100ms
  thor-2 stopped after 5 of 8: no chunk from thor-1 within 100ms
  thor-3 stopped after 2 of 8: failed
  thor-4 stopped after 2 of 8: no chunk from thor-3 within 100ms
  thor-5 stopped after 3 of 8: no chunk from thor-4 within 100ms
checkpoint interval, for a 50 s checkpoint and one failure per machine every 30 days:
  machines  job fails every  checkpoint every  time lost
         1       720.0 hours        268 minutes       0.6%
         5       144.0 hours        120 minutes       1.4%
        64        11.2 hours         34 minutes       5.0%
       512         1.4 hours         12 minutes      14.1%

In the first part every worker ends with the same five sums, starting with the 15 we computed by hand. Each sent 8 chunks, 1.6 gradients’ worth, as the count predicted.

In the second part thor-3 fails after two exchanges. thor-4, on its right, also completed two and stopped at the third, waiting for thor-3. thor-5 completed three, thor-1 four, and thor-2 five. Reading the workers in ring order from the failure, the counts are 2, 3, 4, 5: the stall moved one worker per exchange, as we reasoned. Each worker names the neighbour it was waiting for, and only one of the four names the machine that failed. A person who reads the report from thor-2 alone would suspect thor-1, which is healthy. Finding the worker that every chain of blame leads back to is the same counting exercise as chapter 4’s triage.

Animation 7.1 replays both runs on the ring, following chunk 0 in the first and the count of completed exchanges in the second.

The third part assumes that one machine fails every 30 days and that a checkpoint takes 50 seconds, which is 100 GB written at 2 GB per second. These are assumptions for the arithmetic, and the formula takes whatever numbers a real fleet measures.

A job on one machine should checkpoint every 268 minutes and loses 0.6 percent of its time. Our five Thors together fail every 144 hours, and the best interval is exactly two hours, at a cost of 1.4 percent. At 512 machines the job fails every 1.4 hours on average. It should checkpoint every 12 minutes, and 14.1 percent of its time goes to checkpoints and repeated work. The machines are no less reliable than before. There are more of them in one job, and this is why large training jobs invest in making checkpoints fast.

The complete program

Listing 7.11 The complete program in Go. code/go/allreduce/main.go

// Command allreduce runs one step of distributed training's communication:
// five workers each hold a gradient, and every worker has to end the step
// holding the sum of all five.
//
// The workers are goroutines connected in a ring by channels, and they run
// the ring all-reduce algorithm. The program runs the step three ways: with
// every worker healthy, with one worker failing part of the way through, and
// then as arithmetic, to find how often a job of a given size should save a
// checkpoint.
//
// Run it from code/go with:
//
//	go run ./allreduce
package main

import (
	"errors"
	"fmt"
	"math"
	"strings"
	"sync"
	"time"
)

// message carries one chunk of a gradient from a worker to its right-hand
// neighbour.
type message struct {
	// chunk is the index of the chunk within the gradient.
	chunk int
	// value is the chunk's contents: one number stands for the whole chunk.
	value int
}

// Worker is one machine's part in the step.
type Worker struct {
	// id is the worker's position in the ring, from 0.
	id int
	// gradient holds one value per chunk. It starts as this worker's own
	// gradient and ends as the sum over all workers.
	gradient []int
	// in receives from the left-hand neighbour, and out sends to the right.
	in  <-chan message
	out chan<- message
	// failAfter is the number of messages the worker exchanges before it
	// fails, or -1 for a worker that stays healthy.
	failAfter int
}

// errFailed is the result of the worker that the run makes fail.
var errFailed = errors.New("failed")

// patience is how long a worker waits for its neighbour before giving up.
const patience = 100 * time.Millisecond


// name returns the host name of the worker at a ring position.
func name(id int) string { return fmt.Sprintf("thor-%d", id+1) }

// mod returns a modulo n as a value from 0 to n-1, for negative a as well.
func mod(a, n int) int { return ((a % n) + n) % n }

// Run performs the worker's 2(n-1) exchanges. In each one it sends a chunk to
// the right and receives a chunk from the left. During the first n-1
// exchanges it adds what it receives to its own chunk. During the last n-1
// the chunks it receives are finished sums, and it stores them. Run returns
// the number of exchanges completed and the reason it stopped early, if any.
func (w *Worker) Run(n int) (int, error) {
	for step := range 2 * (n - 1) {
		if step == w.failAfter {
			return step, errFailed
		}

		// The chunk to send moves one place back around the gradient each
		// step, so after the first step it is the chunk just received.
		chunk := mod(w.id-step, n)
		w.out <- message{chunk, w.gradient[chunk]}
		reducing := step < n-1

		select {
		case m := <-w.in:
			if reducing {
				w.gradient[m.chunk] += m.value
			} else {
				w.gradient[m.chunk] = m.value
			}
		case <-time.After(patience):
			return step, fmt.Errorf("no chunk from %s within %v", name(mod(w.id-1, n)), patience)
		}
	}

	return 2 * (n - 1), nil
}


// result is what one worker reports when it stops.
type result struct {
	steps int
	err   error
}

// ring connects n workers in a circle, runs them all at once, and returns the
// workers and their results. The worker at position failing stops after
// failAfter exchanges; pass -1 for a run with no failure.
func ring(n, failing, failAfter int) ([]*Worker, []result) {
	// Each link holds every message a worker can send in one step, so a send
	// goes through even when the receiver has stopped.
	links := make([]chan message, n)

	for i := range links {
		links[i] = make(chan message, 2*(n-1))
	}

	workers := make([]*Worker, n)

	for i := range workers {
		gradient := make([]int, n)

		for chunk := range gradient {
			gradient[chunk] = i + 1 + 10*chunk
		}

		workers[i] = &Worker{id: i, gradient: gradient, in: links[i], out: links[mod(i+1, n)], failAfter: -1}
	}

	if failing >= 0 {
		workers[failing].failAfter = failAfter
	}

	results := make([]result, n)
	var wg sync.WaitGroup

	for i, w := range workers {
		wg.Go(func() {
			steps, err := w.Run(n)
			results[i] = result{steps, err}
		})
	}

	wg.Wait()

	return workers, results
}


// show formats a gradient as its values separated by spaces.
func show(gradient []int) string {
	return strings.Trim(fmt.Sprint(gradient), "[]")
}

// Costs and failure rates for the checkpoint arithmetic, in seconds.
const (
	// checkpointCost is the time one checkpoint takes to write.
	checkpointCost = 50.0
	// machineMTBF is the mean time between failures of one machine: 30 days.
	machineMTBF = 30 * 24 * 3600.0
)

// interval returns the time between checkpoints that loses the least work,
// by Young's formula: the square root of twice the checkpoint cost times the
// job's mean time between failures.
func interval(jobMTBF float64) float64 {
	return math.Sqrt(2 * checkpointCost * jobMTBF)
}


func main() {
	const n = 5

	fmt.Printf("ring all-reduce over %d workers, %d chunks each:\n", n, n)
	workers, _ := ring(n, -1, 0)

	for i := range n {
		own := make([]int, n)

		for chunk := range own {
			own[chunk] = i + 1 + 10*chunk
		}

		fmt.Printf("  %s starts with %-17s and ends with %s\n", name(i), show(own), show(workers[i].gradient))
	}

	sent := 2 * (n - 1)
	fmt.Printf("  each worker sent %d chunks, %.1f times its gradient; %d messages in all\n",
		sent, float64(sent)/n, sent*n)

	fmt.Printf("the same step, with %s failing after 2 of its %d exchanges:\n", name(2), sent)
	_, results := ring(n, 2, 2)

	for i, r := range results {
		fmt.Printf("  %s stopped after %d of %d: %v\n", name(i), r.steps, sent, r.err)
	}

	fmt.Println("checkpoint interval, for a 50 s checkpoint and one failure per machine every 30 days:")
	fmt.Println("  machines  job fails every  checkpoint every  time lost")

	for _, machines := range []int{1, 5, 64, 512} {
		jobMTBF := machineMTBF / float64(machines)
		every := interval(jobMTBF)
		lost := checkpointCost/every + every/(2*jobMTBF)
		fmt.Printf("  %8d  %10.1f hours  %9.0f minutes  %8.1f%%\n", machines, jobMTBF/3600, every/60, lost*100)
	}
}

Listing 7.12 The complete program in Rust. code/rust/allreduce/src/main.rs

//! Runs one step of distributed training's communication: five workers each
//! hold a gradient, and every worker has to end the step holding the sum of
//! all five.
//!
//! The workers are threads connected in a ring by channels, and they run the
//! ring all-reduce algorithm. The program runs the step three ways: with every
//! worker healthy, with one worker failing part of the way through, and then
//! as arithmetic, to find how often a job of a given size should save a
//! checkpoint.
//!
//! Run it from code/rust with:
//!
//! ```text
//! cargo run -p allreduce
//! ```

use std::sync::mpsc::{self, Receiver, Sender};
use std::thread;
use std::time::Duration;

/// Carries one chunk of a gradient from a worker to its right-hand neighbour.
struct Message {
    /// The index of the chunk within the gradient.
    chunk: usize,
    /// The chunk's contents: one number stands for the whole chunk.
    value: i64,
}

/// One machine's part in the step.
struct Worker {
    /// The worker's position in the ring, from 0.
    id: usize,
    /// One value per chunk. It starts as this worker's own gradient and ends
    /// as the sum over all workers.
    gradient: Vec<i64>,
    /// Receives from the left-hand neighbour.
    inbox: Receiver<Message>,
    /// Sends to the right-hand neighbour.
    out: Sender<Message>,
    /// The number of messages the worker exchanges before it fails, if the
    /// run makes it fail.
    fail_after: Option<usize>,
}

/// How long a worker waits for its neighbour before giving up.
const PATIENCE: Duration = Duration::from_millis(100);

/// Returns the host name of the worker at a ring position.
fn name(id: usize) -> String {
    format!("thor-{}", id + 1)
}

impl Worker {
    /// Performs the worker's 2(n-1) exchanges. In each one it sends a chunk to
    /// the right and receives a chunk from the left. During the first n-1
    /// exchanges it adds what it receives to its own chunk. During the last
    /// n-1 the chunks it receives are finished sums, and it stores them.
    /// Returns the worker with the number of exchanges completed and the
    /// reason it stopped early, if any.
    fn run(mut self, n: usize) -> (Self, usize, Option<String>) {
        for step in 0..2 * (n - 1) {
            if self.fail_after == Some(step) {
                return (self, step, Some("failed".to_owned()));
            }

            // The chunk to send moves one place back around the gradient each
            // step, so after the first step it is the chunk just received.
            // Adding 2n keeps the subtraction from going below zero.
            let chunk = (self.id + 2 * n - step) % n;
            // The right-hand neighbour may have stopped; its channel still accepts.
            let _ = self.out.send(Message {
                chunk,
                value: self.gradient[chunk],
            });
            let reducing = step < n - 1;

            match self.inbox.recv_timeout(PATIENCE) {
                Ok(m) if reducing => self.gradient[m.chunk] += m.value,
                Ok(m) => self.gradient[m.chunk] = m.value,
                Err(_) => {
                    let why = format!(
                        "no chunk from {} within {PATIENCE:?}",
                        name((self.id + n - 1) % n)
                    );

                    return (self, step, Some(why));
                }
            }
        }

        (self, 2 * (n - 1), None)
    }
}

/// Connects n workers in a circle, runs them all at once, and returns each
/// worker with its result. The worker named in `failing` stops after that
/// many exchanges.
fn ring(n: usize, failing: Option<(usize, usize)>) -> Vec<(Worker, usize, Option<String>)> {
    // One channel per link. `senders[i]` leads to worker i.
    let (senders, receivers): (Vec<_>, Vec<_>) = (0..n).map(|_| mpsc::channel()).unzip();
    let workers: Vec<Worker> = receivers
        .into_iter()
        .enumerate()
        .map(|(id, inbox)| Worker {
            id,
            gradient: own_gradient(id, n),
            inbox,
            out: senders[(id + 1) % n].clone(),
            fail_after: failing
                .filter(|(who, _)| *who == id)
                .map(|(_, after)| after),
        })
        .collect();

    let handles: Vec<_> = workers
        .into_iter()
        .map(|w| thread::spawn(move || w.run(n)))
        .collect();

    handles
        .into_iter()
        .map(|h| h.join().expect("a worker panicked"))
        .collect()
}

/// Returns the gradient a worker starts with: a different value in each chunk.
fn own_gradient(id: usize, n: usize) -> Vec<i64> {
    (0..n).map(|chunk| (id + 1 + 10 * chunk) as i64).collect()
}

/// Formats a gradient as its values separated by spaces.
fn show(gradient: &[i64]) -> String {
    gradient
        .iter()
        .map(i64::to_string)
        .collect::<Vec<_>>()
        .join(" ")
}

/// The time one checkpoint takes to write, in seconds.
const CHECKPOINT_COST: f64 = 50.0;
/// The mean time between failures of one machine, in seconds: 30 days.
const MACHINE_MTBF: f64 = 30.0 * 24.0 * 3600.0;

/// Returns the time between checkpoints that loses the least work, by Young's
/// formula: the square root of twice the checkpoint cost times the job's mean
/// time between failures.
fn interval(job_mtbf: f64) -> f64 {
    (2.0 * CHECKPOINT_COST * job_mtbf).sqrt()
}

fn main() {
    const N: usize = 5;

    println!("ring all-reduce over {N} workers, {N} chunks each:");

    for (worker, _, _) in ring(N, None) {
        println!(
            "  {} starts with {:<17} and ends with {}",
            name(worker.id),
            show(&own_gradient(worker.id, N)),
            show(&worker.gradient)
        );
    }

    let sent = 2 * (N - 1);
    println!(
        "  each worker sent {sent} chunks, {:.1} times its gradient; {} messages in all",
        sent as f64 / N as f64,
        sent * N
    );

    println!(
        "the same step, with {} failing after 2 of its {sent} exchanges:",
        name(2)
    );

    for (worker, steps, why) in ring(N, Some((2, 2))) {
        let why = why.unwrap_or_else(|| "finished".to_owned());
        println!(
            "  {} stopped after {steps} of {sent}: {why}",
            name(worker.id)
        );
    }

    println!(
        "checkpoint interval, for a 50 s checkpoint and one failure per machine every 30 days:"
    );
    println!("  machines  job fails every  checkpoint every  time lost");

    for machines in [1, 5, 64, 512] {
        let job_mtbf = MACHINE_MTBF / f64::from(machines);
        let every = interval(job_mtbf);
        let lost = CHECKPOINT_COST / every + every / (2.0 * job_mtbf);
        println!(
            "  {machines:8}  {:10.1} hours  {:9.0} minutes  {:8.1}%",
            job_mtbf / 3600.0,
            every / 60.0,
            lost * 100.0
        );
    }
}

References

• NVIDIA, NCCL documentation: https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/index.html
• A. Gibiansky, “Bringing HPC Techniques to Deep Learning”, Baidu Research, 2017, which describes the ring all-reduce used here.
• P. Patarasuk and X. Yuan, “Bandwidth optimal all-reduce algorithms for clusters of workstations”, Journal of Parallel and Distributed Computing, 2009.
• J. W. Young, “A first order approximation to the optimum checkpoint interval”, Communications of the ACM, 17(9), 1974.
• MPI Forum, MPI: A Message-Passing Interface Standard, chapter on collective communication: https://www.mpi-forum.org/docs/

8. Placing jobs on GPUs

Chapter 7 gave us two facts about a training job that spans several machines. It makes progress only when all of its workers are running, and its speed is set by the slowest network link among them. In this chapter, we will see what those facts ask of the scheduler, the component that chooses a machine for each pod.

We will first look at how the Kubernetes scheduler places one pod, and at what goes wrong when a job of several pods is placed that way. Then we will write gangsched, in Go and in Rust: a model that hands out the eight GPUs of our fleet one pod at a time, then one whole job at a time, and then chooses a job’s machines by the links between them. Finally we will map the model onto the pieces of a real cluster: Kueue, pod groups, and topology labels.

How one pod is placed

The scheduler is a controller in the sense of chapter 1. It watches for pods that have no machine yet, which it recognises by an empty spec.nodeName, and it handles them one at a time, in three stages.

1. Filter. It discards every node that cannot run the pod: nodes without enough free CPU, memory, or nvidia.com/gpu, nodes that are cordoned, as in chapter 4, and nodes whose labels do not satisfy the pod’s nodeSelector, as in chapter 5.
2. Score. It gives each remaining node a number by a set of rules, such as preferring nodes that already have the pod’s image.
3. Bind. It writes the name of the best node into the pod’s spec.nodeName. The kubelet on that node sees the pod through its own watch and starts the containers.

Each pod is a separate decision. The scheduler has no notion that five pods belong together, and that is the source of the first problem.

Two jobs, each with part of what it needs

Suppose two training jobs arrive at the same moment, each made of five worker pods that ask for one GPU, and our fleet has eight GPUs. Ten pods wait, and the scheduler takes them in turn. The first eight are placed and the last two wait.

If the pods of the two jobs are interleaved in the queue, each job ends up with four of its five workers. Neither can run a single step, because an all-reduce needs every worker. Each holds four GPUs and waits for a fifth, and the only way a fifth can become free is for the other job to give one up. Nothing in the system will make either do so. All eight GPUs are allocated, the dashboards show a cluster in full use, and no training is happening. Figure 8.1 shows the state.

The remedy is to make the job the unit of the decision. A job is admitted when all of its workers can be placed, and all of them are placed together. Until then it holds nothing. This is called gang scheduling, the gang being the group of pods that start together or stay queued together.

With gang scheduling the same two jobs behave differently. The first is admitted and takes five GPUs. The second needs five, finds three, and waits without taking any. One job runs at full speed, the three free GPUs remain available to smaller jobs, and when the first job finishes the second starts.

Which machines

Admitting a job settles when it starts. Where it runs matters as much, because of the links.

Our eight machines are connected in three groups, which we will call domains. The five Thors share the 25 GbE switch on the shelf. The two Sparks share the 200 GbE cable on the desk. The Orin NX sits on the bench with 1 GbE. Between domains, the only path is the house network that all machines use for SSH and for reaching the API server, at 1 GbE. Figure 8.2 draws this as a tree.

Chapter 7 showed that a ring all-reduce runs at the speed of its slowest link and moves 2(N − 1)/N gradients per worker. For the 14 GB gradient of that chapter, two workers that share the Sparks’ cable need 14 / 25 = 0.56 seconds for each all-reduce. The same two workers placed in different domains exchange their gradients over the house network and need 14 / 0.125 = 112 seconds. The job is the same and the GPUs are equally fast. The placement makes it two hundred times slower.

So the placement rule has two parts.

• Keep a job inside one domain when one can hold it. Then the job’s slowest link is the domain’s own.
• Among the domains that can hold it, choose the one with the fewest free machines. A two-worker job fits on the desk and on the shelf. Putting it on the shelf leaves three Thors, and a later job that needs four of them no longer fits anywhere. Putting it on the desk uses the desk completely and leaves the shelf whole.

The second part is a defence against fragmentation: free GPUs scattered across domains in groups too small for the jobs that are waiting. The rule is a heuristic. It makes good choices for the jobs it has seen and cannot know what will arrive next.

Building gangsched

The fleet as the scheduler sees it

Each machine has one GPU, so a node is a name, a domain, and the job that holds its GPU. The machines are listed in name order, which is the order in which a listing of Node objects from the API server arrives.

Listing 8.1 Nodes, links, and jobs, in Go. code/go/gangsched/main.go

// Node is one machine with one GPU.
type Node struct {
	// Name is the machine's host name.
	Name string
	// Domain names the group of machines that share a fast link.
	Domain string
	// Job is the job that holds the machine's GPU, or "" when it is free.
	Job string
}

// Link is the network path a group of machines shares.
type Link struct {
	// Label is how the link is printed.
	Label string
	// GBps is what the link carries, in gigabytes per second.
	GBps float64
}

// links maps each domain to its link.
var links = map[string]Link{
	"shelf": {"the 25 GbE switch", 3.125},
	"desk":  {"the 200 GbE cable", 25},
	"bench": {"1 GbE", 0.125},
}

// houseNetwork is the link between machines in different domains.
var houseNetwork = Link{"the 1 GbE house network", 0.125}

// newFleet returns the machines in name order, all free.
func newFleet() []*Node {
	return []*Node{
		{Name: "orin-nx-1", Domain: "bench"},
		{Name: "spark-1", Domain: "desk"}, {Name: "spark-2", Domain: "desk"},
		{Name: "thor-1", Domain: "shelf"}, {Name: "thor-2", Domain: "shelf"}, {Name: "thor-3", Domain: "shelf"},
		{Name: "thor-4", Domain: "shelf"}, {Name: "thor-5", Domain: "shelf"},
	}
}

// Job is a group of workers that have to run at the same time, one GPU each.
type Job struct {
	Name    string
	Workers int
}

Listing 8.2 Nodes, links, and jobs, in Rust. code/rust/gangsched/src/main.rs

/// One machine with one GPU.
struct Node {
    /// The machine's host name.
    name: &'static str,
    /// The group of machines that share a fast link.
    domain: Domain,
    /// The job that holds the machine's GPU, if any.
    job: Option<&'static str>,
}

/// A group of machines that share a fast link.
#[derive(Clone, Copy, PartialEq)]
enum Domain {
    /// The Orin NX, alone on 1 GbE.
    Bench,
    /// The two Sparks, on a 200 GbE cable.
    Desk,
    /// The five Thors, on a 25 GbE switch.
    Shelf,
}

/// The network path a group of machines shares.
struct Link {
    /// How the link is printed.
    label: &'static str,
    /// What the link carries, in gigabytes per second.
    gbps: f64,
}

impl Domain {
    /// Returns the link the machines of this domain share.
    fn link(self) -> Link {
        match self {
            Domain::Shelf => Link {
                label: "the 25 GbE switch",
                gbps: 3.125,
            },
            Domain::Desk => Link {
                label: "the 200 GbE cable",
                gbps: 25.0,
            },
            Domain::Bench => Link {
                label: "1 GbE",
                gbps: 0.125,
            },
        }
    }
}

/// The link between machines in different domains.
const HOUSE_NETWORK: Link = Link {
    label: "the 1 GbE house network",
    gbps: 0.125,
};

/// Returns the machines in name order, all free.
fn new_fleet() -> Vec<Node> {
    [
        ("orin-nx-1", Domain::Bench),
        ("spark-1", Domain::Desk),
        ("spark-2", Domain::Desk),
        ("thor-1", Domain::Shelf),
        ("thor-2", Domain::Shelf),
        ("thor-3", Domain::Shelf),
        ("thor-4", Domain::Shelf),
        ("thor-5", Domain::Shelf),
    ]
    .into_iter()
    .map(|(name, domain)| Node {
        name,
        domain,
        job: None,
    })
    .collect()
}

/// A group of workers that have to run at the same time, one GPU each.
#[derive(Clone, Copy)]
struct Job {
    name: &'static str,
    workers: usize,
}

The Go version names a domain with a string and looks its link up in a map. The Rust version makes the domain an enum and the lookup a match, so a domain with no link is a compile error.

One pod at a time

podAtATime plays the scheduler without gangs. It goes round the jobs, giving each one free GPU per turn, until a whole round passes in which nobody could be given anything.

Listing 8.3 Handing out GPUs one pod at a time, in Go. code/go/gangsched/main.go

// podAtATime gives each job one GPU in turn until none are left, the way a
// scheduler that sees only single pods would, and reports what each job
// ended up holding.
func podAtATime(fleet []*Node, jobs []Job) {
	held := map[string][]*Node{}

	for progress := true; progress; {
		progress = false

		for _, job := range jobs {
			open := free(fleet)

			if len(held[job.Name]) < job.Workers && len(open) > 0 {
				open[0].Job = job.Name
				held[job.Name] = append(held[job.Name], open[0])
				progress = true
			}
		}
	}

	running := 0

	for _, job := range jobs {
		fmt.Printf("  %-8s holds %d of %d GPUs: %s\n", job.Name, len(held[job.Name]), job.Workers, names(held[job.Name]))

		if len(held[job.Name]) == job.Workers {
			running++
		}
	}

	fmt.Printf("  %d of %d GPUs are held and %d jobs are running\n", len(fleet)-len(free(fleet)), len(fleet), running)
}

Listing 8.4 Handing out GPUs one pod at a time, in Rust. code/rust/gangsched/src/main.rs

/// Gives each job one GPU in turn until none are left, the way a scheduler
/// that sees only single pods would, and reports what each job ended up
/// holding.
fn pod_at_a_time(mut fleet: Vec<Node>, jobs: &[Job]) {
    let held = |fleet: &[Node], job: &Job| -> Vec<usize> {
        (0..fleet.len())
            .filter(|&i| fleet[i].job == Some(job.name))
            .collect()
    };
    let mut progress = true;

    while progress {
        progress = false;

        for job in jobs {
            if held(&fleet, job).len() < job.workers
                && let Some(&first) = free(&fleet).first()
            {
                fleet[first].job = Some(job.name);
                progress = true;
            }
        }
    }

    let mut running = 0;

    for job in jobs {
        let mine = held(&fleet, job);
        println!(
            "  {:<8} holds {} of {} GPUs: {}",
            job.name,
            mine.len(),
            job.workers,
            names(&fleet, &mine)
        );

        if mine.len() == job.workers {
            running += 1;
        }
    }

    let taken = fleet.len() - free(&fleet).len();
    println!(
        "  {taken} of {} GPUs are held and {running} jobs are running",
        fleet.len()
    );
}

Two ways to choose machines

A placer is a function that takes the free machines and a number of workers and returns the machines to use, or nothing. Making it a function lets admit work with either policy.

firstFree takes the first free machines in the list. It is what we get when nobody has thought about placement.

byTopology groups the free machines by domain, keeps the domains with enough room, and takes the smallest of those. When no single domain can hold the job, it falls back to firstFree, and the job runs across domains at the speed of the house network. A real scheduler would let the job’s owner choose between that and waiting.

Listing 8.5 The two placers, in Go. code/go/gangsched/main.go

// A placer chooses machines for a job from the free ones, or returns nil when
// it cannot place every worker.
type placer func(open []*Node, workers int) []*Node

// firstFree takes the first free machines in name order.
func firstFree(open []*Node, workers int) []*Node {
	if len(open) < workers {
		return nil
	}

	return open[:workers]
}

// byTopology looks for one domain that can hold the whole job and, among
// those, takes the one with the fewest free machines, so that large domains
// stay whole for large jobs. A job that fits in no single domain is spread
// over the first free machines.
func byTopology(open []*Node, workers int) []*Node {
	byDomain := map[string][]*Node{}
	var order []string

	for _, n := range open {
		if byDomain[n.Domain] == nil {
			order = append(order, n.Domain)
		}

		byDomain[n.Domain] = append(byDomain[n.Domain], n)
	}

	best := ""

	for _, domain := range order {
		size := len(byDomain[domain])

		if size >= workers && (best == "" || size < len(byDomain[best])) {
			best = domain
		}
	}

	if best == "" {
		return firstFree(open, workers)
	}

	return byDomain[best][:workers]
}

Listing 8.6 The two placers, in Rust. code/rust/gangsched/src/main.rs

/// Chooses machines for a job from the free ones, or returns `None` when it
/// cannot place every worker.
type Placer = fn(fleet: &[Node], open: &[usize], workers: usize) -> Option<Vec<usize>>;

/// Takes the first free machines in name order.
fn first_free(_fleet: &[Node], open: &[usize], workers: usize) -> Option<Vec<usize>> {
    open.get(..workers).map(<[usize]>::to_vec)
}

/// Looks for one domain that can hold the whole job and, among those, takes
/// the one with the fewest free machines, so that large domains stay whole
/// for large jobs. A job that fits in no single domain is spread over the
/// first free machines.
fn by_topology(fleet: &[Node], open: &[usize], workers: usize) -> Option<Vec<usize>> {
    let in_domain = |domain: Domain| -> Vec<usize> {
        open.iter()
            .copied()
            .filter(|&i| fleet[i].domain == domain)
            .collect()
    };
    [Domain::Bench, Domain::Desk, Domain::Shelf]
        .into_iter()
        .map(in_domain)
        .filter(|members| members.len() >= workers)
        .min_by_key(Vec::len)
        .map(|members| members[..workers].to_vec())
        .or_else(|| first_free(fleet, open, workers))
}

What a placement costs

To compare placements we need a number. slowest finds the link a set of machines is limited by: their domain’s link if they share one, the house network if they do not. exchangeSeconds applies chapter 7’s count to it.

Listing 8.7 The slowest link and the time of one all-reduce, in Go. code/go/gangsched/main.go

// gradientGB is the size of the gradient every worker contributes per step.
const gradientGB = 14.0

// slowest returns the link a job's all-reduce is limited by: the domain's
// link when every machine is in one domain, and the house network otherwise.
func slowest(nodes []*Node) Link {
	for _, n := range nodes {
		if n.Domain != nodes[0].Domain {
			return houseNetwork
		}
	}

	return links[nodes[0].Domain]
}

// exchangeSeconds is the time one ring all-reduce takes: each worker sends
// 2(n-1)/n gradients over the slowest link.
func exchangeSeconds(nodes []*Node) float64 {
	n := float64(len(nodes))

	return 2 * (n - 1) / n * gradientGB / slowest(nodes).GBps
}

Listing 8.8 The slowest link and the time of one all-reduce, in Rust. code/rust/gangsched/src/main.rs

/// The size of the gradient every worker contributes per step, in gigabytes.
const GRADIENT_GB: f64 = 14.0;

/// Returns the link a job's all-reduce is limited by: the domain's link when
/// every machine is in one domain, and the house network otherwise.
fn slowest(fleet: &[Node], nodes: &[usize]) -> Link {
    let first = fleet[nodes[0]].domain;
    if nodes.iter().all(|&i| fleet[i].domain == first) {
        first.link()
    } else {
        HOUSE_NETWORK
    }
}

/// Returns the time one ring all-reduce takes: each worker sends 2(n-1)/n
/// gradients over the slowest link.
fn exchange_seconds(fleet: &[Node], nodes: &[usize]) -> f64 {
    let n = nodes.len() as f64;

    2.0 * (n - 1.0) / n * GRADIENT_GB / slowest(fleet, nodes).gbps
}

Admitting a whole job

admit is gang scheduling in a dozen lines. It asks the placer for machines for every worker. If the placer returns nothing, the job waits and the fleet is untouched. Only when every worker has a machine does admit mark any of them as taken.

Listing 8.9 All of a job's workers, or none, in Go. code/go/gangsched/main.go

// admit starts a job only when every one of its workers has a machine. It
// reports the placement and what the job's all-reduce will cost, and returns
// false when the job has to wait.
func admit(fleet []*Node, job Job, place placer) bool {
	open := free(fleet)
	nodes := place(open, job.Workers)

	if nodes == nil {
		fmt.Printf("  %-8s waits: it needs %d GPUs and %d are free\n", job.Name, job.Workers, len(open))

		return false
	}

	for _, n := range nodes {
		n.Job = job.Name
	}

	fmt.Printf("  %-8s runs on %s\n", job.Name, names(nodes))

	if len(nodes) > 1 {
		fmt.Printf("           over %s: %.1f s for each all-reduce\n", slowest(nodes).Label, exchangeSeconds(nodes))
	}

	return true
}

// release frees every GPU a job holds.
func release(fleet []*Node, job string) {
	for _, n := range fleet {
		if n.Job == job {
			n.Job = ""
		}
	}

	fmt.Printf("  %-8s finishes\n", job)
}

Listing 8.10 All of a job's workers, or none, in Rust. code/rust/gangsched/src/main.rs

/// Starts a job only when every one of its workers has a machine. It reports
/// the placement and what the job's all-reduce will cost, and returns false
/// when the job has to wait.
fn admit(fleet: &mut [Node], job: Job, place: Placer) -> bool {
    let open = free(fleet);
    let Some(nodes) = place(fleet, &open, job.workers) else {
        println!(
            "  {:<8} waits: it needs {} GPUs and {} are free",
            job.name,
            job.workers,
            open.len()
        );

        return false;
    };

    for &i in &nodes {
        fleet[i].job = Some(job.name);
    }

    println!("  {:<8} runs on {}", job.name, names(fleet, &nodes));

    if nodes.len() > 1 {
        let link = slowest(fleet, &nodes);
        println!(
            "           over {}: {:.1} s for each all-reduce",
            link.label,
            exchange_seconds(fleet, &nodes)
        );
    }

    true
}

/// Frees every GPU a job holds.
fn release(fleet: &mut [Node], job: &str) {
    for node in fleet.iter_mut().filter(|n| n.job == Some(job)) {
        node.job = None;
    }

    println!("  {job:<8} finishes");
}

The run

Listing 8.11 Four experiments on the same fleet, in Go. code/go/gangsched/main.go

func main() {
	two := []Job{{"train-a", 5}, {"train-b", 5}}

	fmt.Println("one pod at a time, two jobs of five workers:")
	podAtATime(newFleet(), two)

	fmt.Println("whole jobs at a time:")
	fleet := newFleet()
	admit(fleet, two[0], byTopology)
	admit(fleet, two[1], byTopology)
	release(fleet, "train-a")
	admit(fleet, two[1], byTopology)

	three := []Job{{"pair", 2}, {"quad", 4}, {"single", 1}}
	fmt.Println("three jobs, on the first free machines:")
	fleet = newFleet()

	for _, job := range three {
		admit(fleet, job, firstFree)
	}

	fmt.Println("the same three jobs, placed by topology:")
	fleet = newFleet()

	for _, job := range three {
		admit(fleet, job, byTopology)
	}
}

Listing 8.12 Four experiments on the same fleet, in Rust. code/rust/gangsched/src/main.rs

fn main() {
    let two = [
        Job {
            name: "train-a",
            workers: 5,
        },
        Job {
            name: "train-b",
            workers: 5,
        },
    ];

    println!("one pod at a time, two jobs of five workers:");
    pod_at_a_time(new_fleet(), &two);

    println!("whole jobs at a time:");
    let mut fleet = new_fleet();
    admit(&mut fleet, two[0], by_topology);
    admit(&mut fleet, two[1], by_topology);
    release(&mut fleet, "train-a");
    admit(&mut fleet, two[1], by_topology);

    let three = [
        Job {
            name: "pair",
            workers: 2,
        },
        Job {
            name: "quad",
            workers: 4,
        },
        Job {
            name: "single",
            workers: 1,
        },
    ];
    println!("three jobs, on the first free machines:");
    let mut fleet = new_fleet();

    for job in three {
        admit(&mut fleet, job, first_free);
    }

    println!("the same three jobs, placed by topology:");
    let mut fleet = new_fleet();

    for job in three {
        admit(&mut fleet, job, by_topology);
    }
}

$ go run ./gangsched
one pod at a time, two jobs of five workers:
  train-a  holds 4 of 5 GPUs: orin-nx-1 spark-2 thor-2 thor-4
  train-b  holds 4 of 5 GPUs: spark-1 thor-1 thor-3 thor-5
  8 of 8 GPUs are held and 0 jobs are running
whole jobs at a time:
  train-a  runs on thor-1 thor-2 thor-3 thor-4 thor-5
           over the 25 GbE switch: 7.2 s for each all-reduce
  train-b  waits: it needs 5 GPUs and 3 are free
  train-a  finishes
  train-b  runs on thor-1 thor-2 thor-3 thor-4 thor-5
           over the 25 GbE switch: 7.2 s for each all-reduce
three jobs, on the first free machines:
  pair     runs on orin-nx-1 spark-1
           over the 1 GbE house network: 112.0 s for each all-reduce
  quad     runs on spark-2 thor-1 thor-2 thor-3
           over the 1 GbE house network: 168.0 s for each all-reduce
  single   runs on thor-4
the same three jobs, placed by topology:
  pair     runs on spark-1 spark-2
           over the 200 GbE cable: 0.6 s for each all-reduce
  quad     runs on thor-1 thor-2 thor-3 thor-4
           over the 25 GbE switch: 6.7 s for each all-reduce
  single   runs on orin-nx-1

The first experiment is the deadlock of Figure 8.1: train-a and train-b each hold four GPUs, all eight are taken, and zero jobs run.

In the second, train-a is admitted whole and placed on the five Thors, the only domain that can hold five workers. Its all-reduce takes 7.2 seconds over one 25 GbE lane, the figure from chapter 7’s table. train-b waits with three GPUs free, and starts on the Thors when train-a finishes.

The third experiment places three jobs on the first free machines. pair gets orin-nx-1 and spark-1, two machines in different domains, and each of its all-reduces takes 112 seconds over the house network. quad gets the other Spark and three Thors and takes 168 seconds per all-reduce.

The fourth places the same jobs by topology. pair fits on the desk and on the shelf, and the desk is the tighter fit, so it gets the two Sparks and 0.6 seconds per all-reduce. quad fits only on the shelf and gets four Thors at 6.7 seconds. single fits on the bench and on the one Thor that is left. Both have a single free machine, the bench is listed first, and it gets the Orin NX, which leaves the last Thor free. Every job runs inside one domain, and the communication of pair is about two hundred times faster than in the third experiment.

Animation 8.1 plays the four experiments on the fleet.

On a real cluster

The Kubernetes scheduler places single pods, so gang scheduling and topology are added around it. Three pieces do the work.

Kueue is a job queueing controller from the Kubernetes project. It decides when a job may start. A job managed by Kueue is created with spec.suspend: true, a field of the built-in Job object that tells the Job controller to create no pods yet. Kueue keeps such jobs in a queue and holds a quota of each resource for each team, in an object called a ClusterQueue. When the whole job fits within the quota, Kueue admits it by setting suspend to false, and only then do its pods exist and reach the scheduler. This is the admit of our model with quota in place of free machines. Kueue’s quotas are also where fairness between teams lives: each team has a share of the GPUs it is guaranteed, and a queue may borrow what another team is not using.

A pod group handles the case where quota is available and the machines are not. The scheduler can be extended with plugins, and a coscheduling plugin lets a set of pods declare that they form a group with a minimum size. The scheduler finds a node for each pod as usual and then, before binding, holds the pod until enough members of its group have found nodes too. If they do not within a timeout, all of them are released and retried. This removes the partial placements of our first experiment at the level of individual machines.

Topology labels tell the scheduler which machines share a link. A node carries labels, as we saw in chapter 5, and a cluster’s administrator adds one for each level of the network: the rack, the switch, the block of racks. For our fleet one label is enough:

$ kubectl label node thor-1 thor-2 thor-3 thor-4 thor-5 gpucloud.dev/domain=shelf
$ kubectl label node spark-1 spark-2 gpucloud.dev/domain=desk
$ kubectl label node orin-nx-1 gpucloud.dev/domain=bench

A job then asks for all of its pods to land on nodes with the same value of that label. Kueue’s topology-aware scheduling reads the levels from an object that lists the label keys from widest to narrowest, and picks the narrowest level whose free capacity holds the job, which is the rule byTopology applies to our three domains.

The complete program

Listing 8.13 The complete program in Go. code/go/gangsched/main.go

// Command gangsched models two decisions a scheduler makes for a training job
// that needs several GPUs at once: whether to start the job, and which
// machines to give it.
//
// The fleet is the eight machines of chapter 4, each with one GPU. The
// program first hands out GPUs one pod at a time and shows two jobs blocking
// each other. It then admits whole jobs, and finally compares two ways of
// choosing a job's machines by the network link the job would have to use.
//
// Run it from code/go with:
//
//	go run ./gangsched
package main

import (
	"fmt"
	"strings"
)

// Node is one machine with one GPU.
type Node struct {
	// Name is the machine's host name.
	Name string
	// Domain names the group of machines that share a fast link.
	Domain string
	// Job is the job that holds the machine's GPU, or "" when it is free.
	Job string
}

// Link is the network path a group of machines shares.
type Link struct {
	// Label is how the link is printed.
	Label string
	// GBps is what the link carries, in gigabytes per second.
	GBps float64
}

// links maps each domain to its link.
var links = map[string]Link{
	"shelf": {"the 25 GbE switch", 3.125},
	"desk":  {"the 200 GbE cable", 25},
	"bench": {"1 GbE", 0.125},
}

// houseNetwork is the link between machines in different domains.
var houseNetwork = Link{"the 1 GbE house network", 0.125}

// newFleet returns the machines in name order, all free.
func newFleet() []*Node {
	return []*Node{
		{Name: "orin-nx-1", Domain: "bench"},
		{Name: "spark-1", Domain: "desk"}, {Name: "spark-2", Domain: "desk"},
		{Name: "thor-1", Domain: "shelf"}, {Name: "thor-2", Domain: "shelf"}, {Name: "thor-3", Domain: "shelf"},
		{Name: "thor-4", Domain: "shelf"}, {Name: "thor-5", Domain: "shelf"},
	}
}

// Job is a group of workers that have to run at the same time, one GPU each.
type Job struct {
	Name    string
	Workers int
}


// free returns the machines whose GPU nobody holds, in name order.
func free(fleet []*Node) []*Node {
	var out []*Node

	for _, n := range fleet {
		if n.Job == "" {
			out = append(out, n)
		}
	}

	return out
}

// names joins the names of some machines with spaces.
func names(nodes []*Node) string {
	parts := make([]string, len(nodes))

	for i, n := range nodes {
		parts[i] = n.Name
	}

	return strings.Join(parts, " ")
}

// podAtATime gives each job one GPU in turn until none are left, the way a
// scheduler that sees only single pods would, and reports what each job
// ended up holding.
func podAtATime(fleet []*Node, jobs []Job) {
	held := map[string][]*Node{}

	for progress := true; progress; {
		progress = false

		for _, job := range jobs {
			open := free(fleet)

			if len(held[job.Name]) < job.Workers && len(open) > 0 {
				open[0].Job = job.Name
				held[job.Name] = append(held[job.Name], open[0])
				progress = true
			}
		}
	}

	running := 0

	for _, job := range jobs {
		fmt.Printf("  %-8s holds %d of %d GPUs: %s\n", job.Name, len(held[job.Name]), job.Workers, names(held[job.Name]))

		if len(held[job.Name]) == job.Workers {
			running++
		}
	}

	fmt.Printf("  %d of %d GPUs are held and %d jobs are running\n", len(fleet)-len(free(fleet)), len(fleet), running)
}


// A placer chooses machines for a job from the free ones, or returns nil when
// it cannot place every worker.
type placer func(open []*Node, workers int) []*Node

// firstFree takes the first free machines in name order.
func firstFree(open []*Node, workers int) []*Node {
	if len(open) < workers {
		return nil
	}

	return open[:workers]
}

// byTopology looks for one domain that can hold the whole job and, among
// those, takes the one with the fewest free machines, so that large domains
// stay whole for large jobs. A job that fits in no single domain is spread
// over the first free machines.
func byTopology(open []*Node, workers int) []*Node {
	byDomain := map[string][]*Node{}
	var order []string

	for _, n := range open {
		if byDomain[n.Domain] == nil {
			order = append(order, n.Domain)
		}

		byDomain[n.Domain] = append(byDomain[n.Domain], n)
	}

	best := ""

	for _, domain := range order {
		size := len(byDomain[domain])

		if size >= workers && (best == "" || size < len(byDomain[best])) {
			best = domain
		}
	}

	if best == "" {
		return firstFree(open, workers)
	}

	return byDomain[best][:workers]
}


// gradientGB is the size of the gradient every worker contributes per step.
const gradientGB = 14.0

// slowest returns the link a job's all-reduce is limited by: the domain's
// link when every machine is in one domain, and the house network otherwise.
func slowest(nodes []*Node) Link {
	for _, n := range nodes {
		if n.Domain != nodes[0].Domain {
			return houseNetwork
		}
	}

	return links[nodes[0].Domain]
}

// exchangeSeconds is the time one ring all-reduce takes: each worker sends
// 2(n-1)/n gradients over the slowest link.
func exchangeSeconds(nodes []*Node) float64 {
	n := float64(len(nodes))

	return 2 * (n - 1) / n * gradientGB / slowest(nodes).GBps
}


// admit starts a job only when every one of its workers has a machine. It
// reports the placement and what the job's all-reduce will cost, and returns
// false when the job has to wait.
func admit(fleet []*Node, job Job, place placer) bool {
	open := free(fleet)
	nodes := place(open, job.Workers)

	if nodes == nil {
		fmt.Printf("  %-8s waits: it needs %d GPUs and %d are free\n", job.Name, job.Workers, len(open))

		return false
	}

	for _, n := range nodes {
		n.Job = job.Name
	}

	fmt.Printf("  %-8s runs on %s\n", job.Name, names(nodes))

	if len(nodes) > 1 {
		fmt.Printf("           over %s: %.1f s for each all-reduce\n", slowest(nodes).Label, exchangeSeconds(nodes))
	}

	return true
}

// release frees every GPU a job holds.
func release(fleet []*Node, job string) {
	for _, n := range fleet {
		if n.Job == job {
			n.Job = ""
		}
	}

	fmt.Printf("  %-8s finishes\n", job)
}


func main() {
	two := []Job{{"train-a", 5}, {"train-b", 5}}

	fmt.Println("one pod at a time, two jobs of five workers:")
	podAtATime(newFleet(), two)

	fmt.Println("whole jobs at a time:")
	fleet := newFleet()
	admit(fleet, two[0], byTopology)
	admit(fleet, two[1], byTopology)
	release(fleet, "train-a")
	admit(fleet, two[1], byTopology)

	three := []Job{{"pair", 2}, {"quad", 4}, {"single", 1}}
	fmt.Println("three jobs, on the first free machines:")
	fleet = newFleet()

	for _, job := range three {
		admit(fleet, job, firstFree)
	}

	fmt.Println("the same three jobs, placed by topology:")
	fleet = newFleet()

	for _, job := range three {
		admit(fleet, job, byTopology)
	}
}

Listing 8.14 The complete program in Rust. code/rust/gangsched/src/main.rs

//! Models two decisions a scheduler makes for a training job that needs
//! several GPUs at once: whether to start the job, and which machines to give
//! it.
//!
//! The fleet is the eight machines of chapter 4, each with one GPU. The
//! program first hands out GPUs one pod at a time and shows two jobs blocking
//! each other. It then admits whole jobs, and finally compares two ways of
//! choosing a job's machines by the network link the job would have to use.
//!
//! Run it from code/rust with:
//!
//! ```text
//! cargo run -p gangsched
//! ```

/// One machine with one GPU.
struct Node {
    /// The machine's host name.
    name: &'static str,
    /// The group of machines that share a fast link.
    domain: Domain,
    /// The job that holds the machine's GPU, if any.
    job: Option<&'static str>,
}

/// A group of machines that share a fast link.
#[derive(Clone, Copy, PartialEq)]
enum Domain {
    /// The Orin NX, alone on 1 GbE.
    Bench,
    /// The two Sparks, on a 200 GbE cable.
    Desk,
    /// The five Thors, on a 25 GbE switch.
    Shelf,
}

/// The network path a group of machines shares.
struct Link {
    /// How the link is printed.
    label: &'static str,
    /// What the link carries, in gigabytes per second.
    gbps: f64,
}

impl Domain {
    /// Returns the link the machines of this domain share.
    fn link(self) -> Link {
        match self {
            Domain::Shelf => Link {
                label: "the 25 GbE switch",
                gbps: 3.125,
            },
            Domain::Desk => Link {
                label: "the 200 GbE cable",
                gbps: 25.0,
            },
            Domain::Bench => Link {
                label: "1 GbE",
                gbps: 0.125,
            },
        }
    }
}

/// The link between machines in different domains.
const HOUSE_NETWORK: Link = Link {
    label: "the 1 GbE house network",
    gbps: 0.125,
};

/// Returns the machines in name order, all free.
fn new_fleet() -> Vec<Node> {
    [
        ("orin-nx-1", Domain::Bench),
        ("spark-1", Domain::Desk),
        ("spark-2", Domain::Desk),
        ("thor-1", Domain::Shelf),
        ("thor-2", Domain::Shelf),
        ("thor-3", Domain::Shelf),
        ("thor-4", Domain::Shelf),
        ("thor-5", Domain::Shelf),
    ]
    .into_iter()
    .map(|(name, domain)| Node {
        name,
        domain,
        job: None,
    })
    .collect()
}

/// A group of workers that have to run at the same time, one GPU each.
#[derive(Clone, Copy)]
struct Job {
    name: &'static str,
    workers: usize,
}

/// Returns the indexes of the machines whose GPU nobody holds, in name order.
fn free(fleet: &[Node]) -> Vec<usize> {
    (0..fleet.len())
        .filter(|&i| fleet[i].job.is_none())
        .collect()
}

/// Joins the names of some machines with spaces.
fn names(fleet: &[Node], nodes: &[usize]) -> String {
    nodes
        .iter()
        .map(|&i| fleet[i].name)
        .collect::<Vec<_>>()
        .join(" ")
}

/// Gives each job one GPU in turn until none are left, the way a scheduler
/// that sees only single pods would, and reports what each job ended up
/// holding.
fn pod_at_a_time(mut fleet: Vec<Node>, jobs: &[Job]) {
    let held = |fleet: &[Node], job: &Job| -> Vec<usize> {
        (0..fleet.len())
            .filter(|&i| fleet[i].job == Some(job.name))
            .collect()
    };
    let mut progress = true;

    while progress {
        progress = false;

        for job in jobs {
            if held(&fleet, job).len() < job.workers
                && let Some(&first) = free(&fleet).first()
            {
                fleet[first].job = Some(job.name);
                progress = true;
            }
        }
    }

    let mut running = 0;

    for job in jobs {
        let mine = held(&fleet, job);
        println!(
            "  {:<8} holds {} of {} GPUs: {}",
            job.name,
            mine.len(),
            job.workers,
            names(&fleet, &mine)
        );

        if mine.len() == job.workers {
            running += 1;
        }
    }

    let taken = fleet.len() - free(&fleet).len();
    println!(
        "  {taken} of {} GPUs are held and {running} jobs are running",
        fleet.len()
    );
}

/// Chooses machines for a job from the free ones, or returns `None` when it
/// cannot place every worker.
type Placer = fn(fleet: &[Node], open: &[usize], workers: usize) -> Option<Vec<usize>>;

/// Takes the first free machines in name order.
fn first_free(_fleet: &[Node], open: &[usize], workers: usize) -> Option<Vec<usize>> {
    open.get(..workers).map(<[usize]>::to_vec)
}

/// Looks for one domain that can hold the whole job and, among those, takes
/// the one with the fewest free machines, so that large domains stay whole
/// for large jobs. A job that fits in no single domain is spread over the
/// first free machines.
fn by_topology(fleet: &[Node], open: &[usize], workers: usize) -> Option<Vec<usize>> {
    let in_domain = |domain: Domain| -> Vec<usize> {
        open.iter()
            .copied()
            .filter(|&i| fleet[i].domain == domain)
            .collect()
    };
    [Domain::Bench, Domain::Desk, Domain::Shelf]
        .into_iter()
        .map(in_domain)
        .filter(|members| members.len() >= workers)
        .min_by_key(Vec::len)
        .map(|members| members[..workers].to_vec())
        .or_else(|| first_free(fleet, open, workers))
}

/// The size of the gradient every worker contributes per step, in gigabytes.
const GRADIENT_GB: f64 = 14.0;

/// Returns the link a job's all-reduce is limited by: the domain's link when
/// every machine is in one domain, and the house network otherwise.
fn slowest(fleet: &[Node], nodes: &[usize]) -> Link {
    let first = fleet[nodes[0]].domain;
    if nodes.iter().all(|&i| fleet[i].domain == first) {
        first.link()
    } else {
        HOUSE_NETWORK
    }
}

/// Returns the time one ring all-reduce takes: each worker sends 2(n-1)/n
/// gradients over the slowest link.
fn exchange_seconds(fleet: &[Node], nodes: &[usize]) -> f64 {
    let n = nodes.len() as f64;

    2.0 * (n - 1.0) / n * GRADIENT_GB / slowest(fleet, nodes).gbps
}

/// Starts a job only when every one of its workers has a machine. It reports
/// the placement and what the job's all-reduce will cost, and returns false
/// when the job has to wait.
fn admit(fleet: &mut [Node], job: Job, place: Placer) -> bool {
    let open = free(fleet);
    let Some(nodes) = place(fleet, &open, job.workers) else {
        println!(
            "  {:<8} waits: it needs {} GPUs and {} are free",
            job.name,
            job.workers,
            open.len()
        );

        return false;
    };

    for &i in &nodes {
        fleet[i].job = Some(job.name);
    }

    println!("  {:<8} runs on {}", job.name, names(fleet, &nodes));

    if nodes.len() > 1 {
        let link = slowest(fleet, &nodes);
        println!(
            "           over {}: {:.1} s for each all-reduce",
            link.label,
            exchange_seconds(fleet, &nodes)
        );
    }

    true
}

/// Frees every GPU a job holds.
fn release(fleet: &mut [Node], job: &str) {
    for node in fleet.iter_mut().filter(|n| n.job == Some(job)) {
        node.job = None;
    }

    println!("  {job:<8} finishes");
}

fn main() {
    let two = [
        Job {
            name: "train-a",
            workers: 5,
        },
        Job {
            name: "train-b",
            workers: 5,
        },
    ];

    println!("one pod at a time, two jobs of five workers:");
    pod_at_a_time(new_fleet(), &two);

    println!("whole jobs at a time:");
    let mut fleet = new_fleet();
    admit(&mut fleet, two[0], by_topology);
    admit(&mut fleet, two[1], by_topology);
    release(&mut fleet, "train-a");
    admit(&mut fleet, two[1], by_topology);

    let three = [
        Job {
            name: "pair",
            workers: 2,
        },
        Job {
            name: "quad",
            workers: 4,
        },
        Job {
            name: "single",
            workers: 1,
        },
    ];
    println!("three jobs, on the first free machines:");
    let mut fleet = new_fleet();

    for job in three {
        admit(&mut fleet, job, first_free);
    }

    println!("the same three jobs, placed by topology:");
    let mut fleet = new_fleet();

    for job in three {
        admit(&mut fleet, job, by_topology);
    }
}

References

• Kubernetes, Scheduling framework: https://kubernetes.io/docs/concepts/scheduling-eviction/scheduling-framework/
• Kueue documentation: https://kueue.sigs.k8s.io/docs/overview/
• Kueue, Topology aware scheduling: https://kueue.sigs.k8s.io/docs/concepts/topology_aware_scheduling/
• Kubernetes scheduler plugins, Coscheduling: https://github.com/kubernetes-sigs/scheduler-plugins/tree/master/pkg/coscheduling
• Kubernetes, Jobs, suspending a Job: https://kubernetes.io/docs/concepts/workloads/controllers/job/#suspending-a-job

9. Counting across replicas

A GPU platform counts all day. It counts the GPU time each team’s jobs have used, so that the time can be billed or shared out fairly, and it counts the GPUs each team has in use right now, so that a quota can be enforced. Inside one cluster, counting is not hard: the API server is the single place every write passes through. A platform that spans several sites has several places where things happen, and they cannot always reach each other.

In this chapter, we will count across three replicas of a usage service, one beside each group of machines in our fleet. We will first see why the obvious ways of keeping the replicas in step lose counts or count twice. Then we will write counter, in Go and in Rust, which builds a count that survives a network partition, extends it to a count that goes down as well as up, and finally shows the one thing such a count cannot do, which is to enforce a limit, and what does that job in its place.

One authority, or several

Suppose each group of machines in our fleet has a small service beside it that records usage: one on the shelf with the Thors, one on the desk with the Sparks, one on the bench with the Orin NX. Jobs report to the nearest one. Each of these is a replica of the usage service, and all three are meant to hold the same totals.

The first design gives the count a single home. One replica, or a database that all three write to, holds the number, and every report becomes a write to it. etcd is a database of this kind: it stores each write on a majority of its members before it acknowledges the write, so that all members agree on the order of writes. The count is always exact.

The cost appears when the network between the sites fails. Our three groups are joined only by the house network. If it goes down, the bench is alone. A replica that cannot reach a majority has two choices, and both are bad for a usage counter:

• It can refuse reports until the network returns. The jobs near it keep running, and their usage for that period is not recorded.
• It can accept reports into a count of its own. Now there are two counts, and the design gave us no rule for combining them afterwards.

A split like this is called a network partition, and the choice is a general one. During a partition, a replica that is cut off can stay consistent, by refusing work it cannot coordinate, or stay available, by doing the work alone. It cannot do both. For a count of usage we want availability: a report should be accepted wherever it arrives. So we need the thing the second choice lacked, a rule for combining counts that were kept apart.

Two rules that do not work

Let each replica keep one number, its total. After some jobs have reported, the shelf holds 300 GPU-seconds, the desk 120, and the bench 40. The true total is 460. Now the replicas exchange numbers and each has to decide what its total should be.

Keep the larger. Each replica replaces its number with the larger of its own and the one it received. All three end at 300. The 160 GPU-seconds recorded at the desk and the bench are gone, because the rule treats two different counts as two versions of the same count.

Add the other’s number. Each replica adds what it receives to its own. After the shelf and the desk exchange once, both say 420, which is right for those two. But exchanges repeat, and messages are sometimes delivered twice. When the desk’s 120 arrives at the shelf a second time, the shelf says 540. A rule based on adding counts every report again each time the replicas talk.

Both rules fail for the same reason. A single number does not say whose reports it contains, so a replica that receives one cannot tell which part of it is news.

One entry per replica

The fix is to keep that information. Each replica holds a row with one entry for every replica, itself included. Two rules govern the row:

• A replica records a report by adding to its own entry. It leaves every other entry as it received it.
• When a replica receives another’s row, it goes through the entries one by one and keeps the larger of its own value and the received one.

The count is the sum of the row.

“Keep the larger” failed for totals and works here, and the difference is what the two numbers being compared mean. An entry is written by exactly one replica and only ever increases. So when two values of the same entry differ, they really are an older and a newer version of the same thing, and the larger is the newer.

Let us run the earlier example. After the first reports, written as rows of (shelf, desk, bench):

• the shelf holds (300, 0, 0),
• the desk holds (0, 120, 0),
• the bench holds (0, 0, 40).

The shelf receives the desk’s row and takes the larger in each position: (300, 120, 0), which sums to 420. If the same row arrives again, nothing changes, because 120 is not larger than 120. When the bench’s row arrives, the shelf holds (300, 120, 40), and the sum is 460.

This structure is called a grow-only counter, or G-counter. Figure 9.1 follows three of them through a partition.

What the merge guarantees

The merge has three properties, and each one removes a requirement from the network between the replicas.

• Merging the same row twice gives the same result as merging it once. So a message may be delivered twice, and replicas may exchange rows as often as they like.
• Merging row A and then row B gives the same result as B and then A. So messages may arrive in any order.
• Merging A with the result of merging B and C gives the same result as merging A and B first and then C. So a row may travel through other replicas on its way, picking up their entries, and nothing is lost or doubled.

Together they mean that any two replicas that have received the same reports, by any route, hold the same row. Nothing has to be coordinated, and a replica that was cut off for an hour catches up with one exchange. Data structures with a merge of this kind are called conflict-free replicated data types (CRDTs).

Counting down

A count of GPUs in use has to go down when a job ends, and “keep the larger” cannot represent a decrease: a replica that lowered its entry would have the old, larger value merged straight back in.

The way round this is to use two grow-only counters. One counts every increase and the other counts every decrease, and the value is the first minus the second. Both halves only grow, so both merge safely. This is a PN-counter, for positive and negative.

What a merged count cannot do

A quota is a limit: this team may use at most six GPUs at once. Enforcing it means refusing a job when the count is at the limit, and here the merged count lets us down.

Suppose the shelf and the desk both know that four of the six GPUs are in use, and then lose contact. A two-GPU job arrives at each. Each replica looks at its own count, sees four in use, and admits the job, correctly by everything it knows. When contact returns and the counts merge, eight GPUs are in use against a quota of six.

No cleverer merge fixes this. A limit is a statement about the sum of what all replicas have done, and a replica that is cut off cannot know what the others are doing at this moment. What we can do is divide the limit beforehand, while the replicas can still talk. Each replica is given a share of the quota and admits jobs only out of its own share. The shares add up to the quota, so whatever the replicas do while apart, the total stays within it. When they are in contact, a replica with unused share can hand some to one that has run out.

The price is that a job can be refused while capacity exists. If the shelf has one GPU of its share left and the desk has one, a two-GPU job is refused at both, even though two GPUs are free in total, until the replicas can talk and move a share across. The partition forces a choice here too, and for a limit we chose to refuse.

Building counter

The grow-only counter

A GCounter is an array with one entry per replica. Add touches one entry, Merge takes the larger value of each, and Value sums them.

Listing 9.1 The grow-only counter, in Go. code/go/counter/main.go

// GCounter is a count that only grows. It keeps one entry per replica, and
// each replica adds only to its own entry.
type GCounter [replicas]int

// Add records n more at the given replica.
func (c *GCounter) Add(replica, n int) { c[replica] += n }

// Merge takes, entry by entry, the larger of this counter's value and the
// other's. Merging is safe to repeat and to do in any order.
func (c *GCounter) Merge(other GCounter) {
	for i := range c {
		c[i] = max(c[i], other[i])
	}
}

// Value is the count: the sum of every replica's entry.
func (c GCounter) Value() int {
	total := 0

	for _, n := range c {
		total += n
	}

	return total
}

Listing 9.2 The grow-only counter, in Rust. code/rust/counter/src/main.rs

/// A count that only grows. It keeps one entry per replica, and each replica
/// adds only to its own entry.
#[derive(Clone, Copy, Default)]
struct GCounter([i64; REPLICAS]);

impl GCounter {
    /// Records n more at the given replica.
    fn add(&mut self, replica: usize, n: i64) {
        self.0[replica] += n;
    }

    /// Takes, entry by entry, the larger of this counter's value and the
    /// other's. Merging is safe to repeat and to do in any order.
    fn merge(&mut self, other: GCounter) {
        for (mine, theirs) in self.0.iter_mut().zip(other.0) {
            *mine = (*mine).max(theirs);
        }
    }

    /// The count: the sum of every replica's entry.
    fn value(self) -> i64 {
        self.0.iter().sum()
    }
}

The counter is a fixed-size array in both languages, which makes it a plain value: assigning it copies it, and “send my row to another replica” in the program is passing the array to Merge.

The counter that goes down

PNCounter holds two grow-only counters. A positive change goes into Up, a negative one into Down with its sign removed.

Listing 9.3 The up-and-down counter, in Go. code/go/counter/main.go

// PNCounter is a count that goes up and down, built from two counters that
// only grow: one for the increments and one for the decrements.
type PNCounter struct {
	Up, Down GCounter
}

// Add records a change of n, which may be negative, at the given replica.
func (c *PNCounter) Add(replica, n int) {
	if n >= 0 {
		c.Up.Add(replica, n)
	} else {
		c.Down.Add(replica, -n)
	}
}

// Merge merges both halves.
func (c *PNCounter) Merge(other PNCounter) {
	c.Up.Merge(other.Up)
	c.Down.Merge(other.Down)
}

// Value is the increments minus the decrements.
func (c PNCounter) Value() int { return c.Up.Value() - c.Down.Value() }

Listing 9.4 The up-and-down counter, in Rust. code/rust/counter/src/main.rs

/// A count that goes up and down, built from two counters that only grow: one
/// for the increments and one for the decrements.
#[derive(Clone, Copy, Default)]
struct PnCounter {
    up: GCounter,
    down: GCounter,
}

impl PnCounter {
    /// Records a change of n, which may be negative, at the given replica.
    fn add(&mut self, replica: usize, n: i64) {
        if n >= 0 {
            self.up.add(replica, n);
        } else {
            self.down.add(replica, -n);
        }
    }

    /// Merges both halves.
    fn merge(&mut self, other: PnCounter) {
        self.up.merge(other.up);
        self.down.merge(other.down);
    }

    /// The increments minus the decrements.
    fn value(self) -> i64 {
        self.up.value() - self.down.value()
    }
}

One number per replica

The first experiment is the rule that fails: one total per replica, synchronised by keeping the larger.

Listing 9.5 Keeping the larger of two totals, in Go. code/go/counter/main.go

// oneNumber counts with a single number per replica and synchronises by
// keeping the larger of two numbers.
func oneNumber() {
	totals := [replicas]int{}
	adds := [replicas]int{300, 120, 40}
	sum := 0

	for r, n := range adds {
		totals[r] += n
		sum += n
		fmt.Printf("  %-5s records %d GPU-seconds\n", names[r], n)
	}

	largest := max(totals[shelf], totals[desk], totals[bench])
	fmt.Printf("  after synchronising, every replica says %d; the jobs used %d\n", largest, sum)
}

Listing 9.6 Keeping the larger of two totals, in Rust. code/rust/counter/src/main.rs

/// Counts with a single number per replica and synchronises by keeping the
/// larger of two numbers.
fn one_number() {
    let adds = [300, 120, 40];

    for (name, n) in NAMES.iter().zip(adds) {
        println!("  {name:<5} records {n} GPU-seconds");
    }

    let largest = adds.iter().max().expect("three replicas");
    let sum: i64 = adds.iter().sum();
    println!("  after synchronising, every replica says {largest}; the jobs used {sum}");
}

Through a partition

The second experiment gives each replica a GCounter. The bench is cut off, the shelf and the desk merge with each other, and reports keep arriving on both sides. After the partition the rows are merged in an order chosen to be untidy, with one merge repeated.

Listing 9.7 Grow-only counters, a partition, and an untidy recovery, in Go. code/go/counter/main.go

// growOnly counts the same usage with a GCounter at each replica, through a
// partition that cuts the bench off from the other two.
func growOnly() {
	var at [replicas]GCounter

	for r, n := range [replicas]int{300, 120, 40} {
		at[r].Add(r, n)
	}

	fmt.Println("  the bench is cut off; shelf and desk exchange counters")
	at[shelf].Merge(at[desk])
	at[desk].Merge(at[shelf])
	at[shelf].Add(shelf, 60)
	at[bench].Add(bench, 25)

	for r := range at {
		show(r, at[r])
	}

	fmt.Println("  the network heals; counters are exchanged in no particular order, some twice")
	at[desk].Merge(at[bench])
	at[shelf].Merge(at[desk])
	at[bench].Merge(at[shelf])
	at[desk].Merge(at[shelf])
	at[desk].Merge(at[shelf])

	for r := range at {
		show(r, at[r])
	}
}

Listing 9.8 Grow-only counters, a partition, and an untidy recovery, in Rust. code/rust/counter/src/main.rs

/// Counts the same usage with a `GCounter` at each replica, through a
/// partition that cuts the bench off from the other two.
fn grow_only() {
    let mut at = [GCounter::default(); REPLICAS];

    for (r, n) in [300, 120, 40].into_iter().enumerate() {
        at[r].add(r, n);
    }

    println!("  the bench is cut off; shelf and desk exchange counters");
    at[SHELF].merge(at[DESK]);
    at[DESK].merge(at[SHELF]);
    at[SHELF].add(SHELF, 60);
    at[BENCH].add(BENCH, 25);

    for (r, c) in at.into_iter().enumerate() {
        show(r, c);
    }

    println!("  the network heals; counters are exchanged in no particular order, some twice");
    at[DESK].merge(at[BENCH]);
    at[SHELF].merge(at[DESK]);
    at[BENCH].merge(at[SHELF]);
    at[DESK].merge(at[SHELF]);
    at[DESK].merge(at[SHELF]);

    for (r, c) in at.into_iter().enumerate() {
        show(r, c);
    }
}

GPUs in use

The third experiment counts jobs starting and ending. One job’s end is reported to a different replica from the one that recorded its start, which is ordinary: a job’s machine and the replica that hears of its end need not be in the same place.

Listing 9.9 Starts and ends recorded at different replicas, in Go. code/go/counter/main.go

// upAndDown counts GPUs in use. A job's start and its end may be recorded at
// different replicas.
func upAndDown() {
	var at [replicas]PNCounter
	at[shelf].Add(shelf, 3)  // three jobs start on the Thors
	at[desk].Add(desk, 2)    // two start on the Sparks
	at[shelf].Add(shelf, -1) // one Thor job ends
	at[bench].Add(bench, -1) // the bench replica hears that a Spark job ended

	for r := range at {
		for other := range at {
			at[r].Merge(at[other])
		}
	}

	c := at[bench]
	fmt.Printf("  started %v, ended %v: %d GPUs in use\n", [replicas]int(c.Up), [replicas]int(c.Down), c.Value())
}

Listing 9.10 Starts and ends recorded at different replicas, in Rust. code/rust/counter/src/main.rs

/// Counts GPUs in use. A job's start and its end may be recorded at different
/// replicas.
fn up_and_down() {
    let mut at = [PnCounter::default(); REPLICAS];
    at[SHELF].add(SHELF, 3); // three jobs start on the Thors
    at[DESK].add(DESK, 2); // two start on the Sparks
    at[SHELF].add(SHELF, -1); // one Thor job ends
    at[BENCH].add(BENCH, -1); // the bench replica hears that a Spark job ended

    for r in 0..REPLICAS {
        for other in 0..REPLICAS {
            let theirs = at[other];
            at[r].merge(theirs);
        }
    }

    let c = at[BENCH];
    println!(
        "  started {}, ended {}: {} GPUs in use",
        row(c.up.0),
        row(c.down.0),
        c.value()
    );
}

A quota, two ways

The last two experiments enforce a quota of six GPUs. sharedLimit lets each replica check its own merged count. allowances gives each replica a share.

Listing 9.11 A limit checked against a local count, and a limit divided into shares, in Go. code/go/counter/main.go

// quota is the number of GPUs one team may use at once.
const quota = 6

// sharedLimit lets each replica admit a job whenever its own view of the
// count leaves room under the quota.
func sharedLimit() {
	var at [replicas]PNCounter
	at[shelf].Add(shelf, 3)
	at[desk].Add(desk, 1)
	at[shelf].Merge(at[desk])
	at[desk].Merge(at[shelf])
	fmt.Printf("  both replicas see %d of %d in use; then they lose contact\n", at[shelf].Value(), quota)

	for _, r := range []int{shelf, desk} {
		if at[r].Value()+2 <= quota {
			at[r].Add(r, 2)
			fmt.Printf("  %-5s admits a 2-GPU job: by its count, %d of %d are in use\n", names[r], at[r].Value(), quota)
		}
	}

	at[shelf].Merge(at[desk])
	fmt.Printf("  after they merge: %d of %d in use\n", at[shelf].Value(), quota)
}

// allowances gives each replica a share of the quota. A replica admits a job
// only from its own share, and shares move between replicas when they can
// reach each other.
func allowances() {
	share := [replicas]int{4, 2, 0}
	used := [replicas]int{3, 1, 0}
	fmt.Printf("  shares %v, in use %v; then the replicas lose contact\n", share, used)
	admit := func(r, gpus int) {
		if left := share[r] - used[r]; left < gpus {
			fmt.Printf("  %-5s refuses a %d-GPU job: its share has %d left\n", names[r], gpus, left)

			return
		}

		used[r] += gpus
		fmt.Printf("  %-5s admits a %d-GPU job\n", names[r], gpus)
	}
	admit(shelf, 2)
	admit(desk, 2)
	fmt.Println("  contact returns; desk gives its unused GPU to shelf")
	share[desk]--
	share[shelf]++
	admit(shelf, 2)
	fmt.Printf("  shares %v, in use %v: %d of %d\n", share, used, used[shelf]+used[desk]+used[bench], quota)
}

Listing 9.12 A limit checked against a local count, and a limit divided into shares, in Rust. code/rust/counter/src/main.rs

/// The number of GPUs one team may use at once.
const QUOTA: i64 = 6;

/// Lets each replica admit a job whenever its own view of the count leaves
/// room under the quota.
fn shared_limit() {
    let mut at = [PnCounter::default(); REPLICAS];
    at[SHELF].add(SHELF, 3);
    at[DESK].add(DESK, 1);
    let (shelf, desk) = (at[SHELF], at[DESK]);
    at[SHELF].merge(desk);
    at[DESK].merge(shelf);
    println!(
        "  both replicas see {} of {QUOTA} in use; then they lose contact",
        at[SHELF].value()
    );

    for r in [SHELF, DESK] {
        if at[r].value() + 2 <= QUOTA {
            at[r].add(r, 2);
            println!(
                "  {:<5} admits a 2-GPU job: by its count, {} of {QUOTA} are in use",
                NAMES[r],
                at[r].value()
            );
        }
    }

    let desk = at[DESK];
    at[SHELF].merge(desk);
    println!(
        "  after they merge: {} of {QUOTA} in use",
        at[SHELF].value()
    );
}

/// Gives each replica a share of the quota. A replica admits a job only from
/// its own share, and shares move between replicas when they can reach each
/// other.
fn allowances() {
    let mut share = [4, 2, 0];
    let mut used = [3, 1, 0];
    println!(
        "  shares {}, in use {}; then the replicas lose contact",
        row(share),
        row(used)
    );
    let admit = |share: &[i64; REPLICAS], used: &mut [i64; REPLICAS], r: usize, gpus: i64| {
        let left = share[r] - used[r];

        if left < gpus {
            println!(
                "  {:<5} refuses a {gpus}-GPU job: its share has {left} left",
                NAMES[r]
            );

            return;
        }

        used[r] += gpus;
        println!("  {:<5} admits a {gpus}-GPU job", NAMES[r]);
    };
    admit(&share, &mut used, SHELF, 2);
    admit(&share, &mut used, DESK, 2);
    println!("  contact returns; desk gives its unused GPU to shelf");
    share[DESK] -= 1;
    share[SHELF] += 1;
    admit(&share, &mut used, SHELF, 2);
    println!(
        "  shares {}, in use {}: {} of {QUOTA}",
        row(share),
        row(used),
        used.iter().sum::<i64>()
    );
}

The run

$ go run ./counter
one number per replica:
  shelf records 300 GPU-seconds
  desk  records 120 GPU-seconds
  bench records 40 GPU-seconds
  after synchronising, every replica says 300; the jobs used 460
one entry per replica, merged by taking the larger:
  the bench is cut off; shelf and desk exchange counters
  shelf holds [360 120 0] = 480
  desk  holds [300 120 0] = 420
  bench holds [0 0 65] = 65
  the network heals; counters are exchanged in no particular order, some twice
  shelf holds [360 120 65] = 545
  desk  holds [360 120 65] = 545
  bench holds [360 120 65] = 545
a count that goes down as well:
  started [3 2 0], ended [1 0 1]: 3 GPUs in use
a quota checked against each replica's own count:
  both replicas see 4 of 6 in use; then they lose contact
  shelf admits a 2-GPU job: by its count, 6 of 6 are in use
  desk  admits a 2-GPU job: by its count, 6 of 6 are in use
  after they merge: 8 of 6 in use
a quota divided into shares:
  shares [4 2 0], in use [3 1 0]; then the replicas lose contact
  shelf refuses a 2-GPU job: its share has 1 left
  desk  refuses a 2-GPU job: its share has 1 left
  contact returns; desk gives its unused GPU to shelf
  shelf admits a 2-GPU job
  shares [5 1 0], in use [5 1 0]: 6 of 6

With one number per replica, every replica says 300 after synchronising, and the jobs used 460.

With one entry per replica, the two sides of the partition disagree while it lasts: the shelf counts 480, the desk 420, and the bench 65. That is expected, since each is missing reports the others hold. After the untidy recovery all three hold the row [360 120 65] and the count 545, which is every report added once: 300, 120, 40, 60, and 25.

Animation 9.1 plays the second experiment, one merge at a time.

The count of GPUs in use comes to 3: five jobs started, and two ended, one of them recorded at the bench.

In the fourth experiment both replicas admit a job at “6 of 6”, and the merged count is 8 of 6. In the fifth, with shares of 4 and 2 and usage of 3 and 1, both replicas refuse the two-GPU job while apart. When contact returns, the desk gives its unused GPU to the shelf, the shelf admits the job, and the total in use is exactly 6. Figure 9.2 shows the shares before and after.

Where Kubernetes stands

Inside one cluster, Kubernetes takes the first design. A ResourceQuota object limits what a namespace may use, and the API server checks it when a pod is created. The check is exact because the API server is the one authority and etcd puts every write in a single order: the resource version of chapter 1 is that order made visible. A cluster whose control plane is cut off from part of itself stops accepting changes from that part, which is the consistent choice.

The techniques of this chapter apply above that level: usage that is gathered from many clusters into one bill, a team’s quota that spans clusters in several regions, a rate limit enforced by many copies of an API gateway. In each of these, no single API server sees every event, and the count has to be assembled from replicas that are sometimes apart.

The complete program

Listing 9.13 The complete program in Go. code/go/counter/main.go

// Command counter shows how replicas of a service that cannot always reach
// each other can keep one count.
//
// Three replicas of a usage service run beside the three groups of machines
// in our fleet: the shelf, the desk, and the bench. Each records the GPU time
// that jobs near it report. The program counts in four ways: with one number
// per replica, with a grow-only counter, with a counter that also goes down,
// and with a quota divided into allowances.
//
// Run it from code/go with:
//
//	go run ./counter
package main

import "fmt"

// The replicas, by index.
const (
	shelf = iota
	desk
	bench
	replicas
)

// names lists the replicas in index order.
var names = [replicas]string{"shelf", "desk", "bench"}

// GCounter is a count that only grows. It keeps one entry per replica, and
// each replica adds only to its own entry.
type GCounter [replicas]int

// Add records n more at the given replica.
func (c *GCounter) Add(replica, n int) { c[replica] += n }

// Merge takes, entry by entry, the larger of this counter's value and the
// other's. Merging is safe to repeat and to do in any order.
func (c *GCounter) Merge(other GCounter) {
	for i := range c {
		c[i] = max(c[i], other[i])
	}
}

// Value is the count: the sum of every replica's entry.
func (c GCounter) Value() int {
	total := 0

	for _, n := range c {
		total += n
	}

	return total
}


// PNCounter is a count that goes up and down, built from two counters that
// only grow: one for the increments and one for the decrements.
type PNCounter struct {
	Up, Down GCounter
}

// Add records a change of n, which may be negative, at the given replica.
func (c *PNCounter) Add(replica, n int) {
	if n >= 0 {
		c.Up.Add(replica, n)
	} else {
		c.Down.Add(replica, -n)
	}
}

// Merge merges both halves.
func (c *PNCounter) Merge(other PNCounter) {
	c.Up.Merge(other.Up)
	c.Down.Merge(other.Down)
}

// Value is the increments minus the decrements.
func (c PNCounter) Value() int { return c.Up.Value() - c.Down.Value() }


// oneNumber counts with a single number per replica and synchronises by
// keeping the larger of two numbers.
func oneNumber() {
	totals := [replicas]int{}
	adds := [replicas]int{300, 120, 40}
	sum := 0

	for r, n := range adds {
		totals[r] += n
		sum += n
		fmt.Printf("  %-5s records %d GPU-seconds\n", names[r], n)
	}

	largest := max(totals[shelf], totals[desk], totals[bench])
	fmt.Printf("  after synchronising, every replica says %d; the jobs used %d\n", largest, sum)
}


// show prints one replica's counter and its value.
func show(replica int, c GCounter) {
	fmt.Printf("  %-5s holds %v = %d\n", names[replica], [replicas]int(c), c.Value())
}

// growOnly counts the same usage with a GCounter at each replica, through a
// partition that cuts the bench off from the other two.
func growOnly() {
	var at [replicas]GCounter

	for r, n := range [replicas]int{300, 120, 40} {
		at[r].Add(r, n)
	}

	fmt.Println("  the bench is cut off; shelf and desk exchange counters")
	at[shelf].Merge(at[desk])
	at[desk].Merge(at[shelf])
	at[shelf].Add(shelf, 60)
	at[bench].Add(bench, 25)

	for r := range at {
		show(r, at[r])
	}

	fmt.Println("  the network heals; counters are exchanged in no particular order, some twice")
	at[desk].Merge(at[bench])
	at[shelf].Merge(at[desk])
	at[bench].Merge(at[shelf])
	at[desk].Merge(at[shelf])
	at[desk].Merge(at[shelf])

	for r := range at {
		show(r, at[r])
	}
}


// upAndDown counts GPUs in use. A job's start and its end may be recorded at
// different replicas.
func upAndDown() {
	var at [replicas]PNCounter
	at[shelf].Add(shelf, 3)  // three jobs start on the Thors
	at[desk].Add(desk, 2)    // two start on the Sparks
	at[shelf].Add(shelf, -1) // one Thor job ends
	at[bench].Add(bench, -1) // the bench replica hears that a Spark job ended

	for r := range at {
		for other := range at {
			at[r].Merge(at[other])
		}
	}

	c := at[bench]
	fmt.Printf("  started %v, ended %v: %d GPUs in use\n", [replicas]int(c.Up), [replicas]int(c.Down), c.Value())
}


// quota is the number of GPUs one team may use at once.
const quota = 6

// sharedLimit lets each replica admit a job whenever its own view of the
// count leaves room under the quota.
func sharedLimit() {
	var at [replicas]PNCounter
	at[shelf].Add(shelf, 3)
	at[desk].Add(desk, 1)
	at[shelf].Merge(at[desk])
	at[desk].Merge(at[shelf])
	fmt.Printf("  both replicas see %d of %d in use; then they lose contact\n", at[shelf].Value(), quota)

	for _, r := range []int{shelf, desk} {
		if at[r].Value()+2 <= quota {
			at[r].Add(r, 2)
			fmt.Printf("  %-5s admits a 2-GPU job: by its count, %d of %d are in use\n", names[r], at[r].Value(), quota)
		}
	}

	at[shelf].Merge(at[desk])
	fmt.Printf("  after they merge: %d of %d in use\n", at[shelf].Value(), quota)
}

// allowances gives each replica a share of the quota. A replica admits a job
// only from its own share, and shares move between replicas when they can
// reach each other.
func allowances() {
	share := [replicas]int{4, 2, 0}
	used := [replicas]int{3, 1, 0}
	fmt.Printf("  shares %v, in use %v; then the replicas lose contact\n", share, used)
	admit := func(r, gpus int) {
		if left := share[r] - used[r]; left < gpus {
			fmt.Printf("  %-5s refuses a %d-GPU job: its share has %d left\n", names[r], gpus, left)

			return
		}

		used[r] += gpus
		fmt.Printf("  %-5s admits a %d-GPU job\n", names[r], gpus)
	}
	admit(shelf, 2)
	admit(desk, 2)
	fmt.Println("  contact returns; desk gives its unused GPU to shelf")
	share[desk]--
	share[shelf]++
	admit(shelf, 2)
	fmt.Printf("  shares %v, in use %v: %d of %d\n", share, used, used[shelf]+used[desk]+used[bench], quota)
}


func main() {
	fmt.Println("one number per replica:")
	oneNumber()
	fmt.Println("one entry per replica, merged by taking the larger:")
	growOnly()
	fmt.Println("a count that goes down as well:")
	upAndDown()
	fmt.Println("a quota checked against each replica's own count:")
	sharedLimit()
	fmt.Println("a quota divided into shares:")
	allowances()
}

Listing 9.14 The complete program in Rust. code/rust/counter/src/main.rs

//! Shows how replicas of a service that cannot always reach each other can
//! keep one count.
//!
//! Three replicas of a usage service run beside the three groups of machines
//! in our fleet: the shelf, the desk, and the bench. Each records the GPU time
//! that jobs near it report. The program counts in four ways: with one number
//! per replica, with a grow-only counter, with a counter that also goes down,
//! and with a quota divided into allowances.
//!
//! Run it from code/rust with:
//!
//! ```text
//! cargo run -p counter
//! ```

/// The number of replicas.
const REPLICAS: usize = 3;
/// The replicas, by index.
const SHELF: usize = 0;
const DESK: usize = 1;
const BENCH: usize = 2;
/// The replicas' names, in index order.
const NAMES: [&str; REPLICAS] = ["shelf", "desk", "bench"];

/// Formats a row of numbers as Go prints an array: `[1 2 3]`.
fn row(values: [i64; REPLICAS]) -> String {
    format!("[{}]", values.map(|v| v.to_string()).join(" "))
}

/// A count that only grows. It keeps one entry per replica, and each replica
/// adds only to its own entry.
#[derive(Clone, Copy, Default)]
struct GCounter([i64; REPLICAS]);

impl GCounter {
    /// Records n more at the given replica.
    fn add(&mut self, replica: usize, n: i64) {
        self.0[replica] += n;
    }

    /// Takes, entry by entry, the larger of this counter's value and the
    /// other's. Merging is safe to repeat and to do in any order.
    fn merge(&mut self, other: GCounter) {
        for (mine, theirs) in self.0.iter_mut().zip(other.0) {
            *mine = (*mine).max(theirs);
        }
    }

    /// The count: the sum of every replica's entry.
    fn value(self) -> i64 {
        self.0.iter().sum()
    }
}

/// A count that goes up and down, built from two counters that only grow: one
/// for the increments and one for the decrements.
#[derive(Clone, Copy, Default)]
struct PnCounter {
    up: GCounter,
    down: GCounter,
}

impl PnCounter {
    /// Records a change of n, which may be negative, at the given replica.
    fn add(&mut self, replica: usize, n: i64) {
        if n >= 0 {
            self.up.add(replica, n);
        } else {
            self.down.add(replica, -n);
        }
    }

    /// Merges both halves.
    fn merge(&mut self, other: PnCounter) {
        self.up.merge(other.up);
        self.down.merge(other.down);
    }

    /// The increments minus the decrements.
    fn value(self) -> i64 {
        self.up.value() - self.down.value()
    }
}

/// Counts with a single number per replica and synchronises by keeping the
/// larger of two numbers.
fn one_number() {
    let adds = [300, 120, 40];

    for (name, n) in NAMES.iter().zip(adds) {
        println!("  {name:<5} records {n} GPU-seconds");
    }

    let largest = adds.iter().max().expect("three replicas");
    let sum: i64 = adds.iter().sum();
    println!("  after synchronising, every replica says {largest}; the jobs used {sum}");
}

/// Prints one replica's counter and its value.
fn show(replica: usize, c: GCounter) {
    println!("  {:<5} holds {} = {}", NAMES[replica], row(c.0), c.value());
}

/// Counts the same usage with a `GCounter` at each replica, through a
/// partition that cuts the bench off from the other two.
fn grow_only() {
    let mut at = [GCounter::default(); REPLICAS];

    for (r, n) in [300, 120, 40].into_iter().enumerate() {
        at[r].add(r, n);
    }

    println!("  the bench is cut off; shelf and desk exchange counters");
    at[SHELF].merge(at[DESK]);
    at[DESK].merge(at[SHELF]);
    at[SHELF].add(SHELF, 60);
    at[BENCH].add(BENCH, 25);

    for (r, c) in at.into_iter().enumerate() {
        show(r, c);
    }

    println!("  the network heals; counters are exchanged in no particular order, some twice");
    at[DESK].merge(at[BENCH]);
    at[SHELF].merge(at[DESK]);
    at[BENCH].merge(at[SHELF]);
    at[DESK].merge(at[SHELF]);
    at[DESK].merge(at[SHELF]);

    for (r, c) in at.into_iter().enumerate() {
        show(r, c);
    }
}

/// Counts GPUs in use. A job's start and its end may be recorded at different
/// replicas.
fn up_and_down() {
    let mut at = [PnCounter::default(); REPLICAS];
    at[SHELF].add(SHELF, 3); // three jobs start on the Thors
    at[DESK].add(DESK, 2); // two start on the Sparks
    at[SHELF].add(SHELF, -1); // one Thor job ends
    at[BENCH].add(BENCH, -1); // the bench replica hears that a Spark job ended

    for r in 0..REPLICAS {
        for other in 0..REPLICAS {
            let theirs = at[other];
            at[r].merge(theirs);
        }
    }

    let c = at[BENCH];
    println!(
        "  started {}, ended {}: {} GPUs in use",
        row(c.up.0),
        row(c.down.0),
        c.value()
    );
}

/// The number of GPUs one team may use at once.
const QUOTA: i64 = 6;

/// Lets each replica admit a job whenever its own view of the count leaves
/// room under the quota.
fn shared_limit() {
    let mut at = [PnCounter::default(); REPLICAS];
    at[SHELF].add(SHELF, 3);
    at[DESK].add(DESK, 1);
    let (shelf, desk) = (at[SHELF], at[DESK]);
    at[SHELF].merge(desk);
    at[DESK].merge(shelf);
    println!(
        "  both replicas see {} of {QUOTA} in use; then they lose contact",
        at[SHELF].value()
    );

    for r in [SHELF, DESK] {
        if at[r].value() + 2 <= QUOTA {
            at[r].add(r, 2);
            println!(
                "  {:<5} admits a 2-GPU job: by its count, {} of {QUOTA} are in use",
                NAMES[r],
                at[r].value()
            );
        }
    }

    let desk = at[DESK];
    at[SHELF].merge(desk);
    println!(
        "  after they merge: {} of {QUOTA} in use",
        at[SHELF].value()
    );
}

/// Gives each replica a share of the quota. A replica admits a job only from
/// its own share, and shares move between replicas when they can reach each
/// other.
fn allowances() {
    let mut share = [4, 2, 0];
    let mut used = [3, 1, 0];
    println!(
        "  shares {}, in use {}; then the replicas lose contact",
        row(share),
        row(used)
    );
    let admit = |share: &[i64; REPLICAS], used: &mut [i64; REPLICAS], r: usize, gpus: i64| {
        let left = share[r] - used[r];

        if left < gpus {
            println!(
                "  {:<5} refuses a {gpus}-GPU job: its share has {left} left",
                NAMES[r]
            );

            return;
        }

        used[r] += gpus;
        println!("  {:<5} admits a {gpus}-GPU job", NAMES[r]);
    };
    admit(&share, &mut used, SHELF, 2);
    admit(&share, &mut used, DESK, 2);
    println!("  contact returns; desk gives its unused GPU to shelf");
    share[DESK] -= 1;
    share[SHELF] += 1;
    admit(&share, &mut used, SHELF, 2);
    println!(
        "  shares {}, in use {}: {} of {QUOTA}",
        row(share),
        row(used),
        used.iter().sum::<i64>()
    );
}

fn main() {
    println!("one number per replica:");
    one_number();
    println!("one entry per replica, merged by taking the larger:");
    grow_only();
    println!("a count that goes down as well:");
    up_and_down();
    println!("a quota checked against each replica's own count:");
    shared_limit();
    println!("a quota divided into shares:");
    allowances();
}

References

• M. Shapiro, N. Preguiça, C. Baquero, M. Zawirski, “A comprehensive study of Convergent and Commutative Replicated Data Types”, INRIA research report 7506, 2011.
• V. Balegas and others, “Extending Eventually Consistent Cloud Databases for Enforcing Numeric Invariants”, 34th IEEE Symposium on Reliable Distributed Systems (SRDS), 2015, which describes the bounded counter with transferable shares.
• S. Gilbert and N. Lynch, “Brewer’s conjecture and the feasibility of consistent, available, partition-tolerant web services”, ACM SIGACT News, 33(2), 2002.
• Kubernetes, Resource quotas: https://kubernetes.io/docs/concepts/policy/resource-quotas/
• etcd, KV API guarantees: https://etcd.io/docs/latest/learning/api_guarantees/

10. A content-addressed artifact store

Before a job can run on a machine, the machine needs two large things: the container image with the job’s programs, and usually a model’s weights. A Thor or a Spark has room for a model of a hundred gigabytes, and a training run produces a new version of such a model every few hours. Something has to hold all of these versions, deliver them to the machines, and do so without copying a hundred gigabytes every time one layer of a model changes.

In this chapter, we will build the storage scheme that container registries and model stores share. We will first see what is wrong with naming a file by a name someone chose, and what changes when a file is named by its own contents. Then we will write cas, in Go and in Rust, a store that keeps artifacts as chunks named by their hashes, and use it to push two versions of a model, pull one to a machine that already has the other, detect a corrupted chunk and repair it, and reclaim space.

What a chosen name does not promise

Suppose we store model weights in files named by hand: llm-v1.bin, llm-v2.bin, and llm-latest.bin. A machine that is told to load llm-latest.bin and has a file of that name faces three questions that the name cannot answer.

• Is it the right version? latest pointed at one model yesterday and points at another today. Two machines in the same training job can hold different files under the same name, and the job’s results depend on which machine a worker landed on.
• Is it intact? A file that was cut short by a full disk or changed by a failing memory chip still has its name.
• Do I already have it? If llm-v2.bin is llm-v1.bin with one layer retrained, the two files are almost the same bytes, and nothing in their names says so. The machine downloads all of it.

All three come from one cause: the name and the contents are unrelated, so knowing one tells us nothing about the other.

A name computed from the contents

A cryptographic hash function takes any sequence of bytes and computes a short, fixed-size value from it. SHA-256, the one we will use, produces 256 bits, written as 64 hexadecimal digits. Two properties make it useful here. The same bytes always give the same value. And although different inputs that give the same value must exist, nobody knows how to find two, so in practice a value identifies one particular sequence of bytes. The value is called a digest.

A content-addressed store uses the digest as the name. To store a blob of bytes, we compute its digest and file the blob under it. To fetch a blob, we ask for a digest. This one change answers all three questions.

• A digest names exactly one sequence of bytes, for ever. Two machines that load the same digest load the same model.
• Anyone who receives a blob can compute its digest and compare it with the name they asked for. A blob that was cut short or altered does not match.
• Two blobs with the same contents have the same name, so storing the second one is a lookup that finds the first. A machine that is asked for a digest it holds has nothing to download.

The same scheme on our node

The container images on our cluster are stored this way already. containerd keeps the parts of every image in a directory of files named by their SHA-256 digests. On our kind node:

$ docker exec gpu-lab-control-plane sh -c \
    'ls /var/lib/containerd/io.containerd.content.v1.content/blobs/sha256 | head -3'
0286df745178d6f77c036891bed8ddc93a624c1fa9c054d134674a096ee882e6
069d1e267530c2e681fbd4d481553b4d05f98082b18fafac86e7f12996dddd0b
0771a8a05619f4bb6f4c8484f942f509b21a04ebe56ee26d17a3006e250bfd46
$ docker exec gpu-lab-control-plane sh -c \
    'ls /var/lib/containerd/io.containerd.content.v1.content/blobs/sha256 | wc -l'
168

The node holds 168 blobs. The images we built in earlier chapters appear with the first digits of their digests when we list them:

$ docker exec gpu-lab-control-plane crictl images --digests | grep -E 'IMAGE|gpu-'
IMAGE                            TAG   DIGEST          IMAGE ID        SIZE
docker.io/library/gpu-check      1     3714eb7096522   07ee9035514b4   73.3MB
docker.io/library/gpu-sharing    1     413391f612a31   4212abdefc4b5   74.3MB

gpu-check and gpu-sharing were built from the same base image, so most of their 73 MB are the same blobs, stored once. A pod’s spec can name an image by its digest, gpu-check@sha256:..., and the kubelet then runs exactly those bytes on every node.

Chunks and manifests

Naming a whole 100 GB file by one digest answers the first two questions and does little for the third: change one byte and the digest changes, and the whole file is a new blob. To share the unchanged parts between versions, we store a large artifact in pieces.

• The artifact is cut into chunks of a fixed size. Each chunk is a blob, stored under its own digest.
• A manifest is a small text file that lists the chunks’ digests in order. It is a blob too, and its digest names the whole artifact: a different chunk anywhere gives a different line in the manifest and so a different manifest digest.
• A tag is a name a person chose, such as llm:v2, that points at a manifest’s digest. Tags are the only part of the store that can change.

Figure 10.1 shows two versions of a model stored this way.

Version 2 is version 1 with its last quarter retrained. Its first six chunks have the same bytes as version 1’s, therefore the same digests, therefore they are the same blobs. Pushing version 2 stores two new chunks and one new manifest.

Container images have had this shape for years. An image is a manifest that lists layers, each layer is a blob named by its digest, and two images built from the same base share the base’s layers. The format is standardised as the OCI image specification, and registries that store images this way are now also used to store model weights as OCI artifacts.

Building cas

Our store keeps its blobs in memory, in a map from digest to bytes, so that the program has no files to clean up. A store on disk is the same map with one file per blob, in a directory layout like the one containerd showed us.

Blobs by digest

Put computes the digest, and stores the bytes only if that digest is not already a key. Get does more than a lookup: it hashes the bytes again and compares the result with the name it was asked for. Every read is therefore a check of the data.

Listing 10.1 Digests, the store, and verified reads, in Go. code/go/cas/main.go

// Digest names a blob: "sha256:" followed by the hash of its bytes in hex.
type Digest string

// digestOf computes the name of some bytes.
func digestOf(data []byte) Digest {
	sum := sha256.Sum256(data)

	return Digest("sha256:" + hex.EncodeToString(sum[:]))
}

// short returns the first twelve hex digits, for printing.
func (d Digest) short() string { return string(d)[:19] }

// Store holds blobs by digest and tags by name.
type Store struct {
	// blobs maps each digest to the bytes it names.
	blobs map[Digest][]byte
	// tags maps a name people use, such as "llm:v1", to a manifest's digest.
	tags map[string]Digest
}

// NewStore returns an empty store.
func NewStore() *Store {
	return &Store{blobs: map[Digest][]byte{}, tags: map[string]Digest{}}
}

// Put stores a blob and returns its digest and whether it was new. Bytes the
// store already holds are not stored a second time.
func (s *Store) Put(data []byte) (Digest, bool) {
	d := digestOf(data)

	if _, held := s.blobs[d]; held {
		return d, false
	}

	s.blobs[d] = data

	return d, true
}

// ErrCorrupt is returned when a blob's bytes no longer match its name.
var ErrCorrupt = errors.New("contents do not match the digest")

// Get returns a blob after checking that its bytes still hash to its name.
func (s *Store) Get(d Digest) ([]byte, error) {
	data, held := s.blobs[d]

	if !held {
		return nil, fmt.Errorf("%s: not found", d.short())
	}

	if digestOf(data) != d {
		return nil, fmt.Errorf("%s: %w", d.short(), ErrCorrupt)
	}

	return data, nil
}

Listing 10.2 Digests, the store, and verified reads, in Rust. code/rust/cas/src/main.rs

/// Names a blob: "sha256:" followed by the hash of its bytes in hex.
#[derive(Clone, PartialEq, Eq, PartialOrd, Ord)]
struct Digest(String);

impl Digest {
    /// Computes the name of some bytes.
    fn of(data: &[u8]) -> Self {
        let mut name = String::from("sha256:");

        for byte in Sha256::digest(data) {
            write!(name, "{byte:02x}").expect("writing to a String");
        }

        Digest(name)
    }

    /// Returns the first twelve hex digits, for printing.
    fn short(&self) -> &str {
        &self.0[..19]
    }
}

/// The ways reading from a store can fail.
#[derive(Debug)]
enum StoreError {
    /// The store holds no blob of this name.
    NotFound(String),
    /// A blob's bytes no longer match its name.
    Corrupt(String),
    /// A chunk of an artifact could not be read.
    Chunk(usize, Box<StoreError>),
}

impl std::fmt::Display for StoreError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            StoreError::NotFound(d) => write!(f, "{d}: not found"),
            StoreError::Corrupt(d) => write!(f, "{d}: contents do not match the digest"),
            StoreError::Chunk(i, e) => write!(f, "chunk {i}: {e}"),
        }
    }
}

/// Holds blobs by digest and tags by name.
#[derive(Default)]
struct Store {
    /// Maps each digest to the bytes it names.
    blobs: BTreeMap<Digest, Vec<u8>>,
    /// Maps a name people use, such as "llm:v1", to a manifest's digest.
    tags: BTreeMap<&'static str, Digest>,
}

impl Store {
    /// Stores a blob and returns its digest and whether it was new. Bytes the
    /// store already holds are not stored a second time.
    fn put(&mut self, data: &[u8]) -> (Digest, bool) {
        let d = Digest::of(data);

        if self.blobs.contains_key(&d) {
            return (d, false);
        }

        self.blobs.insert(d.clone(), data.to_vec());

        (d, true)
    }

    /// Returns a blob after checking that its bytes still hash to its name.
    fn get(&self, d: &Digest) -> Result<&[u8], StoreError> {
        let data = self
            .blobs
            .get(d)
            .ok_or_else(|| StoreError::NotFound(d.short().to_owned()))?;
        if Digest::of(data) != *d {
            return Err(StoreError::Corrupt(d.short().to_owned()));
        }

        Ok(data)
    }
}

The Go standard library includes SHA-256 in crypto/sha256. Rust’s standard library has no hash functions of this kind, and the Rust version uses the sha2 crate. The two programs print the same digests, which is itself a check: both computed the same function over the same bytes.

Pushing an artifact

Push walks the artifact in steps of the chunk size, stores each chunk with Put, and writes one line of the manifest for it. It counts how many chunks Put reported as new, which is the amount of data a push to a remote registry would have to send. The manifest is then stored like any other blob, and the tag is pointed at it.

Listing 10.3 Chunking an artifact and writing its manifest, in Go. code/go/cas/main.go

// chunkSize is the size of the pieces an artifact is cut into.
const chunkSize = 8 * 1024

// Push cuts an artifact into chunks, stores each chunk, stores a manifest
// that lists the chunks' digests in order, and points the tag at the
// manifest. It returns how many chunks were new to the store.
func (s *Store) Push(tag string, artifact []byte) (manifest Digest, chunks, added int) {
	var lines strings.Builder

	for start := 0; start < len(artifact); start += chunkSize {
		chunk := artifact[start:min(start+chunkSize, len(artifact))]
		d, fresh := s.Put(chunk)

		if fresh {
			added++
		}

		chunks++
		fmt.Fprintf(&lines, "%s %d\n", d, len(chunk))
	}

	manifest, _ = s.Put([]byte(lines.String()))
	s.tags[tag] = manifest

	return manifest, chunks, added
}

// chunksOf reads a manifest and returns the digests it lists.
func (s *Store) chunksOf(manifest Digest) ([]Digest, error) {
	data, err := s.Get(manifest)

	if err != nil {
		return nil, err
	}

	var out []Digest

	for line := range strings.Lines(string(data)) {
		digest, _, _ := strings.Cut(line, " ")
		out = append(out, Digest(digest))
	}

	return out, nil
}

Listing 10.4 Chunking an artifact and writing its manifest, in Rust. code/rust/cas/src/main.rs

impl Store {
    /// Cuts an artifact into chunks, stores each chunk, stores a manifest that
    /// lists the chunks' digests in order, and points the tag at the manifest.
    /// Returns the manifest's digest, the number of chunks, and how many were
    /// new to the store.
    fn push(&mut self, tag: &'static str, artifact: &[u8]) -> (Digest, usize, usize) {
        let mut lines = String::new();
        let (mut chunks, mut added) = (0, 0);

        for chunk in artifact.chunks(CHUNK_SIZE) {
            let (d, fresh) = self.put(chunk);
            added += usize::from(fresh);
            chunks += 1;
            writeln!(lines, "{} {}", d.0, chunk.len()).expect("writing to a String");
        }

        let (manifest, _) = self.put(lines.as_bytes());
        self.tags.insert(tag, manifest.clone());

        (manifest, chunks, added)
    }

    /// Reads a manifest and returns the digests it lists.
    fn chunks_of(&self, manifest: &Digest) -> Result<Vec<Digest>, StoreError> {
        let text = String::from_utf8_lossy(self.get(manifest)?);
        Ok(text
            .lines()
            .filter_map(|line| line.split(' ').next())
            .map(|d| Digest(d.to_owned()))
            .collect())
    }
}

Our chunks are 8 KiB so that the example stays small. Real stores use chunks of a few megabytes: large enough that the list of digests stays short, small enough that a change touches few of them.

Pulling and reading

Pull is how a machine gets an artifact. It reads the manifest from the remote store, and for each chunk it lists, fetches the chunk only if the local store has no blob of that name. Read reassembles the artifact from local chunks, and since it goes through Get, a damaged chunk stops the read with an error that says which chunk it was.

Listing 10.5 Fetching what is missing, and reading with verification, in Go. code/go/cas/main.go

// Pull copies a tagged artifact from a remote store into this one, fetching
// only the chunks this store does not hold. It returns the number of chunks
// and bytes fetched.
func (s *Store) Pull(remote *Store, tag string) (fetched, bytes int, err error) {
	manifest := remote.tags[tag]
	chunks, err := remote.chunksOf(manifest)

	if err != nil {
		return 0, 0, err
	}

	for _, d := range chunks {
		if _, held := s.blobs[d]; held {
			continue
		}

		data, err := remote.Get(d)

		if err != nil {
			return fetched, bytes, err
		}

		s.blobs[d] = data
		fetched++
		bytes += len(data)
	}

	data, _ := remote.Get(manifest)
	s.blobs[manifest] = data
	s.tags[tag] = manifest

	return fetched, bytes, nil
}

// Read reassembles a tagged artifact, verifying every chunk on the way.
func (s *Store) Read(tag string) ([]byte, error) {
	chunks, err := s.chunksOf(s.tags[tag])

	if err != nil {
		return nil, err
	}

	var artifact []byte

	for i, d := range chunks {
		data, err := s.Get(d)

		if err != nil {
			return nil, fmt.Errorf("chunk %d: %w", i, err)
		}

		artifact = append(artifact, data...)
	}

	return artifact, nil
}

Listing 10.6 Fetching what is missing, and reading with verification, in Rust. code/rust/cas/src/main.rs

impl Store {
    /// Copies a tagged artifact from a remote store into this one, fetching
    /// only the chunks this store does not hold. Returns the number of chunks
    /// and bytes fetched.
    fn pull(&mut self, remote: &Store, tag: &'static str) -> Result<(usize, usize), StoreError> {
        let manifest = &remote.tags[tag];
        let (mut fetched, mut bytes) = (0, 0);

        for d in remote.chunks_of(manifest)? {
            if self.blobs.contains_key(&d) {
                continue;
            }

            let data = remote.get(&d)?;
            fetched += 1;
            bytes += data.len();
            self.blobs.insert(d, data.to_vec());
        }

        self.blobs
            .insert(manifest.clone(), remote.get(manifest)?.to_vec());
        self.tags.insert(tag, manifest.clone());

        Ok((fetched, bytes))
    }

    /// Reassembles a tagged artifact, verifying every chunk on the way.
    fn read(&self, tag: &str) -> Result<Vec<u8>, StoreError> {
        let mut artifact = Vec::new();

        for (i, d) in self.chunks_of(&self.tags[tag])?.iter().enumerate() {
            let data = self.get(d).map_err(|e| StoreError::Chunk(i, Box::new(e)))?;
            artifact.extend_from_slice(data);
        }

        Ok(artifact)
    }
}

Reclaiming space

A content-addressed store grows with every push, and blobs cannot be deleted one artifact at a time, because a chunk may be listed by several manifests. The store removes blobs by garbage collection in two phases. The mark phase starts from the tags, which are the only names people use, and notes every manifest they point at and every chunk those manifests list. The sweep phase removes every blob that was not noted.

Listing 10.7 Mark from the tags, then sweep, in Go. code/go/cas/main.go

// Collect removes every blob that no tag leads to. It first marks the
// manifests the tags name and the chunks those manifests list, then sweeps
// the rest. It returns the number of blobs and bytes removed.
func (s *Store) Collect() (removed, bytes int) {
	live := map[Digest]bool{}

	for _, manifest := range s.tags {
		live[manifest] = true
		chunks, _ := s.chunksOf(manifest)

		for _, d := range chunks {
			live[d] = true
		}
	}

	for d, data := range s.blobs {
		if !live[d] {
			removed++
			bytes += len(data)
			delete(s.blobs, d)
		}
	}

	return removed, bytes
}

Listing 10.8 Mark from the tags, then sweep, in Rust. code/rust/cas/src/main.rs

impl Store {
    /// Removes every blob that no tag leads to. It first marks the manifests
    /// the tags name and the chunks those manifests list, then sweeps the
    /// rest. Returns the number of blobs and bytes removed.
    fn collect(&mut self) -> (usize, usize) {
        let mut live = BTreeSet::new();

        for manifest in self.tags.values() {
            live.extend(self.chunks_of(manifest).unwrap_or_default());
            live.insert(manifest.clone());
        }

        let (before, bytes_before) = self.size();
        self.blobs.retain(|d, _| live.contains(d));
        let (after, bytes_after) = self.size();

        (before - after, bytes_before - bytes_after)
    }
}

The run

The model’s weights are 64 KiB from a fixed pseudo-random sequence, which stand in for gigabytes. Version 2 keeps the first 48 KiB and replaces the last 16.

Listing 10.9 Two versions, a pull, a damaged chunk, and a collection, in Go. code/go/cas/main.go

func main() {
	// Version 1 is 64 KiB. Version 2 is a fine-tuned copy: the first three
	// quarters are unchanged and the last quarter is new.
	v1 := weights(64*1024, 1)
	v2 := append(append([]byte{}, v1[:48*1024]...), weights(16*1024, 2)...)

	registry := NewStore()
	m1, chunks, added := registry.Push("llm:v1", v1)
	fmt.Printf("push llm:v1: %d chunks, %d new; manifest %s\n", chunks, added, m1.short())
	m2, chunks, added := registry.Push("llm:v2", v2)
	fmt.Printf("push llm:v2: %d chunks, %d new; manifest %s\n", chunks, added, m2.short())
	blobs, bytes := registry.size()
	fmt.Printf("the registry holds %d blobs and %d bytes for %d bytes of artifacts\n", blobs, bytes, len(v1)+len(v2))

	node := NewStore()
	fetched, n, _ := node.Pull(registry, "llm:v1")
	fmt.Printf("thor-1 pulls llm:v1: %d chunks fetched, %d bytes\n", fetched, n)
	fetched, n, _ = node.Pull(registry, "llm:v2")
	fmt.Printf("thor-1 pulls llm:v2: %d chunks fetched, %d bytes\n", fetched, n)

	fmt.Println("one bit flips in a chunk stored on thor-1:")
	chunkList, _ := node.chunksOf(node.tags["llm:v1"])
	damaged := append([]byte{}, node.blobs[chunkList[3]]...)
	damaged[100] ^= 1
	node.blobs[chunkList[3]] = damaged

	if _, err := node.Read("llm:v1"); err != nil {
		fmt.Printf("  read llm:v1: %v\n", err)
	}

	// The bad chunk is dropped and pulled again, by name, from the registry.
	delete(node.blobs, chunkList[3])
	fetched, n, _ = node.Pull(registry, "llm:v1")
	fmt.Printf("  thor-1 drops the chunk and pulls llm:v1 again: %d chunk fetched, %d bytes\n", fetched, n)

	if got, err := node.Read("llm:v1"); err == nil {
		fmt.Printf("  read llm:v1: %d bytes, digest %s\n", len(got), digestOf(got).short())
	}

	fmt.Println("the tag llm:v1 is deleted from the registry:")
	delete(registry.tags, "llm:v1")
	removed, n := registry.Collect()
	blobs, bytes = registry.size()
	fmt.Printf("  collected %d blobs, %d bytes; %d blobs and %d bytes remain\n", removed, n, blobs, bytes)
}

Listing 10.10 Two versions, a pull, a damaged chunk, and a collection, in Rust. code/rust/cas/src/main.rs

fn main() -> Result<(), StoreError> {
    // Version 1 is 64 KiB. Version 2 is a fine-tuned copy: the first three
    // quarters are unchanged and the last quarter is new.
    let v1 = weights(64 * 1024, 1);
    let v2 = [&v1[..48 * 1024], &weights(16 * 1024, 2)].concat();

    let mut registry = Store::default();
    let (m1, chunks, added) = registry.push("llm:v1", &v1);
    println!(
        "push llm:v1: {chunks} chunks, {added} new; manifest {}",
        m1.short()
    );
    let (m2, chunks, added) = registry.push("llm:v2", &v2);
    println!(
        "push llm:v2: {chunks} chunks, {added} new; manifest {}",
        m2.short()
    );
    let (blobs, bytes) = registry.size();
    println!(
        "the registry holds {blobs} blobs and {bytes} bytes for {} bytes of artifacts",
        v1.len() + v2.len()
    );

    let mut node = Store::default();
    let (fetched, n) = node.pull(&registry, "llm:v1")?;
    println!("thor-1 pulls llm:v1: {fetched} chunks fetched, {n} bytes");
    let (fetched, n) = node.pull(&registry, "llm:v2")?;
    println!("thor-1 pulls llm:v2: {fetched} chunks fetched, {n} bytes");

    println!("one bit flips in a chunk stored on thor-1:");
    let fourth = node.chunks_of(&node.tags["llm:v1"])?[3].clone();
    node.blobs.get_mut(&fourth).expect("the chunk was pulled")[100] ^= 1;

    if let Err(e) = node.read("llm:v1") {
        println!("  read llm:v1: {e}");
    }

    // The bad chunk is dropped and pulled again, by name, from the registry.
    node.blobs.remove(&fourth);
    let (fetched, n) = node.pull(&registry, "llm:v1")?;
    println!("  thor-1 drops the chunk and pulls llm:v1 again: {fetched} chunk fetched, {n} bytes");
    let got = node.read("llm:v1")?;
    println!(
        "  read llm:v1: {} bytes, digest {}",
        got.len(),
        Digest::of(&got).short()
    );

    println!("the tag llm:v1 is deleted from the registry:");
    registry.tags.remove("llm:v1");
    let (removed, n) = registry.collect();
    let (blobs, bytes) = registry.size();
    println!("  collected {removed} blobs, {n} bytes; {blobs} blobs and {bytes} bytes remain");

    Ok(())
}

$ go run ./cas
push llm:v1: 8 chunks, 8 new; manifest sha256:eea2b9042521
push llm:v2: 8 chunks, 2 new; manifest sha256:7df8dd2bef7e
the registry holds 12 blobs and 83152 bytes for 131072 bytes of artifacts
thor-1 pulls llm:v1: 8 chunks fetched, 65536 bytes
thor-1 pulls llm:v2: 2 chunks fetched, 16384 bytes
one bit flips in a chunk stored on thor-1:
  read llm:v1: chunk 3: sha256:156fd5154cde: contents do not match the digest
  thor-1 drops the chunk and pulls llm:v1 again: 1 chunk fetched, 8192 bytes
  read llm:v1: 65536 bytes, digest sha256:a6a98d7f309f
the tag llm:v1 is deleted from the registry:
  collected 3 blobs, 17000 bytes; 9 blobs and 66152 bytes remain

Pushing llm:v1 stores eight chunks. Pushing llm:v2 cuts the new artifact into eight chunks as well and finds six of them already in the store, so it adds two. The registry holds 12 blobs, ten chunks and two manifests, in 83,152 bytes, for two artifacts that total 131,072 bytes.

thor-1 starts empty and fetches all eight chunks of version 1. For version 2 it fetches two chunks, 16,384 bytes, a quarter of the artifact. On a 100 GB model with the same proportion changed, that is 25 GB where a store of whole files would send 100, and at the 3.125 GB per second of one 25 GbE lane it is 8 seconds where the whole file takes 32.

Then one bit of one chunk on thor-1 flips. The next read of llm:v1 stops at chunk 3 and names the digest whose contents no longer match. The repair is as local as the damage: the machine drops that one blob and pulls again, the pull finds seven chunks present and one missing, and 8,192 bytes come from the registry. The digest printed for the reassembled artifact, sha256:a6a98d7f309f, is the digest of the whole 64 KiB of version 1.

Finally the tag llm:v1 is deleted. Collection marks the manifest of version 2 and its eight chunks, and sweeps three blobs: the manifest of version 1 and the two chunks that only version 1 listed. The six shared chunks stay, because version 2 still leads to them.

Animation 10.1 plays the run, with the registry’s chunks on the left and those on thor-1 on the right.

What the control plane gains

Content addressing is useful to every part of the system built in the earlier chapters.

A GpuJob that names its image and its model by digest runs the same bytes on every worker of a gang, which chapter 7’s synchronous step depends on. The checkpoint store of chapter 4 can be one of these: a checkpoint is an artifact, its chunks that did not change since the last checkpoint are not written again, and the epoch check sits on the tag, which is the only thing that is overwritten. And a machine that returns from a reimage, the last rung of chapter 4’s ladder, refills its local store by pulling, and the verification on every chunk doubles as a test of its new disk.

The complete program

Listing 10.11 The complete program in Go. code/go/cas/main.go

// Command cas builds a small content-addressed store for model artifacts and
// runs it through the situations such a store exists for.
//
// A blob is stored under the SHA-256 digest of its bytes. An artifact is cut
// into chunks, each chunk is a blob, and a manifest lists the chunks. The
// program pushes two versions of a model that differ in their last quarter,
// pulls the second to a machine that has the first, corrupts one stored
// chunk, and removes what nothing refers to any more.
//
// Run it from code/go with:
//
//	go run ./cas
package main

import (
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
	"strings"
)

// Digest names a blob: "sha256:" followed by the hash of its bytes in hex.
type Digest string

// digestOf computes the name of some bytes.
func digestOf(data []byte) Digest {
	sum := sha256.Sum256(data)

	return Digest("sha256:" + hex.EncodeToString(sum[:]))
}

// short returns the first twelve hex digits, for printing.
func (d Digest) short() string { return string(d)[:19] }

// Store holds blobs by digest and tags by name.
type Store struct {
	// blobs maps each digest to the bytes it names.
	blobs map[Digest][]byte
	// tags maps a name people use, such as "llm:v1", to a manifest's digest.
	tags map[string]Digest
}

// NewStore returns an empty store.
func NewStore() *Store {
	return &Store{blobs: map[Digest][]byte{}, tags: map[string]Digest{}}
}

// Put stores a blob and returns its digest and whether it was new. Bytes the
// store already holds are not stored a second time.
func (s *Store) Put(data []byte) (Digest, bool) {
	d := digestOf(data)

	if _, held := s.blobs[d]; held {
		return d, false
	}

	s.blobs[d] = data

	return d, true
}

// ErrCorrupt is returned when a blob's bytes no longer match its name.
var ErrCorrupt = errors.New("contents do not match the digest")

// Get returns a blob after checking that its bytes still hash to its name.
func (s *Store) Get(d Digest) ([]byte, error) {
	data, held := s.blobs[d]

	if !held {
		return nil, fmt.Errorf("%s: not found", d.short())
	}

	if digestOf(data) != d {
		return nil, fmt.Errorf("%s: %w", d.short(), ErrCorrupt)
	}

	return data, nil
}


// chunkSize is the size of the pieces an artifact is cut into.
const chunkSize = 8 * 1024

// Push cuts an artifact into chunks, stores each chunk, stores a manifest
// that lists the chunks' digests in order, and points the tag at the
// manifest. It returns how many chunks were new to the store.
func (s *Store) Push(tag string, artifact []byte) (manifest Digest, chunks, added int) {
	var lines strings.Builder

	for start := 0; start < len(artifact); start += chunkSize {
		chunk := artifact[start:min(start+chunkSize, len(artifact))]
		d, fresh := s.Put(chunk)

		if fresh {
			added++
		}

		chunks++
		fmt.Fprintf(&lines, "%s %d\n", d, len(chunk))
	}

	manifest, _ = s.Put([]byte(lines.String()))
	s.tags[tag] = manifest

	return manifest, chunks, added
}

// chunksOf reads a manifest and returns the digests it lists.
func (s *Store) chunksOf(manifest Digest) ([]Digest, error) {
	data, err := s.Get(manifest)

	if err != nil {
		return nil, err
	}

	var out []Digest

	for line := range strings.Lines(string(data)) {
		digest, _, _ := strings.Cut(line, " ")
		out = append(out, Digest(digest))
	}

	return out, nil
}


// Pull copies a tagged artifact from a remote store into this one, fetching
// only the chunks this store does not hold. It returns the number of chunks
// and bytes fetched.
func (s *Store) Pull(remote *Store, tag string) (fetched, bytes int, err error) {
	manifest := remote.tags[tag]
	chunks, err := remote.chunksOf(manifest)

	if err != nil {
		return 0, 0, err
	}

	for _, d := range chunks {
		if _, held := s.blobs[d]; held {
			continue
		}

		data, err := remote.Get(d)

		if err != nil {
			return fetched, bytes, err
		}

		s.blobs[d] = data
		fetched++
		bytes += len(data)
	}

	data, _ := remote.Get(manifest)
	s.blobs[manifest] = data
	s.tags[tag] = manifest

	return fetched, bytes, nil
}

// Read reassembles a tagged artifact, verifying every chunk on the way.
func (s *Store) Read(tag string) ([]byte, error) {
	chunks, err := s.chunksOf(s.tags[tag])

	if err != nil {
		return nil, err
	}

	var artifact []byte

	for i, d := range chunks {
		data, err := s.Get(d)

		if err != nil {
			return nil, fmt.Errorf("chunk %d: %w", i, err)
		}

		artifact = append(artifact, data...)
	}

	return artifact, nil
}


// Collect removes every blob that no tag leads to. It first marks the
// manifests the tags name and the chunks those manifests list, then sweeps
// the rest. It returns the number of blobs and bytes removed.
func (s *Store) Collect() (removed, bytes int) {
	live := map[Digest]bool{}

	for _, manifest := range s.tags {
		live[manifest] = true
		chunks, _ := s.chunksOf(manifest)

		for _, d := range chunks {
			live[d] = true
		}
	}

	for d, data := range s.blobs {
		if !live[d] {
			removed++
			bytes += len(data)
			delete(s.blobs, d)
		}
	}

	return removed, bytes
}


// size returns the number of blobs in a store and their total size.
func (s *Store) size() (blobs, bytes int) {
	for _, data := range s.blobs {
		bytes += len(data)
	}

	return len(s.blobs), bytes
}

// weights returns n bytes that stand for a model's weights. The bytes come
// from a fixed sequence, so every run and both languages produce the same.
func weights(n int, seed uint32) []byte {
	out := make([]byte, n)
	x := seed

	for i := range out {
		x = x*1664525 + 1013904223
		out[i] = byte(x >> 24)
	}

	return out
}

func main() {
	// Version 1 is 64 KiB. Version 2 is a fine-tuned copy: the first three
	// quarters are unchanged and the last quarter is new.
	v1 := weights(64*1024, 1)
	v2 := append(append([]byte{}, v1[:48*1024]...), weights(16*1024, 2)...)

	registry := NewStore()
	m1, chunks, added := registry.Push("llm:v1", v1)
	fmt.Printf("push llm:v1: %d chunks, %d new; manifest %s\n", chunks, added, m1.short())
	m2, chunks, added := registry.Push("llm:v2", v2)
	fmt.Printf("push llm:v2: %d chunks, %d new; manifest %s\n", chunks, added, m2.short())
	blobs, bytes := registry.size()
	fmt.Printf("the registry holds %d blobs and %d bytes for %d bytes of artifacts\n", blobs, bytes, len(v1)+len(v2))

	node := NewStore()
	fetched, n, _ := node.Pull(registry, "llm:v1")
	fmt.Printf("thor-1 pulls llm:v1: %d chunks fetched, %d bytes\n", fetched, n)
	fetched, n, _ = node.Pull(registry, "llm:v2")
	fmt.Printf("thor-1 pulls llm:v2: %d chunks fetched, %d bytes\n", fetched, n)

	fmt.Println("one bit flips in a chunk stored on thor-1:")
	chunkList, _ := node.chunksOf(node.tags["llm:v1"])
	damaged := append([]byte{}, node.blobs[chunkList[3]]...)
	damaged[100] ^= 1
	node.blobs[chunkList[3]] = damaged

	if _, err := node.Read("llm:v1"); err != nil {
		fmt.Printf("  read llm:v1: %v\n", err)
	}

	// The bad chunk is dropped and pulled again, by name, from the registry.
	delete(node.blobs, chunkList[3])
	fetched, n, _ = node.Pull(registry, "llm:v1")
	fmt.Printf("  thor-1 drops the chunk and pulls llm:v1 again: %d chunk fetched, %d bytes\n", fetched, n)

	if got, err := node.Read("llm:v1"); err == nil {
		fmt.Printf("  read llm:v1: %d bytes, digest %s\n", len(got), digestOf(got).short())
	}

	fmt.Println("the tag llm:v1 is deleted from the registry:")
	delete(registry.tags, "llm:v1")
	removed, n := registry.Collect()
	blobs, bytes = registry.size()
	fmt.Printf("  collected %d blobs, %d bytes; %d blobs and %d bytes remain\n", removed, n, blobs, bytes)
}

Listing 10.12 The complete program in Rust. code/rust/cas/src/main.rs

//! Builds a small content-addressed store for model artifacts and runs it
//! through the situations such a store exists for.
//!
//! A blob is stored under the SHA-256 digest of its bytes. An artifact is cut
//! into chunks, each chunk is a blob, and a manifest lists the chunks. The
//! program pushes two versions of a model that differ in their last quarter,
//! pulls the second to a machine that has the first, corrupts one stored
//! chunk, and removes what nothing refers to any more.
//!
//! Run it from code/rust with:
//!
//! ```text
//! cargo run -p cas
//! ```

use std::collections::{BTreeMap, BTreeSet};
use std::fmt::Write as _;

use sha2::{Digest as _, Sha256};

/// Names a blob: "sha256:" followed by the hash of its bytes in hex.
#[derive(Clone, PartialEq, Eq, PartialOrd, Ord)]
struct Digest(String);

impl Digest {
    /// Computes the name of some bytes.
    fn of(data: &[u8]) -> Self {
        let mut name = String::from("sha256:");

        for byte in Sha256::digest(data) {
            write!(name, "{byte:02x}").expect("writing to a String");
        }

        Digest(name)
    }

    /// Returns the first twelve hex digits, for printing.
    fn short(&self) -> &str {
        &self.0[..19]
    }
}

/// The ways reading from a store can fail.
#[derive(Debug)]
enum StoreError {
    /// The store holds no blob of this name.
    NotFound(String),
    /// A blob's bytes no longer match its name.
    Corrupt(String),
    /// A chunk of an artifact could not be read.
    Chunk(usize, Box<StoreError>),
}

impl std::fmt::Display for StoreError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            StoreError::NotFound(d) => write!(f, "{d}: not found"),
            StoreError::Corrupt(d) => write!(f, "{d}: contents do not match the digest"),
            StoreError::Chunk(i, e) => write!(f, "chunk {i}: {e}"),
        }
    }
}

/// Holds blobs by digest and tags by name.
#[derive(Default)]
struct Store {
    /// Maps each digest to the bytes it names.
    blobs: BTreeMap<Digest, Vec<u8>>,
    /// Maps a name people use, such as "llm:v1", to a manifest's digest.
    tags: BTreeMap<&'static str, Digest>,
}

impl Store {
    /// Stores a blob and returns its digest and whether it was new. Bytes the
    /// store already holds are not stored a second time.
    fn put(&mut self, data: &[u8]) -> (Digest, bool) {
        let d = Digest::of(data);

        if self.blobs.contains_key(&d) {
            return (d, false);
        }

        self.blobs.insert(d.clone(), data.to_vec());

        (d, true)
    }

    /// Returns a blob after checking that its bytes still hash to its name.
    fn get(&self, d: &Digest) -> Result<&[u8], StoreError> {
        let data = self
            .blobs
            .get(d)
            .ok_or_else(|| StoreError::NotFound(d.short().to_owned()))?;
        if Digest::of(data) != *d {
            return Err(StoreError::Corrupt(d.short().to_owned()));
        }

        Ok(data)
    }
}

impl Store {
    /// Cuts an artifact into chunks, stores each chunk, stores a manifest that
    /// lists the chunks' digests in order, and points the tag at the manifest.
    /// Returns the manifest's digest, the number of chunks, and how many were
    /// new to the store.
    fn push(&mut self, tag: &'static str, artifact: &[u8]) -> (Digest, usize, usize) {
        let mut lines = String::new();
        let (mut chunks, mut added) = (0, 0);

        for chunk in artifact.chunks(CHUNK_SIZE) {
            let (d, fresh) = self.put(chunk);
            added += usize::from(fresh);
            chunks += 1;
            writeln!(lines, "{} {}", d.0, chunk.len()).expect("writing to a String");
        }

        let (manifest, _) = self.put(lines.as_bytes());
        self.tags.insert(tag, manifest.clone());

        (manifest, chunks, added)
    }

    /// Reads a manifest and returns the digests it lists.
    fn chunks_of(&self, manifest: &Digest) -> Result<Vec<Digest>, StoreError> {
        let text = String::from_utf8_lossy(self.get(manifest)?);
        Ok(text
            .lines()
            .filter_map(|line| line.split(' ').next())
            .map(|d| Digest(d.to_owned()))
            .collect())
    }
}

impl Store {
    /// Copies a tagged artifact from a remote store into this one, fetching
    /// only the chunks this store does not hold. Returns the number of chunks
    /// and bytes fetched.
    fn pull(&mut self, remote: &Store, tag: &'static str) -> Result<(usize, usize), StoreError> {
        let manifest = &remote.tags[tag];
        let (mut fetched, mut bytes) = (0, 0);

        for d in remote.chunks_of(manifest)? {
            if self.blobs.contains_key(&d) {
                continue;
            }

            let data = remote.get(&d)?;
            fetched += 1;
            bytes += data.len();
            self.blobs.insert(d, data.to_vec());
        }

        self.blobs
            .insert(manifest.clone(), remote.get(manifest)?.to_vec());
        self.tags.insert(tag, manifest.clone());

        Ok((fetched, bytes))
    }

    /// Reassembles a tagged artifact, verifying every chunk on the way.
    fn read(&self, tag: &str) -> Result<Vec<u8>, StoreError> {
        let mut artifact = Vec::new();

        for (i, d) in self.chunks_of(&self.tags[tag])?.iter().enumerate() {
            let data = self.get(d).map_err(|e| StoreError::Chunk(i, Box::new(e)))?;
            artifact.extend_from_slice(data);
        }

        Ok(artifact)
    }
}

impl Store {
    /// Removes every blob that no tag leads to. It first marks the manifests
    /// the tags name and the chunks those manifests list, then sweeps the
    /// rest. Returns the number of blobs and bytes removed.
    fn collect(&mut self) -> (usize, usize) {
        let mut live = BTreeSet::new();

        for manifest in self.tags.values() {
            live.extend(self.chunks_of(manifest).unwrap_or_default());
            live.insert(manifest.clone());
        }

        let (before, bytes_before) = self.size();
        self.blobs.retain(|d, _| live.contains(d));
        let (after, bytes_after) = self.size();

        (before - after, bytes_before - bytes_after)
    }
}

impl Store {
    /// Returns the number of blobs in the store and their total size.
    fn size(&self) -> (usize, usize) {
        (self.blobs.len(), self.blobs.values().map(Vec::len).sum())
    }
}

/// The size of the pieces an artifact is cut into.
const CHUNK_SIZE: usize = 8 * 1024;

/// Returns n bytes that stand for a model's weights. The bytes come from a
/// fixed sequence, so every run and both languages produce the same.
fn weights(n: usize, seed: u32) -> Vec<u8> {
    let mut x = seed;
    (0..n)
        .map(|_| {
            x = x.wrapping_mul(1_664_525).wrapping_add(1_013_904_223);
            (x >> 24) as u8
        })
        .collect()
}

fn main() -> Result<(), StoreError> {
    // Version 1 is 64 KiB. Version 2 is a fine-tuned copy: the first three
    // quarters are unchanged and the last quarter is new.
    let v1 = weights(64 * 1024, 1);
    let v2 = [&v1[..48 * 1024], &weights(16 * 1024, 2)].concat();

    let mut registry = Store::default();
    let (m1, chunks, added) = registry.push("llm:v1", &v1);
    println!(
        "push llm:v1: {chunks} chunks, {added} new; manifest {}",
        m1.short()
    );
    let (m2, chunks, added) = registry.push("llm:v2", &v2);
    println!(
        "push llm:v2: {chunks} chunks, {added} new; manifest {}",
        m2.short()
    );
    let (blobs, bytes) = registry.size();
    println!(
        "the registry holds {blobs} blobs and {bytes} bytes for {} bytes of artifacts",
        v1.len() + v2.len()
    );

    let mut node = Store::default();
    let (fetched, n) = node.pull(&registry, "llm:v1")?;
    println!("thor-1 pulls llm:v1: {fetched} chunks fetched, {n} bytes");
    let (fetched, n) = node.pull(&registry, "llm:v2")?;
    println!("thor-1 pulls llm:v2: {fetched} chunks fetched, {n} bytes");

    println!("one bit flips in a chunk stored on thor-1:");
    let fourth = node.chunks_of(&node.tags["llm:v1"])?[3].clone();
    node.blobs.get_mut(&fourth).expect("the chunk was pulled")[100] ^= 1;

    if let Err(e) = node.read("llm:v1") {
        println!("  read llm:v1: {e}");
    }

    // The bad chunk is dropped and pulled again, by name, from the registry.
    node.blobs.remove(&fourth);
    let (fetched, n) = node.pull(&registry, "llm:v1")?;
    println!("  thor-1 drops the chunk and pulls llm:v1 again: {fetched} chunk fetched, {n} bytes");
    let got = node.read("llm:v1")?;
    println!(
        "  read llm:v1: {} bytes, digest {}",
        got.len(),
        Digest::of(&got).short()
    );

    println!("the tag llm:v1 is deleted from the registry:");
    registry.tags.remove("llm:v1");
    let (removed, n) = registry.collect();
    let (blobs, bytes) = registry.size();
    println!("  collected {removed} blobs, {n} bytes; {blobs} blobs and {bytes} bytes remain");

    Ok(())
}

References

• Open Container Initiative, Image Format Specification: https://github.com/opencontainers/image-spec/blob/main/spec.md
• Open Container Initiative, Distribution Specification: https://github.com/opencontainers/distribution-spec/blob/main/spec.md
• containerd, Content flow: https://github.com/containerd/containerd/blob/main/docs/content-flow.md
• Kubernetes, Images: https://kubernetes.io/docs/concepts/containers/images/
• NIST, FIPS 180-4, Secure Hash Standard: https://csrc.nist.gov/pubs/fips/180-4/upd1/final

11. A command-line tool for Windows, macOS, and Linux

Everything we have built so far runs inside the cluster or next to it, on Linux. The people who use the cluster sit somewhere else. A researcher submits jobs from a MacBook, an operator checks nodes from a Windows laptop, and a build pipeline does both from a Linux container. What they all run is a command-line tool, and it is the one piece of a GPU platform that has to work on three operating systems.

In this chapter, we will write that tool. gpuctl is small: it lists the nodes of a cluster with their GPUs, reports where it keeps its files, and prints its version. We will first decide how a command-line tool should behave so that both people and scripts can use it, then write it in Go and in Rust, and then build it for the three systems from our Jetson and see where each language makes that easy and where it does not.

What a command line promises

A command-line tool has two kinds of user. A person reads what it prints. A script runs it, captures what it prints, and decides what to do next. A tool that serves both keeps three channels apart.

• Standard output carries the result and nothing else. A script that runs gpuctl nodes -o json and feeds the output to another program must receive only JSON, with no progress messages or warnings mixed in.
• Standard error carries everything meant for a person: errors, warnings, and help text shown after a mistake.
• The exit status is a small number the tool returns to whatever started it. By convention 0 means success. A script checks the status before it trusts the output.

gpuctl uses three statuses. 0 is success. 1 means the command was understood and failed, for example because the cluster could not be reached. 2 means the command line itself was wrong. The distinction lets a script retry a failure of the first kind and stop at once on the second, since a mistyped command stays mistyped on the next attempt.

Figure 11.1 shows the tool’s shape.

Where a tool’s files live

A tool that remembers settings needs a place to keep them, and each operating system has its own rule for where that is.

Two details in this table catch programs that were written on one system and run on another. The directory differs, so a path built by joining the home directory and .config is wrong on two of the three. And the separator between the parts of a path differs: Windows uses a backslash. Both languages provide a path type that joins parts with the right separator, filepath.Join in Go and PathBuf::join in Rust, and a portable program builds every path with it and writes no separator of its own.

Go’s standard library knows the table: os.UserConfigDir returns the right directory on each system. Rust’s standard library has no such function, so the Rust version writes the three cases out. cfg!(windows) and cfg!(target_os = "macos") are decided when the program is compiled, so each build contains only its own case.

Listing 11.1 The configuration file and the kubeconfig, in Go. code/go/gpuctl/main.go

// configFile returns the path of gpuctl's own configuration file. The
// directory is the one each operating system sets aside for a user's
// settings.
func configFile() (string, error) {
	dir, err := os.UserConfigDir()

	if err != nil {
		return "", err
	}

	return filepath.Join(dir, "gpuctl", "config.toml"), nil
}

// kubeconfig returns the file that names the cluster and the credentials:
// the KUBECONFIG environment variable if it is set, and .kube/config in the
// user's home directory otherwise.
func kubeconfig() string {
	return clientcmd.NewDefaultClientConfigLoadingRules().GetDefaultFilename()
}

Listing 11.2 The configuration file and the kubeconfig, in Rust. code/rust/gpuctl/src/main.rs

/// Returns the path of gpuctl's own configuration file. The directory is the
/// one each operating system sets aside for a user's settings.
fn config_file() -> Option<PathBuf> {
    let dir = if cfg!(windows) {
        PathBuf::from(env::var_os("AppData")?)
    } else if cfg!(target_os = "macos") {
        env::home_dir()?.join("Library").join("Application Support")
    } else {
        match env::var_os("XDG_CONFIG_HOME") {
            Some(dir) if !dir.is_empty() => PathBuf::from(dir),
            _ => env::home_dir()?.join(".config"),
        }
    };

    Some(dir.join("gpuctl").join("config.toml"))
}

/// Returns the file that names the cluster and the credentials: the
/// KUBECONFIG environment variable if it is set, and .kube/config in the
/// user's home directory otherwise.
fn kubeconfig() -> Option<PathBuf> {
    match env::var_os("KUBECONFIG") {
        Some(path) if !path.is_empty() => Some(PathBuf::from(path)),
        _ => Some(env::home_dir()?.join(".kube").join("config")),
    }
}

The second path is the kubeconfig, the file from which kubectl and every other client reads the address of the API server and the user’s credentials. Its location follows one rule on all three systems: the KUBECONFIG environment variable if it is set, and .kube/config under the home directory otherwise. gpuctl reads the same file, so it talks to whichever cluster kubectl talks to. The Go version asks client-go’s loading rules for the path. The Rust version applies the rule itself.

Reading the nodes

gpuctl nodes makes one request, a list of Node objects, which we met in chapter 1, and keeps five fields of each: the name, the allocatable nvidia.com/gpu count from chapter 5, and the architecture, operating system, and kubelet version that the kubelet reports about its machine.

Listing 11.3 Listing the nodes and keeping five fields, in Go. code/go/gpuctl/main.go

// NodeInfo is one row of "gpuctl nodes".
type NodeInfo struct {
	Name    string `json:"name"`
	GPUs    int64  `json:"gpus"`
	Arch    string `json:"arch"`
	OS      string `json:"os"`
	Kubelet string `json:"kubelet"`
}

// listNodes asks the API server for every node and keeps the fields the
// command prints.
func listNodes(ctx context.Context) ([]NodeInfo, error) {
	rules := clientcmd.NewDefaultClientConfigLoadingRules()
	config, err := clientcmd.NewNonInteractiveDeferredLoadingClientConfig(rules, nil).ClientConfig()

	if err != nil {
		return nil, err
	}

	clientset, err := kubernetes.NewForConfig(config)

	if err != nil {
		return nil, err
	}

	nodes, err := clientset.CoreV1().Nodes().List(ctx, metav1.ListOptions{})

	if err != nil {
		return nil, err
	}

	rows := make([]NodeInfo, 0, len(nodes.Items))

	for _, n := range nodes.Items {
		gpus := n.Status.Allocatable["nvidia.com/gpu"]
		rows = append(rows, NodeInfo{
			Name:    n.Name,
			GPUs:    gpus.Value(),
			Arch:    n.Status.NodeInfo.Architecture,
			OS:      n.Status.NodeInfo.OperatingSystem,
			Kubelet: n.Status.NodeInfo.KubeletVersion,
		})
	}

	return rows, nil
}

Listing 11.4 Listing the nodes and keeping five fields, in Rust. code/rust/gpuctl/src/main.rs

/// One row of "gpuctl nodes".
#[derive(Serialize)]
struct NodeInfo {
    name: String,
    gpus: i64,
    arch: String,
    os: String,
    kubelet: String,
}

/// Asks the API server for every node and keeps the fields the command
/// prints.
async fn list_nodes() -> Result<Vec<NodeInfo>, kube::Error> {
    let client = Client::try_default().await?;
    let nodes = Api::<Node>::all(client)
        .list(&ListParams::default())
        .await?;
    Ok(nodes
        .into_iter()
        .map(|node| {
            let status = node.status.unwrap_or_default();
            let info = status.node_info.unwrap_or_default();
            let gpus = status
                .allocatable
                .and_then(|resources| resources.get("nvidia.com/gpu")?.0.parse().ok())
                .unwrap_or(0);
            NodeInfo {
                name: node.metadata.name.unwrap_or_default(),
                gpus,
                arch: info.architecture,
                os: info.operating_system,
                kubelet: info.kubelet_version,
            }
        })
        .collect())
}

The Go version builds its client from the kubeconfig with clientcmd, the package of client-go that implements the loading rule. Earlier programs used controller-runtime’s GetConfigOrDie for this. A command-line tool wants the error back as a value, so that it can print it and choose the exit status. In the Rust version, Client::try_default does the same loading and returns a Result.

Kubernetes stores a resource amount as a quantity, a number with an optional unit suffix such as 16031232Ki. A GPU count is a plain whole number. The Go version reads it with the quantity’s Value method. In k8s-openapi a quantity is a wrapper around the string the API server sent, so the Rust version parses that string as an integer, and a node with no GPUs, which has no such entry at all, counts as zero.

One result, two forms

The rows are printed in one of two forms, chosen by the -o flag.

Listing 11.5 A table or JSON, in Go. code/go/gpuctl/main.go

// printNodes writes the rows as a table for a person or as JSON for a program.
func printNodes(out io.Writer, rows []NodeInfo, format string) error {
	if format == "json" {
		enc := json.NewEncoder(out)
		enc.SetIndent("", "  ")

		return enc.Encode(rows)
	}

	table := tabwriter.NewWriter(out, 0, 0, 2, ' ', 0)
	fmt.Fprintln(table, "NAME\tGPUS\tARCH\tOS\tKUBELET")

	for _, r := range rows {
		fmt.Fprintf(table, "%s\t%d\t%s\t%s\t%s\n", r.Name, r.GPUs, r.Arch, r.OS, r.Kubelet)
	}

	return table.Flush()
}

Listing 11.6 A table or JSON, in Rust. code/rust/gpuctl/src/main.rs

/// Writes the rows as a table for a person or as JSON for a program.
fn print_nodes(rows: &[NodeInfo], json: bool) -> serde_json::Result<()> {
    if json {
        println!("{}", serde_json::to_string_pretty(rows)?);

        return Ok(());
    }

    let header = ["NAME", "GPUS", "ARCH", "OS", "KUBELET"].map(String::from);
    let table: Vec<[String; 5]> = std::iter::once(header)
        .chain(rows.iter().map(|r| {
            [
                r.name.clone(),
                r.gpus.to_string(),
                r.arch.clone(),
                r.os.clone(),
                r.kubelet.clone(),
            ]
        }))
        .collect();
    // Each column is as wide as its longest cell, plus two spaces.
    let widths: [usize; 5] =
        std::array::from_fn(|col| table.iter().map(|row| row[col].len()).max().unwrap_or(0));

    for row in &table {
        let line: String = (0..5)
            .map(|col| format!("{:<width$}", row[col], width = widths[col] + 2))
            .collect();
        println!("{}", line.trim_end());
    }

    Ok(())
}

The table is for reading. Its columns are padded so that they line up, which Go’s text/tabwriter does from tab-separated cells and the Rust version does by measuring each column’s widest cell.

The JSON is for programs, and it is part of the tool’s interface in a way the table is not. A script that reads gpus from each object keeps working when a later version adds a column to the table or changes its spacing. For that reason the field names are written out in both versions, as struct tags in Go and as the field names of a Serialize struct in Rust, and they would not be renamed lightly.

Commands and exit statuses

run ties the pieces together. It looks at the first argument, carries out the command, and returns the status. Every path through it ends in one of the three statuses, and every message for a person goes to standard error.

Listing 11.7 Choosing the command and returning the status, in Go. code/go/gpuctl/main.go

// Exit statuses. A script that calls gpuctl tells a failed command from a
// mistyped one by the status alone.
const (
	exitOK      = 0
	exitFailure = 1
	exitUsage   = 2
)

const usage = `usage:
  gpuctl nodes [-o table|json]
  gpuctl config
  gpuctl version
`

Listing 11.8 Choosing the command and returning the status, in Rust. code/rust/gpuctl/src/main.rs

/// The exit status for a command line that is wrong. A script that calls
/// gpuctl tells a failed command from a mistyped one by the status alone:
/// success is 0 and failure is 1, which `ExitCode` already names.
const EXIT_USAGE: u8 = 2;

const USAGE: &str = "usage:
  gpuctl nodes [-o table|json]
  gpuctl config
  gpuctl version
";

Listing 11.9 The run function and main, in Go. code/go/gpuctl/main.go

// run carries out one command line and returns the exit status. It takes its
// arguments and its two output streams as parameters, so a test can call it
// without starting a process.
func run(args []string, out, errOut io.Writer) int {
	if len(args) == 0 {
		fmt.Fprint(errOut, usage)

		return exitUsage
	}

	switch args[0] {
	case "version":
		fmt.Fprintf(out, "gpuctl %s %s/%s\n", version, runtime.GOOS, runtime.GOARCH)
	case "config":
		file, err := configFile()

		if err != nil {
			fmt.Fprintln(errOut, "gpuctl:", err)

			return exitFailure
		}

		fmt.Fprintf(out, "config file  %s\nkubeconfig   %s\n", file, kubeconfig())
	case "nodes":
		flags := flag.NewFlagSet("nodes", flag.ContinueOnError)
		flags.SetOutput(io.Discard)
		format := flags.String("o", "table", "output format: table or json")

		if err := flags.Parse(args[1:]); err != nil || (*format != "table" && *format != "json") {
			fmt.Fprint(errOut, usage)

			return exitUsage
		}

		rows, err := listNodes(context.Background())

		if err == nil {
			err = printNodes(out, rows, *format)
		}

		if err != nil {
			fmt.Fprintln(errOut, "gpuctl:", err)

			return exitFailure
		}
	default:
		fmt.Fprintf(errOut, "gpuctl: unknown command %q\n%s", args[0], usage)

		return exitUsage
	}

	return exitOK
}

func main() {
	os.Exit(run(os.Args[1:], os.Stdout, os.Stderr))
}

Listing 11.10 The run function and main, in Rust. code/rust/gpuctl/src/main.rs

/// Carries out one command line and returns the exit status.
async fn run(args: &[String]) -> ExitCode {
    let args: Vec<&str> = args.iter().map(String::as_str).collect();

    match args.as_slice() {
        ["version"] => println!("gpuctl {} {}", env!("CARGO_PKG_VERSION"), platform()),
        ["config"] => {
            let (Some(file), Some(kube)) = (config_file(), kubeconfig()) else {
                eprintln!("gpuctl: cannot find the home directory");

                return ExitCode::FAILURE;
            };
            println!(
                "config file  {}\nkubeconfig   {}",
                file.display(),
                kube.display()
            );
        }
        ["nodes"] | ["nodes", "-o", "table" | "json"] => {
            let json = args.last() == Some(&"json");
            let rows = match list_nodes().await {
                Ok(rows) => rows,
                Err(e) => {
                    eprintln!("gpuctl: {e}");

                    return ExitCode::FAILURE;
                }
            };

            if let Err(e) = print_nodes(&rows, json) {
                eprintln!("gpuctl: {e}");

                return ExitCode::FAILURE;
            }
        }
        [] | ["nodes", ..] => {
            eprint!("{USAGE}");

            return ExitCode::from(EXIT_USAGE);
        }
        [other, ..] => {
            eprint!("gpuctl: unknown command \"{other}\"\n{USAGE}");

            return ExitCode::from(EXIT_USAGE);
        }
    }

    ExitCode::SUCCESS
}

#[tokio::main]
async fn main() -> ExitCode {
    let args: Vec<String> = env::args().skip(1).collect();

    run(&args).await
}

The Go version parses the -o flag with the standard flag package and tells it to stay silent, so that run prints one consistent usage text for every kind of mistake. The Rust version matches the whole argument list against patterns. The pattern ["nodes", "-o", "table" | "json"] accepts exactly the two valid forms, and any other list that starts with nodes falls through to the usage arm.

Here is the tool on our cluster. Both versions print exactly this:

$ gpuctl nodes
NAME                   GPUS  ARCH   OS     KUBELET
gpu-lab-control-plane  1     arm64  linux  v1.37.0
$ gpuctl nodes -o json
[
  {
    "name": "gpu-lab-control-plane",
    "gpus": 1,
    "arch": "arm64",
    "os": "linux",
    "kubelet": "v1.37.0"
  }
]
$ gpuctl config
config file  /home/jetson/.config/gpuctl/config.toml
kubeconfig   /home/jetson/.kube/config
$ gpuctl version
gpuctl 0.1.0 linux/arm64
$ gpuctl pods
gpuctl: unknown command "pods"
usage:
  gpuctl nodes [-o table|json]
  gpuctl config
  gpuctl version
$ echo $?
2

The table and the JSON hold the same row: our one node, with its one GPU, on arm64. The unknown command prints its message and the usage text on standard error and exits with 2. When the cluster cannot be reached, each version prints the error its client library produced and exits with 1. The Go version prints:

$ KUBECONFIG=/nonexistent gpuctl nodes
gpuctl: invalid configuration: no configuration has been provided, try setting KUBERNETES_MASTER environment variable
$ echo $?
1

Animation 11.1 shows three of these runs, each with its two streams and its status.

Building for three systems

A program that runs on our Jetson is a file of ARM machine code in the format Linux loads, called ELF. A Windows laptop with an Intel processor needs different machine code in a different format, PE, and a MacBook needs ARM code in a third, Mach-O. Producing a program for a system other than the one the compiler runs on is called cross-compilation. Figure 11.2 shows what we are after.

In Go

The Go compiler and linker are written in Go and contain the code generators for every platform Go supports. Two environment variables select the target: GOOS names the operating system and GOARCH the processor. On our Jetson:

$ GOOS=windows GOARCH=amd64 go build -o gpuctl-windows-amd64 ./gpuctl
$ GOOS=darwin  GOARCH=arm64 go build -o gpuctl-darwin-arm64  ./gpuctl
$ GOOS=linux   GOARCH=amd64 go build -o gpuctl-linux-amd64   ./gpuctl
$ file gpuctl-*
gpuctl-darwin-arm64:  Mach-O 64-bit arm64 executable, flags:<|DYLDLINK|PIE>
gpuctl-linux-amd64:   ELF 64-bit LSB executable, x86-64, version 1 (SYSV), statically linked
gpuctl-windows-amd64: PE32+ executable (console) x86-64, for MS Windows

Three commands on an ARM Linux board produced programs for Windows, macOS, and Intel Linux. Each file is between 37 and 39 MB and needs nothing else installed: the Go runtime and all of client-go are inside it. This works because gpuctl and everything it imports are written in Go. A Go program that calls C code through cgo loses this property and needs a C compiler for each target.

In Rust

Rust selects a target with a target triple, such as x86_64-unknown-linux-gnu or aarch64-apple-darwin, and rustup target add installs the standard library for one. Our Jetson has these installed:

$ rustup target list --installed
aarch64-apple-darwin
aarch64-unknown-linux-gnu
x86_64-unknown-linux-gnu

The Rust compiler can generate code for all three. The build stops for another reason:

$ cargo check -p gpuctl --target x86_64-unknown-linux-gnu
warning: ring@0.17.14: Compiler family detection failed due to error: ToolNotFound:
failed to find tool "x86_64-linux-gnu-gcc": No such file or directory (os error 2)
error: failed to run custom build command for `ring v0.17.14`

ring is the cryptography library that kube-rs uses, through its TLS library, for its encrypted connection to the API server. Parts of it are written in C and assembly, and its build compiles them with a C compiler for the target. Our board has a C compiler for ARM Linux and none for Intel Linux or for macOS. The same would happen at the last step of any Rust build for another system: Rust calls the target system’s linker to produce the final file, and that linker has to be installed.

So cross-compiling Rust means supplying a C toolchain for each target. Projects handle this in one of two ways. Most build each target on a machine that runs it: a pipeline with one Linux, one macOS, and one Windows runner, each running a plain cargo build --release and uploading its binary. Others install cross toolchains, either one by one or through a tool such as cross, which runs the build inside a container image that has the right toolchain for the target.

A Rust program with no C in its dependencies gets further. fleethealth from chapter 4 uses only the standard library, and cargo check -p fleethealth --target aarch64-apple-darwin passes on our Jetson, which shows that the compiler itself has no difficulty with the target. The difficulty belongs to one dependency with C in it, the same condition under which Go’s cross-compilation stops working.

The complete program

Listing 11.11 The complete program in Go. code/go/gpuctl/main.go

// Command gpuctl is a small command-line client for a GPU cluster, written to
// build and behave the same on Linux, macOS, and Windows.
//
//	gpuctl nodes [-o table|json]   list the cluster's nodes and their GPUs
//	gpuctl config                  print where gpuctl looks for its files
//	gpuctl version                 print the version and the platform
//
// Data goes to standard output and messages to standard error. The exit
// status is 0 on success, 1 when a command fails, and 2 when the command line
// is wrong.
//
// Run it from code/go with:
//
//	go run ./gpuctl nodes
package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"runtime"
	"text/tabwriter"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/tools/clientcmd"
)

// version is the tool's version, printed by "gpuctl version".
const version = "0.1.0"

// Exit statuses. A script that calls gpuctl tells a failed command from a
// mistyped one by the status alone.
const (
	exitOK      = 0
	exitFailure = 1
	exitUsage   = 2
)

const usage = `usage:
  gpuctl nodes [-o table|json]
  gpuctl config
  gpuctl version
`


// configFile returns the path of gpuctl's own configuration file. The
// directory is the one each operating system sets aside for a user's
// settings.
func configFile() (string, error) {
	dir, err := os.UserConfigDir()

	if err != nil {
		return "", err
	}

	return filepath.Join(dir, "gpuctl", "config.toml"), nil
}

// kubeconfig returns the file that names the cluster and the credentials:
// the KUBECONFIG environment variable if it is set, and .kube/config in the
// user's home directory otherwise.
func kubeconfig() string {
	return clientcmd.NewDefaultClientConfigLoadingRules().GetDefaultFilename()
}


// NodeInfo is one row of "gpuctl nodes".
type NodeInfo struct {
	Name    string `json:"name"`
	GPUs    int64  `json:"gpus"`
	Arch    string `json:"arch"`
	OS      string `json:"os"`
	Kubelet string `json:"kubelet"`
}

// listNodes asks the API server for every node and keeps the fields the
// command prints.
func listNodes(ctx context.Context) ([]NodeInfo, error) {
	rules := clientcmd.NewDefaultClientConfigLoadingRules()
	config, err := clientcmd.NewNonInteractiveDeferredLoadingClientConfig(rules, nil).ClientConfig()

	if err != nil {
		return nil, err
	}

	clientset, err := kubernetes.NewForConfig(config)

	if err != nil {
		return nil, err
	}

	nodes, err := clientset.CoreV1().Nodes().List(ctx, metav1.ListOptions{})

	if err != nil {
		return nil, err
	}

	rows := make([]NodeInfo, 0, len(nodes.Items))

	for _, n := range nodes.Items {
		gpus := n.Status.Allocatable["nvidia.com/gpu"]
		rows = append(rows, NodeInfo{
			Name:    n.Name,
			GPUs:    gpus.Value(),
			Arch:    n.Status.NodeInfo.Architecture,
			OS:      n.Status.NodeInfo.OperatingSystem,
			Kubelet: n.Status.NodeInfo.KubeletVersion,
		})
	}

	return rows, nil
}


// printNodes writes the rows as a table for a person or as JSON for a program.
func printNodes(out io.Writer, rows []NodeInfo, format string) error {
	if format == "json" {
		enc := json.NewEncoder(out)
		enc.SetIndent("", "  ")

		return enc.Encode(rows)
	}

	table := tabwriter.NewWriter(out, 0, 0, 2, ' ', 0)
	fmt.Fprintln(table, "NAME\tGPUS\tARCH\tOS\tKUBELET")

	for _, r := range rows {
		fmt.Fprintf(table, "%s\t%d\t%s\t%s\t%s\n", r.Name, r.GPUs, r.Arch, r.OS, r.Kubelet)
	}

	return table.Flush()
}


// run carries out one command line and returns the exit status. It takes its
// arguments and its two output streams as parameters, so a test can call it
// without starting a process.
func run(args []string, out, errOut io.Writer) int {
	if len(args) == 0 {
		fmt.Fprint(errOut, usage)

		return exitUsage
	}

	switch args[0] {
	case "version":
		fmt.Fprintf(out, "gpuctl %s %s/%s\n", version, runtime.GOOS, runtime.GOARCH)
	case "config":
		file, err := configFile()

		if err != nil {
			fmt.Fprintln(errOut, "gpuctl:", err)

			return exitFailure
		}

		fmt.Fprintf(out, "config file  %s\nkubeconfig   %s\n", file, kubeconfig())
	case "nodes":
		flags := flag.NewFlagSet("nodes", flag.ContinueOnError)
		flags.SetOutput(io.Discard)
		format := flags.String("o", "table", "output format: table or json")

		if err := flags.Parse(args[1:]); err != nil || (*format != "table" && *format != "json") {
			fmt.Fprint(errOut, usage)

			return exitUsage
		}

		rows, err := listNodes(context.Background())

		if err == nil {
			err = printNodes(out, rows, *format)
		}

		if err != nil {
			fmt.Fprintln(errOut, "gpuctl:", err)

			return exitFailure
		}
	default:
		fmt.Fprintf(errOut, "gpuctl: unknown command %q\n%s", args[0], usage)

		return exitUsage
	}

	return exitOK
}

func main() {
	os.Exit(run(os.Args[1:], os.Stdout, os.Stderr))
}

Listing 11.12 The complete program in Rust. code/rust/gpuctl/src/main.rs

//! A small command-line client for a GPU cluster, written to build and behave
//! the same on Linux, macOS, and Windows.
//!
//! ```text
//! gpuctl nodes [-o table|json]   list the cluster's nodes and their GPUs
//! gpuctl config                  print where gpuctl looks for its files
//! gpuctl version                 print the version and the platform
//! ```
//!
//! Data goes to standard output and messages to standard error. The exit
//! status is 0 on success, 1 when a command fails, and 2 when the command line
//! is wrong.
//!
//! Run it from code/rust with:
//!
//! ```text
//! cargo run -p gpuctl -- nodes
//! ```

use std::env;
use std::path::PathBuf;
use std::process::ExitCode;

use k8s_openapi::api::core::v1::Node;
use kube::{Api, Client, api::ListParams};
use serde::Serialize;

/// The exit status for a command line that is wrong. A script that calls
/// gpuctl tells a failed command from a mistyped one by the status alone:
/// success is 0 and failure is 1, which `ExitCode` already names.
const EXIT_USAGE: u8 = 2;

const USAGE: &str = "usage:
  gpuctl nodes [-o table|json]
  gpuctl config
  gpuctl version
";

/// Returns the path of gpuctl's own configuration file. The directory is the
/// one each operating system sets aside for a user's settings.
fn config_file() -> Option<PathBuf> {
    let dir = if cfg!(windows) {
        PathBuf::from(env::var_os("AppData")?)
    } else if cfg!(target_os = "macos") {
        env::home_dir()?.join("Library").join("Application Support")
    } else {
        match env::var_os("XDG_CONFIG_HOME") {
            Some(dir) if !dir.is_empty() => PathBuf::from(dir),
            _ => env::home_dir()?.join(".config"),
        }
    };

    Some(dir.join("gpuctl").join("config.toml"))
}

/// Returns the file that names the cluster and the credentials: the
/// KUBECONFIG environment variable if it is set, and .kube/config in the
/// user's home directory otherwise.
fn kubeconfig() -> Option<PathBuf> {
    match env::var_os("KUBECONFIG") {
        Some(path) if !path.is_empty() => Some(PathBuf::from(path)),
        _ => Some(env::home_dir()?.join(".kube").join("config")),
    }
}

/// One row of "gpuctl nodes".
#[derive(Serialize)]
struct NodeInfo {
    name: String,
    gpus: i64,
    arch: String,
    os: String,
    kubelet: String,
}

/// Asks the API server for every node and keeps the fields the command
/// prints.
async fn list_nodes() -> Result<Vec<NodeInfo>, kube::Error> {
    let client = Client::try_default().await?;
    let nodes = Api::<Node>::all(client)
        .list(&ListParams::default())
        .await?;
    Ok(nodes
        .into_iter()
        .map(|node| {
            let status = node.status.unwrap_or_default();
            let info = status.node_info.unwrap_or_default();
            let gpus = status
                .allocatable
                .and_then(|resources| resources.get("nvidia.com/gpu")?.0.parse().ok())
                .unwrap_or(0);
            NodeInfo {
                name: node.metadata.name.unwrap_or_default(),
                gpus,
                arch: info.architecture,
                os: info.operating_system,
                kubelet: info.kubelet_version,
            }
        })
        .collect())
}

/// Writes the rows as a table for a person or as JSON for a program.
fn print_nodes(rows: &[NodeInfo], json: bool) -> serde_json::Result<()> {
    if json {
        println!("{}", serde_json::to_string_pretty(rows)?);

        return Ok(());
    }

    let header = ["NAME", "GPUS", "ARCH", "OS", "KUBELET"].map(String::from);
    let table: Vec<[String; 5]> = std::iter::once(header)
        .chain(rows.iter().map(|r| {
            [
                r.name.clone(),
                r.gpus.to_string(),
                r.arch.clone(),
                r.os.clone(),
                r.kubelet.clone(),
            ]
        }))
        .collect();
    // Each column is as wide as its longest cell, plus two spaces.
    let widths: [usize; 5] =
        std::array::from_fn(|col| table.iter().map(|row| row[col].len()).max().unwrap_or(0));

    for row in &table {
        let line: String = (0..5)
            .map(|col| format!("{:<width$}", row[col], width = widths[col] + 2))
            .collect();
        println!("{}", line.trim_end());
    }

    Ok(())
}

/// Returns the platform in the form Go prints it, such as `linux/arm64`.
/// Rust calls the architectures `aarch64` and `x86_64`.
fn platform() -> String {
    let arch = match env::consts::ARCH {
        "aarch64" => "arm64",
        "x86_64" => "amd64",
        other => other,
    };

    format!("{}/{arch}", env::consts::OS)
}

/// Carries out one command line and returns the exit status.
async fn run(args: &[String]) -> ExitCode {
    let args: Vec<&str> = args.iter().map(String::as_str).collect();

    match args.as_slice() {
        ["version"] => println!("gpuctl {} {}", env!("CARGO_PKG_VERSION"), platform()),
        ["config"] => {
            let (Some(file), Some(kube)) = (config_file(), kubeconfig()) else {
                eprintln!("gpuctl: cannot find the home directory");

                return ExitCode::FAILURE;
            };
            println!(
                "config file  {}\nkubeconfig   {}",
                file.display(),
                kube.display()
            );
        }
        ["nodes"] | ["nodes", "-o", "table" | "json"] => {
            let json = args.last() == Some(&"json");
            let rows = match list_nodes().await {
                Ok(rows) => rows,
                Err(e) => {
                    eprintln!("gpuctl: {e}");

                    return ExitCode::FAILURE;
                }
            };

            if let Err(e) = print_nodes(&rows, json) {
                eprintln!("gpuctl: {e}");

                return ExitCode::FAILURE;
            }
        }
        [] | ["nodes", ..] => {
            eprint!("{USAGE}");

            return ExitCode::from(EXIT_USAGE);
        }
        [other, ..] => {
            eprint!("gpuctl: unknown command \"{other}\"\n{USAGE}");

            return ExitCode::from(EXIT_USAGE);
        }
    }

    ExitCode::SUCCESS
}

#[tokio::main]
async fn main() -> ExitCode {
    let args: Vec<String> = env::args().skip(1).collect();

    run(&args).await
}

References

• Go, os.UserConfigDir: https://pkg.go.dev/os#UserConfigDir
• Go, Installing Go from source, the list of GOOS and GOARCH values: https://go.dev/doc/install/source#environment
• The rustc book, Platform support: https://doc.rust-lang.org/rustc/platform-support.html
• The Cargo book, cargo build --target: https://doc.rust-lang.org/cargo/commands/cargo-build.html
• Kubernetes, Organizing cluster access using kubeconfig files: https://kubernetes.io/docs/concepts/configuration/organize-cluster-access-kubeconfig/
• XDG Base Directory Specification: https://specifications.freedesktop.org/basedir-spec/latest/
• cross, zero-setup cross compilation for Rust: https://github.com/cross-rs/cross

Appendix A. A local GPU cluster with kind

Our controller in section 1.1 runs pods that ask for GPUs. In this appendix, we will set up a kind cluster on a Jetson so that those pods really run on the board’s GPU. We will first follow a GPU request all the way from a pod to the hardware, so that we know every piece involved and why it is there. Then we will look at a fresh kind node, see which of those pieces it is missing, and add them one at a time, checking our work after each step. Finally we will run a CUDA program in a pod and watch it use the GPU. All the files are in code/jetson-gpu.

We tested every step on this system:

$ tr -d '\0' < /proc/device-tree/model; echo
NVIDIA Jetson Orin NX Engineering Reference Developer Kit Super
$ head -1 /etc/nv_tegra_release
# R36 (release), REVISION: 4.3, GCID: 38968081, BOARD: generic, EABI: aarch64, DATE: Wed Jan  8 01:49:37 UTC 2025
$ nvidia-ctk --version | head -1
NVIDIA Container Toolkit CLI version 1.17.3
$ kind version
kind v0.33.0 go1.26.7 linux/arm64

Release R36.4.3 of Jetson Linux is part of JetPack 6.2, which ships CUDA 12.6 and the NVIDIA Container Toolkit.

How a program reaches the GPU

The GPU in a Jetson sits on the same chip as the CPU and shares its memory. A program on the board reaches it in the same two ways it would reach any device on Linux.

The first is through device files under /dev. The kernel’s GPU driver creates files such as /dev/nvhost-gpu and /dev/nvmap, and a program talks to the driver by opening these files and making system calls on them.

The second is through the user-space driver libraries: code that runs inside the program and knows how to talk to those device files. The most important one is libcuda.so, the CUDA driver library. When a CUDA program calls a function such as cudaMalloc, the call ends up in libcuda.so, which turns it into operations on the device files. JetPack installs these libraries under /usr/lib/aarch64-linux-gnu/nvidia.

A program running directly on the board has both of these. A program in a container has neither.

Why a container cannot see the GPU

A container’s files come from its image, and its /dev holds only a handful of harmless devices such as /dev/null. The GPU’s device files are not there, and the host’s driver libraries are not in the image. The image cannot simply bring its own copy of the driver libraries, either, because they have to match the kernel driver running on this particular board exactly.

So something has to add the device files and the host’s libraries to each GPU container, and it has to do so at the moment the container is created. To see where that can happen, we need to know how containers are created.

When Kubernetes starts a container, the kubelet (the Kubernetes agent on each machine) asks the machine’s container engine, containerd in our case, to create it. containerd prepares the container’s filesystem and writes a JSON document called the OCI runtime specification, which lists everything about the container: its mounts, its device files, its environment variables, and the command to run. containerd then hands that document to a low-level runtime, usually a program called runc, which creates the container exactly as the document says.

The NVIDIA Container Runtime slots into that hand-off. NVIDIA’s documentation describes it as “a shim for OCI-compliant low-level runtimes such as runc”. In practice, containerd is configured to call nvidia-container-runtime instead of runc. The NVIDIA runtime reads the specification, adds the GPU’s device files and mounts for the driver libraries, and then calls runc with the edited document. The container starts with the GPU already in place.

How the runtime knows what to add

On a data center machine with PCIe GPUs, the runtime asks the driver which GPUs exist and which files each one needs. On a Jetson it reads that information from plain text files instead. NVIDIA’s runtime documentation calls this CSV mode: “CSV files at /etc/nvidia-container-runtime/host-files-for-container.d define the devices and mounts that are to be injected”, and the mode “is primarily targeted at Tegra-based systems without NVML available”. NVML is the management library that data center drivers provide; Tegra is the family name of NVIDIA’s system-on-chip products, including the Orin in our Jetson.

Each line of a CSV file gives a type and a path. Here are the lines for libcuda:

$ grep -h "libcuda" /etc/nvidia-container-runtime/host-files-for-container.d/drivers.csv
lib, /usr/lib/aarch64-linux-gnu/nvidia/libcuda.so.1.1
sym, /usr/lib/aarch64-linux-gnu/libcuda.so
sym, /usr/lib/aarch64-linux-gnu/nvidia/libcuda.so
sym, /usr/lib/aarch64-linux-gnu/nvidia/libcuda.so.1
$ wc -l /etc/nvidia-container-runtime/host-files-for-container.d/*.csv
   40 /etc/nvidia-container-runtime/host-files-for-container.d/devices.csv
  285 /etc/nvidia-container-runtime/host-files-for-container.d/drivers.csv
  325 total

The type tells the runtime what to do with the path:

• A dev line names a device file to add to the container.
• A lib line names a file to bind-mount from the host: the same file appears inside the container at the same path, without being copied.
• A sym line names a symbolic link. The runtime recreates the link inside the container, pointing at the same target as on the host.

The libcuda lines show why the last type exists. The real library is libcuda.so.1.1, and programs look for it under other names, such as libcuda.so.1, which are links to the real file. The runtime mounts the real file once and recreates the links around it.

The lists contain the driver and nothing else. The CUDA toolkit, libcudart and the rest of /usr/local/cuda, is not in them, so a container has to bring its own CUDA runtime library, either in its image or linked into its program. We will do the latter for our test program.

One question remains: how does the runtime know which containers should get a GPU? It looks for an environment variable, NVIDIA_VISIBLE_DEVICES. If a container’s specification sets it, the runtime adds the GPU. If not, the runtime passes the specification to runc untouched, and the container starts like any other. The next piece of the puzzle is who sets that variable.

How a pod asks for a GPU

Kubernetes does not know what a GPU is. It knows about CPU and memory, and beyond those it supports extended resources: named quantities, such as nvidia.com/gpu, that a machine advertises and pods request. The scheduler, the component that chooses a machine for each new pod, counts extended resources the same way it counts CPU and memory, without knowing what they are. Extended resources come in whole units only, and a machine never hands out more units than it advertises.

A machine advertises an extended resource through a device plugin. A device plugin is a small program, usually running as a pod on the machine, that speaks a protocol defined by Kubernetes to the kubelet. The conversation goes like this:

1. The plugin starts a gRPC server on a Unix socket in /var/lib/kubelet/device-plugins. It then connects to the kubelet’s own socket in the same directory, kubelet.sock, and registers: “I serve the resource nvidia.com/gpu, and you can reach me at this socket.”
2. The kubelet calls the plugin’s ListAndWatch method. The plugin answers with a list of device IDs and their health, and keeps the call open to report changes, such as a GPU becoming unhealthy. The kubelet adds the number of healthy devices to the machine’s advertised capacity, and from that moment the scheduler sees, for example, nvidia.com/gpu: 1 on our node.
3. When the scheduler places a pod that asks for nvidia.com/gpu: 1 on the machine, the kubelet chooses a free device ID and calls the plugin’s Allocate method with it. The plugin answers with what the container needs in order to use that device: environment variables, device files, or mounts. The kubelet adds them to the container’s configuration before asking containerd to create it.
4. The kubelet writes every assignment of a device to a container into a checkpoint file, kubelet_internal_checkpoint in the same directory, so that it still knows which devices are in use after it restarts.

NVIDIA’s device plugin can answer Allocate in several ways. Its README lists the options for the setting DEVICE_LIST_STRATEGY and describes the default: “envvar (default): the NVIDIA_VISIBLE_DEVICES environment variable … is used to select the devices that are to be injected by the NVIDIA Container Runtime.” That is the missing link: the plugin sets the variable, and the runtime acts on it.

Putting the whole path together:

pod: limits: nvidia.com/gpu: 1
 |
 |  scheduler: places the pod on a machine that advertises a free nvidia.com/gpu
 v
kubelet ---- Allocate(device) ----> NVIDIA device plugin
 |       <--- NVIDIA_VISIBLE_DEVICES=<device> ---
 v
containerd: writes the OCI specification, with that variable in the environment
 |
 v
nvidia-container-runtime: sees the variable, reads the CSV lists,
 |                        adds the device files and library mounts to the
 |                        specification, and registers hooks that recreate
 |                        the symlinks and refresh the library cache
 v
runc: creates the container

Figure A.1 draws the same path, with the CSV lists that the runtime reads beside it.

For this path to work, containerd must call the NVIDIA runtime for every container, since it cannot know in advance which ones will carry the variable. The plugin’s README says so directly: it requires “nvidia-container-runtime configured as the default low-level runtime” on each GPU machine. It also mentions an alternative, a RuntimeClass that names the nvidia runtime, but then every GPU pod would have to name that RuntimeClass too.

Step 1: create the cluster

Now we can build our cluster. A kind cluster with a single node is enough, and the node image kindest/node:v1.37.0 gives us Kubernetes 1.37:

$ kind create cluster --name gpu-lab --image kindest/node:v1.37.0
Creating cluster "gpu-lab" ...
...
Set kubectl context to "kind-gpu-lab"
$ kubectl get nodes
NAME                    STATUS   ROLES           AGE   VERSION
gpu-lab-control-plane   Ready    control-plane   22s   v1.37.0

kind names the node’s container after the cluster, gpu-lab-control-plane. In the steps below, we use that name with docker exec to run commands inside the node.

What our kind node is missing

A kind node is a Docker container that runs systemd, containerd, and the kubelet, much like a small virtual machine. kind starts it with Docker’s --privileged option, and a privileged container sees all of the host’s device files. So the GPU’s device files are already there:

$ docker exec gpu-lab-control-plane ls -l /dev/nvhost-gpu /dev/nvmap
crw-rw---- 1 root video 487,   1 Sep 26 13:59 /dev/nvhost-gpu
crw-rw---- 1 root video  10, 124 Sep 26 13:59 /dev/nvmap

Everything else on our path is missing. The node’s files come from the kindest/node image, a Debian system that has none of the Jetson’s driver libraries, none of the CSV lists, and no NVIDIA Container Toolkit. Its containerd calls plain runc. And no device plugin runs on it, so it advertises no nvidia.com/gpu, and every pod that asks for one waits forever.

On a data center cluster, NVIDIA’s GPU Operator would install all of these pieces for us. It cannot help here: its platform support page states that “NVIDIA Jetson, or other embedded products with integrated GPUs, are not supported.” So we will install the same pieces by hand.

Step 2: give the node the driver and the runtime

Our script, enable-gpu-in-kind.sh, copies the files the node lacks and then configures containerd:

Listing A.1 Giving the kind node the driver and the runtime. code/jetson-gpu/enable-gpu-in-kind.sh

#!/usr/bin/env bash
# Makes the Jetson's integrated GPU usable from pods in a kind cluster.
#
# A kind node is a privileged container, so it already sees the GPU's device
# files. It lacks three things that the host has: the driver libraries listed
# in the Tegra CSV files, the NVIDIA Container Toolkit binaries, and a
# containerd configuration that uses nvidia-container-runtime. This script
# copies the first two into every node and then uses nvidia-ctk, inside the
# node, to configure the runtime as NVIDIA's documentation describes.
#
# Usage: ./enable-gpu-in-kind.sh [cluster-name]
set -euo pipefail

cluster="${1:-gpu-lab}"
csv_dir=/etc/nvidia-container-runtime/host-files-for-container.d
toolkit=(/usr/bin/nvidia-container-runtime /usr/bin/nvidia-container-runtime-hook
         /usr/bin/nvidia-ctk /usr/bin/nvidia-cdi-hook
         /etc/nvidia-container-runtime/config.toml)

# Every file named in the CSV lists that exists on this board. Device nodes
# are left out, because the privileged node already has them.
mapfile -t host_files < <(cut -d, -f2 "$csv_dir"/*.csv | tr -d ' ' | sort -u |
    grep -v '^/dev/' | while read -r f; do [[ -e $f ]] && echo "$f"; done)

for node in $(kind get nodes --name "$cluster"); do
    echo "$node: copying ${#host_files[@]} driver files and the container toolkit"
    # The runtime bind-mounts the CSV "lib" entries and recreates the "sym"
    # entries as links, so symlinks must arrive as symlinks. Paths that are
    # links on the host are removed first, in case an earlier copy left a
    # regular file or directory there.
    for f in "${host_files[@]}"; do [[ -L $f ]] && echo "$f"; done |
        docker exec -i "$node" xargs rm -rf
    tar -cf - --hard-dereference "${host_files[@]}" "$csv_dir" "${toolkit[@]}" 2>/dev/null |
        docker exec -i "$node" tar -xf - -C /

    # The host runs Ubuntu, where the real ldconfig is /sbin/ldconfig.real.
    # The kind node runs Debian, where it is /sbin/ldconfig.
    docker exec "$node" nvidia-ctk config --in-place \
        --set nvidia-container-cli.ldconfig=@/sbin/ldconfig

    if docker exec "$node" grep -q 'default_runtime_name = "nvidia"' /etc/containerd/config.toml; then
        echo "$node: containerd already uses the nvidia runtime by default"
        continue
    fi
    echo "$node: configuring containerd and restarting it"
    docker exec "$node" nvidia-ctk runtime configure --runtime=containerd --set-as-default
    docker exec "$node" systemctl restart containerd
done

The script builds its list of files from the CSV lists themselves, so it copies exactly what the runtime will later try to mount. It skips the dev lines, because the device files are already in the node, and it skips paths that do not exist on this particular board.

Two details of the copy matter, and we ran into both while building this setup. First, sym entries must arrive as symbolic links. Suppose a link were copied as a regular file. The runtime would bind-mount that file into the container, and then its hook would try to replace the mounted file with a link and fail with “device or resource busy”. So the script copies links as links, and before copying it deletes any path that is a link on the host, in case an earlier copy left a regular file or a directory there. Second, some firmware files are hard links, two names for the same file, and --hard-dereference makes tar store each name as a separate file so that both arrive intact.

The host’s configuration file for the runtime, config.toml, travels along with the binaries. One setting in it names the ldconfig program, which the runtime runs inside each GPU container to refresh the list of libraries the container can load. On the host, which runs Ubuntu, that program is /sbin/ldconfig.real. On the kind node, which runs Debian, it is /sbin/ldconfig. The script corrects the path with nvidia-ctk config, the toolkit’s own tool for editing that file.

Finally, the script configures containerd with the command from NVIDIA’s installation guide, adding --set-as-default because the device plugin requires the NVIDIA runtime to be the default. Then it restarts containerd. You might expect that to stop every running pod, but it does not: each container is supervised by its own small shim process, which keeps running while containerd restarts and reconnects to it afterwards. The check at the top of that block skips the whole step when containerd is already configured, so the script is safe to run again.

We run it with the name of our cluster:

$ cd code/jetson-gpu
$ ./enable-gpu-in-kind.sh gpu-lab
gpu-lab-control-plane: copying 283 driver files and the container toolkit
gpu-lab-control-plane: configuring containerd and restarting it
time="2026-09-26T14:00:01Z" level=info msg="Wrote updated config to /etc/containerd/config.toml"
time="2026-09-26T14:00:01Z" level=info msg="It is recommended that containerd daemon be restarted."

Version 1.17.3 of the toolkit edits config.toml directly. Newer versions write a separate file, /etc/containerd/conf.d/99-nvidia.toml, and add it to the imports list in config.toml, as NVIDIA’s installation guide describes. The result is the same.

Let us check our work. A sym entry should be a link inside the node, the NVIDIA runtime should be registered and set as the default, and the node should still be Ready:

$ docker exec gpu-lab-control-plane ls -l /usr/lib/aarch64-linux-gnu/nvidia/libcuda.so
lrwxrwxrwx 1 root root 14 Jan  8  2025 /usr/lib/aarch64-linux-gnu/nvidia/libcuda.so -> libcuda.so.1.1
$ docker exec gpu-lab-control-plane grep -A5 'runtimes.nvidia\]' /etc/containerd/config.toml
        [plugins."io.containerd.grpc.v1.cri".containerd.runtimes.nvidia]
          base_runtime_spec = "/etc/containerd/cri-base.json"
          runtime_type = "io.containerd.runc.v2"

          [plugins."io.containerd.grpc.v1.cri".containerd.runtimes.nvidia.options]
            BinaryName = "/usr/bin/nvidia-container-runtime"
$ docker exec gpu-lab-control-plane grep default_runtime_name /etc/containerd/config.toml
      default_runtime_name = "nvidia"
$ kubectl get nodes
NAME                    STATUS   ROLES           AGE   VERSION
gpu-lab-control-plane   Ready    control-plane   30s   v1.37.0

Step 3: run the device plugin

The device plugin runs as a DaemonSet, a Kubernetes workload that runs one pod on every node, which is exactly what a per-machine agent needs:

Listing A.2 The device plugin. code/jetson-gpu/device-plugin.yaml

# NVIDIA device plugin configured for a Jetson. The integrated GPU is found
# from the Tegra CSV file lists under the driver root. For each pod that
# requests nvidia.com/gpu, the plugin sets NVIDIA_VISIBLE_DEVICES, and the
# node's nvidia-container-runtime mounts the devices and libraries.
apiVersion: apps/v1
kind: DaemonSet
metadata:
  name: nvidia-device-plugin
  namespace: kube-system
spec:
  selector:
    matchLabels:
      app: nvidia-device-plugin
  template:
    metadata:
      labels:
        app: nvidia-device-plugin
    spec:
      priorityClassName: system-node-critical
      tolerations:
        - key: nvidia.com/gpu
          operator: Exists
          effect: NoSchedule
      containers:
        - name: plugin
          image: nvcr.io/nvidia/k8s-device-plugin:v0.20.1
          env:
            - name: DEVICE_DISCOVERY_STRATEGY
              value: tegra
            - name: NVIDIA_DRIVER_ROOT      # driver root as the node sees it
              value: /
            - name: DRIVER_ROOT_CTR_PATH    # the same root inside this container
              value: /driver-root
            - name: FAIL_ON_INIT_ERROR
              value: "true"
          volumeMounts:
            - name: device-plugins
              mountPath: /var/lib/kubelet/device-plugins
            - name: driver-root
              mountPath: /driver-root
              readOnly: true
      volumes:
        - name: device-plugins
          hostPath:
            path: /var/lib/kubelet/device-plugins
        - name: driver-root
          hostPath:
            path: /

The plugin needs to read the same CSV lists and driver files as the runtime, so it mounts the node’s whole root filesystem, read-only, at /driver-root inside its own container. The two driver-root settings tell it where to find the driver from both points of view: NVIDIA_DRIVER_ROOT is where the driver is on the node, and DRIVER_ROOT_CTR_PATH is where that same root appears inside the plugin’s container.

DEVICE_DISCOVERY_STRATEGY=tegra tells the plugin to discover the GPU from the CSV lists. This setting is not in the plugin’s README; it is defined among the plugin’s command-line flags in its source code (cmd/nvidia-device-plugin/main.go), which accepts auto, nvml, or tegra. The default, auto, inspects the system and would most likely choose the same, but we prefer to be explicit. We leave DEVICE_LIST_STRATEGY at its default, envvar, the mechanism we traced above.

The plugin also offers newer strategies based on CDI, the Container Device Interface, a standard file format that describes devices to container engines. They do not work on this board: with DEVICE_LIST_STRATEGY=cdi-cri, version 0.20.1 stops at startup with CDI --device-list-strategy options are only supported on NVML-based systems.

We apply the DaemonSet and read the plugin’s log:

$ kubectl apply -f device-plugin.yaml
$ kubectl -n kube-system rollout status ds/nvidia-device-plugin
daemon set "nvidia-device-plugin" successfully rolled out
$ kubectl -n kube-system logs ds/nvidia-device-plugin | tail -6
I0926 14:00:19.469889  1 main.go:391] Retrieving plugins.
I0926 14:00:19.469935  1 plugin-manager.go:96] Using requested device discovery strategy: tegra
I0926 14:00:19.469986  1 plugin-manager.go:50] Using device discovery strategy: tegra
I0926 14:00:19.470717  1 server.go:198] Starting GRPC server for 'nvidia.com/gpu'
I0926 14:00:19.472172  1 server.go:142] Starting to serve 'nvidia.com/gpu' on /var/lib/kubelet/device-plugins/nvidia-gpu.sock
I0926 14:00:19.477135  1 server.go:149] Registered device plugin for 'nvidia.com/gpu' with Kubelet
$ kubectl get node -o jsonpath='{.items[0].status.allocatable.nvidia\.com/gpu}'; echo
1

The last three log lines are the first step of the conversation we described earlier: the plugin starts its gRPC server, listens on nvidia-gpu.sock, and registers with the kubelet. The node now advertises one GPU. If you check within a few seconds of the plugin starting, you may still see 0, because the kubelet updates the node’s status shortly after the plugin registers.

Step 4: run a CUDA program in a pod

To prove the whole path works, we need a program that uses the GPU and tells us so. Ours adds two vectors of a million elements on the GPU and prints the device it ran on. If no GPU is visible, or the result is wrong, it exits with a non-zero status, so Kubernetes marks the pod as failed.

Listing A.3 The kernel: one addition per thread. code/jetson-gpu/gpu_check.cu

// add writes a[i] + b[i] to out[i]. Each thread handles one element, and
// threads past the end of the vectors do nothing.
__global__ void add(const float *a, const float *b, float *out, int n) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        out[i] = a[i] + b[i];
    }
}

Listing A.4 Allocating, launching, and checking the result. code/jetson-gpu/gpu_check.cu

// main prints the device, runs the addition, and prints the kernel's status
// with the first and last results. It exits with status 1 if no GPU is
// visible, the kernel failed, or the result is wrong.
int main() {
    cudaDeviceProp prop;
    if (cudaGetDeviceProperties(&prop, 0) != cudaSuccess) {
        std::printf("no CUDA device visible\n");
        return 1;
    }
    std::printf("device: %s, compute %d.%d, %d SMs\n", prop.name, prop.major, prop.minor,
                prop.multiProcessorCount);

    float *a, *b, *out;
    cudaMallocManaged(&a, kElements * sizeof(float));
    cudaMallocManaged(&b, kElements * sizeof(float));
    cudaMallocManaged(&out, kElements * sizeof(float));
    for (int i = 0; i < kElements; i++) {
        a[i] = 1.0f;
        b[i] = 2.0f;
    }

    add<<<(kElements + kThreads - 1) / kThreads, kThreads>>>(a, b, out, kElements);
    cudaError_t err = cudaDeviceSynchronize();
    std::printf("kernel: %s, out[0] = %.1f, out[n-1] = %.1f\n", cudaGetErrorString(err),
                out[0], out[kElements - 1]);

    bool ok = err == cudaSuccess && out[kElements - 1] == 3.0f;
    cudaFree(a);
    cudaFree(b);
    cudaFree(out);
    return ok ? 0 : 1;
}

We build it on the board with the CUDA compiler from JetPack. -arch=sm_87 targets the Orin’s GPU, whose compute capability, NVIDIA’s version number for a GPU’s feature set, is 8.7. -cudart static links the CUDA runtime library into the program itself, so our image does not need to contain it. The only library the program needs at run time is the driver library, libcuda, and that is one of the files the NVIDIA runtime mounts.

$ /usr/local/cuda/bin/nvcc -O2 -arch=sm_87 -cudart static -o gpu_check gpu_check.cu
$ ./gpu_check
device: Orin, compute 8.7, 8 SMs
kernel: no error, out[0] = 3.0, out[n-1] = 3.0

The image only has to hold the program and the C++ standard library that Ubuntu provides:

Listing A.5 The image for the test program. code/jetson-gpu/Dockerfile

FROM ubuntu:22.04
COPY gpu_check /usr/local/bin/gpu_check
CMD ["gpu_check"]

A kind node has its own image store and cannot see the images in the host’s Docker, so we build the image and then copy it into the node:

$ docker build -t gpu-check:1 .
$ kind load docker-image gpu-check:1 --name gpu-lab

We run the image in two pods. The first asks for a GPU and the second does not, so the second one tells us whether containers that did not ask for a GPU really stay without one.

Listing A.6 A pod that runs `gpu_check` with one GPU. code/jetson-gpu/gpu-check-pod.yaml

apiVersion: v1
kind: Pod
metadata:
  name: gpu-check
spec:
  restartPolicy: Never
  containers:
    - name: check
      image: gpu-check:1
      imagePullPolicy: Never
      resources:
        limits:
          nvidia.com/gpu: 1

Listing A.7 The same program in a pod that asks for no GPU. code/jetson-gpu/no-gpu-pod.yaml

apiVersion: v1
kind: Pod
metadata:
  name: no-gpu
spec:
  restartPolicy: Never
  containers:
    - name: check
      image: gpu-check:1
      imagePullPolicy: Never

imagePullPolicy: Never tells the kubelet to use the image that kind load placed on the node and not to look for it in a registry.

$ kubectl apply -f gpu-check-pod.yaml -f no-gpu-pod.yaml
$ kubectl get pods gpu-check no-gpu
NAME        READY   STATUS      RESTARTS   AGE
gpu-check   0/1     Completed   0          15s
no-gpu      0/1     Error       0          15s
$ kubectl logs gpu-check
device: Orin, compute 8.7, 8 SMs
kernel: no error, out[0] = 3.0, out[n-1] = 3.0
$ kubectl logs no-gpu
no CUDA device visible
$ kubectl delete -f gpu-check-pod.yaml -f no-gpu-pod.yaml

The first pod ran our kernel on the GPU. The second ran the same image, through the same default runtime, and saw no GPU, because the device plugin was never asked to allocate one for it and so never set NVIDIA_VISIBLE_DEVICES.

We can also look inside a GPU container while it runs, to see every piece of the path at work. This pod asks for a GPU and then just sleeps for five minutes, which gives us time to look around inside it:

Listing A.8 A pod that sets the variable itself. code/jetson-gpu/gpu-env-pod.yaml

apiVersion: v1
kind: Pod
metadata:
  name: gpu-env
spec:
  restartPolicy: Never
  containers:
    - name: check
      image: gpu-check:1
      imagePullPolicy: Never
      command: ["sleep", "300"]
      resources:
        limits:
          nvidia.com/gpu: 1

We start it, wait until it is running, and look at its environment, its device files, and its driver libraries:

$ kubectl apply -f gpu-env-pod.yaml
$ kubectl wait --for=condition=Ready pod/gpu-env
$ kubectl exec gpu-env -- sh -c 'env | grep NVIDIA; ls /dev | grep -E "nvhost-gpu|nvmap"'
NVIDIA_VISIBLE_DEVICES=tegra
nvhost-gpu
nvmap
$ kubectl exec gpu-env -- sh -c 'ls /usr/lib/aarch64-linux-gnu/nvidia/libcuda.so*'
/usr/lib/aarch64-linux-gnu/nvidia/libcuda.so
/usr/lib/aarch64-linux-gnu/nvidia/libcuda.so.1
/usr/lib/aarch64-linux-gnu/nvidia/libcuda.so.1.1

The device plugin advertises our single GPU under the device ID tegra, and its answer to Allocate set NVIDIA_VISIBLE_DEVICES=tegra in the container. The NVIDIA runtime saw the variable and added the device files, the real library libcuda.so.1.1, and the two links to it. The kubelet’s checkpoint file on the node shows the same allocation from the kubelet’s side:

$ docker exec gpu-lab-control-plane cat /var/lib/kubelet/device-plugins/kubelet_internal_checkpoint
{"Data":{"PodDeviceEntries":[{"PodUID":"bf52b2cd-b611-4e79-86b1-280b283185c3","ContainerName":"check","ResourceName":"nvidia.com/gpu","DeviceIDs":{"-1":["tegra"]},"AllocResp":"Ch8KFk5WSURJQV9WSVNJQkxFX0RFVklDRVMSBXRlZ3Jh"}],"RegisteredDevices":{"nvidia.com/gpu":["tegra"]}},"Checksum":3660592824}

RegisteredDevices is the device list the plugin reported through ListAndWatch, and the entry under PodDeviceEntries records that the container check of our pod holds the device tegra. AllocResp is the plugin’s answer to Allocate, stored in its binary Protocol Buffers encoding; decoded, it contains the environment variable NVIDIA_VISIBLE_DEVICES=tegra. When we are done looking, we delete the pod so that it releases the GPU:

$ kubectl delete -f gpu-env-pod.yaml

Because the plugin advertises the board’s one GPU as one unit, only one pod that asks for a GPU can run at a time. Any other GPU pod waits in Pending until the first one finishes, and section 1.1 shows a job in exactly that state.

When a step fails

We ran into each of these problems while building this setup.

The plugin exits with CDI --device-list-strategy options are only supported on NVML-based systems. DEVICE_LIST_STRATEGY is set to one of the CDI strategies. Remove the setting, so that the plugin uses envvar.

A GPU pod fails with failed to create symlink ... device or resource busy. A sym entry reached the node as a regular file, so the runtime mounted it and its hook could not replace it with a link. Run enable-gpu-in-kind.sh again; it removes such paths before copying.

A pod that asked for a GPU prints no CUDA device visible. The container was not created through the NVIDIA runtime. Check default_runtime_name in the node’s containerd configuration, and make sure containerd was restarted after the change.

A GPU pod stays Pending with Insufficient nvidia.com/gpu. Either the plugin has not registered, which its log and the node’s allocatable resources will show, or another pod is holding the GPU.

GPU pods fail after the cluster was deleted and created again. A new node starts from the kindest/node image, without any of our changes. Run the script and apply the DaemonSet again. Restarting the node’s container, or Docker itself, keeps the changes, since they live in the container’s filesystem.

The complete test program

Listing A.9 `gpu_check`, complete. code/jetson-gpu/gpu_check.cu

// gpu_check proves that a container can use the GPU: it prints the device it
// sees, adds two vectors of a million elements on the GPU, and checks the
// result.
//
// How it runs, end to end:
//   1. nvcc compiles this file on the host for the Orin's GPU (-arch=sm_87) and
//      links the CUDA runtime library statically (-cudart static), so the
//      binary needs only the CUDA driver library, libcuda, at run time.
//   2. The binary is copied into the gpu-check:1 image and run in a pod.
//   3. If the pod requested nvidia.com/gpu, the device plugin sets
//      NVIDIA_VISIBLE_DEVICES and the NVIDIA container runtime mounts the GPU's
//      device files and libcuda into the container. If it did not, the
//      container has neither, and the first CUDA call fails.
//   4. The vectors live in managed memory, which the CPU and the GPU can both
//      address, so the program fills them on the CPU, adds them on the GPU,
//      and reads the result on the CPU without explicit copies.
//
// Build: nvcc -O2 -arch=sm_87 -cudart static -o gpu_check gpu_check.cu
#include <cstdio>
#include <cuda_runtime.h>

// kElements is the length of each vector.
constexpr int kElements = 1 << 20;

// kThreads is the number of threads in each block of the add kernel.
constexpr int kThreads = 256;

// add writes a[i] + b[i] to out[i]. Each thread handles one element, and
// threads past the end of the vectors do nothing.
__global__ void add(const float *a, const float *b, float *out, int n) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        out[i] = a[i] + b[i];
    }
}

// main prints the device, runs the addition, and prints the kernel's status
// with the first and last results. It exits with status 1 if no GPU is
// visible, the kernel failed, or the result is wrong.
int main() {
    cudaDeviceProp prop;
    if (cudaGetDeviceProperties(&prop, 0) != cudaSuccess) {
        std::printf("no CUDA device visible\n");
        return 1;
    }
    std::printf("device: %s, compute %d.%d, %d SMs\n", prop.name, prop.major, prop.minor,
                prop.multiProcessorCount);

    float *a, *b, *out;
    cudaMallocManaged(&a, kElements * sizeof(float));
    cudaMallocManaged(&b, kElements * sizeof(float));
    cudaMallocManaged(&out, kElements * sizeof(float));
    for (int i = 0; i < kElements; i++) {
        a[i] = 1.0f;
        b[i] = 2.0f;
    }

    add<<<(kElements + kThreads - 1) / kThreads, kThreads>>>(a, b, out, kElements);
    cudaError_t err = cudaDeviceSynchronize();
    std::printf("kernel: %s, out[0] = %.1f, out[n-1] = %.1f\n", cudaGetErrorString(err),
                out[0], out[kElements - 1]);

    bool ok = err == cudaSuccess && out[kElements - 1] == 3.0f;
    cudaFree(a);
    cudaFree(b);
    cudaFree(out);
    return ok ? 0 : 1;
}

References

• NVIDIA Container Runtime, modes and configuration: https://github.com/NVIDIA/nvidia-container-toolkit/blob/main/cmd/nvidia-container-runtime/README.md
• NVIDIA Container Toolkit, architecture overview: https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/arch-overview.html
• NVIDIA Container Toolkit, installation guide, “Configuring containerd (for Kubernetes)”: https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html
• NVIDIA device plugin for Kubernetes, README and source for v0.20.1: https://github.com/NVIDIA/k8s-device-plugin/tree/v0.20.1
• NVIDIA GPU Operator, platform support: https://docs.nvidia.com/datacenter/cloud-native/gpu-operator/latest/platform-support.html
• JetPack 6.2 release notes: https://docs.nvidia.com/jetson/jetpack/6.2/release-notes/index.html
• NVIDIA Jetson Linux Developer Guide, R36.4.3: https://docs.nvidia.com/jetson/archives/r36.4.3/DeveloperGuide/index.html
• Kubernetes, device plugins: https://kubernetes.io/docs/concepts/extend-kubernetes/compute-storage-net/device-plugins/
• kind, configuration: https://kind.sigs.k8s.io/docs/user/configuration/