Skip to content

Commit 70ce9ae

Browse files
authored
snapshot-agent: add Kubernetes-style feature gates; gate the direct_memory backend (#149)
1 parent 64dac67 commit 70ce9ae

11 files changed

Lines changed: 645 additions & 10 deletions

File tree

cmd/snapshot-agent/main.go

Lines changed: 21 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ import (
2323

2424
"github.com/llm-d-incubation/llm-d-rl-time-slicing/pkg/logging"
2525
"github.com/llm-d-incubation/llm-d-rl-time-slicing/pkg/snapshot-agent/backends"
26+
"github.com/llm-d-incubation/llm-d-rl-time-slicing/pkg/snapshot-agent/features"
2627
"github.com/llm-d-incubation/llm-d-rl-time-slicing/pkg/snapshot-agent/server"
2728
)
2829

@@ -34,6 +35,8 @@ func main() {
3435

3536
port := flag.Int("port", 9001, "The port to listen on")
3637
deploymentMode := flag.String("deployment-mode", "standalone", "Deployment mode ('standalone' or 'k8s')")
38+
featureGatesSpec := flag.String("feature-gates", "",
39+
"Comma-separated list of Name=bool pairs selecting experimental features, e.g. 'DirectMemoryBackend=true'")
3740
flag.Parse()
3841

3942
depMode := *deploymentMode
@@ -57,6 +60,19 @@ func main() {
5760
slog.Error("Invalid deployment mode, must be 'standalone' or 'k8s'", "mode", depMode)
5861
os.Exit(1)
5962
}
63+
64+
// FEATURE_GATES overrides the flag, mirroring DEPLOYMENT_MODE and
65+
// AGENT_PORT: the Helm chart configures the agent through env vars.
66+
gatesSpec := *featureGatesSpec
67+
if envGates := os.Getenv("FEATURE_GATES"); envGates != "" {
68+
gatesSpec = envGates
69+
}
70+
featureGates, err := features.Parse(gatesSpec)
71+
if err != nil {
72+
slog.Error("Invalid feature gates", "value", gatesSpec, "error", err)
73+
os.Exit(1)
74+
}
75+
6076
ctx := context.Background()
6177

6278
// The channel registry is shared between the app-channel backend and the
@@ -69,8 +85,11 @@ func main() {
6985
backends.BackendAppChannel: backends.NewAppChannelBackend(channelRegistry),
7086
}
7187

72-
slog.InfoContext(ctx, "Starting Snapshot Agent", "port", listenPort, "deploymentMode", depMode)
73-
if err := server.StartServer(ctx, listenPort, registeredBackends, backends.BackendCuda, depMode, channelRegistry); err != nil {
88+
slog.InfoContext(ctx, "Starting Snapshot Agent",
89+
"port", listenPort, "deploymentMode", depMode, "featureGates", featureGates.String())
90+
err = server.StartServer(
91+
ctx, listenPort, registeredBackends, backends.BackendCuda, depMode, channelRegistry, featureGates)
92+
if err != nil {
7493
slog.ErrorContext(ctx, "Failed to start server", "error", err)
7594
os.Exit(1)
7695
}

deploy/snapshot-agent/templates/daemonset.yaml

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,17 @@ spec:
4242
fieldPath: spec.nodeName
4343
- name: AGENT_PORT
4444
value: {{ .Values.port | quote }}
45+
{{- if .Values.featureGates }}
46+
{{- $gates := list }}
47+
{{- range $name, $enabled := .Values.featureGates }}
48+
{{- if not (kindIs "bool" $enabled) }}
49+
{{- fail (printf "featureGates.%s must be true or false, got %q" $name (toString $enabled)) }}
50+
{{- end }}
51+
{{- $gates = append $gates (printf "%s=%t" $name $enabled) }}
52+
{{- end }}
53+
- name: FEATURE_GATES
54+
value: {{ join "," $gates | quote }}
55+
{{- end }}
4556
{{- if .Values.nvidia.driver.enabled }}
4657
- name: LD_LIBRARY_PATH
4758
value: {{ join ":" .Values.nvidia.driver.libraryPaths }}

deploy/snapshot-agent/values.yaml

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,18 @@ nvidia:
4747
hostPath: /dev
4848
mountPath: /dev
4949

50+
# Kubernetes-style feature gates (map of Name: bool), rendered into the
51+
# agent's FEATURE_GATES env var. Experimental capabilities are off by
52+
# default; requests selecting a gated backend fail with FAILED_PRECONDITION
53+
# until its gate is enabled.
54+
# Known gates:
55+
# DirectMemoryBackend — the experimental direct_memory backend (driven by
56+
# GPU-CR). Default: false.
57+
# Example:
58+
# featureGates:
59+
# DirectMemoryBackend: true
60+
featureGates: {}
61+
5062
rbac:
5163
create: true
5264

pkg/snapshot-agent/api/v1alpha1/snapshot_agent.pb.go

Lines changed: 8 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

pkg/snapshot-agent/api/v1alpha1/snapshot_agent.proto

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,12 @@ message CudaBackendConfig {
5151
}
5252

5353
// Configuration for the Direct Memory (process-level) backend.
54+
//
55+
// Experimental: this backend is driven by GPU-CR, which is itself
56+
// experimental — expect deployment requirements and operational caveats.
57+
// It is disabled by default: requests selecting it fail with
58+
// FAILED_PRECONDITION unless the agent runs with
59+
// --feature-gates=DirectMemoryBackend=true (or the FEATURE_GATES env var).
5460
message DirectMemoryBackendConfig {
5561
// The target processes to checkpoint/restore.
5662
ProcessTarget explicit_target = 1;
@@ -120,6 +126,8 @@ message BackendConfig {
120126
CudaBackendConfig cuda = 1;
121127
AppEndpointConfig app_endpoint = 2;
122128
AppChannelConfig app_channel = 3;
129+
// Experimental: disabled by default behind the DirectMemoryBackend
130+
// feature gate; see DirectMemoryBackendConfig.
123131
DirectMemoryBackendConfig direct_memory = 4;
124132
}
125133
}
Lines changed: 95 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,95 @@
1+
// Package features implements Kubernetes-style feature gates for the
2+
// snapshot agent: experimental capabilities are explicit per-agent opt-ins
3+
// selected with --feature-gates=Name=bool,... (or the FEATURE_GATES env
4+
// var), with alpha gates off by default.
5+
package features
6+
7+
import (
8+
"fmt"
9+
"sort"
10+
"strconv"
11+
"strings"
12+
)
13+
14+
// Feature is the name of a feature gate.
15+
type Feature string
16+
17+
// DirectMemoryBackend gates the direct_memory backend. It is driven by
18+
// GPU-CR, which is itself experimental and carries deployment requirements
19+
// and operational caveats.
20+
const DirectMemoryBackend Feature = "DirectMemoryBackend"
21+
22+
// defaults registers every known gate and its default state. Alpha gates
23+
// default to false; flipping a default here is the promotion path
24+
// (alpha → beta → GA), as in Kubernetes. Adding a gate is one const plus
25+
// one entry here.
26+
var defaults = map[Feature]bool{
27+
DirectMemoryBackend: false,
28+
}
29+
30+
// Gates holds the explicitly configured gate values. The zero value (nil)
31+
// is valid and yields every gate's default — callers that don't care pass
32+
// nil.
33+
type Gates map[Feature]bool
34+
35+
// Parse parses a comma-separated list of Name=bool pairs, tolerating
36+
// whitespace around names, values, and separators. An empty spec yields
37+
// the defaults. Unknown gate names and non-boolean values are errors so
38+
// that a typo never runs silently ignored. A gate repeated within the
39+
// spec takes its last value, matching Kubernetes --feature-gates
40+
// behavior.
41+
func Parse(spec string) (Gates, error) {
42+
gates := Gates{}
43+
if strings.TrimSpace(spec) == "" {
44+
return gates, nil
45+
}
46+
for _, pair := range strings.Split(spec, ",") {
47+
pair = strings.TrimSpace(pair)
48+
if pair == "" {
49+
continue
50+
}
51+
name, value, found := strings.Cut(pair, "=")
52+
if !found {
53+
return nil, fmt.Errorf("invalid feature gate %q: must be Name=bool", pair)
54+
}
55+
feature := Feature(strings.TrimSpace(name))
56+
if _, known := defaults[feature]; !known {
57+
return nil, fmt.Errorf("unknown feature gate %q; known gates: %s", strings.TrimSpace(name), knownGates())
58+
}
59+
enabled, err := strconv.ParseBool(strings.TrimSpace(value))
60+
if err != nil {
61+
return nil, fmt.Errorf("invalid value %q for feature gate %q: must be a boolean", strings.TrimSpace(value), feature)
62+
}
63+
gates[feature] = enabled
64+
}
65+
return gates, nil
66+
}
67+
68+
// Enabled reports whether the gate is on: the explicitly configured value
69+
// if set, else the registry default.
70+
func (g Gates) Enabled(f Feature) bool {
71+
if enabled, ok := g[f]; ok {
72+
return enabled
73+
}
74+
return defaults[f]
75+
}
76+
77+
// String returns every known gate with its effective value as a sorted,
78+
// comma-separated Name=bool list, for startup logging.
79+
func (g Gates) String() string {
80+
pairs := make([]string, 0, len(defaults))
81+
for feature := range defaults {
82+
pairs = append(pairs, fmt.Sprintf("%s=%t", feature, g.Enabled(feature)))
83+
}
84+
sort.Strings(pairs)
85+
return strings.Join(pairs, ",")
86+
}
87+
88+
func knownGates() string {
89+
names := make([]string, 0, len(defaults))
90+
for feature := range defaults {
91+
names = append(names, string(feature))
92+
}
93+
sort.Strings(names)
94+
return strings.Join(names, ", ")
95+
}

0 commit comments

Comments
 (0)