From e938ae36a4f058629f4f903680887d26c40373d0 Mon Sep 17 00:00:00 2001 From: Jingxiang Zhang Date: Wed, 24 Jun 2026 10:21:32 -0700 Subject: [PATCH] fix: degrade dcgm query timeouts Signed-off-by: Jingxiang Zhang --- .../pkg/nvidia-query/dcgm/query_validator.go | 4 +-- .../nvidia-query/dcgm/query_validator_test.go | 28 +++++++++++++++++++ 2 files changed, 29 insertions(+), 3 deletions(-) diff --git a/third_party/fleet-intelligence-sdk/pkg/nvidia-query/dcgm/query_validator.go b/third_party/fleet-intelligence-sdk/pkg/nvidia-query/dcgm/query_validator.go index 7463946d..7626fc9a 100644 --- a/third_party/fleet-intelligence-sdk/pkg/nvidia-query/dcgm/query_validator.go +++ b/third_party/fleet-intelligence-sdk/pkg/nvidia-query/dcgm/query_validator.go @@ -210,9 +210,7 @@ func IsUnhealthyAPIError(err error) bool { case dcgm.DCGM_ST_NVML_ERROR, dcgm.DCGM_ST_GPU_IS_LOST, dcgm.DCGM_ST_RESET_REQUIRED, - dcgm.DCGM_ST_GPU_NOT_SUPPORTED, - dcgm.DCGM_ST_TIMEOUT, - dcgm.DCGM_ST_NVML_DRIVER_TIMEOUT: + dcgm.DCGM_ST_GPU_NOT_SUPPORTED: return true default: return false diff --git a/third_party/fleet-intelligence-sdk/pkg/nvidia-query/dcgm/query_validator_test.go b/third_party/fleet-intelligence-sdk/pkg/nvidia-query/dcgm/query_validator_test.go index e07b2089..c167a3cb 100644 --- a/third_party/fleet-intelligence-sdk/pkg/nvidia-query/dcgm/query_validator_test.go +++ b/third_party/fleet-intelligence-sdk/pkg/nvidia-query/dcgm/query_validator_test.go @@ -16,7 +16,10 @@ package dcgm import ( + "fmt" "testing" + + dcgm "github.com/NVIDIA/go-dcgm/pkg/dcgm" ) // Note: Tests for CheckSentinel() and CheckSentinelV2() require actual DCGM library @@ -66,3 +69,28 @@ func TestSentinelType_String(t *testing.T) { }) } } + +func TestIsUnhealthyAPIError(t *testing.T) { + tests := []struct { + name string + code int32 + want bool + }{ + {"nvml error is unhealthy", dcgm.DCGM_ST_NVML_ERROR, true}, + {"gpu lost is unhealthy", dcgm.DCGM_ST_GPU_IS_LOST, true}, + {"reset required is unhealthy", dcgm.DCGM_ST_RESET_REQUIRED, true}, + {"gpu not supported is unhealthy", dcgm.DCGM_ST_GPU_NOT_SUPPORTED, true}, + {"dcgm timeout is degraded", dcgm.DCGM_ST_TIMEOUT, false}, + {"nvml driver timeout is degraded", dcgm.DCGM_ST_NVML_DRIVER_TIMEOUT, false}, + {"stale data is not unhealthy", dcgm.DCGM_ST_STALE_DATA, false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := fmt.Errorf("dcgm request failed with error code %d", tt.code) + if got := IsUnhealthyAPIError(err); got != tt.want { + t.Errorf("IsUnhealthyAPIError() = %v, want %v", got, tt.want) + } + }) + } +}