Files
bee/audit/cmd/bee/main.go
T

743 lines
22 KiB
Go

package main
import (
"context"
"errors"
"flag"
"fmt"
"io"
"log/slog"
"os"
"os/exec"
"runtime/debug"
"strconv"
"strings"
"bee/audit/internal/app"
"bee/audit/internal/platform"
"bee/audit/internal/runtimeenv"
"bee/audit/internal/webui"
)
var Version = "dev"
var BuildCommit = "unknown"
func buildLabel() string {
label := strings.TrimSpace(Version)
if label == "" {
return "dev"
}
return label
}
func main() {
os.Exit(run(os.Args[1:], os.Stdout, os.Stderr))
}
func run(args []string, stdout, stderr io.Writer) (exitCode int) {
slog.SetDefault(slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{
Level: slog.LevelInfo,
})))
defer func() {
if rec := recover(); rec != nil {
slog.Error("fatal panic",
"panic", fmt.Sprint(rec),
"stack", string(debug.Stack()),
)
exitCode = 1
}
}()
if len(args) == 0 {
printRootUsage(stderr)
return 2
}
switch args[0] {
case "help", "--help", "-h":
if len(args) > 1 {
return runHelp(args[1:], stdout, stderr)
}
printRootUsage(stdout)
return 0
case "audit":
return runAudit(args[1:], stdout, stderr)
case "export":
return runExport(args[1:], stdout, stderr)
case "preflight":
return runPreflight(args[1:], stdout, stderr)
case "install-to-ram":
return runInstallToRAM(args[1:], stdout, stderr)
case "support-bundle":
return runSupportBundle(args[1:], stdout, stderr)
case "web":
return runWeb(args[1:], stdout, stderr)
case "blackbox":
return runBlackbox(args[1:], stdout, stderr)
case "sat":
return runSAT(args[1:], stdout, stderr)
case "run":
return runScenario(args[1:], stdout, stderr)
case "scenario":
// "bee scenario run <arg>" is a longer alias for "bee run <arg>";
// strip the "run" verb (if present) and delegate to the same code.
rest := args[1:]
if len(rest) > 0 && rest[0] == "run" {
rest = rest[1:]
}
return runScenario(rest, stdout, stderr)
case "benchmark":
return runBenchmark(args[1:], stdout, stderr)
case "bee-worker":
return runBeeWorker(args[1:], stdout, stderr)
case "gpu-bandwidth-groups":
return runGPUBandwidthGroups(args[1:], stdout, stderr)
case "version", "--version", "-version":
if len(args) > 1 && args[1] == "--commit" {
fmt.Fprintln(stdout, BuildCommit)
return 0
}
fmt.Fprintln(stdout, Version)
return 0
default:
fmt.Fprintf(stderr, "bee: unknown command %q\n\n", args[0])
printRootUsage(stderr)
return 2
}
}
func printRootUsage(w io.Writer) {
fmt.Fprintln(w, `bee commands:
bee audit --runtime auto|local|livecd --output stdout|file:<path>
bee preflight --output stdout|file:<path>
bee install-to-ram
bee export --target <device>
bee support-bundle --output stdout|file:<path>
bee web --listen :80 [--audit-path `+app.DefaultAuditJSONPath+`]
bee blackbox --export-dir `+app.DefaultExportDir+` [--state-file `+app.DefaultBlackboxStatePath+`]
bee sat nvidia|memory|storage|cpu [--duration <seconds>]
bee run <file.json|name> (bare name is looked up as scenarios/<name>.json on removable media)
bee benchmark nvidia [--profile standard|stability|overnight]
bee bee-worker --export-dir `+app.DefaultExportDir+` --task-id TASK-001
bee gpu-bandwidth-groups [--stage <label>]
bee version
bee help [command]`)
}
func runHelp(args []string, stdout, stderr io.Writer) int {
switch args[0] {
case "audit":
return runAudit([]string{"--help"}, stdout, stdout)
case "export":
return runExport([]string{"--help"}, stdout, stdout)
case "preflight":
return runPreflight([]string{"--help"}, stdout, stdout)
case "install-to-ram":
return runInstallToRAM([]string{"--help"}, stdout, stdout)
case "support-bundle":
return runSupportBundle([]string{"--help"}, stdout, stdout)
case "web":
return runWeb([]string{"--help"}, stdout, stdout)
case "blackbox":
return runBlackbox([]string{"--help"}, stdout, stdout)
case "sat":
return runSAT([]string{"--help"}, stdout, stderr)
case "run", "scenario":
return runScenario([]string{"--help"}, stdout, stderr)
case "benchmark":
return runBenchmark([]string{"--help"}, stdout, stderr)
case "bee-worker":
return runBeeWorker([]string{"--help"}, stdout, stderr)
case "gpu-bandwidth-groups":
return runGPUBandwidthGroups([]string{"--help"}, stdout, stderr)
case "version":
fmt.Fprintln(stdout, "usage: bee version")
return 0
default:
fmt.Fprintf(stderr, "bee help: unknown command %q\n\n", args[0])
printRootUsage(stderr)
return 2
}
}
func runAudit(args []string, stdout, stderr io.Writer) int {
fs := flag.NewFlagSet("audit", flag.ContinueOnError)
fs.SetOutput(stderr)
output := fs.String("output", "stdout", "output destination: stdout or file:<path>")
runtimeFlag := fs.String("runtime", "auto", "runtime environment: auto, local, livecd")
showVersion := fs.Bool("version", false, "print version and exit")
fs.Usage = func() {
fmt.Fprintln(stderr, "usage: bee audit [--runtime auto|local|livecd] [--output stdout|file:<path>]")
fs.PrintDefaults()
}
if err := fs.Parse(args); err != nil {
if err == flag.ErrHelp {
return 0
}
return 2
}
if fs.NArg() != 0 {
fs.Usage()
return 2
}
if *showVersion {
fmt.Fprintln(stdout, Version)
return 0
}
runtimeInfo, err := runtimeenv.Detect(*runtimeFlag)
if err != nil {
slog.Error("resolve runtime", "err", err)
return 1
}
slog.Info("runtime resolved", "mode", runtimeInfo.Mode, "reason", runtimeInfo.Reason)
application := app.New(platform.New())
path, err := application.RunAudit(runtimeInfo.Mode, *output)
if err != nil {
slog.Error("run audit", "err", err)
return 1
}
if path != "stdout" {
slog.Info("audit output written", "path", path)
}
return 0
}
func runExport(args []string, stdout, stderr io.Writer) int {
fs := flag.NewFlagSet("export", flag.ContinueOnError)
fs.SetOutput(stderr)
targetDevice := fs.String("target", "", "removable device path, e.g. /dev/sdb1")
fs.Usage = func() {
fmt.Fprintln(stderr, "usage: bee export --target <device>")
fs.PrintDefaults()
}
if err := fs.Parse(args); err != nil {
if err == flag.ErrHelp {
return 0
}
return 2
}
if fs.NArg() != 0 {
fs.Usage()
return 2
}
if strings.TrimSpace(*targetDevice) == "" {
fmt.Fprintln(stderr, "bee export: --target is required")
fs.Usage()
return 2
}
application := app.New(platform.New())
targets, err := application.ListRemovableTargets()
if err != nil {
slog.Error("list removable targets", "err", err)
return 1
}
for _, target := range targets {
if target.Device == *targetDevice {
path, err := application.ExportLatestAudit(target)
if err != nil {
slog.Error("export latest audit", "err", err)
return 1
}
slog.Info("audit exported", "path", path)
return 0
}
}
slog.Error("target device not found among removable filesystems", "device", *targetDevice)
return 1
}
func runPreflight(args []string, stdout, stderr io.Writer) int {
fs := flag.NewFlagSet("preflight", flag.ContinueOnError)
fs.SetOutput(stderr)
output := fs.String("output", "stdout", "output destination: stdout or file:<path>")
fs.Usage = func() {
fmt.Fprintf(stderr, "usage: bee preflight [--output stdout|file:%s]\n", app.DefaultRuntimeJSONPath)
fs.PrintDefaults()
}
if err := fs.Parse(args); err != nil {
if err == flag.ErrHelp {
return 0
}
return 2
}
if fs.NArg() != 0 {
fs.Usage()
return 2
}
application := app.New(platform.New())
path, err := application.RunRuntimePreflight(*output)
if err != nil {
slog.Error("run preflight", "err", err)
return 1
}
if path != "stdout" {
slog.Info("runtime health written", "path", path)
}
return 0
}
func runInstallToRAM(args []string, stdout, stderr io.Writer) int {
fs := flag.NewFlagSet("install-to-ram", flag.ContinueOnError)
fs.SetOutput(stderr)
fs.Usage = func() {
fmt.Fprintln(stderr, "usage: bee install-to-ram")
}
if err := fs.Parse(args); err != nil {
if err == flag.ErrHelp {
return 0
}
return 2
}
if fs.NArg() != 0 {
fs.Usage()
return 2
}
application := app.New(platform.New())
logLine := func(s string) { fmt.Fprintln(stdout, s) }
if err := application.RunInstallToRAM(context.Background(), logLine); err != nil {
slog.Error("run install-to-ram", "err", err)
return 1
}
return 0
}
func runSupportBundle(args []string, stdout, stderr io.Writer) int {
fs := flag.NewFlagSet("support-bundle", flag.ContinueOnError)
fs.SetOutput(stderr)
output := fs.String("output", "stdout", "output destination: stdout or file:<path>")
fs.Usage = func() {
fmt.Fprintln(stderr, "usage: bee support-bundle [--output stdout|file:<path>]")
fs.PrintDefaults()
}
if err := fs.Parse(args); err != nil {
if err == flag.ErrHelp {
return 0
}
return 2
}
if fs.NArg() != 0 {
fs.Usage()
return 2
}
path, err := app.BuildSupportBundle(app.DefaultExportDir)
if err != nil {
slog.Error("build support bundle", "err", err)
return 1
}
defer os.Remove(path)
raw, err := os.ReadFile(path)
if err != nil {
slog.Error("read support bundle", "err", err)
return 1
}
switch {
case *output == "stdout":
if _, err := stdout.Write(raw); err != nil {
slog.Error("write support bundle stdout", "err", err)
return 1
}
case strings.HasPrefix(*output, "file:"):
dst := strings.TrimPrefix(*output, "file:")
if err := os.WriteFile(dst, raw, 0644); err != nil {
slog.Error("write support bundle", "err", err)
return 1
}
slog.Info("support bundle written", "path", dst)
default:
fmt.Fprintln(stderr, "bee support-bundle: unknown output destination")
fs.Usage()
return 2
}
return 0
}
func runWeb(args []string, stdout, stderr io.Writer) int {
fs := flag.NewFlagSet("web", flag.ContinueOnError)
fs.SetOutput(stderr)
listenAddr := fs.String("listen", ":8080", "listen address, e.g. :80")
auditPath := fs.String("audit-path", "", "optional path to the latest audit JSON snapshot")
exportDir := fs.String("export-dir", app.DefaultExportDir, "directory with logs, SAT results, and support bundles")
title := fs.String("title", "Bee Hardware Audit", "page title")
fs.Usage = func() {
fmt.Fprintf(stderr, "usage: bee web [--listen :80] [--audit-path %s] [--export-dir %s] [--title \"Bee Hardware Audit\"]\n", app.DefaultAuditJSONPath, app.DefaultExportDir)
fs.PrintDefaults()
}
if err := fs.Parse(args); err != nil {
if err == flag.ErrHelp {
return 0
}
return 2
}
if fs.NArg() != 0 {
fs.Usage()
return 2
}
slog.Info("starting bee web", "listen", *listenAddr, "audit_path", *auditPath)
runtimeInfo, err := runtimeenv.Detect("auto")
if err != nil {
slog.Warn("resolve runtime for web", "err", err)
}
if err := webui.ListenAndServe(*listenAddr, webui.HandlerOptions{
Title: *title,
BuildLabel: buildLabel(),
AuditPath: *auditPath,
ExportDir: *exportDir,
App: app.New(platform.New()),
RuntimeMode: runtimeInfo.Mode,
}); err != nil {
slog.Error("run web", "err", err)
return 1
}
return 0
}
func runBlackbox(args []string, stdout, stderr io.Writer) int {
fs := flag.NewFlagSet("blackbox", flag.ContinueOnError)
fs.SetOutput(stderr)
exportDir := fs.String("export-dir", app.DefaultExportDir, "directory with logs, SAT results, and support bundles")
statePath := fs.String("state-file", app.DefaultBlackboxStatePath, "blackbox state file")
fs.Usage = func() {
fmt.Fprintf(stderr, "usage: bee blackbox [--export-dir %s] [--state-file %s]\n", app.DefaultExportDir, app.DefaultBlackboxStatePath)
fs.PrintDefaults()
}
if err := fs.Parse(args); err != nil {
if err == flag.ErrHelp {
return 0
}
return 2
}
if fs.NArg() != 0 {
fs.Usage()
return 2
}
slog.Info("starting bee blackbox", "export_dir", *exportDir, "state_file", *statePath)
if err := app.RunBlackbox(context.Background(), *exportDir, *statePath, platform.New()); err != nil && !errors.Is(err, context.Canceled) {
slog.Error("run blackbox", "err", err)
return 1
}
return 0
}
func runSAT(args []string, stdout, stderr io.Writer) int {
if len(args) == 0 {
fmt.Fprintln(stderr, "usage: bee sat nvidia|memory|storage|cpu [--duration <seconds>]")
return 2
}
if args[0] == "help" || args[0] == "--help" || args[0] == "-h" {
fmt.Fprintln(stdout, "usage: bee sat nvidia|memory|storage|cpu [--duration <seconds>]")
return 0
}
fs := flag.NewFlagSet("sat", flag.ContinueOnError)
fs.SetOutput(stderr)
duration := fs.Int("duration", 0, "stress-ng duration in seconds (cpu only; default: 60)")
diagLevel := fs.Int("diag-level", 0, "DCGM diagnostic level for nvidia (1=quick, 2=medium, 3=targeted stress, 4=extended stress; default: 1)")
if err := fs.Parse(args[1:]); err != nil {
if err == flag.ErrHelp {
return 0
}
return 2
}
if fs.NArg() != 0 {
fmt.Fprintf(stderr, "bee sat: unexpected arguments\n")
return 2
}
target := args[0]
if target != "nvidia" && target != "memory" && target != "storage" && target != "cpu" {
fmt.Fprintf(stderr, "bee sat: unknown target %q\n", target)
fmt.Fprintln(stderr, "usage: bee sat nvidia|memory|storage|cpu [--duration <seconds>] [--diag-level <1-4>]")
return 2
}
application := app.New(platform.New())
var (
archive string
err error
)
logLine := func(s string) { fmt.Fprintln(os.Stderr, s) }
switch target {
case "nvidia":
level := *diagLevel
if level > 0 {
_, err = application.RunNvidiaAcceptancePackWithOptions(context.Background(), "", level, nil, logLine)
} else {
archive, err = application.RunNvidiaAcceptancePack("", logLine)
}
case "memory":
archive, err = application.RunMemoryAcceptancePackCtx(context.Background(), "", 256, 1, logLine)
case "storage":
archive, err = application.RunStorageAcceptancePackCtx(context.Background(), "", false, logLine)
case "cpu":
dur := *duration
if dur <= 0 {
dur = 60
}
archive, err = application.RunCPUAcceptancePackCtx(context.Background(), "", dur, logLine)
}
if err != nil {
slog.Error("run sat", "target", target, "err", err)
return 1
}
slog.Info("sat archive written", "target", target, "path", archive)
return 0
}
const scenarioUsage = `usage: bee run <file.json | scenario-name>
A path (contains "/" or ends in ".json") is read directly from local disk.
A bare name is instead resolved as scenarios/<name>.json, checked first
against the scenarios shipped with this image (always available) and
then against any mounted removable media (e.g. the same USB stick
already plugged in for blackbox) — lets an air-gapped engineer author a
scenario on another machine and drop it there without needing a network
path onto the host.
See ParseScenarioJSON in audit/internal/platform/scenario.go for the file
format and a worked example (per-GPU/all-GPU load with concurrent IPMI/
nvidia-smi sampling).`
// runScenario implements both "bee run <arg>" and "bee scenario run <arg>"
// (the latter kept as a longer alias). args is just the scenario
// file/name — no leading verb.
func runScenario(args []string, stdout, stderr io.Writer) int {
if len(args) > 0 && (args[0] == "help" || args[0] == "--help" || args[0] == "-h") {
fmt.Fprintln(stdout, scenarioUsage)
return 0
}
fs := flag.NewFlagSet("run", flag.ContinueOnError)
fs.SetOutput(stderr)
if err := fs.Parse(args); err != nil {
if err == flag.ErrHelp {
return 0
}
return 2
}
if fs.NArg() != 1 {
fmt.Fprintln(stderr, scenarioUsage)
return 2
}
arg := fs.Arg(0)
sys := platform.New()
var (
data []byte
err error
)
if strings.Contains(arg, "/") || strings.HasSuffix(arg, ".json") {
data, err = os.ReadFile(arg)
} else {
data, err = sys.ReadScenario(arg)
}
if err != nil {
fmt.Fprintf(stderr, "bee run: %v\n", err)
return 1
}
spec, err := platform.ParseScenarioJSON(data)
if err != nil {
fmt.Fprintf(stderr, "bee run: %v\n", err)
return 1
}
// Run through the App (not sys.RunScenario directly) for two reasons:
// it wires the same blackbox kick/sync-bracket hooks the SAT job runner
// uses (see app.New()), AND App.RunScenario defaults the base dir to
// DefaultSATBaseDir (under DefaultExportDir) — the tree the blackbox
// worker actually mirrors to removable media. Calling sys.RunScenario
// with an empty base dir would write to /var/log/bee-sat instead, which
// blackbox does not mirror, so a `bee run` scenario would never reach the
// USB stick even with the hooks firing.
a := app.New(sys)
logLine := func(s string) { fmt.Fprintln(stderr, s) }
runDir, err := a.RunScenario(context.Background(), "", spec, logLine)
if err != nil {
fmt.Fprintf(stderr, "bee run: %v\n", err)
return 1
}
fmt.Fprintln(stdout, runDir)
return 0
}
func runBenchmark(args []string, stdout, stderr io.Writer) int {
if len(args) == 0 {
fmt.Fprintln(stderr, "usage: bee benchmark nvidia [--profile standard|stability|overnight] [--devices 0,1] [--exclude 2,3] [--size-mb N] [--skip-nccl]")
return 2
}
if args[0] == "help" || args[0] == "--help" || args[0] == "-h" {
fmt.Fprintln(stdout, "usage: bee benchmark nvidia [--profile standard|stability|overnight] [--devices 0,1] [--exclude 2,3] [--size-mb N] [--skip-nccl]")
return 0
}
target := args[0]
if target != "nvidia" {
fmt.Fprintf(stderr, "bee benchmark: unknown target %q\n", target)
fmt.Fprintln(stderr, "usage: bee benchmark nvidia [--profile standard|stability|overnight] [--devices 0,1] [--exclude 2,3] [--size-mb N] [--skip-nccl]")
return 2
}
fs := flag.NewFlagSet("benchmark", flag.ContinueOnError)
fs.SetOutput(stderr)
profile := fs.String("profile", platform.NvidiaBenchmarkProfileStandard, "benchmark profile: standard, stability, overnight")
devices := fs.String("devices", "", "comma-separated GPU indices to include")
exclude := fs.String("exclude", "", "comma-separated GPU indices to exclude")
sizeMB := fs.Int("size-mb", 0, "per-GPU benchmark buffer size in MB (0 = auto)")
skipNCCL := fs.Bool("skip-nccl", false, "skip multi-GPU NCCL interconnect benchmark")
if err := fs.Parse(args[1:]); err != nil {
if err == flag.ErrHelp {
return 0
}
return 2
}
if fs.NArg() != 0 {
fmt.Fprintf(stderr, "bee benchmark: unexpected arguments\n")
return 2
}
includeIndices, err := parseBenchmarkIndexCSV(*devices)
if err != nil {
fmt.Fprintf(stderr, "bee benchmark: invalid --devices: %v\n", err)
return 2
}
excludeIndices, err := parseBenchmarkIndexCSV(*exclude)
if err != nil {
fmt.Fprintf(stderr, "bee benchmark: invalid --exclude: %v\n", err)
return 2
}
application := app.New(platform.New())
logLine := func(s string) { fmt.Fprintln(os.Stderr, s) }
archive, err := application.RunNvidiaBenchmark("", platform.NvidiaBenchmarkOptions{
Profile: *profile,
SizeMB: *sizeMB,
GPUIndices: includeIndices,
ExcludeGPUIndices: excludeIndices,
RunNCCL: !*skipNCCL,
}, logLine)
if err != nil {
slog.Error("run benchmark", "target", target, "err", err)
return 1
}
slog.Info("benchmark archive written", "target", target, "path", archive)
return 0
}
// runGPUBandwidthGroups discovers this host's GPU-to-socket layout from
// "nvidia-smi topo -m" and prints a progressive multi-GPU test plan: a pair
// within each socket, one cross-socket pair, then every GPU. gpu_indices in
// a scenario JSON file is host-specific (see scenarios/README.md) — this
// exists so a scenario can stage a GPU-count escalation ("same socket, other
// socket, cross-socket, all") without baking any host's specific GPU indices
// into the file at all.
//
// With --stage, prints just that stage's comma-joined GPU indices (for a
// scenario job's cmd to substitute directly, e.g.
// `dcgmi diag -r nvbandwidth -i "$(bee gpu-bandwidth-groups --stage all)"`)
// and exits nonzero with nothing on stdout if that stage doesn't apply on
// this host (e.g. only one socket present, so there's no "cross-socket"
// stage) — callers should treat that as "skip this stage here", not a
// hard failure.
// Without --stage, lists every applicable stage as "<label>\t<gpu,indices>".
func runGPUBandwidthGroups(args []string, stdout, stderr io.Writer) int {
fs := flag.NewFlagSet("gpu-bandwidth-groups", flag.ContinueOnError)
fs.SetOutput(stderr)
stage := fs.String("stage", "", "print only this stage's GPU indices (same-socket-1, same-socket-2, ..., cross-socket, all)")
fs.Usage = func() {
fmt.Fprintln(stderr, "usage: bee gpu-bandwidth-groups [--stage <label>]")
fs.PrintDefaults()
}
if err := fs.Parse(args); err != nil {
if err == flag.ErrHelp {
return 0
}
return 2
}
if fs.NArg() != 0 {
fs.Usage()
return 2
}
raw, err := exec.Command("nvidia-smi", "topo", "-m").CombinedOutput()
if err != nil {
fmt.Fprintf(stderr, "bee gpu-bandwidth-groups: nvidia-smi topo -m: %v\n", err)
return 1
}
groups := platform.NvidiaProgressiveBandwidthGroups(platform.ParseNvidiaSocketGroups(string(raw)))
if *stage == "" {
for _, g := range groups {
fmt.Fprintf(stdout, "%s\t%s\n", g.Label, joinInts(g.GPUIndices))
}
return 0
}
for _, g := range groups {
if g.Label == *stage {
fmt.Fprintln(stdout, joinInts(g.GPUIndices))
return 0
}
}
fmt.Fprintf(stderr, "bee gpu-bandwidth-groups: stage %q does not apply on this host\n", *stage)
return 1
}
func joinInts(vals []int) string {
parts := make([]string, len(vals))
for i, v := range vals {
parts[i] = strconv.Itoa(v)
}
return strings.Join(parts, ",")
}
func runBeeWorker(args []string, stdout, stderr io.Writer) int {
fs := flag.NewFlagSet("bee-worker", flag.ContinueOnError)
fs.SetOutput(stderr)
exportDir := fs.String("export-dir", app.DefaultExportDir, "directory with task state and artifacts")
taskID := fs.String("task-id", "", "task identifier, e.g. TASK-001")
fs.Usage = func() {
fmt.Fprintf(stderr, "usage: bee bee-worker --export-dir %s --task-id TASK-001\n", app.DefaultExportDir)
fs.PrintDefaults()
}
if err := fs.Parse(args); err != nil {
if err == flag.ErrHelp {
return 0
}
return 2
}
if fs.NArg() != 0 {
fs.Usage()
return 2
}
return webui.RunPersistedTask(*exportDir, *taskID, stdout, stderr)
}
func parseBenchmarkIndexCSV(raw string) ([]int, error) {
raw = strings.TrimSpace(raw)
if raw == "" {
return nil, nil
}
var indices []int
for _, part := range strings.Split(raw, ",") {
part = strings.TrimSpace(part)
if part == "" {
continue
}
value, err := strconv.Atoi(part)
if err != nil || value < 0 {
return nil, fmt.Errorf("bad gpu index %q", part)
}
indices = append(indices, value)
}
return indices, nil
}