fix(snapshot-agent): filter out 0-VRAM daemon processes during PID di… (#148)

* fix(snapshot-agent): filter out 0-VRAM daemon processes during PID discovery to prevent cuda-checkpoint initialization errors

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* fix lint: avoid uint64->int64 conversion flagged by gosec G115

Compare UsedGpuMemory against the NVML_VALUE_NOT_AVAILABLE sentinel in
uint64 space instead of casting.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* chart(snapshot-agent): tolerate the timeslice.io/shared taint by default

The snapshot-agent must run on the shared GPU nodes that the timeslice
stack itself taints, so the chart's DaemonSet should tolerate the taint
out of the box.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

---------

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

Author: jessicaochen

deploy/snapshot-agent/values.yaml

@@ -25,6 +25,9 @@ tolerations:
     operator: "Equal"
     value: "present"
     effect: "NoSchedule"
+  - key: "timeslice.io/shared"
+    operator: "Exists"
+    effect: "NoSchedule"
 
 affinity: {}
 

pkg/snapshot-agent/utils/pod-utils.go

@@ -41,6 +41,10 @@ var (
 	IsPIDInPodCgroupFunc = isPIDInPodCgroup
 )
 
+// nvmlValueNotAvailable is nvml.VALUE_NOT_AVAILABLE (-1) as it appears in
+// unsigned fields such as ProcessInfo.UsedGpuMemory.
+const nvmlValueNotAvailable = ^uint64(0)
+
 type DeviceInterface interface {
 	GetComputeRunningProcesses() ([]nvml.ProcessInfo, nvml.Return)
 	GetGraphicsRunningProcesses() ([]nvml.ProcessInfo, nvml.Return)
@@ -150,6 +154,11 @@ func getPodPIDsInternal(ctx context.Context, podName, namespace string) ([]int,
 		}
 
 		for _, proc := range procs {
+			// Skip non-computing parent/daemon processes that hold 0 VRAM; they have no
+			// active CUDA context, so cuda-checkpoint fails to lock them.
+			if proc.UsedGpuMemory == 0 || proc.UsedGpuMemory == nvmlValueNotAvailable {
+				continue
+			}
 			pid := int(proc.Pid)
 			if seenPIDs[pid] {
 				continue
@@ -170,6 +179,11 @@ func getPodPIDsInternal(ctx context.Context, podName, namespace string) ([]int,
 		graphicsProcs, ret := device.GetGraphicsRunningProcesses()
 		if ret == nvml.SUCCESS {
 			for _, proc := range graphicsProcs {
+				// Skip non-computing parent/daemon processes that hold 0 VRAM; they have no
+				// active CUDA context, so cuda-checkpoint fails to lock them.
+				if proc.UsedGpuMemory == 0 || proc.UsedGpuMemory == nvmlValueNotAvailable {
+					continue
+				}
 				pid := int(proc.Pid)
 				if seenPIDs[pid] {
 					continue

pkg/snapshot-agent/utils/pod-utils_test.go

@@ -77,8 +77,8 @@ func TestGetPodPIDs(t *testing.T) {
 				snapshotutils.NvmlInit = func() nvml.Return { return nvml.SUCCESS }
 				snapshotutils.NvmlDeviceGetCount = func() (int, nvml.Return) { return 1, nvml.SUCCESS }
 				device := &mockDevice{
-					computeProcs:  []nvml.ProcessInfo{{Pid: 100}, {Pid: 200}},
-					graphicsProcs: []nvml.ProcessInfo{{Pid: 200}, {Pid: 300}},
+					computeProcs:  []nvml.ProcessInfo{{Pid: 100, UsedGpuMemory: 1024}, {Pid: 200, UsedGpuMemory: 1024}},
+					graphicsProcs: []nvml.ProcessInfo{{Pid: 200, UsedGpuMemory: 1024}, {Pid: 300, UsedGpuMemory: 1024}},
 				}
 				snapshotutils.NvmlDeviceGetHandleByIndex = func(index int) (snapshotutils.DeviceInterface, nvml.Return) {
 					return device, nvml.SUCCESS
@@ -90,6 +90,39 @@ func TestGetPodPIDs(t *testing.T) {
 			expectedPIDs: []int{100, 300},
 			expectError:  false,
 		},
+		{
+			name:      "Ignore Zero VRAM Processes",
+			podName:   "test-pod",
+			namespace: "test-ns",
+			setupMocks: func() {
+				podUID := "test-uid"
+				pod := &corev1.Pod{
+					ObjectMeta: metav1.ObjectMeta{
+						Name:      "test-pod",
+						Namespace: "test-ns",
+						UID:       types.UID(podUID),
+					},
+				}
+				snapshotutils.GetK8sClient = func() (kubernetes.Interface, error) {
+					return fake.NewSimpleClientset(pod), nil
+				}
+				snapshotutils.NvmlInit = func() nvml.Return { return nvml.SUCCESS }
+				snapshotutils.NvmlDeviceGetCount = func() (int, nvml.Return) { return 1, nvml.SUCCESS }
+				device := &mockDevice{
+					computeProcs: []nvml.ProcessInfo{{Pid: 100, UsedGpuMemory: 1024}, {Pid: 200, UsedGpuMemory: 0}},
+					// ^uint64(0) is NVML_VALUE_NOT_AVAILABLE (-1) as reported in the uint64 field.
+					graphicsProcs: []nvml.ProcessInfo{{Pid: 300, UsedGpuMemory: ^uint64(0)}},
+				}
+				snapshotutils.NvmlDeviceGetHandleByIndex = func(index int) (snapshotutils.DeviceInterface, nvml.Return) {
+					return device, nvml.SUCCESS
+				}
+				snapshotutils.IsPIDInPodCgroupFunc = func(pid int, uid string) (bool, error) {
+					return true, nil
+				}
+			},
+			expectedPIDs: []int{100},
+			expectError:  false,
+		},
 		{
 			name:      "GetPodUID Failure",
 			podName:   "pod",
@@ -149,7 +182,7 @@ func TestGetPodPIDs(t *testing.T) {
 				snapshotutils.NvmlInit = func() nvml.Return { return nvml.SUCCESS }
 				snapshotutils.NvmlDeviceGetCount = func() (int, nvml.Return) { return 1, nvml.SUCCESS }
 				device := &mockDevice{
-					computeProcs: []nvml.ProcessInfo{{Pid: 400}, {Pid: 500}},
+					computeProcs: []nvml.ProcessInfo{{Pid: 400, UsedGpuMemory: 1024}, {Pid: 500, UsedGpuMemory: 1024}},
 				}
 				snapshotutils.NvmlDeviceGetHandleByIndex = func(index int) (snapshotutils.DeviceInterface, nvml.Return) {
 					return device, nvml.SUCCESS
@@ -180,7 +213,7 @@ func TestGetPodPIDs(t *testing.T) {
 				snapshotutils.NvmlInit = func() nvml.Return { return nvml.SUCCESS }
 				snapshotutils.NvmlDeviceGetCount = func() (int, nvml.Return) { return 1, nvml.SUCCESS }
 				device := &mockDevice{
-					computeProcs: []nvml.ProcessInfo{{Pid: 600}},
+					computeProcs: []nvml.ProcessInfo{{Pid: 600, UsedGpuMemory: 1024}},
 				}
 				snapshotutils.NvmlDeviceGetHandleByIndex = func(index int) (snapshotutils.DeviceInterface, nvml.Return) {
 					return device, nvml.SUCCESS