llmgw/internal/agent/discovery/membw.go

package discovery

import (
	"bufio"
	"os"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"time"
)

// How fast this machine reads its memory, which on a CPU is what decides how
// fast a model answers: every token of a reply reads all the weights it uses
// once, so a node that reads 40 GB/s writes a 4 GB model's tokens about ten
// times a second, whatever its core count. Measured once when the agent
// starts, in a fraction of a second, by every CPU summing a share of a buffer
// too large for any cache — a VM sharing its host's memory with others gets
// its share, which is the number that matters for it.

// bandwidthBuffer is what is read at most; the least is eight times the
// processor's last cache, and never under minBuffer. Smaller, it would be
// the cache that was measured — some hold a gigabyte — and the figure kept
// would be too high.
const (
	bandwidthBuffer = 2 << 30
	minBuffer       = 256 << 20
)

// Bandwidth is how fast this machine's memory reads, in GB/s: with every
// CPU reading, and with one alone. A replica given a few threads of a large
// machine reads at a few times the second, not at the first: on a server of
// many memory channels one core cannot keep more than a fraction of them
// busy.
type Bandwidth struct {
	All, Core float64
}

// coreBuffer is what the one-core passes read: past any cache, and short
// enough that one core reads it in a fraction of a second.
const coreBuffer = 512 << 20

// MeasureMemoryBandwidth is the memory's read rate, the best of a few passes
// each way; zero when the machine has too little free memory to measure.
func MeasureMemoryBandwidth() Bandwidth {
	size := int64(bandwidthBuffer)
	// Never more than a quarter of what is free: a node already holding
	// models must not be pushed into swap by a measurement. Without room
	// for a buffer past the cache, none is measured: the hourly retry
	// measures it when there is.
	if avail := memAvailable("/proc/meminfo"); avail > 0 {
		size = min(size, avail/4)
	}
	if size < max(minBuffer, 8*lastCacheBytes("/sys/devices/system/cpu")) {
		return Bandwidth{}
	}
	buf := make([]uint64, size/8)
	for i := range buf {
		buf[i] = uint64(i)
	}
	workers := runtime.NumCPU()
	chunk := (len(buf) + workers - 1) / workers
	best := 0.0
	for pass := 0; pass < 4; pass++ {
		var wg sync.WaitGroup
		start := time.Now()
		for w := 0; w < workers; w++ {
			lo, hi := w*chunk, min((w+1)*chunk, len(buf))
			if lo >= hi {
				continue
			}
			wg.Add(1)
			go func(part []uint64) {
				defer wg.Done()
				sink(sum(part))
			}(buf[lo:hi])
		}
		wg.Wait()
		if s := time.Since(start).Seconds(); s > 0 {
			best = max(best, float64(size)/s/1e9)
		}
	}
	// And one core alone, over a part of the same buffer.
	part := buf[:min(len(buf), coreBuffer/8)]
	core := 0.0
	for pass := 0; pass < 3; pass++ {
		start := time.Now()
		sink(sum(part))
		if s := time.Since(start).Seconds(); s > 0 {
			core = max(core, float64(len(part)*8)/s/1e9)
		}
	}
	return Bandwidth{All: best, Core: min(core, best)}
}

// sum reads every word, eight at a time so the loop is not what limits it.
func sum(p []uint64) uint64 {
	var a, b, c, d, e, f, g, h uint64
	i := 0
	for ; i+8 <= len(p); i += 8 {
		a += p[i]
		b += p[i+1]
		c += p[i+2]
		d += p[i+3]
		e += p[i+4]
		f += p[i+5]
		g += p[i+6]
		h += p[i+7]
	}
	for ; i < len(p); i++ {
		a += p[i]
	}
	return a + b + c + d + e + f + g + h
}

// sinkValue keeps the sums alive, so the reads are not optimised away.
var sinkValue struct {
	sync.Mutex
	v uint64
}

func sink(v uint64) {
	sinkValue.Lock()
	sinkValue.v += v
	sinkValue.Unlock()
}

// memAvailable is MemAvailable in bytes, 0 when unknown.
func memAvailable(path string) int64 {
	f, err := os.Open(path) //nolint:gosec // a fixed system path
	if err != nil {
		return 0
	}
	defer f.Close()
	sc := bufio.NewScanner(f)
	for sc.Scan() {
		fields := strings.Fields(sc.Text())
		if len(fields) >= 2 && fields[0] == "MemAvailable:" {
			kb, _ := strconv.ParseInt(fields[1], 10, 64)
			return kb << 10
		}
	}
	return 0
}

// lastCacheBytes is the largest cache any CPU reports, 0 when unknown.
func lastCacheBytes(root string) int64 {
	paths, _ := filepath.Glob(filepath.Join(root, "cpu*", "cache", "index*", "size"))
	var most int64
	for _, p := range paths {
		b, err := os.ReadFile(p) //nolint:gosec // under /sys
		if err != nil {
			continue
		}
		v := strings.TrimSpace(string(b))
		mult := int64(1)
		switch {
		case strings.HasSuffix(v, "K"):
			mult, v = 1<<10, strings.TrimSuffix(v, "K")
		case strings.HasSuffix(v, "M"):
			mult, v = 1<<20, strings.TrimSuffix(v, "M")
		}
		if n, err := strconv.ParseInt(v, 10, 64); err == nil {
			most = max(most, n*mult)
		}
	}
	return most
}