mirror of
https://github.com/mulgadc/spinifex.git
synced 2026-10-15 09:01:31 +00:00
feat: VFIO GPU passthrough (#188)
* feat: add GPU discovery package (mulga-992) Adds spinifex/gpu package with sysfs-based GPU enumeration, IOMMU group member lookup, and a PCI device ID model table. Vendor-specific enrichment via nvidia-smi is attempted opportunistically and falls back to the built-in table when unavailable. All discovery and IOMMU logic is covered by unit tests using mock sysfs trees. * feat: add VFIO bind/unbind and GPU pool manager (mulga-993) Adds gpu.BindVFIO/UnbindVFIO using sysfs driver_override for per-device driver switching, and gpu.Manager which owns the claim/release lifecycle for the GPU pool including full IOMMU group binding, rollback on partial failure, and marking entries unavailable when rebind fails after VM termination. * feat: add G5 instance type generator with GPU model support (mulga-994) - Add GPUModel type, NVIDIAa10g constant, and GPUModelForVendorDevice lookup - Add generateGPUTypes to emit g5.* types with GpuInfo populated - Add IsGPUType helper for consumers filtering GPU-capable types - Update DetectAndGenerate to accept gpuModels []GPUModel parameter * feat: integrate GPU passthrough into daemon lifecycle (mulga-995) - Add gpu_passthrough config flag; discover GPUs and build manager at startup - Gate g5 instance types on GPU pool availability via canAllocateCount - Claim GPU (VFIO bind) in LaunchInstance; inject vfio-pci QEMU device - Release GPU (VFIO unbind, rebind to host driver) in stopInstance - Re-claim GPU on daemon restart when QEMU process is still alive - Add UseUEFI / OVMF support for x86_64 VMs (optional, falls back to SeaBIOS) - Fix Claim to use stored OriginalDriver when device is already vfio-pci bound * feat: add GPU fields to NodeStatus and spx top nodes (mulga-996) - Add TotalGPUs/AllocGPUs to NodeStatusResponse (zero on non-GPU nodes) - Populate from gpuManager in handleNodeStatus - Show GPU (used/total) column in spx top nodes; "-" when no GPUs present * feat: Ubuntu 24.04 GPU AMI build pipeline (mulga-997) Extends build-system-image.sh to support Ubuntu alongside Alpine. Key differences handled: partition table (parted resizepart p1), bind mounts for DKMS (/proc /sys /dev), apt vs apk, systemctl vs rc-update. DISTRO defaults to "alpine" for backward compatibility. Adds ubuntu-gpu.conf manifest and ubuntu-gpu-setup.sh chroot script that installs nvidia-headless-550-server with DKMS pre-built against a pinned linux-image-generic kernel, blacklists nouveau, and prevents unattended-upgrades from invalidating the pre-built module. * Fixed GPU passthrough and added test script * feat: always-probe GPU at startup, surface capability in NodeStatusResponse - Add gpuProbeResult struct and probeGPU() which runs at daemon startup regardless of gpu_passthrough config. Checks sysfs for GPUs, IOMMU groups, and vfio_pci module presence. - Store gpuProbe on Daemon struct; split NewDaemon into probe (always) and activate (opt-in) phases. - Populate GPUCapable, GPUPassthrough, GPUModels on NodeStatusResponse so operators can see hardware state before enabling passthrough. - Update 'spx top' GPU column: '-' = no hardware, '0/N*' = capable but passthrough not active, 'used/total' = passthrough active. * feat: SIGHUP hot-reload for GPU passthrough config (mulga-1001) - Export GenerateGPUTypes for use by the reload path - Add ResourceManager.reloadGPUTypes: mutates instanceTypes map in-place (preserves references held by instanceService) and updates NATS subscriptions without full daemon restart - Separate SIGHUP from SIGTERM/SIGINT: setupReload() listens for SIGHUP in a context-aware loop; setupShutdown no longer intercepts SIGHUP - Add reloadConfig() + applyGPUConfig(): re-reads spinifex.toml on signal, handles false→true (re-probe + activate) and true→false (refuse if instances running, otherwise deactivate) transitions * feat: spx admin gpu status/enable/disable commands (mulga-1000) - admin.SetGPUPassthrough: Go implementation of the idempotent TOML write (replaces the inline Python heredoc in spx-test-gpu.sh) - spx admin gpu status: shows GPU hardware, IOMMU/vfio state, passthrough activation, and advertised g5 instance types via NATS NodeStatusResponse - spx admin gpu enable: checks prerequisites, writes TOML, sends SIGHUP, polls until daemon confirms GPUPassthrough=true - spx admin gpu disable: refuses while GPU instances are running, otherwise writes TOML, sends SIGHUP, polls until confirmed - spx-test-gpu.sh: replace Python TOML heredoc with spx admin gpu enable * fixed admin * fix: filter Intel iGPU from passthrough pool, use lspci for consumer GPU names Discover() now skips non-NVIDIA/AMD vendors so Intel iGPUs never enter the passthrough pool. enrichFromLspci() uses the system PCI ID database via lspci -vmm to populate the model name for any card not in the hardcoded table, without needing nvidia-smi or driver tools. * fix: grant CAP_SYS_ADMIN and vfio device access to spinifex-daemon GPU passthrough requires writing to /sys/bus/pci/drivers/vfio-pci/unbind on instance stop. The kernel gates this on CAP_SYS_ADMIN regardless of file permissions. QEMU subprocesses also need char-vfio rw access to open /dev/vfio/N for the passed-through device. * fixing gpu bining permissions * working state of gpu passthrough on phyiscal box * installer grub image fixes * made gpu-passthrough more user facing * fixed banner * fixed banner * banner fix * gpu passthrough docs * preflight tweak * vpcd fixes to set br-wan to promisc for public IP DHCP processes * fixed vpcd br-wan promisc and veth pair * ssh keygen fix * fixes after running /review-branch * feat: re-wire GPU claim and device injection after vm.Manager refactor GPU passthrough claim (VFIO bind) and QEMU device injection were lost when dev's vm.Manager refactor (#185/#186) moved StartInstance/ CreateQMPClient into vm/lifecycle.go. Re-add both in their correct new homes: - daemon_handlers_instance.go: claim GPU before vmMgr.Run() in both RunInstances and StartStoppedInstance handlers; release on failure - vm/lifecycle.go startQEMU(): inject vfio-pci device with x-vga flag derived from instance.GPUXVGAEnabled when GPUPCIAddress is set * fix: Create spinifex user ~/.ssh and keys in firstboot --------- Co-authored-by: Ben Duncan <5178243+benduncan@users.noreply.github.com> Co-authored-by: Julian Sommer <115516918+juliansommer@users.noreply.github.com>
This commit is contained in:
Binary file not shown.
|
After Width: | Height: | Size: 525 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 490 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 754 KiB |
@@ -0,0 +1,12 @@
|
||||
[Unit]
|
||||
Description=Regenerate SSH host keys
|
||||
ConditionPathExists=!/etc/ssh/ssh_host_ed25519_key
|
||||
Before=ssh.service
|
||||
|
||||
[Service]
|
||||
Type=oneshot
|
||||
ExecStart=/usr/bin/ssh-keygen -A
|
||||
RemainAfterExit=yes
|
||||
|
||||
[Install]
|
||||
WantedBy=multi-user.target
|
||||
@@ -17,12 +17,18 @@ Environment=SPINIFEX_CONFIG_PATH=/etc/spinifex/spinifex.toml
|
||||
Environment=SPINIFEX_BASE_DIR=/var/lib/spinifex/spinifex/
|
||||
Environment=SPINIFEX_WAL_DIR=/var/lib/spinifex/spinifex/
|
||||
|
||||
# Filesystem access — NoNewPrivileges omitted (daemon uses sudo for ip, ovs-vsctl)
|
||||
# Filesystem access — NoNewPrivileges omitted (daemon uses sudo for ip, ovs-vsctl, dhcpcd)
|
||||
# CAP_SYS_ADMIN: required for GPU passthrough sysfs operations (vfio group management).
|
||||
# CAP_DAC_OVERRIDE: sysfs driver bind/unbind paths (/sys/bus/pci/drivers/vfio-pci/bind,
|
||||
# unbind) are root:root 0200; non-root processes need CAP_DAC_OVERRIDE to open them.
|
||||
AmbientCapabilities=CAP_SYS_ADMIN CAP_DAC_OVERRIDE
|
||||
ProtectSystem=full
|
||||
PrivateDevices=no
|
||||
DevicePolicy=closed
|
||||
DeviceAllow=/dev/kvm rw
|
||||
DeviceAllow=/dev/net/tun rw
|
||||
DeviceAllow=char-vfio rw
|
||||
DeviceAllow=/dev/vfio/vfio rw
|
||||
ReadOnlyPaths=/etc/spinifex/spinifex.toml /etc/spinifex/server.pem /etc/spinifex/server.key /etc/spinifex/ca.pem
|
||||
ReadWritePaths=/var/lib/spinifex/spinifex /var/log/spinifex /run/spinifex /run/spinifex/nbd
|
||||
|
||||
@@ -50,6 +56,9 @@ RemoveIPC=yes
|
||||
LimitCORE=0
|
||||
LimitNOFILE=65536
|
||||
LimitNPROC=4096
|
||||
# VFIO GPU passthrough requires locking all guest RAM into physical memory.
|
||||
# Without this, QEMU fails with RLIMIT_MEMLOCK exceeded when pinning guest pages.
|
||||
LimitMEMLOCK=infinity
|
||||
|
||||
[Install]
|
||||
WantedBy=spinifex.target
|
||||
|
||||
@@ -0,0 +1,5 @@
|
||||
# Allow spinifex group to access VFIO devices for GPU passthrough.
|
||||
# QEMU (launched by spinifex-daemon, group spinifex) must open /dev/vfio/N
|
||||
# to use vfio-pci passthrough devices. The kernel creates these nodes as
|
||||
# root:root 0600 by default; this rule grants group-level read-write access.
|
||||
SUBSYSTEM=="vfio", GROUP="spinifex", MODE="0660"
|
||||
@@ -0,0 +1,172 @@
|
||||
// dhcptest is a diagnostic tool that runs a DHCP DORA handshake using the
|
||||
// same nclient4 library and options that spinifex-vpcd uses for gateway pool
|
||||
// leases. Run it on a node to verify DHCP works at the library level before
|
||||
// debugging the full spinifex stack.
|
||||
//
|
||||
// Usage (must run as root — needs AF_PACKET):
|
||||
//
|
||||
// sudo dhcptest --iface=br-wan
|
||||
// sudo dhcptest --iface=br-wan --mac=02:00:00:7c:b7:b9 --promisc
|
||||
// sudo dhcptest --iface=br-wan --spinifex-id=gateway-wan1
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"flag"
|
||||
"fmt"
|
||||
"log"
|
||||
"net"
|
||||
"os"
|
||||
"os/exec"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/insomniacslk/dhcp/dhcpv4"
|
||||
"github.com/insomniacslk/dhcp/dhcpv4/nclient4"
|
||||
)
|
||||
|
||||
func main() {
|
||||
iface := flag.String("iface", "br-wan", "Interface/bridge to bind the DHCP client to")
|
||||
macStr := flag.String("mac", "", "Hardware address for DORA (chaddr + AF_PACKET filter). If empty, uses the real interface MAC.")
|
||||
spinifexID := flag.String("spinifex-id", "", "Generate a spinifex virtual MAC from this resource ID (e.g. 'gateway-wan1'). Overrides --mac.")
|
||||
clientID := flag.String("client-id", "", "DHCP option 61 client identifier. Defaults to --spinifex-id or --mac value.")
|
||||
hostname := flag.String("hostname", "", "DHCP option 12 hostname")
|
||||
vendorClass := flag.String("vendor-class", "mulga-spinifex-gw", "DHCP option 60 vendor class identifier")
|
||||
promisc := flag.Bool("promisc", false, "Enable promiscuous mode on the interface before the test (mirrors vpcd's setBridgePromisc)")
|
||||
timeout := flag.Duration("timeout", 10*time.Second, "Per-attempt DHCP timeout")
|
||||
retries := flag.Int("retries", 3, "Number of DORA attempts before giving up")
|
||||
verbose := flag.Bool("verbose", false, "Print full DHCP packet contents")
|
||||
flag.Parse()
|
||||
|
||||
// Resolve the hardware address to use.
|
||||
var hwAddr net.HardwareAddr
|
||||
var err error
|
||||
|
||||
switch {
|
||||
case *spinifexID != "":
|
||||
hwAddr, err = net.ParseMAC(generateMAC(*spinifexID))
|
||||
if err != nil {
|
||||
log.Fatalf("generateMAC(%q): %v", *spinifexID, err)
|
||||
}
|
||||
fmt.Printf("spinifex virtual MAC for %q: %s\n", *spinifexID, hwAddr)
|
||||
case *macStr != "":
|
||||
hwAddr, err = net.ParseMAC(*macStr)
|
||||
if err != nil {
|
||||
log.Fatalf("parse --mac %q: %v", *macStr, err)
|
||||
}
|
||||
default:
|
||||
hwAddr, err = ifaceMAC(*iface)
|
||||
if err != nil {
|
||||
log.Fatalf("get MAC for %s: %v", *iface, err)
|
||||
}
|
||||
fmt.Printf("using real interface MAC: %s\n", hwAddr)
|
||||
}
|
||||
|
||||
// Resolve client-id: default to spinifex-id, then MAC string.
|
||||
cid := *clientID
|
||||
if cid == "" {
|
||||
if *spinifexID != "" {
|
||||
cid = *spinifexID
|
||||
} else {
|
||||
cid = hwAddr.String()
|
||||
}
|
||||
}
|
||||
|
||||
fmt.Printf("iface: %s\n", *iface)
|
||||
fmt.Printf("hw-addr: %s\n", hwAddr)
|
||||
fmt.Printf("client-id: %s\n", cid)
|
||||
fmt.Printf("hostname: %s\n", *hostname)
|
||||
fmt.Printf("vendor-class: %s\n", *vendorClass)
|
||||
fmt.Printf("promisc: %v\n", *promisc)
|
||||
fmt.Printf("timeout: %s retries: %d\n", *timeout, *retries)
|
||||
fmt.Println()
|
||||
|
||||
if *promisc {
|
||||
if err := setPromisc(*iface, true); err != nil {
|
||||
log.Printf("WARNING: failed to set promisc on %s: %v", *iface, err)
|
||||
} else {
|
||||
fmt.Printf("set %s promisc ON\n\n", *iface)
|
||||
}
|
||||
}
|
||||
|
||||
// Build nclient4 options — identical to spinifex/services/vpcd/dhcp/nclient4.go.
|
||||
opts := []nclient4.ClientOpt{
|
||||
nclient4.WithHWAddr(hwAddr),
|
||||
nclient4.WithTimeout(*timeout),
|
||||
nclient4.WithRetry(*retries),
|
||||
}
|
||||
if *verbose {
|
||||
opts = append(opts, nclient4.WithLogger(nclient4.ShortSummaryLogger{Printfer: log.New(os.Stderr, "[nclient4] ", 0)}))
|
||||
}
|
||||
|
||||
client, err := nclient4.New(*iface, opts...)
|
||||
if err != nil {
|
||||
log.Fatalf("nclient4.New(%s): %v", *iface, err)
|
||||
}
|
||||
defer client.Close()
|
||||
|
||||
// Identity modifiers — identical to identityModifiers() in vpcd dhcp package.
|
||||
var mods []dhcpv4.Modifier
|
||||
if cid != "" {
|
||||
mods = append(mods, dhcpv4.WithOption(dhcpv4.OptClientIdentifier([]byte(cid))))
|
||||
}
|
||||
if *hostname != "" {
|
||||
mods = append(mods, dhcpv4.WithOption(dhcpv4.OptHostName(*hostname)))
|
||||
}
|
||||
if *vendorClass != "" {
|
||||
mods = append(mods, dhcpv4.WithOption(dhcpv4.OptClassIdentifier(*vendorClass)))
|
||||
}
|
||||
|
||||
fmt.Println("Sending DHCP Discover...")
|
||||
start := time.Now()
|
||||
ctx, cancel := context.WithTimeout(context.Background(), *timeout*time.Duration(*retries+1))
|
||||
defer cancel()
|
||||
|
||||
lease, err := client.Request(ctx, mods...)
|
||||
if err != nil {
|
||||
fmt.Printf("\nFAILED after %s: %v\n", time.Since(start).Round(time.Millisecond), err)
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
fmt.Printf("SUCCESS in %s\n\n", time.Since(start).Round(time.Millisecond))
|
||||
ack := lease.ACK
|
||||
fmt.Printf(" Your IP: %s\n", ack.YourIPAddr)
|
||||
fmt.Printf(" Subnet: %s\n", ack.SubnetMask())
|
||||
fmt.Printf(" Gateway: %v\n", ack.Router())
|
||||
fmt.Printf(" DNS: %v\n", ack.DNS())
|
||||
fmt.Printf(" Server ID: %s\n", ack.ServerIdentifier())
|
||||
fmt.Printf(" Lease time: %s\n", ack.IPAddressLeaseTime(24*time.Hour))
|
||||
if *verbose {
|
||||
fmt.Printf("\n--- Full ACK ---\n%s\n", ack.Summary())
|
||||
}
|
||||
}
|
||||
|
||||
// generateMAC is the same deterministic hash used by vpcd's topology.go.
|
||||
// It produces a locally-administered MAC (02:00:00:xx:xx:xx) from a resource ID.
|
||||
func generateMAC(resourceID string) string {
|
||||
h := uint32(0)
|
||||
for _, c := range resourceID {
|
||||
h = h*31 + uint32(c)
|
||||
}
|
||||
return fmt.Sprintf("02:00:00:%02x:%02x:%02x", (h>>16)&0xff, (h>>8)&0xff, h&0xff)
|
||||
}
|
||||
|
||||
func ifaceMAC(name string) (net.HardwareAddr, error) {
|
||||
ifc, err := net.InterfaceByName(name)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return ifc.HardwareAddr, nil
|
||||
}
|
||||
|
||||
func setPromisc(iface string, on bool) error {
|
||||
state := "on"
|
||||
if !on {
|
||||
state = "off"
|
||||
}
|
||||
out, err := exec.Command("ip", "link", "set", iface, "promisc", state).CombinedOutput()
|
||||
if err != nil {
|
||||
return fmt.Errorf("%w: %s", err, strings.TrimSpace(string(out)))
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -134,6 +134,10 @@ func buildConfig() (*install.Config, error) {
|
||||
}
|
||||
}
|
||||
|
||||
if os.Getenv("SPINIFEX_GPU_PASSTHROUGH") == "1" {
|
||||
cfg.GPUPassthrough = true
|
||||
}
|
||||
|
||||
role := strings.ToLower(os.Getenv("SPINIFEX_ROLE"))
|
||||
if role == "" {
|
||||
role = "init"
|
||||
|
||||
@@ -43,6 +43,9 @@ type Config struct {
|
||||
// the headless path. Passed to `spx admin init --email=<value>` when set;
|
||||
// omitted entirely when empty.
|
||||
Email string
|
||||
// GPUPassthrough enables VFIO GPU passthrough by passing --gpu-passthrough
|
||||
// to `spx admin init`, which writes gpu_passthrough = true in the daemon config.
|
||||
GPUPassthrough bool
|
||||
}
|
||||
|
||||
// Write drops the firstboot script and systemd unit into root, which should be
|
||||
@@ -146,6 +149,18 @@ SETUP_STAGES=fixown /usr/local/share/spinifex/setup.sh
|
||||
# setup.sh's create_directories — it resolves to /etc/spinifex, so
|
||||
# {BaseDir}/config/master.key automatically points at /etc/spinifex/master.key.
|
||||
|
||||
# Bootstrap the spinifex user's SSH directory with a generated key pair so
|
||||
# ~/.ssh exists and the operator can immediately use ssh-keygen / ssh-copy-id
|
||||
# without hitting "No such file or directory" errors documented in the setup guide.
|
||||
if [ ! -f /home/spinifex/.ssh/id_ed25519 ]; then
|
||||
mkdir -p /home/spinifex/.ssh
|
||||
chmod 700 /home/spinifex/.ssh
|
||||
ssh-keygen -t ed25519 -f /home/spinifex/.ssh/id_ed25519 -N "" -C "spinifex@$(hostname)"
|
||||
chown -R spinifex:spinifex /home/spinifex/.ssh
|
||||
chmod 600 /home/spinifex/.ssh/id_ed25519
|
||||
chmod 644 /home/spinifex/.ssh/id_ed25519.pub
|
||||
fi
|
||||
|
||||
# Copy AWS credentials to the spinifex user's home directory.
|
||||
# spx admin init runs with HOME=/root (set by the systemd unit), so credentials
|
||||
# land in /root/.aws/. Copy them to the spinifex user's home so the operator
|
||||
@@ -210,11 +225,15 @@ func buildClusterCmd(cfg Config) string {
|
||||
// validator already rejects whitespace and @-chains.
|
||||
emailFlag = " --email=" + shellEscapeSingle(cfg.Email)
|
||||
}
|
||||
gpuFlag := ""
|
||||
if cfg.GPUPassthrough {
|
||||
gpuFlag = " --gpu-passthrough"
|
||||
}
|
||||
switch cfg.ClusterRole {
|
||||
case "join":
|
||||
return fmt.Sprintf("spx admin join --node %s --host %s%s", cfg.Hostname, cfg.JoinAddr, emailFlag)
|
||||
default:
|
||||
return fmt.Sprintf("spx admin init --node %s --nodes 1%s", cfg.Hostname, emailFlag)
|
||||
return fmt.Sprintf("spx admin init --node %s --nodes 1%s%s", cfg.Hostname, emailFlag, gpuFlag)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -67,4 +67,9 @@ type Config struct {
|
||||
// Required on interactive installs; may be empty on headless/CI installs
|
||||
// when SPINIFEX_EMAIL is not supplied on the kernel cmdline.
|
||||
Email string
|
||||
|
||||
// GPUPassthrough enables VFIO GPU passthrough on this node.
|
||||
// Set via SPINIFEX_GPU_PASSTHROUGH=1 on headless installs.
|
||||
// Passes --gpu-passthrough to `spx admin init` during firstboot.
|
||||
GPUPassthrough bool
|
||||
}
|
||||
|
||||
+93
-11
@@ -45,6 +45,7 @@ import (
|
||||
"github.com/mulgadc/spinifex/spinifex/admin"
|
||||
"github.com/mulgadc/spinifex/spinifex/config"
|
||||
"github.com/mulgadc/spinifex/spinifex/formation"
|
||||
"github.com/mulgadc/spinifex/spinifex/gpu"
|
||||
handlers_ec2_vpc "github.com/mulgadc/spinifex/spinifex/handlers/ec2/vpc"
|
||||
handlers_iam "github.com/mulgadc/spinifex/spinifex/handlers/iam"
|
||||
"github.com/mulgadc/spinifex/spinifex/utils"
|
||||
@@ -274,6 +275,7 @@ func init() {
|
||||
adminInitCmd.Flags().String("gateway-ip", "", "OVN gateway router's external IP for SNAT (default: pool range_start for pool mode, required for nat mode without DHCP)")
|
||||
adminInitCmd.Flags().Int("external-prefix-len", 24, "External pool subnet prefix length (auto-detected)")
|
||||
adminInitCmd.Flags().Bool("no-external", false, "Disable external networking (overlay-only, no internet for VMs)")
|
||||
adminInitCmd.Flags().Bool("gpu-passthrough", false, "Enable VFIO GPU passthrough (sets gpu_passthrough = true in daemon config)")
|
||||
|
||||
// Flags for admin join
|
||||
adminJoinCmd.Flags().String("region", "ap-southeast-2", "Region for this node")
|
||||
@@ -660,6 +662,7 @@ func runAdminInit(cmd *cobra.Command, args []string) {
|
||||
externalPrefixLen, _ := cmd.Flags().GetInt("external-prefix-len")
|
||||
gatewayIP, _ := cmd.Flags().GetString("gateway-ip")
|
||||
noExternal, _ := cmd.Flags().GetBool("no-external")
|
||||
gpuPassthrough, _ := cmd.Flags().GetBool("gpu-passthrough")
|
||||
|
||||
// Fire telemetry in background (completes during init work, waited at end)
|
||||
noTelemetry, _ := cmd.Flags().GetBool("no-telemetry")
|
||||
@@ -1072,6 +1075,8 @@ func runAdminInit(cmd *cobra.Command, args []string) {
|
||||
BootstrapIgwId: bootstrapIgwId,
|
||||
BootstrapCidr: handlers_ec2_vpc.DefaultVPCCidr,
|
||||
BootstrapSubnetCidr: handlers_ec2_vpc.DefaultSubnetCidr,
|
||||
|
||||
GPUPassthrough: gpuPassthrough,
|
||||
}
|
||||
|
||||
// Print external networking summary
|
||||
@@ -2331,17 +2336,17 @@ func runAdminBanner(cmd *cobra.Command, _ []string) {
|
||||
}
|
||||
|
||||
banner := fmt.Sprintf(`
|
||||
+----------------------------------------------------+
|
||||
| Spinifex — Mulga Defense Corporation |
|
||||
+----------------------------------------------------+
|
||||
| Node: %-39s|
|
||||
| Login: %-39s|
|
||||
| Dashboard: %-39s|
|
||||
| API: %-39s|
|
||||
| SSH: %-39s|
|
||||
+----------------------------------------------------+
|
||||
| AWS credentials: cat ~/.aws/credentials |
|
||||
+----------------------------------------------------+
|
||||
+--------------------------------------------------------------+
|
||||
| Spinifex — Mulga Defense Corporation |
|
||||
+--------------------------------------------------------------+
|
||||
| Node: %-49s|
|
||||
| Login: %-49s|
|
||||
| Dashboard: %-49s|
|
||||
| API: %-49s|
|
||||
| SSH: %-49s|
|
||||
+--------------------------------------------------------------+
|
||||
| AWS credentials: cat ~/.aws/credentials |
|
||||
+--------------------------------------------------------------+
|
||||
|
||||
`,
|
||||
hostname,
|
||||
@@ -2351,6 +2356,8 @@ func runAdminBanner(cmd *cobra.Command, _ []string) {
|
||||
"spinifex@"+currentIP,
|
||||
)
|
||||
|
||||
banner += gpuBannerSection()
|
||||
|
||||
// Write to /etc/issue — displayed on the console before the login prompt.
|
||||
// Overwrite entirely; this is a purpose-built appliance so we own this file.
|
||||
if err := os.WriteFile("/etc/issue", []byte(banner), 0o644); err != nil {
|
||||
@@ -2390,6 +2397,81 @@ func appendBannerToMotd(banner string) error {
|
||||
return os.WriteFile(motdPath, []byte(base+sentinel+banner), 0o644)
|
||||
}
|
||||
|
||||
// gpuBannerSection returns an optional banner box section describing GPU state.
|
||||
// Returns "" when no GPU hardware is detected. Safe to call at boot before the
|
||||
// daemon starts — all checks are sysfs/file reads, no NATS required.
|
||||
func gpuBannerSection() string {
|
||||
devices, err := gpu.Discover()
|
||||
if err != nil || len(devices) == 0 {
|
||||
return ""
|
||||
}
|
||||
|
||||
iommuEntries, _ := os.ReadDir("/sys/kernel/iommu_groups/")
|
||||
iommuActive := len(iommuEntries) > 0
|
||||
|
||||
_, vfioErr := os.Stat("/sys/module/vfio_pci")
|
||||
vfioPresent := vfioErr == nil
|
||||
|
||||
passthroughEnabled := false
|
||||
cfgPath := DefaultConfigFile()
|
||||
if cfg, err := config.LoadConfig(cfgPath); err == nil {
|
||||
if nodeCfg, ok := cfg.Nodes[cfg.Node]; ok {
|
||||
passthroughEnabled = nodeCfg.Daemon.GPUPassthrough
|
||||
}
|
||||
}
|
||||
|
||||
models := gpuModelSummary(devices)
|
||||
|
||||
var statusLine, hintLine string
|
||||
switch {
|
||||
case passthroughEnabled:
|
||||
statusLine = "Passthrough enabled"
|
||||
case iommuActive && vfioPresent:
|
||||
statusLine = "Ready to enable"
|
||||
hintLine = "sudo spx admin gpu enable"
|
||||
default:
|
||||
statusLine = "Setup required"
|
||||
hintLine = "sudo spx admin gpu setup"
|
||||
}
|
||||
|
||||
const (
|
||||
sep = " +--------------------------------------------------------------+\n"
|
||||
maxVal = 55
|
||||
)
|
||||
if len([]rune(models)) > maxVal {
|
||||
models = string([]rune(models)[:maxVal-3]) + "..."
|
||||
}
|
||||
|
||||
section := sep +
|
||||
fmt.Sprintf(" | GPU: %-55s|\n", models) +
|
||||
fmt.Sprintf(" | %-55s|\n", statusLine)
|
||||
if hintLine != "" {
|
||||
section += fmt.Sprintf(" | %-55s|\n", hintLine)
|
||||
}
|
||||
section += sep + "\n"
|
||||
return section
|
||||
}
|
||||
|
||||
func gpuModelSummary(devices []gpu.GPUDevice) string {
|
||||
counts := make(map[string]int)
|
||||
var order []string
|
||||
for _, d := range devices {
|
||||
if counts[d.Model] == 0 {
|
||||
order = append(order, d.Model)
|
||||
}
|
||||
counts[d.Model]++
|
||||
}
|
||||
var parts []string
|
||||
for _, m := range order {
|
||||
if n := counts[m]; n > 1 {
|
||||
parts = append(parts, fmt.Sprintf("%dx %s", n, m))
|
||||
} else {
|
||||
parts = append(parts, m)
|
||||
}
|
||||
}
|
||||
return strings.Join(parts, ", ")
|
||||
}
|
||||
|
||||
// parseNodeConf reads a KEY=VALUE shell-format file and returns a map.
|
||||
// Lines starting with # and blank lines are ignored.
|
||||
func parseNodeConf(path string) map[string]string {
|
||||
|
||||
@@ -0,0 +1,447 @@
|
||||
package cmd
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/mulgadc/spinifex/spinifex/admin"
|
||||
"github.com/mulgadc/spinifex/spinifex/config"
|
||||
"github.com/mulgadc/spinifex/spinifex/gpu"
|
||||
"github.com/mulgadc/spinifex/spinifex/types"
|
||||
"github.com/spf13/cobra"
|
||||
"github.com/spf13/viper"
|
||||
)
|
||||
|
||||
var adminGpuCmd = &cobra.Command{
|
||||
Use: "gpu",
|
||||
Short: "Manage GPU passthrough for a node",
|
||||
}
|
||||
|
||||
var adminGpuStatusCmd = &cobra.Command{
|
||||
Use: "status",
|
||||
Short: "Show GPU hardware and passthrough state",
|
||||
Run: runAdminGpuStatus,
|
||||
}
|
||||
|
||||
var adminGpuEnableCmd = &cobra.Command{
|
||||
Use: "enable",
|
||||
Short: "Enable GPU passthrough on a node",
|
||||
Run: runAdminGpuEnable,
|
||||
}
|
||||
|
||||
var adminGpuDisableCmd = &cobra.Command{
|
||||
Use: "disable",
|
||||
Short: "Disable GPU passthrough on a node (blocked while GPU instances are running)",
|
||||
Run: runAdminGpuDisable,
|
||||
}
|
||||
|
||||
var adminGpuSetupCmd = &cobra.Command{
|
||||
Use: "setup",
|
||||
Short: "Configure this host for GPU passthrough (IOMMU, vfio-pci, nouveau blacklist)",
|
||||
Long: `Idempotent host setup for GPU passthrough. Run once before a reboot to configure
|
||||
GRUB, vfio-pci early binding, and the nouveau blacklist; then run again after the
|
||||
reboot to verify bindings and enable GPU passthrough in the daemon.
|
||||
|
||||
Must be run as root.`,
|
||||
Run: runAdminGpuSetup,
|
||||
}
|
||||
|
||||
func init() {
|
||||
adminCmd.AddCommand(adminGpuCmd)
|
||||
adminGpuCmd.AddCommand(adminGpuStatusCmd)
|
||||
adminGpuCmd.AddCommand(adminGpuEnableCmd)
|
||||
adminGpuCmd.AddCommand(adminGpuDisableCmd)
|
||||
adminGpuCmd.AddCommand(adminGpuSetupCmd)
|
||||
|
||||
adminGpuStatusCmd.Flags().String("node", "", "Target node name (default: local node)")
|
||||
// enable, disable, and setup write to the local spinifex.toml and signal the
|
||||
// local daemon — they must be run directly on the target host.
|
||||
}
|
||||
|
||||
// gpuNodeStatus queries NATS and returns the NodeStatusResponse for the target node.
|
||||
func gpuNodeStatus(targetNode string) (*types.NodeStatusResponse, error) {
|
||||
cfg, nc, err := loadConfigAndConnect()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer nc.Close()
|
||||
|
||||
if targetNode == "" {
|
||||
targetNode = cfg.Node
|
||||
}
|
||||
|
||||
responses, err := collectResponses(nc, "spinifex.node.status", 3*time.Second)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("collect node status: %w", err)
|
||||
}
|
||||
for _, raw := range responses {
|
||||
var resp types.NodeStatusResponse
|
||||
if err := json.Unmarshal(raw, &resp); err != nil {
|
||||
continue
|
||||
}
|
||||
if resp.Node == targetNode {
|
||||
return &resp, nil
|
||||
}
|
||||
}
|
||||
return nil, fmt.Errorf("node %q not found or not responding", targetNode)
|
||||
}
|
||||
|
||||
func runAdminGpuStatus(cmd *cobra.Command, _ []string) {
|
||||
targetNode, _ := cmd.Flags().GetString("node")
|
||||
resp, err := gpuNodeStatus(targetNode)
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "Error: %v\n", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
iommu := "unknown"
|
||||
vfio := "unknown"
|
||||
if resp.GPUCapable || len(resp.GPUModels) > 0 {
|
||||
iommu = "active"
|
||||
vfio = "loaded"
|
||||
}
|
||||
|
||||
fmt.Printf("Node: %s\n", resp.Node)
|
||||
if len(resp.GPUModels) > 0 {
|
||||
fmt.Printf("GPU hardware: %s\n", strings.Join(resp.GPUModels, ", "))
|
||||
} else {
|
||||
fmt.Printf("GPU hardware: none detected\n")
|
||||
}
|
||||
fmt.Printf("IOMMU: %s\n", iommu)
|
||||
fmt.Printf("vfio-pci: %s\n", vfio)
|
||||
|
||||
if resp.GPUPassthrough {
|
||||
fmt.Printf("Passthrough: enabled\n")
|
||||
fmt.Printf("GPU pool: %d/%d allocated\n", resp.AllocGPUs, resp.TotalGPUs)
|
||||
|
||||
var g5Types []string
|
||||
for _, cap := range resp.InstanceTypes {
|
||||
if strings.HasPrefix(cap.Name, "g5.") {
|
||||
g5Types = append(g5Types, cap.Name)
|
||||
}
|
||||
}
|
||||
if len(g5Types) > 0 {
|
||||
fmt.Printf("Instance types: %s\n", strings.Join(g5Types, " "))
|
||||
}
|
||||
} else if resp.GPUCapable {
|
||||
fmt.Printf("Passthrough: disabled\n")
|
||||
fmt.Printf("Instance types: (none — run 'spx admin gpu enable' to activate g5.*)\n")
|
||||
} else {
|
||||
fmt.Printf("Passthrough: disabled\n")
|
||||
fmt.Printf("Instance types: (prerequisites not met)\n")
|
||||
}
|
||||
}
|
||||
|
||||
// gpuToggle writes the TOML setting, sends SIGHUP, and polls until the daemon confirms the new state.
|
||||
// enable/disable always operate on the local node; run the command directly on the target host.
|
||||
func gpuToggle(_ *cobra.Command, enable bool) {
|
||||
cfgPath := viper.GetString("config")
|
||||
if cfgPath == "" {
|
||||
cfgPath = DefaultConfigFile()
|
||||
}
|
||||
cfg, err := config.LoadConfig(cfgPath)
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "Error loading config: %v\n", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
localNode := cfg.Node
|
||||
|
||||
// Check current state first.
|
||||
resp, err := gpuNodeStatus(localNode)
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "Error: %v\n", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
if enable {
|
||||
if resp.GPUPassthrough {
|
||||
fmt.Println("GPU passthrough is already enabled.")
|
||||
return
|
||||
}
|
||||
if !resp.GPUCapable {
|
||||
fmt.Fprintln(os.Stderr, "Error: prerequisites not met on this node.")
|
||||
fmt.Fprintln(os.Stderr, " Run 'sudo spx admin gpu setup' to configure the host.")
|
||||
os.Exit(1)
|
||||
}
|
||||
} else {
|
||||
if !resp.GPUPassthrough {
|
||||
fmt.Println("GPU passthrough is already disabled.")
|
||||
return
|
||||
}
|
||||
if resp.AllocGPUs > 0 {
|
||||
fmt.Fprintf(os.Stderr, "Error: %d GPU instance(s) are running. Terminate them first.\n", resp.AllocGPUs)
|
||||
os.Exit(1)
|
||||
}
|
||||
}
|
||||
|
||||
// Write the TOML setting.
|
||||
if err := admin.SetGPUPassthrough(cfgPath, localNode, enable); err != nil {
|
||||
fmt.Fprintf(os.Stderr, "Error writing config: %v\n", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
// Signal the daemon to reload.
|
||||
out, err := exec.Command("systemctl", "kill", "-s", "HUP", "spinifex-daemon").CombinedOutput()
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "Error sending SIGHUP: %v\n%s\n", err, out)
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
// Poll until the daemon confirms the new state.
|
||||
action := "enable"
|
||||
if !enable {
|
||||
action = "disable"
|
||||
}
|
||||
fmt.Printf("Waiting for daemon to %s GPU passthrough", action)
|
||||
deadline := time.Now().Add(30 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
time.Sleep(1 * time.Second)
|
||||
fmt.Print(".")
|
||||
r, err := gpuNodeStatus(localNode)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
if r.GPUPassthrough == enable {
|
||||
fmt.Println(" done.")
|
||||
runAdminGpuStatus(adminGpuStatusCmd, nil)
|
||||
return
|
||||
}
|
||||
}
|
||||
fmt.Println()
|
||||
fmt.Fprintln(os.Stderr, "Timed out waiting for daemon. Check: journalctl -u spinifex-daemon -n 50")
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
func runAdminGpuEnable(cmd *cobra.Command, _ []string) {
|
||||
gpuToggle(cmd, true)
|
||||
}
|
||||
|
||||
func runAdminGpuDisable(cmd *cobra.Command, _ []string) {
|
||||
gpuToggle(cmd, false)
|
||||
}
|
||||
|
||||
func runAdminGpuSetup(_ *cobra.Command, _ []string) {
|
||||
if os.Getuid() != 0 {
|
||||
fmt.Fprintln(os.Stderr, "Error: spx admin gpu setup must be run as root.")
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
fmt.Println("==> Detecting GPU")
|
||||
devices, err := gpu.Discover()
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "GPU discovery failed: %v\n", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
if len(devices) == 0 {
|
||||
fmt.Fprintln(os.Stderr, "No NVIDIA or AMD GPU found — is the card seated?")
|
||||
os.Exit(1)
|
||||
}
|
||||
for _, d := range devices {
|
||||
fmt.Printf(" %s - %s\n", d.PCIAddress, d.Model)
|
||||
}
|
||||
|
||||
fmt.Println("==> Collecting PCI IDs for vfio-pci")
|
||||
ids := gpuSetupCollectVFIOIDs(devices)
|
||||
idsCSV := strings.Join(ids, ",")
|
||||
fmt.Printf(" IDs: %s\n", idsCSV)
|
||||
|
||||
rebootNeeded := false
|
||||
|
||||
fmt.Println("==> Checking IOMMU")
|
||||
iommuEntries, _ := os.ReadDir("/sys/kernel/iommu_groups/")
|
||||
if len(iommuEntries) > 0 {
|
||||
fmt.Println(" Active")
|
||||
} else {
|
||||
params := gpuSetupIOMMUParams()
|
||||
changed, err := gpuSetupAddGRUBParams(params)
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "Failed to update GRUB: %v\n", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
if changed {
|
||||
fmt.Printf(" GRUB updated with: %s\n", params)
|
||||
} else {
|
||||
fmt.Println(" GRUB already has IOMMU params — refreshing")
|
||||
}
|
||||
ug := exec.Command("update-grub")
|
||||
ug.Stdout, ug.Stderr = os.Stdout, os.Stderr
|
||||
if err := ug.Run(); err != nil {
|
||||
fmt.Fprintf(os.Stderr, "update-grub failed: %v\n", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
rebootNeeded = true
|
||||
}
|
||||
|
||||
fmt.Println("==> Configuring vfio udev rule")
|
||||
const vfioUdevRule = "/etc/udev/rules.d/99-spinifex-vfio.rules"
|
||||
if _, err := os.Stat(vfioUdevRule); os.IsNotExist(err) {
|
||||
rule := "SUBSYSTEM==\"vfio\", GROUP=\"spinifex\", MODE=\"0660\"\n"
|
||||
if err := os.WriteFile(vfioUdevRule, []byte(rule), 0o644); err != nil {
|
||||
fmt.Fprintf(os.Stderr, "Failed to write vfio udev rule: %v\n", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
fmt.Println(" Rule installed")
|
||||
_ = exec.Command("udevadm", "control", "--reload-rules").Run()
|
||||
_ = exec.Command("udevadm", "trigger", "--subsystem-match=vfio").Run()
|
||||
} else {
|
||||
fmt.Println(" Rule already present")
|
||||
}
|
||||
|
||||
fmt.Println("==> Configuring nouveau blacklist")
|
||||
const nouveauConf = "/etc/modprobe.d/blacklist-nouveau.conf"
|
||||
if _, err := os.Stat(nouveauConf); os.IsNotExist(err) {
|
||||
if err := os.WriteFile(nouveauConf, []byte("blacklist nouveau\noptions nouveau modeset=0\n"), 0o644); err != nil {
|
||||
fmt.Fprintf(os.Stderr, "Failed to write nouveau blacklist: %v\n", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
fmt.Println(" nouveau blacklisted")
|
||||
rebootNeeded = true
|
||||
} else {
|
||||
fmt.Println(" Already blacklisted")
|
||||
}
|
||||
|
||||
fmt.Println("==> Configuring vfio-pci early binding")
|
||||
const vfioPCIConf = "/etc/modprobe.d/vfio-pci.conf"
|
||||
existing, _ := os.ReadFile(vfioPCIConf)
|
||||
if !strings.Contains(string(existing), "ids="+idsCSV) {
|
||||
if err := os.WriteFile(vfioPCIConf, []byte("options vfio-pci ids="+idsCSV+"\n"), 0o644); err != nil {
|
||||
fmt.Fprintf(os.Stderr, "Failed to write vfio-pci.conf: %v\n", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
fmt.Printf(" vfio-pci.conf written (ids=%s)\n", idsCSV)
|
||||
rebootNeeded = true
|
||||
} else {
|
||||
fmt.Printf(" Already configured for %s\n", idsCSV)
|
||||
}
|
||||
|
||||
const initramfsModules = "/etc/initramfs-tools/modules"
|
||||
modData, _ := os.ReadFile(initramfsModules)
|
||||
if !strings.Contains(string(modData), "vfio_pci") {
|
||||
f, err := os.OpenFile(initramfsModules, os.O_APPEND|os.O_WRONLY|os.O_CREATE, 0o644)
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "Failed to open initramfs modules: %v\n", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
_, werr := fmt.Fprint(f, "\n# vfio early binding for GPU passthrough\nvfio\nvfio_iommu_type1\nvfio_pci\n")
|
||||
f.Close()
|
||||
if werr != nil {
|
||||
fmt.Fprintf(os.Stderr, "Failed to write initramfs modules: %v\n", werr)
|
||||
os.Exit(1)
|
||||
}
|
||||
fmt.Println(" vfio modules added to initramfs")
|
||||
rebootNeeded = true
|
||||
} else {
|
||||
fmt.Println(" vfio modules already in initramfs")
|
||||
}
|
||||
|
||||
if rebootNeeded {
|
||||
fmt.Println("==> Updating initramfs")
|
||||
ui := exec.Command("update-initramfs", "-u")
|
||||
ui.Stdout, ui.Stderr = os.Stdout, os.Stderr
|
||||
if err := ui.Run(); err != nil {
|
||||
fmt.Fprintf(os.Stderr, "update-initramfs failed: %v\n", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
fmt.Println("\nSetup complete — reboot required.")
|
||||
fmt.Println(" sudo reboot")
|
||||
fmt.Println(" Then re-run: sudo spx admin gpu setup")
|
||||
return
|
||||
}
|
||||
|
||||
fmt.Println("==> Verifying vfio-pci module")
|
||||
if _, err := os.Stat("/sys/module/vfio_pci"); os.IsNotExist(err) {
|
||||
fmt.Fprintln(os.Stderr, "vfio_pci not loaded — check: journalctl -b | grep vfio")
|
||||
os.Exit(1)
|
||||
}
|
||||
fmt.Println(" Loaded")
|
||||
|
||||
fmt.Println("==> Verifying GPU driver binding")
|
||||
for _, d := range devices {
|
||||
driverLink := "/sys/bus/pci/devices/" + d.PCIAddress + "/driver"
|
||||
target, err := os.Readlink(driverLink)
|
||||
driver := filepath.Base(target)
|
||||
if err != nil || driver != "vfio-pci" {
|
||||
fmt.Fprintf(os.Stderr, "%s is bound to %q, not vfio-pci\n", d.PCIAddress, driver)
|
||||
fmt.Fprintf(os.Stderr, " Check: ls -la %s\n", driverLink)
|
||||
os.Exit(1)
|
||||
}
|
||||
fmt.Printf(" %s - vfio-pci\n", d.PCIAddress)
|
||||
}
|
||||
|
||||
fmt.Println("==> Enabling GPU passthrough")
|
||||
gpuToggle(nil, true)
|
||||
}
|
||||
|
||||
// gpuSetupCollectVFIOIDs collects vendor:device PCI ID pairs for each GPU and
|
||||
// all sibling PCI functions on the same device slot (e.g. GPU + HDMI audio).
|
||||
// All siblings must bind to vfio-pci together for IOMMU group isolation.
|
||||
func gpuSetupCollectVFIOIDs(devices []gpu.GPUDevice) []string {
|
||||
seen := make(map[string]bool)
|
||||
var ids []string
|
||||
for _, d := range devices {
|
||||
dot := strings.LastIndex(d.PCIAddress, ".")
|
||||
if dot < 0 {
|
||||
continue
|
||||
}
|
||||
prefix := d.PCIAddress[:dot]
|
||||
matches, _ := filepath.Glob("/sys/bus/pci/devices/" + prefix + ".*")
|
||||
for _, m := range matches {
|
||||
vb, err1 := os.ReadFile(m + "/vendor")
|
||||
db, err2 := os.ReadFile(m + "/device")
|
||||
if err1 != nil || err2 != nil {
|
||||
continue
|
||||
}
|
||||
v := strings.TrimPrefix(strings.ToLower(strings.TrimSpace(string(vb))), "0x")
|
||||
dv := strings.TrimPrefix(strings.ToLower(strings.TrimSpace(string(db))), "0x")
|
||||
id := v + ":" + dv
|
||||
if !seen[id] {
|
||||
seen[id] = true
|
||||
ids = append(ids, id)
|
||||
}
|
||||
}
|
||||
}
|
||||
return ids
|
||||
}
|
||||
|
||||
func gpuSetupIOMMUParams() string {
|
||||
data, _ := os.ReadFile("/proc/cpuinfo")
|
||||
if strings.Contains(string(data), "GenuineIntel") {
|
||||
return "intel_iommu=on iommu=pt"
|
||||
}
|
||||
return "amd_iommu=on iommu=pt"
|
||||
}
|
||||
|
||||
// gpuSetupAddGRUBParams appends iommu kernel params to GRUB_CMDLINE_LINUX_DEFAULT.
|
||||
// Returns (true, nil) if the file was modified, (false, nil) if params were already present.
|
||||
func gpuSetupAddGRUBParams(params string) (bool, error) {
|
||||
const grubFile = "/etc/default/grub"
|
||||
data, err := os.ReadFile(grubFile)
|
||||
if err != nil {
|
||||
return false, fmt.Errorf("read %s: %w", grubFile, err)
|
||||
}
|
||||
content := string(data)
|
||||
if strings.Contains(content, "intel_iommu=on") || strings.Contains(content, "amd_iommu=on") {
|
||||
return false, nil
|
||||
}
|
||||
re := regexp.MustCompile(`(GRUB_CMDLINE_LINUX_DEFAULT="[^"]*)"`)
|
||||
updated := re.ReplaceAllStringFunc(content, func(match string) string {
|
||||
groups := re.FindStringSubmatch(match)
|
||||
if len(groups) < 2 {
|
||||
return match
|
||||
}
|
||||
return groups[1] + " " + params + `"`
|
||||
})
|
||||
if updated == content {
|
||||
return false, fmt.Errorf("GRUB_CMDLINE_LINUX_DEFAULT not found in %s", grubFile)
|
||||
}
|
||||
if err := os.WriteFile(grubFile, []byte(updated), 0o644); err != nil {
|
||||
return false, fmt.Errorf("write %s: %w", grubFile, err)
|
||||
}
|
||||
return true, nil
|
||||
}
|
||||
@@ -67,6 +67,9 @@ dev_networking = false
|
||||
{{- else}}
|
||||
dev_networking = true
|
||||
{{- end}}
|
||||
{{- if .GPUPassthrough}}
|
||||
gpu_passthrough = true
|
||||
{{- end}}
|
||||
|
||||
[nodes.{{.Node}}.awsgw]
|
||||
host = "{{.BindIP}}:9999"
|
||||
|
||||
@@ -88,7 +88,7 @@ func runTopNodes(cmd *cobra.Command, args []string) {
|
||||
sort.Strings(nodeNames)
|
||||
|
||||
nodeTable := pterm.TableData{
|
||||
{"NAME", "CPU (used/total)", "MEM (used/total)", "VMs"},
|
||||
{"NAME", "CPU (used/total)", "MEM (used/total)", "GPU (used/total)", "VMs"},
|
||||
}
|
||||
|
||||
// Aggregate instance type capacity across all nodes
|
||||
@@ -96,10 +96,17 @@ func runTopNodes(cmd *cobra.Command, args []string) {
|
||||
|
||||
for _, name := range nodeNames {
|
||||
if resp, ok := respondedNodes[name]; ok {
|
||||
gpuCol := "-"
|
||||
if resp.GPUPassthrough {
|
||||
gpuCol = fmt.Sprintf("%d/%d", resp.AllocGPUs, resp.TotalGPUs)
|
||||
} else if resp.GPUCapable {
|
||||
gpuCol = fmt.Sprintf("0/%d*", len(resp.GPUModels))
|
||||
}
|
||||
nodeTable = append(nodeTable, []string{
|
||||
resp.Node,
|
||||
fmt.Sprintf("%d/%d", resp.AllocVCPU, resp.TotalVCPU),
|
||||
fmt.Sprintf("%s/%s", formatMemGB(resp.AllocMemGB), formatMemGB(resp.TotalMemGB)),
|
||||
gpuCol,
|
||||
strconv.Itoa(resp.VMCount),
|
||||
})
|
||||
|
||||
@@ -120,6 +127,7 @@ func runTopNodes(cmd *cobra.Command, args []string) {
|
||||
"-",
|
||||
"-",
|
||||
"-",
|
||||
"-",
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,47 @@
|
||||
# Configuring GPU Passthrough for Spinifex Nodes
|
||||
|
||||
|
||||
Spinifex is designed for bare-metal hardware, edge nodes and data-centre use. Follow this guide to configure your Spinifex node to enable GPU Passthrough to its VM instances if hardware is available.
|
||||
|
||||
>**Note:** this tutorial is for x86 architecture.
|
||||
|
||||
|
||||
|
||||
### Getting Started - GPU Setup
|
||||
By default, Spinifex will probe for the presence of a GPU on its machine. After booting, the presence and availability of GPU hardware is printed in the info banner:
|
||||

|
||||
|
||||
To enable GPU passthrough, you may first need to run:
|
||||
```bash
|
||||
sudo spx admin gpu setup
|
||||
```
|
||||
Note that this command will blacklist the GPU's drivers from the host machine such that they can be used by VM instances Therefore the GPU will **NOT** be available to the host machine until after the VM instance It will also require rebooting the host machine before the changes take effect.
|
||||
|
||||
### Enable Passthrough
|
||||
After reboot, the banner message should reflect that GPU passthrough is **ready** but not **enabled**.
|
||||
|
||||

|
||||
|
||||
This will require you to run:
|
||||
```bash
|
||||
sudo spx admin gpu enable
|
||||
```
|
||||
After running this command, you should see:
|
||||
|
||||

|
||||
|
||||
The host machine is now configured to provide GPU passthrough to EC2 instances on the host machine.
|
||||
|
||||
Lastly, import the Spinifex GPU enabled AMI: [link]
|
||||
|
||||
As a final check, spin up an EC2 instance, SSH into it, and check that the host machine's GPU is visible from the EC2 instance.
|
||||
|
||||
### Multiple GPUs
|
||||
Spinifex is able to detect multiple GPUs. After running the setup commands, the total number of allocatable GPUs is printed.
|
||||
|
||||
### Mapping GPUs to EC2 instances
|
||||
By defualt, any consumer grade GPUs are mapped to the g5 EC2 instance class, which is the most basic "Accelerated Computing" Ec2 instance provided by AWS. If using a consumer GPU, running:
|
||||
```bash
|
||||
aws ec2 describe-instance-types
|
||||
```
|
||||
will return the various sizes of g5 class. All other AWS "Accelerated Computing" EC2 instance classes map directly to their associated GPU, so the Spinifex host machine will need to have the specified GPU if a specific EC2 instance class is desired.
|
||||
+144
-45
@@ -1,10 +1,11 @@
|
||||
#!/bin/bash
|
||||
set -euo pipefail
|
||||
|
||||
# build-system-image.sh — Build a minimal Alpine AMI from a manifest
|
||||
# build-system-image.sh — Build a minimal system AMI from a manifest
|
||||
#
|
||||
# Creates a pre-baked Alpine Linux image with custom packages, binaries,
|
||||
# and setup scripts installed, ready for import as a Spinifex AMI.
|
||||
# Supports Alpine Linux and Ubuntu 24.04 as base distros.
|
||||
# Creates a pre-baked image with custom packages, binaries, and setup scripts
|
||||
# installed, ready for import as a Spinifex AMI.
|
||||
#
|
||||
# Requirements: qemu-nbd, qemu-img, sudo (for mount/chroot), curl
|
||||
# Usage: ./scripts/build-system-image.sh <manifest.conf> [--import]
|
||||
@@ -46,8 +47,11 @@ done
|
||||
# shellcheck source=/dev/null
|
||||
source "$MANIFEST"
|
||||
|
||||
# Default distro is alpine for backward compatibility
|
||||
DISTRO="${DISTRO:-alpine}"
|
||||
|
||||
# Validate required manifest fields
|
||||
for field in IMAGE_NAME ALPINE_VERSION IMAGE_SIZE; do
|
||||
for field in IMAGE_NAME IMAGE_SIZE; do
|
||||
if [[ -z "${!field:-}" ]]; then
|
||||
echo "ERROR: Manifest missing required field: $field"
|
||||
exit 1
|
||||
@@ -56,16 +60,48 @@ done
|
||||
|
||||
# Derived paths
|
||||
BUILD_DIR="/tmp/${IMAGE_NAME}-image-build"
|
||||
ALPINE_IMAGE="generic_alpine-${ALPINE_VERSION}-x86_64-bios-cloudinit-r0.qcow2"
|
||||
ALPINE_URL="https://dl-cdn.alpinelinux.org/alpine/v${ALPINE_VERSION%.*}/releases/cloud/${ALPINE_IMAGE}"
|
||||
OUTPUT_IMAGE="${BUILD_DIR}/${IMAGE_NAME}-alpine.qcow2"
|
||||
OUTPUT_RAW="${BUILD_DIR}/${IMAGE_NAME}-alpine.raw"
|
||||
NBD_DEV="/dev/nbd0"
|
||||
MOUNT_DIR="${BUILD_DIR}/mnt"
|
||||
|
||||
case "$DISTRO" in
|
||||
alpine)
|
||||
if [[ -z "${ALPINE_VERSION:-}" ]]; then
|
||||
echo "ERROR: Manifest missing required field: ALPINE_VERSION"
|
||||
exit 1
|
||||
fi
|
||||
SOURCE_IMAGE="generic_alpine-${ALPINE_VERSION}-x86_64-bios-cloudinit-r0.qcow2"
|
||||
SOURCE_URL="https://dl-cdn.alpinelinux.org/alpine/v${ALPINE_VERSION%.*}/releases/cloud/${SOURCE_IMAGE}"
|
||||
OUTPUT_IMAGE="${BUILD_DIR}/${IMAGE_NAME}-alpine.qcow2"
|
||||
OUTPUT_RAW="${BUILD_DIR}/${IMAGE_NAME}-alpine.raw"
|
||||
DISTRO_VERSION="${ALPINE_VERSION}"
|
||||
ROOT_PART="${NBD_DEV}"
|
||||
;;
|
||||
ubuntu)
|
||||
if [[ -z "${UBUNTU_VERSION:-}" ]]; then
|
||||
echo "ERROR: Manifest missing required field: UBUNTU_VERSION"
|
||||
exit 1
|
||||
fi
|
||||
SOURCE_IMAGE="ubuntu-${UBUNTU_VERSION}-minimal-cloudimg-amd64.img"
|
||||
SOURCE_URL="https://cloud-images.ubuntu.com/minimal/releases/noble/release/${SOURCE_IMAGE}"
|
||||
OUTPUT_IMAGE="${BUILD_DIR}/${IMAGE_NAME}-ubuntu.qcow2"
|
||||
OUTPUT_RAW="${BUILD_DIR}/${IMAGE_NAME}-ubuntu.raw"
|
||||
DISTRO_VERSION="${UBUNTU_VERSION}"
|
||||
ROOT_PART="${NBD_DEV}p1"
|
||||
;;
|
||||
*)
|
||||
echo "ERROR: Unknown DISTRO: $DISTRO (supported: alpine, ubuntu)"
|
||||
exit 1
|
||||
;;
|
||||
esac
|
||||
|
||||
cleanup() {
|
||||
echo "Cleaning up..."
|
||||
sudo umount "${MOUNT_DIR}" 2>/dev/null || true
|
||||
# Ubuntu chroot has bind mounts that must be released first
|
||||
sudo umount "${MOUNT_DIR}/dev/pts" 2>/dev/null || true
|
||||
sudo umount "${MOUNT_DIR}/dev" 2>/dev/null || true
|
||||
sudo umount "${MOUNT_DIR}/sys" 2>/dev/null || true
|
||||
sudo umount "${MOUNT_DIR}/proc" 2>/dev/null || true
|
||||
sudo umount "${MOUNT_DIR}" 2>/dev/null || true
|
||||
sudo qemu-nbd --disconnect "${NBD_DEV}" 2>/dev/null || true
|
||||
exec 9>&- 2>/dev/null || true # release nbd lock if held
|
||||
echo "Done."
|
||||
@@ -80,7 +116,7 @@ fi
|
||||
|
||||
echo "=== System Image Builder ==="
|
||||
echo "Image: ${IMAGE_NAME} — ${IMAGE_DESCRIPTION:-}"
|
||||
echo "Alpine: ${ALPINE_VERSION}"
|
||||
echo "Distro: ${DISTRO} ${DISTRO_VERSION}"
|
||||
echo "Size: ${IMAGE_SIZE}"
|
||||
echo "Build: ${BUILD_DIR}"
|
||||
echo ""
|
||||
@@ -96,6 +132,11 @@ if ! command -v qemu-img &>/dev/null; then
|
||||
exit 1
|
||||
fi
|
||||
|
||||
if [[ "$DISTRO" == "ubuntu" ]] && ! command -v parted &>/dev/null; then
|
||||
echo "ERROR: parted not found. Install parted."
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Build binaries if BUILD_COMMANDS is set
|
||||
if [[ -n "${BUILD_COMMANDS:-}" ]]; then
|
||||
echo "Building binaries: ${BUILD_COMMANDS}"
|
||||
@@ -105,7 +146,7 @@ if [[ -n "${BUILD_COMMANDS:-}" ]]; then
|
||||
fi
|
||||
fi
|
||||
|
||||
# Verify binaries exist and are statically linked
|
||||
# Verify binaries exist; Alpine also requires static linking (uses musl)
|
||||
if [[ -n "${INSTALL_BINARIES:-}" ]]; then
|
||||
IFS=' ' read -ra BINARY_PAIRS <<< "$INSTALL_BINARIES"
|
||||
for pair in "${BINARY_PAIRS[@]}"; do
|
||||
@@ -115,7 +156,7 @@ if [[ -n "${INSTALL_BINARIES:-}" ]]; then
|
||||
echo "ERROR: Binary not found: $src_path"
|
||||
exit 1
|
||||
fi
|
||||
if ! file "$src_path" | grep -q "statically linked"; then
|
||||
if [[ "$DISTRO" == "alpine" ]] && ! file "$src_path" | grep -q "statically linked"; then
|
||||
echo "ERROR: $src_path is not statically linked (Alpine uses musl — dynamic glibc binaries will fail)"
|
||||
echo " Rebuild with: CGO_ENABLED=0 go build ..."
|
||||
exit 1
|
||||
@@ -147,7 +188,7 @@ if [[ -f "$OUTPUT_RAW" ]] && [[ $(( $(date +%s) - $(stat -c %Y "$OUTPUT_RAW") ))
|
||||
|
||||
if [[ "$DO_IMPORT" == true ]]; then
|
||||
echo "Importing as AMI..."
|
||||
IMPORT_ARGS=(--file "$OUTPUT_RAW" --distro alpine --version "${ALPINE_VERSION}" --arch x86_64)
|
||||
IMPORT_ARGS=(--file "$OUTPUT_RAW" --distro "${DISTRO}" --version "${DISTRO_VERSION}" --arch x86_64)
|
||||
if [[ -n "${SYSTEM_TAG:-}" ]]; then
|
||||
IMPORT_ARGS+=(--tag "$SYSTEM_TAG")
|
||||
fi
|
||||
@@ -158,20 +199,20 @@ fi
|
||||
|
||||
mkdir -p "$BUILD_DIR" "$MOUNT_DIR"
|
||||
|
||||
# Step 1: Download Alpine cloud image
|
||||
if [[ -f "${BUILD_DIR}/${ALPINE_IMAGE}" ]]; then
|
||||
echo "Alpine image already downloaded."
|
||||
# Step 1: Download base cloud image
|
||||
if [[ -f "${BUILD_DIR}/${SOURCE_IMAGE}" ]]; then
|
||||
echo "Base image already downloaded."
|
||||
else
|
||||
echo "Downloading Alpine ${ALPINE_VERSION} cloud image..."
|
||||
if ! curl --fail -L -o "${BUILD_DIR}/${ALPINE_IMAGE}" "$ALPINE_URL"; then
|
||||
rm -f "${BUILD_DIR}/${ALPINE_IMAGE}"
|
||||
echo "ERROR: Failed to download Alpine image from $ALPINE_URL"
|
||||
echo "Downloading ${DISTRO} ${DISTRO_VERSION} cloud image..."
|
||||
if ! curl --fail -L -o "${BUILD_DIR}/${SOURCE_IMAGE}" "$SOURCE_URL"; then
|
||||
rm -f "${BUILD_DIR}/${SOURCE_IMAGE}"
|
||||
echo "ERROR: Failed to download image from $SOURCE_URL"
|
||||
exit 1
|
||||
fi
|
||||
# Verify the download is a valid qcow2 image
|
||||
if ! qemu-img info "${BUILD_DIR}/${ALPINE_IMAGE}" &>/dev/null; then
|
||||
rm -f "${BUILD_DIR}/${ALPINE_IMAGE}"
|
||||
echo "ERROR: Downloaded file is not a valid qcow2 image"
|
||||
# Verify the download is a valid qcow2/img
|
||||
if ! qemu-img info "${BUILD_DIR}/${SOURCE_IMAGE}" &>/dev/null; then
|
||||
rm -f "${BUILD_DIR}/${SOURCE_IMAGE}"
|
||||
echo "ERROR: Downloaded file is not a valid disk image"
|
||||
exit 1
|
||||
fi
|
||||
fi
|
||||
@@ -179,9 +220,9 @@ fi
|
||||
# Step 2: Copy image for customization
|
||||
rm -f "$OUTPUT_IMAGE"
|
||||
echo "Copying image for customization..."
|
||||
cp "${BUILD_DIR}/${ALPINE_IMAGE}" "$OUTPUT_IMAGE"
|
||||
cp "${BUILD_DIR}/${SOURCE_IMAGE}" "$OUTPUT_IMAGE"
|
||||
|
||||
# Resize the image (Alpine cloud images are ~200MB, need room for packages)
|
||||
# Resize the image to provide room for packages
|
||||
qemu-img resize "$OUTPUT_IMAGE" "$IMAGE_SIZE"
|
||||
|
||||
# Step 3: Connect via qemu-nbd
|
||||
@@ -194,7 +235,7 @@ fi
|
||||
sudo qemu-nbd --disconnect "${NBD_DEV}" 2>/dev/null || true
|
||||
sudo qemu-nbd --connect="${NBD_DEV}" "$OUTPUT_IMAGE"
|
||||
|
||||
# Wait for the nbd device to be ready (kernel needs time to scan the block device)
|
||||
# Wait for the nbd device to be ready
|
||||
for i in $(seq 1 10); do
|
||||
if sudo blockdev --getsize64 "${NBD_DEV}" &>/dev/null; then
|
||||
break
|
||||
@@ -203,32 +244,57 @@ for i in $(seq 1 10); do
|
||||
done
|
||||
sleep 1
|
||||
|
||||
# Ubuntu images have a partition table; wait for partition devices then resize
|
||||
# the partition to fill the expanded image before resizing the filesystem.
|
||||
if [[ "$DISTRO" == "ubuntu" ]]; then
|
||||
for i in $(seq 1 10); do
|
||||
if [[ -e "${ROOT_PART}" ]]; then break; fi
|
||||
sleep 1
|
||||
done
|
||||
echo "Resizing partition..."
|
||||
# After qemu-img resize the GPT backup header sits at the old end of disk.
|
||||
# Move it to the new end before parted tries to resize, otherwise parted
|
||||
# refuses with "Unable to satisfy all constraints on the partition."
|
||||
sudo sgdisk --move-second-header "${NBD_DEV}" 2>/dev/null || true
|
||||
sudo parted --script "${NBD_DEV}" resizepart 1 100%
|
||||
fi
|
||||
|
||||
# Alpine cloud images have ext4 directly on the block device (no partition table).
|
||||
# Resize the filesystem to fill the resized image.
|
||||
echo "Checking filesystem..."
|
||||
sudo e2fsck -f -y "${NBD_DEV}" || {
|
||||
sudo e2fsck -f -y "${ROOT_PART}" || {
|
||||
ec=$?
|
||||
if [[ $ec -gt 1 ]]; then
|
||||
echo "ERROR: e2fsck failed with exit code $ec on ${NBD_DEV}"
|
||||
echo "ERROR: e2fsck failed with exit code $ec on ${ROOT_PART}"
|
||||
exit 1
|
||||
fi
|
||||
}
|
||||
|
||||
echo "Resizing filesystem..."
|
||||
if ! sudo resize2fs "${NBD_DEV}"; then
|
||||
echo "ERROR: resize2fs failed on ${NBD_DEV}"
|
||||
if ! sudo resize2fs "${ROOT_PART}"; then
|
||||
echo "ERROR: resize2fs failed on ${ROOT_PART}"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Step 4: Mount and customize
|
||||
echo "Mounting root filesystem..."
|
||||
sudo mount "${NBD_DEV}" "$MOUNT_DIR"
|
||||
sudo mount "${ROOT_PART}" "$MOUNT_DIR"
|
||||
|
||||
# Set up resolv.conf for DNS inside chroot
|
||||
# Set up resolv.conf for DNS inside chroot.
|
||||
# Ubuntu cloud images symlink /etc/resolv.conf → /run/systemd/resolve/stub-resolv.conf
|
||||
# which doesn't exist yet. Remove the symlink and write a real file.
|
||||
sudo rm -f "${MOUNT_DIR}/etc/resolv.conf"
|
||||
sudo cp /etc/resolv.conf "${MOUNT_DIR}/etc/resolv.conf"
|
||||
|
||||
# Ubuntu chroot requires /proc /sys /dev bind mounts for systemd and DKMS
|
||||
if [[ "$DISTRO" == "ubuntu" ]]; then
|
||||
sudo mount --bind /proc "${MOUNT_DIR}/proc"
|
||||
sudo mount --bind /sys "${MOUNT_DIR}/sys"
|
||||
sudo mount --bind /dev "${MOUNT_DIR}/dev"
|
||||
sudo mount --bind /dev/pts "${MOUNT_DIR}/dev/pts"
|
||||
fi
|
||||
|
||||
# Install packages
|
||||
if [[ -n "${APK_PACKAGES:-}" ]]; then
|
||||
if [[ "$DISTRO" == "alpine" ]] && [[ -n "${APK_PACKAGES:-}" ]]; then
|
||||
echo "Installing packages: ${APK_PACKAGES}..."
|
||||
sudo chroot "$MOUNT_DIR" /bin/sh -c "
|
||||
set -e
|
||||
@@ -244,16 +310,30 @@ fi
|
||||
apk update
|
||||
apk add ${APK_PACKAGES}
|
||||
"
|
||||
elif [[ "$DISTRO" == "ubuntu" ]] && [[ -n "${APT_PACKAGES:-}" ]]; then
|
||||
echo "Installing packages: ${APT_PACKAGES}..."
|
||||
sudo chroot "$MOUNT_DIR" /bin/bash -c "
|
||||
export DEBIAN_FRONTEND=noninteractive
|
||||
apt-get update
|
||||
apt-get install -y --no-install-recommends ${APT_PACKAGES}
|
||||
"
|
||||
fi
|
||||
|
||||
# Enable OpenRC services
|
||||
# Enable services
|
||||
if [[ -n "${ENABLE_SERVICES:-}" ]]; then
|
||||
echo "Enabling services: ${ENABLE_SERVICES}..."
|
||||
IFS=' ' read -ra SERVICES <<< "$ENABLE_SERVICES"
|
||||
for svc in "${SERVICES[@]}"; do
|
||||
if ! sudo chroot "$MOUNT_DIR" /bin/sh -c "rc-update add ${svc} default"; then
|
||||
echo "ERROR: Failed to enable service '${svc}' — does it exist in the image?"
|
||||
exit 1
|
||||
if [[ "$DISTRO" == "alpine" ]]; then
|
||||
if ! sudo chroot "$MOUNT_DIR" /bin/sh -c "rc-update add ${svc} default"; then
|
||||
echo "ERROR: Failed to enable service '${svc}' — does it exist in the image?"
|
||||
exit 1
|
||||
fi
|
||||
else
|
||||
if ! sudo chroot "$MOUNT_DIR" /bin/bash -c "systemctl enable ${svc}"; then
|
||||
echo "ERROR: Failed to enable service '${svc}'"
|
||||
exit 1
|
||||
fi
|
||||
fi
|
||||
done
|
||||
fi
|
||||
@@ -282,21 +362,40 @@ if [[ -n "${SETUP_SCRIPT:-}" ]]; then
|
||||
echo "Running setup script: ${SETUP_SCRIPT}..."
|
||||
sudo cp "$setup_path" "${MOUNT_DIR}/tmp/setup.sh"
|
||||
sudo chmod 755 "${MOUNT_DIR}/tmp/setup.sh"
|
||||
sudo chroot "$MOUNT_DIR" /tmp/setup.sh
|
||||
if [[ "$DISTRO" == "ubuntu" ]]; then
|
||||
sudo chroot "$MOUNT_DIR" /bin/bash /tmp/setup.sh
|
||||
else
|
||||
sudo chroot "$MOUNT_DIR" /tmp/setup.sh
|
||||
fi
|
||||
sudo rm -f "${MOUNT_DIR}/tmp/setup.sh"
|
||||
fi
|
||||
|
||||
# Step 6: Clean up and unmount
|
||||
echo "Cleaning up image..."
|
||||
sudo chroot "$MOUNT_DIR" /bin/sh -c '
|
||||
if [[ "$DISTRO" == "alpine" ]]; then
|
||||
sudo chroot "$MOUNT_DIR" /bin/sh -c '
|
||||
apk cache clean 2>/dev/null || true
|
||||
rm -rf /var/cache/apk/* /tmp/*
|
||||
'
|
||||
else
|
||||
sudo chroot "$MOUNT_DIR" /bin/bash -c '
|
||||
export DEBIAN_FRONTEND=noninteractive
|
||||
apt-get clean
|
||||
rm -rf /var/lib/apt/lists/* /tmp/*
|
||||
'
|
||||
fi
|
||||
|
||||
# Restore original resolv.conf (cloud-init will set it on boot)
|
||||
# Restore the systemd-resolved symlink (Ubuntu default); cloud-init sets DNS on boot.
|
||||
sudo rm -f "${MOUNT_DIR}/etc/resolv.conf"
|
||||
sudo ln -sf /run/systemd/resolve/stub-resolv.conf "${MOUNT_DIR}/etc/resolv.conf"
|
||||
|
||||
echo "Unmounting..."
|
||||
if [[ "$DISTRO" == "ubuntu" ]]; then
|
||||
sudo umount "${MOUNT_DIR}/dev/pts"
|
||||
sudo umount "${MOUNT_DIR}/dev"
|
||||
sudo umount "${MOUNT_DIR}/sys"
|
||||
sudo umount "${MOUNT_DIR}/proc"
|
||||
fi
|
||||
sudo umount "$MOUNT_DIR"
|
||||
sudo qemu-nbd --disconnect "${NBD_DEV}"
|
||||
|
||||
@@ -318,7 +417,7 @@ echo ""
|
||||
|
||||
if [[ "$DO_IMPORT" == true ]]; then
|
||||
echo "Importing as AMI..."
|
||||
IMPORT_ARGS=(--file "$OUTPUT_RAW" --distro alpine --version "${ALPINE_VERSION}" --arch x86_64)
|
||||
IMPORT_ARGS=(--file "$OUTPUT_RAW" --distro "${DISTRO}" --version "${DISTRO_VERSION}" --arch x86_64)
|
||||
if [[ -n "${SYSTEM_TAG:-}" ]]; then
|
||||
IMPORT_ARGS+=(--tag "$SYSTEM_TAG")
|
||||
fi
|
||||
@@ -328,9 +427,9 @@ else
|
||||
echo " cd $PROJECT_DIR && ./bin/spx admin images import \\"
|
||||
echo " --file $OUTPUT_RAW \\"
|
||||
if [[ -n "${SYSTEM_TAG:-}" ]]; then
|
||||
echo " --distro alpine --version ${ALPINE_VERSION} --arch x86_64 \\"
|
||||
echo " --distro ${DISTRO} --version ${DISTRO_VERSION} --arch x86_64 \\"
|
||||
echo " --tag ${SYSTEM_TAG}"
|
||||
else
|
||||
echo " --distro alpine --version ${ALPINE_VERSION} --arch x86_64"
|
||||
echo " --distro ${DISTRO} --version ${DISTRO_VERSION} --arch x86_64"
|
||||
fi
|
||||
fi
|
||||
|
||||
@@ -33,6 +33,8 @@ mkdir -p "$STAGING/systemd"
|
||||
cp "$PROJECT_ROOT/build/systemd/"* "$STAGING/systemd/"
|
||||
mkdir -p "$STAGING/scripts"
|
||||
cp "$PROJECT_ROOT/build/scripts/"* "$STAGING/scripts/"
|
||||
mkdir -p "$STAGING/udev"
|
||||
cp "$PROJECT_ROOT/build/udev/"* "$STAGING/udev/"
|
||||
cp "$PROJECT_ROOT/build/logrotate/spinifex" "$STAGING/logrotate-spinifex"
|
||||
tar czf /tmp/spinifex-local.tar.gz -C "$STAGING" .
|
||||
|
||||
|
||||
Executable
+82
@@ -0,0 +1,82 @@
|
||||
#!/bin/bash
|
||||
set -euo pipefail
|
||||
|
||||
# ubuntu-gpu-setup.sh — Chroot setup for Ubuntu 24.04 GPU AMI
|
||||
#
|
||||
# Installs NVIDIA headless server driver with DKMS module pre-built against
|
||||
# a pinned kernel. Blacklists nouveau. Prevents unattended-upgrades from
|
||||
# invalidating the pre-built module.
|
||||
#
|
||||
# Must run inside an Ubuntu 24.04 chroot with /proc /sys /dev bind-mounted
|
||||
# (required by DKMS during the nvidia package post-install hook).
|
||||
|
||||
export DEBIAN_FRONTEND=noninteractive
|
||||
|
||||
apt-get update
|
||||
|
||||
# Install a pinned kernel and its headers so DKMS has a stable build target.
|
||||
# linux-image-generic and linux-headers-generic pull in the current HWE stack.
|
||||
apt-get install -y --no-install-recommends \
|
||||
linux-image-generic \
|
||||
linux-headers-generic
|
||||
|
||||
# Capture the installed kernel version before adding the driver.
|
||||
KVER=$(ls /boot/vmlinuz-* 2>/dev/null | sort -V | tail -1 | sed 's|/boot/vmlinuz-||')
|
||||
if [[ -z "${KVER}" ]]; then
|
||||
echo "ERROR: No kernel found under /boot after linux-image-generic install"
|
||||
exit 1
|
||||
fi
|
||||
echo "Target kernel: ${KVER}"
|
||||
|
||||
# Install the NVIDIA headless server driver.
|
||||
# The nvidia-dkms-550-server post-install hook runs `dkms autoinstall`, which
|
||||
# builds the module against every kernel found under /lib/modules/ — including
|
||||
# ${KVER} whose headers we just installed above.
|
||||
apt-get install -y --no-install-recommends \
|
||||
nvidia-headless-550-server \
|
||||
nvidia-utils-550-server
|
||||
|
||||
# Verify the DKMS module was registered and built.
|
||||
if dkms status 2>/dev/null | grep -q "nvidia"; then
|
||||
echo "NVIDIA DKMS module status: $(dkms status 2>/dev/null | grep nvidia | head -1)"
|
||||
else
|
||||
echo "WARNING: nvidia DKMS module not found in dkms status — manual verification required"
|
||||
fi
|
||||
|
||||
# Pin the kernel so unattended-upgrades cannot replace it and invalidate
|
||||
# the pre-built DKMS module.
|
||||
apt-mark hold \
|
||||
linux-image-generic \
|
||||
linux-headers-generic \
|
||||
"linux-image-${KVER}" \
|
||||
"linux-headers-${KVER}"
|
||||
|
||||
# Pin all installed NVIDIA packages for the same reason.
|
||||
NVIDIA_PKGS=$(dpkg --get-selections 'nvidia-*' 2>/dev/null | awk '/install$/{print $1}' || true)
|
||||
LIBNVIDIA_PKGS=$(dpkg --get-selections 'libnvidia-*' 2>/dev/null | awk '/install$/{print $1}' || true)
|
||||
# shellcheck disable=SC2086
|
||||
[[ -n "${NVIDIA_PKGS}" ]] && apt-mark hold ${NVIDIA_PKGS}
|
||||
# shellcheck disable=SC2086
|
||||
[[ -n "${LIBNVIDIA_PKGS}" ]] && apt-mark hold ${LIBNVIDIA_PKGS}
|
||||
|
||||
# Blacklist nouveau so it does not conflict with the NVIDIA driver on boot.
|
||||
cat > /etc/modprobe.d/blacklist-nouveau.conf << 'EOF'
|
||||
blacklist nouveau
|
||||
options nouveau modeset=0
|
||||
EOF
|
||||
|
||||
# Prevent unattended-upgrades from updating kernel or NVIDIA packages,
|
||||
# which would break the pre-built DKMS module.
|
||||
mkdir -p /etc/apt/apt.conf.d
|
||||
cat > /etc/apt/apt.conf.d/99-gpu-ami << 'EOF'
|
||||
Unattended-Upgrade::Package-Blacklist {
|
||||
"linux-";
|
||||
"nvidia-";
|
||||
"libnvidia-";
|
||||
};
|
||||
EOF
|
||||
|
||||
# Rebuild initramfs with nouveau blacklisted and NVIDIA module included.
|
||||
update-initramfs -u -k "${KVER}"
|
||||
|
||||
echo "GPU AMI setup complete: kernel=${KVER}, NVIDIA 550 driver installed, DKMS module pre-built"
|
||||
@@ -0,0 +1,12 @@
|
||||
IMAGE_NAME="ubuntu-gpu"
|
||||
IMAGE_DESCRIPTION="Ubuntu 24.04 GPU-ready AMI (NVIDIA A10G / G5 family)"
|
||||
DISTRO="ubuntu"
|
||||
UBUNTU_VERSION="24.04"
|
||||
|
||||
# 12G provides headroom for kernel headers, NVIDIA driver, and DKMS build artefacts
|
||||
IMAGE_SIZE="12G"
|
||||
|
||||
# NVIDIA driver install and DKMS pre-build are handled in the setup script
|
||||
SETUP_SCRIPT="scripts/images/ubuntu-gpu-setup.sh"
|
||||
|
||||
SYSTEM_TAG="gpu"
|
||||
+26
-4
@@ -15,7 +15,7 @@
|
||||
# VERBOSE Set to 1 to echo "[setup] <stage>" before each top-level step.
|
||||
# SETUP_STAGES Comma-separated subset of stages to run:
|
||||
# deps, aws, users, sudoers, files, directories,
|
||||
# env, systemd, logrotate, fixown, migrations
|
||||
# env, systemd, logrotate, udev, fixown, migrations
|
||||
# Unset = run every stage appropriate for the current mode.
|
||||
|
||||
set -e
|
||||
@@ -192,6 +192,7 @@ spinifex-vpcd ALL=(root) NOPASSWD: /usr/sbin/dhcpcd
|
||||
spinifex-vpcd ALL=(root) NOPASSWD: /usr/bin/ovs-vsctl, /usr/bin/ovs-appctl
|
||||
spinifex-vpcd ALL=(root) NOPASSWD: /usr/bin/ovn-nbctl, /usr/bin/ovn-sbctl
|
||||
spinifex-vpcd ALL=(root) NOPASSWD: /usr/bin/systemctl is-active --quiet ovn-controller
|
||||
spinifex-vpcd ALL=(root) NOPASSWD: /sbin/ip, /usr/sbin/ip
|
||||
SUDOERS
|
||||
$SUDO chmod 0440 /etc/sudoers.d/spinifex-network
|
||||
$SUDO visudo -cf /etc/sudoers.d/spinifex-network || fatal "Invalid sudoers syntax in spinifex-network"
|
||||
@@ -213,8 +214,9 @@ install_apt_deps() {
|
||||
|
||||
DEBIAN_FRONTEND=noninteractive $SUDO apt-get install -y -qq \
|
||||
nbdkit \
|
||||
$QEMU_PACKAGES qemu-utils qemu-kvm less \
|
||||
$QEMU_PACKAGES qemu-utils qemu-kvm ovmf less \
|
||||
libvirt-daemon-system libvirt-clients \
|
||||
pciutils \
|
||||
jq curl iproute2 netcat-openbsd wget unzip xz-utils file \
|
||||
ovn-central ovn-host openvswitch-switch dhcpcd-base \
|
||||
> /dev/null
|
||||
@@ -608,6 +610,24 @@ install_logrotate() {
|
||||
info "Logrotate config installed"
|
||||
}
|
||||
|
||||
# --- Install udev rules ---
|
||||
install_udev() {
|
||||
stage "installing udev rules"
|
||||
if [ ! -d "$EXTRACT_DIR/udev" ]; then
|
||||
return
|
||||
fi
|
||||
$SUDO install -d /etc/udev/rules.d
|
||||
for rule in "$EXTRACT_DIR"/udev/*; do
|
||||
$SUDO install -m 0644 "$rule" "/etc/udev/rules.d/$(basename "$rule")"
|
||||
info " /etc/udev/rules.d/$(basename "$rule")"
|
||||
done
|
||||
if [ "${ISO_BUILD:-0}" != "1" ]; then
|
||||
$SUDO udevadm control --reload-rules
|
||||
$SUDO udevadm trigger --subsystem-match=vfio 2>/dev/null || true
|
||||
fi
|
||||
info "udev rules installed"
|
||||
}
|
||||
|
||||
# --- Upgrade handling ---
|
||||
handle_upgrade() {
|
||||
if $SUDO systemctl is-active --quiet spinifex.target 2>/dev/null; then
|
||||
@@ -688,10 +708,11 @@ main() {
|
||||
handle_upgrade
|
||||
fi
|
||||
|
||||
# Stages that need EXTRACT_DIR: files, directories, systemd, logrotate.
|
||||
# Stages that need EXTRACT_DIR: files, directories, systemd, logrotate, udev.
|
||||
# Only download/extract when at least one such stage is enabled.
|
||||
if stage_enabled files || stage_enabled directories \
|
||||
|| stage_enabled systemd || stage_enabled logrotate; then
|
||||
|| stage_enabled systemd || stage_enabled logrotate \
|
||||
|| stage_enabled udev; then
|
||||
download_spinifex
|
||||
fi
|
||||
|
||||
@@ -705,6 +726,7 @@ main() {
|
||||
stage_enabled fixown && fix_file_ownership
|
||||
stage_enabled systemd && install_systemd
|
||||
stage_enabled logrotate && install_logrotate
|
||||
stage_enabled udev && install_udev
|
||||
|
||||
# Migrations are only safe on a live system (need a running NATS and a
|
||||
# persisted config file). Skip under ISO_BUILD and under any explicit
|
||||
|
||||
+77
-2
@@ -6,8 +6,23 @@
|
||||
#
|
||||
# Assumes services are running and AWS_PROFILE=spinifex is configured
|
||||
# by a prior 'spx admin init'.
|
||||
#
|
||||
# GPU passthrough test (opt-in):
|
||||
# TEST_GPU=1 scripts/smoke-test.sh
|
||||
# TEST_GPU=1 GPU_VENDOR_ID=10de GPU_DEVICE_ID=2236 scripts/smoke-test.sh
|
||||
#
|
||||
# GPU env vars:
|
||||
# TEST_GPU=1 enable GPU passthrough checks
|
||||
# GPU_VENDOR_ID PCI vendor ID (default: auto-detect first NVIDIA VGA/3D device)
|
||||
# GPU_DEVICE_ID PCI device ID (default: auto-detect)
|
||||
# GPU_FAMILY instance family to test (default: g5)
|
||||
set -euo pipefail
|
||||
|
||||
TEST_GPU="${TEST_GPU:-0}"
|
||||
GPU_VENDOR_ID="${GPU_VENDOR_ID:-}"
|
||||
GPU_DEVICE_ID="${GPU_DEVICE_ID:-}"
|
||||
GPU_FAMILY="${GPU_FAMILY:-g5}"
|
||||
|
||||
export AWS_PROFILE=spinifex
|
||||
|
||||
# --- Wait for EC2 daemon to subscribe to NATS ---
|
||||
@@ -42,6 +57,41 @@ echo "==> Importing SSH key"
|
||||
--public-key-material "fileb://$SSH_KEY.pub"
|
||||
aws ec2 describe-key-pairs
|
||||
|
||||
# --- GPU pre-flight (TEST_GPU=1 only) ---
|
||||
if [[ "$TEST_GPU" == "1" ]]; then
|
||||
echo "==> [GPU] Checking IOMMU"
|
||||
if ! dmesg | grep -qi "DMAR.*IOMMU enabled\|iommu.*Translated\|Adding to iommu group\|AMD-Vi.*enabled"; then
|
||||
echo "❌ IOMMU not detected — add intel_iommu=on (or amd_iommu=on) and iommu=pt to kernel cmdline, then reboot"
|
||||
exit 1
|
||||
fi
|
||||
echo " IOMMU active"
|
||||
|
||||
echo "==> [GPU] Detecting NVIDIA GPU"
|
||||
if [[ -z "$GPU_VENDOR_ID" || -z "$GPU_DEVICE_ID" ]]; then
|
||||
GPU_LINE=$(lspci -nn | grep -i nvidia | grep -iE "VGA|3D" | head -1)
|
||||
if [[ -z "$GPU_LINE" ]]; then
|
||||
echo "❌ No NVIDIA GPU found via lspci — is the card seated and detectable?"
|
||||
lspci | grep -i nvidia >&2 || true
|
||||
exit 1
|
||||
fi
|
||||
echo " Detected: $GPU_LINE"
|
||||
PCI_IDS=$(echo "$GPU_LINE" | grep -oP '\[\K[0-9a-f]{4}:[0-9a-f]{4}(?=\])')
|
||||
GPU_VENDOR_ID=$(echo "$PCI_IDS" | cut -d: -f1)
|
||||
GPU_DEVICE_ID=$(echo "$PCI_IDS" | cut -d: -f2)
|
||||
fi
|
||||
echo " GPU PCI ID: ${GPU_VENDOR_ID}:${GPU_DEVICE_ID}"
|
||||
|
||||
echo "==> [GPU] Verifying ${GPU_FAMILY} instance types advertised"
|
||||
G5_TYPES=$(aws ec2 describe-instance-types \
|
||||
--query "InstanceTypes[?starts_with(InstanceType, \`${GPU_FAMILY}\`)].InstanceType" \
|
||||
--output text 2>/dev/null || true)
|
||||
if [[ -z "$G5_TYPES" ]]; then
|
||||
echo "❌ No ${GPU_FAMILY}.* instance types returned by describe-instance-types"
|
||||
exit 1
|
||||
fi
|
||||
echo " ${GPU_FAMILY} types available: $G5_TYPES"
|
||||
fi
|
||||
|
||||
# --- Import AMI ---
|
||||
# --- Import AMI ---
|
||||
echo "==> Importing AMI"
|
||||
@@ -79,7 +129,9 @@ fi
|
||||
|
||||
# --- Launch smoke-test instance ---
|
||||
echo "==> Launching smoke-test instance"
|
||||
if grep -q 'AuthenticAMD' /proc/cpuinfo; then
|
||||
if [[ "$TEST_GPU" == "1" ]]; then
|
||||
INSTANCE_TYPE="${GPU_FAMILY}.xlarge"
|
||||
elif grep -q 'AuthenticAMD' /proc/cpuinfo; then
|
||||
INSTANCE_TYPE="t3a.small"
|
||||
else
|
||||
INSTANCE_TYPE="t3.small"
|
||||
@@ -195,4 +247,27 @@ if ! echo "$SSH_OUT" | grep -q "ec2-user"; then
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo "✅ Smoke test passed — instance $INSTANCE_ID launched, running, and SSH-verified"
|
||||
# --- GPU guest verification (TEST_GPU=1 only) ---
|
||||
if [[ "$TEST_GPU" == "1" ]]; then
|
||||
echo "==> [GPU] Verifying GPU visible inside guest"
|
||||
GPU_VISIBLE=0
|
||||
for _i in $(seq 1 12); do
|
||||
if ssh -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null \
|
||||
-o ConnectTimeout=5 -o BatchMode=yes \
|
||||
-p "$SSH_INST_PORT" -i "$SSH_KEY" \
|
||||
ec2-user@"$SSH_INST_HOST" \
|
||||
'lspci 2>/dev/null | grep -i nvidia' 2>/dev/null; then
|
||||
GPU_VISIBLE=1
|
||||
break
|
||||
fi
|
||||
sleep 5
|
||||
done
|
||||
if [[ $GPU_VISIBLE -eq 0 ]]; then
|
||||
echo "❌ NVIDIA GPU not visible in guest via lspci"
|
||||
exit 1
|
||||
fi
|
||||
echo " GPU visible in guest"
|
||||
echo "✅ Smoke test passed (GPU passthrough) — instance $INSTANCE_ID launched with GPU and verified"
|
||||
else
|
||||
echo "✅ Smoke test passed — instance $INSTANCE_ID launched, running, and SSH-verified"
|
||||
fi
|
||||
|
||||
@@ -15,6 +15,7 @@ import (
|
||||
"os"
|
||||
"os/user"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"strconv"
|
||||
"strings"
|
||||
"text/template"
|
||||
@@ -99,6 +100,10 @@ type ConfigSettings struct {
|
||||
BootstrapIgwId string
|
||||
BootstrapCidr string
|
||||
BootstrapSubnetCidr string
|
||||
|
||||
// GPUPassthrough enables VFIO GPU passthrough in the daemon config.
|
||||
// Sets gpu_passthrough = true under [nodes.<node>.daemon].
|
||||
GPUPassthrough bool
|
||||
}
|
||||
|
||||
// PredastoreNodeConfig describes a single Predastore node for multi-node config generation.
|
||||
@@ -931,3 +936,58 @@ func ParsePredastoreNodeIDFromConfig(tomlContent string, ip string) int {
|
||||
}
|
||||
return FindNodeIDByIP(cfg.DB, ip)
|
||||
}
|
||||
|
||||
// SetGPUPassthrough idempotently writes gpu_passthrough = <enabled> for the
|
||||
// given node into spinifex.toml, preserving all other content and comments.
|
||||
// Returns nil without touching the file if the setting is already correct.
|
||||
func SetGPUPassthrough(tomlPath, node string, enabled bool) error {
|
||||
raw, err := os.ReadFile(tomlPath)
|
||||
if err != nil {
|
||||
return fmt.Errorf("read %s: %w", tomlPath, err)
|
||||
}
|
||||
text := string(raw)
|
||||
value := "false"
|
||||
if enabled {
|
||||
value = "true"
|
||||
}
|
||||
|
||||
sectionHeader := "[nodes." + node + ".daemon]"
|
||||
|
||||
sectionStart := strings.Index(text, sectionHeader)
|
||||
if sectionStart < 0 {
|
||||
// No daemon section — append one.
|
||||
text = strings.TrimRight(text, "\n") +
|
||||
"\n\n[nodes." + node + ".daemon]\ngpu_passthrough = " + value + "\n"
|
||||
return os.WriteFile(tomlPath, []byte(text), 0640) //nolint:gosec // spinifex.toml is root:spinifex 0640 so the daemon can read it
|
||||
}
|
||||
|
||||
// Extract body of just this section (stops at the next section header).
|
||||
bodyStart := sectionStart + len(sectionHeader)
|
||||
rest := text[bodyStart:]
|
||||
nextSection := strings.Index(rest, "\n[")
|
||||
var body, suffix string
|
||||
if nextSection < 0 {
|
||||
body = rest
|
||||
} else {
|
||||
body = rest[:nextSection]
|
||||
suffix = rest[nextSection:]
|
||||
}
|
||||
|
||||
// Already correct — no-op.
|
||||
if regexp.MustCompile(`gpu_passthrough\s*=\s*` + value).MatchString(body) {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Key exists with wrong value — flip within this section only.
|
||||
flipRe := regexp.MustCompile(`gpu_passthrough\s*=\s*(?:true|false)`)
|
||||
var newBody string
|
||||
if flipRe.MatchString(body) {
|
||||
newBody = flipRe.ReplaceAllString(body, "gpu_passthrough = "+value)
|
||||
} else {
|
||||
// Key absent — insert right after section header.
|
||||
newBody = "\ngpu_passthrough = " + value + body
|
||||
}
|
||||
|
||||
text = text[:bodyStart] + newBody + suffix
|
||||
return os.WriteFile(tomlPath, []byte(text), 0640) //nolint:gosec // spinifex.toml is root:spinifex 0640 so the daemon can read it
|
||||
}
|
||||
|
||||
@@ -862,3 +862,72 @@ func TestChownRecursive_NonExistentPath(t *testing.T) {
|
||||
}
|
||||
ChownRecursive("/tmp/nonexistent-path-12345", currentUser)
|
||||
}
|
||||
|
||||
// --- SetGPUPassthrough ---
|
||||
|
||||
func writeToml(t *testing.T, content string) string {
|
||||
t.Helper()
|
||||
f, err := os.CreateTemp(t.TempDir(), "spinifex*.toml")
|
||||
require.NoError(t, err)
|
||||
_, err = f.WriteString(content)
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, f.Close())
|
||||
return f.Name()
|
||||
}
|
||||
|
||||
func readToml(t *testing.T, path string) string {
|
||||
t.Helper()
|
||||
b, err := os.ReadFile(path)
|
||||
require.NoError(t, err)
|
||||
return string(b)
|
||||
}
|
||||
|
||||
func TestSetGPUPassthrough_NoOp(t *testing.T) {
|
||||
toml := "[nodes.node1.daemon]\ngpu_passthrough = true\n"
|
||||
path := writeToml(t, toml)
|
||||
require.NoError(t, SetGPUPassthrough(path, "node1", true))
|
||||
assert.Equal(t, toml, readToml(t, path))
|
||||
}
|
||||
|
||||
func TestSetGPUPassthrough_Flip(t *testing.T) {
|
||||
path := writeToml(t, "[nodes.node1.daemon]\ngpu_passthrough = false\n")
|
||||
require.NoError(t, SetGPUPassthrough(path, "node1", true))
|
||||
assert.Contains(t, readToml(t, path), "gpu_passthrough = true")
|
||||
}
|
||||
|
||||
func TestSetGPUPassthrough_AddKeyToSection(t *testing.T) {
|
||||
path := writeToml(t, "[nodes.node1.daemon]\nsome_other = true\n")
|
||||
require.NoError(t, SetGPUPassthrough(path, "node1", true))
|
||||
got := readToml(t, path)
|
||||
assert.Contains(t, got, "gpu_passthrough = true")
|
||||
assert.Contains(t, got, "some_other = true")
|
||||
}
|
||||
|
||||
func TestSetGPUPassthrough_AppendSection(t *testing.T) {
|
||||
path := writeToml(t, "[nodes.node1.network]\ncidr = \"10.0.0.0/24\"\n")
|
||||
require.NoError(t, SetGPUPassthrough(path, "node1", true))
|
||||
got := readToml(t, path)
|
||||
assert.Contains(t, got, "[nodes.node1.daemon]")
|
||||
assert.Contains(t, got, "gpu_passthrough = true")
|
||||
}
|
||||
|
||||
func TestSetGPUPassthrough_DisableNoOp(t *testing.T) {
|
||||
toml := "[nodes.node1.daemon]\ngpu_passthrough = false\n"
|
||||
path := writeToml(t, toml)
|
||||
require.NoError(t, SetGPUPassthrough(path, "node1", false))
|
||||
assert.Equal(t, toml, readToml(t, path))
|
||||
}
|
||||
|
||||
func TestSetGPUPassthrough_ReadError(t *testing.T) {
|
||||
err := SetGPUPassthrough("/nonexistent/path/spinifex.toml", "node1", true)
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
func TestSetGPUPassthrough_SectionBoundary(t *testing.T) {
|
||||
// gpu_passthrough = true exists but in a DIFFERENT node's section; should still write.
|
||||
path := writeToml(t, "[nodes.node2.daemon]\ngpu_passthrough = true\n[nodes.node1.daemon]\nsome_key = 1\n")
|
||||
require.NoError(t, SetGPUPassthrough(path, "node1", true))
|
||||
got := readToml(t, path)
|
||||
// node1 section should now have the key
|
||||
assert.Contains(t, got, "[nodes.node1.daemon]\ngpu_passthrough = true")
|
||||
}
|
||||
|
||||
@@ -123,13 +123,30 @@ type PredastoreConfig struct {
|
||||
NodeID int `json:"NodeID" mapstructure:"node_id"`
|
||||
}
|
||||
|
||||
// GPUModelOverride maps a PCI vendor/device ID to a GPU instance family for
|
||||
// dev or test nodes that carry consumer GPUs not in the production model list.
|
||||
// Add entries under [[nodes.<node>.daemon.gpu_model_overrides]] in spinifex.toml.
|
||||
type GPUModelOverride struct {
|
||||
VendorID string `json:"VendorID" mapstructure:"vendor_id"`
|
||||
DeviceID string `json:"DeviceID" mapstructure:"device_id"`
|
||||
Family string `json:"Family" mapstructure:"family"`
|
||||
Manufacturer string `json:"Manufacturer" mapstructure:"manufacturer"`
|
||||
Name string `json:"Name" mapstructure:"name"`
|
||||
MemoryMiB int64 `json:"MemoryMiB" mapstructure:"memory_mib"`
|
||||
// XVGAOff forces x-vga=off for this GPU in QEMU passthrough, overriding the
|
||||
// default (on for consumer GPUs, off for known datacenter cards).
|
||||
XVGAOff bool `json:"XVGAOff" mapstructure:"xvga_off"`
|
||||
}
|
||||
|
||||
// DaemonConfig holds the daemon configuration
|
||||
type DaemonConfig struct {
|
||||
Host string `json:"Host" mapstructure:"host"`
|
||||
TLSKey string `json:"TLSKey" mapstructure:"tlskey"`
|
||||
TLSCert string `json:"TLSCert" mapstructure:"tlscert"`
|
||||
DevNetworking bool `json:"DevNetworking" mapstructure:"dev_networking"` // VPC instances get both TAP + hostfwd for SSH dev access
|
||||
MgmtBridge string `json:"MgmtBridge" mapstructure:"mgmt_bridge"` // Linux bridge for system instance control plane (default "br-mgmt")
|
||||
Host string `json:"Host" mapstructure:"host"`
|
||||
TLSKey string `json:"TLSKey" mapstructure:"tlskey"`
|
||||
TLSCert string `json:"TLSCert" mapstructure:"tlscert"`
|
||||
DevNetworking bool `json:"DevNetworking" mapstructure:"dev_networking"` // VPC instances get both TAP + hostfwd for SSH dev access
|
||||
MgmtBridge string `json:"MgmtBridge" mapstructure:"mgmt_bridge"` // Linux bridge for system instance control plane (default "br-mgmt")
|
||||
GPUPassthrough bool `json:"GPUPassthrough" mapstructure:"gpu_passthrough"` // Enable VFIO GPU passthrough for g5.* instance types
|
||||
GPUModelOverrides []GPUModelOverride `json:"GPUModelOverrides" mapstructure:"gpu_model_overrides"` // Dev/test GPU mappings not in the production model list
|
||||
}
|
||||
|
||||
// NATSConfig holds the NATS configuration
|
||||
|
||||
+256
-10
@@ -9,6 +9,7 @@ import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"maps"
|
||||
"net"
|
||||
"os"
|
||||
"os/exec"
|
||||
@@ -30,6 +31,7 @@ import (
|
||||
"github.com/mulgadc/spinifex/spinifex/admin"
|
||||
"github.com/mulgadc/spinifex/spinifex/awserrors"
|
||||
"github.com/mulgadc/spinifex/spinifex/config"
|
||||
"github.com/mulgadc/spinifex/spinifex/gpu"
|
||||
handlers_ec2_account "github.com/mulgadc/spinifex/spinifex/handlers/ec2/account"
|
||||
handlers_ec2_eigw "github.com/mulgadc/spinifex/spinifex/handlers/ec2/eigw"
|
||||
handlers_ec2_eip "github.com/mulgadc/spinifex/spinifex/handlers/ec2/eip"
|
||||
@@ -92,6 +94,7 @@ type ResourceManager struct {
|
||||
allocatedVCPU int
|
||||
allocatedMem float64
|
||||
instanceTypes map[string]*ec2.InstanceTypeInfo
|
||||
gpuManager *gpu.Manager // nil if GPU passthrough is disabled or no GPUs present
|
||||
|
||||
// Dynamic instance-type subscription management
|
||||
subsMu sync.Mutex
|
||||
@@ -163,6 +166,14 @@ type Daemon struct {
|
||||
// management NIC. Set when AWSGW binds to a specific IP (multi-node).
|
||||
mgmtRouteVia string
|
||||
|
||||
// gpuProbe holds the result of the always-on startup GPU hardware probe.
|
||||
// Populated regardless of whether gpu_passthrough is enabled in config.
|
||||
gpuProbe gpuProbeResult
|
||||
|
||||
// gpuManager handles VFIO bind/unbind lifecycle for GPU passthrough.
|
||||
// Nil when GPUPassthrough is disabled in config or no recognised GPUs are found.
|
||||
gpuManager *gpu.Manager
|
||||
|
||||
// shuttingDown is set to true during coordinated cluster shutdown (GATE phase)
|
||||
// or during SIGTERM-based shutdown. When true, the daemon rejects new work,
|
||||
// crash handlers bail out, and setupShutdown skips redundant VM stops.
|
||||
@@ -224,7 +235,9 @@ func getSystemMemory() (float64, error) {
|
||||
// Also returns an error if the host is too small to satisfy the daemon's
|
||||
// reserve — clamping silently would defeat the reserve and look like a
|
||||
// runtime bug.
|
||||
func NewResourceManager() (*ResourceManager, error) {
|
||||
// gpuModels is the list of recognised GPU models present on the host; pass nil if
|
||||
// GPU passthrough is disabled or no GPUs were found.
|
||||
func NewResourceManager(gpuModels []instancetypes.GPUModel, gpuMgr *gpu.Manager) (*ResourceManager, error) {
|
||||
// Get system CPU cores
|
||||
numCPU := runtime.NumCPU()
|
||||
|
||||
@@ -248,8 +261,9 @@ func NewResourceManager() (*ResourceManager, error) {
|
||||
arch = "arm64"
|
||||
}
|
||||
|
||||
// Detect CPU generation and generate matching instance types
|
||||
instanceTypes := instancetypes.DetectAndGenerate(instancetypes.HostCPU{}, arch)
|
||||
// Detect CPU generation and generate matching instance types (including GPU
|
||||
// families when gpuModels is non-nil).
|
||||
instanceTypes := instancetypes.DetectAndGenerate(instancetypes.HostCPU{}, arch, gpuModels)
|
||||
|
||||
slog.Info("System resources detected",
|
||||
"hostVCPU", numCPU, "hostMemGB", totalMemGB,
|
||||
@@ -263,6 +277,7 @@ func NewResourceManager() (*ResourceManager, error) {
|
||||
reservedVCPU: reservedVCPU,
|
||||
reservedMem: reservedMem,
|
||||
instanceTypes: instanceTypes,
|
||||
gpuManager: gpuMgr,
|
||||
}, nil
|
||||
}
|
||||
|
||||
@@ -318,11 +333,17 @@ func (rm *ResourceManager) GetAvailableInstanceTypeInfos(showCapacity bool) []*e
|
||||
continue
|
||||
}
|
||||
|
||||
requiresGPU := instancetypes.IsGPUType(it)
|
||||
availGPU := 0
|
||||
if rm.gpuManager != nil && requiresGPU {
|
||||
availGPU = rm.gpuManager.Available()
|
||||
}
|
||||
count := canAllocateCount(
|
||||
rm.hostVCPU-rm.reservedVCPU, rm.allocatedVCPU,
|
||||
rm.hostMemGB-rm.reservedMem, rm.allocatedMem,
|
||||
vCPUs, memMiB,
|
||||
1<<30, // effectively unlimited — let resources be the constraint
|
||||
availGPU, requiresGPU,
|
||||
)
|
||||
|
||||
if showCapacity {
|
||||
@@ -389,14 +410,36 @@ func NewDaemon(cfg *config.ClusterConfig) (*Daemon, error) {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
|
||||
// If WalDir is not set, use BaseDir
|
||||
config := cfg.Nodes[cfg.Node]
|
||||
nodeCfg := cfg.Nodes[cfg.Node]
|
||||
if cfg.Nodes[cfg.Node].WalDir == "" {
|
||||
config.WalDir = config.BaseDir
|
||||
|
||||
cfg.Nodes[cfg.Node] = config
|
||||
nodeCfg.WalDir = nodeCfg.BaseDir
|
||||
cfg.Nodes[cfg.Node] = nodeCfg
|
||||
}
|
||||
|
||||
rm, err := NewResourceManager()
|
||||
// Phase 1: always probe GPU hardware (no side effects, no config required).
|
||||
gpuProbe := probeGPU()
|
||||
|
||||
// Phase 2: activate GPU passthrough only when the operator has opted in.
|
||||
var gpuModels []instancetypes.GPUModel
|
||||
var gpuMgr *gpu.Manager
|
||||
if nodeCfg.Daemon.GPUPassthrough {
|
||||
if !gpuProbe.Capable {
|
||||
slog.Warn("GPU passthrough enabled in config but prerequisites not met",
|
||||
"iommu", gpuProbe.IOMMUActive, "vfio", gpuProbe.VFIOPresent,
|
||||
"gpus", len(gpuProbe.Devices))
|
||||
} else {
|
||||
for _, dev := range gpuProbe.Devices {
|
||||
gpuModels = append(gpuModels, resolveGPUModel(dev, nodeCfg.Daemon.GPUModelOverrides))
|
||||
}
|
||||
gpuMgr = gpu.NewManager(gpuProbe.Devices)
|
||||
slog.Info("GPU passthrough enabled", "gpus", len(gpuProbe.Devices), "knownModels", len(gpuModels))
|
||||
}
|
||||
} else if gpuProbe.Capable {
|
||||
slog.Info("GPU hardware detected, passthrough not enabled",
|
||||
"gpus", len(gpuProbe.Devices), "hint", "run 'spx admin gpu enable' to activate")
|
||||
}
|
||||
|
||||
rm, err := NewResourceManager(gpuModels, gpuMgr)
|
||||
if err != nil {
|
||||
cancel()
|
||||
return nil, fmt.Errorf("initialize resource manager: %w", err)
|
||||
@@ -405,8 +448,10 @@ func NewDaemon(cfg *config.ClusterConfig) (*Daemon, error) {
|
||||
return &Daemon{
|
||||
node: cfg.Node,
|
||||
clusterConfig: cfg,
|
||||
config: &config,
|
||||
config: &nodeCfg,
|
||||
resourceMgr: rm,
|
||||
gpuProbe: gpuProbe,
|
||||
gpuManager: gpuMgr,
|
||||
ctx: ctx,
|
||||
cancel: cancel,
|
||||
vmMgr: vm.NewManager(),
|
||||
@@ -876,6 +921,7 @@ func (d *Daemon) Start() error {
|
||||
d.ready.Store(true)
|
||||
slog.Info("Daemon fully initialized", "node", d.node, "startupTime", time.Since(d.startTime).Round(time.Second))
|
||||
|
||||
d.setupReload()
|
||||
d.setupShutdown()
|
||||
d.awaitShutdown()
|
||||
|
||||
@@ -1295,10 +1341,86 @@ func (d *Daemon) SendQMPCommand(q *qmp.QMPClient, cmd qmp.QMPCommand, instanceId
|
||||
}
|
||||
}
|
||||
|
||||
// setupReload registers a SIGHUP handler that reloads GPU config without restarting.
|
||||
func (d *Daemon) setupReload() {
|
||||
d.shutdownWg.Go(func() {
|
||||
sigChan := make(chan os.Signal, 1)
|
||||
signal.Notify(sigChan, syscall.SIGHUP)
|
||||
defer signal.Stop(sigChan)
|
||||
for {
|
||||
select {
|
||||
case <-d.ctx.Done():
|
||||
return
|
||||
case <-sigChan:
|
||||
slog.Info("SIGHUP received — reloading GPU config")
|
||||
d.reloadConfig()
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// reloadConfig re-reads spinifex.toml and applies any GPU passthrough changes.
|
||||
func (d *Daemon) reloadConfig() {
|
||||
if d.configPath == "" {
|
||||
slog.Warn("SIGHUP: no config path set, cannot reload")
|
||||
return
|
||||
}
|
||||
newCfg, err := config.LoadConfig(d.configPath)
|
||||
if err != nil {
|
||||
slog.Error("SIGHUP: config reload failed", "err", err)
|
||||
return
|
||||
}
|
||||
newNodeCfg := newCfg.Nodes[d.node]
|
||||
d.applyGPUConfig(newNodeCfg.Daemon.GPUPassthrough)
|
||||
}
|
||||
|
||||
// applyGPUConfig activates or deactivates GPU passthrough at runtime.
|
||||
// Transition false→true: re-probes hardware, initialises gpuManager, adds g5 types.
|
||||
// Transition true→false: refused when GPU instances are running; otherwise tears down.
|
||||
func (d *Daemon) applyGPUConfig(enabled bool) {
|
||||
d.mu.Lock()
|
||||
defer d.mu.Unlock()
|
||||
|
||||
wasEnabled := d.gpuManager != nil
|
||||
if enabled == wasEnabled {
|
||||
slog.Debug("GPU passthrough state unchanged on reload", "passthrough", enabled)
|
||||
return
|
||||
}
|
||||
|
||||
if enabled {
|
||||
probe := probeGPU()
|
||||
d.gpuProbe = probe
|
||||
if !probe.Capable {
|
||||
slog.Warn("GPU passthrough enable failed: prerequisites not met",
|
||||
"iommu", probe.IOMMUActive, "vfio", probe.VFIOPresent, "gpus", len(probe.Devices))
|
||||
return
|
||||
}
|
||||
var models []instancetypes.GPUModel
|
||||
for _, dev := range probe.Devices {
|
||||
models = append(models, resolveGPUModel(dev, d.config.Daemon.GPUModelOverrides))
|
||||
}
|
||||
mgr := gpu.NewManager(probe.Devices)
|
||||
d.gpuManager = mgr
|
||||
d.resourceMgr.reloadGPUTypes(models, mgr)
|
||||
slog.Info("GPU passthrough enabled via config reload", "gpus", len(probe.Devices))
|
||||
return
|
||||
}
|
||||
|
||||
// true → false: refuse if instances are running
|
||||
if d.gpuManager.AllocatedCount() > 0 {
|
||||
slog.Warn("GPU passthrough disable refused: GPU instances are running",
|
||||
"allocated", d.gpuManager.AllocatedCount())
|
||||
return
|
||||
}
|
||||
d.gpuManager = nil
|
||||
d.resourceMgr.reloadGPUTypes(nil, nil)
|
||||
slog.Info("GPU passthrough disabled via config reload")
|
||||
}
|
||||
|
||||
func (d *Daemon) setupShutdown() {
|
||||
d.shutdownWg.Go(func() {
|
||||
sigChan := make(chan os.Signal, 1)
|
||||
signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM, syscall.SIGHUP)
|
||||
signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
|
||||
|
||||
<-sigChan
|
||||
slog.Info("Received shutdown signal, cleaning up...")
|
||||
@@ -1390,12 +1512,19 @@ func (rm *ResourceManager) canAllocate(instanceType *ec2.InstanceTypeInfo, count
|
||||
rm.mu.RLock()
|
||||
defer rm.mu.RUnlock()
|
||||
|
||||
requiresGPU := instancetypes.IsGPUType(instanceType)
|
||||
availGPU := 0
|
||||
if rm.gpuManager != nil && requiresGPU {
|
||||
availGPU = rm.gpuManager.Available()
|
||||
}
|
||||
|
||||
return canAllocateCount(
|
||||
rm.hostVCPU-rm.reservedVCPU, rm.allocatedVCPU,
|
||||
rm.hostMemGB-rm.reservedMem, rm.allocatedMem,
|
||||
instanceTypeVCPUs(instanceType),
|
||||
instanceTypeMemoryMiB(instanceType),
|
||||
count,
|
||||
availGPU, requiresGPU,
|
||||
)
|
||||
}
|
||||
|
||||
@@ -1434,6 +1563,30 @@ func (rm *ResourceManager) deallocate(instanceType *ec2.InstanceTypeInfo) {
|
||||
|
||||
// initSubscriptions sets up dynamic per-instance-type NATS subscriptions.
|
||||
// Called once during daemon startup after NATS is connected.
|
||||
// reloadGPUTypes replaces GPU instance types in-place and updates NATS subscriptions.
|
||||
// Called on SIGHUP when gpu_passthrough is toggled. Mutates the existing map so that
|
||||
// all holders of the map reference (e.g. instanceService) see the updated types.
|
||||
func (rm *ResourceManager) reloadGPUTypes(models []instancetypes.GPUModel, mgr *gpu.Manager) {
|
||||
arch := "x86_64"
|
||||
if runtime.GOARCH == "arm64" {
|
||||
arch = "arm64"
|
||||
}
|
||||
|
||||
rm.mu.Lock()
|
||||
for name, it := range rm.instanceTypes {
|
||||
if instancetypes.IsGPUType(it) {
|
||||
delete(rm.instanceTypes, name)
|
||||
}
|
||||
}
|
||||
if len(models) > 0 {
|
||||
maps.Copy(rm.instanceTypes, instancetypes.GenerateGPUTypes(models, arch))
|
||||
}
|
||||
rm.gpuManager = mgr
|
||||
rm.mu.Unlock()
|
||||
|
||||
rm.updateInstanceSubscriptions()
|
||||
}
|
||||
|
||||
func (rm *ResourceManager) initSubscriptions(nc *nats.Conn, handler nats.MsgHandler, nodeID string) {
|
||||
rm.natsConn = nc
|
||||
rm.handler = handler
|
||||
@@ -1608,3 +1761,96 @@ func (d *Daemon) wireLBAgentConfig() {
|
||||
d.elbv2Service.MgmtBridgeIP = d.mgmtBridgeIP
|
||||
d.elbv2Service.AdvertiseIP = advertiseIP
|
||||
}
|
||||
|
||||
// resolveGPUModel maps a discovered GPU to an instance type model.
|
||||
// Overrides take priority, then the production model list, then a g5 default.
|
||||
// Any GPU device that reaches the default path is treated as a g5 instance,
|
||||
// so consumer GPUs used for testing work without explicit config entries.
|
||||
func resolveGPUModel(dev gpu.GPUDevice, overrides []config.GPUModelOverride) instancetypes.GPUModel {
|
||||
for i := range overrides {
|
||||
o := &overrides[i]
|
||||
if o.VendorID == dev.VendorID && o.DeviceID == dev.DeviceID {
|
||||
return instancetypes.GPUModel{
|
||||
VendorID: o.VendorID,
|
||||
DeviceID: o.DeviceID,
|
||||
Family: o.Family,
|
||||
Manufacturer: o.Manufacturer,
|
||||
Name: o.Name,
|
||||
MemoryMiB: o.MemoryMiB,
|
||||
}
|
||||
}
|
||||
}
|
||||
if m := instancetypes.GPUModelForVendorDevice(dev.VendorID, dev.DeviceID); m != nil {
|
||||
return *m
|
||||
}
|
||||
// Default: any discovered GPU device maps to g5 using its detected specs.
|
||||
name := dev.Model
|
||||
if name == "" {
|
||||
name = fmt.Sprintf("GPU %s:%s", dev.VendorID, dev.DeviceID)
|
||||
}
|
||||
return instancetypes.GPUModel{
|
||||
VendorID: dev.VendorID,
|
||||
DeviceID: dev.DeviceID,
|
||||
Family: "g5",
|
||||
Manufacturer: gpuVendorDisplayName(dev.Vendor),
|
||||
Name: name,
|
||||
MemoryMiB: dev.MemoryMiB,
|
||||
}
|
||||
}
|
||||
|
||||
// gpuXVGAEnabled returns true when the QEMU device should include x-vga=on.
|
||||
// Consumer GPUs default to true; known datacenter/compute cards default to false.
|
||||
// A GPUModelOverride with XVGAOff=true forces false regardless of the model table.
|
||||
func gpuXVGAEnabled(dev *gpu.GPUDevice, overrides []config.GPUModelOverride) bool {
|
||||
for _, o := range overrides {
|
||||
if o.VendorID == dev.VendorID && o.DeviceID == dev.DeviceID {
|
||||
return !o.XVGAOff
|
||||
}
|
||||
}
|
||||
return !gpu.IsComputeGPU(dev.VendorID, dev.DeviceID)
|
||||
}
|
||||
|
||||
func gpuVendorDisplayName(v gpu.Vendor) string {
|
||||
switch v {
|
||||
case gpu.VendorNVIDIA:
|
||||
return "NVIDIA"
|
||||
case gpu.VendorAMD:
|
||||
return "AMD"
|
||||
case gpu.VendorIntel:
|
||||
return "Intel"
|
||||
default:
|
||||
return "Unknown"
|
||||
}
|
||||
}
|
||||
|
||||
// gpuProbeResult holds the outcome of the always-on startup hardware probe.
|
||||
// Populated before any config-gated activation logic runs.
|
||||
type gpuProbeResult struct {
|
||||
Capable bool // true when Devices, IOMMUActive, and VFIOPresent are all satisfied
|
||||
IOMMUActive bool
|
||||
VFIOPresent bool
|
||||
Devices []gpu.GPUDevice
|
||||
}
|
||||
|
||||
// probeGPU discovers GPU hardware and checks passthrough prerequisites.
|
||||
// It is read-only and has no side effects.
|
||||
func probeGPU() gpuProbeResult {
|
||||
var r gpuProbeResult
|
||||
|
||||
devices, err := gpu.Discover()
|
||||
if err != nil {
|
||||
slog.Debug("GPU probe: discover failed", "err", err)
|
||||
}
|
||||
r.Devices = devices
|
||||
|
||||
// IOMMU is active when the kernel has populated iommu_groups in sysfs.
|
||||
groups, err := os.ReadDir("/sys/kernel/iommu_groups")
|
||||
r.IOMMUActive = err == nil && len(groups) > 0
|
||||
|
||||
// vfio_pci module is present when its sysfs module directory exists.
|
||||
_, err = os.Stat("/sys/module/vfio_pci")
|
||||
r.VFIOPresent = err == nil
|
||||
|
||||
r.Capable = len(r.Devices) > 0 && r.IOMMUActive && r.VFIOPresent
|
||||
return r
|
||||
}
|
||||
|
||||
@@ -162,22 +162,38 @@ func (d *Daemon) handleNodeStatus(msg *nats.Msg) {
|
||||
}
|
||||
})
|
||||
|
||||
totalGPUs, allocGPUs := 0, 0
|
||||
if d.gpuManager != nil {
|
||||
totalGPUs = d.gpuManager.TotalCount()
|
||||
allocGPUs = d.gpuManager.AllocatedCount()
|
||||
}
|
||||
|
||||
var gpuModelNames []string
|
||||
for _, dev := range d.gpuProbe.Devices {
|
||||
gpuModelNames = append(gpuModelNames, dev.Model)
|
||||
}
|
||||
|
||||
resp := types.NodeStatusResponse{
|
||||
Node: d.node,
|
||||
Status: "Ready",
|
||||
Host: d.daemonIP(),
|
||||
Region: d.config.Region,
|
||||
AZ: d.config.AZ,
|
||||
Uptime: int64(time.Since(d.startTime).Seconds()),
|
||||
Services: d.config.GetServices(),
|
||||
TotalVCPU: totalVCPU,
|
||||
TotalMemGB: totalMemGB,
|
||||
ReservedVCPU: reservedVCPU,
|
||||
ReservedMemGB: reservedMemGB,
|
||||
AllocVCPU: allocVCPU,
|
||||
AllocMemGB: allocMemGB,
|
||||
VMCount: vmCount,
|
||||
InstanceTypes: caps,
|
||||
Node: d.node,
|
||||
Status: "Ready",
|
||||
Host: d.daemonIP(),
|
||||
Region: d.config.Region,
|
||||
AZ: d.config.AZ,
|
||||
Uptime: int64(time.Since(d.startTime).Seconds()),
|
||||
Services: d.config.GetServices(),
|
||||
TotalVCPU: totalVCPU,
|
||||
TotalMemGB: totalMemGB,
|
||||
ReservedVCPU: reservedVCPU,
|
||||
ReservedMemGB: reservedMemGB,
|
||||
AllocVCPU: allocVCPU,
|
||||
AllocMemGB: allocMemGB,
|
||||
TotalGPUs: totalGPUs,
|
||||
AllocGPUs: allocGPUs,
|
||||
GPUCapable: d.gpuProbe.Capable,
|
||||
GPUPassthrough: d.gpuManager != nil,
|
||||
GPUModels: gpuModelNames,
|
||||
VMCount: vmCount,
|
||||
InstanceTypes: caps,
|
||||
}
|
||||
|
||||
// Query service roles concurrently to halve worst-case latency (500ms vs 1s).
|
||||
|
||||
@@ -13,6 +13,7 @@ import (
|
||||
"github.com/aws/aws-sdk-go/service/ec2"
|
||||
"github.com/mulgadc/spinifex/spinifex/awserrors"
|
||||
"github.com/mulgadc/spinifex/spinifex/filterutil"
|
||||
"github.com/mulgadc/spinifex/spinifex/instancetypes"
|
||||
"github.com/mulgadc/spinifex/spinifex/types"
|
||||
"github.com/mulgadc/spinifex/spinifex/utils"
|
||||
"github.com/mulgadc/spinifex/spinifex/vm"
|
||||
@@ -381,10 +382,29 @@ func (d *Daemon) handleEC2RunInstances(msg *nats.Msg) {
|
||||
instance.Instance.BlockDeviceMappings = append(instance.Instance.BlockDeviceMappings, mapping)
|
||||
}
|
||||
|
||||
// Claim GPU for GPU instance types — binds the full IOMMU group to vfio-pci.
|
||||
if d.gpuManager != nil && instancetypes.IsGPUType(instanceType) {
|
||||
gpuDevice, gpuErr := d.gpuManager.Claim(instance.ID)
|
||||
if gpuErr != nil {
|
||||
slog.Error("handleEC2RunInstances: GPU claim failed", "instanceId", instance.ID, "err", gpuErr)
|
||||
d.vmMgr.MarkFailed(instance, "gpu_claim_failed")
|
||||
continue
|
||||
}
|
||||
instance.GPUPCIAddress = gpuDevice.PCIAddress
|
||||
instance.GPUXVGAEnabled = gpuXVGAEnabled(gpuDevice, d.config.Daemon.GPUModelOverrides)
|
||||
slog.Info("GPU claimed for instance", "instanceId", instance.ID, "gpu", gpuDevice.PCIAddress, "xvga", instance.GPUXVGAEnabled)
|
||||
}
|
||||
|
||||
// Launch the instance infrastructure (QEMU, QMP, NATS subscriptions)
|
||||
err = d.vmMgr.Run(instance)
|
||||
if err != nil {
|
||||
slog.Error("handleEC2RunInstances vmMgr.Run failed", "instanceId", instance.ID, "err", err)
|
||||
if instance.GPUPCIAddress != "" && d.gpuManager != nil {
|
||||
if releaseErr := d.gpuManager.Release(instance.ID); releaseErr != nil {
|
||||
slog.Error("handleEC2RunInstances: GPU release failed after launch failure",
|
||||
"instanceId", instance.ID, "err", releaseErr)
|
||||
}
|
||||
}
|
||||
d.vmMgr.MarkFailed(instance, "launch_failed")
|
||||
continue
|
||||
}
|
||||
@@ -868,11 +888,32 @@ func (d *Daemon) handleEC2StartStoppedInstance(msg *nats.Msg) {
|
||||
instance.Attributes = types.EC2CommandAttributes{StartInstance: true}
|
||||
d.vmMgr.Insert(instance)
|
||||
|
||||
// Claim GPU for GPU instance types — binds the full IOMMU group to vfio-pci.
|
||||
if d.gpuManager != nil && instancetypes.IsGPUType(instanceType) {
|
||||
gpuDevice, gpuErr := d.gpuManager.Claim(instance.ID)
|
||||
if gpuErr != nil {
|
||||
slog.Error("handleEC2StartStoppedInstance: GPU claim failed", "instanceId", req.InstanceID, "err", gpuErr)
|
||||
d.resourceMgr.deallocate(instanceType)
|
||||
d.vmMgr.Delete(instance.ID)
|
||||
respondWithError(msg, awserrors.ErrorInsufficientInstanceCapacity)
|
||||
return
|
||||
}
|
||||
instance.GPUPCIAddress = gpuDevice.PCIAddress
|
||||
instance.GPUXVGAEnabled = gpuXVGAEnabled(gpuDevice, d.config.Daemon.GPUModelOverrides)
|
||||
slog.Info("GPU claimed for instance", "instanceId", req.InstanceID, "gpu", gpuDevice.PCIAddress, "xvga", instance.GPUXVGAEnabled)
|
||||
}
|
||||
|
||||
// Launch the instance infrastructure (QEMU, QMP, NATS subscriptions)
|
||||
err = d.vmMgr.Run(instance)
|
||||
if err != nil {
|
||||
slog.Error("handleEC2StartStoppedInstance: vmMgr.Run failed", "instanceId", req.InstanceID, "err", err)
|
||||
// Rollback: deallocate resources and remove from local map
|
||||
// Rollback: release GPU if claimed, deallocate resources, remove from local map
|
||||
if instance.GPUPCIAddress != "" && d.gpuManager != nil {
|
||||
if releaseErr := d.gpuManager.Release(instance.ID); releaseErr != nil {
|
||||
slog.Error("handleEC2StartStoppedInstance: GPU release failed after launch failure",
|
||||
"instanceId", req.InstanceID, "err", releaseErr)
|
||||
}
|
||||
}
|
||||
if ok {
|
||||
d.resourceMgr.deallocate(instanceType)
|
||||
}
|
||||
|
||||
@@ -26,6 +26,7 @@ import (
|
||||
"github.com/aws/aws-sdk-go/service/ec2"
|
||||
awss3 "github.com/aws/aws-sdk-go/service/s3"
|
||||
"github.com/mulgadc/spinifex/spinifex/config"
|
||||
"github.com/mulgadc/spinifex/spinifex/gpu"
|
||||
handlers_ec2_instance "github.com/mulgadc/spinifex/spinifex/handlers/ec2/instance"
|
||||
handlers_ec2_volume "github.com/mulgadc/spinifex/spinifex/handlers/ec2/volume"
|
||||
"github.com/mulgadc/spinifex/spinifex/instancetypes"
|
||||
@@ -128,7 +129,7 @@ func createTestDaemon(t *testing.T, natsURL string) *Daemon {
|
||||
// getTestInstanceType returns a valid instance type for testing based on the system's CPU
|
||||
func getTestInstanceType(t *testing.T) string {
|
||||
t.Helper()
|
||||
rm, err := NewResourceManager()
|
||||
rm, err := NewResourceManager(nil, nil)
|
||||
require.NoError(t, err)
|
||||
// Find any .micro instance type
|
||||
for key := range rm.instanceTypes {
|
||||
@@ -411,7 +412,7 @@ func TestDaemon_Initialization(t *testing.T) {
|
||||
|
||||
// TestResourceManager tests resource manager functionality
|
||||
func TestResourceManager(t *testing.T) {
|
||||
rm, err := NewResourceManager()
|
||||
rm, err := NewResourceManager(nil, nil)
|
||||
require.NoError(t, err)
|
||||
|
||||
require.NotNil(t, rm)
|
||||
@@ -461,7 +462,7 @@ func TestResourceManager(t *testing.T) {
|
||||
// Test canAllocate with count parameter
|
||||
t.Run("canAllocate_with_count", func(t *testing.T) {
|
||||
// Fresh resource manager for predictable testing
|
||||
rm, err := NewResourceManager()
|
||||
rm, err := NewResourceManager(nil, nil)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Find a .micro instance type
|
||||
@@ -498,7 +499,7 @@ func TestResourceManager(t *testing.T) {
|
||||
|
||||
// TestGetInstanceTypeInfos tests the GetInstanceTypeInfos method
|
||||
func TestGetInstanceTypeInfos(t *testing.T) {
|
||||
rm, err := NewResourceManager()
|
||||
rm, err := NewResourceManager(nil, nil)
|
||||
require.NoError(t, err)
|
||||
|
||||
infos := rm.GetInstanceTypeInfos()
|
||||
@@ -526,7 +527,7 @@ func TestGetInstanceTypeInfos(t *testing.T) {
|
||||
|
||||
// TestGetAvailableInstanceTypeInfos_ResourceFiltering tests that instance types are filtered by available resources
|
||||
func TestGetAvailableInstanceTypeInfos_ResourceFiltering(t *testing.T) {
|
||||
rm, err := NewResourceManager()
|
||||
rm, err := NewResourceManager(nil, nil)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Get initial count of all available types
|
||||
@@ -1034,7 +1035,7 @@ func createValidRunInstancesInput(t *testing.T) *ec2.RunInstancesInput {
|
||||
// TestCanAllocate_CountEdgeCases tests edge cases for canAllocate with count parameter
|
||||
func TestCanAllocate_CountEdgeCases(t *testing.T) {
|
||||
t.Run("MinCount_equals_MaxCount", func(t *testing.T) {
|
||||
rm, err := NewResourceManager()
|
||||
rm, err := NewResourceManager(nil, nil)
|
||||
require.NoError(t, err)
|
||||
|
||||
var microType *ec2.InstanceTypeInfo
|
||||
@@ -1053,7 +1054,7 @@ func TestCanAllocate_CountEdgeCases(t *testing.T) {
|
||||
})
|
||||
|
||||
t.Run("Request_exceeds_capacity", func(t *testing.T) {
|
||||
rm, err := NewResourceManager()
|
||||
rm, err := NewResourceManager(nil, nil)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Find the largest instance type to exhaust resources faster
|
||||
@@ -1076,7 +1077,7 @@ func TestCanAllocate_CountEdgeCases(t *testing.T) {
|
||||
})
|
||||
|
||||
t.Run("Capacity_decreases_after_allocation", func(t *testing.T) {
|
||||
rm, err := NewResourceManager()
|
||||
rm, err := NewResourceManager(nil, nil)
|
||||
require.NoError(t, err)
|
||||
|
||||
var microType *ec2.InstanceTypeInfo
|
||||
@@ -1126,7 +1127,7 @@ func TestCanAllocate_CountEdgeCases(t *testing.T) {
|
||||
})
|
||||
|
||||
t.Run("Mixed_instance_types", func(t *testing.T) {
|
||||
rm, err := NewResourceManager()
|
||||
rm, err := NewResourceManager(nil, nil)
|
||||
require.NoError(t, err)
|
||||
|
||||
var microType, mediumType *ec2.InstanceTypeInfo
|
||||
@@ -1159,7 +1160,7 @@ func TestCanAllocate_CountEdgeCases(t *testing.T) {
|
||||
})
|
||||
|
||||
t.Run("Zero_and_negative_counts", func(t *testing.T) {
|
||||
rm, err := NewResourceManager()
|
||||
rm, err := NewResourceManager(nil, nil)
|
||||
require.NoError(t, err)
|
||||
|
||||
var microType *ec2.InstanceTypeInfo
|
||||
@@ -1455,7 +1456,7 @@ func TestInstanceTypeSubscriptions(t *testing.T) {
|
||||
|
||||
t.Run("InitialSubscriptions", func(t *testing.T) {
|
||||
// A fresh ResourceManager should subscribe to all instance types that fit
|
||||
rm, err := NewResourceManager()
|
||||
rm, err := NewResourceManager(nil, nil)
|
||||
require.NoError(t, err)
|
||||
nc, err := nats.Connect(natsURL)
|
||||
require.NoError(t, err)
|
||||
@@ -1489,7 +1490,7 @@ func TestInstanceTypeSubscriptions(t *testing.T) {
|
||||
})
|
||||
|
||||
t.Run("UnsubscribesWhenFull", func(t *testing.T) {
|
||||
rm, err := NewResourceManager()
|
||||
rm, err := NewResourceManager(nil, nil)
|
||||
require.NoError(t, err)
|
||||
nc, err := nats.Connect(natsURL)
|
||||
require.NoError(t, err)
|
||||
@@ -1514,7 +1515,7 @@ func TestInstanceTypeSubscriptions(t *testing.T) {
|
||||
})
|
||||
|
||||
t.Run("ResubscribesWhenFreed", func(t *testing.T) {
|
||||
rm, err := NewResourceManager()
|
||||
rm, err := NewResourceManager(nil, nil)
|
||||
require.NoError(t, err)
|
||||
nc, err := nats.Connect(natsURL)
|
||||
require.NoError(t, err)
|
||||
@@ -1545,7 +1546,7 @@ func TestInstanceTypeSubscriptions(t *testing.T) {
|
||||
})
|
||||
|
||||
t.Run("PartialCapacity", func(t *testing.T) {
|
||||
rm, err := NewResourceManager()
|
||||
rm, err := NewResourceManager(nil, nil)
|
||||
require.NoError(t, err)
|
||||
nc, err := nats.Connect(natsURL)
|
||||
require.NoError(t, err)
|
||||
@@ -1574,7 +1575,7 @@ func TestInstanceTypeSubscriptions(t *testing.T) {
|
||||
})
|
||||
|
||||
t.Run("AllocateTriggersSubs", func(t *testing.T) {
|
||||
rm, err := NewResourceManager()
|
||||
rm, err := NewResourceManager(nil, nil)
|
||||
require.NoError(t, err)
|
||||
nc, err := nats.Connect(natsURL)
|
||||
require.NoError(t, err)
|
||||
@@ -1618,7 +1619,7 @@ func TestInstanceTypeSubscriptions(t *testing.T) {
|
||||
})
|
||||
|
||||
t.Run("NoRespondersWhenFull", func(t *testing.T) {
|
||||
rm, err := NewResourceManager()
|
||||
rm, err := NewResourceManager(nil, nil)
|
||||
require.NoError(t, err)
|
||||
nc, err := nats.Connect(natsURL)
|
||||
require.NoError(t, err)
|
||||
@@ -1647,7 +1648,7 @@ func TestInstanceTypeSubscriptions(t *testing.T) {
|
||||
|
||||
// TestResourceManager_ConcurrentAccess tests thread safety of resource manager
|
||||
func TestResourceManager_ConcurrentAccess(t *testing.T) {
|
||||
rm, err := NewResourceManager()
|
||||
rm, err := NewResourceManager(nil, nil)
|
||||
require.NoError(t, err)
|
||||
|
||||
var microType *ec2.InstanceTypeInfo
|
||||
@@ -3638,3 +3639,58 @@ func generateTestCert(t *testing.T) (certPEM, keyPEM []byte) {
|
||||
|
||||
return certBuf.Bytes(), keyBuf.Bytes()
|
||||
}
|
||||
|
||||
func TestResolveGPUModel_ProductionModel(t *testing.T) {
|
||||
dev := gpu.GPUDevice{VendorID: "10de", DeviceID: "2236", Vendor: gpu.VendorNVIDIA}
|
||||
m := resolveGPUModel(dev, nil)
|
||||
assert.Equal(t, "g5", m.Family)
|
||||
assert.Equal(t, "A10G", m.Name)
|
||||
}
|
||||
|
||||
func TestResolveGPUModel_ConsumerGPUDefaultsToG5(t *testing.T) {
|
||||
// Unknown PCI ID auto-maps to g5 using discovered specs — no config needed.
|
||||
dev := gpu.GPUDevice{
|
||||
VendorID: "10de",
|
||||
DeviceID: "2487",
|
||||
Vendor: gpu.VendorNVIDIA,
|
||||
Model: "NVIDIA GeForce RTX 3060",
|
||||
MemoryMiB: 12288,
|
||||
}
|
||||
m := resolveGPUModel(dev, nil)
|
||||
assert.Equal(t, "g5", m.Family)
|
||||
assert.Equal(t, "NVIDIA GeForce RTX 3060", m.Name)
|
||||
assert.Equal(t, int64(12288), m.MemoryMiB)
|
||||
assert.Equal(t, "NVIDIA", m.Manufacturer)
|
||||
}
|
||||
|
||||
func TestResolveGPUModel_ConsumerGPUFallbackName(t *testing.T) {
|
||||
// No Model field: name falls back to "GPU <vendor>:<device>".
|
||||
dev := gpu.GPUDevice{VendorID: "dead", DeviceID: "beef", Vendor: gpu.VendorUnknown}
|
||||
m := resolveGPUModel(dev, nil)
|
||||
assert.Equal(t, "g5", m.Family)
|
||||
assert.Equal(t, "GPU dead:beef", m.Name)
|
||||
assert.Equal(t, "Unknown", m.Manufacturer)
|
||||
}
|
||||
|
||||
func TestResolveGPUModel_OverrideShadowsProduction(t *testing.T) {
|
||||
// An override for a known production PCI ID shadows the built-in entry.
|
||||
overrides := []config.GPUModelOverride{
|
||||
{VendorID: "10de", DeviceID: "2236", Family: "g6", Manufacturer: "NVIDIA", Name: "Custom", MemoryMiB: 999},
|
||||
}
|
||||
dev := gpu.GPUDevice{VendorID: "10de", DeviceID: "2236", Vendor: gpu.VendorNVIDIA}
|
||||
m := resolveGPUModel(dev, overrides)
|
||||
assert.Equal(t, "g6", m.Family)
|
||||
assert.Equal(t, "Custom", m.Name)
|
||||
}
|
||||
|
||||
func TestResolveGPUModel_OverrideCustomisesConsumerGPU(t *testing.T) {
|
||||
// Override can pin specific name/memory for a consumer GPU that would
|
||||
// otherwise be auto-mapped with nvidia-smi-discovered or zero values.
|
||||
overrides := []config.GPUModelOverride{
|
||||
{VendorID: "10de", DeviceID: "2487", Family: "g5", Manufacturer: "NVIDIA", Name: "RTX 3060", MemoryMiB: 12288},
|
||||
}
|
||||
dev := gpu.GPUDevice{VendorID: "10de", DeviceID: "2487", Vendor: gpu.VendorNVIDIA}
|
||||
m := resolveGPUModel(dev, overrides)
|
||||
assert.Equal(t, "RTX 3060", m.Name)
|
||||
assert.Equal(t, int64(12288), m.MemoryMiB)
|
||||
}
|
||||
|
||||
@@ -13,7 +13,7 @@ import (
|
||||
|
||||
// TestBuildHeartbeat verifies that buildHeartbeat populates the struct from daemon state.
|
||||
func TestBuildHeartbeat(t *testing.T) {
|
||||
rm, err := NewResourceManager()
|
||||
rm, err := NewResourceManager(nil, nil)
|
||||
require.NoError(t, err)
|
||||
|
||||
d := &Daemon{
|
||||
@@ -46,7 +46,7 @@ func TestBuildHeartbeat(t *testing.T) {
|
||||
|
||||
// TestHeartbeatReflectsAllocation verifies that allocating resources changes the heartbeat values.
|
||||
func TestHeartbeatReflectsAllocation(t *testing.T) {
|
||||
rm, err := NewResourceManager()
|
||||
rm, err := NewResourceManager(nil, nil)
|
||||
require.NoError(t, err)
|
||||
|
||||
d := &Daemon{
|
||||
|
||||
@@ -61,8 +61,17 @@ func applyHostReserve(host hostReserve, totalVCPU int, totalMemGB float64) (vcpu
|
||||
// slog.Error when remaining capacity is negative — that condition would
|
||||
// otherwise be silently clamped to zero, hiding a misconfigured reserve
|
||||
// or allocation accounting drift.
|
||||
//
|
||||
// availGPU is the number of free GPUs in the pool; requiresGPU indicates
|
||||
// that this instance type needs one. When requiresGPU is true and availGPU
|
||||
// is zero, the result is always 0.
|
||||
func canAllocateCount(availVCPU, allocVCPU int, availMem, allocMem float64,
|
||||
vCPUs int64, memMiB int64, maxCount int) int {
|
||||
vCPUs int64, memMiB int64, maxCount int,
|
||||
availGPU int, requiresGPU bool) int {
|
||||
if requiresGPU && availGPU == 0 {
|
||||
return 0
|
||||
}
|
||||
|
||||
remainingVCPU := availVCPU - allocVCPU
|
||||
remainingMem := availMem - allocMem
|
||||
if remainingVCPU < 0 || remainingMem < 0 {
|
||||
@@ -83,6 +92,9 @@ func canAllocateCount(availVCPU, allocVCPU int, availMem, allocMem float64,
|
||||
}
|
||||
|
||||
result := min(countByMem, countByCPU)
|
||||
if requiresGPU {
|
||||
result = min(result, availGPU)
|
||||
}
|
||||
result = min(result, maxCount)
|
||||
return max(result, 0)
|
||||
}
|
||||
|
||||
@@ -112,12 +112,44 @@ func TestCanAllocateCount(t *testing.T) {
|
||||
|
||||
for _, tc := range tests {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
got := canAllocateCount(tc.availVCPU, tc.allocVCPU, tc.availMem, tc.allocMem, tc.vCPUs, tc.memMiB, tc.maxCount)
|
||||
got := canAllocateCount(tc.availVCPU, tc.allocVCPU, tc.availMem, tc.allocMem, tc.vCPUs, tc.memMiB, tc.maxCount, 0, false)
|
||||
assert.Equal(t, tc.want, got)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestCanAllocateCount_GPU(t *testing.T) {
|
||||
// Base resource params: plenty of CPU and memory, 1 GPU available.
|
||||
const (
|
||||
avCPU = 32
|
||||
alCPU = 0
|
||||
avMem = 256.0
|
||||
alMem = 0.0
|
||||
vCPUs = 4 // g5.xlarge
|
||||
memMiB = 16384
|
||||
)
|
||||
|
||||
// GPU available: normal CPU/mem-constrained allocation with GPU cap.
|
||||
got := canAllocateCount(avCPU, alCPU, avMem, alMem, vCPUs, memMiB, 10, 1, true)
|
||||
assert.Equal(t, 1, got, "with 1 GPU available, max 1 instance despite spare CPU/mem")
|
||||
|
||||
// 3 GPUs available: cap is 3.
|
||||
got = canAllocateCount(avCPU, alCPU, avMem, alMem, vCPUs, memMiB, 10, 3, true)
|
||||
assert.Equal(t, 3, got, "with 3 GPUs, GPU pool caps the count")
|
||||
|
||||
// 0 GPUs: must return 0 regardless of CPU/mem headroom.
|
||||
got = canAllocateCount(avCPU, alCPU, avMem, alMem, vCPUs, memMiB, 10, 0, true)
|
||||
assert.Equal(t, 0, got, "no GPUs available must return 0")
|
||||
|
||||
// Non-GPU type is unaffected by availGPU=0.
|
||||
got = canAllocateCount(avCPU, alCPU, avMem, alMem, vCPUs, memMiB, 10, 0, false)
|
||||
assert.Equal(t, 8, got, "non-GPU type: CPU 32/4=8; memory and GPU not limiting")
|
||||
|
||||
// maxCount caps GPU result.
|
||||
got = canAllocateCount(avCPU, alCPU, avMem, alMem, vCPUs, memMiB, 2, 5, true)
|
||||
assert.Equal(t, 2, got, "maxCount=2 caps GPU result")
|
||||
}
|
||||
|
||||
func TestResourceStatsForType(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
|
||||
@@ -473,6 +473,19 @@ func (d *Daemon) onInstanceUpHook() func(*vm.VM) error {
|
||||
return fmt.Errorf("subscribe ec2.%s.GetConsoleOutput: %w", instance.ID, err)
|
||||
}
|
||||
d.natsSubscriptions[consoleSubKey] = consoleSub
|
||||
|
||||
// Re-claim GPU after a daemon restart with a still-running QEMU
|
||||
// process: the manager's reconnect path fires OnInstanceUp without
|
||||
// going through the handler-side Claim, so the GPU pool would
|
||||
// otherwise treat the slot as free. ReclaimByAddress is a no-op
|
||||
// when the same instance already owns the slot, so the launch and
|
||||
// start-stopped paths (which Claim before Run) are unaffected.
|
||||
if d.gpuManager != nil && instance.GPUPCIAddress != "" {
|
||||
if err := d.gpuManager.ReclaimByAddress(instance.GPUPCIAddress, instance.ID); err != nil {
|
||||
slog.Warn("Failed to re-claim GPU on instance up",
|
||||
"gpu", instance.GPUPCIAddress, "instanceId", instance.ID, "err", err)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
}
|
||||
@@ -692,3 +705,18 @@ func (a *instanceCleanerAdapter) RemoveFromPlacementGroup(instance *vm.VM) {
|
||||
"instanceId", instance.ID, "groupName", instance.PlacementGroupName, "err", err)
|
||||
}
|
||||
}
|
||||
|
||||
// ReleaseGPU unbinds the instance's GPU from vfio-pci and rebinds to its
|
||||
// original host driver. No-op for instances without a GPU allocation or
|
||||
// when GPU passthrough is disabled.
|
||||
func (a *instanceCleanerAdapter) ReleaseGPU(instance *vm.VM) {
|
||||
if a.d.gpuManager == nil || instance.GPUPCIAddress == "" {
|
||||
return
|
||||
}
|
||||
if err := a.d.gpuManager.Release(instance.ID); err != nil {
|
||||
slog.Error("Failed to release GPU on stop, device may need manual rebind",
|
||||
"gpu", instance.GPUPCIAddress, "instanceId", instance.ID, "err", err)
|
||||
return
|
||||
}
|
||||
slog.Info("GPU released", "gpu", instance.GPUPCIAddress, "instanceId", instance.ID)
|
||||
}
|
||||
|
||||
@@ -3,6 +3,7 @@ package daemon
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/mulgadc/spinifex/spinifex/gpu"
|
||||
"github.com/mulgadc/spinifex/spinifex/vm"
|
||||
"github.com/nats-io/nats.go"
|
||||
"github.com/stretchr/testify/assert"
|
||||
@@ -95,6 +96,42 @@ func TestOnInstanceDownHook_NoOpWhenAbsent(t *testing.T) {
|
||||
assert.Empty(t, d.natsSubscriptions)
|
||||
}
|
||||
|
||||
// When the daemon's gpuManager is unset, the hook must still register NATS
|
||||
// subscriptions and ignore the GPUPCIAddress on the VM.
|
||||
func TestOnInstanceUpHook_NoGPUManager_SkipsReclaim(t *testing.T) {
|
||||
d, _ := newHookTestDaemon(t)
|
||||
instance := &vm.VM{ID: "i-up-nogpu", GPUPCIAddress: "0000:01:00.0"}
|
||||
|
||||
require.NoError(t, d.onInstanceUpHook()(instance))
|
||||
require.Contains(t, d.natsSubscriptions, instance.ID)
|
||||
}
|
||||
|
||||
// With a gpuManager configured but a CPU-only instance, the hook must not
|
||||
// touch the GPU pool. We use the AllocatedCount as the observable: an
|
||||
// unintended Reclaim would bump it.
|
||||
func TestOnInstanceUpHook_NoGPUAddress_SkipsReclaim(t *testing.T) {
|
||||
d, _ := newHookTestDaemon(t)
|
||||
d.gpuManager = gpu.NewManager(nil)
|
||||
instance := &vm.VM{ID: "i-up-cpu"}
|
||||
|
||||
require.NoError(t, d.onInstanceUpHook()(instance))
|
||||
assert.Equal(t, 0, d.gpuManager.AllocatedCount(),
|
||||
"hook must not call Reclaim for instances without a GPUPCIAddress")
|
||||
}
|
||||
|
||||
// With a gpuManager that has no entries, calling Reclaim for an instance
|
||||
// with a GPUPCIAddress will fail inside the manager. The hook logs a warning
|
||||
// and returns nil — the NATS subscriptions must still register so the
|
||||
// reconnect path doesn't roll back.
|
||||
func TestOnInstanceUpHook_GPUReclaimError_DoesNotPropagate(t *testing.T) {
|
||||
d, _ := newHookTestDaemon(t)
|
||||
d.gpuManager = gpu.NewManager(nil)
|
||||
instance := &vm.VM{ID: "i-up-gpu-missing", GPUPCIAddress: "0000:99:00.0"}
|
||||
|
||||
require.NoError(t, d.onInstanceUpHook()(instance))
|
||||
require.Contains(t, d.natsSubscriptions, instance.ID)
|
||||
}
|
||||
|
||||
func TestOnInstanceDownHook_OnlyRemovesTargetedInstance(t *testing.T) {
|
||||
d, _ := newHookTestDaemon(t)
|
||||
keep := &vm.VM{ID: "i-keep"}
|
||||
|
||||
@@ -0,0 +1,92 @@
|
||||
package gpu
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
)
|
||||
|
||||
// BindVFIO unbinds addr from its current driver and binds it to vfio-pci
|
||||
// using driver_override. Returns the original driver name (empty if unbound).
|
||||
// Idempotent: returns ("vfio-pci", nil) if already bound to vfio-pci.
|
||||
func BindVFIO(addr string) (string, error) {
|
||||
return bindVFIO("/sys", addr)
|
||||
}
|
||||
|
||||
// UnbindVFIO unbinds addr from vfio-pci, clears driver_override, and rebinds
|
||||
// to originalDriver. If originalDriver is empty the device is left unbound.
|
||||
func UnbindVFIO(addr, originalDriver string) error {
|
||||
return unbindVFIO("/sys", addr, originalDriver)
|
||||
}
|
||||
|
||||
// CurrentDriver returns the name of the driver currently bound to addr,
|
||||
// or "" if the device is unbound.
|
||||
func CurrentDriver(addr string) (string, error) {
|
||||
return readDriver(filepath.Join("/sys/bus/pci/devices", addr))
|
||||
}
|
||||
|
||||
func bindVFIO(sysfsRoot, addr string) (string, error) {
|
||||
devPath := filepath.Join(sysfsRoot, "bus/pci/devices", addr)
|
||||
|
||||
current, err := readDriver(devPath)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("read driver for %s: %w", addr, err)
|
||||
}
|
||||
if current == "vfio-pci" {
|
||||
return "vfio-pci", nil
|
||||
}
|
||||
|
||||
if current != "" {
|
||||
unbindPath := filepath.Join(sysfsRoot, "bus/pci/drivers", current, "unbind")
|
||||
if err := writeSysfs(unbindPath, addr); err != nil {
|
||||
return "", fmt.Errorf("unbind %s from %s: %w", addr, current, err)
|
||||
}
|
||||
}
|
||||
|
||||
if err := writeSysfs(filepath.Join(devPath, "driver_override"), "vfio-pci"); err != nil {
|
||||
return "", fmt.Errorf("set driver_override for %s: %w", addr, err)
|
||||
}
|
||||
|
||||
if err := writeSysfs(filepath.Join(sysfsRoot, "bus/pci/drivers/vfio-pci/bind"), addr); err != nil {
|
||||
return "", fmt.Errorf("bind %s to vfio-pci: %w", addr, err)
|
||||
}
|
||||
|
||||
return current, nil
|
||||
}
|
||||
|
||||
func unbindVFIO(sysfsRoot, addr, originalDriver string) error {
|
||||
devPath := filepath.Join(sysfsRoot, "bus/pci/devices", addr)
|
||||
|
||||
current, err := readDriver(devPath)
|
||||
if err != nil {
|
||||
return fmt.Errorf("read driver for %s: %w", addr, err)
|
||||
}
|
||||
if current != "vfio-pci" {
|
||||
// Already unbound from vfio-pci (e.g. released by QEMU before Release was called).
|
||||
return nil
|
||||
}
|
||||
|
||||
if err := writeSysfs(filepath.Join(sysfsRoot, "bus/pci/drivers/vfio-pci/unbind"), addr); err != nil {
|
||||
return fmt.Errorf("unbind %s from vfio-pci: %w", addr, err)
|
||||
}
|
||||
|
||||
if err := writeSysfs(filepath.Join(devPath, "driver_override"), ""); err != nil {
|
||||
return fmt.Errorf("clear driver_override for %s: %w", addr, err)
|
||||
}
|
||||
|
||||
if originalDriver != "" {
|
||||
rebindPath := filepath.Join(sysfsRoot, "bus/pci/drivers", originalDriver, "bind")
|
||||
if err := writeSysfs(rebindPath, addr); err != nil {
|
||||
return fmt.Errorf("rebind %s to %s: %w", addr, originalDriver, err)
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// writeSysfs writes value to a sysfs attribute path.
|
||||
// Uses O_WRONLY|O_CREATE|O_TRUNC so it works against both real sysfs
|
||||
// (file exists, kernel handles the write) and test temp dirs (file created).
|
||||
func writeSysfs(path, value string) error {
|
||||
return os.WriteFile(path, []byte(value), 0o600)
|
||||
}
|
||||
@@ -0,0 +1,164 @@
|
||||
package gpu
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// makeSysfsDriverDir creates a driver directory under root/bus/pci/drivers/<name>.
|
||||
func makeSysfsDriverDir(t *testing.T, root, driver string) string {
|
||||
t.Helper()
|
||||
path := filepath.Join(root, "bus/pci/drivers", driver)
|
||||
must(t, os.MkdirAll(path, 0o755))
|
||||
return path
|
||||
}
|
||||
|
||||
// readSysfsTestFile reads a file written by writeSysfs in a test sysfs tree.
|
||||
func readSysfsTestFile(t *testing.T, path string) string {
|
||||
t.Helper()
|
||||
b, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
t.Fatalf("read %s: %v", path, err)
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
|
||||
func TestBindVFIO_BoundToDriver(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
buildSysfsDevice(t, root, "0000:03:00.0", "0x030200", "0x10de", "0x2236", "nvidia", 7)
|
||||
makeSysfsDriverDir(t, root, "vfio-pci")
|
||||
|
||||
orig, err := bindVFIO(root, "0000:03:00.0")
|
||||
if err != nil {
|
||||
t.Fatalf("bindVFIO: %v", err)
|
||||
}
|
||||
if orig != "nvidia" {
|
||||
t.Errorf("original driver = %q, want nvidia", orig)
|
||||
}
|
||||
|
||||
devPath := filepath.Join(root, "bus/pci/devices/0000:03:00.0")
|
||||
|
||||
if got := readSysfsTestFile(t, filepath.Join(root, "bus/pci/drivers/nvidia/unbind")); got != "0000:03:00.0" {
|
||||
t.Errorf("nvidia/unbind = %q, want 0000:03:00.0", got)
|
||||
}
|
||||
if got := readSysfsTestFile(t, filepath.Join(devPath, "driver_override")); got != "vfio-pci" {
|
||||
t.Errorf("driver_override = %q, want vfio-pci", got)
|
||||
}
|
||||
if got := readSysfsTestFile(t, filepath.Join(root, "bus/pci/drivers/vfio-pci/bind")); got != "0000:03:00.0" {
|
||||
t.Errorf("vfio-pci/bind = %q, want 0000:03:00.0", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBindVFIO_AlreadyBoundToVFIO(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
buildSysfsDevice(t, root, "0000:03:00.0", "0x030200", "0x10de", "0x2236", "vfio-pci", 7)
|
||||
|
||||
orig, err := bindVFIO(root, "0000:03:00.0")
|
||||
if err != nil {
|
||||
t.Fatalf("bindVFIO: %v", err)
|
||||
}
|
||||
if orig != "vfio-pci" {
|
||||
t.Errorf("expected vfio-pci for idempotent call, got %q", orig)
|
||||
}
|
||||
|
||||
// No vfio-pci/bind file should have been written (idempotent).
|
||||
bindPath := filepath.Join(root, "bus/pci/drivers/vfio-pci/bind")
|
||||
if _, err := os.Stat(bindPath); !os.IsNotExist(err) {
|
||||
t.Error("vfio-pci/bind should not be written for already-bound device")
|
||||
}
|
||||
}
|
||||
|
||||
func TestBindVFIO_Unbound(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
// No driver symlink — device is unbound.
|
||||
buildSysfsDevice(t, root, "0000:03:00.0", "0x030200", "0x10de", "0x2236", "", 7)
|
||||
makeSysfsDriverDir(t, root, "vfio-pci")
|
||||
|
||||
orig, err := bindVFIO(root, "0000:03:00.0")
|
||||
if err != nil {
|
||||
t.Fatalf("bindVFIO: %v", err)
|
||||
}
|
||||
if orig != "" {
|
||||
t.Errorf("original driver = %q, want empty for unbound device", orig)
|
||||
}
|
||||
|
||||
// No unbind file should have been written (nothing to unbind).
|
||||
nvidiaUnbind := filepath.Join(root, "bus/pci/drivers/nvidia/unbind")
|
||||
if _, err := os.Stat(nvidiaUnbind); !os.IsNotExist(err) {
|
||||
t.Error("no driver unbind should be written when device was unbound")
|
||||
}
|
||||
|
||||
devPath := filepath.Join(root, "bus/pci/devices/0000:03:00.0")
|
||||
if got := readSysfsTestFile(t, filepath.Join(devPath, "driver_override")); got != "vfio-pci" {
|
||||
t.Errorf("driver_override = %q, want vfio-pci", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUnbindVFIO_WithOriginalDriver(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
buildSysfsDevice(t, root, "0000:03:00.0", "0x030200", "0x10de", "0x2236", "vfio-pci", 7)
|
||||
makeSysfsDriverDir(t, root, "vfio-pci")
|
||||
makeSysfsDriverDir(t, root, "nvidia")
|
||||
|
||||
err := unbindVFIO(root, "0000:03:00.0", "nvidia")
|
||||
if err != nil {
|
||||
t.Fatalf("unbindVFIO: %v", err)
|
||||
}
|
||||
|
||||
devPath := filepath.Join(root, "bus/pci/devices/0000:03:00.0")
|
||||
|
||||
if got := readSysfsTestFile(t, filepath.Join(root, "bus/pci/drivers/vfio-pci/unbind")); got != "0000:03:00.0" {
|
||||
t.Errorf("vfio-pci/unbind = %q, want 0000:03:00.0", got)
|
||||
}
|
||||
if got := readSysfsTestFile(t, filepath.Join(devPath, "driver_override")); got != "" {
|
||||
t.Errorf("driver_override = %q, want empty after unbind", got)
|
||||
}
|
||||
if got := readSysfsTestFile(t, filepath.Join(root, "bus/pci/drivers/nvidia/bind")); got != "0000:03:00.0" {
|
||||
t.Errorf("nvidia/bind = %q, want 0000:03:00.0", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBindVFIO_VfioBindFails(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
buildSysfsDevice(t, root, "0000:03:00.0", "0x030200", "0x10de", "0x2236", "nvidia", 7)
|
||||
// nvidia dir is created by buildSysfsDevice; deliberately omit vfio-pci dir so bind fails.
|
||||
|
||||
_, err := bindVFIO(root, "0000:03:00.0")
|
||||
if err == nil {
|
||||
t.Error("want error when vfio-pci bind directory is missing, got nil")
|
||||
}
|
||||
}
|
||||
|
||||
func TestUnbindVFIO_ClearDriverOverrideFails(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
buildSysfsDevice(t, root, "0000:03:00.0", "0x030200", "0x10de", "0x2236", "vfio-pci", 7)
|
||||
makeSysfsDriverDir(t, root, "vfio-pci")
|
||||
makeSysfsDriverDir(t, root, "nvidia")
|
||||
|
||||
// Pre-create driver_override as read-only so the clear write fails.
|
||||
devPath := filepath.Join(root, "bus/pci/devices/0000:03:00.0")
|
||||
must(t, os.WriteFile(filepath.Join(devPath, "driver_override"), []byte("vfio-pci"), 0o444))
|
||||
|
||||
err := unbindVFIO(root, "0000:03:00.0", "nvidia")
|
||||
if err == nil {
|
||||
t.Error("want error when driver_override clear fails, got nil")
|
||||
}
|
||||
}
|
||||
|
||||
func TestUnbindVFIO_NoOriginalDriver(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
buildSysfsDevice(t, root, "0000:03:00.0", "0x030200", "0x10de", "0x2236", "vfio-pci", 7)
|
||||
makeSysfsDriverDir(t, root, "vfio-pci")
|
||||
|
||||
err := unbindVFIO(root, "0000:03:00.0", "")
|
||||
if err != nil {
|
||||
t.Fatalf("unbindVFIO: %v", err)
|
||||
}
|
||||
|
||||
// With no original driver there should be no rebind file written.
|
||||
matches, _ := filepath.Glob(filepath.Join(root, "bus/pci/drivers/*/bind"))
|
||||
if len(matches) > 0 {
|
||||
t.Errorf("no driver rebind expected when originalDriver is empty, got: %v", matches)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,209 @@
|
||||
package gpu
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// Discover enumerates GPUs on the host via sysfs and returns a GPUDevice for
|
||||
// each display-class PCI device found. nvidia-smi is used to enrich NVIDIA
|
||||
// entries if available; falls back to the built-in model table otherwise.
|
||||
func Discover() ([]GPUDevice, error) {
|
||||
return discover("/sys")
|
||||
}
|
||||
|
||||
func discover(sysfsRoot string) ([]GPUDevice, error) {
|
||||
devicesPath := filepath.Join(sysfsRoot, "bus/pci/devices")
|
||||
entries, err := os.ReadDir(devicesPath)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("read pci devices: %w", err)
|
||||
}
|
||||
|
||||
var gpus []GPUDevice
|
||||
for _, entry := range entries {
|
||||
devPath := filepath.Join(devicesPath, entry.Name())
|
||||
|
||||
class, err := readSysfsString(filepath.Join(devPath, "class"))
|
||||
if err != nil || !isDisplayClass(class) {
|
||||
continue
|
||||
}
|
||||
|
||||
vendorRaw, err := readSysfsString(filepath.Join(devPath, "vendor"))
|
||||
if err != nil {
|
||||
slog.Debug("gpu discover: skip device, cannot read vendor", "addr", entry.Name(), "err", err)
|
||||
continue
|
||||
}
|
||||
deviceRaw, err := readSysfsString(filepath.Join(devPath, "device"))
|
||||
if err != nil {
|
||||
slog.Debug("gpu discover: skip device, cannot read device id", "addr", entry.Name(), "err", err)
|
||||
continue
|
||||
}
|
||||
|
||||
vendorID := normHex(vendorRaw)
|
||||
deviceID := normHex(deviceRaw)
|
||||
|
||||
// Only NVIDIA and AMD GPUs are eligible for VFIO passthrough.
|
||||
// Intel iGPUs are the host's primary display and must stay on the host driver.
|
||||
if vendorID != "10de" && vendorID != "1002" {
|
||||
slog.Debug("gpu discover: skipping non-passthrough vendor", "addr", entry.Name(), "vendor", vendorID)
|
||||
continue
|
||||
}
|
||||
|
||||
group, err := readIOMMUGroup(devPath)
|
||||
if err != nil {
|
||||
slog.Debug("gpu discover: IOMMU group unavailable", "addr", entry.Name(), "err", err)
|
||||
group = -1
|
||||
}
|
||||
|
||||
driver, err := readDriver(devPath)
|
||||
if err != nil {
|
||||
driver = ""
|
||||
}
|
||||
|
||||
gpu := GPUDevice{
|
||||
PCIAddress: entry.Name(),
|
||||
Vendor: vendorFromID(vendorID),
|
||||
VendorID: vendorID,
|
||||
DeviceID: deviceID,
|
||||
IOMMUGroup: group,
|
||||
OriginalDriver: driver,
|
||||
}
|
||||
|
||||
// Enrich model name and VRAM from the built-in table, then try
|
||||
// lspci (PCI ID database) for cards not in the table, then
|
||||
// nvidia-smi for precise VRAM on NVIDIA cards.
|
||||
gpu.Model, gpu.MemoryMiB = lookupModel(vendorID, deviceID)
|
||||
if gpu.Model == "" {
|
||||
enrichFromLspci(&gpu)
|
||||
}
|
||||
if gpu.Model == "" {
|
||||
gpu.Model = fmt.Sprintf("Unknown GPU %s:%s", vendorID, deviceID)
|
||||
}
|
||||
|
||||
if gpu.Vendor == VendorNVIDIA {
|
||||
enrichNVIDIA(&gpu)
|
||||
}
|
||||
|
||||
gpus = append(gpus, gpu)
|
||||
}
|
||||
|
||||
slog.Debug("gpu discover: found GPUs", "count", len(gpus))
|
||||
return gpus, nil
|
||||
}
|
||||
|
||||
// enrichFromLspci attempts to populate Model from the system PCI ID database
|
||||
// via lspci -vmm. This covers any card known to the database without needing
|
||||
// vendor-specific tools or a hardcoded table.
|
||||
func enrichFromLspci(g *GPUDevice) {
|
||||
out, err := exec.Command("lspci", "-vmm", "-s", g.PCIAddress).Output()
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
for line := range strings.SplitSeq(string(out), "\n") {
|
||||
name, ok := strings.CutPrefix(line, "Device:\t")
|
||||
if ok && strings.TrimSpace(name) != "" {
|
||||
g.Model = strings.TrimSpace(name)
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// enrichNVIDIA attempts to populate Model and MemoryMiB from nvidia-smi.
|
||||
// Leaves existing values untouched if nvidia-smi is unavailable or fails.
|
||||
func enrichNVIDIA(gpu *GPUDevice) {
|
||||
out, err := exec.Command(
|
||||
"nvidia-smi",
|
||||
"--query-gpu=gpu_name,memory.total",
|
||||
"--format=csv,noheader",
|
||||
fmt.Sprintf("--id=%s", gpu.PCIAddress),
|
||||
).Output()
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
line := strings.TrimSpace(string(out))
|
||||
parts := strings.SplitN(line, ", ", 2)
|
||||
if len(parts) != 2 {
|
||||
return
|
||||
}
|
||||
|
||||
name := strings.TrimSpace(parts[0])
|
||||
memStr := strings.TrimSpace(parts[1])
|
||||
memStr = strings.TrimSuffix(memStr, " MiB")
|
||||
memMiB, err := strconv.ParseInt(memStr, 10, 64)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
gpu.Model = name
|
||||
gpu.MemoryMiB = memMiB
|
||||
}
|
||||
|
||||
// isDisplayClass reports whether a PCI class string (e.g. "0x030200") is a
|
||||
// display controller (class byte 0x03).
|
||||
func isDisplayClass(class string) bool {
|
||||
s := strings.ToLower(strings.TrimSpace(class))
|
||||
s = strings.TrimPrefix(s, "0x")
|
||||
return len(s) >= 2 && s[:2] == "03"
|
||||
}
|
||||
|
||||
// readIOMMUGroup reads the IOMMU group number from the iommu_group symlink.
|
||||
// Returns an error if IOMMU is not active (symlink absent).
|
||||
func readIOMMUGroup(devPath string) (int, error) {
|
||||
target, err := os.Readlink(filepath.Join(devPath, "iommu_group"))
|
||||
if err != nil {
|
||||
return -1, err
|
||||
}
|
||||
n, err := strconv.Atoi(filepath.Base(target))
|
||||
if err != nil {
|
||||
return -1, fmt.Errorf("parse iommu group from %q: %w", target, err)
|
||||
}
|
||||
return n, nil
|
||||
}
|
||||
|
||||
// readDriver returns the basename of the driver currently bound to devPath,
|
||||
// or "" if no driver is bound.
|
||||
func readDriver(devPath string) (string, error) {
|
||||
target, err := os.Readlink(filepath.Join(devPath, "driver"))
|
||||
if os.IsNotExist(err) {
|
||||
return "", nil
|
||||
}
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return filepath.Base(target), nil
|
||||
}
|
||||
|
||||
// vendorFromID maps a PCI vendor hex ID to a Vendor constant.
|
||||
func vendorFromID(vendorID string) Vendor {
|
||||
switch strings.ToLower(vendorID) {
|
||||
case "10de":
|
||||
return VendorNVIDIA
|
||||
case "1002":
|
||||
return VendorAMD
|
||||
case "8086":
|
||||
return VendorIntel
|
||||
default:
|
||||
return VendorUnknown
|
||||
}
|
||||
}
|
||||
|
||||
// readSysfsString reads a sysfs attribute file and returns its trimmed content.
|
||||
func readSysfsString(path string) (string, error) {
|
||||
b, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return strings.TrimSpace(string(b)), nil
|
||||
}
|
||||
|
||||
// normHex strips the "0x" prefix and lowercases a hex string from sysfs.
|
||||
func normHex(s string) string {
|
||||
s = strings.ToLower(strings.TrimSpace(s))
|
||||
return strings.TrimPrefix(s, "0x")
|
||||
}
|
||||
@@ -0,0 +1,229 @@
|
||||
package gpu
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// buildSysfsDevice creates a fake sysfs PCI device under root.
|
||||
// driver and iommuGroup are optional: pass "" / -1 to omit.
|
||||
func buildSysfsDevice(t *testing.T, root, pciAddr, class, vendor, device, driver string, iommuGroup int) {
|
||||
t.Helper()
|
||||
devPath := filepath.Join(root, "bus/pci/devices", pciAddr)
|
||||
must(t, os.MkdirAll(devPath, 0755))
|
||||
writeFile(t, filepath.Join(devPath, "class"), class)
|
||||
writeFile(t, filepath.Join(devPath, "vendor"), vendor)
|
||||
writeFile(t, filepath.Join(devPath, "device"), device)
|
||||
|
||||
if iommuGroup >= 0 {
|
||||
groupPath := filepath.Join(root, "kernel/iommu_groups", strconv.Itoa(iommuGroup))
|
||||
must(t, os.MkdirAll(filepath.Join(groupPath, "devices"), 0755))
|
||||
// symlink: <devPath>/iommu_group -> <groupPath>
|
||||
must(t, os.Symlink(groupPath, filepath.Join(devPath, "iommu_group")))
|
||||
// symlink: <groupPath>/devices/<pciAddr> -> <devPath>
|
||||
must(t, os.Symlink(devPath, filepath.Join(groupPath, "devices", pciAddr)))
|
||||
}
|
||||
|
||||
if driver != "" {
|
||||
driverPath := filepath.Join(root, "bus/pci/drivers", driver)
|
||||
must(t, os.MkdirAll(driverPath, 0755))
|
||||
must(t, os.Symlink(driverPath, filepath.Join(devPath, "driver")))
|
||||
}
|
||||
}
|
||||
|
||||
func writeFile(t *testing.T, path, content string) {
|
||||
t.Helper()
|
||||
must(t, os.WriteFile(path, []byte(content+"\n"), 0444))
|
||||
}
|
||||
|
||||
func must(t *testing.T, err error) {
|
||||
t.Helper()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDiscover_NVIDIA3DController(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
buildSysfsDevice(t, root, "0000:03:00.0", "0x030200", "0x10de", "0x2236", "nvidia", 7)
|
||||
|
||||
gpus, err := discover(root)
|
||||
if err != nil {
|
||||
t.Fatalf("discover: %v", err)
|
||||
}
|
||||
if len(gpus) != 1 {
|
||||
t.Fatalf("want 1 GPU, got %d", len(gpus))
|
||||
}
|
||||
|
||||
g := gpus[0]
|
||||
if g.PCIAddress != "0000:03:00.0" {
|
||||
t.Errorf("PCIAddress = %q, want 0000:03:00.0", g.PCIAddress)
|
||||
}
|
||||
if g.Vendor != VendorNVIDIA {
|
||||
t.Errorf("Vendor = %q, want nvidia", g.Vendor)
|
||||
}
|
||||
if g.VendorID != "10de" {
|
||||
t.Errorf("VendorID = %q, want 10de", g.VendorID)
|
||||
}
|
||||
if g.DeviceID != "2236" {
|
||||
t.Errorf("DeviceID = %q, want 2236", g.DeviceID)
|
||||
}
|
||||
if g.IOMMUGroup != 7 {
|
||||
t.Errorf("IOMMUGroup = %d, want 7", g.IOMMUGroup)
|
||||
}
|
||||
if g.OriginalDriver != "nvidia" {
|
||||
t.Errorf("OriginalDriver = %q, want nvidia", g.OriginalDriver)
|
||||
}
|
||||
if g.Model != "NVIDIA A10" {
|
||||
t.Errorf("Model = %q, want NVIDIA A10", g.Model)
|
||||
}
|
||||
if g.MemoryMiB != 23028 {
|
||||
t.Errorf("MemoryMiB = %d, want 23028", g.MemoryMiB)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDiscover_AMDVGAController(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
buildSysfsDevice(t, root, "0000:05:00.0", "0x030000", "0x1002", "0x744c", "amdgpu", 3)
|
||||
|
||||
gpus, err := discover(root)
|
||||
if err != nil {
|
||||
t.Fatalf("discover: %v", err)
|
||||
}
|
||||
if len(gpus) != 1 {
|
||||
t.Fatalf("want 1 GPU, got %d", len(gpus))
|
||||
}
|
||||
|
||||
g := gpus[0]
|
||||
if g.Vendor != VendorAMD {
|
||||
t.Errorf("Vendor = %q, want amd", g.Vendor)
|
||||
}
|
||||
if g.OriginalDriver != "amdgpu" {
|
||||
t.Errorf("OriginalDriver = %q, want amdgpu", g.OriginalDriver)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDiscover_SkipsNonDisplay(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
// NVMe — class 0x010802 (storage controller)
|
||||
buildSysfsDevice(t, root, "0000:01:00.0", "0x010802", "0x144d", "0xa809", "nvme", 1)
|
||||
// USB controller — class 0x0c0330
|
||||
buildSysfsDevice(t, root, "0000:00:14.0", "0x0c0330", "0x8086", "0x9d2f", "xhci_hcd", 2)
|
||||
|
||||
gpus, err := discover(root)
|
||||
if err != nil {
|
||||
t.Fatalf("discover: %v", err)
|
||||
}
|
||||
if len(gpus) != 0 {
|
||||
t.Errorf("want 0 GPUs, got %d", len(gpus))
|
||||
}
|
||||
}
|
||||
|
||||
func TestDiscover_NoIOMMU(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
// GPU with no iommu_group symlink (IOMMU disabled)
|
||||
buildSysfsDevice(t, root, "0000:03:00.0", "0x030200", "0x10de", "0x2236", "nvidia", -1)
|
||||
|
||||
gpus, err := discover(root)
|
||||
if err != nil {
|
||||
t.Fatalf("discover: %v", err)
|
||||
}
|
||||
if len(gpus) != 1 {
|
||||
t.Fatalf("want 1 GPU, got %d", len(gpus))
|
||||
}
|
||||
if gpus[0].IOMMUGroup != -1 {
|
||||
t.Errorf("IOMMUGroup = %d, want -1 when IOMMU inactive", gpus[0].IOMMUGroup)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDiscover_UnboundDriver(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
// No driver symlink
|
||||
buildSysfsDevice(t, root, "0000:03:00.0", "0x030200", "0x10de", "0x2236", "", 7)
|
||||
|
||||
gpus, err := discover(root)
|
||||
if err != nil {
|
||||
t.Fatalf("discover: %v", err)
|
||||
}
|
||||
if gpus[0].OriginalDriver != "" {
|
||||
t.Errorf("OriginalDriver = %q, want empty for unbound device", gpus[0].OriginalDriver)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDiscover_UnknownDevice(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
// NVIDIA device not in the model table
|
||||
buildSysfsDevice(t, root, "0000:03:00.0", "0x030200", "0x10de", "0xffff", "nvidia", 7)
|
||||
|
||||
gpus, err := discover(root)
|
||||
if err != nil {
|
||||
t.Fatalf("discover: %v", err)
|
||||
}
|
||||
if len(gpus) != 1 {
|
||||
t.Fatalf("want 1 GPU, got %d", len(gpus))
|
||||
}
|
||||
if gpus[0].Model == "" {
|
||||
t.Error("Model should be non-empty fallback for unknown device IDs")
|
||||
}
|
||||
if gpus[0].MemoryMiB != 0 {
|
||||
t.Errorf("MemoryMiB = %d, want 0 for unknown device", gpus[0].MemoryMiB)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDiscover_MultipleDevices(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
buildSysfsDevice(t, root, "0000:03:00.0", "0x030200", "0x10de", "0x2236", "nvidia", 7)
|
||||
buildSysfsDevice(t, root, "0000:04:00.0", "0x030200", "0x10de", "0x2236", "nvidia", 8)
|
||||
buildSysfsDevice(t, root, "0000:01:00.0", "0x010802", "0x144d", "0xa809", "nvme", 1) // non-GPU
|
||||
|
||||
gpus, err := discover(root)
|
||||
if err != nil {
|
||||
t.Fatalf("discover: %v", err)
|
||||
}
|
||||
if len(gpus) != 2 {
|
||||
t.Errorf("want 2 GPUs, got %d", len(gpus))
|
||||
}
|
||||
}
|
||||
|
||||
func TestDiscover_ReadDirError(t *testing.T) {
|
||||
// Pass a sysfsRoot with no bus/pci/devices directory.
|
||||
_, err := discover(t.TempDir())
|
||||
if err == nil {
|
||||
t.Error("want error when devices directory is missing, got nil")
|
||||
}
|
||||
}
|
||||
|
||||
func TestDiscover_SkipsIntelGPU(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
// Intel iGPU (vendor 0x8086) must be skipped even though it is a display class.
|
||||
buildSysfsDevice(t, root, "0000:00:02.0", "0x030000", "0x8086", "0x46a8", "i915", 1)
|
||||
|
||||
gpus, err := discover(root)
|
||||
if err != nil {
|
||||
t.Fatalf("discover: %v", err)
|
||||
}
|
||||
if len(gpus) != 0 {
|
||||
t.Errorf("want Intel GPU skipped, got %d GPUs", len(gpus))
|
||||
}
|
||||
}
|
||||
|
||||
func TestIsDisplayClass(t *testing.T) {
|
||||
cases := []struct {
|
||||
class string
|
||||
want bool
|
||||
}{
|
||||
{"0x030200", true}, // 3D controller
|
||||
{"0x030000", true}, // VGA
|
||||
{"0x038000", true}, // display controller
|
||||
{"0x010802", false}, // NVMe
|
||||
{"0x0c0330", false}, // USB
|
||||
{"0x020000", false}, // Ethernet
|
||||
}
|
||||
for _, c := range cases {
|
||||
if got := isDisplayClass(c.class); got != c.want {
|
||||
t.Errorf("isDisplayClass(%q) = %v, want %v", c.class, got, c.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,46 @@
|
||||
package gpu
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
)
|
||||
|
||||
// GroupMembers returns all PCI devices in the given IOMMU group.
|
||||
// All members must be bound to vfio-pci together for passthrough.
|
||||
func GroupMembers(groupNum int) ([]IOMMUGroupMember, error) {
|
||||
return groupMembers("/sys", groupNum)
|
||||
}
|
||||
|
||||
func groupMembers(sysfsRoot string, groupNum int) ([]IOMMUGroupMember, error) {
|
||||
devicesPath := filepath.Join(sysfsRoot, "kernel/iommu_groups", strconv.Itoa(groupNum), "devices")
|
||||
entries, err := os.ReadDir(devicesPath)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("read iommu group %d: %w", groupNum, err)
|
||||
}
|
||||
|
||||
members := make([]IOMMUGroupMember, 0, len(entries))
|
||||
for _, entry := range entries {
|
||||
pciAddr := entry.Name()
|
||||
// Read device attributes from the canonical pci devices path to avoid
|
||||
// chasing relative symlinks inside the iommu_groups tree.
|
||||
devPath := filepath.Join(sysfsRoot, "bus/pci/devices", pciAddr)
|
||||
|
||||
member := IOMMUGroupMember{PCIAddress: pciAddr}
|
||||
|
||||
if v, err := readSysfsString(filepath.Join(devPath, "vendor")); err == nil {
|
||||
member.VendorID = normHex(v)
|
||||
}
|
||||
if d, err := readSysfsString(filepath.Join(devPath, "device")); err == nil {
|
||||
member.DeviceID = normHex(d)
|
||||
}
|
||||
if c, err := readSysfsString(filepath.Join(devPath, "class")); err == nil {
|
||||
member.Class = c
|
||||
}
|
||||
|
||||
members = append(members, member)
|
||||
}
|
||||
|
||||
return members, nil
|
||||
}
|
||||
@@ -0,0 +1,72 @@
|
||||
package gpu
|
||||
|
||||
import (
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestGroupMembers_GPUWithAudioCompanion(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
// GPU: 3D controller
|
||||
buildSysfsDevice(t, root, "0000:03:00.0", "0x030200", "0x10de", "0x2236", "nvidia", 7)
|
||||
// Companion audio device (HDMI audio on the same card)
|
||||
buildSysfsDevice(t, root, "0000:03:00.1", "0x040300", "0x10de", "0x1aef", "snd_hda_intel", 7)
|
||||
|
||||
members, err := groupMembers(root, 7)
|
||||
if err != nil {
|
||||
t.Fatalf("groupMembers: %v", err)
|
||||
}
|
||||
if len(members) != 2 {
|
||||
t.Fatalf("want 2 members, got %d", len(members))
|
||||
}
|
||||
|
||||
byAddr := make(map[string]IOMMUGroupMember, len(members))
|
||||
for _, m := range members {
|
||||
byAddr[m.PCIAddress] = m
|
||||
}
|
||||
|
||||
gpu, ok := byAddr["0000:03:00.0"]
|
||||
if !ok {
|
||||
t.Fatal("GPU member 0000:03:00.0 not found")
|
||||
}
|
||||
if gpu.VendorID != "10de" {
|
||||
t.Errorf("GPU VendorID = %q, want 10de", gpu.VendorID)
|
||||
}
|
||||
if gpu.DeviceID != "2236" {
|
||||
t.Errorf("GPU DeviceID = %q, want 2236", gpu.DeviceID)
|
||||
}
|
||||
if gpu.Class != "0x030200" {
|
||||
t.Errorf("GPU Class = %q, want 0x030200", gpu.Class)
|
||||
}
|
||||
|
||||
audio, ok := byAddr["0000:03:00.1"]
|
||||
if !ok {
|
||||
t.Fatal("audio member 0000:03:00.1 not found")
|
||||
}
|
||||
if audio.Class != "0x040300" {
|
||||
t.Errorf("audio Class = %q, want 0x040300", audio.Class)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGroupMembers_SingleDevice(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
buildSysfsDevice(t, root, "0000:05:00.0", "0x030200", "0x10de", "0x2330", "nvidia", 12)
|
||||
|
||||
members, err := groupMembers(root, 12)
|
||||
if err != nil {
|
||||
t.Fatalf("groupMembers: %v", err)
|
||||
}
|
||||
if len(members) != 1 {
|
||||
t.Fatalf("want 1 member, got %d", len(members))
|
||||
}
|
||||
if members[0].PCIAddress != "0000:05:00.0" {
|
||||
t.Errorf("PCIAddress = %q, want 0000:05:00.0", members[0].PCIAddress)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGroupMembers_MissingGroup(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
_, err := groupMembers(root, 99)
|
||||
if err == nil {
|
||||
t.Error("want error for non-existent IOMMU group, got nil")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,238 @@
|
||||
package gpu
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"sync"
|
||||
)
|
||||
|
||||
type gpuEntry struct {
|
||||
Device GPUDevice
|
||||
InstanceID string // empty = available
|
||||
Available bool // false = error state requiring operator action
|
||||
groupMembers []IOMMUGroupMember
|
||||
memberDrivers map[string]string // PCI addr -> driver held before VFIO takeover
|
||||
}
|
||||
|
||||
// Manager is a thread-safe pool of passthrough-eligible GPUs.
|
||||
// It owns the VFIO bind/unbind lifecycle for each claimed device.
|
||||
type Manager struct {
|
||||
mu sync.Mutex
|
||||
pool []gpuEntry
|
||||
sysfsRoot string
|
||||
}
|
||||
|
||||
// NewManager creates a Manager from the given discovered GPUs.
|
||||
// All entries start as available.
|
||||
func NewManager(devices []GPUDevice) *Manager {
|
||||
pool := make([]gpuEntry, len(devices))
|
||||
for i, d := range devices {
|
||||
pool[i] = gpuEntry{Device: d, Available: true}
|
||||
}
|
||||
return &Manager{pool: pool, sysfsRoot: "/sys"}
|
||||
}
|
||||
|
||||
// Claim binds the first available GPU (and all its IOMMU group members) to
|
||||
// vfio-pci and associates it with instanceID. Returns the claimed device.
|
||||
func (m *Manager) Claim(instanceID string) (*GPUDevice, error) {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
|
||||
idx := -1
|
||||
for i := range m.pool {
|
||||
if m.pool[i].Available && m.pool[i].InstanceID == "" {
|
||||
idx = i
|
||||
break
|
||||
}
|
||||
}
|
||||
if idx == -1 {
|
||||
return nil, errors.New("no GPU available")
|
||||
}
|
||||
|
||||
entry := &m.pool[idx]
|
||||
if entry.Device.IOMMUGroup < 0 {
|
||||
return nil, fmt.Errorf("GPU %s has no IOMMU group (is IOMMU enabled?)", entry.Device.PCIAddress)
|
||||
}
|
||||
|
||||
members, err := groupMembers(m.sysfsRoot, entry.Device.IOMMUGroup)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("enumerate IOMMU group %d: %w", entry.Device.IOMMUGroup, err)
|
||||
}
|
||||
|
||||
drivers := make(map[string]string, len(members))
|
||||
var bound []string
|
||||
|
||||
for _, member := range members {
|
||||
orig, err := bindVFIO(m.sysfsRoot, member.PCIAddress)
|
||||
if err != nil {
|
||||
// Rollback any members already bound in this attempt.
|
||||
for _, addr := range bound {
|
||||
if rbErr := unbindVFIO(m.sysfsRoot, addr, drivers[addr]); rbErr != nil {
|
||||
slog.Error("GPU claim rollback failed", "addr", addr, "err", rbErr)
|
||||
}
|
||||
}
|
||||
return nil, fmt.Errorf("bind IOMMU group member %s: %w", member.PCIAddress, err)
|
||||
}
|
||||
// bindVFIO returns "vfio-pci" when the device was already bound (idempotent
|
||||
// path — e.g. pre-bound by spx-test-gpu, or after a daemon restart with a
|
||||
// live GPU instance). Preserve "vfio-pci" as the recorded original so
|
||||
// Release knows to skip the unbind/rebind cycle for these devices.
|
||||
if orig == "vfio-pci" {
|
||||
if member.PCIAddress == entry.Device.PCIAddress {
|
||||
orig = entry.Device.OriginalDriver
|
||||
}
|
||||
// else: companion was already vfio-pci bound; keep orig="vfio-pci"
|
||||
// so Release leaves it bound rather than unbinding without a rebind target.
|
||||
}
|
||||
drivers[member.PCIAddress] = orig
|
||||
bound = append(bound, member.PCIAddress)
|
||||
}
|
||||
|
||||
entry.InstanceID = instanceID
|
||||
entry.groupMembers = members
|
||||
entry.memberDrivers = drivers
|
||||
|
||||
slog.Info("GPU claimed", "instance", instanceID,
|
||||
"gpu", entry.Device.PCIAddress, "model", entry.Device.Model)
|
||||
return &entry.Device, nil
|
||||
}
|
||||
|
||||
// Release returns the GPU held by instanceID to the available pool.
|
||||
// For devices that were already bound to vfio-pci before the daemon started
|
||||
// (e.g. by spx-test-gpu), the device is left vfio-pci bound — QEMU has
|
||||
// already released the group fd on exit, so the next Claim works immediately.
|
||||
// For devices that were bound by Claim itself, Release unbinds vfio-pci and
|
||||
// rebinds to the original driver; if that fails the entry is marked unavailable.
|
||||
func (m *Manager) Release(instanceID string) error {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
|
||||
idx := -1
|
||||
for i := range m.pool {
|
||||
if m.pool[i].InstanceID == instanceID {
|
||||
idx = i
|
||||
break
|
||||
}
|
||||
}
|
||||
if idx == -1 {
|
||||
return fmt.Errorf("no GPU claimed by instance %s", instanceID)
|
||||
}
|
||||
|
||||
entry := &m.pool[idx]
|
||||
var firstErr error
|
||||
|
||||
for _, member := range entry.groupMembers {
|
||||
orig := entry.memberDrivers[member.PCIAddress]
|
||||
if orig == "vfio-pci" {
|
||||
// GPU was already bound to vfio-pci before this instance (e.g. by
|
||||
// spx-test-gpu at node setup). QEMU released the group fd on exit;
|
||||
// the device stays vfio-pci bound and is immediately available for
|
||||
// the next claim. Skipping unbind avoids a needless round-trip and
|
||||
// eliminates a brief window where the device is unbound.
|
||||
continue
|
||||
}
|
||||
if err := unbindVFIO(m.sysfsRoot, member.PCIAddress, orig); err != nil {
|
||||
slog.Error("GPU release failed for IOMMU group member",
|
||||
"instance", instanceID, "pci", member.PCIAddress, "err", err)
|
||||
if firstErr == nil {
|
||||
firstErr = err
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
entry.InstanceID = ""
|
||||
entry.groupMembers = nil
|
||||
entry.memberDrivers = nil
|
||||
|
||||
if firstErr != nil {
|
||||
entry.Available = false
|
||||
slog.Error("GPU marked unavailable after failed release — operator action required",
|
||||
"gpu", entry.Device.PCIAddress)
|
||||
return firstErr
|
||||
}
|
||||
|
||||
slog.Info("GPU released", "instance", instanceID, "gpu", entry.Device.PCIAddress)
|
||||
return nil
|
||||
}
|
||||
|
||||
// Available returns the count of GPUs that can be claimed right now.
|
||||
func (m *Manager) Available() int {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
n := 0
|
||||
for _, e := range m.pool {
|
||||
if e.Available && e.InstanceID == "" {
|
||||
n++
|
||||
}
|
||||
}
|
||||
return n
|
||||
}
|
||||
|
||||
// AllocatedCount returns the number of GPUs currently claimed by instances.
|
||||
func (m *Manager) AllocatedCount() int {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
n := 0
|
||||
for _, e := range m.pool {
|
||||
if e.InstanceID != "" {
|
||||
n++
|
||||
}
|
||||
}
|
||||
return n
|
||||
}
|
||||
|
||||
// TotalCount returns the total number of GPUs in the pool regardless of state.
|
||||
func (m *Manager) TotalCount() int {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
return len(m.pool)
|
||||
}
|
||||
|
||||
// ReclaimByAddress marks the GPU at addr as claimed by instanceID without
|
||||
// performing sysfs writes. Used on daemon restart to re-register GPUs that
|
||||
// are still bound to vfio-pci from a previous run.
|
||||
func (m *Manager) ReclaimByAddress(addr, instanceID string) error {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
|
||||
for i := range m.pool {
|
||||
if m.pool[i].Device.PCIAddress != addr {
|
||||
continue
|
||||
}
|
||||
entry := &m.pool[i]
|
||||
if entry.InstanceID == instanceID {
|
||||
// Already owned by the same instance — OnInstanceUp can fire
|
||||
// for a fresh launch where the handler already Claim'd this
|
||||
// slot. Treat as a no-op rather than a conflict.
|
||||
return nil
|
||||
}
|
||||
if entry.InstanceID != "" {
|
||||
return fmt.Errorf("GPU %s already claimed by %s", addr, entry.InstanceID)
|
||||
}
|
||||
|
||||
members, err := groupMembers(m.sysfsRoot, entry.Device.IOMMUGroup)
|
||||
if err != nil {
|
||||
slog.Warn("Failed to re-discover IOMMU group on restart, teardown may be incomplete",
|
||||
"gpu", addr, "group", entry.Device.IOMMUGroup, "err", err)
|
||||
members = []IOMMUGroupMember{{PCIAddress: addr}}
|
||||
}
|
||||
|
||||
// Use the stored pre-passthrough driver for the primary GPU so Release
|
||||
// can rebind correctly. Companion devices are left unbound on release
|
||||
// (original drivers not persisted across restarts).
|
||||
drivers := make(map[string]string, len(members))
|
||||
for _, member := range members {
|
||||
if member.PCIAddress == addr {
|
||||
drivers[member.PCIAddress] = entry.Device.OriginalDriver
|
||||
}
|
||||
}
|
||||
|
||||
entry.InstanceID = instanceID
|
||||
entry.groupMembers = members
|
||||
entry.memberDrivers = drivers
|
||||
slog.Info("GPU re-claimed after daemon restart", "gpu", addr, "instance", instanceID)
|
||||
return nil
|
||||
}
|
||||
return fmt.Errorf("GPU %s not found in pool", addr)
|
||||
}
|
||||
@@ -0,0 +1,375 @@
|
||||
package gpu
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// newManagerForTest creates a Manager backed by a custom sysfsRoot for unit tests.
|
||||
func newManagerForTest(devices []GPUDevice, sysfsRoot string) *Manager {
|
||||
m := NewManager(devices)
|
||||
m.sysfsRoot = sysfsRoot
|
||||
return m
|
||||
}
|
||||
|
||||
// buildManagerSysfs sets up a complete sysfs mock for a single GPU with a
|
||||
// companion audio device in the same IOMMU group. Creates all driver
|
||||
// directories needed for bind/unbind to succeed.
|
||||
func buildManagerSysfs(t *testing.T) (root string, gpu GPUDevice) {
|
||||
t.Helper()
|
||||
root = t.TempDir()
|
||||
|
||||
buildSysfsDevice(t, root, "0000:03:00.0", "0x030200", "0x10de", "0x2236", "nvidia", 7)
|
||||
buildSysfsDevice(t, root, "0000:03:00.1", "0x040300", "0x10de", "0x1aef", "snd_hda_intel", 7)
|
||||
|
||||
makeSysfsDriverDir(t, root, "vfio-pci")
|
||||
makeSysfsDriverDir(t, root, "nvidia")
|
||||
makeSysfsDriverDir(t, root, "snd_hda_intel")
|
||||
|
||||
gpu = GPUDevice{
|
||||
PCIAddress: "0000:03:00.0",
|
||||
Vendor: VendorNVIDIA,
|
||||
VendorID: "10de",
|
||||
DeviceID: "2236",
|
||||
Model: "NVIDIA A10",
|
||||
MemoryMiB: 23028,
|
||||
IOMMUGroup: 7,
|
||||
OriginalDriver: "nvidia",
|
||||
}
|
||||
return root, gpu
|
||||
}
|
||||
|
||||
func TestManagerClaim_Success(t *testing.T) {
|
||||
root, gpu := buildManagerSysfs(t)
|
||||
m := newManagerForTest([]GPUDevice{gpu}, root)
|
||||
|
||||
if m.Available() != 1 {
|
||||
t.Fatalf("want Available=1 before claim, got %d", m.Available())
|
||||
}
|
||||
|
||||
dev, err := m.Claim("i-001")
|
||||
if err != nil {
|
||||
t.Fatalf("Claim: %v", err)
|
||||
}
|
||||
if dev.PCIAddress != gpu.PCIAddress {
|
||||
t.Errorf("claimed device PCI = %q, want %q", dev.PCIAddress, gpu.PCIAddress)
|
||||
}
|
||||
if m.Available() != 0 {
|
||||
t.Errorf("want Available=0 after claim, got %d", m.Available())
|
||||
}
|
||||
if m.AllocatedCount() != 1 {
|
||||
t.Errorf("want AllocatedCount=1 after claim, got %d", m.AllocatedCount())
|
||||
}
|
||||
if m.TotalCount() != 1 {
|
||||
t.Errorf("want TotalCount=1, got %d", m.TotalCount())
|
||||
}
|
||||
}
|
||||
|
||||
func TestManagerClaim_NoGPUAvailable(t *testing.T) {
|
||||
root, gpu := buildManagerSysfs(t)
|
||||
m := newManagerForTest([]GPUDevice{gpu}, root)
|
||||
|
||||
if _, err := m.Claim("i-001"); err != nil {
|
||||
t.Fatalf("first Claim: %v", err)
|
||||
}
|
||||
|
||||
// Second claim on a pool with only one GPU.
|
||||
_, err := m.Claim("i-002")
|
||||
if err == nil {
|
||||
t.Error("want error when no GPU available, got nil")
|
||||
}
|
||||
}
|
||||
|
||||
func TestManagerClaim_NoIOMMU(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
gpu := GPUDevice{PCIAddress: "0000:03:00.0", IOMMUGroup: -1}
|
||||
m := newManagerForTest([]GPUDevice{gpu}, root)
|
||||
|
||||
_, err := m.Claim("i-001")
|
||||
if err == nil {
|
||||
t.Error("want error for GPU with no IOMMU group, got nil")
|
||||
}
|
||||
}
|
||||
|
||||
func TestManagerRelease_Success(t *testing.T) {
|
||||
root, gpu := buildManagerSysfs(t)
|
||||
m := newManagerForTest([]GPUDevice{gpu}, root)
|
||||
|
||||
if _, err := m.Claim("i-001"); err != nil {
|
||||
t.Fatalf("Claim: %v", err)
|
||||
}
|
||||
|
||||
if err := m.Release("i-001"); err != nil {
|
||||
t.Fatalf("Release: %v", err)
|
||||
}
|
||||
|
||||
if m.Available() != 1 {
|
||||
t.Errorf("want Available=1 after release, got %d", m.Available())
|
||||
}
|
||||
if m.AllocatedCount() != 0 {
|
||||
t.Errorf("want AllocatedCount=0 after release, got %d", m.AllocatedCount())
|
||||
}
|
||||
}
|
||||
|
||||
func TestManagerRelease_UnknownInstance(t *testing.T) {
|
||||
root, gpu := buildManagerSysfs(t)
|
||||
m := newManagerForTest([]GPUDevice{gpu}, root)
|
||||
|
||||
err := m.Release("i-does-not-exist")
|
||||
if err == nil {
|
||||
t.Error("want error releasing unknown instance, got nil")
|
||||
}
|
||||
}
|
||||
|
||||
func TestManagerRelease_FailureMarksUnavailable(t *testing.T) {
|
||||
root, gpu := buildManagerSysfs(t)
|
||||
m := newManagerForTest([]GPUDevice{gpu}, root)
|
||||
|
||||
if _, err := m.Claim("i-001"); err != nil {
|
||||
t.Fatalf("Claim: %v", err)
|
||||
}
|
||||
|
||||
// Simulate what the kernel does after vfio-pci binding: update the driver
|
||||
// symlink for each IOMMU group member to point to vfio-pci.
|
||||
vfioDriverPath := filepath.Join(root, "bus/pci/drivers/vfio-pci")
|
||||
for _, addr := range []string{"0000:03:00.0", "0000:03:00.1"} {
|
||||
devPath := filepath.Join(root, "bus/pci/devices", addr)
|
||||
_ = os.Remove(filepath.Join(devPath, "driver"))
|
||||
must(t, os.Symlink(vfioDriverPath, filepath.Join(devPath, "driver")))
|
||||
}
|
||||
|
||||
// Make the vfio-pci driver directory non-writable so the unbind write fails.
|
||||
vfioDir := filepath.Join(root, "bus/pci/drivers/vfio-pci")
|
||||
must(t, os.Chmod(vfioDir, 0o555))
|
||||
t.Cleanup(func() { os.Chmod(vfioDir, 0o755) })
|
||||
|
||||
err := m.Release("i-001")
|
||||
if err == nil {
|
||||
t.Error("want error when vfio-pci unbind fails, got nil")
|
||||
}
|
||||
|
||||
// Pool entry must be marked unavailable so the broken GPU isn't re-offered.
|
||||
if m.Available() != 0 {
|
||||
t.Errorf("want Available=0 after failed release, got %d", m.Available())
|
||||
}
|
||||
if m.TotalCount() != 1 {
|
||||
t.Errorf("TotalCount should still reflect physical GPU count, got %d", m.TotalCount())
|
||||
}
|
||||
}
|
||||
|
||||
func TestManagerClaim_RollbackOnPartialFailure(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
|
||||
buildSysfsDevice(t, root, "0000:03:00.0", "0x030200", "0x10de", "0x2236", "nvidia", 7)
|
||||
buildSysfsDevice(t, root, "0000:03:00.1", "0x040300", "0x10de", "0x1aef", "snd_hda_intel", 7)
|
||||
|
||||
makeSysfsDriverDir(t, root, "vfio-pci")
|
||||
makeSysfsDriverDir(t, root, "nvidia")
|
||||
makeSysfsDriverDir(t, root, "snd_hda_intel")
|
||||
|
||||
// Pre-create the audio device's driver_override as read-only so the bind
|
||||
// sequence fails when it tries to write "vfio-pci" into it.
|
||||
audioDevPath := filepath.Join(root, "bus/pci/devices/0000:03:00.1")
|
||||
must(t, os.WriteFile(filepath.Join(audioDevPath, "driver_override"), []byte(""), 0o444))
|
||||
|
||||
gpu := GPUDevice{
|
||||
PCIAddress: "0000:03:00.0", Vendor: VendorNVIDIA,
|
||||
VendorID: "10de", DeviceID: "2236", IOMMUGroup: 7, OriginalDriver: "nvidia",
|
||||
}
|
||||
m := newManagerForTest([]GPUDevice{gpu}, root)
|
||||
|
||||
_, err := m.Claim("i-001")
|
||||
if err == nil {
|
||||
t.Error("want error when partial IOMMU group bind fails, got nil")
|
||||
}
|
||||
|
||||
// GPU must be back in the pool after rollback.
|
||||
if m.Available() != 1 {
|
||||
t.Errorf("want Available=1 after rollback, got %d", m.Available())
|
||||
}
|
||||
if m.AllocatedCount() != 0 {
|
||||
t.Errorf("want AllocatedCount=0 after rollback, got %d", m.AllocatedCount())
|
||||
}
|
||||
}
|
||||
|
||||
func TestManagerClaim_GroupMembersError(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
// IOMMU group 99 has no devices directory — groupMembers will fail.
|
||||
gpu := GPUDevice{PCIAddress: "0000:03:00.0", IOMMUGroup: 99, OriginalDriver: "nvidia"}
|
||||
m := newManagerForTest([]GPUDevice{gpu}, root)
|
||||
|
||||
_, err := m.Claim("i-001")
|
||||
if err == nil {
|
||||
t.Error("want error when IOMMU group directory is missing, got nil")
|
||||
}
|
||||
}
|
||||
|
||||
// TestManagerRelease_PreBound verifies that when a GPU was already bound to
|
||||
// vfio-pci before the daemon started (OriginalDriver == "vfio-pci"), Release
|
||||
// skips the sysfs unbind/rebind and just returns the device to the pool.
|
||||
// This is the normal path for nodes set up with spx-test-gpu.
|
||||
func TestManagerRelease_PreBound(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
|
||||
// GPU already bound to vfio-pci (simulating spx-test-gpu setup).
|
||||
buildSysfsDevice(t, root, "0000:03:00.0", "0x030200", "0x10de", "0x2487", "vfio-pci", 14)
|
||||
buildSysfsDevice(t, root, "0000:03:00.1", "0x040300", "0x10de", "0x228b", "vfio-pci", 14)
|
||||
makeSysfsDriverDir(t, root, "vfio-pci")
|
||||
|
||||
gpu := GPUDevice{
|
||||
PCIAddress: "0000:03:00.0",
|
||||
Vendor: VendorNVIDIA,
|
||||
VendorID: "10de",
|
||||
DeviceID: "2487",
|
||||
IOMMUGroup: 14,
|
||||
OriginalDriver: "vfio-pci",
|
||||
}
|
||||
m := newManagerForTest([]GPUDevice{gpu}, root)
|
||||
|
||||
if _, err := m.Claim("i-001"); err != nil {
|
||||
t.Fatalf("Claim: %v", err)
|
||||
}
|
||||
|
||||
// Make vfio-pci driver dir read-only to prove Release doesn't write to it.
|
||||
vfioDir := filepath.Join(root, "bus/pci/drivers/vfio-pci")
|
||||
must(t, os.Chmod(vfioDir, 0o555))
|
||||
t.Cleanup(func() { os.Chmod(vfioDir, 0o755) })
|
||||
|
||||
if err := m.Release("i-001"); err != nil {
|
||||
t.Fatalf("Release on pre-bound GPU should not write sysfs, got err: %v", err)
|
||||
}
|
||||
|
||||
if m.Available() != 1 {
|
||||
t.Errorf("want Available=1 after release, got %d", m.Available())
|
||||
}
|
||||
if m.AllocatedCount() != 0 {
|
||||
t.Errorf("want AllocatedCount=0 after release, got %d", m.AllocatedCount())
|
||||
}
|
||||
}
|
||||
|
||||
func TestManagerCounts_MultiGPU(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
|
||||
for i, addr := range []string{"0000:03:00.0", "0000:04:00.0"} {
|
||||
buildSysfsDevice(t, root, addr, "0x030200", "0x10de", "0x2236", "nvidia", i+7)
|
||||
makeSysfsDriverDir(t, root, "nvidia")
|
||||
makeSysfsDriverDir(t, root, "vfio-pci")
|
||||
}
|
||||
|
||||
gpus := []GPUDevice{
|
||||
{PCIAddress: "0000:03:00.0", IOMMUGroup: 7, OriginalDriver: "nvidia"},
|
||||
{PCIAddress: "0000:04:00.0", IOMMUGroup: 8, OriginalDriver: "nvidia"},
|
||||
}
|
||||
m := newManagerForTest(gpus, root)
|
||||
|
||||
if m.TotalCount() != 2 {
|
||||
t.Fatalf("want TotalCount=2, got %d", m.TotalCount())
|
||||
}
|
||||
|
||||
if _, err := m.Claim("i-001"); err != nil {
|
||||
t.Fatalf("first Claim: %v", err)
|
||||
}
|
||||
if m.Available() != 1 || m.AllocatedCount() != 1 {
|
||||
t.Errorf("after first claim: Available=%d AllocatedCount=%d, want 1/1",
|
||||
m.Available(), m.AllocatedCount())
|
||||
}
|
||||
|
||||
if _, err := m.Claim("i-002"); err != nil {
|
||||
t.Fatalf("second Claim: %v", err)
|
||||
}
|
||||
if m.Available() != 0 || m.AllocatedCount() != 2 {
|
||||
t.Errorf("after second claim: Available=%d AllocatedCount=%d, want 0/2",
|
||||
m.Available(), m.AllocatedCount())
|
||||
}
|
||||
|
||||
if err := m.Release("i-001"); err != nil {
|
||||
t.Fatalf("Release: %v", err)
|
||||
}
|
||||
if m.Available() != 1 || m.AllocatedCount() != 1 {
|
||||
t.Errorf("after release: Available=%d AllocatedCount=%d, want 1/1",
|
||||
m.Available(), m.AllocatedCount())
|
||||
}
|
||||
}
|
||||
|
||||
func TestManagerReclaim_Success(t *testing.T) {
|
||||
root, gpu := buildManagerSysfs(t)
|
||||
m := newManagerForTest([]GPUDevice{gpu}, root)
|
||||
|
||||
if err := m.ReclaimByAddress(gpu.PCIAddress, "i-001"); err != nil {
|
||||
t.Fatalf("ReclaimByAddress: %v", err)
|
||||
}
|
||||
if m.AllocatedCount() != 1 {
|
||||
t.Errorf("want AllocatedCount=1 after reclaim, got %d", m.AllocatedCount())
|
||||
}
|
||||
if m.Available() != 0 {
|
||||
t.Errorf("want Available=0 after reclaim, got %d", m.Available())
|
||||
}
|
||||
}
|
||||
|
||||
func TestManagerReclaim_AlreadyClaimed(t *testing.T) {
|
||||
root, gpu := buildManagerSysfs(t)
|
||||
m := newManagerForTest([]GPUDevice{gpu}, root)
|
||||
|
||||
if _, err := m.Claim("i-001"); err != nil {
|
||||
t.Fatalf("Claim: %v", err)
|
||||
}
|
||||
|
||||
err := m.ReclaimByAddress(gpu.PCIAddress, "i-002")
|
||||
if err == nil {
|
||||
t.Error("want error when GPU already claimed, got nil")
|
||||
}
|
||||
}
|
||||
|
||||
// OnInstanceUp can fire on a freshly-launched instance whose handler-side
|
||||
// Claim already populated the slot with the same instance ID. Reclaim must
|
||||
// short-circuit instead of returning an "already claimed" error so callers
|
||||
// don't log a spurious warning on every successful Run/Start.
|
||||
func TestManagerReclaim_SameInstanceIsNoop(t *testing.T) {
|
||||
root, gpu := buildManagerSysfs(t)
|
||||
m := newManagerForTest([]GPUDevice{gpu}, root)
|
||||
|
||||
if _, err := m.Claim("i-001"); err != nil {
|
||||
t.Fatalf("Claim: %v", err)
|
||||
}
|
||||
|
||||
if err := m.ReclaimByAddress(gpu.PCIAddress, "i-001"); err != nil {
|
||||
t.Fatalf("ReclaimByAddress for same instance: got %v, want nil", err)
|
||||
}
|
||||
if m.AllocatedCount() != 1 {
|
||||
t.Errorf("AllocatedCount after same-id reclaim: got %d, want 1", m.AllocatedCount())
|
||||
}
|
||||
}
|
||||
|
||||
func TestManagerReclaim_NotFound(t *testing.T) {
|
||||
root, gpu := buildManagerSysfs(t)
|
||||
m := newManagerForTest([]GPUDevice{gpu}, root)
|
||||
|
||||
err := m.ReclaimByAddress("0000:ff:00.0", "i-001")
|
||||
if err == nil {
|
||||
t.Error("want error for unknown PCI address, got nil")
|
||||
}
|
||||
}
|
||||
|
||||
func TestManagerReclaim_GroupMembersError(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
// GPU with IOMMU group 99 but no group directory — groupMembers falls back to single member.
|
||||
buildSysfsDevice(t, root, "0000:03:00.0", "0x030200", "0x10de", "0x2236", "vfio-pci", -1)
|
||||
gpu := GPUDevice{
|
||||
PCIAddress: "0000:03:00.0",
|
||||
Vendor: VendorNVIDIA,
|
||||
VendorID: "10de",
|
||||
DeviceID: "2236",
|
||||
IOMMUGroup: 99,
|
||||
OriginalDriver: "nvidia",
|
||||
}
|
||||
m := newManagerForTest([]GPUDevice{gpu}, root)
|
||||
|
||||
// Must succeed even though IOMMU group can't be read (fallback to single member).
|
||||
if err := m.ReclaimByAddress(gpu.PCIAddress, "i-001"); err != nil {
|
||||
t.Fatalf("ReclaimByAddress with missing IOMMU group: %v", err)
|
||||
}
|
||||
if m.AllocatedCount() != 1 {
|
||||
t.Errorf("want AllocatedCount=1, got %d", m.AllocatedCount())
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,71 @@
|
||||
package gpu
|
||||
|
||||
type modelInfo struct {
|
||||
name string
|
||||
memoryMiB int64
|
||||
}
|
||||
|
||||
// knownModels maps "vendorID:deviceID" to human-readable model info.
|
||||
// Used as a fallback when nvidia-smi/rocm-smi is unavailable.
|
||||
var knownModels = map[string]modelInfo{
|
||||
// NVIDIA datacenter
|
||||
"10de:2236": {"NVIDIA A10", 23028},
|
||||
"10de:20b2": {"NVIDIA A100 SXM 40GB", 40960},
|
||||
"10de:20b5": {"NVIDIA A100 SXM 80GB", 81920},
|
||||
"10de:20f1": {"NVIDIA A100 PCIe 40GB", 40960},
|
||||
"10de:20f3": {"NVIDIA A100 PCIe 80GB", 81920},
|
||||
"10de:2331": {"NVIDIA H100 SXM", 81920},
|
||||
"10de:2330": {"NVIDIA H100 PCIe", 81920},
|
||||
"10de:233a": {"NVIDIA H100 NVL", 94208},
|
||||
"10de:2233": {"NVIDIA A30", 24576},
|
||||
"10de:26b5": {"NVIDIA L40", 46068},
|
||||
"10de:26b9": {"NVIDIA L40S", 46068},
|
||||
// NVIDIA consumer
|
||||
"10de:2684": {"NVIDIA GeForce RTX 4090", 24576},
|
||||
"10de:2782": {"NVIDIA GeForce RTX 4070", 12288},
|
||||
"10de:2204": {"NVIDIA GeForce RTX 3090", 24576},
|
||||
// AMD datacenter
|
||||
"1002:7448": {"AMD Instinct MI300X", 196608},
|
||||
"1002:740c": {"AMD Instinct MI250X", 65536},
|
||||
// AMD consumer
|
||||
"1002:744c": {"AMD Radeon RX 7900 XTX", 24576},
|
||||
"1002:73bf": {"AMD Radeon RX 6900 XT", 16384},
|
||||
}
|
||||
|
||||
// computeModels is the set of "vendorID:deviceID" keys for datacenter/compute-only
|
||||
// GPUs that have no display output and must use x-vga=off for QEMU passthrough.
|
||||
// Consumer GPUs not in this set default to x-vga=on.
|
||||
var computeModels = map[string]bool{
|
||||
// NVIDIA datacenter
|
||||
"10de:2236": true, // A10
|
||||
"10de:20b2": true, // A100 SXM 40GB
|
||||
"10de:20b5": true, // A100 SXM 80GB
|
||||
"10de:20f1": true, // A100 PCIe 40GB
|
||||
"10de:20f3": true, // A100 PCIe 80GB
|
||||
"10de:2331": true, // H100 SXM
|
||||
"10de:2330": true, // H100 PCIe
|
||||
"10de:233a": true, // H100 NVL
|
||||
"10de:2233": true, // A30
|
||||
"10de:26b5": true, // L40
|
||||
"10de:26b9": true, // L40S
|
||||
// AMD datacenter
|
||||
"1002:7448": true, // MI300X
|
||||
"1002:740c": true, // MI250X
|
||||
}
|
||||
|
||||
// IsComputeGPU reports whether vendorID:deviceID identifies a headless compute GPU
|
||||
// that should use x-vga=off in QEMU. Consumer/display GPUs return false.
|
||||
// Both IDs must be lowercase hex without a "0x" prefix.
|
||||
func IsComputeGPU(vendorID, deviceID string) bool {
|
||||
return computeModels[vendorID+":"+deviceID]
|
||||
}
|
||||
|
||||
// lookupModel returns the model name and VRAM for a known vendorID:deviceID pair.
|
||||
// Both IDs are expected as lowercase hex without a "0x" prefix.
|
||||
// Returns zero values if the pair is not in the table.
|
||||
func lookupModel(vendorID, deviceID string) (name string, memoryMiB int64) {
|
||||
if info, ok := knownModels[vendorID+":"+deviceID]; ok {
|
||||
return info.name, info.memoryMiB
|
||||
}
|
||||
return "", 0
|
||||
}
|
||||
@@ -0,0 +1,40 @@
|
||||
package gpu
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestLookupModel_KnownNVIDIA(t *testing.T) {
|
||||
name, mem := lookupModel("10de", "2236")
|
||||
if name != "NVIDIA A10" {
|
||||
t.Errorf("name = %q, want NVIDIA A10", name)
|
||||
}
|
||||
if mem != 23028 {
|
||||
t.Errorf("memoryMiB = %d, want 23028", mem)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLookupModel_KnownAMD(t *testing.T) {
|
||||
name, mem := lookupModel("1002", "744c")
|
||||
if name != "AMD Radeon RX 7900 XTX" {
|
||||
t.Errorf("name = %q, want AMD Radeon RX 7900 XTX", name)
|
||||
}
|
||||
if mem != 24576 {
|
||||
t.Errorf("memoryMiB = %d, want 24576", mem)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLookupModel_Unknown(t *testing.T) {
|
||||
name, mem := lookupModel("10de", "ffff")
|
||||
if name != "" {
|
||||
t.Errorf("name = %q, want empty for unknown device", name)
|
||||
}
|
||||
if mem != 0 {
|
||||
t.Errorf("memoryMiB = %d, want 0 for unknown device", mem)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLookupModel_UnknownVendor(t *testing.T) {
|
||||
name, mem := lookupModel("dead", "beef")
|
||||
if name != "" || mem != 0 {
|
||||
t.Errorf("want empty for completely unknown vendor:device, got name=%q mem=%d", name, mem)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,31 @@
|
||||
package gpu
|
||||
|
||||
// Vendor identifies the GPU manufacturer.
|
||||
type Vendor string
|
||||
|
||||
const (
|
||||
VendorNVIDIA Vendor = "nvidia"
|
||||
VendorAMD Vendor = "amd"
|
||||
VendorIntel Vendor = "intel"
|
||||
VendorUnknown Vendor = "unknown"
|
||||
)
|
||||
|
||||
// GPUDevice describes a physical GPU discovered on the host.
|
||||
type GPUDevice struct {
|
||||
PCIAddress string // "0000:03:00.0"
|
||||
Vendor Vendor
|
||||
VendorID string // hex, e.g. "10de"
|
||||
DeviceID string // hex, e.g. "2236"
|
||||
Model string // e.g. "NVIDIA A10"
|
||||
MemoryMiB int64 // VRAM in MiB; 0 if undiscoverable
|
||||
IOMMUGroup int // -1 if IOMMU not active
|
||||
OriginalDriver string // driver bound before VFIO takeover, e.g. "nvidia"
|
||||
}
|
||||
|
||||
// IOMMUGroupMember describes one PCI device within an IOMMU group.
|
||||
type IOMMUGroupMember struct {
|
||||
PCIAddress string
|
||||
VendorID string
|
||||
DeviceID string
|
||||
Class string // e.g. "0x030200"
|
||||
}
|
||||
@@ -2,6 +2,62 @@ package instancetypes
|
||||
|
||||
import "slices"
|
||||
|
||||
// GPUModel describes a GPU device model that maps to a specific GPU instance family.
|
||||
type GPUModel struct {
|
||||
VendorID string // PCI vendor ID, lowercase hex, e.g. "10de"
|
||||
DeviceID string // PCI device ID, lowercase hex, e.g. "2236"
|
||||
Family string // instance family prefix, e.g. "g5"
|
||||
Manufacturer string // e.g. "NVIDIA"
|
||||
Name string // e.g. "A10G"
|
||||
MemoryMiB int64
|
||||
}
|
||||
|
||||
var (
|
||||
// NVIDIA GPU models
|
||||
NVIDIAa10g = GPUModel{"10de", "2236", "g5", "NVIDIA", "A10G", 24576}
|
||||
NVIDIAt4 = GPUModel{"10de", "1eb8", "g4dn", "NVIDIA", "T4", 16384}
|
||||
NVIDIAl4 = GPUModel{"10de", "27b8", "g6", "NVIDIA", "L4", 24576} // gr6 uses identical hardware; requires operator config
|
||||
NVIDIAl40s = GPUModel{"10de", "26b9", "g6e", "NVIDIA", "L40S", 49152}
|
||||
NVIDIAv100sxm16 = GPUModel{"10de", "1db1", "p3", "NVIDIA", "V100", 16384} // SXM2 16 GiB
|
||||
NVIDIAv100sxm32 = GPUModel{"10de", "1db3", "p3dn", "NVIDIA", "V100", 32768} // SXM2 32 GiB
|
||||
NVIDIAv100pcie = GPUModel{"10de", "1dba", "p3", "NVIDIA", "V100", 16384} // PCIe 16 GiB
|
||||
NVIDIAa100sxm40 = GPUModel{"10de", "20b0", "p4d", "NVIDIA", "A100", 40960} // SXM4 40 GiB
|
||||
NVIDIAa100sxm80 = GPUModel{"10de", "20b5", "p4de", "NVIDIA", "A100", 81920} // SXM4 80 GiB
|
||||
NVIDIAh100sxm = GPUModel{"10de", "2330", "p5", "NVIDIA", "H100", 81920} // SXM5 80 GiB
|
||||
NVIDIAh100pcie = GPUModel{"10de", "2331", "p5", "NVIDIA", "H100", 81920} // PCIe 80 GiB
|
||||
NVIDIAh200sxm = GPUModel{"10de", "2335", "p5e", "NVIDIA", "H200", 144384} // SXM5 141 GiB
|
||||
|
||||
// AMD GPU models
|
||||
AMDradeonV520 = GPUModel{"1002", "7362", "g4ad", "AMD", "Radeon Pro V520", 8192}
|
||||
)
|
||||
|
||||
var knownGPUModels = []GPUModel{
|
||||
NVIDIAa10g,
|
||||
NVIDIAt4,
|
||||
NVIDIAl4,
|
||||
NVIDIAl40s,
|
||||
NVIDIAv100sxm16,
|
||||
NVIDIAv100sxm32,
|
||||
NVIDIAv100pcie,
|
||||
NVIDIAa100sxm40,
|
||||
NVIDIAa100sxm80,
|
||||
NVIDIAh100sxm,
|
||||
NVIDIAh100pcie,
|
||||
NVIDIAh200sxm,
|
||||
AMDradeonV520,
|
||||
}
|
||||
|
||||
// GPUModelForVendorDevice returns the GPUModel for a PCI vendor/device ID pair,
|
||||
// or nil if the device is not a recognized GPU model.
|
||||
func GPUModelForVendorDevice(vendorID, deviceID string) *GPUModel {
|
||||
for i := range knownGPUModels {
|
||||
if knownGPUModels[i].VendorID == vendorID && knownGPUModels[i].DeviceID == deviceID {
|
||||
return &knownGPUModels[i]
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// cpuGeneration represents a specific CPU microarchitecture generation
|
||||
// and the AWS instance families it maps to.
|
||||
type cpuGeneration struct {
|
||||
@@ -101,6 +157,91 @@ var memorySizes = []instanceSize{
|
||||
// memorySizesSmall is memorySizes without 12xlarge and 24xlarge (older/ARM families).
|
||||
var memorySizesSmall = slices.Clone(memorySizes[:6])
|
||||
|
||||
// g4dnSizes are the single-GPU G4dn instance sizes (1x NVIDIA T4 each).
|
||||
var g4dnSizes = []instanceSize{
|
||||
{"xlarge", 4, 16},
|
||||
{"2xlarge", 8, 32},
|
||||
{"4xlarge", 16, 64},
|
||||
{"8xlarge", 32, 128},
|
||||
{"16xlarge", 64, 256},
|
||||
}
|
||||
|
||||
// g4adSizes are the single-GPU G4ad instance sizes (1x AMD Radeon Pro V520 each).
|
||||
var g4adSizes = []instanceSize{
|
||||
{"xlarge", 4, 16},
|
||||
{"2xlarge", 8, 32},
|
||||
{"4xlarge", 16, 64},
|
||||
}
|
||||
|
||||
// g5Sizes are the single-GPU G5 instance sizes (1x NVIDIA A10G each).
|
||||
var g5Sizes = []instanceSize{
|
||||
{"xlarge", 4, 16},
|
||||
{"2xlarge", 8, 32},
|
||||
{"4xlarge", 16, 64},
|
||||
{"8xlarge", 32, 128},
|
||||
{"16xlarge", 64, 256},
|
||||
}
|
||||
|
||||
// g6Sizes are the single-GPU G6 instance sizes (1x NVIDIA L4 each).
|
||||
var g6Sizes = []instanceSize{
|
||||
{"xlarge", 4, 16},
|
||||
{"2xlarge", 8, 32},
|
||||
{"4xlarge", 16, 64},
|
||||
{"8xlarge", 32, 128},
|
||||
{"16xlarge", 64, 256},
|
||||
}
|
||||
|
||||
// gr6Sizes are the single-GPU Gr6 instance sizes (1x NVIDIA L4, memory-optimized).
|
||||
var gr6Sizes = []instanceSize{
|
||||
{"4xlarge", 16, 128},
|
||||
{"8xlarge", 32, 256},
|
||||
}
|
||||
|
||||
// g6eSizes are the single-GPU G6e instance sizes (1x NVIDIA L40S each).
|
||||
var g6eSizes = []instanceSize{
|
||||
{"xlarge", 4, 32},
|
||||
{"2xlarge", 8, 64},
|
||||
{"4xlarge", 16, 128},
|
||||
{"8xlarge", 32, 256},
|
||||
{"16xlarge", 64, 512},
|
||||
}
|
||||
|
||||
// p3Sizes are the single-GPU P3 instance sizes (1x NVIDIA V100 16 GiB each).
|
||||
// P3 8xlarge and larger have multiple GPUs and are excluded.
|
||||
var p3Sizes = []instanceSize{
|
||||
{"2xlarge", 8, 61},
|
||||
}
|
||||
|
||||
// p3dnSizes are the single-GPU P3dn instance sizes (1x NVIDIA V100 32 GiB each).
|
||||
// P3dn 24xlarge has 8 GPUs and is excluded.
|
||||
var p3dnSizes = []instanceSize{
|
||||
{"2xlarge", 8, 61},
|
||||
}
|
||||
|
||||
// p4dSizes are the single-GPU P4d instance sizes (1x NVIDIA A100 SXM4 40 GiB each).
|
||||
// P4d 24xlarge has 8 GPUs and is excluded.
|
||||
var p4dSizes = []instanceSize{
|
||||
{"xlarge", 4, 32},
|
||||
}
|
||||
|
||||
// p4deSizes are the single-GPU P4de instance sizes (1x NVIDIA A100 SXM4 80 GiB each).
|
||||
// P4de 24xlarge has 8 GPUs and is excluded.
|
||||
var p4deSizes = []instanceSize{
|
||||
{"xlarge", 4, 32},
|
||||
}
|
||||
|
||||
// p5Sizes are the single-GPU P5 instance sizes (1x NVIDIA H100 SXM5 80 GiB each).
|
||||
// P5 48xlarge has 8 GPUs and is excluded.
|
||||
var p5Sizes = []instanceSize{
|
||||
{"4xlarge", 16, 256},
|
||||
}
|
||||
|
||||
// p5eSizes are the single-GPU P5e instance sizes (1x NVIDIA H200 SXM5 141 GiB each).
|
||||
// P5e 48xlarge has 8 GPUs and is excluded.
|
||||
var p5eSizes = []instanceSize{
|
||||
{"4xlarge", 16, 256},
|
||||
}
|
||||
|
||||
// systemSizes defines internal-only instance types for system VMs (LB, NAT GW, etc.).
|
||||
// These are registered in the type map for allocation but excluded from DescribeInstanceTypes.
|
||||
var systemSizes = []instanceSize{
|
||||
@@ -117,8 +258,9 @@ var systemSizes = []instanceSize{
|
||||
// r5n, r5dn, r5b, c6gd, c6gn, c6id, c6in, m6gd, m6id, m6idn, m6in, r6gd, r6id, r6idn, r6in,
|
||||
// c7gd, c7gn, c7i-flex, m7gd, m7i-flex, r7gd, r7iz, c8gd, c8gn, c8i-flex, m8gd, m8i-flex,
|
||||
// r8gd, r8gn, r8gb, r8i-flex — require NVMe instance storage or enhanced networking
|
||||
// - GPU/accelerator: g2-g6, g6e, g6f, gr6, gr6f, p2-p6, inf1-inf2, trn1-trn2, dl1, dl2q — require
|
||||
// GPU, Inferentia, Trainium, or other accelerator hardware. note inf1 and trn1 are not supported since its aws only hardware.
|
||||
// - Multi-GPU only: g4dn 12xlarge, p3 8/16xlarge, p3dn 24xlarge, p4d/p4de 24xlarge, p5/p5e 48xlarge
|
||||
// - AWS-proprietary accelerators: inf1, inf2, trn1, trn2, dl1, dl2q — Inferentia/Trainium chips, AWS-only hardware
|
||||
// - Unsupported GPU: g2, g6f, gr6f, p2, p6 — too old (g2/p2), future/unannounced (p6), or insufficient hardware availability (g6f/gr6f)
|
||||
// - Storage optimized: d2, d3, d3en, h1, i2-i8g, i7ie, i8ge, im4gn, is4gen — require dense HDD/NVMe
|
||||
// - FPGA: f1, f2 — require FPGA hardware
|
||||
// - High memory: u-*, u7i-*, x1, x1e, x2gd, x2idn, x2iedn, x2iezn, x8g — require TB-scale memory
|
||||
@@ -164,6 +306,20 @@ var instanceFamilyDefs = []instanceFamilyDef{
|
||||
// System (internal-only, not exposed via DescribeInstanceTypes)
|
||||
{name: "sys", sizes: systemSizes, currentGen: true},
|
||||
|
||||
// GPU Accelerated — not emitted by generateForGeneration; used by generateGPUTypes.
|
||||
{name: "g4dn", sizes: g4dnSizes, currentGen: true},
|
||||
{name: "g4ad", sizes: g4adSizes, currentGen: true},
|
||||
{name: "g5", sizes: g5Sizes, currentGen: true},
|
||||
{name: "g6", sizes: g6Sizes, currentGen: true},
|
||||
{name: "gr6", sizes: gr6Sizes, currentGen: true}, // same L4 GPU as g6; requires operator config (PCI ID indistinguishable)
|
||||
{name: "g6e", sizes: g6eSizes, currentGen: true},
|
||||
{name: "p3", sizes: p3Sizes, currentGen: false},
|
||||
{name: "p3dn", sizes: p3dnSizes, currentGen: false},
|
||||
{name: "p4d", sizes: p4dSizes, currentGen: true},
|
||||
{name: "p4de", sizes: p4deSizes, currentGen: true},
|
||||
{name: "p5", sizes: p5Sizes, currentGen: true},
|
||||
{name: "p5e", sizes: p5eSizes, currentGen: true},
|
||||
|
||||
// Memory Optimized (1:8 vCPU:memory)
|
||||
{name: "r4", sizes: memorySizesSmall, currentGen: false},
|
||||
{name: "r5", sizes: memorySizes, currentGen: true},
|
||||
|
||||
@@ -219,14 +219,14 @@ func TestDetectARMGeneration(t *testing.T) {
|
||||
|
||||
func TestDetectAndGenerate_Intel(t *testing.T) {
|
||||
cpu := &mockCPU{vendorID: cpuid.Intel, family: 6, model: 143}
|
||||
types := DetectAndGenerate(cpu, "x86_64")
|
||||
types := DetectAndGenerate(cpu, "x86_64", nil)
|
||||
assert.NotEmpty(t, types)
|
||||
assert.True(t, hasFamily(types, "c7i."), "Sapphire Rapids should include c7i")
|
||||
}
|
||||
|
||||
func TestDetectAndGenerate_NormalizesAmd64(t *testing.T) {
|
||||
cpu := &mockCPU{vendorID: cpuid.Intel, family: 6, model: 106}
|
||||
types := DetectAndGenerate(cpu, "amd64")
|
||||
types := DetectAndGenerate(cpu, "amd64", nil)
|
||||
assert.NotEmpty(t, types)
|
||||
assert.True(t, hasFamily(types, "c6i."), "amd64 should normalize to x86_64 and produce Ice Lake types")
|
||||
}
|
||||
|
||||
@@ -95,8 +95,79 @@ func generateSystemTypes(arch string) map[string]*ec2.InstanceTypeInfo {
|
||||
return types
|
||||
}
|
||||
|
||||
// generateGPUTypes generates instance types for the discovered GPU models.
|
||||
// For each unique GPU family, it emits all sizes for that family with GpuInfo populated.
|
||||
// Duplicates (multiple GPUs of the same model) produce only one set of types.
|
||||
// GenerateGPUTypes returns InstanceTypeInfo entries for each GPU model with GpuInfo populated.
|
||||
// It is exported for use by the daemon's hot-reload path.
|
||||
func GenerateGPUTypes(models []GPUModel, arch string) map[string]*ec2.InstanceTypeInfo {
|
||||
return generateGPUTypes(models, arch)
|
||||
}
|
||||
|
||||
func generateGPUTypes(models []GPUModel, arch string) map[string]*ec2.InstanceTypeInfo {
|
||||
types := make(map[string]*ec2.InstanceTypeInfo)
|
||||
seen := make(map[string]bool)
|
||||
|
||||
for _, model := range models {
|
||||
if seen[model.Family] {
|
||||
continue
|
||||
}
|
||||
seen[model.Family] = true
|
||||
|
||||
for _, def := range instanceFamilyDefs {
|
||||
if def.name != model.Family {
|
||||
continue
|
||||
}
|
||||
for _, size := range def.sizes {
|
||||
name := fmt.Sprintf("%s.%s", def.name, size.suffix)
|
||||
types[name] = &ec2.InstanceTypeInfo{
|
||||
InstanceType: aws.String(name),
|
||||
VCpuInfo: &ec2.VCpuInfo{
|
||||
DefaultVCpus: aws.Int64(int64(size.vcpus)),
|
||||
},
|
||||
MemoryInfo: &ec2.MemoryInfo{
|
||||
SizeInMiB: aws.Int64(int64(size.memoryGB * 1024)),
|
||||
},
|
||||
ProcessorInfo: &ec2.ProcessorInfo{
|
||||
SupportedArchitectures: []*string{aws.String(arch)},
|
||||
},
|
||||
GpuInfo: &ec2.GpuInfo{
|
||||
Gpus: []*ec2.GpuDeviceInfo{{
|
||||
Count: aws.Int64(1),
|
||||
Manufacturer: aws.String(model.Manufacturer),
|
||||
Name: aws.String(model.Name),
|
||||
MemoryInfo: &ec2.GpuDeviceMemoryInfo{
|
||||
SizeInMiB: aws.Int64(model.MemoryMiB),
|
||||
},
|
||||
}},
|
||||
TotalGpuMemoryInMiB: aws.Int64(model.MemoryMiB),
|
||||
},
|
||||
CurrentGeneration: aws.Bool(def.currentGen),
|
||||
BurstablePerformanceSupported: aws.Bool(false),
|
||||
Hypervisor: aws.String("kvm"),
|
||||
SupportedVirtualizationTypes: []*string{aws.String("hvm")},
|
||||
SupportedRootDeviceTypes: []*string{aws.String("ebs")},
|
||||
PlacementGroupInfo: &ec2.PlacementGroupInfo{
|
||||
SupportedStrategies: []*string{
|
||||
aws.String("cluster"),
|
||||
aws.String("spread"),
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return types
|
||||
}
|
||||
|
||||
// IsGPUType returns true if the instance type has GPU resources.
|
||||
func IsGPUType(info *ec2.InstanceTypeInfo) bool {
|
||||
return info.GpuInfo != nil && len(info.GpuInfo.Gpus) > 0
|
||||
}
|
||||
|
||||
// DetectAndGenerate detects the host CPU generation and generates matching instance types.
|
||||
func DetectAndGenerate(cpu CPUInfo, arch string) map[string]*ec2.InstanceTypeInfo {
|
||||
// gpuModels is the list of GPU models discovered on the host; pass nil if no GPUs are present.
|
||||
func DetectAndGenerate(cpu CPUInfo, arch string, gpuModels []GPUModel) map[string]*ec2.InstanceTypeInfo {
|
||||
// Normalize Go's "amd64" to the Linux/AWS convention "x86_64".
|
||||
if arch == "amd64" {
|
||||
arch = "x86_64"
|
||||
@@ -108,6 +179,11 @@ func DetectAndGenerate(cpu CPUInfo, arch string) map[string]*ec2.InstanceTypeInf
|
||||
// Merge in system types (always available regardless of CPU generation).
|
||||
maps.Copy(types, generateSystemTypes(arch))
|
||||
|
||||
// Merge in GPU instance types if the host has recognized GPUs.
|
||||
if len(gpuModels) > 0 {
|
||||
maps.Copy(types, generateGPUTypes(gpuModels, arch))
|
||||
}
|
||||
|
||||
if len(types) == 0 {
|
||||
slog.Error("No instance types generated, daemon will be unable to run VMs",
|
||||
"generation", gen.name, "arch", arch)
|
||||
|
||||
@@ -5,6 +5,8 @@ import (
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
cpuid "github.com/klauspost/cpuid/v2"
|
||||
|
||||
"github.com/aws/aws-sdk-go/service/ec2"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
@@ -366,3 +368,86 @@ func TestGenerateInstanceTypes_PlacementGroupInfo(t *testing.T) {
|
||||
assert.True(t, strategies["spread"], "%s should support spread", name)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGPUModelForVendorDevice_Known(t *testing.T) {
|
||||
m := GPUModelForVendorDevice("10de", "2236")
|
||||
require.NotNil(t, m, "NVIDIA A10G should be a known GPU model")
|
||||
assert.Equal(t, "g5", m.Family)
|
||||
assert.Equal(t, "NVIDIA", m.Manufacturer)
|
||||
assert.Equal(t, "A10G", m.Name)
|
||||
assert.Equal(t, int64(24576), m.MemoryMiB)
|
||||
}
|
||||
|
||||
func TestGPUModelForVendorDevice_Unknown(t *testing.T) {
|
||||
assert.Nil(t, GPUModelForVendorDevice("dead", "beef"), "unknown PCI IDs should return nil")
|
||||
assert.Nil(t, GPUModelForVendorDevice("10de", "0000"), "wrong device ID should return nil")
|
||||
}
|
||||
|
||||
func TestGenerateGPUTypes_NVIDIAa10g(t *testing.T) {
|
||||
types := generateGPUTypes([]GPUModel{NVIDIAa10g}, "x86_64")
|
||||
|
||||
// g5 has 5 single-GPU sizes: xlarge, 2xlarge, 4xlarge, 8xlarge, 16xlarge
|
||||
assert.Len(t, types, 5)
|
||||
for _, name := range []string{"g5.xlarge", "g5.2xlarge", "g5.4xlarge", "g5.8xlarge", "g5.16xlarge"} {
|
||||
assert.True(t, hasFamily(types, name), "expected %s", name)
|
||||
}
|
||||
|
||||
it, ok := types["g5.xlarge"]
|
||||
require.True(t, ok)
|
||||
assert.Equal(t, int64(4), *it.VCpuInfo.DefaultVCpus)
|
||||
assert.Equal(t, int64(16384), *it.MemoryInfo.SizeInMiB)
|
||||
require.NotNil(t, it.GpuInfo)
|
||||
require.Len(t, it.GpuInfo.Gpus, 1)
|
||||
assert.Equal(t, int64(1), *it.GpuInfo.Gpus[0].Count)
|
||||
assert.Equal(t, "NVIDIA", *it.GpuInfo.Gpus[0].Manufacturer)
|
||||
assert.Equal(t, "A10G", *it.GpuInfo.Gpus[0].Name)
|
||||
assert.Equal(t, int64(24576), *it.GpuInfo.Gpus[0].MemoryInfo.SizeInMiB)
|
||||
assert.Equal(t, int64(24576), *it.GpuInfo.TotalGpuMemoryInMiB)
|
||||
assert.False(t, *it.BurstablePerformanceSupported)
|
||||
assert.True(t, *it.CurrentGeneration)
|
||||
}
|
||||
|
||||
func TestGenerateGPUTypes_DeduplicatesSameFamily(t *testing.T) {
|
||||
// Two GPUs of the same model (same family) should produce only one set of types.
|
||||
types := generateGPUTypes([]GPUModel{NVIDIAa10g, NVIDIAa10g}, "x86_64")
|
||||
assert.Len(t, types, 5, "duplicate GPU model should not double the type count")
|
||||
}
|
||||
|
||||
func TestGenerateGPUTypes_EmptyModels(t *testing.T) {
|
||||
types := generateGPUTypes(nil, "x86_64")
|
||||
assert.Empty(t, types)
|
||||
}
|
||||
|
||||
func TestIsGPUType(t *testing.T) {
|
||||
gpuTypes := generateGPUTypes([]GPUModel{NVIDIAa10g}, "x86_64")
|
||||
for name, info := range gpuTypes {
|
||||
assert.True(t, IsGPUType(info), "%s should be a GPU type", name)
|
||||
}
|
||||
|
||||
cpuTypes := generateForGeneration(genIntelSkylake, "x86_64")
|
||||
for name, info := range cpuTypes {
|
||||
assert.False(t, IsGPUType(info), "%s should not be a GPU type", name)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDetectAndGenerate_WithGPUModels(t *testing.T) {
|
||||
cpu := &mockCPU{vendorID: cpuid.Intel, family: 6, model: 143} // Sapphire Rapids
|
||||
types := DetectAndGenerate(cpu, "x86_64", []GPUModel{NVIDIAa10g})
|
||||
|
||||
assert.True(t, hasFamily(types, "c7i."), "should include CPU types")
|
||||
assert.True(t, hasFamily(types, "g5."), "should include GPU types when models provided")
|
||||
}
|
||||
|
||||
func TestDetectAndGenerate_WithoutGPUModels(t *testing.T) {
|
||||
cpu := &mockCPU{vendorID: cpuid.Intel, family: 6, model: 143}
|
||||
types := DetectAndGenerate(cpu, "x86_64", nil)
|
||||
|
||||
assert.True(t, hasFamily(types, "c7i."), "should include CPU types")
|
||||
assert.False(t, hasFamily(types, "g5."), "should not include GPU types when no models provided")
|
||||
}
|
||||
|
||||
func TestGenerateGPUTypes_CPUFamiliesNotIncluded(t *testing.T) {
|
||||
// GPU generation must not contaminate CPU-only generation output.
|
||||
cpuTypes := generateForGeneration(genIntelSapphireRapids, "x86_64")
|
||||
assert.False(t, hasFamily(cpuTypes, "g5."), "CPU generation must not emit g5 types")
|
||||
}
|
||||
|
||||
@@ -48,6 +48,15 @@ type DHCPManager struct {
|
||||
jitterFraction float64
|
||||
rng *rand.Rand
|
||||
|
||||
// promiscOn/Off are called to toggle IFF_PROMISC on the DHCP bridge
|
||||
// around each Acquire. nil = no-op (non-veth modes, or tests).
|
||||
// promiscCount tracks in-flight acquires so the bridge is not cleared
|
||||
// while a concurrent DORA is still waiting for a reply.
|
||||
promiscOn func()
|
||||
promiscOff func()
|
||||
promiscMu sync.Mutex
|
||||
promiscCount int
|
||||
|
||||
mu sync.Mutex
|
||||
leases map[string]*managedLease // keyed by ClientID
|
||||
|
||||
@@ -99,6 +108,17 @@ func WithDHCPRand(r *rand.Rand) DHCPManagerOption {
|
||||
return func(m *DHCPManager) { m.rng = r }
|
||||
}
|
||||
|
||||
// WithDHCPPromiscFuncs wires IFF_PROMISC toggling around each Acquire. on is
|
||||
// called before the DORA and off is called after (even on error). Both are
|
||||
// reference-counted so concurrent acquires keep the bridge promisc until the
|
||||
// last one completes. Pass nil to disable (default).
|
||||
func WithDHCPPromiscFuncs(on, off func()) DHCPManagerOption {
|
||||
return func(m *DHCPManager) {
|
||||
m.promiscOn = on
|
||||
m.promiscOff = off
|
||||
}
|
||||
}
|
||||
|
||||
// NewDHCPManager creates a Manager wired against js for KV persistence and
|
||||
// nc for NATS request/reply. The caller still needs to call Bootstrap (to
|
||||
// rehydrate persisted leases) and Subscribe (to start servicing requests).
|
||||
@@ -258,6 +278,9 @@ func (m *DHCPManager) handleAcquire(msg *nats.Msg) {
|
||||
return
|
||||
}
|
||||
|
||||
m.enablePromisc()
|
||||
defer m.disablePromisc()
|
||||
|
||||
ctx, cancel := context.WithTimeout(m.stopCtx, m.acquireTimeout)
|
||||
defer cancel()
|
||||
|
||||
@@ -333,6 +356,34 @@ func (m *DHCPManager) handleRelease(msg *nats.Msg) {
|
||||
m.replyRelease(msg, "")
|
||||
}
|
||||
|
||||
// enablePromisc increments the in-flight acquire count and enables IFF_PROMISC
|
||||
// on the first call.
|
||||
func (m *DHCPManager) enablePromisc() {
|
||||
if m.promiscOn == nil {
|
||||
return
|
||||
}
|
||||
m.promiscMu.Lock()
|
||||
defer m.promiscMu.Unlock()
|
||||
if m.promiscCount == 0 {
|
||||
m.promiscOn()
|
||||
}
|
||||
m.promiscCount++
|
||||
}
|
||||
|
||||
// disablePromisc decrements the count and disables IFF_PROMISC when it reaches
|
||||
// zero (i.e. no more in-flight acquires need it).
|
||||
func (m *DHCPManager) disablePromisc() {
|
||||
if m.promiscOff == nil {
|
||||
return
|
||||
}
|
||||
m.promiscMu.Lock()
|
||||
defer m.promiscMu.Unlock()
|
||||
m.promiscCount--
|
||||
if m.promiscCount == 0 {
|
||||
m.promiscOff()
|
||||
}
|
||||
}
|
||||
|
||||
// trackAndRenew registers the lease in the in-memory map and starts the
|
||||
// per-lease renewal goroutine. Replaces any prior tracker for the same
|
||||
// ClientID.
|
||||
|
||||
+20
-15
@@ -12,21 +12,26 @@ type NodeDiscoverResponse struct {
|
||||
// TotalVCPU - ReservedVCPU - AllocVCPU (and the equivalent for memory);
|
||||
// per-type counts in InstanceTypes already account for the reserve.
|
||||
type NodeStatusResponse struct {
|
||||
Node string `json:"node"`
|
||||
Status string `json:"status"`
|
||||
Host string `json:"host"`
|
||||
Region string `json:"region"`
|
||||
AZ string `json:"az"`
|
||||
Uptime int64 `json:"uptime"`
|
||||
Services []string `json:"services"`
|
||||
TotalVCPU int `json:"total_vcpu"`
|
||||
TotalMemGB float64 `json:"total_mem_gb"`
|
||||
ReservedVCPU int `json:"reserved_vcpu"`
|
||||
ReservedMemGB float64 `json:"reserved_mem_gb"`
|
||||
AllocVCPU int `json:"alloc_vcpu"`
|
||||
AllocMemGB float64 `json:"alloc_mem_gb"`
|
||||
VMCount int `json:"vm_count"`
|
||||
InstanceTypes []InstanceTypeCap `json:"instance_types"`
|
||||
Node string `json:"node"`
|
||||
Status string `json:"status"`
|
||||
Host string `json:"host"`
|
||||
Region string `json:"region"`
|
||||
AZ string `json:"az"`
|
||||
Uptime int64 `json:"uptime"`
|
||||
Services []string `json:"services"`
|
||||
TotalVCPU int `json:"total_vcpu"`
|
||||
TotalMemGB float64 `json:"total_mem_gb"`
|
||||
ReservedVCPU int `json:"reserved_vcpu"`
|
||||
ReservedMemGB float64 `json:"reserved_mem_gb"`
|
||||
AllocVCPU int `json:"alloc_vcpu"`
|
||||
AllocMemGB float64 `json:"alloc_mem_gb"`
|
||||
TotalGPUs int `json:"total_gpus"`
|
||||
AllocGPUs int `json:"alloc_gpus"`
|
||||
GPUCapable bool `json:"gpu_capable,omitempty"`
|
||||
GPUPassthrough bool `json:"gpu_passthrough,omitempty"`
|
||||
GPUModels []string `json:"gpu_models,omitempty"`
|
||||
VMCount int `json:"vm_count"`
|
||||
InstanceTypes []InstanceTypeCap `json:"instance_types"`
|
||||
|
||||
// Leader roles for clustered services (empty string = service not running or not clustered)
|
||||
NATSRole string `json:"nats_role,omitempty"` // "leader" or "follower"
|
||||
|
||||
@@ -148,6 +148,7 @@ type recordingInstanceCleaner struct {
|
||||
releasePublicIP []string
|
||||
detachAndDeleteENI []string
|
||||
removeFromPlacement []string
|
||||
releaseGPU []string
|
||||
}
|
||||
|
||||
func (c *recordingInstanceCleaner) DeleteVolumes(v *VM) {
|
||||
@@ -180,6 +181,12 @@ func (c *recordingInstanceCleaner) RemoveFromPlacementGroup(v *VM) {
|
||||
c.removeFromPlacement = append(c.removeFromPlacement, v.ID)
|
||||
}
|
||||
|
||||
func (c *recordingInstanceCleaner) ReleaseGPU(v *VM) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
c.releaseGPU = append(c.releaseGPU, v.ID)
|
||||
}
|
||||
|
||||
func (c *recordingInstanceCleaner) deleteVolumesCount() int {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
|
||||
@@ -32,4 +32,9 @@ type InstanceCleaner interface {
|
||||
// RemoveFromPlacementGroup removes the instance from its placement
|
||||
// group, if one is set. Called only on Terminate.
|
||||
RemoveFromPlacementGroup(v *VM)
|
||||
|
||||
// ReleaseGPU unbinds the instance's GPU from vfio-pci and rebinds to
|
||||
// its original host driver. Called on both Stop and Terminate. No-op
|
||||
// when the instance has no GPU allocation.
|
||||
ReleaseGPU(v *VM)
|
||||
}
|
||||
|
||||
@@ -253,6 +253,18 @@ func (m *Manager) startQEMU(instance *VM) error {
|
||||
|
||||
instance.Config.Devices = append(instance.Config.Devices, Device{Value: "virtio-rng-pci"})
|
||||
|
||||
if instance.GPUPCIAddress != "" {
|
||||
xvga := "off"
|
||||
if instance.GPUXVGAEnabled {
|
||||
xvga = "on"
|
||||
}
|
||||
instance.Config.Devices = append(instance.Config.Devices, Device{
|
||||
Value: fmt.Sprintf("vfio-pci,host=%s,id=gpu0,x-vga=%s", instance.GPUPCIAddress, xvga),
|
||||
})
|
||||
slog.Info("GPU passthrough device configured",
|
||||
"pci", instance.GPUPCIAddress, "instanceId", instance.ID, "xvga", xvga)
|
||||
}
|
||||
|
||||
qmpSocket, err := utils.GenerateSocketFile(fmt.Sprintf("qmp-%s", instance.ID))
|
||||
if err != nil {
|
||||
slog.Error("Failed to generate QMP socket", "err", err)
|
||||
|
||||
@@ -222,6 +222,9 @@ func (m *Manager) finalizeTerminated(instance *VM) error {
|
||||
func (m *Manager) stopCleanup(instance *VM) {
|
||||
m.shutdownAndUnmount(instance)
|
||||
m.cleanupTapDevices(instance)
|
||||
if m.deps.InstanceCleaner != nil {
|
||||
m.deps.InstanceCleaner.ReleaseGPU(instance)
|
||||
}
|
||||
m.deallocateResources(instance)
|
||||
}
|
||||
|
||||
@@ -238,6 +241,7 @@ func (m *Manager) terminateCleanup(instance *VM) {
|
||||
m.cleanupTapDevices(instance)
|
||||
|
||||
if m.deps.InstanceCleaner != nil {
|
||||
m.deps.InstanceCleaner.ReleaseGPU(instance)
|
||||
m.deps.InstanceCleaner.ReleasePublicIP(instance)
|
||||
m.deps.InstanceCleaner.DetachAndDeleteENI(instance)
|
||||
m.deps.InstanceCleaner.RemoveFromPlacementGroup(instance)
|
||||
|
||||
@@ -389,3 +389,50 @@ func (r *reclaimingStateStore) WriteStoppedInstance(id string, v *VM) error {
|
||||
r.m.Insert(&VM{ID: id, Status: StateRunning})
|
||||
return nil
|
||||
}
|
||||
|
||||
func TestStopCleanup_InvokesReleaseGPU(t *testing.T) {
|
||||
cleaner := &recordingInstanceCleaner{}
|
||||
m := NewManagerWithDeps(Deps{InstanceCleaner: cleaner})
|
||||
instance := &VM{ID: "i-stop", GPUPCIAddress: "0000:01:00.0"}
|
||||
|
||||
m.stopCleanup(instance)
|
||||
|
||||
if got := cleaner.releaseGPU; len(got) != 1 || got[0] != "i-stop" {
|
||||
t.Fatalf("ReleaseGPU on stopCleanup: got %v, want [i-stop]", got)
|
||||
}
|
||||
if len(cleaner.deleteVolumes) != 0 || len(cleaner.releasePublicIP) != 0 || len(cleaner.detachAndDeleteENI) != 0 || len(cleaner.removeFromPlacement) != 0 {
|
||||
t.Fatalf("stopCleanup leaked terminate-only calls: delete=%v pubip=%v eni=%v pg=%v",
|
||||
cleaner.deleteVolumes, cleaner.releasePublicIP, cleaner.detachAndDeleteENI, cleaner.removeFromPlacement)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTerminateCleanup_InvokesReleaseGPU(t *testing.T) {
|
||||
cleaner := &recordingInstanceCleaner{}
|
||||
m := NewManagerWithDeps(Deps{InstanceCleaner: cleaner})
|
||||
instance := &VM{ID: "i-term", GPUPCIAddress: "0000:01:00.0"}
|
||||
|
||||
m.terminateCleanup(instance)
|
||||
|
||||
if got := cleaner.releaseGPU; len(got) != 1 || got[0] != "i-term" {
|
||||
t.Fatalf("ReleaseGPU on terminateCleanup: got %v, want [i-term]", got)
|
||||
}
|
||||
if got := cleaner.deleteVolumes; len(got) != 1 || got[0] != "i-term" {
|
||||
t.Fatalf("DeleteVolumes on terminateCleanup: got %v, want [i-term]", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCleanup_NoGPU_StillInvokesReleaseGPU(t *testing.T) {
|
||||
// The adapter no-ops for GPU-less instances; the manager must still
|
||||
// invoke the method so the adapter owns that decision rather than the
|
||||
// manager second-guessing it.
|
||||
cleaner := &recordingInstanceCleaner{}
|
||||
m := NewManagerWithDeps(Deps{InstanceCleaner: cleaner})
|
||||
instance := &VM{ID: "i-cpu"}
|
||||
|
||||
m.stopCleanup(instance)
|
||||
m.terminateCleanup(instance)
|
||||
|
||||
if got := len(cleaner.releaseGPU); got != 2 {
|
||||
t.Fatalf("ReleaseGPU calls across stop+terminate: got %d, want 2", got)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -117,6 +117,13 @@ type VM struct {
|
||||
// VM (e.g. "elbv2"). Empty for customer-launched instances. The UI
|
||||
// filters out tagged VMs from customer-facing listings.
|
||||
ManagedBy string `json:"managed_by,omitempty"`
|
||||
|
||||
// GPUPCIAddress is the PCI address of the GPU bound to this instance via VFIO
|
||||
// (e.g. "0000:03:00.0"). Empty for non-GPU instances.
|
||||
GPUPCIAddress string `json:"gpu_pci_address,omitempty"`
|
||||
// GPUXVGAEnabled controls whether x-vga=on is passed to QEMU for this instance's
|
||||
// GPU device. True for consumer GPUs; false for headless datacenter cards.
|
||||
GPUXVGAEnabled bool `json:"gpu_xvga_enabled,omitempty"`
|
||||
}
|
||||
|
||||
// ResetNodeLocalState zeroes out fields that are specific to the daemon node
|
||||
@@ -173,6 +180,10 @@ type Config struct {
|
||||
// InstanceType is a friendly name (e.g., t3.micro, t4g.micro)
|
||||
InstanceType string `json:"instance_type"`
|
||||
Architecture string `json:"architecture"`
|
||||
|
||||
// UseUEFI requests OVMF firmware instead of the default SeaBIOS for x86_64 VMs.
|
||||
// Falls back silently to SeaBIOS if OVMF is not installed on the host.
|
||||
UseUEFI bool `json:"use_uefi,omitempty"`
|
||||
}
|
||||
|
||||
func (cfg *Config) Execute() (*exec.Cmd, error) {
|
||||
@@ -304,6 +315,14 @@ func (cfg *Config) Execute() (*exec.Cmd, error) {
|
||||
}
|
||||
} else if cfg.MachineType != "" {
|
||||
args = append(args, "-M", cfg.MachineType)
|
||||
if cfg.Architecture == "x86_64" && cfg.UseUEFI {
|
||||
uefiPath := "/usr/share/ovmf/OVMF.fd"
|
||||
if _, err := os.Stat(uefiPath); err == nil {
|
||||
args = append(args, "-bios", uefiPath)
|
||||
} else {
|
||||
slog.Warn("OVMF firmware not found, falling back to SeaBIOS", "path", uefiPath)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
slog.Info("Executing QEMU command:", "cmd", qemuArchitecture, "args", args)
|
||||
|
||||
Reference in New Issue
Block a user