Implement Direct Memory (process-level) backend in snapshot-agent

- Add BackendDirectMemory constant to BackendType
- Implement DirectMemory backend struct using cr_client binary
- Support direct_memory backend in server logic and PID discovery for k8s and standalone modes
- Add BuildDirectMemoryConfig and ExtractDirectMemoryPIDStrings helper functions
- Register DirectMemory backend in cmd/snapshot-agent
- Add unit tests for DirectMemory backend and server integration

Author: Edwinhr716

cmd/snapshot-agent/main.go

@@ -63,10 +63,11 @@ func main() {
 	// server's WorkloadChannel RPC handler.
 	channelRegistry := backends.NewChannelRegistry()
 	registeredBackends := map[backends.BackendType]backends.Backend{
-		backends.BackendCuda:        backends.NewCudaCheckpoint(),
-		backends.BackendNoop:        backends.NewNoopBackend(),
-		backends.BackendAppEndpoint: backends.NewAppEndpointBackend(),
-		backends.BackendAppChannel:  backends.NewAppChannelBackend(channelRegistry),
+		backends.BackendCuda:         backends.NewCudaCheckpoint(),
+		backends.BackendNoop:         backends.NewNoopBackend(),
+		backends.BackendAppEndpoint:  backends.NewAppEndpointBackend(),
+		backends.BackendAppChannel:   backends.NewAppChannelBackend(channelRegistry),
+		backends.BackendDirectMemory: backends.NewDirectMemory(),
 	}
 
 	slog.InfoContext(ctx, "Starting Snapshot Agent", "port", listenPort, "deploymentMode", depMode)

pkg/snapshot-agent/backends/checkpoint.go

@@ -20,6 +20,8 @@ const (
 	// BackendAppChannel suspends/resumes application-aware workloads through
 	// their registered workload channels (see the WorkloadChannel RPC).
 	BackendAppChannel BackendType = "app-channel"
+	// BackendDirectMemory is the Direct Memory (process-level) checkpointing backend.
+	BackendDirectMemory BackendType = "direct-memory"
 )
 
 // Request carries one backend invocation: the job it targets and the

pkg/snapshot-agent/backends/direct_memory.go

@@ -0,0 +1,165 @@
+package backends
+
+import (
+	"context"
+	"fmt"
+	"log/slog"
+	"os"
+	"os/exec"
+	"strconv"
+	"sync"
+	"time"
+
+	pb "github.com/llm-d-incubation/llm-d-rl-time-slicing/pkg/snapshot-agent/api/v1alpha1"
+)
+
+// DirectMemory implements the Backend interface using cr_client.
+type DirectMemory struct {
+	mu          sync.Mutex
+	execCommand func(ctx context.Context, name string, args ...string) ([]byte, error)
+	lookPath    func(string) (string, error)
+	statFunc    func(string) (os.FileInfo, error)
+}
+
+// NewDirectMemory creates a new DirectMemory backend.
+func NewDirectMemory() *DirectMemory {
+	return &DirectMemory{
+		execCommand: func(ctx context.Context, name string, args ...string) ([]byte, error) {
+			return exec.CommandContext(ctx, name, args...).CombinedOutput()
+		},
+		lookPath: exec.LookPath,
+		statFunc: os.Stat,
+	}
+}
+
+// Snapshot triggers a snapshot of the target processes for a job using cr_client.
+func (d *DirectMemory) Snapshot(ctx context.Context, req Request) error {
+	pids := ExtractDirectMemoryPIDStrings(req.Config)
+	if len(pids) == 0 {
+		return fmt.Errorf("at least one PID is required for Direct Memory snapshot")
+	}
+
+	d.mu.Lock()
+	defer d.mu.Unlock()
+
+	slog.InfoContext(ctx, "Snapshotting PIDs using Direct Memory", "pids", pids)
+
+	t0 := time.Now()
+	for _, pid := range pids {
+		if err := d.checkpointPID(ctx, pid); err != nil {
+			return fmt.Errorf("cr_client checkpoint failed for PID %s: %w", pid, err)
+		}
+	}
+	slog.InfoContext(ctx, "cr_client checkpoint took", "duration", time.Since(t0))
+	return nil
+}
+
+// Restore triggers a restoration of the target processes for a job using cr_client.
+func (d *DirectMemory) Restore(ctx context.Context, req Request) error {
+	pids := ExtractDirectMemoryPIDStrings(req.Config)
+	if len(pids) == 0 {
+		return fmt.Errorf("at least one PID is required for Direct Memory restore")
+	}
+
+	d.mu.Lock()
+	defer d.mu.Unlock()
+
+	slog.InfoContext(ctx, "Restoring PIDs using Direct Memory", "pids", pids)
+	t0 := time.Now()
+	for _, pid := range pids {
+		if err := d.restorePID(ctx, pid); err != nil {
+			return fmt.Errorf("cr_client restore failed for PID %s: %w", pid, err)
+		}
+	}
+	slog.InfoContext(ctx, "cr_client restore took", "duration", time.Since(t0), "pids", pids)
+	return nil
+}
+
+func (d *DirectMemory) getCrClientPath() string {
+	for _, p := range []string{
+		"cr_client",
+		"/usr/bin/cr_client",
+		"/bin/cr_client",
+		"/opt/bin/cr_client",
+		"/usr/local/bin/cr_client",
+	} {
+		if path, err := d.lookPath(p); err == nil {
+			return path
+		}
+		if _, err := d.statFunc(p); err == nil {
+			return p
+		}
+	}
+	return "/opt/bin/cr_client"
+}
+
+func (d *DirectMemory) runCommand(ctx context.Context, name string, args ...string) error {
+	if out, err := d.execCommand(ctx, name, args...); err != nil {
+		return fmt.Errorf("command failed: %w, output: %s", err, string(out))
+	}
+	return nil
+}
+
+func (d *DirectMemory) checkpointPID(ctx context.Context, pid string) error {
+	binaryPath := d.getCrClientPath()
+	if err := d.runCommand(ctx, binaryPath, "-c", "-p", pid); err != nil {
+		return err
+	}
+	return nil
+}
+
+func (d *DirectMemory) restorePID(ctx context.Context, pid string) error {
+	binaryPath := d.getCrClientPath()
+	if err := d.runCommand(ctx, binaryPath, "-r", "-p", pid); err != nil {
+		return err
+	}
+	return nil
+}
+
+// HealthCheck checks if the Direct Memory backend is healthy.
+func (d *DirectMemory) HealthCheck(ctx context.Context) error {
+	binaryPath := d.getCrClientPath()
+	if _, err := d.lookPath(binaryPath); err != nil {
+		if _, errStat := d.statFunc(binaryPath); errStat != nil {
+			return fmt.Errorf("cr_client executable not found: %w", err)
+		}
+	}
+	return nil
+}
+
+// ExtractDirectMemoryPIDStrings extracts PID strings from a DirectMemory BackendConfig.
+func ExtractDirectMemoryPIDStrings(config *pb.BackendConfig) []string {
+	if config == nil {
+		return nil
+	}
+	dm := config.GetDirectMemory()
+	if dm == nil {
+		return nil
+	}
+	target := dm.GetExplicitTarget()
+	if target == nil {
+		return nil
+	}
+	pids := make([]string, 0, len(target.GetPids()))
+	for _, pid := range target.GetPids() {
+		pids = append(pids, strconv.Itoa(int(pid)))
+	}
+	return pids
+}
+
+// BuildDirectMemoryConfig wraps PID strings into a DirectMemory BackendConfig.
+func BuildDirectMemoryConfig(pidStrings []string) *pb.BackendConfig {
+	pids := make([]int32, 0, len(pidStrings))
+	for _, s := range pidStrings {
+		if pid, err := strconv.ParseInt(s, 10, 32); err == nil {
+			pids = append(pids, int32(pid))
+		}
+	}
+	return &pb.BackendConfig{
+		Backend: &pb.BackendConfig_DirectMemory{
+			DirectMemory: &pb.DirectMemoryBackendConfig{
+				ExplicitTarget: &pb.ProcessTarget{Pids: pids},
+			},
+		},
+	}
+}

pkg/snapshot-agent/backends/direct_memory_test.go

@@ -0,0 +1,207 @@
+package backends_test
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"reflect"
+	"testing"
+
+	pb "github.com/llm-d-incubation/llm-d-rl-time-slicing/pkg/snapshot-agent/api/v1alpha1"
+	"github.com/llm-d-incubation/llm-d-rl-time-slicing/pkg/snapshot-agent/backends"
+)
+
+func directMemoryConfig(pids ...int32) *pb.BackendConfig {
+	return &pb.BackendConfig{
+		Backend: &pb.BackendConfig_DirectMemory{
+			DirectMemory: &pb.DirectMemoryBackendConfig{
+				ExplicitTarget: &pb.ProcessTarget{Pids: pids},
+			},
+		},
+	}
+}
+
+func TestNewDirectMemory(t *testing.T) {
+	dm := backends.NewDirectMemory()
+	if dm == nil {
+		t.Fatal("NewDirectMemory returned nil")
+	}
+}
+
+func TestDirectMemorySnapshot(t *testing.T) {
+	tests := []struct {
+		name        string
+		config      *pb.BackendConfig
+		execErr     error
+		expectedErr bool
+		expectArgs  [][]string
+	}{
+		{
+			name:   "SuccessMultiplePIDs",
+			config: directMemoryConfig(123, 456),
+			expectArgs: [][]string{
+				{"-c", "-p", "123"},
+				{"-c", "-p", "456"},
+			},
+		},
+		{
+			name:        "ExecFailure",
+			config:      directMemoryConfig(123),
+			execErr:     fmt.Errorf("exec error"),
+			expectedErr: true,
+			expectArgs: [][]string{
+				{"-c", "-p", "123"},
+			},
+		},
+		{
+			name:        "NoPIDs",
+			config:      directMemoryConfig(),
+			expectedErr: true,
+		},
+		{
+			name:        "NilConfig",
+			config:      nil,
+			expectedErr: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			dm := backends.NewDirectMemory()
+			var calledArgs [][]string
+			dm.SetExecCommand(func(_ context.Context, _ string, args ...string) ([]byte, error) {
+				calledArgs = append(calledArgs, args)
+				return nil, tt.execErr
+			})
+
+			err := dm.Snapshot(context.Background(), backends.Request{JobID: "test-job", Config: tt.config})
+			if (err != nil) != tt.expectedErr {
+				t.Errorf("Snapshot() error = %v, expectedErr %v", err, tt.expectedErr)
+			}
+			if !tt.expectedErr && !reflect.DeepEqual(calledArgs, tt.expectArgs) {
+				t.Errorf("Snapshot() calledArgs = %v, expected %v", calledArgs, tt.expectArgs)
+			}
+		})
+	}
+}
+
+func TestDirectMemoryRestore(t *testing.T) {
+	tests := []struct {
+		name        string
+		config      *pb.BackendConfig
+		execErr     error
+		expectedErr bool
+		expectArgs  [][]string
+	}{
+		{
+			name:   "SuccessMultiplePIDs",
+			config: directMemoryConfig(123, 456),
+			expectArgs: [][]string{
+				{"-r", "-p", "123"},
+				{"-r", "-p", "456"},
+			},
+		},
+		{
+			name:        "NoPIDs",
+			config:      directMemoryConfig(),
+			expectedErr: true,
+		},
+		{
+			name:        "NilConfig",
+			config:      nil,
+			expectedErr: true,
+		},
+		{
+			name:        "ExecFailure",
+			config:      directMemoryConfig(123),
+			execErr:     fmt.Errorf("exec error"),
+			expectedErr: true,
+			expectArgs: [][]string{
+				{"-r", "-p", "123"},
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			dm := backends.NewDirectMemory()
+			var calledArgs [][]string
+			dm.SetExecCommand(func(_ context.Context, _ string, args ...string) ([]byte, error) {
+				calledArgs = append(calledArgs, args)
+				return nil, tt.execErr
+			})
+
+			err := dm.Restore(context.Background(), backends.Request{JobID: "test-job", Config: tt.config})
+			if (err != nil) != tt.expectedErr {
+				t.Errorf("Restore() error = %v, expectedErr %v", err, tt.expectedErr)
+			}
+			if !tt.expectedErr && !reflect.DeepEqual(calledArgs, tt.expectArgs) {
+				t.Errorf("Restore() calledArgs = %v, expected %v", calledArgs, tt.expectArgs)
+			}
+		})
+	}
+}
+
+func TestDirectMemoryHealthCheck(t *testing.T) {
+	tests := []struct {
+		name        string
+		lookErr     error
+		statErr     error
+		expectedErr bool
+	}{
+		{
+			name:        "SuccessInPath",
+			lookErr:     nil,
+			statErr:     nil,
+			expectedErr: false,
+		},
+		{
+			name:        "SuccessViaStat",
+			lookErr:     fmt.Errorf("not in path"),
+			statErr:     nil,
+			expectedErr: false,
+		},
+		{
+			name:        "NotFoundAnywhere",
+			lookErr:     fmt.Errorf("not in path"),
+			statErr:     fmt.Errorf("no stat"),
+			expectedErr: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			dm := backends.NewDirectMemory()
+			dm.SetLookPath(func(path string) (string, error) {
+				if tt.lookErr != nil {
+					return "", tt.lookErr
+				}
+				return path, nil
+			})
+			dm.SetStatFunc(func(path string) (os.FileInfo, error) {
+				if tt.statErr != nil {
+					return nil, tt.statErr
+				}
+				return os.Stat(".")
+			})
+
+			err := dm.HealthCheck(context.Background())
+			if (err != nil) != tt.expectedErr {
+				t.Errorf("HealthCheck() error = %v, expectedErr %v", err, tt.expectedErr)
+			}
+		})
+	}
+}
+
+func TestDirectMemoryConfigHelpers(t *testing.T) {
+	pids := []string{"100", "200"}
+	cfg := backends.BuildDirectMemoryConfig(pids)
+	extracted := backends.ExtractDirectMemoryPIDStrings(cfg)
+	if !reflect.DeepEqual(extracted, pids) {
+		t.Errorf("ExtractDirectMemoryPIDStrings() = %v, want %v", extracted, pids)
+	}
+
+	if len(backends.ExtractDirectMemoryPIDStrings(nil)) != 0 {
+		t.Errorf("Expected nil when extracting from nil config")
+	}
+}