llmgw/internal/agent/discovery/perftest.go

package discovery

import (
	"context"
	"errors"
	"fmt"
	"os/exec"
	"regexp"
	"strconv"
	"strings"
	"time"
)

// WriteBW is one side of an ib_write_bw run, from perftest: an RDMA write
// bandwidth test between two adapters, which is what a collective spanning
// nodes rides on.
type WriteBW struct {
	HCA  string
	Port int
	// GIDIndex is the RoCE v2 GID to use; negative on InfiniBand, where the
	// port's LID is used.
	GIDIndex int
	// Peer is the server's address, for the client; empty makes the server.
	Peer    string
	TCPPort int
	Seconds int
}

// Args are the command line: a duration rather than a count, so the test
// takes as long on a slow link as a fast one, and results in Gb/s.
func (w WriteBW) Args() []string {
	args := []string{"-d", w.HCA, "-i", strconv.Itoa(max(w.Port, 1)),
		"-p", strconv.Itoa(w.TCPPort), "-D", strconv.Itoa(max(w.Seconds, 1)),
		"-F", "--report_gbits"}
	if w.GIDIndex >= 0 {
		args = append(args, "-x", strconv.Itoa(w.GIDIndex))
	}
	if w.Peer != "" {
		args = append(args, w.Peer)
	}
	return args
}

// RunWriteBW runs one side and returns the average bandwidth the output
// reports, in Gb/s, and the output. The client retries its connection for a
// while: the server is started first, but a few seconds may pass before it
// listens.
func RunWriteBW(ctx context.Context, w WriteBW) (float64, string, error) {
	bin, err := exec.LookPath("ib_write_bw")
	if err != nil {
		return 0, "", errors.New("ib_write_bw is not installed (the perftest package)")
	}
	limit := time.Duration(w.Seconds+60) * time.Second
	var out []byte
	for attempt := 0; ; attempt++ {
		cctx, cancel := context.WithTimeout(ctx, limit)
		out, err = exec.CommandContext(cctx, bin, w.Args()...).CombinedOutput() //nolint:gosec // fixed tool
		cancel()
		if err == nil || w.Peer == "" || !refused(string(out)) || attempt >= 10 {
			break
		}
		select {
		case <-ctx.Done():
			return 0, string(out), ctx.Err()
		case <-time.After(time.Second):
		}
	}
	text := string(out)
	if err != nil {
		return 0, text, fmt.Errorf("ib_write_bw: %v: %s", err, lastNonEmpty(text))
	}
	gbps, perr := ParseWriteBW(text)
	if perr != nil && w.Peer != "" {
		// Only the client is sure to print the result; a server that does not
		// is not a failure.
		return 0, text, perr
	}
	return gbps, text, nil
}

func refused(out string) bool {
	for _, s := range []string{"Couldn't connect", "Unable to init the socket connection", "Connection refused"} {
		if strings.Contains(out, s) {
			return true
		}
	}
	return false
}

var columnSplit = regexp.MustCompile(`\s{2,}`)

// ParseWriteBW finds the average bandwidth in perftest's result table:
//
//	#bytes     #iterations    BW peak[Gb/sec]    BW average[Gb/sec]   MsgRate[Mpps]
//	65536      887036           0.00               96.93               0.184877
//
// by the heading's position rather than a fixed column, since versions add
// columns and the peak reads 0.00 when the test runs for a duration.
func ParseWriteBW(out string) (float64, error) {
	lines := strings.Split(out, "\n")
	for i, l := range lines {
		if !strings.Contains(l, "BW average") {
			continue
		}
		heads := columnSplit.Split(strings.TrimSpace(l), -1)
		col := -1
		for j, h := range heads {
			if strings.HasPrefix(h, "BW average") {
				col = j
			}
		}
		for _, d := range lines[i+1:] {
			f := strings.Fields(d)
			if len(f) == 0 || strings.HasPrefix(f[0], "-") {
				continue
			}
			if len(f) != len(heads) || col < 0 {
				return 0, fmt.Errorf("perftest's result table has a shape this does not know: %q / %q", l, d)
			}
			v, err := strconv.ParseFloat(f[col], 64)
			if err != nil {
				return 0, fmt.Errorf("reading the average bandwidth %q: %w", f[col], err)
			}
			return v, nil
		}
	}
	return 0, errors.New("perftest printed no result")
}

func lastNonEmpty(s string) string {
	lines := strings.Split(strings.TrimSpace(s), "\n")
	return strings.TrimSpace(lines[len(lines)-1])
}