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