[Snapshot-Agent Memory-Regions 1/N]: Adding proto and feature flag (#163)

* snapshot-agent: define the memory-regions backend in the API and route it

Add MemoryRegion and MemoryRegionsBackendConfig to the BackendConfig
oneof (field 5) with regenerated Go and Python bindings: selective
checkpoint/restore of explicit device memory ranges of a running
process, with snapshot_name naming the snapshot slot so multiple
snapshots of one process can coexist (job_id fallback; deliberately not
the request group, which identifies related jobs for the orchestrator
and does not name agent-side storage).

Route memory_regions configs to the new BackendMemoryRegions type in
getSnapshotBackendType. No implementation registers this backend yet,
so a memory_regions request fails cleanly at backend lookup; every
other config routes exactly as before.

Signed-off-by: Edwinhr716 <edandres249@gmail.com>

* snapshot-agent: gate the memory_regions backend behind MemoryRegionsBackend

Register a MemoryRegionsBackend feature gate (alpha, default off) and
reject memory_regions configs in checkFeatureGates when it is disabled,
mirroring DirectMemoryBackend: the backend is driven by GPU-CR and
requires the workload to run under the GPU-CR preloader with a shared
checkpoint directory, so selecting it must be an explicit per-agent
opt-in. The check runs before routing, so a gated config never falls
through to another backend. Gate-independence and error-text cases are
pinned in the server gate tests.

Signed-off-by: Edwinhr716 <edandres249@gmail.com>

* tests: include MemoryRegionsBackend in the gate-registry String expectations

Signed-off-by: Edwinhr716 <edandres249@gmail.com>

---------

Signed-off-by: Edwinhr716 <edandres249@gmail.com>

Author: Edwinhr716

pkg/client/python/timeslice/snapshot_agent/snapshot_agent_pb2.py

@@ -121,6 +121,31 @@ message AppChannelConfig {
   repeated string tags = 2;
 }
 
+// MemoryRegion identifies one device-memory range owned by a process.
+message MemoryRegion {
+  // PID of the process that owns the region.
+  int32 pid = 1;
+  // Device virtual address of the start of the region.
+  uint64 address = 2;
+  // Length of the region in bytes.
+  uint64 size_bytes = 3;
+}
+
+// Configuration for the memory-regions backend: selective checkpoint and
+// restore of specific device memory ranges of a running process (GPU-CR
+// cr_client with a `-s addr:size,...` spec). Regions are required on both
+// Snapshot and Restore; the agent performs no discovery for this backend.
+message MemoryRegionsBackendConfig {
+  repeated MemoryRegion regions = 1;
+  // snapshot_name names the snapshot slot the regions are saved to
+  // (Snapshot) or loaded from (Restore). Distinct names let multiple
+  // snapshots of the same process coexist and be swapped on demand.
+  // If empty, the request's job_id is used. Note this is deliberately NOT
+  // the request's `group`: group identifies a set of related jobs for the
+  // orchestrator and does not name agent-side storage.
+  string snapshot_name = 2;
+}
+
 message BackendConfig {
   oneof backend {
     CudaBackendConfig cuda = 1;
@@ -129,6 +154,7 @@ message BackendConfig {
     // Experimental: disabled by default behind the DirectMemoryBackend
     // feature gate; see DirectMemoryBackendConfig.
     DirectMemoryBackendConfig direct_memory = 4;
+    MemoryRegionsBackendConfig memory_regions = 5;
   }
 }
 

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

@@ -20,6 +20,9 @@ const (
 	// BackendAppChannel suspends/resumes application-aware workloads through
 	// their registered workload channels (see the WorkloadChannel RPC).
 	BackendAppChannel BackendType = "app-channel"
+	// BackendMemoryRegions selectively checkpoints/restores explicit device
+	// memory regions of a running process via the GPU-CR cr_client.
+	BackendMemoryRegions BackendType = "memory-regions"
 )
 
 // Request carries one backend invocation: the job it targets and the

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

@@ -19,12 +19,19 @@ type Feature string
 // and operational caveats.
 const DirectMemoryBackend Feature = "DirectMemoryBackend"
 
+// MemoryRegionsBackend gates the memory_regions backend. Like
+// direct_memory it is driven by GPU-CR and additionally requires the
+// target workload to run under the GPU-CR preloader with a shared
+// checkpoint directory.
+const MemoryRegionsBackend Feature = "MemoryRegionsBackend"
+
 // defaults registers every known gate and its default state. Alpha gates
 // default to false; flipping a default here is the promotion path
 // (alpha → beta → GA), as in Kubernetes. Adding a gate is one const plus
 // one entry here.
 var defaults = map[Feature]bool{
-	DirectMemoryBackend: false,
+	DirectMemoryBackend:  false,
+	MemoryRegionsBackend: false,
 }
 
 // Gates holds the explicitly configured gate values. The zero value (nil)

pkg/snapshot-agent/backends/checkpoint.go

@@ -180,12 +180,13 @@ func TestGates_String(t *testing.T) {
 	withGate(t, "AAATestGate", false)
 
 	got := Gates{DirectMemoryBackend: true}.String()
-	want := "AAATestGate=false,DirectMemoryBackend=true"
+	want := "AAATestGate=false,DirectMemoryBackend=true,MemoryRegionsBackend=false"
 	if got != want {
 		t.Errorf("String() = %q, want %q", got, want)
 	}
 
-	if got := Gates(nil).String(); !strings.Contains(got, "DirectMemoryBackend=false") {
+	if got := Gates(nil).String(); !strings.Contains(got, "DirectMemoryBackend=false") ||
+		!strings.Contains(got, "MemoryRegionsBackend=false") {
 		t.Errorf("nil Gates String() should show defaults, got %q", got)
 	}
 }

pkg/snapshot-agent/features/features.go

@@ -65,6 +65,13 @@ func (s *Server) checkFeatureGates(config *pb.BackendConfig) error {
 				"with --feature-gates=%s=true (or the FEATURE_GATES env var) to enable it",
 			features.DirectMemoryBackend)
 	}
+	if config.GetMemoryRegions() != nil && !s.featureGates.Enabled(features.MemoryRegionsBackend) {
+		return status.Errorf(codes.FailedPrecondition,
+			"the memory_regions backend is experimental (driven by GPU-CR, which carries deployment "+
+				"requirements and operational caveats) and is disabled by default; restart the agent "+
+				"with --feature-gates=%s=true (or the FEATURE_GATES env var) to enable it",
+			features.MemoryRegionsBackend)
+	}
 	return nil
 }
 
@@ -118,9 +125,9 @@ func (s *Server) getSnapshotBackendType(config *pb.BackendConfig) backends.Backe
 	if config.GetAppChannel() != nil {
 		return backends.BackendAppChannel
 	}
-	// NOTE: direct_memory is not routed yet and falls through to the
-	// default backend. It is unreachable unless the DirectMemoryBackend
-	// feature gate is enabled (checkFeatureGates runs before routing).
+	if config.GetMemoryRegions() != nil {
+		return backends.BackendMemoryRegions
+	}
 	return s.defaultBackend
 }
 

pkg/snapshot-agent/features/features_internal_test.go

@@ -27,12 +27,24 @@ func directMemoryConfig() *pb.BackendConfig {
 	}
 }
 
+func memoryRegionsGatedConfig() *pb.BackendConfig {
+	return &pb.BackendConfig{
+		Backend: &pb.BackendConfig_MemoryRegions{
+			MemoryRegions: &pb.MemoryRegionsBackendConfig{
+				Regions: []*pb.MemoryRegion{{Pid: 123, Address: 0x7f0000000000, SizeBytes: 4096}},
+			},
+		},
+	}
+}
+
 func TestServer_CheckFeatureGates(t *testing.T) {
 	tests := []struct {
 		name     string
 		gates    features.Gates
 		config   *pb.BackendConfig
 		wantCode codes.Code
+		// wantGate is the gate a FailedPrecondition error must name.
+		wantGate features.Feature
 	}{
 		{
 			name:     "nil config passes with nil gates",
@@ -53,12 +65,48 @@ func TestServer_CheckFeatureGates(t *testing.T) {
 			gates:    nil,
 			config:   directMemoryConfig(),
 			wantCode: codes.FailedPrecondition,
+			wantGate: features.DirectMemoryBackend,
 		},
 		{
 			name:     "gated config rejected with explicit false",
 			gates:    features.Gates{features.DirectMemoryBackend: false},
 			config:   directMemoryConfig(),
 			wantCode: codes.FailedPrecondition,
+			wantGate: features.DirectMemoryBackend,
+		},
+		{
+			name:     "memory-regions config rejected by default",
+			gates:    nil,
+			config:   memoryRegionsGatedConfig(),
+			wantCode: codes.FailedPrecondition,
+			wantGate: features.MemoryRegionsBackend,
+		},
+		{
+			name:     "memory-regions config rejected with explicit false",
+			gates:    features.Gates{features.MemoryRegionsBackend: false},
+			config:   memoryRegionsGatedConfig(),
+			wantCode: codes.FailedPrecondition,
+			wantGate: features.MemoryRegionsBackend,
+		},
+		{
+			name:     "memory-regions config passes with its gate enabled",
+			gates:    features.Gates{features.MemoryRegionsBackend: true},
+			config:   memoryRegionsGatedConfig(),
+			wantCode: codes.OK,
+		},
+		{
+			name:     "direct-memory gate does not enable memory-regions",
+			gates:    features.Gates{features.DirectMemoryBackend: true},
+			config:   memoryRegionsGatedConfig(),
+			wantCode: codes.FailedPrecondition,
+			wantGate: features.MemoryRegionsBackend,
+		},
+		{
+			name:     "memory-regions gate does not enable direct-memory",
+			gates:    features.Gates{features.MemoryRegionsBackend: true},
+			config:   directMemoryConfig(),
+			wantCode: codes.FailedPrecondition,
+			wantGate: features.DirectMemoryBackend,
 		},
 		{
 			name:     "gated config passes with its gate enabled",
@@ -85,8 +133,8 @@ func TestServer_CheckFeatureGates(t *testing.T) {
 				t.Fatalf("checkFeatureGates() = %v, want code %v", err, tc.wantCode)
 			}
 			if tc.wantCode == codes.FailedPrecondition {
-				if !strings.Contains(err.Error(), string(features.DirectMemoryBackend)) {
-					t.Errorf("Expected error to name the gate, got: %v", err)
+				if !strings.Contains(err.Error(), string(tc.wantGate)) {
+					t.Errorf("Expected error to name gate %s, got: %v", tc.wantGate, err)
 				}
 				if !strings.Contains(err.Error(), "--feature-gates") {
 					t.Errorf("Expected error to name the --feature-gates flag, got: %v", err)