Tutorial: Writing a Preflight Check

1
package main
2
3
import (
4
	"context"
5
	"fmt"
6
	"log/slog"
7
	"os"
8
	"os/exec"
9
	"strconv"
10
	"strings"
11
	"time"
12
13
	pb "github.com/nvidia/nvsentinel/data-models/pkg/protos"
14
	"google.golang.org/grpc"
15
	"google.golang.org/grpc/credentials/insecure"
16
	"google.golang.org/protobuf/types/known/timestamppb"
17
)
18
19
const (
20
	agentName      = "preflight-demo-gpu-visible"
21
	componentClass = "Node"
22
	checkName      = "GPUVisibility"
23
24
	exitSuccess        = 0
25
	exitTestFailed     = 1
26
	exitConfigError    = 2
27
	exitSendEventError = 3
28
)
29
30
func main() \{
31
	os.Exit(run())
32
\}
33
34
func run() int \{
35
	ctx := context.Background()
36
37
	minGPUs, err := parsePositiveInt("MIN_GPU_COUNT", 1)
38
	if err != nil \{
39
		slog.Error("config error", "error", err)
40
		return exitConfigError
41
	\}
42
43
	nodeName, err := requireEnv("NODE_NAME")
44
	if err != nil \{
45
		return exitConfigError
46
	\}
47
48
	socket, err := requireEnv("PLATFORM_CONNECTOR_SOCKET")
49
	if err != nil \{
50
		return exitConfigError
51
	\}
52
53
	strategy, err := parseProcessingStrategy()
54
	if err != nil \{
55
		return exitConfigError
56
	\}
57
58
	gpuLines, err := listGPUs(ctx)
59
	healthy := err == nil && len(gpuLines) >= minGPUs
60
61
	var message string
62
	var errorCode string
63
	if err != nil \{
64
		message = fmt.Sprintf("nvidia-smi failed: %v", err)
65
		errorCode = "NVIDIA_SMI_ERROR"
66
	\} else if len(gpuLines) < minGPUs \{
67
		message = fmt.Sprintf("expected >= %d GPUs, found %d", minGPUs, len(gpuLines))
68
		errorCode = "NO_GPUS_VISIBLE"
69
	\} else \{
70
		message = fmt.Sprintf("found %d GPU(s): %s", len(gpuLines), strings.Join(gpuLines, "; "))
71
	\}
72
73
	exitCode := exitSuccess
74
	if !healthy \{
75
		exitCode = exitTestFailed
76
	\}
77
78
	if sendErr := sendEvent(ctx, socket, nodeName, strategy, healthy, healthy == false, message, errorCode); sendErr != nil \{
79
		slog.Error("failed to send health event", "error", sendErr)
80
		return exitSendEventError
81
	\}
82
83
	if exitCode != exitSuccess && strategy == pb.ProcessingStrategy_STORE_ONLY \{
84
		slog.Warn("check failed (STORE_ONLY — not blocking pod)")
85
		return exitSuccess
86
	\}
87
88
	return exitCode
89
\}
90
91
func listGPUs(ctx context.Context) ([]string, error) \{
92
	ctx, cancel := context.WithTimeout(ctx, 30*time.Second)
93
	defer cancel()
94
95
	cmd := exec.CommandContext(ctx, "nvidia-smi", "-L")
96
	out, err := cmd.CombinedOutput()
97
	if err != nil \{
98
		return nil, fmt.Errorf("%w: %s", err, strings.TrimSpace(string(out)))
99
	\}
100
101
	lines := strings.Split(strings.TrimSpace(string(out)), "\n")
102
	var gpus []string
103
	for _, line := range lines \{
104
		line = strings.TrimSpace(line)
105
		if line != "" \{
106
			gpus = append(gpus, line)
107
		\}
108
	\}
109
	return gpus, nil
110
\}
111
112
func sendEvent(ctx context.Context, socket, nodeName string, strategy pb.ProcessingStrategy,
113
	isHealthy, isFatal bool, message, errorCode string) error \{
114
115
	recommendedAction := pb.RecommendedAction_NONE
116
	if !isHealthy \{
117
		recommendedAction = pb.RecommendedAction_CONTACT_SUPPORT
118
	\}
119
120
	var codes []string
121
	if errorCode != "" \{
122
		codes = []string\{errorCode\}
123
	\}
124
125
	event := &pb.HealthEvent\{
126
		Version:            1,
127
		Agent:              agentName,
128
		ComponentClass:     componentClass,
129
		CheckName:          checkName,
130
		IsHealthy:          isHealthy,
131
		IsFatal:            isFatal,
132
		Message:            message,
133
		RecommendedAction:  recommendedAction,
134
		ErrorCode:          codes,
135
		GeneratedTimestamp: timestamppb.Now(),
136
		NodeName:           nodeName,
137
		ProcessingStrategy: strategy,
138
	\}
139
140
	events := &pb.HealthEvents\{Version: 1, Events: []*pb.HealthEvent\{event\}\}
141
142
	target := strings.TrimPrefix(socket, "unix://")
143
	conn, err := grpc.NewClient("unix://"+target, grpc.WithTransportCredentials(insecure.NewCredentials()))
144
	if err != nil \{
145
		return err
146
	\}
147
	defer conn.Close()
148
149
	client := pb.NewPlatformConnectorClient(conn)
150
	ctx, cancel := context.WithTimeout(ctx, 30*time.Second)
151
	defer cancel()
152
153
	_, err = client.HealthEventOccurredV1(ctx, events)
154
	return err
155
\}
156
157
func parseProcessingStrategy() (pb.ProcessingStrategy, error) \{
158
	raw := os.Getenv("PROCESSING_STRATEGY")
159
	if raw == "" \{
160
		raw = "EXECUTE_REMEDIATION"
161
	\}
162
	v, ok := pb.ProcessingStrategy_value[raw]
163
	if !ok \{
164
		return 0, fmt.Errorf("invalid PROCESSING_STRATEGY: %s", raw)
165
	\}
166
	return pb.ProcessingStrategy(v), nil
167
\}
168
169
func parsePositiveInt(key string, defaultVal int) (int, error) \{
170
	raw := os.Getenv(key)
171
	if raw == "" \{
172
		return defaultVal, nil
173
	\}
174
	i, err := strconv.Atoi(raw)
175
	if err != nil || i <= 0 \{
176
		return 0, fmt.Errorf("%s must be a positive integer", key)
177
	\}
178
	return i, nil
179
\}
180
181
func requireEnv(key string) (string, error) \{
182
	v := os.Getenv(key)
183
	if v == "" \{
184
		return "", fmt.Errorf("%s is required", key)
185
	\}
186
	return v, nil
187
\}