Help us shape the future of Go! Take the 2026 Go Developer Survey

agent

package
v0.6.0 Latest Latest
Warning

This package is not in the latest version of its module.

Go to latest
Published: Oct 28, 2026 License: Apache-2.0 Imports: 13 Imported by: 0

Documentation

Overview

Package agent implements the Mokka Node Agent: a level-triggered reconciler and supervisor that compiles desired GPU simulation state and fans it out to per-component Simulators in parallel.

Index

Constants

View Source
const DefaultRootComplexID = "pci0000:00"

DefaultRootComplexID is the host bridge a synthesized layout hangs every device off, matching what a single-socket profile declares explicitly.

Variables

This section is empty.

Functions

func ValidBDF

func ValidBDF(s string) bool

ValidBDF reports whether s is a PCI address in the form the kernel names sysfs entries with, DDDD:BB:DD.F, lowercase hex.

The check is a boundary, not a nicety: gpu.customConfig authors bus_id by hand, and the value becomes both a path component under the rendered tree and the basis for a root-complex ID. Anything else is dropped rather than placed somewhere it does not belong — a string that is not an address is one no consumer can look up either. The 8-digit domain NVML reports through nvmlPciInfo_t.busId is not this form; the profile field carries the 4-digit one, as busIdLegacy does.

Exported for the simulators that put a device's address somewhere a hand-authored string must not reach on trust. The kernel log is the sharpest of those: it is node-wide, unnamespaced, and read by health agents, so an address carrying a newline would let whoever can edit the profile write kernel lines of their own.

Types

type Agent

type Agent struct {
	// contains filtered or unexported fields
}

Agent is the reconciler and supervisor for all simulators.

func New

func New(cfg Config) *Agent

New returns an Agent from cfg.

func (*Agent) Liveness

func (a *Agent) Liveness() health.Probe

Liveness passes while the last Stage wave completed without error, naming the simulators that failed when it did. It recovers once Stage succeeds again.

func (*Agent) Readiness

func (a *Agent) Readiness() health.Probe

Readiness aggregates every simulator's readiness and attributes the result per simulator, so a red /readyz names which one is not serving.

func (*Agent) Run

func (a *Agent) Run(ctx context.Context) error

Run starts the agent and blocks until ctx is cancelled or a required component fails. On return it executes a best-effort Revoke → Discard teardown.

func (*Agent) Staged added in v0.5.0

func (a *Agent) Staged() health.Probe

Staged passes once the latest Stage wave has written the tree NRI injects. It closes during restaging and teardown, and stays closed after a failed Stage. It ignores Apply and daemon readiness, including an intentional fabricmanager initialization delay.

type Applier

type Applier interface {
	// Apply publishes artifacts after all Stage calls have completed (barrier).
	Apply(ctx context.Context, state *State) error
	// Revoke is the inverse of Apply and runs on shutdown before Discard.
	Revoke(ctx context.Context) error
}

Applier is implemented by simulators whose staged artifacts something outside the node acts on (containerd, NFD, the GPU Operator validator). Only gpudriver, pcibus, and cri implement this interface.

type Config

type Config struct {
	Simulators      []Simulator
	Source          StateSource
	Log             *zap.Logger
	ShutdownTimeout time.Duration
	StagingLockPath string
}

Config carries Agent constructor arguments.

type Daemon

type Daemon interface {
	// Run blocks until ctx is cancelled. ctx is the agent's lifetime, so a Run
	// survives every subsequent state update.
	Run(ctx context.Context) error
	// Reload delivers later States to the running daemon, so a profile edit
	// applies without a pod restart. No-op for daemons holding no staged state.
	Reload(ctx context.Context, state *State) error
}

Daemon is implemented by simulators that supervise a long-running background process (fabricmanager marker loop, mock-ib server). The agent launches Run on the Stage barrier of the first successful reconcile, so a daemon may assume its own Stage artifacts exist.

type DeviceSpec

type DeviceSpec struct {
	Index            int
	UUID             string
	MinorNumber      int
	Serial           string
	PCIBusID         string
	Name             string
	Architecture     string
	ComputeCapMajor  int
	ComputeCapMinor  int
	MemoryTotalBytes uint64
	PCIDeviceID      uint32
	PCISubsystemID   uint32
	// VGPUGuest is set when the GPU is a vGPU as a VM sees it, and carries the
	// host driver the hypervisor writes into the vGPU's PCI config space.
	VGPUGuest *VGPUGuestSpec
}

DeviceSpec carries per-GPU identity and static hardware properties.

type FabricState

type FabricState struct {
	// Profile declares NVLink; does not imply fabricmanager runs.
	Enabled bool
	// Fabricmanager readiness marker directory; empty when the daemon is off.
	ManagerStateDir      string
	ClusterUUID          string
	CliqueID             uint32
	LinksPerGPU          int
	BandwidthPerLinkMbps int
}

FabricState describes the NVLink / NVSwitch fabric configuration.

type IMEXState

type IMEXState struct {
	Enabled   bool
	IMEXMajor int
	CapsMajor int
}

IMEXState describes the IMEX capability surface for the DRA compute-domain plugin.

type MIGComputeInstance added in v0.5.0

type MIGComputeInstance struct {
	ID   uint32
	UUID string
}

MIGComputeInstance is one compute instance, which is one MIG device as a consumer sees it. UUID is what NVML reports for the partition, and the container runtime resolves an allocated partition by that name, so it is carried alongside the ID the capability names are keyed by.

type MIGGPU added in v0.5.0

type MIGGPU struct {
	Minor        int
	GPUInstances []MIGGPUInstance
}

MIGGPU is one partitioned GPU. Minor is its device-node minor, which is how the capability names identify it.

type MIGGPUInstance added in v0.5.0

type MIGGPUInstance struct {
	ID               uint32
	ComputeInstances []MIGComputeInstance
}

MIGGPUInstance is one GPU instance and the compute instances inside it.

type MIGState added in v0.5.0

type MIGState struct {
	// CapsMajor is the char-device major for /dev/nvidia-caps.
	CapsMajor int
	GPUs      []MIGGPU
}

MIGState describes the MIG capability surface: which GPUs boot partitioned and how, so the agent can stage the cap device nodes and the mig-minors table a consumer needs to reach a MIG device it found through NVML.

The instance IDs are compiled from the same profile the mock NVML library loads, via engine.DeclaredMIGLayout, so the two cannot name a partition differently.

func (MIGState) Caps added in v0.5.0

func (m MIGState) Caps() []migcaps.Cap

Caps allocates the node's MIG capability table.

It lives on the state because two simulators consume it and they must agree exactly: one stages the chardevs and the mig-minors table, the other names those same chardevs in the CDI spec that delivers a partition into a container. Allocating twice from the same input would work until the two walked it in different orders, at which point a container would receive the node guarding a different partition than the one it was allocated.

func (MIGState) Partitioned added in v0.5.0

func (m MIGState) Partitioned() bool

Partitioned reports whether any GPU boots with MIG partitions, which is what decides whether the MIG capability surface exists on this node at all.

type NetworkShape

type NetworkShape struct {
	IBEnabled bool
	// HCACount is already resolved from hca_count or hcas_per_gpu * NumGPUs.
	HCACount         int
	HCAType          string
	FWVersion        string
	HWRev            string
	BoardID          string
	NodeDescTemplate string
	LinkLayer        string
	RateGbps         int
	PortState        string
	PhysState        string
	GUIDPrefix       string
}

NetworkShape describes the simulated InfiniBand HCAs attached to the node. Every field is resolved at compile time so simulators never re-derive defaults. All fields are comparable, so == detects a real IB change.

type NodeMeta

type NodeMeta struct {
	NodeName string
	Hostname string
	HostRoot string
}

NodeMeta carries node identity fields.

type NodeShape

type NodeShape struct {
	NumGPUs  int
	Topology PCIeTopology
	Network  NetworkShape
}

NodeShape describes the simulated node's topology.

type PCIeTopology

type PCIeTopology struct {
	RootComplexes []RootComplex
	CoresPerNUMA  int
}

PCIeTopology describes the PCIe root-complex / NUMA layout.

type Requeuer added in v0.6.0

type Requeuer interface {
	Requeue() <-chan struct{}
}

Requeuer is implemented by simulators whose Stage publishes work they finish in the background, such as a download. Each value on Requeue reconciles the latest State again, so that work is staged under the same staging lock as the rest of the tree instead of being written while NRI decides.

type RootComplex

type RootComplex struct {
	ID          string
	NUMANode    int
	DeviceBDFs  []string
	CPUAffinity string
}

RootComplex is one PCI host bridge with its attached GPU BDFs and NUMA node.

type Simulator

type Simulator interface {
	Name() string
	// Stage materializes artifacts that are not yet externally visible.
	// Re-runs on every State change and must be idempotent.
	Stage(ctx context.Context, state *State) error
	// Discard is the inverse of Stage and runs only on shutdown.
	Discard(ctx context.Context) error
	// Ready reports whether the simulator is successfully done staging and publishing its artifacts.
	Ready() bool
}

Simulator owns one simulated component's reconciliation lifecycle. Stage and Discard are exact inverses.

type SoftwareVersions

type SoftwareVersions struct {
	DriverVersion string
	NVMLVersion   string
	CUDAVersion   string
}

SoftwareVersions holds driver / NVML / CUDA version strings.

type State

type State struct {
	Generation int64
	Node       NodeMeta
	Software   SoftwareVersions
	NodeShape  NodeShape
	Devices    []DeviceSpec
	Switches   []SwitchSpec
	Fabric     FabricState
	IMEX       IMEXState
	MIG        MIGState
	// ConfigRaw holds the raw YAML profile bytes so gpudriver can write the
	// engine config without re-deriving it from the narrower State fields.
	// TODO(https://github.com/NVIDIA/k8s-test-infra/issues/717): replace with Profile/Runtime *config.YAMLConfig split — Profile carries
	// the static hardware model, Runtime carries operational overrides (GPU_COUNT,
	// faults, fabric) written by GPU_COUNT env, MEP-0001 CP, and nvml-mock-ctl.
	ConfigRaw []byte
	// TopologyRaw holds the cluster ComputeDomain topology document verbatim.
	// It stays raw because the document describes every node and the mock NVML
	// engine picks this node's entry by NODE_NAME at load time, so there is no
	// per-node view for the agent to compile.
	TopologyRaw []byte
	// MIGProfilesRaw holds the board's MIG partition table verbatim, empty on
	// a board that has none. It is carried so gpudriver can stage it beside
	// every config it writes: a consumer process resolves the table relative
	// to the config it was given, and the device plugin is handed no
	// environment that could name it anywhere else.
	MIGProfilesRaw []byte
}

State is the compiled desired simulation state the agent reconciles toward. It maps 1:1 to what MEP-0001's Control Plane will emit.

func (*State) HasPCITopology

func (s *State) HasPCITopology() bool

HasPCITopology reports whether a PCI sysfs tree will be rendered for this state. Whoever serves that tree gates on this, so it answers from the same reconciliation the renderer reads rather than from a filesystem probe: a bind mount whose source is missing fails container creation for the whole pod.

func (*State) PCITopology

func (s *State) PCITopology() []RootComplex

PCITopology returns the root complexes to render, reconciled against the devices that exist at runtime rather than taken from the profile as written. Devices are the authority because they are what NVML reports, and a served tree that disagrees with NVML is worse than no tree: a consumer resolves a GPU in one and not the other.

The GPUs are not the whole tree. An HGX baseboard's NVSwitches sit on the node's PCIe bus alongside them, so Switches are rendered too and a consumer enumerating the bus finds the fabric silicon real hardware would show it.

The profile's pcie_topology outlives its device list — GPU_COUNT truncates Devices and leaves the layout whole — so declared BDFs nothing claims are dropped, along with any root that empties out. A device no root claims is rendered under a root its own address implies, never folded into a declared one: locality is what a consumer reads this tree for, so a GPU carrying the NUMA node of a root the profile never put it in is worse than one carrying none. See adopt.

Returns nil when no device carries a BDF, which is the signal to render nothing.

type StateSource

type StateSource interface {
	Watch(ctx context.Context) <-chan Update
	Close() error
}

StateSource emits State observations. Watch sends the current State immediately on subscribe and pushes on change. A closed channel means the source is terminally done; the agent stops.

type SwitchSpec added in v0.5.0

type SwitchSpec struct {
	PCIBusID       string
	PCIDeviceID    uint32
	PCISubsystemID uint32
}

SwitchSpec is one NVSwitch that sits on this node's own PCIe bus, as an HGX baseboard's do. It carries only a PCI identity because that is the whole of what the mock can simulate: NVML models an NVSwitch as the far end of a GPU's NVLink and offers no per-switch API, and the nvmlUnit* chassis calls stay stubbed the way they are on real DGX/HGX nodes.

A rack-scale platform has NVSwitches that are not on the node's bus at all — on GB200/GB300 NVL they live in their own switch trays, reachable over NVLink and invisible to the compute tray's lspci. Those profiles declare the switches for NVLink topology and no PCI identity, so nothing is compiled here for them.

type Update

type Update struct {
	State *State
	Err   error
	At    time.Time
}

Update is one observation from a StateSource. Exactly one of State or Err is set. Err leaves the last good State in force; the agent does not reconcile.

type VGPUGuestSpec added in v0.6.0

type VGPUGuestSpec struct {
	HostDriverVersion string
	HostDriverBranch  string
}

VGPUGuestSpec is the host driver a vGPU guest reads from PCI config space.

Directories

Path Synopsis
Package cpulocality decides which host CPUs are local to a PCI root complex.
Package cpulocality decides which host CPUs are local to a PCI root complex.
Package cri implements the container runtime simulator: it owns what the node's container runtime needs from Mokka to inject mock GPU devices into workload containers.
Package cri implements the container runtime simulator: it owns what the node's container runtime needs from Mokka to inject mock GPU devices into workload containers.
Package fabricmanager simulates nv-fabricmanager, the service that registers GPUs with the NVSwitches on an NVSwitch platform before they are usable.
Package fabricmanager simulates nv-fabricmanager, the service that registers GPUs with the NVSwitches on an NVSwitch platform before they are usable.
Package gpudriver implements the GPU driver footprint simulator: chardevs, NVML/CUDA shims, nvidia-smi, procfs entries, engine config, and a bind mount of that tree onto /run/nvidia/driver, the path GPU Operator's own operands default to.
Package gpudriver implements the GPU driver footprint simulator: chardevs, NVML/CUDA shims, nvidia-smi, procfs entries, engine config, and a bind mount of that tree onto /run/nvidia/driver, the path GPU Operator's own operands default to.
Package host abstracts on-host filesystem roots so simulators can retarget to t.TempDir() in tests without mounting a real /host.
Package host abstracts on-host filesystem roots so simulators can retarget to t.TempDir() in tests without mounting a real /host.
Package ib implements the InfiniBand HCA simulator: a fake /sys/class/infiniband tree, the LD_PRELOAD shims and IB CLI tools that read it, and the mock-ib daemon serving UMAD and verbs over a Unix socket.
Package ib implements the InfiniBand HCA simulator: a fake /sys/class/infiniband tree, the LD_PRELOAD shims and IB CLI tools that read it, and the mock-ib daemon serving UMAD and verbs over a Unix socket.
Package imex implements the IMEX simulator: the capability surface the NVIDIA kernel module provides (channel character devices, a /proc/devices overlay and the fabric-imex-mgmt capability file), read by the NVIDIA DRA compute-domain kubelet plugin, and the IMEX node software (nvidia-imex and nvidia-imex-ctl) its ComputeDomain daemon runs.
Package imex implements the IMEX simulator: the capability surface the NVIDIA kernel module provides (channel character devices, a /proc/devices overlay and the fabric-imex-mgmt capability file), read by the NVIDIA DRA compute-domain kubelet plugin, and the IMEX node software (nvidia-imex and nvidia-imex-ctl) its ComputeDomain daemon runs.
Package kernellog simulates the kernel side of an Xid: the printk a real driver emits when a GPU faults.
Package kernellog simulates the kernel side of an Xid: the printk a real driver emits when a GPU faults.
Package migcaps implements the MIG capability surface simulator: the /dev/nvidia-caps character devices, the capabilities/gpu*/mig procfs tree, and the mig-minors table that maps between them.
Package migcaps implements the MIG capability surface simulator: the /dev/nvidia-caps character devices, the capabilities/gpu*/mig procfs tree, and the mig-minors table that maps between them.
Package nvlink stages the ComputeDomain topology document, which gives each node its NVLink fabric identity.
Package nvlink stages the ComputeDomain topology document, which gives each node its NVLink fabric identity.
Package pcibus implements the PCI bus simulator: it renders a fake /sys/bus/pci/devices tree and stages libmockfs.so so that lspci and topology-aware schedulers see mock GPU BDFs.
Package pcibus implements the PCI bus simulator: it renders a fake /sys/bus/pci/devices tree and stages libmockfs.so so that lspci and topology-aware schedulers see mock GPU BDFs.
Package source provides StateSource implementations for the Mokka Node Agent.
Package source provides StateSource implementations for the Mokka Node Agent.

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL