feat(nvidia/xid): allow to set reboot threshold (default remains 2)

cmd/gpud/command/command.go

@@ -22,6 +22,7 @@ import (
 	cmdstatus "github.com/leptonai/gpud/cmd/gpud/status"
 	cmdup "github.com/leptonai/gpud/cmd/gpud/up"
 	cmdupdate "github.com/leptonai/gpud/cmd/gpud/update"
+	componentsxid "github.com/leptonai/gpud/components/accelerator/nvidia/xid"
 	pkgconfig "github.com/leptonai/gpud/pkg/config"
 	pkgcustomplugins "github.com/leptonai/gpud/pkg/custom-plugins"
 	pkgupdate "github.com/leptonai/gpud/pkg/update"
@@ -189,6 +190,11 @@ sudo rm /etc/systemd/system/gpud.service
 					Name:  "nfs-checker-configs",
 					Usage: "set the NFS checker group configs in JSON (leave empty for default, useful for testing)",
 				},
+				&cli.IntFlag{
+					Name:  "xid-reboot-threshold",
+					Usage: fmt.Sprintf("set the allowed reboot attempts for XID errors before escalation (defaults to %d)", componentsxid.DefaultRebootThreshold),
+					Value: componentsxid.DefaultRebootThreshold,
+				},
 
 				cli.StringFlag{
 					Name:   "infiniband-class-root-dir",
@@ -461,6 +467,11 @@ sudo rm /etc/systemd/system/gpud.service
 					Name:  "nfs-checker-configs",
 					Usage: "set the NFS checker group configs in JSON (leave empty for default, useful for testing)",
 				},
+				&cli.IntFlag{
+					Name:  "xid-reboot-threshold",
+					Usage: fmt.Sprintf("set the allowed reboot attempts for XID errors before escalation (defaults to %d)", componentsxid.DefaultRebootThreshold),
+					Value: componentsxid.DefaultRebootThreshold,
+				},
 				cli.StringFlag{
 					Name:   "infiniband-class-root-dir",
 					Usage:  "sets the infiniband class root directory (leave empty for default)",

cmd/gpud/run/command.go

@@ -20,6 +20,7 @@ import (
 	gpudcomponents "github.com/leptonai/gpud/components"
 	componentsnvidiagpucounts "github.com/leptonai/gpud/components/accelerator/nvidia/gpu-counts"
 	componentsinfiniband "github.com/leptonai/gpud/components/accelerator/nvidia/infiniband"
+	componentsxid "github.com/leptonai/gpud/components/accelerator/nvidia/xid"
 	componentsnfs "github.com/leptonai/gpud/components/nfs"
 	"github.com/leptonai/gpud/pkg/config"
 	gpudmanager "github.com/leptonai/gpud/pkg/gpud-manager"
@@ -67,6 +68,7 @@ func Command(cliContext *cli.Context) error {
 	gpuCount := cliContext.Int("gpu-count")
 	infinibandExpectedPortStates := cliContext.String("infiniband-expected-port-states")
 	nfsCheckerConfigs := cliContext.String("nfs-checker-configs")
+	xidRebootThreshold := cliContext.Int("xid-reboot-threshold")
 
 	if gpuCount > 0 {
 		componentsnvidiagpucounts.SetDefaultExpectedGPUCounts(componentsnvidiagpucounts.ExpectedGPUCounts{
@@ -96,6 +98,17 @@ func Command(cliContext *cli.Context) error {
 		log.Logger.Infow("set nfs checker group configs", "groupConfigs", groupConfigs)
 	}
 
+	if cliContext.IsSet("xid-reboot-threshold") {
+		if xidRebootThreshold > 0 {
+			componentsxid.SetDefaultRebootThreshold(componentsxid.RebootThreshold{
+				Threshold: xidRebootThreshold,
+			})
+			log.Logger.Infow("set xid reboot threshold", "xidRebootThreshold", xidRebootThreshold)
+		} else {
+			log.Logger.Warnw("ignoring xid reboot threshold override, value must be positive", "xidRebootThreshold", xidRebootThreshold)
+		}
+	}
+
 	gpuUUIDsWithRowRemappingPendingRaw := cliContext.String("gpu-uuids-with-row-remapping-pending")
 	gpuUUIDsWithRowRemappingPending := common.ParseGPUUUIDs(gpuUUIDsWithRowRemappingPendingRaw)
 

cmd/gpud/scan/command.go

@@ -12,6 +12,7 @@ import (
 	gpudcomponents "github.com/leptonai/gpud/components"
 	componentsnvidiagpucounts "github.com/leptonai/gpud/components/accelerator/nvidia/gpu-counts"
 	componentsinfiniband "github.com/leptonai/gpud/components/accelerator/nvidia/infiniband"
+	componentsxid "github.com/leptonai/gpud/components/accelerator/nvidia/xid"
 	componentsnfs "github.com/leptonai/gpud/components/nfs"
 	"github.com/leptonai/gpud/pkg/log"
 	pkgnfschecker "github.com/leptonai/gpud/pkg/nfs-checker"
@@ -32,6 +33,8 @@ func CreateCommand() func(*cli.Context) error {
 			cliContext.String("gpu-uuids-with-hw-slowdown"),
 			cliContext.String("gpu-uuids-with-hw-slowdown-thermal"),
 			cliContext.String("gpu-uuids-with-hw-slowdown-power-brake"),
+			cliContext.Int("xid-reboot-threshold"),
+			cliContext.IsSet("xid-reboot-threshold"),
 		)
 	}
 }
@@ -47,6 +50,8 @@ func cmdScan(
 	gpuUUIDsWithHWSlowdownRaw string,
 	gpuUUIDsWithHWSlowdownThermalRaw string,
 	gpuUUIDsWithHWSlowdownPowerBrakeRaw string,
+	xidRebootThreshold int,
+	xidRebootThresholdIsSet bool,
 ) error {
 	zapLvl, err := log.ParseLogLevel(logLevel)
 	if err != nil {
@@ -84,6 +89,17 @@ func cmdScan(
 		log.Logger.Infow("set nfs checker group configs", "groupConfigs", groupConfigs)
 	}
 
+	if xidRebootThresholdIsSet {
+		if xidRebootThreshold > 0 {
+			componentsxid.SetDefaultRebootThreshold(componentsxid.RebootThreshold{
+				Threshold: xidRebootThreshold,
+			})
+			log.Logger.Infow("set xid reboot threshold", "xidRebootThreshold", xidRebootThreshold)
+		} else {
+			log.Logger.Warnw("ignoring xid reboot threshold override, value must be positive", "xidRebootThreshold", xidRebootThreshold)
+		}
+	}
+
 	gpuUUIDsWithRowRemappingPending := common.ParseGPUUUIDs(gpuUUIDsWithRowRemappingPendingRaw)
 	gpuUUIDsWithRowRemappingFailed := common.ParseGPUUUIDs(gpuUUIDsWithRowRemappingFailedRaw)
 	gpuUUIDsWithHWSlowdown := common.ParseGPUUUIDs(gpuUUIDsWithHWSlowdownRaw)

components/accelerator/nvidia/xid/component.go

@@ -52,7 +52,8 @@ type component struct {
 	nvmlInstance nvidianvml.Instance
 	devices      map[string]device.Device
 
-	getTimeNowFunc func() time.Time
+	getTimeNowFunc   func() time.Time
+	getThresholdFunc func() RebootThreshold
 
 	rebootEventStore pkghost.RebootEventStore
 	eventBucket      eventstore.Bucket
@@ -78,6 +79,7 @@ func New(gpudInstance *components.GPUdInstance) (components.Component, error) {
 		getTimeNowFunc: func() time.Time {
 			return time.Now().UTC()
 		},
+		getThresholdFunc: GetDefaultRebootThreshold,
 
 		rebootEventStore: gpudInstance.RebootEventStore,
 		extraEventCh:     make(chan *eventstore.Event, 256),
@@ -501,6 +503,7 @@ func (c *component) updateCurrentState() error {
 	}
 
 	now := c.getTimeNowFunc()
+	rebootThreshold := c.getThresholdFunc()
 
 	var rebootErr string
 	rebootEvents, err := c.rebootEventStore.GetRebootEvents(c.ctx, now.Add(-DefaultRetentionPeriod))
@@ -517,7 +520,7 @@ func (c *component) updateCurrentState() error {
 	events := mergeEvents(rebootEvents, localEvents)
 
 	c.mu.Lock()
-	c.currState = evolveHealthyState(events, c.devices)
+	c.currState = evolveHealthyState(events, c.devices, rebootThreshold.Threshold)
 	if rebootErr != "" {
 		c.currState.Error = fmt.Sprintf("%s\n%s", rebootErr, c.currState.Error)
 	}

components/accelerator/nvidia/xid/health_state.go

@@ -36,11 +36,9 @@ func translateToStateHealth(health int) apiv1.HealthStateType {
 	}
 }
 
-const rebootThreshold = 2
-
 // evolveHealthyState resolves the state of the XID error component.
 // note: assume events are sorted by time in descending order
-func evolveHealthyState(events eventstore.Events, devices map[string]device.Device) (ret apiv1.HealthState) {
+func evolveHealthyState(events eventstore.Events, devices map[string]device.Device, rebootThreshold int) (ret apiv1.HealthState) {
 	defer func() {
 		log.Logger.Debugf("EvolveHealthyState: %v", ret)
 	}()

components/accelerator/nvidia/xid/health_state_test.go

@@ -38,7 +38,7 @@ func createXidEvent(eventTime time.Time, xid uint64, eventType apiv1.EventType,
 
 func TestStateUpdateBasedOnEvents(t *testing.T) {
 	t.Run("no event found", func(t *testing.T) {
-		state := evolveHealthyState(eventstore.Events{}, nil)
+		state := evolveHealthyState(eventstore.Events{}, nil, DefaultRebootThreshold)
 		assert.Equal(t, apiv1.HealthStateTypeHealthy, state.Health)
 		assert.Equal(t, "XIDComponent is healthy", state.Reason)
 	})
@@ -49,7 +49,7 @@ func TestStateUpdateBasedOnEvents(t *testing.T) {
 		events := eventstore.Events{
 			createXidEvent(time.Time{}, 123, apiv1.EventTypeCritical, apiv1.RepairActionTypeRebootSystem),
 		}
-		state := evolveHealthyState(events, map[string]device.Device{"GPU-b850f46d-d5ea-c752-ddf3-c4453e44d3f7": mockDevice})
+		state := evolveHealthyState(events, map[string]device.Device{"GPU-b850f46d-d5ea-c752-ddf3-c4453e44d3f7": mockDevice}, DefaultRebootThreshold)
 		assert.Equal(t, apiv1.HealthStateTypeDegraded, state.Health)
 		assert.Equal(t, "XID 123 (SPI PMU RPC Write Failure) detected on GPU PCI:0000:9b:00 UUID:GPU-b850f46d-d5ea-c752-ddf3-c4453e44d3f7", state.Reason)
 	})
@@ -58,7 +58,7 @@ func TestStateUpdateBasedOnEvents(t *testing.T) {
 		events := eventstore.Events{
 			createXidEvent(time.Time{}, 456, apiv1.EventTypeFatal, apiv1.RepairActionTypeRebootSystem),
 		}
-		state := evolveHealthyState(events, map[string]device.Device{"GPU-b850f46d-d5ea-c752-ddf3-c4453e44d3f7": mockDevice})
+		state := evolveHealthyState(events, map[string]device.Device{"GPU-b850f46d-d5ea-c752-ddf3-c4453e44d3f7": mockDevice}, DefaultRebootThreshold)
 		assert.Equal(t, apiv1.HealthStateTypeUnhealthy, state.Health)
 		assert.Equal(t, "XID 456 detected on GPU PCI:0000:9b:00 UUID:GPU-b850f46d-d5ea-c752-ddf3-c4453e44d3f7", state.Reason)
 	})
@@ -68,7 +68,7 @@ func TestStateUpdateBasedOnEvents(t *testing.T) {
 			{Name: "reboot"},
 			createXidEvent(time.Time{}, 789, apiv1.EventTypeCritical, apiv1.RepairActionTypeRebootSystem),
 		}
-		state := evolveHealthyState(events, nil)
+		state := evolveHealthyState(events, nil, DefaultRebootThreshold)
 		assert.Equal(t, apiv1.HealthStateTypeHealthy, state.Health)
 	})
 
@@ -81,7 +81,7 @@ func TestStateUpdateBasedOnEvents(t *testing.T) {
 			createXidEvent(time.Time{}, 94, apiv1.EventTypeCritical, apiv1.RepairActionTypeRebootSystem),
 			createXidEvent(time.Time{}, 31, apiv1.EventTypeWarning, apiv1.RepairActionTypeCheckUserAppAndGPU),
 		}
-		state := evolveHealthyState(events, nil)
+		state := evolveHealthyState(events, nil, DefaultRebootThreshold)
 		assert.Equal(t, apiv1.HealthStateTypeDegraded, state.Health)
 		assert.Equal(t, apiv1.RepairActionTypeHardwareInspection, state.SuggestedActions.RepairActions[0])
 	})
@@ -91,7 +91,7 @@ func TestStateUpdateBasedOnEvents(t *testing.T) {
 			{Name: "SetHealthy"},
 			createXidEvent(time.Time{}, 789, apiv1.EventTypeFatal, apiv1.RepairActionTypeRebootSystem),
 		}
-		state := evolveHealthyState(events, nil)
+		state := evolveHealthyState(events, nil, DefaultRebootThreshold)
 		assert.Equal(t, apiv1.HealthStateTypeHealthy, state.Health)
 		assert.Nil(t, state.SuggestedActions)
 	})
@@ -104,7 +104,7 @@ func TestStateUpdateBasedOnEvents(t *testing.T) {
 				ExtraInfo: map[string]string{EventKeyErrorXidData: "invalid json"},
 			},
 		}
-		state := evolveHealthyState(events, nil)
+		state := evolveHealthyState(events, nil, DefaultRebootThreshold)
 		assert.Equal(t, apiv1.HealthStateTypeHealthy, state.Health)
 	})
 }

components/accelerator/nvidia/xid/threshold.go

@@ -0,0 +1,40 @@
+package xid
+
+import (
+	"sync"
+
+	"github.com/leptonai/gpud/pkg/log"
+)
+
+// RebootThreshold configures the expected reboot threshold.
+type RebootThreshold struct {
+	// Threshold is the expected number of reboot events within the evaluation window.
+	// If not set, it defaults to 2.
+	Threshold int `json:"threshold"`
+}
+
+var (
+	defaultRebootThresholdMu sync.RWMutex
+	defaultRebootThreshold   = RebootThreshold{
+		Threshold: DefaultRebootThreshold,
+	}
+)
+
+const (
+	// DefaultRebootThreshold is the default reboot threshold.
+	DefaultRebootThreshold = 2
+)
+
+func GetDefaultRebootThreshold() RebootThreshold {
+	defaultRebootThresholdMu.RLock()
+	defer defaultRebootThresholdMu.RUnlock()
+	return defaultRebootThreshold
+}
+
+func SetDefaultRebootThreshold(threshold RebootThreshold) {
+	log.Logger.Infow("setting default reboot threshold", "threshold", threshold.Threshold)
+
+	defaultRebootThresholdMu.Lock()
+	defer defaultRebootThresholdMu.Unlock()
+	defaultRebootThreshold = threshold
+}

components/accelerator/nvidia/xid/threshold_test.go

@@ -0,0 +1,22 @@
+package xid
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestDefaultExpectedPortStates(t *testing.T) {
+	// Test default values
+	defaultRebootThreshold := GetDefaultRebootThreshold()
+	assert.Equal(t, DefaultRebootThreshold, defaultRebootThreshold.Threshold)
+
+	// Test setting new values
+	newRebootThreshold := RebootThreshold{
+		Threshold: 4,
+	}
+	SetDefaultRebootThreshold(newRebootThreshold)
+
+	updatedRebootThreshold := GetDefaultRebootThreshold()
+	assert.Equal(t, newRebootThreshold.Threshold, updatedRebootThreshold.Threshold)
+}

pkg/session/serve_update_config.go

@@ -7,6 +7,7 @@ import (
 
 	componentsnvidiagpucounts "github.com/leptonai/gpud/components/accelerator/nvidia/gpu-counts"
 	componentsnvidiainfiniband "github.com/leptonai/gpud/components/accelerator/nvidia/infiniband"
+	componentsxid "github.com/leptonai/gpud/components/accelerator/nvidia/xid"
 	componentsnfs "github.com/leptonai/gpud/components/nfs"
 	"github.com/leptonai/gpud/pkg/log"
 	pkgnfschecker "github.com/leptonai/gpud/pkg/nfs-checker"
@@ -47,6 +48,18 @@ func (s *Session) processUpdateConfig(configMap map[string]string, resp *Respons
 				s.setDefaultGPUCountsFunc(updateCfg)
 			}
 
+		case componentsxid.Name:
+			setComponents[componentName] = struct{}{}
+			var updateCfg componentsxid.RebootThreshold
+			if err := json.Unmarshal([]byte(value), &updateCfg); err != nil {
+				log.Logger.Warnw("failed to unmarshal xid config", "error", err)
+				resp.Error = err.Error()
+				return
+			}
+			if s.setDefaultXIDRebootThresholdFunc != nil {
+				s.setDefaultXIDRebootThresholdFunc(updateCfg)
+			}
+
 		case componentsnfs.Name:
 			setComponents[componentName] = struct{}{}
 			var updateCfgs pkgnfschecker.Configs
@@ -76,7 +89,7 @@ func (s *Session) processUpdateConfig(configMap map[string]string, resp *Respons
 		}
 	}
 
-	// fallback to default empty if the component is not set
+	// fallback to default if the component is not set
 	if _, ok := setComponents[componentsnvidiainfiniband.Name]; !ok && s.setDefaultIbExpectedPortStatesFunc != nil {
 		log.Logger.Infow("falling back to default empty infiniband config")
 		s.setDefaultIbExpectedPortStatesFunc(infiniband.ExpectedPortStates{})
@@ -89,4 +102,8 @@ func (s *Session) processUpdateConfig(configMap map[string]string, resp *Respons
 		log.Logger.Infow("falling back to default empty nfs config")
 		s.setDefaultNFSGroupConfigsFunc(pkgnfschecker.Configs{})
 	}
+	if _, ok := setComponents[componentsxid.Name]; !ok && s.setDefaultXIDRebootThresholdFunc != nil {
+		log.Logger.Infow("falling back to default xid config")
+		s.setDefaultXIDRebootThresholdFunc(componentsxid.RebootThreshold{Threshold: componentsxid.DefaultRebootThreshold})
+	}
 }