Skip to main content

GPU Detection with an Agent Extension

Machine learning workloads are often scheduled onto a mixed fleet where only some nodes carry accelerators. This example builds an agent extension that reports the GPU hardware present on the node, so a workload's SVID can record what it was actually running on.

This guide adds custom attributes using the Agent Workload Attestation Extension, an executable that runs on every agent host. To add attributes from the trust-domain server instead of from the agent, see Add Custom Attributes. To compare all four extension points, see Which extension do you need?.

What you'll build​

A Go executable that runs on each agent host, queries nvidia-smi, and returns the node's GPU inventory as custom attributes. Workloads on accelerator nodes receive custom.node_type=gpu-enabled along with the GPU count, model, and memory. Workloads elsewhere receive custom.node_type=cpu-only.

The hardware inventory only exists on the node itself, which is what makes the agent the right place for this rather than a server webhook.

Prerequisites​

  • A Kubernetes cluster registered with Defakto, with the Agent installed
  • Go 1.22 or later on a machine that can build for the agent nodes' architecture
  • Write access to the agent nodes, or an image build pipeline that can place a binary on them
  • NVIDIA drivers and nvidia-smi present on the accelerator nodes

Write the extension​

The executable reads one JSON request per line from stdin, writes one JSON response per line to stdout, and stays running for the lifetime of the agent.

package main

import (
	"bufio"
	"encoding/json"
	"fmt"
	"os"
	"os/exec"
	"strings"
)

// Request is the JSON the agent writes to stdin, one object per line.
type Request struct {
	Meta map[string]string `json:"_meta"`
	PID  string            `json:"pid"`
}

// Response is the JSON written back to stdout. Every field other than Error
// becomes a custom.* attribute on the workload.
type Response struct {
	NodeType  string `json:"node_type,omitempty"`
	GPUCount  string `json:"gpu_count,omitempty"`
	GPUModel  string `json:"gpu_model,omitempty"`
	GPUMemory string `json:"gpu_memory,omitempty"`
	Error     string `json:"error,omitempty"`
}

func main() {
	scanner := bufio.NewScanner(os.Stdin)
	encoder := json.NewEncoder(os.Stdout)

	for scanner.Scan() {
		var req Request
		if err := json.Unmarshal(scanner.Bytes(), &req); err != nil {
			// Report the failure through the protocol rather than exiting.
			encoder.Encode(Response{Error: fmt.Sprintf("parse error: %v", err)})
			continue
		}

		if err := encoder.Encode(detectGPUs()); err != nil {
			// stdout is reserved for responses, so log to stderr.
			fmt.Fprintf(os.Stderr, "encode error: %v\n", err)
		}
	}

	if err := scanner.Err(); err != nil {
		fmt.Fprintf(os.Stderr, "read error: %v\n", err)
	}
}

func detectGPUs() Response {
	cmd := exec.Command("nvidia-smi",
		"--query-gpu=count,name,memory.total",
		"--format=csv,noheader,nounits")

	output, err := cmd.Output()
	if err != nil {
		// nvidia-smi is absent or reported no device. A CPU-only node is a
		// valid answer, not an error.
		return Response{NodeType: "cpu-only"}
	}

	lines := strings.Split(strings.TrimSpace(string(output)), "\n")
	if len(lines) == 0 || lines[0] == "" {
		return Response{NodeType: "cpu-only"}
	}

	parts := strings.Split(lines[0], ",")
	if len(parts) < 3 {
		return Response{NodeType: "cpu-only"}
	}

	return Response{
		NodeType:  "gpu-enabled",
		GPUCount:  strings.TrimSpace(parts[0]),
		GPUModel:  strings.TrimSpace(parts[1]),
		GPUMemory: strings.TrimSpace(parts[2]),
	}
}

Absent hardware returns cpu-only rather than an error. An error response fails workload attestation, so reserve it for cases where the extension genuinely cannot answer.

Build the binary and record its checksum:

CGO_ENABLED=0 go build -o gpu-detector main.go
sha256sum gpu-detector

Install it at /usr/local/bin/gpu-detector on the agent nodes.

Wire it into Defakto​

Add the extension to the spirl-system Helm values, using the checksum from the previous step:

agent:
  extensionWorkloadAttestation:
    cmd: "/usr/local/bin/gpu-detector"
    timeout: "200ms"
    checksum: "sha256:a1b2c3d4e5f6..."

nvidia-smi typically responds in tens of milliseconds, which fits inside the 200ms timeout. Raise the timeout if the fleet includes nodes with many devices.

To use the attributes in SPIFFE IDs, reference them in a SPIFFE ID path template:

/{{custom.node_type}}/{{kubernetes.pod.namespace}}/{{kubernetes.pod.service_account}}

Verify​

Test the executable by hand before relying on it. Pipe a single request into stdin:

echo '{"_meta":{"version":"1.0"},"kubernetes":{"pod":{"name":"ml-training"}},"pid":"5678"}' | /usr/local/bin/gpu-detector

On an accelerator node:

{"node_type":"gpu-enabled","gpu_count":"4","gpu_model":"Tesla V100","gpu_memory":"32768"}

On a node without a GPU:

{"node_type":"cpu-only"}

Then confirm the agent is using it. Agent logs record the process start and each collection:

INFO Extension process started {"cmd": "/usr/local/bin/gpu-detector"}
INFO Custom attributes collected {"count": 4}

Finally, request an SVID from a workload on an accelerator node and confirm the attributes reached the Trust Domain Server. If they are missing, check that Attribute Redaction is not filtering custom.*.

Next steps​