feat(xid/sxid): allow custom reboot thresholds (#1244)

Add per-XID and per-SXID reboot threshold maps for gpud run.

```bash
gpud run \
  --xid-thresholds '{"overrides":{"94":{"rebootThreshold":1000}}}' \
  --sxid-thresholds '{"overrides":{"11004":{"rebootThreshold":7}}}'
```

Author: gyuho

cmd/gpud/command/command.go

@@ -253,6 +253,14 @@ sudo rm /etc/systemd/system/gpud.service
 					Usage: fmt.Sprintf("set the allowed reboot attempts for XID errors before escalation (defaults to %d)", componentsxid.DefaultRebootThreshold),
 					Value: componentsxid.DefaultRebootThreshold,
 				},
+				&cli.StringFlag{
+					Name:  "xid-thresholds",
+					Usage: `set per-XID thresholds in JSON, e.g. '{"overrides":{"94":{"rebootThreshold":1000}}}'`,
+				},
+				&cli.StringFlag{
+					Name:  "sxid-thresholds",
+					Usage: `set per-SXID thresholds in JSON, e.g. '{"overrides":{"11004":{"rebootThreshold":7}}}'`,
+				},
 				&cli.DurationFlag{
 					Name:  "xid-lookback-period",
 					Usage: "set the lookback period for XID errors",

cmd/gpud/run/command.go

@@ -142,6 +142,8 @@ func Command(cliContext *cli.Context) error {
 	nvlinkExpectedLinkStates := cliContext.String("nvlink-expected-link-states")
 	nfsCheckerConfigs := cliContext.String("nfs-checker-configs")
 	xidRebootThreshold := cliContext.Int("xid-reboot-threshold")
+	xidThresholds := cliContext.String("xid-thresholds")
+	sxidThresholds := cliContext.String("sxid-thresholds")
 	temperatureMarginThresholdCelsius := cliContext.Int("threshold-celsius-slowdown-margin")
 
 	if len(infinibandExpectedPortStates) > 0 {
@@ -174,16 +176,42 @@ func Command(cliContext *cli.Context) error {
 		log.Logger.Infow("set nfs checker group configs", "groupConfigs", groupConfigs)
 	}
 
+	xidThresholdsChanged := false
 	if cliContext.IsSet("xid-reboot-threshold") {
 		if xidRebootThreshold > 0 {
-			componentsxid.SetDefaultRebootThreshold(componentsxid.RebootThreshold{
-				Threshold: xidRebootThreshold,
-			})
-			log.Logger.Infow("set xid reboot threshold", "xidRebootThreshold", xidRebootThreshold)
+			componentsxid.SetDefaultRebootThreshold(xidRebootThreshold)
+			xidThresholdsChanged = true
 		} else {
 			log.Logger.Warnw("ignoring xid reboot threshold override, value must be positive", "xidRebootThreshold", xidRebootThreshold)
 		}
 	}
+	if strings.TrimSpace(xidThresholds) != "" {
+		thresholds, err := parseXIDThresholds(xidThresholds)
+		if err != nil {
+			return err
+		}
+		componentsxid.SetDefaultThresholds(thresholds)
+		xidThresholdsChanged = true
+	}
+	if xidThresholdsChanged {
+		xidThresholdConfig := componentsxid.GetDefaultThresholds()
+		log.Logger.Infow(
+			"set xid thresholds",
+			"xidDefaultRebootThreshold",
+			componentsxid.GetDefaultRebootThreshold(),
+			"xidOverrides",
+			xidThresholdConfig.Overrides,
+		)
+	}
+
+	if strings.TrimSpace(sxidThresholds) != "" {
+		thresholds, err := parseSXIDThresholds(sxidThresholds)
+		if err != nil {
+			return err
+		}
+		componentssxid.SetDefaultThresholds(thresholds)
+		log.Logger.Infow("set sxid thresholds", "sxidOverrides", thresholds.Overrides)
+	}
 
 	if eventsRetentionPeriod > 0 && !cliContext.IsSet("xid-lookback-period") {
 		componentsxid.SetLookbackPeriod(eventsRetentionPeriod)
@@ -438,6 +466,79 @@ func parseRetentionPeriods(cliContext *cli.Context) (metricsRetentionPeriod, eve
 	return metricsRetentionPeriod, eventsRetentionPeriod
 }
 
+type thresholdOverrideJSON struct {
+	RebootThreshold int `json:"rebootThreshold"`
+}
+
+type thresholdsJSON struct {
+	Overrides map[int]thresholdOverrideJSON `json:"overrides"`
+}
+
+func parseXIDThresholds(raw string) (componentsxid.Thresholds, error) {
+	thresholds, err := parseThresholds(raw, "xid thresholds")
+	if err != nil {
+		return componentsxid.Thresholds{}, err
+	}
+
+	ret := make(map[int]componentsxid.ThresholdOverride, len(thresholds))
+	for xid, threshold := range thresholds {
+		ret[xid] = componentsxid.ThresholdOverride{
+			RebootThreshold: threshold.RebootThreshold,
+		}
+	}
+	return componentsxid.Thresholds{Overrides: ret}, nil
+}
+
+func parseSXIDThresholds(raw string) (componentssxid.Thresholds, error) {
+	thresholds, err := parseThresholds(raw, "sxid thresholds")
+	if err != nil {
+		return componentssxid.Thresholds{}, err
+	}
+
+	ret := make(map[int]componentssxid.ThresholdOverride, len(thresholds))
+	for sxid, threshold := range thresholds {
+		ret[sxid] = componentssxid.ThresholdOverride{
+			RebootThreshold: threshold.RebootThreshold,
+		}
+	}
+	return componentssxid.Thresholds{Overrides: ret}, nil
+}
+
+func parseThresholds(raw string, name string) (map[int]thresholdOverrideJSON, error) {
+	raw = strings.TrimSpace(raw)
+	if raw == "" {
+		return nil, nil
+	}
+
+	var fields map[string]json.RawMessage
+	if err := json.Unmarshal([]byte(raw), &fields); err != nil {
+		return nil, fmt.Errorf("invalid %s: %w", name, err)
+	}
+	if fields == nil {
+		return nil, fmt.Errorf("invalid %s: expected JSON object", name)
+	}
+	for field := range fields {
+		if field != "overrides" {
+			return nil, fmt.Errorf("invalid %s: unknown field %q", name, field)
+		}
+	}
+
+	var thresholds thresholdsJSON
+	if err := json.Unmarshal([]byte(raw), &thresholds); err != nil {
+		return nil, fmt.Errorf("invalid %s: %w", name, err)
+	}
+
+	for id, threshold := range thresholds.Overrides {
+		if id < 0 {
+			return nil, fmt.Errorf("invalid %s for %d: event ID must be non-negative", name, id)
+		}
+		if threshold.RebootThreshold <= 0 {
+			return nil, fmt.Errorf("invalid %s for %d: rebootThreshold must be positive", name, id)
+		}
+	}
+	return thresholds.Overrides, nil
+}
+
 func parseInfinibandExcludeDevices(s string) []string {
 	if s == "" {
 		return nil

cmd/gpud/run/mock_command_exec_test.go

@@ -63,6 +63,8 @@ func newTestCLIContext(t *testing.T, values cliFlagValues) *cli.Context {
 	set.String("nvlink-expected-link-states", "", "")
 	set.String("nfs-checker-configs", "", "")
 	set.Int("xid-reboot-threshold", 0, "")
+	set.String("xid-thresholds", "", "")
+	set.String("sxid-thresholds", "", "")
 	set.Duration("xid-lookback-period", 0, "")
 	set.Duration("sxid-lookback-period", 0, "")
 	set.Int("threshold-celsius-slowdown-margin", 0, "")
@@ -275,6 +277,62 @@ func TestCommand_SetLookbackPeriods(t *testing.T) {
 	})
 }
 
+func TestCommand_SetThresholds(t *testing.T) {
+	tmpDir := t.TempDir()
+
+	originalXidRebootThreshold := componentsxid.GetDefaultRebootThreshold()
+	originalXidThresholds := componentsxid.GetDefaultThresholds()
+	originalSxidThresholds := componentssxid.GetDefaultThresholds()
+	t.Cleanup(func() {
+		componentsxid.SetDefaultRebootThreshold(originalXidRebootThreshold)
+		componentsxid.SetDefaultThresholds(originalXidThresholds)
+		componentssxid.SetDefaultThresholds(originalSxidThresholds)
+	})
+
+	ctx := newTestCLIContext(t, cliFlagValues{
+		stringFlags: map[string]string{
+			"log-level":       "info",
+			"data-dir":        tmpDir,
+			"xid-thresholds":  `{"overrides":{"94":{"rebootThreshold":2000},"95":{"rebootThreshold":3}}}`,
+			"sxid-thresholds": `{"overrides":{"11004":{"rebootThreshold":7}}}`,
+		},
+		intFlags: map[string]int{
+			"xid-reboot-threshold": 6,
+		},
+	})
+
+	mockey.PatchConvey("Command sets xid and sxid thresholds", t, func() {
+		mockey.Mock(gpudmanager.New).To(func(dataDir string) (*gpudmanager.Manager, error) {
+			return &gpudmanager.Manager{}, nil
+		}).Build()
+		mockey.Mock((*gpudmanager.Manager).Start).To(func(_ *gpudmanager.Manager, _ context.Context) error {
+			return nil
+		}).Build()
+		mockey.Mock(gpudserver.New).To(func(_ context.Context, _ log.AuditLogger, _ *config.Config, _ *gpudmanager.Manager) (*gpudserver.Server, error) {
+			return &gpudserver.Server{}, nil
+		}).Build()
+		mockey.Mock(gpudserver.HandleSignals).To(func(_ context.Context, _ context.CancelFunc, _ chan os.Signal, _ chan gpudserver.ServerStopper, _ func(context.Context) error) chan struct{} {
+			done := make(chan struct{})
+			close(done)
+			return done
+		}).Build()
+		mockey.Mock(pkgsystemd.SystemctlExists).To(func() bool {
+			return false
+		}).Build()
+
+		err := Command(ctx)
+		require.NoError(t, err)
+
+		xidThresholdConfig := componentsxid.GetDefaultThresholds()
+		assert.Equal(t, 6, componentsxid.GetDefaultRebootThreshold())
+		assert.Equal(t, 2000, xidThresholdConfig.Overrides[94].RebootThreshold)
+		assert.Equal(t, 3, xidThresholdConfig.Overrides[95].RebootThreshold)
+
+		sxidThresholds := componentssxid.GetDefaultThresholds()
+		assert.Equal(t, 7, sxidThresholds.Overrides[11004].RebootThreshold)
+	})
+}
+
 func TestCommand_UsesEventsRetentionForXidAndSxidLookback(t *testing.T) {
 	tmpDir := t.TempDir()
 

cmd/gpud/run/thresholds_test.go

@@ -0,0 +1,30 @@
+package run
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestParseThresholds(t *testing.T) {
+	xidThresholds, err := parseXIDThresholds(`{"overrides":{"94":{"rebootThreshold":1000}}}`)
+	require.NoError(t, err)
+	assert.Equal(t, 1000, xidThresholds.Overrides[94].RebootThreshold)
+
+	sxidThresholds, err := parseSXIDThresholds(`{"overrides":{"11004":{"rebootThreshold":7}}}`)
+	require.NoError(t, err)
+	assert.Equal(t, 7, sxidThresholds.Overrides[11004].RebootThreshold)
+
+	_, err = parseXIDThresholds(`{"overrides":{"94":{"rebootThreshold":0}}}`)
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "rebootThreshold must be positive")
+
+	_, err = parseXIDThresholds(`{"94":{"rebootThreshold":1000}}`)
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), `unknown field "94"`)
+
+	_, err = parseSXIDThresholds(`{not-valid-json}`)
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "sxid thresholds")
+}

cmd/gpud/scan/command.go

@@ -142,9 +142,7 @@ func cmdScan(
 
 	if xidRebootThresholdIsSet {
 		if xidRebootThreshold > 0 {
-			componentsxid.SetDefaultRebootThreshold(componentsxid.RebootThreshold{
-				Threshold: xidRebootThreshold,
-			})
+			componentsxid.SetDefaultRebootThreshold(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)

components/accelerator/nvidia/sxid/health_state.go

@@ -32,14 +32,17 @@ func translateToStateHealth(health int) apiv1.HealthStateType {
 	}
 }
 
-const rebootThreshold = 2
-
 // evolveHealthyState resolves the state of the SXID error component.
 // note: assume events are sorted by time in descending order
 func evolveHealthyState(events eventstore.Events) (ret apiv1.HealthState) {
+	return evolveHealthyStateWithThresholds(events, DefaultRebootThreshold, GetDefaultThresholds())
+}
+
+func evolveHealthyStateWithThresholds(events eventstore.Events, defaultRebootThreshold int, thresholds Thresholds) (ret apiv1.HealthState) {
 	defer func() {
 		log.Logger.Debugf("EvolveHealthyState: %v", ret)
 	}()
+	thresholds = normalizeThresholds(thresholds)
 	var lastSuggestedAction *apiv1.SuggestedActions
 	var lastSXidErr *sxidErrorEventDetail
 	lastHealth := healthStateHealthy
@@ -68,9 +71,10 @@ func evolveHealthyState(events eventstore.Events) (ret apiv1.HealthState) {
 			lastSXidErr = &currSXidErr
 			if currSXidErr.SuggestedActionsByGPUd != nil && len(currSXidErr.SuggestedActionsByGPUd.RepairActions) > 0 {
 				if currSXidErr.SuggestedActionsByGPUd.RepairActions[0] == apiv1.RepairActionTypeRebootSystem {
+					effectiveRebootThreshold := rebootThresholdForSXID(currSXidErr.SXid, defaultRebootThreshold, thresholds)
 					if count, ok := sxidRebootMap[currSXidErr.SXid]; !ok {
 						sxidRebootMap[currSXidErr.SXid] = 0
-					} else if count >= rebootThreshold {
+					} else if count >= effectiveRebootThreshold {
 						currSXidErr.SuggestedActionsByGPUd.RepairActions[0] = apiv1.RepairActionTypeHardwareInspection
 					}
 				}

components/accelerator/nvidia/sxid/health_state_test.go

@@ -82,6 +82,26 @@ func TestStateUpdateBasedOnEvents(t *testing.T) {
 		assert.Equal(t, apiv1.RepairActionTypeHardwareInspection, state.SuggestedActions.RepairActions[0])
 	})
 
+	t.Run("custom sxid reboot threshold override prevents hardware inspection", func(t *testing.T) {
+		events := eventstore.Events{
+			createSXidEvent(time.Time{}, 94, apiv1.EventTypeFatal, apiv1.RepairActionTypeRebootSystem),
+			{Name: "reboot"},
+			createSXidEvent(time.Time{}, 94, apiv1.EventTypeFatal, apiv1.RepairActionTypeRebootSystem),
+			{Name: "reboot"},
+			createSXidEvent(time.Time{}, 94, apiv1.EventTypeFatal, apiv1.RepairActionTypeRebootSystem),
+		}
+
+		state := evolveHealthyStateWithThresholds(events, DefaultRebootThreshold, Thresholds{
+			Overrides: map[int]ThresholdOverride{
+				94: {RebootThreshold: 1000},
+			},
+		})
+		assert.Equal(t, apiv1.HealthStateTypeUnhealthy, state.Health)
+		require.NotNil(t, state.SuggestedActions)
+		require.NotEmpty(t, state.SuggestedActions.RepairActions)
+		assert.Equal(t, apiv1.RepairActionTypeRebootSystem, state.SuggestedActions.RepairActions[0])
+	})
+
 	t.Run("EmptyEvents_ReturnsHealthy", func(t *testing.T) {
 		events := eventstore.Events{}
 		state := evolveHealthyState(events)