feat: set nodegroup and computezone on agent

cmd/fleetint/enroll.go

@@ -30,7 +30,7 @@ import (
 )
 
 var (
-	performEnrollWorkflow = enrollment.EnrollWithConfig
+	performEnrollWorkflow = enrollment.EnrollWithConfigAndMetadata
 	fleetintEnvFilePath   = config.DefaultEnvFilePath
 )
 
@@ -84,6 +84,10 @@ func resolveToken(cliContext *cli.Context) (string, error) {
 func enrollCommand(cliContext *cli.Context) error {
 	baseEndpoint := cliContext.String("endpoint")
 	force := cliContext.Bool("force")
+	metadata := &enrollment.EnrollMetadata{
+		NodeGroup:   optionalFlagValue(cliContext, "node-group"),
+		ComputeZone: optionalFlagValue(cliContext, "compute-zone"),
+	}
 
 	sakToken, err := resolveToken(cliContext)
 	if err != nil {
@@ -117,5 +121,13 @@ func enrollCommand(cliContext *cli.Context) error {
 		return fmt.Errorf("failed to configure loop settings from environment variables: %w", err)
 	}
 
-	return performEnrollWorkflow(ctx, baseEndpoint, sakToken, cfg)
+	return performEnrollWorkflow(ctx, baseEndpoint, sakToken, cfg, metadata)
+}
+
+func optionalFlagValue(cliContext *cli.Context, name string) *string {
+	if !cliContext.IsSet(name) {
+		return nil
+	}
+	value := strings.TrimSpace(cliContext.String(name))
+	return &value
 }

cmd/fleetint/enroll_test.go

@@ -28,6 +28,7 @@ import (
 	"github.com/stretchr/testify/require"
 
 	"github.com/NVIDIA/fleet-intelligence-agent/internal/config"
+	"github.com/NVIDIA/fleet-intelligence-agent/internal/enrollment"
 	"github.com/NVIDIA/fleet-intelligence-agent/internal/precheck"
 )
 
@@ -107,8 +108,11 @@ func TestEnrollCommandBlocksOnFailedPrecheck(t *testing.T) {
 			},
 		}, nil
 	}
-	performEnrollWorkflow = func(ctx context.Context, baseEndpoint, sakToken string, cfg *config.Config) error {
+	performEnrollWorkflow = func(ctx context.Context, baseEndpoint, sakToken string, cfg *config.Config, metadata *enrollment.EnrollMetadata) error {
 		enrollmentCalled = true
+		require.NotNil(t, metadata)
+		require.Nil(t, metadata.NodeGroup)
+		require.Nil(t, metadata.ComputeZone)
 		return nil
 	}
 
@@ -142,8 +146,11 @@ func TestEnrollCommandForceBypassesFailedPrecheck(t *testing.T) {
 			},
 		}, nil
 	}
-	performEnrollWorkflow = func(ctx context.Context, baseEndpoint, sakToken string, cfg *config.Config) error {
+	performEnrollWorkflow = func(ctx context.Context, baseEndpoint, sakToken string, cfg *config.Config, metadata *enrollment.EnrollMetadata) error {
 		enrollmentCalled = true
+		require.NotNil(t, metadata)
+		require.Nil(t, metadata.NodeGroup)
+		require.Nil(t, metadata.ComputeZone)
 		return nil
 	}
 
@@ -174,11 +181,14 @@ func TestEnrollCommandPassesTimeoutContext(t *testing.T) {
 		}, nil
 	}
 
-	performEnrollWorkflow = func(ctx context.Context, baseEndpoint, sakToken string, cfg *config.Config) error {
+	performEnrollWorkflow = func(ctx context.Context, baseEndpoint, sakToken string, cfg *config.Config, metadata *enrollment.EnrollMetadata) error {
 		deadline, ok := ctx.Deadline()
 		require.True(t, ok)
 		require.LessOrEqual(t, time.Until(deadline), defaultEnrollTimeout)
 		require.Greater(t, time.Until(deadline), 55*time.Second)
+		require.NotNil(t, metadata)
+		require.Nil(t, metadata.NodeGroup)
+		require.Nil(t, metadata.ComputeZone)
 		return nil
 	}
 
@@ -217,14 +227,17 @@ FLEETINT_ATTESTATION_INTERVAL="6h"
 `), 0o600))
 	fleetintEnvFilePath = envFilePath
 
-	performEnrollWorkflow = func(ctx context.Context, baseEndpoint, sakToken string, cfg *config.Config) error {
+	performEnrollWorkflow = func(ctx context.Context, baseEndpoint, sakToken string, cfg *config.Config, metadata *enrollment.EnrollMetadata) error {
 		require.NotNil(t, cfg)
 		require.NotNil(t, cfg.Inventory)
 		require.False(t, cfg.Inventory.Enabled)
 		require.Equal(t, 15*time.Minute, cfg.Inventory.Interval.Duration)
 		require.NotNil(t, cfg.Attestation)
 		require.True(t, cfg.Attestation.Enabled)
 		require.Equal(t, 6*time.Hour, cfg.Attestation.Interval.Duration)
+		require.NotNil(t, metadata)
+		require.Nil(t, metadata.NodeGroup)
+		require.Nil(t, metadata.ComputeZone)
 		return nil
 	}
 
@@ -234,3 +247,83 @@ FLEETINT_ATTESTATION_INTERVAL="6h"
 	err := app.Run([]string{"fleetint", "enroll", "--endpoint", "https://example.com", "--token", "token"})
 	require.NoError(t, err)
 }
+
+func TestEnrollCommandPassesOptionalMetadata(t *testing.T) {
+	useMissingFleetintEnvFile(t)
+
+	originalRunPrecheck := runPrecheck
+	originalEnrollWorkflow := performEnrollWorkflow
+	t.Cleanup(func() {
+		runPrecheck = originalRunPrecheck
+		performEnrollWorkflow = originalEnrollWorkflow
+	})
+
+	runPrecheck = func() (precheck.Result, error) {
+		return precheck.Result{
+			Checks: []precheck.Check{
+				{Name: "gpu-present", Message: "ok", Passed: true},
+			},
+		}, nil
+	}
+
+	performEnrollWorkflow = func(ctx context.Context, baseEndpoint, sakToken string, cfg *config.Config, metadata *enrollment.EnrollMetadata) error {
+		require.NotNil(t, metadata)
+		require.NotNil(t, metadata.NodeGroup)
+		require.Equal(t, "prod-group", *metadata.NodeGroup)
+		require.NotNil(t, metadata.ComputeZone)
+		require.Equal(t, "us-east-1c", *metadata.ComputeZone)
+		return nil
+	}
+
+	app := App()
+	app.Writer = &bytes.Buffer{}
+
+	err := app.Run([]string{
+		"fleetint", "enroll",
+		"--endpoint", "https://example.com",
+		"--token", "token",
+		"--node-group", "prod-group",
+		"--compute-zone", "us-east-1c",
+	})
+	require.NoError(t, err)
+}
+
+func TestEnrollCommandTreatsExplicitEmptyMetadataAsClear(t *testing.T) {
+	useMissingFleetintEnvFile(t)
+
+	originalRunPrecheck := runPrecheck
+	originalEnrollWorkflow := performEnrollWorkflow
+	t.Cleanup(func() {
+		runPrecheck = originalRunPrecheck
+		performEnrollWorkflow = originalEnrollWorkflow
+	})
+
+	runPrecheck = func() (precheck.Result, error) {
+		return precheck.Result{
+			Checks: []precheck.Check{
+				{Name: "gpu-present", Message: "ok", Passed: true},
+			},
+		}, nil
+	}
+
+	performEnrollWorkflow = func(ctx context.Context, baseEndpoint, sakToken string, cfg *config.Config, metadata *enrollment.EnrollMetadata) error {
+		require.NotNil(t, metadata)
+		require.NotNil(t, metadata.NodeGroup)
+		require.Empty(t, *metadata.NodeGroup)
+		require.NotNil(t, metadata.ComputeZone)
+		require.Empty(t, *metadata.ComputeZone)
+		return nil
+	}
+
+	app := App()
+	app.Writer = &bytes.Buffer{}
+
+	err := app.Run([]string{
+		"fleetint", "enroll",
+		"--endpoint", "https://example.com",
+		"--token", "token",
+		"--node-group=",
+		"--compute-zone=",
+	})
+	require.NoError(t, err)
+}

cmd/fleetint/root.go

@@ -214,6 +214,14 @@ func App() *cli.App {
 					Name:  "force",
 					Usage: "continue enrollment even when precheck fails",
 				},
+				&cli.StringFlag{
+					Name:  "node-group",
+					Usage: "optional node group metadata associated with this node",
+				},
+				&cli.StringFlag{
+					Name:  "compute-zone",
+					Usage: "optional compute zone metadata associated with this node",
+				},
 			},
 		},
 		{

cmd/fleetint/unenroll.go

@@ -71,6 +71,8 @@ func removeEnrollmentMetadata(ctx context.Context, dbRW *sql.DB) error {
 		agentstate.MetadataKeySAKToken,
 		agentstate.MetadataKeyBackendBaseURL,
 		agentstate.MetadataKeyEnrolledAt,
+		agentstate.MetadataKeyNodeGroup,
+		agentstate.MetadataKeyComputeZone,
 		"enroll_endpoint",
 		"metrics_endpoint",
 		"logs_endpoint",

cmd/fleetint/unenroll_test.go

@@ -23,6 +23,8 @@ import (
 	pkgmetadata "github.com/NVIDIA/fleet-intelligence-sdk/pkg/metadata"
 	"github.com/NVIDIA/fleet-intelligence-sdk/pkg/sqlite"
 	"github.com/stretchr/testify/require"
+
+	"github.com/NVIDIA/fleet-intelligence-agent/internal/agentstate"
 )
 
 func TestRemoveEnrollmentMetadata(t *testing.T) {
@@ -36,14 +38,16 @@ func TestRemoveEnrollmentMetadata(t *testing.T) {
 	require.NoError(t, pkgmetadata.CreateTableMetadata(ctx, db))
 
 	for key, value := range map[string]string{
-		pkgmetadata.MetadataKeyToken: "jwt-token",
-		"sak_token":                  "sak-token",
-		"backend_base_url":           "https://backend.example.com",
-		"enroll_endpoint":            "https://backend.example.com/api/v1/enroll",
-		"metrics_endpoint":           "https://backend.example.com/api/v1/health/metrics",
-		"logs_endpoint":              "https://backend.example.com/api/v1/health/logs",
-		"nonce_endpoint":             "https://backend.example.com/api/v1/attest/nonce",
-		"keep_me":                    "still-here",
+		pkgmetadata.MetadataKeyToken:         "jwt-token",
+		agentstate.MetadataKeySAKToken:       "sak-token",
+		agentstate.MetadataKeyBackendBaseURL: "https://backend.example.com",
+		agentstate.MetadataKeyNodeGroup:      "group-a",
+		agentstate.MetadataKeyComputeZone:    "zone-a",
+		"enroll_endpoint":                    "https://backend.example.com/api/v1/enroll",
+		"metrics_endpoint":                   "https://backend.example.com/api/v1/health/metrics",
+		"logs_endpoint":                      "https://backend.example.com/api/v1/health/logs",
+		"nonce_endpoint":                     "https://backend.example.com/api/v1/attest/nonce",
+		"keep_me":                            "still-here",
 	} {
 		require.NoError(t, pkgmetadata.SetMetadata(ctx, db, key, value))
 	}
@@ -52,8 +56,10 @@ func TestRemoveEnrollmentMetadata(t *testing.T) {
 
 	for _, key := range []string{
 		pkgmetadata.MetadataKeyToken,
-		"sak_token",
-		"backend_base_url",
+		agentstate.MetadataKeySAKToken,
+		agentstate.MetadataKeyBackendBaseURL,
+		agentstate.MetadataKeyNodeGroup,
+		agentstate.MetadataKeyComputeZone,
 		"enroll_endpoint",
 		"metrics_endpoint",
 		"logs_endpoint",

docs/usage.md

@@ -128,12 +128,19 @@ One of `--token` or `--token-file` is required.
 
 **Optional Flags:**
 - `--force`: Continue enrollment even if `fleetint precheck` fails
+- `--node-group`: Optional node group metadata persisted in local agent metadata
+- `--compute-zone`: Optional compute zone metadata persisted in local agent metadata
+
+Metadata update behavior for `--node-group` and `--compute-zone`:
+- If the flag is omitted, the existing stored value is preserved.
+- If the flag is provided, the stored value is overwritten with the provided value.
+- Providing an empty value (for example `--node-group=""`) clears the stored value.
 
 **What it does:**
 1. Runs the same prerequisite validation as `fleetint precheck`
 2. Validates the endpoint URL (must be HTTPS)
 3. Makes an enrollment request to exchange the SAK token for a JWT token
-4. Stores the JWT token and backend endpoints (metrics, logs, nonce) in the local metadata database
+4. Stores the JWT token, backend endpoints (metrics, logs, nonce), and optional enrollment metadata (`node_group`, `compute_zone`) in the local metadata database
 5. The stored credentials are used automatically by the agent for data export
 
 **Example output:**

go.mod

@@ -99,12 +99,12 @@ require (
 	go4.org/mem v0.0.0-20220726221520-4f986261bf13 // indirect
 	go4.org/netipx v0.0.0-20231129151722-fdeea329fbba // indirect
 	golang.org/x/arch v0.22.0 // indirect
-	golang.org/x/crypto v0.50.0 // indirect
+	golang.org/x/crypto v0.52.0 // indirect
 	golang.org/x/exp v0.0.0-20240119083558-1b970713d09a // indirect
-	golang.org/x/net v0.53.0 // indirect
+	golang.org/x/net v0.55.0 // indirect
 	golang.org/x/sync v0.20.0 // indirect
-	golang.org/x/sys v0.43.0 // indirect
-	golang.org/x/text v0.36.0 // indirect
+	golang.org/x/sys v0.45.0 // indirect
+	golang.org/x/text v0.37.0 // indirect
 	golang.org/x/time v0.12.0 // indirect
 	golang.zx2c4.com/wireguard/windows v0.5.3 // indirect
 	google.golang.org/genproto/googleapis/rpc v0.0.0-20251202230838-ff82c1b0f217 // indirect

go.sum

@@ -248,20 +248,20 @@ golang.org/x/arch v0.22.0/go.mod h1:dNHoOeKiyja7GTvF9NJS1l3Z2yntpQNzgrjh1cU103A=
 golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
 golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI=
 golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto=
-golang.org/x/crypto v0.50.0 h1:zO47/JPrL6vsNkINmLoo/PH1gcxpls50DNogFvB5ZGI=
-golang.org/x/crypto v0.50.0/go.mod h1:3muZ7vA7PBCE6xgPX7nkzzjiUq87kRItoJQM1Yo8S+Q=
+golang.org/x/crypto v0.52.0 h1:RMs7fP2rXdep0CftQlK8Uf+kibLm7qkCcradZWYz988=
+golang.org/x/crypto v0.52.0/go.mod h1:1QgfPxDqh0T2M/elOJtp9RvuR95kVjir0e6/BvEmGbc=
 golang.org/x/exp v0.0.0-20240119083558-1b970713d09a h1:Q8/wZp0KX97QFTc2ywcOE0YRjZPVIx+MXInMzdvQqcA=
 golang.org/x/exp v0.0.0-20240119083558-1b970713d09a/go.mod h1:idGWGoKP1toJGkd5/ig9ZLuPcZBC3ewk7SzmH0uou08=
 golang.org/x/mod v0.2.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
 golang.org/x/mod v0.3.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA=
-golang.org/x/mod v0.34.0 h1:xIHgNUUnW6sYkcM5Jleh05DvLOtwc6RitGHbDk4akRI=
-golang.org/x/mod v0.34.0/go.mod h1:ykgH52iCZe79kzLLMhyCUzhMci+nQj+0XkbXpNYtVjY=
+golang.org/x/mod v0.35.0 h1:Ww1D637e6Pg+Zb2KrWfHQUnH2dQRLBQyAtpr/haaJeM=
+golang.org/x/mod v0.35.0/go.mod h1:+GwiRhIInF8wPm+4AoT6L0FA1QWAad3OMdTRx4tFYlU=
 golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg=
 golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
 golang.org/x/net v0.0.0-20200226121028-0de0cce0169b/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
 golang.org/x/net v0.0.0-20201021035429-f5854403a974/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU=
-golang.org/x/net v0.53.0 h1:d+qAbo5L0orcWAr0a9JweQpjXF19LMXJE8Ey7hwOdUA=
-golang.org/x/net v0.53.0/go.mod h1:JvMuJH7rrdiCfbeHoo3fCQU24Lf5JJwT9W3sJFulfgs=
+golang.org/x/net v0.55.0 h1:bcvxaJn3e1U6InsFWt1JUq1aSjnRxLzT2rtD2KfkDF8=
+golang.org/x/net v0.55.0/go.mod h1:L5U2KuzuOe1lY7Z+aWVIKK6qEeJXnXV9yzGA+WCHJww=
 golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20190911185100-cd5d95a43a6e/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20201020160332-67f06af15bc9/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
@@ -278,12 +278,12 @@ golang.org/x/sys v0.0.0-20201204225414-ed752295db88/go.mod h1:h1NjWce9XRLGQEsW7w
 golang.org/x/sys v0.1.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.4.1-0.20230131160137-e7d7f63158de/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
-golang.org/x/sys v0.43.0 h1:Rlag2XtaFTxp19wS8MXlJwTvoh8ArU6ezoyFsMyCTNI=
-golang.org/x/sys v0.43.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw=
+golang.org/x/sys v0.45.0 h1:dO4czNzziLiiXplLQgBCEpCvXQ3dnkn0SdaZSYdQ+FY=
+golang.org/x/sys v0.45.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
 golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
-golang.org/x/text v0.36.0 h1:JfKh3XmcRPqZPKevfXVpI1wXPTqbkE5f7JA92a55Yxg=
-golang.org/x/text v0.36.0/go.mod h1:NIdBknypM8iqVmPiuco0Dh6P5Jcdk8lJL0CUebqK164=
+golang.org/x/text v0.37.0 h1:Cqjiwd9eSg8e0QAkyCaQTNHFIIzWtidPahFWR83rTrc=
+golang.org/x/text v0.37.0/go.mod h1:a5sjxXGs9hsn/AJVwuElvCAo9v8QYLzvavO5z2PiM38=
 golang.org/x/time v0.12.0 h1:ScB/8o8olJvc+CQPWrK3fPZNfh7qgwCrY0zJmoEQLSE=
 golang.org/x/time v0.12.0/go.mod h1:CDIdPxbZBQxdj6cxyCIdrNogrJKMJ7pr37NYpMcMDSg=
 golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=

internal/agentstate/sqlite.go

@@ -112,6 +112,22 @@ func (s *sqliteState) SetNodeUUID(ctx context.Context, value string) error {
 	return s.setMetadata(ctx, pkgmetadata.MetadataKeyMachineID, value)
 }
 
+func (s *sqliteState) GetNodeGroup(ctx context.Context) (string, bool, error) {
+	return s.getMetadata(ctx, MetadataKeyNodeGroup)
+}
+
+func (s *sqliteState) SetNodeGroup(ctx context.Context, value string) error {
+	return s.setMetadata(ctx, MetadataKeyNodeGroup, value)
+}
+
+func (s *sqliteState) GetComputeZone(ctx context.Context) (string, bool, error) {
+	return s.getMetadata(ctx, MetadataKeyComputeZone)
+}
+
+func (s *sqliteState) SetComputeZone(ctx context.Context, value string) error {
+	return s.setMetadata(ctx, MetadataKeyComputeZone, value)
+}
+
 func (s *sqliteState) GetEnrollmentTime(ctx context.Context) (time.Time, bool, error) {
 	value, ok, err := s.getMetadata(ctx, MetadataKeyEnrolledAt)
 	if err != nil || !ok {