Tutorial: Writing a Preflight Check

package main
import (
	"context"
	"fmt"
	"log/slog"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"time"
	pb "github.com/nvidia/nvsentinel/data-models/pkg/protos"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials/insecure"
	"google.golang.org/protobuf/types/known/timestamppb"
)
const (
	agentName      = "preflight-demo-gpu-visible"
	componentClass = "Node"
	checkName      = "GPUVisibility"
	exitSuccess        = 0
	exitTestFailed     = 1
	exitConfigError    = 2
	exitSendEventError = 3
)
func main() {
	os.Exit(run())
}
func run() int {
	ctx := context.Background()
	minGPUs, err := parsePositiveInt("MIN_GPU_COUNT", 1)
	if err != nil {
		slog.Error("config error", "error", err)
		return exitConfigError
	}
	nodeName, err := requireEnv("NODE_NAME")
	if err != nil {
		return exitConfigError
	}
	socket, err := requireEnv("PLATFORM_CONNECTOR_SOCKET")
	if err != nil {
		return exitConfigError
	}
	strategy, err := parseProcessingStrategy()
	if err != nil {
		return exitConfigError
	}
	gpuLines, err := listGPUs(ctx)
	healthy := err == nil && len(gpuLines) >= minGPUs
	var message string
	var errorCode string
	if err != nil {
		message = fmt.Sprintf("nvidia-smi failed: %v", err)
		errorCode = "NVIDIA_SMI_ERROR"
	} else if len(gpuLines) < minGPUs {
		message = fmt.Sprintf("expected >= %d GPUs, found %d", minGPUs, len(gpuLines))
		errorCode = "NO_GPUS_VISIBLE"
	} else {
		message = fmt.Sprintf("found %d GPU(s): %s", len(gpuLines), strings.Join(gpuLines, "; "))
	}
	exitCode := exitSuccess
	if !healthy {
		exitCode = exitTestFailed
	}
	if sendErr := sendEvent(ctx, socket, nodeName, strategy, healthy, healthy == false, message, errorCode); sendErr != nil {
		slog.Error("failed to send health event", "error", sendErr)
		return exitSendEventError
	}
	if exitCode != exitSuccess && strategy == pb.ProcessingStrategy_STORE_ONLY {
		slog.Warn("check failed (STORE_ONLY — not blocking pod)")
		return exitSuccess
	}
	return exitCode
}
func listGPUs(ctx context.Context) ([]string, error) {
	ctx, cancel := context.WithTimeout(ctx, 30*time.Second)
	defer cancel()
	cmd := exec.CommandContext(ctx, "nvidia-smi", "-L")
	out, err := cmd.CombinedOutput()
	if err != nil {
		return nil, fmt.Errorf("%w: %s", err, strings.TrimSpace(string(out)))
	}
	lines := strings.Split(strings.TrimSpace(string(out)), "\n")
	var gpus []string
	for _, line := range lines {
		line = strings.TrimSpace(line)
		if line != "" {
			gpus = append(gpus, line)
		}
	}
	return gpus, nil
}
func sendEvent(ctx context.Context, socket, nodeName string, strategy pb.ProcessingStrategy,
	isHealthy, isFatal bool, message, errorCode string) error {
	recommendedAction := pb.RecommendedAction_NONE
	if !isHealthy {
		recommendedAction = pb.RecommendedAction_CONTACT_SUPPORT
	}
	var codes []string
	if errorCode != "" {
		codes = []string{errorCode}
	}
	event := &pb.HealthEvent{
		Version:            1,
		Agent:              agentName,
		ComponentClass:     componentClass,
		CheckName:          checkName,
		IsHealthy:          isHealthy,
		IsFatal:            isFatal,
		Message:            message,
		RecommendedAction:  recommendedAction,
		ErrorCode:          codes,
		GeneratedTimestamp: timestamppb.Now(),
		NodeName:           nodeName,
		ProcessingStrategy: strategy,
	}
	events := &pb.HealthEvents{Version: 1, Events: []*pb.HealthEvent{event}}
	target := strings.TrimPrefix(socket, "unix://")
	conn, err := grpc.NewClient("unix://"+target, grpc.WithTransportCredentials(insecure.NewCredentials()))
	if err != nil {
		return err
	}
	defer conn.Close()
	client := pb.NewPlatformConnectorClient(conn)
	ctx, cancel := context.WithTimeout(ctx, 30*time.Second)
	defer cancel()
	_, err = client.HealthEventOccurredV1(ctx, events)
	return err
}
func parseProcessingStrategy() (pb.ProcessingStrategy, error) {
	raw := os.Getenv("PROCESSING_STRATEGY")
	if raw == "" {
		raw = "EXECUTE_REMEDIATION"
	}
	v, ok := pb.ProcessingStrategy_value[raw]
	if !ok {
		return 0, fmt.Errorf("invalid PROCESSING_STRATEGY: %s", raw)
	}
	return pb.ProcessingStrategy(v), nil
}
func parsePositiveInt(key string, defaultVal int) (int, error) {
	raw := os.Getenv(key)
	if raw == "" {
		return defaultVal, nil
	}
	i, err := strconv.Atoi(raw)
	if err != nil || i <= 0 {
		return 0, fmt.Errorf("%s must be a positive integer", key)
	}
	return i, nil
}
func requireEnv(key string) (string, error) {
	v := os.Getenv(key)
	if v == "" {
		return "", fmt.Errorf("%s is required", key)
	}
	return v, nil
}