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