fix(gms): rewrite GMS checkpoint/restore operator support (#8194)

Co-authored-by: Dmitry Tokarev <dtokarev@nvidia.com>

deploy/operator/internal/checkpoint/checkpoint_test.go

@@ -23,7 +23,7 @@ import (
 
 	nvidiacomv1alpha1 "github.com/ai-dynamo/dynamo/deploy/operator/api/v1alpha1"
 	"github.com/ai-dynamo/dynamo/deploy/operator/internal/consts"
-	gmsruntime "github.com/ai-dynamo/dynamo/deploy/operator/internal/gms"
+	gms "github.com/ai-dynamo/dynamo/deploy/operator/internal/gms"
 	snapshotprotocol "github.com/ai-dynamo/dynamo/deploy/snapshot/protocol"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -183,50 +183,6 @@ func TestCreateOrGetAutoCheckpointSetsDefaultArtifactVersion(t *testing.T) {
 
 // --- InjectCheckpointIntoPodSpec tests ---
 
-func TestEnsurePodInfoVolumeMergesExistingDownwardAPIItems(t *testing.T) {
-	podSpec := &corev1.PodSpec{
-		Volumes: []corev1.Volume{{
-			Name: consts.PodInfoVolumeName,
-			VolumeSource: corev1.VolumeSource{
-				DownwardAPI: &corev1.DownwardAPIVolumeSource{
-					Items: []corev1.DownwardAPIVolumeFile{
-						{
-							Path:     "pod_name",
-							FieldRef: &corev1.ObjectFieldSelector{FieldPath: "metadata.name"},
-						},
-						{
-							Path:     "custom",
-							FieldRef: &corev1.ObjectFieldSelector{FieldPath: "metadata.labels['custom']"},
-						},
-					},
-				},
-			},
-		}},
-	}
-
-	EnsurePodInfoVolume(podSpec)
-
-	require.Len(t, podSpec.Volumes, 1)
-	require.NotNil(t, podSpec.Volumes[0].DownwardAPI)
-
-	fields := map[string]string{}
-	for _, item := range podSpec.Volumes[0].DownwardAPI.Items {
-		if item.FieldRef != nil {
-			fields[item.Path] = item.FieldRef.FieldPath
-		}
-	}
-
-	assert.Equal(t, consts.PodInfoFieldPodName, fields["pod_name"])
-	assert.Equal(t, consts.PodInfoFieldPodUID, fields["pod_uid"])
-	assert.Equal(t, consts.PodInfoFieldPodNamespace, fields["pod_namespace"])
-	assert.Equal(t, "metadata.labels['custom']", fields["custom"])
-	assert.Equal(t, "metadata.labels['"+consts.KubeLabelDynamoNamespace+"']", fields[consts.PodInfoFileDynNamespace])
-	assert.Equal(t, "metadata.labels['"+consts.KubeLabelDynamoWorkerHash+"']", fields[consts.PodInfoFileDynNamespaceWorkerSuffix])
-	assert.Equal(t, "metadata.labels['"+consts.KubeLabelDynamoComponentType+"']", fields[consts.PodInfoFileDynComponent])
-	assert.Equal(t, "metadata.labels['"+consts.KubeLabelDynamoGraphDeploymentName+"']", fields[consts.PodInfoFileDynParentDGDName])
-	assert.Equal(t, consts.PodInfoFieldPodNamespace, fields[consts.PodInfoFileDynParentDGDNamespace])
-}
-
 func TestInjectCheckpointIntoPodSpec(t *testing.T) {
 	t.Run("ready checkpoint injects podinfo and overrides command", func(t *testing.T) {
 		podSpec := testPodSpec()
@@ -263,65 +219,21 @@ func TestInjectCheckpointIntoPodSpec(t *testing.T) {
 		assert.Equal(t, consts.PodInfoMountPath, mountPaths[consts.PodInfoVolumeName])
 	})
 
-	t.Run("ready checkpoint augments existing podinfo volume", func(t *testing.T) {
-		podSpec := testPodSpec()
-		podSpec.Volumes = append(podSpec.Volumes, corev1.Volume{
-			Name: consts.PodInfoVolumeName,
-			VolumeSource: corev1.VolumeSource{
-				DownwardAPI: &corev1.DownwardAPIVolumeSource{
-					Items: []corev1.DownwardAPIVolumeFile{
-						{Path: "pod_name", FieldRef: &corev1.ObjectFieldSelector{FieldPath: consts.PodInfoFieldPodName}},
-						{Path: "pod_uid", FieldRef: &corev1.ObjectFieldSelector{FieldPath: consts.PodInfoFieldPodUID}},
-						{Path: "pod_namespace", FieldRef: &corev1.ObjectFieldSelector{FieldPath: consts.PodInfoFieldPodNamespace}},
-					},
-				},
-			},
-		})
-		info := &CheckpointInfo{Enabled: true, Ready: true, Identity: ptr.To(testIdentity())}
-		reader := fake.NewClientBuilder().WithScheme(testScheme()).WithObjects(testSnapshotAgentDaemonSet()).Build()
-		require.NoError(t, InjectCheckpointIntoPodSpec(context.Background(), reader, testNamespace, podSpec, info))
-
-		var podInfoVolume *corev1.Volume
-		for i := range podSpec.Volumes {
-			if podSpec.Volumes[i].Name == consts.PodInfoVolumeName {
-				podInfoVolume = &podSpec.Volumes[i]
-				break
-			}
-		}
-		require.NotNil(t, podInfoVolume)
-		require.NotNil(t, podInfoVolume.DownwardAPI)
-
-		fields := map[string]string{}
-		for _, item := range podInfoVolume.DownwardAPI.Items {
-			if item.FieldRef != nil {
-				fields[item.Path] = item.FieldRef.FieldPath
-			}
-		}
-		assert.Equal(t, consts.PodInfoFieldPodName, fields["pod_name"])
-		assert.Equal(t, consts.PodInfoFieldPodUID, fields["pod_uid"])
-		assert.Equal(t, consts.PodInfoFieldPodNamespace, fields["pod_namespace"])
-		assert.Equal(t, "metadata.labels['"+consts.KubeLabelDynamoNamespace+"']", fields[consts.PodInfoFileDynNamespace])
-		assert.Equal(t, "metadata.labels['"+consts.KubeLabelDynamoWorkerHash+"']", fields[consts.PodInfoFileDynNamespaceWorkerSuffix])
-		assert.Equal(t, "metadata.labels['"+consts.KubeLabelDynamoComponentType+"']", fields[consts.PodInfoFileDynComponent])
-		assert.Equal(t, "metadata.labels['"+consts.KubeLabelDynamoGraphDeploymentName+"']", fields[consts.PodInfoFileDynParentDGDName])
-		assert.Equal(t, consts.PodInfoFieldPodNamespace, fields[consts.PodInfoFileDynParentDGDNamespace])
-	})
-
 	t.Run("ready checkpoint targets the container named main", func(t *testing.T) {
 		podSpec := &corev1.PodSpec{
 			Containers: []corev1.Container{
+				{Name: "main", Image: "main:latest", Command: []string{"python3"}, Args: []string{"-m", "dynamo.vllm"}},
 				{Name: "sidecar", Image: "sidecar:latest", Command: []string{"sidecar"}, Args: []string{"run"}},
-				{Name: consts.MainContainerName, Image: "main:latest", Command: []string{"python3"}, Args: []string{"-m", "dynamo.vllm"}},
 			},
 		}
 		info := &CheckpointInfo{Enabled: true, Ready: true, Hash: testHash}
 		reader := fake.NewClientBuilder().WithScheme(testScheme()).WithObjects(testSnapshotAgentDaemonSet()).Build()
 
 		require.NoError(t, InjectCheckpointIntoPodSpec(context.Background(), reader, testNamespace, podSpec, info))
-		assert.Equal(t, []string{"sidecar"}, podSpec.Containers[0].Command)
-		assert.Equal(t, []string{"run"}, podSpec.Containers[0].Args)
-		assert.Equal(t, []string{"sleep", "infinity"}, podSpec.Containers[1].Command)
-		assert.Nil(t, podSpec.Containers[1].Args)
+		assert.Equal(t, []string{"sleep", "infinity"}, podSpec.Containers[0].Command)
+		assert.Nil(t, podSpec.Containers[0].Args)
+		assert.Equal(t, []string{"sidecar"}, podSpec.Containers[1].Command)
+		assert.Equal(t, []string{"run"}, podSpec.Containers[1].Args)
 	})
 
 	t.Run("ready gms checkpoint injects restore sidecars and loader mount", func(t *testing.T) {
@@ -331,22 +243,24 @@ func TestInjectCheckpointIntoPodSpec(t *testing.T) {
 		reader := fake.NewClientBuilder().WithScheme(testScheme()).WithObjects(testSnapshotAgentDaemonSet()).Build()
 
 		require.NoError(t, InjectCheckpointIntoPodSpec(context.Background(), reader, testNamespace, podSpec, info))
-		gmsServer := findContainer(podSpec, gmsruntime.ServerContainerName)
+		gmsServer := findContainer(podSpec, gms.ServerContainerName)
 		require.NotNil(t, gmsServer)
 		loader := findContainer(podSpec, GMSLoaderContainer)
 		require.NotNil(t, loader)
 
-		// Restore: gms-server should be a regular container, not an init container
-		assert.Empty(t, podSpec.InitContainers, "restore pods should not have gms-server as init container")
-		assert.Nil(t, gmsServer.RestartPolicy, "restore gms-server should not have RestartPolicy")
+		// Restore: server and loader are init sidecars (restartPolicy=Always)
+		assert.NotNil(t, gmsServer.RestartPolicy, "restore gms-server should have RestartPolicy")
+		assert.Equal(t, corev1.ContainerRestartPolicyAlways, *gmsServer.RestartPolicy)
 		assert.Nil(t, gmsServer.StartupProbe, "restore gms-server should not have StartupProbe")
+		assert.NotNil(t, loader.RestartPolicy, "restore gms-loader should have RestartPolicy")
+		assert.Equal(t, corev1.ContainerRestartPolicyAlways, *loader.RestartPolicy)
 
 		mounts := map[string]string{}
 		for _, mount := range loader.VolumeMounts {
 			mounts[mount.Name] = mount.MountPath
 		}
 		assert.Equal(t, "/checkpoints", mounts[snapshotprotocol.CheckpointVolumeName])
-		assert.Equal(t, gmsruntime.SharedMountPath, mounts[gmsruntime.SharedVolumeName])
+		assert.Equal(t, gms.SharedMountPath, mounts[gms.SharedVolumeName])
 
 		env := map[string]string{}
 		for _, item := range loader.Env {
@@ -366,8 +280,7 @@ func TestInjectCheckpointIntoPodSpec(t *testing.T) {
 			errMsg  string
 		}{
 			{"hash empty and identity nil", testPodSpec(), &CheckpointInfo{Enabled: true}, fake.NewClientBuilder().WithScheme(testScheme()).WithObjects(testSnapshotAgentDaemonSet()).Build(), "identity is nil"},
-			{"no containers", &corev1.PodSpec{}, testInfo(), fake.NewClientBuilder().WithScheme(testScheme()).WithObjects(testSnapshotAgentDaemonSet()).Build(), "no container found"},
-			{"main container missing", &corev1.PodSpec{Containers: []corev1.Container{{Name: "sidecar", Image: "img", Command: []string{"python3"}}}}, testInfo(), fake.NewClientBuilder().WithScheme(testScheme()).WithObjects(testSnapshotAgentDaemonSet()).Build(), "main container not found"},
+			{"no containers", &corev1.PodSpec{}, testInfo(), fake.NewClientBuilder().WithScheme(testScheme()).WithObjects(testSnapshotAgentDaemonSet()).Build(), "no container named"},
 			{"snapshot daemonset missing", testPodSpec(), testInfo(), fake.NewClientBuilder().WithScheme(testScheme()).Build(), "no snapshot-agent daemonset found"},
 		} {
 			t.Run(tc.name, func(t *testing.T) {

deploy/operator/internal/checkpoint/gms.go

@@ -6,15 +6,13 @@
 package checkpoint
 
 import (
-	"context"
 	"fmt"
 	"path/filepath"
 
-	gmsruntime "github.com/ai-dynamo/dynamo/deploy/operator/internal/gms"
+	gms "github.com/ai-dynamo/dynamo/deploy/operator/internal/gms"
 	snapshotprotocol "github.com/ai-dynamo/dynamo/deploy/snapshot/protocol"
-	appsv1 "k8s.io/api/apps/v1"
 	corev1 "k8s.io/api/core/v1"
-	ctrlclient "sigs.k8s.io/controller-runtime/pkg/client"
+	"k8s.io/utils/ptr"
 )
 
 const (
@@ -23,101 +21,80 @@ const (
 
 	gmsCheckpointLoaderModule = "gpu_memory_service.cli.snapshot.loader"
 	gmsCheckpointSaverModule  = "gpu_memory_service.cli.snapshot.saver"
-)
-
-func ResolveGMSCheckpointStorage(
-	ctx context.Context,
-	reader ctrlclient.Reader,
-	namespace string,
-	checkpointID string,
-	artifactVersion string,
-) (snapshotprotocol.Storage, error) {
-	if reader == nil {
-		return snapshotprotocol.Storage{}, fmt.Errorf("checkpoint client is required")
-	}
-
-	daemonSets := &appsv1.DaemonSetList{}
-	if err := reader.List(
-		ctx,
-		daemonSets,
-		ctrlclient.InNamespace(namespace),
-		ctrlclient.MatchingLabels{snapshotprotocol.SnapshotAgentLabelKey: snapshotprotocol.SnapshotAgentLabelValue},
-	); err != nil {
-		return snapshotprotocol.Storage{}, fmt.Errorf("list snapshot-agent daemonsets in %s: %w", namespace, err)
-	}
 
-	storage, err := snapshotprotocol.DiscoverStorageFromDaemonSets(namespace, daemonSets.Items)
-	if err != nil {
-		return snapshotprotocol.Storage{}, err
-	}
-	return snapshotprotocol.ResolveCheckpointStorage(checkpointID, artifactVersion, storage)
-}
+	// envCheckpointDir is the environment variable name for the GMS
+	// checkpoint artifact directory on the snapshot PVC.
+	envCheckpointDir = "GMS_CHECKPOINT_DIR"
+)
 
-// BuildGMSRestoreSidecars prepares GMS infrastructure for a restore pod and
-// returns the additional containers the caller must append to podSpec.Containers.
-//
-// The GMS server runs as a regular container (not init) because the CRIU-restored
-// main process already has GPU memory mapped and does not need sockets at
-// startup. The gms-loader polls for sockets internally via wait_for_weights_socket.
-func BuildGMSRestoreSidecars(
+// EnsureGMSRestoreSidecars adds GMS server + loader containers to the pod spec
+// for a checkpoint restore. The server runs as a regular container (not init)
+// because the CRIU-restored main process already has GPU memory mapped and
+// all containers must start in parallel.
+func EnsureGMSRestoreSidecars(
 	podSpec *corev1.PodSpec,
 	mainContainer *corev1.Container,
 	storage snapshotprotocol.Storage,
-) []corev1.Container {
+) {
 	if podSpec == nil || mainContainer == nil {
-		return nil
+		return
 	}
 
-	// Remove gms-server from initContainers if the DGD-level
-	// applyGPUMemoryService already placed it there. For restore pods the
-	// server runs as a regular container so that all containers start in
-	// parallel — the restored main process does not need sockets at startup.
+	// The DGD path adds the GMS server as an init sidecar (blocks until
+	// sockets are ready). For restore, move it to a regular container so
+	// all containers start in parallel.
 	for i := range podSpec.InitContainers {
-		if podSpec.InitContainers[i].Name == gmsruntime.ServerContainerName {
+		if podSpec.InitContainers[i].Name == gms.ServerContainerName {
 			podSpec.InitContainers = append(podSpec.InitContainers[:i], podSpec.InitContainers[i+1:]...)
 			break
 		}
 	}
+	gms.EnsureSharedVolume(podSpec, mainContainer)
+
+	snapshotprotocol.InjectCheckpointVolume(podSpec, storage.PVCName)
 
-	server := gmsruntime.BuildServerContainer(podSpec, mainContainer)
+	server := gms.Container(gms.ServerContainerName, gms.ServerModule, mainContainer.Image)
+	server.RestartPolicy = ptr.To(corev1.ContainerRestartPolicyAlways)
 
-	loader := gmsCheckpointLoaderContainer(mainContainer.Image)
-	copyGMSDeviceClaims(mainContainer, &loader)
-	ensureCheckpointVolume(podSpec, storage.PVCName)
+	loader := gms.Container(GMSLoaderContainer, gmsCheckpointLoaderModule, mainContainer.Image)
 	loader.VolumeMounts = append(loader.VolumeMounts, corev1.VolumeMount{Name: snapshotprotocol.CheckpointVolumeName, MountPath: storage.BasePath})
-	loader.Env = append(loader.Env, corev1.EnvVar{Name: "GMS_CHECKPOINT_DIR", Value: resolveGMSArtifactDir(storage)})
+	loader.Env = append(loader.Env, corev1.EnvVar{Name: envCheckpointDir, Value: resolveGMSArtifactDir(storage)})
+	loader.RestartPolicy = ptr.To(corev1.ContainerRestartPolicyAlways)
 
-	return []corev1.Container{server, loader}
+	podSpec.InitContainers = append(podSpec.InitContainers, server, loader)
 }
 
-// BuildGMSCheckpointJobSidecars prepares GMS infrastructure for a checkpoint
-// job and returns the additional containers the caller must append to
-// podSpec.Containers.
-func BuildGMSCheckpointJobSidecars(
+// EnsureGMSCheckpointJobSidecars adds GMS server (init) + saver containers
+// to the pod spec for a checkpoint job.
+func EnsureGMSCheckpointJobSidecars(
 	podSpec *corev1.PodSpec,
 	mainContainer *corev1.Container,
 	storage snapshotprotocol.Storage,
-) ([]corev1.Container, error) {
+) error {
 	if podSpec == nil || mainContainer == nil {
-		return nil, nil
+		return nil
 	}
 	if len(mainContainer.Resources.Claims) == 0 {
-		return nil, fmt.Errorf("gms sidecars require main container resource claims")
+		return fmt.Errorf("gms sidecars require main container resource claims (DRA must be enabled)")
 	}
 	if storage.PVCName == "" || storage.BasePath == "" || storage.Location == "" {
-		return nil, fmt.Errorf("gms checkpoint jobs require resolved checkpoint storage")
+		return fmt.Errorf("gms checkpoint jobs require resolved checkpoint storage")
 	}
 
-	gmsruntime.EnsureServerSidecar(podSpec, mainContainer)
-	ensureGMSCheckpointControl(podSpec)
+	gmsArtifactDir := resolveGMSArtifactDir(storage)
 
-	saver := gmsCheckpointSaverContainer(mainContainer.Image)
-	copyGMSDeviceClaims(mainContainer, &saver)
-	ensureCheckpointVolume(podSpec, storage.PVCName)
-	saver.VolumeMounts = append(saver.VolumeMounts, corev1.VolumeMount{Name: snapshotprotocol.CheckpointVolumeName, MountPath: storage.BasePath})
-	saver.Env = append(saver.Env, corev1.EnvVar{Name: "GMS_CHECKPOINT_DIR", Value: resolveGMSArtifactDir(storage)})
+	gms.EnsureServerSidecar(podSpec, mainContainer)
+	snapshotprotocol.InjectCheckpointVolume(podSpec, storage.PVCName)
 
-	return []corev1.Container{saver}, nil
+	saver := gms.Container(GMSSaverContainer, gmsCheckpointSaverModule, mainContainer.Image)
+	saver.VolumeMounts = append(saver.VolumeMounts, corev1.VolumeMount{Name: snapshotprotocol.CheckpointVolumeName, MountPath: storage.BasePath})
+	saver.Env = append(saver.Env, corev1.EnvVar{Name: envCheckpointDir, Value: gmsArtifactDir})
+	// The saver is an init sidecar (restartPolicy=Always) so it doesn't
+	// affect pod Ready (only the worker's probe matters) and doesn't block
+	// Job completion. It saves, then sleeps until the pod terminates.
+	saver.RestartPolicy = ptr.To(corev1.ContainerRestartPolicyAlways)
+	podSpec.InitContainers = append(podSpec.InitContainers, saver)
+	return nil
 }
 
 func resolveGMSArtifactDir(storage snapshotprotocol.Storage) string {
@@ -128,77 +105,3 @@ func resolveGMSArtifactDir(storage snapshotprotocol.Storage) string {
 	checkpointID := filepath.Base(filepath.Dir(filepath.Dir(storage.Location)))
 	return filepath.Join(storage.BasePath, "gms", checkpointID, "versions", artifactVersion)
 }
-
-func gmsCheckpointLoaderContainer(image string) corev1.Container {
-	container := corev1.Container{
-		Name:    GMSLoaderContainer,
-		Image:   image,
-		Command: []string{"python3", "-m", gmsCheckpointLoaderModule},
-		Env: []corev1.EnvVar{
-			{Name: "TMPDIR", Value: gmsruntime.SharedMountPath},
-			{Name: "GMS_SOCKET_DIR", Value: gmsruntime.SharedMountPath},
-		},
-		VolumeMounts: []corev1.VolumeMount{
-			{Name: gmsruntime.SharedVolumeName, MountPath: gmsruntime.SharedMountPath},
-		},
-	}
-	return container
-}
-
-func gmsCheckpointSaverContainer(image string) corev1.Container {
-	container := corev1.Container{
-		Name:    GMSSaverContainer,
-		Image:   image,
-		Command: []string{"python3", "-m", gmsCheckpointSaverModule},
-		Env: []corev1.EnvVar{
-			{Name: "POD_NAME", ValueFrom: &corev1.EnvVarSource{FieldRef: &corev1.ObjectFieldSelector{FieldPath: "metadata.name"}}},
-			{Name: "POD_NAMESPACE", ValueFrom: &corev1.EnvVarSource{FieldRef: &corev1.ObjectFieldSelector{FieldPath: "metadata.namespace"}}},
-			{Name: "TMPDIR", Value: gmsruntime.SharedMountPath},
-			{Name: "GMS_SOCKET_DIR", Value: gmsruntime.SharedMountPath},
-			{Name: "GMS_CONTROL_DIR", Value: gmsruntime.ControlDir},
-		},
-		VolumeMounts: []corev1.VolumeMount{
-			{Name: gmsruntime.SharedVolumeName, MountPath: gmsruntime.SharedMountPath},
-			{Name: gmsruntime.ControlVolumeName, MountPath: gmsruntime.ControlDir},
-		},
-	}
-	return container
-}
-
-// ensureGMSCheckpointControl adds the control volume and injects
-// GMS_CONTROL_DIR into the GMS server container for checkpoint coordination.
-func ensureGMSCheckpointControl(podSpec *corev1.PodSpec) {
-	podSpec.Volumes = append(podSpec.Volumes, corev1.Volume{
-		Name:         gmsruntime.ControlVolumeName,
-		VolumeSource: corev1.VolumeSource{EmptyDir: &corev1.EmptyDirVolumeSource{}},
-	})
-	server := gmsruntime.FindServerContainer(podSpec)
-	if server != nil {
-		server.VolumeMounts = append(server.VolumeMounts, corev1.VolumeMount{Name: gmsruntime.ControlVolumeName, MountPath: gmsruntime.ControlDir})
-		server.Env = append(server.Env, corev1.EnvVar{Name: "GMS_CONTROL_DIR", Value: gmsruntime.ControlDir})
-	}
-}
-
-func copyGMSDeviceClaims(mainContainer *corev1.Container, container *corev1.Container) {
-	if mainContainer == nil || container == nil || len(mainContainer.Resources.Claims) == 0 {
-		return
-	}
-	container.Resources.Claims = append([]corev1.ResourceClaim{}, mainContainer.Resources.Claims...)
-}
-
-func ensureCheckpointVolume(podSpec *corev1.PodSpec, pvcName string) {
-	if pvcName == "" {
-		return
-	}
-	for i := range podSpec.Volumes {
-		if podSpec.Volumes[i].Name == snapshotprotocol.CheckpointVolumeName {
-			return
-		}
-	}
-	podSpec.Volumes = append(podSpec.Volumes, corev1.Volume{
-		Name: snapshotprotocol.CheckpointVolumeName,
-		VolumeSource: corev1.VolumeSource{
-			PersistentVolumeClaim: &corev1.PersistentVolumeClaimVolumeSource{ClaimName: pvcName},
-		},
-	})
-}

deploy/operator/internal/checkpoint/podinfo.go

@@ -9,42 +9,17 @@ import (
 )
 
 func EnsurePodInfoVolume(podSpec *corev1.PodSpec) {
-	for i := range podSpec.Volumes {
-		if podSpec.Volumes[i].Name != commonconsts.PodInfoVolumeName {
-			continue
-		}
-		if podSpec.Volumes[i].DownwardAPI == nil {
-			podSpec.Volumes[i].VolumeSource.DownwardAPI = &corev1.DownwardAPIVolumeSource{}
-		}
-		// Merge required items into existing downwardAPI volume.
-		source := podSpec.Volumes[i].DownwardAPI
-		pathToIndex := make(map[string]int, len(source.Items))
-		for j := range source.Items {
-			pathToIndex[source.Items[j].Path] = j
-		}
-		for _, item := range podInfoItems() {
-			if idx, ok := pathToIndex[item.Path]; ok {
-				source.Items[idx] = item
-				continue
-			}
-			source.Items = append(source.Items, item)
-			pathToIndex[item.Path] = len(source.Items) - 1
-		}
+	for _, volume := range podSpec.Volumes {
+		if volume.Name == commonconsts.PodInfoVolumeName {
 			return
 		}
+	}
 
 	podSpec.Volumes = append(podSpec.Volumes, corev1.Volume{
 		Name: commonconsts.PodInfoVolumeName,
 		VolumeSource: corev1.VolumeSource{
 			DownwardAPI: &corev1.DownwardAPIVolumeSource{
-				Items: podInfoItems(),
-			},
-		},
-	})
-}
-
-func podInfoItems() []corev1.DownwardAPIVolumeFile {
-	return []corev1.DownwardAPIVolumeFile{
+				Items: []corev1.DownwardAPIVolumeFile{
 					{
 						Path: "pod_name",
 						FieldRef: &corev1.ObjectFieldSelector{
@@ -93,7 +68,10 @@ func podInfoItems() []corev1.DownwardAPIVolumeFile {
 							FieldPath: commonconsts.PodInfoFieldPodNamespace,
 						},
 					},
-	}
+				},
+			},
+		},
+	})
 }
 
 func EnsurePodInfoMount(container *corev1.Container) {

deploy/operator/internal/checkpoint/podspec.go

@@ -38,6 +38,19 @@ func ApplyRestorePodMetadata(labels map[string]string, annotations map[string]st
 	snapshotprotocol.ApplyRestoreTargetMetadata(labels, annotations, enabled, hash, artifactVersion)
 }
 
+// resolveMainContainer finds the container named "main" in the pod spec.
+// ExtraPodSpec.PodSpec.Containers can inject user containers before the main
+// container (mergo merge happens before main is appended), so index 0 is
+// not guaranteed to be the main container here.
+func resolveMainContainer(podSpec *corev1.PodSpec) *corev1.Container {
+	for i := range podSpec.Containers {
+		if podSpec.Containers[i].Name == commonconsts.MainContainerName {
+			return &podSpec.Containers[i]
+		}
+	}
+	return nil
+}
+
 func InjectCheckpointIntoPodSpec(
 	ctx context.Context,
 	reader ctrlclient.Reader,
@@ -62,18 +75,9 @@ func InjectCheckpointIntoPodSpec(
 		info.Hash = hash
 	}
 
-	if len(podSpec.Containers) == 0 {
-		return fmt.Errorf("no container found to inject checkpoint config")
-	}
-	var mainContainer *corev1.Container
-	for i := range podSpec.Containers {
-		if podSpec.Containers[i].Name == commonconsts.MainContainerName {
-			mainContainer = &podSpec.Containers[i]
-			break
-		}
-	}
+	mainContainer := resolveMainContainer(podSpec)
 	if mainContainer == nil {
-		return fmt.Errorf("main container not found in pod spec")
+		return fmt.Errorf("no container named %q found in pod spec", commonconsts.MainContainerName)
 	}
 	if reader == nil {
 		return fmt.Errorf("checkpoint client is required")
@@ -94,14 +98,8 @@ func InjectCheckpointIntoPodSpec(
 
 	EnsurePodInfoVolume(podSpec)
 	EnsurePodInfoMount(mainContainer)
-
-	// GMS restore sidecars (server + loader) are only needed when the checkpoint
-	// is ready and the pod will actually be CRIU-restored.
 	if info.Ready && info.GPUMemoryService != nil && info.GPUMemoryService.Enabled {
-		if len(mainContainer.Resources.Claims) == 0 {
-			return fmt.Errorf("gms sidecars require main container resource claims")
-		}
-		storage, err := ResolveGMSCheckpointStorage(
+		storage, err := snapshotprotocol.DiscoverAndResolveStorage(
 			ctx,
 			reader,
 			namespace,
@@ -111,8 +109,7 @@ func InjectCheckpointIntoPodSpec(
 		if err != nil {
 			return err
 		}
-		gmsSidecars := BuildGMSRestoreSidecars(podSpec, mainContainer, storage)
-		podSpec.Containers = append(podSpec.Containers, gmsSidecars...)
+		EnsureGMSRestoreSidecars(podSpec, mainContainer, storage)
 	}
 
 	return nil

deploy/operator/internal/controller/checkpoint_job.go

@@ -6,17 +6,18 @@ package controller
 import (
 	"context"
 	"fmt"
+	"strings"
 
 	configv1alpha1 "github.com/ai-dynamo/dynamo/deploy/operator/api/config/v1alpha1"
 	nvidiacomv1alpha1 "github.com/ai-dynamo/dynamo/deploy/operator/api/v1alpha1"
 	"github.com/ai-dynamo/dynamo/deploy/operator/internal/checkpoint"
 	"github.com/ai-dynamo/dynamo/deploy/operator/internal/consts"
 	"github.com/ai-dynamo/dynamo/deploy/operator/internal/discovery"
+	"github.com/ai-dynamo/dynamo/deploy/operator/internal/dra"
 	"github.com/ai-dynamo/dynamo/deploy/operator/internal/dynamo"
 	snapshotprotocol "github.com/ai-dynamo/dynamo/deploy/snapshot/protocol"
 	batchv1 "k8s.io/api/batch/v1"
 	corev1 "k8s.io/api/core/v1"
-	"k8s.io/apimachinery/pkg/api/resource"
 	ctrlclient "sigs.k8s.io/controller-runtime/pkg/client"
 )
 
@@ -81,10 +82,10 @@ func buildCheckpointJob(
 
 	checkpoint.EnsurePodInfoVolume(&podTemplate.Spec)
 
-	mainContainer, err := snapshotprotocol.ResolveCheckpointWorkerContainer(&podTemplate.Spec)
-	if err != nil {
-		return nil, err
+	if len(podTemplate.Spec.Containers) == 0 {
+		return nil, fmt.Errorf("checkpoint job requires at least one container")
 	}
+	mainContainer := &podTemplate.Spec.Containers[0]
 	mainContainer.Env = dynamo.MergeEnvs(
 		buildCheckpointWorkerDefaultEnv(ckpt, podTemplate),
 		mainContainer.Env,
@@ -105,12 +106,17 @@ func buildCheckpointJob(
 	}
 	mainContainer.LivenessProbe = nil
 	mainContainer.StartupProbe = nil
+
+	// The snapshot agent sends SIGUSR1 to PID 1 of the main container after
 	checkpoint.EnsurePodInfoMount(mainContainer)
 	dynamo.ApplySharedMemoryVolumeAndMount(&podTemplate.Spec, mainContainer, ckpt.Spec.Job.SharedMemory)
 
-	var gmsSidecars []corev1.Container
 	if ckpt.Spec.GPUMemoryService != nil && ckpt.Spec.GPUMemoryService.Enabled {
-		storage, err := checkpoint.ResolveGMSCheckpointStorage(
+		claimTemplateName := dra.ResourceClaimTemplateName("checkpoint-"+hash, "worker")
+		if err := dra.ApplyClaim(&podTemplate.Spec, claimTemplateName); err != nil {
+			return nil, fmt.Errorf("failed to apply DRA claim for GMS checkpoint: %w", err)
+		}
+		storage, err := snapshotprotocol.DiscoverAndResolveStorage(
 			ctx,
 			reader,
 			ckpt.Namespace,
@@ -120,12 +126,13 @@ func buildCheckpointJob(
 		if err != nil {
 			return nil, err
 		}
-		gmsSidecars, err = checkpoint.BuildGMSCheckpointJobSidecars(&podTemplate.Spec, mainContainer, storage)
-		if err != nil {
+		if err := checkpoint.EnsureGMSCheckpointJobSidecars(&podTemplate.Spec, mainContainer, storage); err != nil {
 			return nil, err
 		}
+		// Re-acquire pointer: append in EnsureGMSCheckpointJobSidecars may
+		// have reallocated the Containers slice.
+		mainContainer = &podTemplate.Spec.Containers[0]
 	}
-	podTemplate.Spec.Containers = append(podTemplate.Spec.Containers, gmsSidecars...)
 
 	activeDeadlineSeconds := ckpt.Spec.Job.ActiveDeadlineSeconds
 	if activeDeadlineSeconds == nil {
@@ -133,10 +140,29 @@ func buildCheckpointJob(
 		activeDeadlineSeconds = &defaultDeadline
 	}
 
-	wrapLaunchJob := false
-	if gpus, ok := mainContainer.Resources.Limits[corev1.ResourceName(consts.KubeResourceGPUNvidia)]; ok {
-		wrapLaunchJob = gpus.Cmp(*resource.NewQuantity(1, resource.DecimalSI)) > 0
+	// Wrap with cuda-checkpoint --launch-job for multi-GPU jobs (TP*PP > 1).
+	// Use checkpoint identity (not container limits) because DRA may have
+	// already removed nvidia.com/gpu from the template.
+	tp := ckpt.Spec.Identity.TensorParallelSize
+	pp := ckpt.Spec.Identity.PipelineParallelSize
+	if tp == 0 {
+		tp = 1
 	}
+	if pp == 0 {
+		pp = 1
+	}
+	wrapLaunchJob := tp*pp > 1
+
+	// For single-GPU jobs (no cuda-checkpoint wrapper), unwrap /bin/sh -c so
+	// the actual process is PID 1 and receives SIGUSR1 from the snapshot agent.
+	if !wrapLaunchJob && len(mainContainer.Command) >= 2 &&
+		mainContainer.Command[len(mainContainer.Command)-1] == "-c" &&
+		len(mainContainer.Args) == 1 {
+		parts := strings.Fields(mainContainer.Args[0])
+		mainContainer.Command = parts[:1]
+		mainContainer.Args = parts[1:]
+	}
+
 	ttlSecondsAfterFinish := snapshotprotocol.DefaultCheckpointJobTTLSeconds
 
 	return snapshotprotocol.NewCheckpointJob(podTemplate, snapshotprotocol.CheckpointJobOptions{

deploy/operator/internal/controller/dynamocheckpoint_controller.go

@@ -26,6 +26,7 @@ import (
 	coordinationv1 "k8s.io/api/coordination/v1"
 	corev1 "k8s.io/api/core/v1"
 	rbacv1 "k8s.io/api/rbac/v1"
+	resourcev1 "k8s.io/api/resource/v1"
 	apierrors "k8s.io/apimachinery/pkg/api/errors"
 	"k8s.io/apimachinery/pkg/api/meta"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
@@ -40,8 +41,10 @@ import (
 	configv1alpha1 "github.com/ai-dynamo/dynamo/deploy/operator/api/config/v1alpha1"
 	nvidiacomv1alpha1 "github.com/ai-dynamo/dynamo/deploy/operator/api/v1alpha1"
 	"github.com/ai-dynamo/dynamo/deploy/operator/internal/checkpoint"
+	"github.com/ai-dynamo/dynamo/deploy/operator/internal/consts"
 	commonController "github.com/ai-dynamo/dynamo/deploy/operator/internal/controller_common"
 	"github.com/ai-dynamo/dynamo/deploy/operator/internal/discovery"
+	"github.com/ai-dynamo/dynamo/deploy/operator/internal/dra"
 	snapshotprotocol "github.com/ai-dynamo/dynamo/deploy/snapshot/protocol"
 )
 
@@ -62,6 +65,8 @@ func (r *CheckpointReconciler) GetRecorder() record.EventRecorder {
 // +kubebuilder:rbac:groups=nvidia.com,resources=dynamocheckpoints/status,verbs=get;update;patch
 // +kubebuilder:rbac:groups=nvidia.com,resources=dynamocheckpoints/finalizers,verbs=update
 // +kubebuilder:rbac:groups=batch,resources=jobs,verbs=get;list;watch;create;update;patch;delete
+// +kubebuilder:rbac:groups=resource.k8s.io,resources=resourceclaimtemplates,verbs=get;list;watch;create;update;patch;delete
+// +kubebuilder:rbac:groups=resource.k8s.io,resources=deviceclasses,verbs=get
 // +kubebuilder:rbac:groups=coordination.k8s.io,resources=leases,verbs=get;list;watch
 
 func (r *CheckpointReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) {
@@ -190,6 +195,32 @@ func (r *CheckpointReconciler) handlePending(ctx context.Context, ckpt *nvidiaco
 			return ctrl.Result{}, fmt.Errorf("failed to compute checkpoint identity hash: %w", err)
 		}
 	}
+
+	// Sync DRA ResourceClaimTemplate for GMS-enabled checkpoints.
+	if ckpt.Spec.GPUMemoryService != nil && ckpt.Spec.GPUMemoryService.Enabled {
+		if !r.RuntimeConfig.DRAEnabled {
+			return ctrl.Result{}, fmt.Errorf(
+				"GMS requires DRA (Dynamic Resource Allocation), but the resource.k8s.io API group is not available")
+		}
+		if len(ckpt.Spec.Job.PodTemplateSpec.Spec.Containers) == 0 {
+			return ctrl.Result{}, fmt.Errorf("checkpoint job requires at least one container for GMS")
+		}
+		gpuQty := ckpt.Spec.Job.PodTemplateSpec.Spec.Containers[0].Resources.Limits[corev1.ResourceName(consts.KubeResourceGPUNvidia)]
+		gpuCount := int(gpuQty.Value())
+		deviceClassName := ""
+		if ckpt.Spec.GPUMemoryService != nil {
+			deviceClassName = ckpt.Spec.GPUMemoryService.DeviceClassName
+		}
+		claimTemplateName := dra.ResourceClaimTemplateName("checkpoint-"+hash, "worker")
+		_, _, err := commonController.SyncResource(ctx, r, ckpt, func(ctx context.Context) (*resourcev1.ResourceClaimTemplate, bool, error) {
+			return dra.GenerateResourceClaimTemplate(ctx, r.Client, claimTemplateName, ckpt.Namespace, gpuCount, deviceClassName)
+		})
+		if err != nil {
+			logger.Error(err, "Failed to sync GMS ResourceClaimTemplate for checkpoint")
+			return ctrl.Result{}, fmt.Errorf("failed to sync GMS ResourceClaimTemplate for checkpoint: %w", err)
+		}
+	}
+
 	jobName := snapshotprotocol.GetCheckpointJobName(
 		hash,
 		ckpt.Annotations[snapshotprotocol.CheckpointArtifactVersionAnnotation],

deploy/operator/internal/controller/dynamocheckpoint_controller_test.go

@@ -26,7 +26,7 @@ import (
 	nvidiacomv1alpha1 "github.com/ai-dynamo/dynamo/deploy/operator/api/v1alpha1"
 	"github.com/ai-dynamo/dynamo/deploy/operator/internal/checkpoint"
 	"github.com/ai-dynamo/dynamo/deploy/operator/internal/consts"
-	gmsruntime "github.com/ai-dynamo/dynamo/deploy/operator/internal/gms"
+	gms "github.com/ai-dynamo/dynamo/deploy/operator/internal/gms"
 	snapshotprotocol "github.com/ai-dynamo/dynamo/deploy/snapshot/protocol"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -256,6 +256,8 @@ func TestBuildCheckpointJob(t *testing.T) {
 	assert.Equal(t, int32(0), *job.Spec.BackoffLimit)
 	assert.Equal(t, int32(300), *job.Spec.TTLSecondsAfterFinished)
 
+	// Multi-GPU: wrapping decision uses identity.TensorParallelSize, not container GPU limits.
+	ckpt.Spec.Identity.TensorParallelSize = 2
 	ckpt.Spec.Job.PodTemplateSpec.Spec.Containers[0].Resources = corev1.ResourceRequirements{
 		Limits: corev1.ResourceList{
 			corev1.ResourceName("nvidia.com/gpu"): resource.MustParse("2"),
@@ -267,16 +269,11 @@ func TestBuildCheckpointJob(t *testing.T) {
 	assert.Equal(t, []string{"--launch-job", "python3", "-m", "dynamo.vllm"}, job.Spec.Template.Spec.Containers[0].Args)
 }
 
-func TestBuildCheckpointJobTargetsMainContainerWhenSidecarIsFirst(t *testing.T) {
+func TestBuildCheckpointJobWrapsWithCudaCheckpointForMultiGPU(t *testing.T) {
 	s := checkpointTestScheme()
 	ckpt := makeTestCheckpoint(nvidiacomv1alpha1.DynamoCheckpointPhasePending)
+	ckpt.Spec.Identity.TensorParallelSize = 2
 	ckpt.Spec.Job.PodTemplateSpec.Spec.Containers = []corev1.Container{
-		{
-			Name:    "sidecar",
-			Image:   "sidecar:latest",
-			Command: []string{"sleep"},
-			Args:    []string{"infinity"},
-		},
 		{
 			Name:    consts.MainContainerName,
 			Image:   "test-image:latest",
@@ -288,13 +285,19 @@ func TestBuildCheckpointJobTargetsMainContainerWhenSidecarIsFirst(t *testing.T)
 				},
 			},
 		},
+		{
+			Name:    "sidecar",
+			Image:   "sidecar:latest",
+			Command: []string{"sleep"},
+			Args:    []string{"infinity"},
+		},
 	}
 
 	r := makeCheckpointReconciler(s, ckpt)
 	job, err := buildCheckpointJob(context.Background(), nil, r.Config, ckpt, defaultCheckpointJobName)
 	require.NoError(t, err)
 
-	main := requireCheckpointContainer(t, job.Spec.Template.Spec.Containers, consts.MainContainerName)
+	main := &job.Spec.Template.Spec.Containers[0]
 	assert.Equal(t, []string{"cuda-checkpoint"}, main.Command)
 	assert.Equal(t, []string{"--launch-job", "python3", "-m", "dynamo.vllm"}, main.Args)
 	require.NotNil(t, main.ReadinessProbe)
@@ -362,22 +365,22 @@ func TestBuildCheckpointJobAddsGMSSidecars(t *testing.T) {
 	require.NoError(t, err)
 
 	main := requireCheckpointContainer(t, job.Spec.Template.Spec.Containers, consts.MainContainerName)
-	weightsServer := requireCheckpointContainer(t, job.Spec.Template.Spec.InitContainers, gmsruntime.ServerContainerName)
-	saver := requireCheckpointContainer(t, job.Spec.Template.Spec.Containers, checkpoint.GMSSaverContainer)
+	weightsServer := requireCheckpointContainer(t, job.Spec.Template.Spec.InitContainers, gms.ServerContainerName)
+	saver := requireCheckpointContainer(t, job.Spec.Template.Spec.InitContainers, checkpoint.GMSSaverContainer)
 
 	volNames := map[string]bool{}
 	for _, v := range job.Spec.Template.Spec.Volumes {
 		volNames[v.Name] = true
 	}
-	assert.True(t, volNames[gmsruntime.SharedVolumeName])
-	assert.True(t, volNames[gmsruntime.ControlVolumeName])
+	assert.True(t, volNames[gms.SharedVolumeName])
+	assert.True(t, volNames[snapshotprotocol.CheckpointVolumeName])
 	assert.True(t, volNames[snapshotprotocol.CheckpointVolumeName])
 
 	mainMounts := map[string]string{}
 	for _, m := range main.VolumeMounts {
 		mainMounts[m.Name] = m.MountPath
 	}
-	assert.Equal(t, gmsruntime.SharedMountPath, mainMounts[gmsruntime.SharedVolumeName])
+	assert.Equal(t, gms.SharedMountPath, mainMounts[gms.SharedVolumeName])
 
 	assert.Equal(t, []string{"python3", "-m", "gpu_memory_service.cli.server"}, weightsServer.Command)
 	assert.Equal(t, corev1.ContainerRestartPolicyAlways, *weightsServer.RestartPolicy)

deploy/operator/internal/controller/dynamocomponentdeployment_controller.go

@@ -40,6 +40,7 @@ import (
 	"github.com/ai-dynamo/dynamo/deploy/operator/internal/common"
 	commonconsts "github.com/ai-dynamo/dynamo/deploy/operator/internal/consts"
 	commonController "github.com/ai-dynamo/dynamo/deploy/operator/internal/controller_common"
+	"github.com/ai-dynamo/dynamo/deploy/operator/internal/dra"
 	"github.com/ai-dynamo/dynamo/deploy/operator/internal/dynamo"
 	"github.com/ai-dynamo/dynamo/deploy/operator/internal/observability"
 	snapshotprotocol "github.com/ai-dynamo/dynamo/deploy/snapshot/protocol"
@@ -190,9 +191,11 @@ func (r *DynamoComponentDeploymentReconciler) Reconcile(ctx context.Context, req
 		if serviceName == "" {
 			serviceName = dynamoComponentDeployment.Name
 		}
-		claimTemplateName := dynamo.GMSResourceClaimTemplateName(dynamoComponentDeployment.GetParentGraphDeploymentName(), serviceName)
+		spec := &dynamoComponentDeployment.Spec.DynamoComponentDeploymentSharedSpec
+		gpuCount, deviceClassName := dra.ExtractGPUParams(spec.GPUMemoryService, spec.Resources)
+		claimTemplateName := dra.ResourceClaimTemplateName(dynamoComponentDeployment.GetParentGraphDeploymentName(), serviceName)
 		_, _, err = commonController.SyncResource(ctx, r, dynamoComponentDeployment, func(ctx context.Context) (*resourcev1.ResourceClaimTemplate, bool, error) {
-			return dynamo.GenerateGMSResourceClaimTemplate(ctx, r.Client, claimTemplateName, dynamoComponentDeployment.Namespace, &dynamoComponentDeployment.Spec.DynamoComponentDeploymentSharedSpec)
+			return dra.GenerateResourceClaimTemplate(ctx, r.Client, claimTemplateName, dynamoComponentDeployment.Namespace, gpuCount, deviceClassName)
 		})
 		if err != nil {
 			return ctrl.Result{}, fmt.Errorf("failed to sync GMS ResourceClaimTemplate: %w", err)

deploy/operator/internal/controller/dynamocomponentdeployment_controller_test.go

@@ -29,7 +29,7 @@ import (
 	commonconsts "github.com/ai-dynamo/dynamo/deploy/operator/internal/consts"
 	"github.com/ai-dynamo/dynamo/deploy/operator/internal/controller_common"
 	"github.com/ai-dynamo/dynamo/deploy/operator/internal/dynamo"
-	gmsruntime "github.com/ai-dynamo/dynamo/deploy/operator/internal/gms"
+	gms "github.com/ai-dynamo/dynamo/deploy/operator/internal/gms"
 	snapshotprotocol "github.com/ai-dynamo/dynamo/deploy/snapshot/protocol"
 	"github.com/google/go-cmp/cmp"
 	"github.com/onsi/gomega"
@@ -1424,7 +1424,7 @@ func TestDynamoComponentDeploymentReconciler_generatePodTemplateSpec_RestoreLabe
 			return nil
 		}
 
-		gmsServer := find(gmsruntime.ServerContainerName)
+		gmsServer := find(gms.ServerContainerName)
 		require.NotNil(t, gmsServer)
 		loader := find(checkpoint.GMSLoaderContainer)
 		require.NotNil(t, loader)
@@ -1439,9 +1439,9 @@ func TestDynamoComponentDeploymentReconciler_generatePodTemplateSpec_RestoreLabe
 		if got := gmsServer.Command; len(got) != 3 || got[0] != "python3" || got[1] != "-m" || got[2] != "gpu_memory_service.cli.server" { //nolint:goconst
 			t.Fatalf("expected weights server to run python module, got %#v", got)
 		}
-		// Restore: gms-server should be a regular container, not an init container
-		if gmsServer.RestartPolicy != nil {
-			t.Fatalf("expected restore gms-server to have no RestartPolicy (regular container), got %#v", gmsServer.RestartPolicy)
+		// Restore: gms-server and loader are init sidecars (restartPolicy=Always)
+		if gmsServer.RestartPolicy == nil || *gmsServer.RestartPolicy != corev1.ContainerRestartPolicyAlways {
+			t.Fatalf("expected restore gms-server to have RestartPolicy=Always, got %#v", gmsServer.RestartPolicy)
 		}
 		if gmsServer.StartupProbe != nil {
 			t.Fatalf("expected restore gms-server to have no StartupProbe")
@@ -1487,8 +1487,11 @@ func TestDynamoComponentDeploymentReconciler_generatePodTemplateSpec_RestoreLabe
 			t.Fatalf("generatePodTemplateSpec failed: %v", err)
 		}
 
+		// User's extra sidecar should remain in Containers, unchanged.
+		// GMS loader is now an init sidecar, so the user's container stays
+		// at Containers[0] and main at Containers[1].
 		if got := podTemplateSpec.Spec.Containers[0]; got.Name != "gms-loader" || len(got.Command) != 1 || got.Command[0] != "python3" {
-			t.Fatalf("expected sidecar container to remain unchanged, got %#v", got)
+			t.Fatalf("expected user sidecar container to remain unchanged, got %#v", got)
 		}
 		if got := podTemplateSpec.Spec.Containers[1]; got.Name != commonconsts.MainContainerName || len(got.Command) != 2 || got.Command[0] != "sleep" || got.Command[1] != "infinity" {
 			t.Fatalf("expected main container to be rewritten for restore, got %#v", got)

deploy/operator/internal/controller/dynamographdeployment_controller.go

@@ -54,6 +54,7 @@ import (
 	nvidiacomv1alpha1 "github.com/ai-dynamo/dynamo/deploy/operator/api/v1alpha1"
 	"github.com/ai-dynamo/dynamo/deploy/operator/internal/consts"
 	commoncontroller "github.com/ai-dynamo/dynamo/deploy/operator/internal/controller_common"
+	"github.com/ai-dynamo/dynamo/deploy/operator/internal/dra"
 	"github.com/ai-dynamo/dynamo/deploy/operator/internal/dynamo"
 	"github.com/ai-dynamo/dynamo/deploy/operator/internal/dynamo/epp"
 	"github.com/ai-dynamo/dynamo/deploy/operator/internal/observability"
@@ -664,10 +665,10 @@ func (r *DynamoGraphDeploymentReconciler) reconcileGroveResources(ctx context.Co
 	// Sync ResourceClaimTemplates for GMS-enabled components before creating pods.
 	if r.RuntimeConfig.DRAEnabled {
 		for serviceName, component := range dynamoDeployment.Spec.Services {
-			svcComponent := component
-			claimTemplateName := dynamo.GMSResourceClaimTemplateName(dynamoDeployment.Name, serviceName)
+			gpuCount, deviceClassName := dra.ExtractGPUParams(component.GPUMemoryService, component.Resources)
+			claimTemplateName := dra.ResourceClaimTemplateName(dynamoDeployment.Name, serviceName)
 			_, _, err := commoncontroller.SyncResource(ctx, r, dynamoDeployment, func(ctx context.Context) (*resourcev1.ResourceClaimTemplate, bool, error) {
-				return dynamo.GenerateGMSResourceClaimTemplate(ctx, r.Client, claimTemplateName, dynamoDeployment.Namespace, svcComponent)
+				return dra.GenerateResourceClaimTemplate(ctx, r.Client, claimTemplateName, dynamoDeployment.Namespace, gpuCount, deviceClassName)
 			})
 			if err != nil {
 				logger.Error(err, "failed to sync GMS ResourceClaimTemplate", "service", serviceName)
@@ -1419,10 +1420,14 @@ func (r *DynamoGraphDeploymentReconciler) buildCheckpointJobPodTemplate(
 		return corev1.PodTemplateSpec{}, err
 	}
 
-	// Create a copy of the component spec without checkpoint config
-	// The checkpoint job is CREATING the checkpoint, not restoring from one
+	// Create a copy of the component spec stripped of features that buildCheckpointJob
+	// or the checkpoint controller handle independently. GenerateBasePodSpec would
+	// otherwise apply DGD-specific transforms (DRA claims, GMS server sidecar,
+	// frontend sidecar) that conflict with the checkpoint path's own setup.
 	componentForJob := component.DeepCopy()
 	componentForJob.Checkpoint = nil
+	componentForJob.GPUMemoryService = nil
+	componentForJob.FrontendSidecar = nil
 
 	// Ensure DYN_NAMESPACE is set for checkpoint job using the same logic as regular pods
 	// This is required for service discovery and distributed coordination

deploy/operator/internal/dynamo/gms.go → deploy/operator/internal/dra/dra.go

@@ -3,7 +3,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-package dynamo
+package dra
 
 import (
 	"context"
@@ -13,7 +13,6 @@ import (
 
 	"github.com/ai-dynamo/dynamo/deploy/operator/api/v1alpha1"
 	commonconsts "github.com/ai-dynamo/dynamo/deploy/operator/internal/consts"
-	gmsruntime "github.com/ai-dynamo/dynamo/deploy/operator/internal/gms"
 	corev1 "k8s.io/api/core/v1"
 	resourcev1 "k8s.io/api/resource/v1"
 	apierrors "k8s.io/apimachinery/pkg/api/errors"
@@ -23,147 +22,91 @@ import (
 )
 
 const (
+	// ClaimName is the pod-level DRA ResourceClaim name for shared GPU access.
+	ClaimName = "intrapod-shared-gpu"
+
 	defaultDeviceClassName = "gpu.nvidia.com"
 )
 
-// IsGMSEnabled reports whether GPU Memory Service is requested for the component.
-func IsGMSEnabled(component *v1alpha1.DynamoComponentDeploymentSharedSpec) bool {
-	return component.GPUMemoryService != nil && component.GPUMemoryService.Enabled
-}
-
-// getGPUCount extracts the GPU count from the component resource spec.
-func getGPUCount(component *v1alpha1.DynamoComponentDeploymentSharedSpec) (int, error) {
-	if component.Resources == nil {
-		return 0, fmt.Errorf("resources must be specified when GPU memory service is enabled")
-	}
-
-	gpuStr := ""
-	if component.Resources.Limits != nil && component.Resources.Limits.GPU != "" {
-		gpuStr = component.Resources.Limits.GPU
-	} else if component.Resources.Requests != nil && component.Resources.Requests.GPU != "" {
-		gpuStr = component.Resources.Requests.GPU
-	}
-
-	if gpuStr == "" {
-		return 0, fmt.Errorf("GPU count must be specified when GPU memory service is enabled")
-	}
-
-	count, err := strconv.Atoi(gpuStr)
-	if err != nil {
-		return 0, fmt.Errorf("invalid GPU count %q: %w", gpuStr, err)
-	}
-	if count <= 0 {
-		return 0, fmt.Errorf("GPU count must be greater than 0 when GPU memory service is enabled")
-	}
-	return count, nil
-}
-
-// getDeviceClassName returns the DRA DeviceClass name for the component.
-// It reads from GPUMemoryServiceSpec.DeviceClassName, falling back to the default.
-func getDeviceClassName(component *v1alpha1.DynamoComponentDeploymentSharedSpec) string {
-	if component.GPUMemoryService != nil && component.GPUMemoryService.DeviceClassName != "" {
-		return component.GPUMemoryService.DeviceClassName
-	}
-	return defaultDeviceClassName
-}
-
-// resolveMainContainer finds the container named "main" in the pod spec.
-// Falls back to Containers[0] when there is no container named "main"
-// (e.g. failover pods with engine-0/engine-1 naming).
-func resolveMainContainer(podSpec *corev1.PodSpec) (*corev1.Container, error) {
+// ApplyClaim replaces the first container's nvidia.com/gpu resources with a
+// shared DRA ResourceClaim. Every container that references this claim name
+// will share the same physical GPUs. The function is idempotent — calling it
+// on a pod that already has the claim is a no-op.
+func ApplyClaim(podSpec *corev1.PodSpec, claimTemplateName string) error {
 	if len(podSpec.Containers) == 0 {
-		return nil, fmt.Errorf("pod spec must have at least one container for GPU memory service")
-	}
-	for i := range podSpec.Containers {
-		if podSpec.Containers[i].Name == commonconsts.MainContainerName {
-			return &podSpec.Containers[i], nil
-		}
+		return fmt.Errorf("pod spec must have at least one container for DRA claim")
 	}
-	return &podSpec.Containers[0], nil
-}
 
-// ApplyGPUMemoryService transforms a pod spec to include GMS server sidecars
-// with DRA shared GPU access. The main container's GPU resources are replaced
-// with a DRA ResourceClaim.
-func ApplyGPUMemoryService(
-	podSpec *corev1.PodSpec,
-	component *v1alpha1.DynamoComponentDeploymentSharedSpec,
-	claimTemplateName string,
-) error {
-	gpuCount, err := getGPUCount(component)
-	if err != nil {
-		return err
+	// Skip if the pod-level claim already exists (idempotent).
+	for i := range podSpec.ResourceClaims {
+		if podSpec.ResourceClaims[i].Name == ClaimName {
+			return nil
 		}
-	_ = gpuCount // GPU count is used for DRA claim template; sidecar discovers devices via pynvml
-
-	mainContainer, err := resolveMainContainer(podSpec)
-	if err != nil {
-		return err
 	}
 
-	// Replace GPU resources with DRA claim on main container
-	removeGPUResources(mainContainer)
-	mainContainer.Resources.Claims = append(mainContainer.Resources.Claims, corev1.ResourceClaim{
-		Name: gmsruntime.DRAClaimName,
+	// Replace nvidia.com/gpu with the shared DRA claim.
+	gpuResource := corev1.ResourceName(commonconsts.KubeResourceGPUNvidia)
+	delete(podSpec.Containers[0].Resources.Limits, gpuResource)
+	delete(podSpec.Containers[0].Resources.Requests, gpuResource)
+	podSpec.Containers[0].Resources.Claims = append(podSpec.Containers[0].Resources.Claims, corev1.ResourceClaim{
+		Name: ClaimName,
 	})
 
-	// Add GMS server sidecar, shared volume, and socket env vars.
-	// The sidecar gets DRA claims copied from main automatically.
-	gmsruntime.EnsureServerSidecar(podSpec, mainContainer)
-
-	// GPU nodes are typically tainted with nvidia.com/gpu=NoSchedule. With
-	// traditional scheduling the device-plugin injects the matching toleration,
-	// but DRA bypasses that path. Re-add the toleration explicitly so the pod
-	// can schedule on GPU nodes.
+	// GPU nodes are typically tainted with nvidia.com/gpu=NoSchedule. DRA
+	// bypasses the device-plugin toleration injection, so add it explicitly.
 	podSpec.Tolerations = append(podSpec.Tolerations, corev1.Toleration{
 		Key:      commonconsts.KubeResourceGPUNvidia,
 		Operator: corev1.TolerationOpExists,
 		Effect:   corev1.TaintEffectNoSchedule,
 	})
 
-	// Add pod-level DRA resource claim referencing the ResourceClaimTemplate
 	podSpec.ResourceClaims = append(podSpec.ResourceClaims, corev1.PodResourceClaim{
-		Name:                      gmsruntime.DRAClaimName,
+		Name:                      ClaimName,
 		ResourceClaimTemplateName: &claimTemplateName,
 	})
 
 	return nil
 }
 
-// removeGPUResources strips nvidia.com/gpu from container resource limits and requests.
-// GPU allocation is handled by DRA when GMS is enabled.
-func removeGPUResources(container *corev1.Container) {
-	gpuResource := corev1.ResourceName(commonconsts.KubeResourceGPUNvidia)
-	if container.Resources.Limits != nil {
-		delete(container.Resources.Limits, gpuResource)
-	}
-	if container.Resources.Requests != nil {
-		delete(container.Resources.Requests, gpuResource)
-	}
+// ResourceClaimTemplateName returns the deterministic name for the
+// ResourceClaimTemplate associated with a component.
+func ResourceClaimTemplateName(parentName, serviceName string) string {
+	return fmt.Sprintf("%s-%s-gpu", parentName, strings.ToLower(serviceName))
 }
 
-// GMSResourceClaimTemplateName returns the deterministic name for the
-// ResourceClaimTemplate associated with a GMS-enabled component.
-func GMSResourceClaimTemplateName(parentName, serviceName string) string {
-	return fmt.Sprintf("%s-%s-gpu", parentName, strings.ToLower(serviceName))
+// ExtractGPUParams extracts the GPU count and device class name from API types
+// shared by DGD components and DynamoCheckpoint specs. Returns gpuCount=0 when
+// GMS is not enabled, which tells GenerateResourceClaimTemplate to delete.
+func ExtractGPUParams(gmsSpec *v1alpha1.GPUMemoryServiceSpec, resources *v1alpha1.Resources) (gpuCount int, deviceClassName string) {
+	if gmsSpec == nil || !gmsSpec.Enabled {
+		return 0, ""
+	}
+	deviceClassName = gmsSpec.DeviceClassName
+	if resources != nil {
+		gpuStr := ""
+		if resources.Limits != nil {
+			gpuStr = resources.Limits.GPU
+		}
+		if gpuStr == "" && resources.Requests != nil {
+			gpuStr = resources.Requests.GPU
+		}
+		gpuCount, _ = strconv.Atoi(gpuStr)
+	}
+	return gpuCount, deviceClassName
 }
 
-// GenerateGMSResourceClaimTemplate builds the ResourceClaimTemplate that
-// provides shared GPU access to all containers in a GMS-enabled pod via DRA.
-//
-// claimTemplateName is the deterministic name for the template; callers should
-// compute it via GMSResourceClaimTemplateName.
+// GenerateResourceClaimTemplate builds the ResourceClaimTemplate that provides
+// shared GPU access to all containers in a pod via DRA.
 //
-// When GMS is not enabled for the component, it returns the template skeleton
-// with toDelete=true so that SyncResource cleans up any previously created template.
-//
-// The cl parameter is used to verify the DeviceClass exists before creating
-// the template. Pass nil to skip the DeviceClass check.
-func GenerateGMSResourceClaimTemplate(
+// When gpuCount <= 0 it returns the template skeleton with toDelete=true so
+// that SyncResource cleans up any previously created template. Pass cl=nil to
+// skip the DeviceClass existence check.
+func GenerateResourceClaimTemplate(
 	ctx context.Context,
 	cl client.Client,
 	claimTemplateName, namespace string,
-	component *v1alpha1.DynamoComponentDeploymentSharedSpec,
+	gpuCount int,
+	deviceClassName string,
 ) (*resourcev1.ResourceClaimTemplate, bool, error) {
 	template := &resourcev1.ResourceClaimTemplate{
 		ObjectMeta: metav1.ObjectMeta{
@@ -172,18 +115,14 @@ func GenerateGMSResourceClaimTemplate(
 		},
 	}
 
-	if !IsGMSEnabled(component) {
+	if gpuCount <= 0 {
 		return template, true, nil
 	}
 
-	gpuCount, err := getGPUCount(component)
-	if err != nil {
-		return nil, false, fmt.Errorf("failed to get GPU count for ResourceClaimTemplate: %w", err)
+	if deviceClassName == "" {
+		deviceClassName = defaultDeviceClassName
 	}
 
-	deviceClassName := getDeviceClassName(component)
-
-	// Verify the DeviceClass exists before creating the template
 	if cl != nil {
 		dc := &resourcev1.DeviceClass{}
 		if err := cl.Get(ctx, types.NamespacedName{Name: deviceClassName}, dc); err != nil {

deploy/operator/internal/dra/dra_test.go

+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package dra
+
+import (
+	"context"
+	"strconv"
+	"testing"
+
+	"github.com/ai-dynamo/dynamo/deploy/operator/api/v1alpha1"
+	commonconsts "github.com/ai-dynamo/dynamo/deploy/operator/internal/consts"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	corev1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/api/resource"
+	"k8s.io/apimachinery/pkg/util/intstr"
+)
+
+func basePodSpec() corev1.PodSpec {
+	httpPort := intstr.FromString("system")
+	return corev1.PodSpec{
+		Containers: []corev1.Container{{
+			Name:    "main",
+			Image:   "test-image:latest",
+			Command: []string{"python3", "-m", "dynamo.vllm"},
+			Env: []corev1.EnvVar{
+				{Name: "DYN_SYSTEM_PORT", Value: "9090"},
+			},
+			Ports: []corev1.ContainerPort{
+				{Name: "system", ContainerPort: 9090, Protocol: corev1.ProtocolTCP},
+			},
+			StartupProbe: &corev1.Probe{
+				ProbeHandler: corev1.ProbeHandler{
+					HTTPGet: &corev1.HTTPGetAction{Path: "/health", Port: httpPort},
+				},
+			},
+			Resources: corev1.ResourceRequirements{
+				Limits: corev1.ResourceList{
+					corev1.ResourceName(commonconsts.KubeResourceGPUNvidia): resource.MustParse("2"),
+				},
+			},
+		}},
+	}
+}
+
+func TestApplyClaim_EmptyContainers(t *testing.T) {
+	ps := corev1.PodSpec{}
+	err := ApplyClaim(&ps, "myapp-worker-gpu")
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "at least one container")
+}
+
+func TestApplyClaim_ReplacesGPUWithDRAClaim(t *testing.T) {
+	ps := basePodSpec()
+	err := ApplyClaim(&ps, "myapp-worker-gpu")
+	require.NoError(t, err)
+
+	main := ps.Containers[0]
+
+	gpuResource := corev1.ResourceName(commonconsts.KubeResourceGPUNvidia)
+	_, hasGPU := main.Resources.Limits[gpuResource]
+	assert.False(t, hasGPU)
+
+	require.Len(t, main.Resources.Claims, 1)
+	assert.Equal(t, ClaimName, main.Resources.Claims[0].Name)
+
+	require.Len(t, ps.ResourceClaims, 1)
+	assert.Equal(t, ClaimName, ps.ResourceClaims[0].Name)
+	assert.Equal(t, "myapp-worker-gpu", *ps.ResourceClaims[0].ResourceClaimTemplateName)
+
+	var hasToleration bool
+	for _, tol := range ps.Tolerations {
+		if tol.Key == commonconsts.KubeResourceGPUNvidia && tol.Effect == corev1.TaintEffectNoSchedule {
+			hasToleration = true
+		}
+	}
+	assert.True(t, hasToleration)
+	assert.Empty(t, ps.InitContainers)
+}
+
+func TestApplyClaim_AlwaysTargetsFirstContainer(t *testing.T) {
+	ps := basePodSpec()
+	ps.Containers = append(ps.Containers, corev1.Container{Name: "sidecar", Image: "sidecar:latest"})
+
+	err := ApplyClaim(&ps, "myapp-worker-gpu")
+	require.NoError(t, err)
+
+	require.Len(t, ps.Containers[0].Resources.Claims, 1)
+	assert.Equal(t, ClaimName, ps.Containers[0].Resources.Claims[0].Name)
+	assert.Empty(t, ps.Containers[1].Resources.Claims)
+}
+
+func TestGenerateResourceClaimTemplate_Enabled(t *testing.T) {
+	tmpl, toDelete, err := GenerateResourceClaimTemplate(context.Background(), nil, "myapp-worker-gpu", "default", 4, "")
+	require.NoError(t, err)
+	assert.False(t, toDelete)
+	assert.Equal(t, "myapp-worker-gpu", tmpl.Name)
+	require.Len(t, tmpl.Spec.Spec.Devices.Requests, 1)
+	req := tmpl.Spec.Spec.Devices.Requests[0]
+	assert.Equal(t, defaultDeviceClassName, req.Exactly.DeviceClassName)
+	assert.Equal(t, int64(4), req.Exactly.Count)
+}
+
+func TestGenerateResourceClaimTemplate_CustomDeviceClass(t *testing.T) {
+	tmpl, _, err := GenerateResourceClaimTemplate(context.Background(), nil, "myapp-worker-gpu", "default", 2, "gpu.intel.com/xe")
+	require.NoError(t, err)
+	assert.Equal(t, "gpu.intel.com/xe", tmpl.Spec.Spec.Devices.Requests[0].Exactly.DeviceClassName)
+}
+
+func TestGenerateResourceClaimTemplate_DisabledReturnsDelete(t *testing.T) {
+	tmpl, toDelete, err := GenerateResourceClaimTemplate(context.Background(), nil, "myapp-worker-gpu", "default", 0, "")
+	require.NoError(t, err)
+	assert.True(t, toDelete)
+	assert.Equal(t, "myapp-worker-gpu", tmpl.Name)
+}
+
+func TestResourceClaimTemplateName(t *testing.T) {
+	assert.Equal(t, "myapp-worker-gpu", ResourceClaimTemplateName("myapp", "Worker"))
+	assert.Equal(t, "app-vllmdecodeworker-gpu", ResourceClaimTemplateName("app", "VllmDecodeWorker"))
+}
+
+func TestExtractGPUParams(t *testing.T) {
+	count, dc := ExtractGPUParams(nil, nil)
+	assert.Equal(t, 0, count)
+	assert.Equal(t, "", dc)
+
+	count, dc = ExtractGPUParams(
+		&v1alpha1.GPUMemoryServiceSpec{Enabled: true},
+		&v1alpha1.Resources{Limits: &v1alpha1.ResourceItem{GPU: strconv.Itoa(4)}},
+	)
+	assert.Equal(t, 4, count)
+	assert.Equal(t, "", dc)
+
+	count, dc = ExtractGPUParams(
+		&v1alpha1.GPUMemoryServiceSpec{Enabled: true, DeviceClassName: "gpu.intel.com/xe"},
+		&v1alpha1.Resources{Requests: &v1alpha1.ResourceItem{GPU: "2"}},
+	)
+	assert.Equal(t, 2, count)
+	assert.Equal(t, "gpu.intel.com/xe", dc)
+}

deploy/operator/internal/dynamo/failover_test.go

@@ -12,7 +12,8 @@ import (
 
 	"github.com/ai-dynamo/dynamo/deploy/operator/api/v1alpha1"
 	commonconsts "github.com/ai-dynamo/dynamo/deploy/operator/internal/consts"
-	gmsruntime "github.com/ai-dynamo/dynamo/deploy/operator/internal/gms"
+	"github.com/ai-dynamo/dynamo/deploy/operator/internal/dra"
+	"github.com/ai-dynamo/dynamo/deploy/operator/internal/gms"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	corev1 "k8s.io/api/core/v1"
@@ -36,7 +37,7 @@ func failoverPodSpec() corev1.PodSpec {
 					{Name: "DYN_SYSTEM_USE_ENDPOINT_HEALTH_STATUS", Value: "true"},
 					{Name: "DYN_HEALTH_CHECK_ENABLED", Value: "true"},
 					{Name: commonconsts.DynamoDiscoveryBackendEnvVar, Value: "kubernetes"},
-					{Name: "TMPDIR", Value: gmsruntime.SharedMountPath},
+					{Name: "TMPDIR", Value: gms.SharedMountPath},
 				},
 				Ports: []corev1.ContainerPort{
 					{Name: "system", ContainerPort: 9090, Protocol: corev1.ProtocolTCP},
@@ -57,10 +58,10 @@ func failoverPodSpec() corev1.PodSpec {
 					},
 				},
 				Resources: corev1.ResourceRequirements{
-					Claims: []corev1.ResourceClaim{{Name: gmsruntime.DRAClaimName}},
+					Claims: []corev1.ResourceClaim{{Name: dra.ClaimName}},
 				},
 				VolumeMounts: []corev1.VolumeMount{
-					{Name: gmsruntime.SharedVolumeName, MountPath: gmsruntime.SharedMountPath},
+					{Name: gms.SharedVolumeName, MountPath: gms.SharedMountPath},
 				},
 			},
 			{
@@ -165,7 +166,7 @@ func TestBuildFailoverPod_PreservesDRAClaim(t *testing.T) {
 	for i := range 2 {
 		engine := ps.Containers[i]
 		require.Len(t, engine.Resources.Claims, 1, "engine-%d should retain DRA claim", i)
-		assert.Equal(t, gmsruntime.DRAClaimName, engine.Resources.Claims[0].Name)
+		assert.Equal(t, dra.ClaimName, engine.Resources.Claims[0].Name)
 	}
 }
 
@@ -214,3 +215,11 @@ func TestIsFailoverEnabled(t *testing.T) {
 	}))
 	assert.False(t, isFailoverEnabled(&v1alpha1.DynamoComponentDeploymentSharedSpec{}))
 }
+
+func envToMap(envs []corev1.EnvVar) map[string]string {
+	m := make(map[string]string, len(envs))
+	for _, e := range envs {
+		m[e.Name] = e.Value
+	}
+	return m
+}

deploy/operator/internal/dynamo/gms_test.go

-/*
- * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
- * SPDX-License-Identifier: Apache-2.0
- */
-
-package dynamo
-
-import (
-	"context"
-	"strconv"
-	"testing"
-
-	"github.com/ai-dynamo/dynamo/deploy/operator/api/v1alpha1"
-	commonconsts "github.com/ai-dynamo/dynamo/deploy/operator/internal/consts"
-	gmsruntime "github.com/ai-dynamo/dynamo/deploy/operator/internal/gms"
-	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
-	corev1 "k8s.io/api/core/v1"
-	"k8s.io/apimachinery/pkg/api/resource"
-	"k8s.io/apimachinery/pkg/util/intstr"
-)
-
-func gmsComponent(gpuCount int) *v1alpha1.DynamoComponentDeploymentSharedSpec {
-	return &v1alpha1.DynamoComponentDeploymentSharedSpec{
-		ComponentType:    commonconsts.ComponentTypeWorker,
-		GPUMemoryService: &v1alpha1.GPUMemoryServiceSpec{Enabled: true},
-		Resources: &v1alpha1.Resources{
-			Limits: &v1alpha1.ResourceItem{GPU: strconv.Itoa(gpuCount)},
-		},
-	}
-}
-
-func gmsBasePodSpec() corev1.PodSpec {
-	httpPort := intstr.FromString("system")
-	return corev1.PodSpec{
-		Containers: []corev1.Container{
-			{
-				Name:    "main",
-				Image:   "test-image:latest",
-				Command: []string{"python3", "-m", "dynamo.vllm"},
-				Env: []corev1.EnvVar{
-					{Name: "DYN_SYSTEM_PORT", Value: "9090"},
-					{Name: commonconsts.DynamoDiscoveryBackendEnvVar, Value: "kubernetes"},
-				},
-				Ports: []corev1.ContainerPort{
-					{Name: "system", ContainerPort: 9090, Protocol: corev1.ProtocolTCP},
-				},
-				StartupProbe: &corev1.Probe{
-					ProbeHandler: corev1.ProbeHandler{
-						HTTPGet: &corev1.HTTPGetAction{Path: "/health", Port: httpPort},
-					},
-				},
-				Resources: corev1.ResourceRequirements{
-					Limits: corev1.ResourceList{
-						corev1.ResourceName(commonconsts.KubeResourceGPUNvidia): resource.MustParse("2"),
-					},
-				},
-			},
-		},
-	}
-}
-
-// --- applyGPUMemoryService ---
-
-func TestApplyGPUMemoryService_EmptyContainers(t *testing.T) {
-	ps := corev1.PodSpec{}
-	err := ApplyGPUMemoryService(&ps, gmsComponent(2), "myapp-worker-gpu")
-	require.Error(t, err)
-	assert.Contains(t, err.Error(), "at least one container")
-}
-
-func TestApplyGPUMemoryService_MainContainerTransformed(t *testing.T) {
-	ps := gmsBasePodSpec()
-	err := ApplyGPUMemoryService(&ps, gmsComponent(2), "myapp-worker-gpu")
-	require.NoError(t, err)
-
-	main := ps.Containers[0]
-
-	// GPU resources should be removed
-	gpuResource := corev1.ResourceName(commonconsts.KubeResourceGPUNvidia)
-	_, hasGPU := main.Resources.Limits[gpuResource]
-	assert.False(t, hasGPU, "main container should not have GPU limits")
-
-	// Should have DRA claim
-	require.Len(t, main.Resources.Claims, 1)
-	assert.Equal(t, gmsruntime.DRAClaimName, main.Resources.Claims[0].Name)
-
-	// Should have shared volume mount
-	var hasSharedMount bool
-	for _, vm := range main.VolumeMounts {
-		if vm.Name == gmsruntime.SharedVolumeName && vm.MountPath == gmsruntime.SharedMountPath {
-			hasSharedMount = true
-		}
-	}
-	assert.True(t, hasSharedMount, "main container should have gms-shared volume mount")
-
-	// Should have TMPDIR and GMS_SOCKET_DIR
-	envMap := envToMap(main.Env)
-	assert.Equal(t, gmsruntime.SharedMountPath, envMap["TMPDIR"])
-	assert.Equal(t, gmsruntime.SharedMountPath, envMap["GMS_SOCKET_DIR"])
-}
-
-func TestApplyGPUMemoryService_GMSSidecarInjected(t *testing.T) {
-	ps := gmsBasePodSpec()
-	err := ApplyGPUMemoryService(&ps, gmsComponent(2), "myapp-worker-gpu")
-	require.NoError(t, err)
-
-	require.Len(t, ps.InitContainers, 1)
-	gms := ps.InitContainers[0]
-	assert.Equal(t, gmsruntime.ServerContainerName, gms.Name)
-	assert.Equal(t, "test-image:latest", gms.Image)
-	assert.Equal(t, []string{"python3", "-m", "gpu_memory_service.cli.server"}, gms.Command)
-	assert.NotNil(t, gms.RestartPolicy)
-	assert.Equal(t, corev1.ContainerRestartPolicyAlways, *gms.RestartPolicy)
-	require.NotNil(t, gms.StartupProbe)
-	assert.Equal(t, int32(1), gms.StartupProbe.PeriodSeconds)
-	assert.Equal(t, int32(300), gms.StartupProbe.FailureThreshold)
-
-	// GMS sidecar should have DRA claim copied from main
-	require.Len(t, gms.Resources.Claims, 1)
-	assert.Equal(t, gmsruntime.DRAClaimName, gms.Resources.Claims[0].Name)
-
-	// GMS sidecar should have TMPDIR
-	gmsEnv := envToMap(gms.Env)
-	assert.Equal(t, gmsruntime.SharedMountPath, gmsEnv["TMPDIR"])
-}
-
-func TestApplyGPUMemoryService_SharedVolume(t *testing.T) {
-	ps := gmsBasePodSpec()
-	err := ApplyGPUMemoryService(&ps, gmsComponent(2), "myapp-worker-gpu")
-	require.NoError(t, err)
-
-	var found bool
-	for _, v := range ps.Volumes {
-		if v.Name == gmsruntime.SharedVolumeName {
-			assert.NotNil(t, v.EmptyDir)
-			found = true
-		}
-	}
-	assert.True(t, found, "should have gms-shared volume")
-}
-
-func TestApplyGPUMemoryService_GPUToleration(t *testing.T) {
-	ps := gmsBasePodSpec()
-	err := ApplyGPUMemoryService(&ps, gmsComponent(2), "myapp-worker-gpu")
-	require.NoError(t, err)
-
-	var found bool
-	for _, tol := range ps.Tolerations {
-		if tol.Key == commonconsts.KubeResourceGPUNvidia && tol.Effect == corev1.TaintEffectNoSchedule {
-			assert.Equal(t, corev1.TolerationOpExists, tol.Operator)
-			found = true
-		}
-	}
-	assert.True(t, found, "should have nvidia.com/gpu NoSchedule toleration")
-}
-
-func TestApplyGPUMemoryService_DRAResourceClaim(t *testing.T) {
-	ps := gmsBasePodSpec()
-	err := ApplyGPUMemoryService(&ps, gmsComponent(2), "myapp-worker-gpu")
-	require.NoError(t, err)
-
-	require.Len(t, ps.ResourceClaims, 1)
-	assert.Equal(t, gmsruntime.DRAClaimName, ps.ResourceClaims[0].Name)
-	assert.Equal(t, "myapp-worker-gpu", *ps.ResourceClaims[0].ResourceClaimTemplateName)
-}
-
-func TestApplyGPUMemoryService_PreservesExistingEnv(t *testing.T) {
-	ps := gmsBasePodSpec()
-	err := ApplyGPUMemoryService(&ps, gmsComponent(2), "myapp-worker-gpu")
-	require.NoError(t, err)
-
-	main := ps.Containers[0]
-	envMap := envToMap(main.Env)
-	assert.Equal(t, "kubernetes", envMap[commonconsts.DynamoDiscoveryBackendEnvVar])
-	assert.Equal(t, "9090", envMap["DYN_SYSTEM_PORT"])
-}
-
-func TestApplyGPUMemoryService_SingleContainer(t *testing.T) {
-	ps := gmsBasePodSpec()
-	err := ApplyGPUMemoryService(&ps, gmsComponent(2), "myapp-worker-gpu")
-	require.NoError(t, err)
-
-	assert.Len(t, ps.Containers, 1)
-	assert.Equal(t, "main", ps.Containers[0].Name)
-}
-
-func TestApplyGPUMemoryService_ResolvesMainByName(t *testing.T) {
-	ps := gmsBasePodSpec()
-	// Prepend a sidecar so main is NOT Containers[0]
-	sidecar := corev1.Container{Name: "sidecar", Image: "sidecar:latest"}
-	ps.Containers = append([]corev1.Container{sidecar}, ps.Containers...)
-	require.Equal(t, "sidecar", ps.Containers[0].Name)
-
-	err := ApplyGPUMemoryService(&ps, gmsComponent(2), "myapp-worker-gpu")
-	require.NoError(t, err)
-
-	// Sidecar should be untouched
-	assert.Equal(t, "sidecar", ps.Containers[0].Name)
-	assert.Empty(t, ps.Containers[0].Resources.Claims)
-
-	// Main should have DRA claim
-	main := ps.Containers[1]
-	assert.Equal(t, "main", main.Name)
-	require.Len(t, main.Resources.Claims, 1)
-	assert.Equal(t, gmsruntime.DRAClaimName, main.Resources.Claims[0].Name)
-}
-
-// --- GenerateGMSResourceClaimTemplate ---
-
-func TestGenerateGMSResourceClaimTemplate_Enabled(t *testing.T) {
-	component := gmsComponent(4)
-	tmpl, toDelete, err := GenerateGMSResourceClaimTemplate(context.Background(), nil, "myapp-worker-gpu", "default", component)
-
-	require.NoError(t, err)
-	assert.False(t, toDelete)
-	assert.Equal(t, "myapp-worker-gpu", tmpl.Name)
-	assert.Equal(t, "default", tmpl.Namespace)
-
-	require.Len(t, tmpl.Spec.Spec.Devices.Requests, 1)
-	req := tmpl.Spec.Spec.Devices.Requests[0]
-	assert.Equal(t, "gpus", req.Name)
-	require.NotNil(t, req.Exactly)
-	assert.Equal(t, defaultDeviceClassName, req.Exactly.DeviceClassName)
-	assert.Equal(t, int64(4), req.Exactly.Count)
-}
-
-func TestGenerateGMSResourceClaimTemplate_CustomDeviceClass(t *testing.T) {
-	component := gmsComponent(2)
-	component.GPUMemoryService.DeviceClassName = "gpu.intel.com/xe"
-	tmpl, toDelete, err := GenerateGMSResourceClaimTemplate(context.Background(), nil, "myapp-worker-gpu", "default", component)
-
-	require.NoError(t, err)
-	assert.False(t, toDelete)
-	assert.Equal(t, "gpu.intel.com/xe", tmpl.Spec.Spec.Devices.Requests[0].Exactly.DeviceClassName)
-}
-
-func TestGenerateGMSResourceClaimTemplate_DisabledReturnsDelete(t *testing.T) {
-	component := &v1alpha1.DynamoComponentDeploymentSharedSpec{
-		ComponentType: commonconsts.ComponentTypeWorker,
-	}
-	tmpl, toDelete, err := GenerateGMSResourceClaimTemplate(context.Background(), nil, "myapp-worker-gpu", "default", component)
-
-	require.NoError(t, err)
-	assert.True(t, toDelete)
-	assert.Equal(t, "myapp-worker-gpu", tmpl.Name)
-}
-
-func TestGenerateGMSResourceClaimTemplate_NoGPUCountError(t *testing.T) {
-	component := &v1alpha1.DynamoComponentDeploymentSharedSpec{
-		ComponentType:    commonconsts.ComponentTypeWorker,
-		GPUMemoryService: &v1alpha1.GPUMemoryServiceSpec{Enabled: true},
-	}
-	_, _, err := GenerateGMSResourceClaimTemplate(context.Background(), nil, "myapp-worker-gpu", "default", component)
-	require.Error(t, err)
-	assert.Contains(t, err.Error(), "resources must be specified")
-}
-
-// --- GMSResourceClaimTemplateName ---
-
-func TestGMSResourceClaimTemplateName(t *testing.T) {
-	assert.Equal(t, "myapp-worker-gpu", GMSResourceClaimTemplateName("myapp", "Worker"))
-	assert.Equal(t, "app-vllmdecodeworker-gpu", GMSResourceClaimTemplateName("app", "VllmDecodeWorker"))
-}
-
-// --- isGMSEnabled ---
-
-func TestIsGMSEnabled(t *testing.T) {
-	assert.True(t, IsGMSEnabled(&v1alpha1.DynamoComponentDeploymentSharedSpec{
-		GPUMemoryService: &v1alpha1.GPUMemoryServiceSpec{Enabled: true},
-	}))
-	assert.False(t, IsGMSEnabled(&v1alpha1.DynamoComponentDeploymentSharedSpec{
-		GPUMemoryService: &v1alpha1.GPUMemoryServiceSpec{Enabled: false},
-	}))
-	assert.False(t, IsGMSEnabled(&v1alpha1.DynamoComponentDeploymentSharedSpec{}))
-}
-
-// --- getGPUCount ---
-
-func TestGetGPUCount(t *testing.T) {
-	tests := []struct {
-		name      string
-		component *v1alpha1.DynamoComponentDeploymentSharedSpec
-		want      int
-		wantErr   bool
-	}{
-		{
-			name:      "from limits",
-			component: &v1alpha1.DynamoComponentDeploymentSharedSpec{Resources: &v1alpha1.Resources{Limits: &v1alpha1.ResourceItem{GPU: "4"}}},
-			want:      4,
-		},
-		{
-			name:      "from requests",
-			component: &v1alpha1.DynamoComponentDeploymentSharedSpec{Resources: &v1alpha1.Resources{Requests: &v1alpha1.ResourceItem{GPU: "2"}}},
-			want:      2,
-		},
-		{
-			name:      "no resources",
-			component: &v1alpha1.DynamoComponentDeploymentSharedSpec{},
-			wantErr:   true,
-		},
-		{
-			name:      "invalid GPU string",
-			component: &v1alpha1.DynamoComponentDeploymentSharedSpec{Resources: &v1alpha1.Resources{Limits: &v1alpha1.ResourceItem{GPU: "abc"}}},
-			wantErr:   true,
-		},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			got, err := getGPUCount(tt.component)
-			if tt.wantErr {
-				assert.Error(t, err)
-			} else {
-				require.NoError(t, err)
-				assert.Equal(t, tt.want, got)
-			}
-		})
-	}
-}
-
-// --- getDeviceClassName ---
-
-func TestGetDeviceClassName(t *testing.T) {
-	assert.Equal(t, defaultDeviceClassName, getDeviceClassName(&v1alpha1.DynamoComponentDeploymentSharedSpec{}))
-	assert.Equal(t, defaultDeviceClassName, getDeviceClassName(&v1alpha1.DynamoComponentDeploymentSharedSpec{
-		GPUMemoryService: &v1alpha1.GPUMemoryServiceSpec{Enabled: true},
-	}))
-	assert.Equal(t, "gpu.intel.com/xe", getDeviceClassName(&v1alpha1.DynamoComponentDeploymentSharedSpec{
-		GPUMemoryService: &v1alpha1.GPUMemoryServiceSpec{Enabled: true, DeviceClassName: "gpu.intel.com/xe"},
-	}))
-}
-
-// helpers
-
-func envToMap(envs []corev1.EnvVar) map[string]string {
-	m := make(map[string]string, len(envs))
-	for _, e := range envs {
-		m[e.Name] = e.Value
-	}
-	return m
-}

deploy/operator/internal/dynamo/graph.go

@@ -38,6 +38,8 @@ import (
 	commonconsts "github.com/ai-dynamo/dynamo/deploy/operator/internal/consts"
 	"github.com/ai-dynamo/dynamo/deploy/operator/internal/controller_common"
 	"github.com/ai-dynamo/dynamo/deploy/operator/internal/discovery"
+	"github.com/ai-dynamo/dynamo/deploy/operator/internal/dra"
+	gms "github.com/ai-dynamo/dynamo/deploy/operator/internal/gms"
 	grovev1alpha1 "github.com/ai-dynamo/grove/operator/api/core/v1alpha1"
 	"github.com/imdario/mergo"
 	networkingv1beta1 "istio.io/client-go/pkg/apis/networking/v1beta1"
@@ -1182,12 +1184,13 @@ func GenerateBasePodSpec(
 		}
 	}
 
-	// Inject GMS sidecar with DRA shared GPU access when GPU memory service is enabled.
-	if IsGMSEnabled(component) {
-		claimTemplateName := GMSResourceClaimTemplateName(parentGraphDeploymentName, serviceName)
-		if err := ApplyGPUMemoryService(&podSpec, component, claimTemplateName); err != nil {
-			return nil, fmt.Errorf("failed to apply GPU memory service: %w", err)
+	// GMS: replace nvidia.com/gpu with a shared DRA claim and add the server sidecar.
+	if component.GPUMemoryService != nil && component.GPUMemoryService.Enabled {
+		claimTemplateName := dra.ResourceClaimTemplateName(parentGraphDeploymentName, serviceName)
+		if err := dra.ApplyClaim(&podSpec, claimTemplateName); err != nil {
+			return nil, fmt.Errorf("failed to apply DRA claim for GMS: %w", err)
 		}
+		gms.EnsureServerSidecar(&podSpec, &podSpec.Containers[0])
 	}
 
 	// Clone main container into two engine containers (active + standby) for failover.

deploy/operator/internal/gms/runtime.go → deploy/operator/internal/gms/gms.go

@@ -3,11 +3,14 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+// Package gms provides GMS (GPU Memory Service) server container building
+// for both steady-state DGD pods and checkpoint/restore flows.
 package gms
 
 import (
 	"path/filepath"
 
+	"github.com/ai-dynamo/dynamo/deploy/operator/internal/dra"
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/utils/ptr"
 )
@@ -23,32 +26,25 @@ const (
 	// SharedMountPath is the mount path for the shared GMS socket directory.
 	SharedMountPath = "/shared"
 
-	// DRAClaimName is the pod-level DRA ResourceClaim name used by both the
-	// main container and GMS sidecars.
-	DRAClaimName = "shared-gpu"
+	// EnvSocketDir is the environment variable name for the GMS UDS socket directory.
+	EnvSocketDir = "GMS_SOCKET_DIR"
 
-	// ControlVolumeName is the checkpoint-specific control volume name.
-	ControlVolumeName = "gms-control"
-
-	// ControlDir is the mount path for the checkpoint control volume.
-	ControlDir = "/tmp/gms-control"
+	// ServerModule is the Python module for the GMS server entry point.
+	ServerModule = "gpu_memory_service.cli.server"
 
 	readyFile = "gms-ready"
-
-	serverSidecarModule = "gpu_memory_service.cli.server"
 )
 
 // EnsureServerSidecar adds the GMS server as a restartable init sidecar with a
-// startup probe. Used for checkpoint jobs and steady-state pods where the main
-// container needs GMS sockets before starting.
+// startup probe. Idempotent — safe to call from both the DGD and checkpoint paths.
 func EnsureServerSidecar(podSpec *corev1.PodSpec, mainContainer *corev1.Container) {
 	if podSpec == nil || mainContainer == nil {
 		return
 	}
 
-	ensureSharedVolume(podSpec, mainContainer)
+	EnsureSharedVolume(podSpec, mainContainer)
 
-	sidecar := serverContainer(mainContainer.Image)
+	sidecar := Container(ServerContainerName, ServerModule, mainContainer.Image)
 	sidecar.RestartPolicy = ptr.To(corev1.ContainerRestartPolicyAlways)
 	sidecar.StartupProbe = &corev1.Probe{
 		ProbeHandler: corev1.ProbeHandler{
@@ -59,9 +55,6 @@ func EnsureServerSidecar(podSpec *corev1.PodSpec, mainContainer *corev1.Containe
 		PeriodSeconds:    1,
 		FailureThreshold: 300, // 1s * 300 = 5 min
 	}
-	copyDeviceClaims(mainContainer, &sidecar)
-	// Idempotent — EnsureServerSidecar may be called by both the
-	// steady-state operator path and the checkpoint overlay.
 	for i := range podSpec.InitContainers {
 		if podSpec.InitContainers[i].Name == sidecar.Name {
 			return
@@ -70,41 +63,9 @@ func EnsureServerSidecar(podSpec *corev1.PodSpec, mainContainer *corev1.Containe
 	podSpec.InitContainers = append(podSpec.InitContainers, sidecar)
 }
 
-// BuildServerContainer prepares the shared GMS volume/env and returns a GMS
-// server container suitable for use as a regular sidecar. The caller must
-// append the returned container to podSpec.Containers.
-//
-// Used for restore pods where the main container is CRIU-restored and does not
-// need GMS sockets at startup. The gms-loader polls for sockets internally.
-func BuildServerContainer(podSpec *corev1.PodSpec, mainContainer *corev1.Container) corev1.Container {
-	ensureSharedVolume(podSpec, mainContainer)
-	sidecar := serverContainer(mainContainer.Image)
-	copyDeviceClaims(mainContainer, &sidecar)
-	return sidecar
-}
-
-// FindServerContainer returns a pointer to the GMS server container, checking
-// both init containers and regular containers. Returns nil if not present.
-func FindServerContainer(podSpec *corev1.PodSpec) *corev1.Container {
-	if podSpec == nil {
-		return nil
-	}
-	for i := range podSpec.InitContainers {
-		if podSpec.InitContainers[i].Name == ServerContainerName {
-			return &podSpec.InitContainers[i]
-		}
-	}
-	for i := range podSpec.Containers {
-		if podSpec.Containers[i].Name == ServerContainerName {
-			return &podSpec.Containers[i]
-		}
-	}
-	return nil
-}
-
-// ensureSharedVolume adds the shared GMS socket volume, mounts, and env vars.
-// Idempotent — may be called by both steady-state and checkpoint paths.
-func ensureSharedVolume(podSpec *corev1.PodSpec, mainContainer *corev1.Container) {
+// EnsureSharedVolume adds the GMS UDS socket volume, mount, and GMS_SOCKET_DIR
+// env var to the main container. Idempotent.
+func EnsureSharedVolume(podSpec *corev1.PodSpec, mainContainer *corev1.Container) {
 	hasVolume := false
 	for _, v := range podSpec.Volumes {
 		if v.Name == SharedVolumeName {
@@ -119,8 +80,6 @@ func ensureSharedVolume(podSpec *corev1.PodSpec, mainContainer *corev1.Container
 		})
 	}
 
-	// Mount and env injection checked independently of volume existence —
-	// another code path may have added the volume without configuring main.
 	hasMount := false
 	for _, m := range mainContainer.VolumeMounts {
 		if m.Name == SharedVolumeName {
@@ -134,41 +93,31 @@ func ensureSharedVolume(podSpec *corev1.PodSpec, mainContainer *corev1.Container
 
 	hasEnv := false
 	for _, e := range mainContainer.Env {
-		if e.Name == "GMS_SOCKET_DIR" {
+		if e.Name == EnvSocketDir {
 			hasEnv = true
 			break
 		}
 	}
 	if !hasEnv {
-		mainContainer.Env = append(mainContainer.Env,
-			corev1.EnvVar{Name: "TMPDIR", Value: SharedMountPath},
-			corev1.EnvVar{Name: "GMS_SOCKET_DIR", Value: SharedMountPath},
-		)
+		mainContainer.Env = append(mainContainer.Env, corev1.EnvVar{Name: EnvSocketDir, Value: SharedMountPath})
 	}
 }
 
-// serverContainer builds the base GMS server container without init-specific
-// fields (RestartPolicy, StartupProbe). Callers add those as needed.
-func serverContainer(image string) corev1.Container {
+// Container builds a GMS container with the shared socket volume, env, and
+// DRA claim. Used for the server, loader, and saver.
+func Container(name, module, image string) corev1.Container {
 	return corev1.Container{
-		Name:    ServerContainerName,
+		Name:    name,
 		Image:   image,
-		Command: []string{"python3", "-m", serverSidecarModule},
+		Command: []string{"python3", "-m", module},
 		Env: []corev1.EnvVar{
-			{Name: "TMPDIR", Value: SharedMountPath},
-			{Name: "GMS_SOCKET_DIR", Value: SharedMountPath},
+			{Name: EnvSocketDir, Value: SharedMountPath},
 		},
 		VolumeMounts: []corev1.VolumeMount{
 			{Name: SharedVolumeName, MountPath: SharedMountPath},
 		},
+		Resources: corev1.ResourceRequirements{
+			Claims: []corev1.ResourceClaim{{Name: dra.ClaimName}},
+		},
 	}
 }
-
-func copyDeviceClaims(src *corev1.Container, dst *corev1.Container) {
-	if src == nil || dst == nil || len(src.Resources.Claims) == 0 {
-		return
-	}
-	claims := make([]corev1.ResourceClaim, len(src.Resources.Claims))
-	copy(claims, src.Resources.Claims)
-	dst.Resources.Claims = append(dst.Resources.Claims, claims...)
-}

deploy/operator/internal/gms/runtime_test.go → deploy/operator/internal/gms/gms_test.go

 /*
- * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -9,6 +9,7 @@ import (
 	"path/filepath"
 	"testing"
 
+	"github.com/ai-dynamo/dynamo/deploy/operator/internal/dra"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	corev1 "k8s.io/api/core/v1"
@@ -20,71 +21,29 @@ func TestEnsureServerSidecar(t *testing.T) {
 			Name:  "main",
 			Image: "test-image:latest",
 			Resources: corev1.ResourceRequirements{
-				Claims: []corev1.ResourceClaim{{Name: DRAClaimName}},
+				Claims: []corev1.ResourceClaim{{Name: dra.ClaimName}},
 			},
 		}},
 	}
 
 	EnsureServerSidecar(podSpec, &podSpec.Containers[0])
 
-	require.Len(t, podSpec.Containers, 1)
 	require.Len(t, podSpec.InitContainers, 1)
-
-	main := &podSpec.Containers[0]
 	server := &podSpec.InitContainers[0]
 
 	assert.Equal(t, ServerContainerName, server.Name)
-	assert.Equal(t, []string{"python3", "-m", serverSidecarModule}, server.Command)
-	assert.Equal(t, SharedMountPath, envValue(t, main, "TMPDIR"))
-	assert.Equal(t, SharedMountPath, envValue(t, main, "GMS_SOCKET_DIR"))
-	assert.Equal(t, SharedMountPath, envValue(t, server, "TMPDIR"))
-	assert.Equal(t, SharedMountPath, envValue(t, server, "GMS_SOCKET_DIR"))
-
+	assert.Equal(t, []string{"python3", "-m", ServerModule}, server.Command)
 	assert.Equal(t, corev1.ContainerRestartPolicyAlways, *server.RestartPolicy)
 	require.NotNil(t, server.StartupProbe)
 	assert.Equal(t, []string{"test", "-f", filepath.Join(SharedMountPath, readyFile)},
 		server.StartupProbe.Exec.Command)
-	assert.Equal(t, int32(1), server.StartupProbe.PeriodSeconds)
-	assert.Equal(t, int32(300), server.StartupProbe.FailureThreshold)
 
-	// DRA claim copied from main
+	// DRA claim on server
 	assert.Len(t, server.Resources.Claims, 1)
-	assert.Equal(t, DRAClaimName, server.Resources.Claims[0].Name)
-}
+	assert.Equal(t, dra.ClaimName, server.Resources.Claims[0].Name)
 
-func TestBuildServerContainer(t *testing.T) {
-	podSpec := &corev1.PodSpec{
-		Containers: []corev1.Container{{
-			Name:  "main",
-			Image: "test-image:latest",
-			Resources: corev1.ResourceRequirements{
-				Claims: []corev1.ResourceClaim{{Name: DRAClaimName}},
-			},
-		}},
-	}
-
-	server := BuildServerContainer(podSpec, &podSpec.Containers[0])
-
-	// Should not be added to init containers
-	assert.Empty(t, podSpec.InitContainers)
-
-	assert.Equal(t, ServerContainerName, server.Name)
-	assert.Equal(t, []string{"python3", "-m", serverSidecarModule}, server.Command)
-
-	// No init-specific fields
-	assert.Nil(t, server.RestartPolicy)
-	assert.Nil(t, server.StartupProbe)
-
-	// DRA claim copied from main
-	assert.Len(t, server.Resources.Claims, 1)
-	assert.Equal(t, DRAClaimName, server.Resources.Claims[0].Name)
-
-	// Shared volume and env should be set on main
-	main := &podSpec.Containers[0]
-	assert.Equal(t, SharedMountPath, envValue(t, main, "TMPDIR"))
-	assert.Equal(t, SharedMountPath, envValue(t, main, "GMS_SOCKET_DIR"))
-
-	// Shared volume should exist
+	// Shared volume and env on main
+	assert.Equal(t, SharedMountPath, envValue(t, &podSpec.Containers[0], "GMS_SOCKET_DIR"))
 	var hasVolume bool
 	for _, v := range podSpec.Volumes {
 		if v.Name == SharedVolumeName {
@@ -94,54 +53,27 @@ func TestBuildServerContainer(t *testing.T) {
 	assert.True(t, hasVolume)
 }
 
-func TestEnsureServerSidecarDoesNotAddCheckpointControl(t *testing.T) {
+func TestEnsureServerSidecarIdempotent(t *testing.T) {
 	podSpec := &corev1.PodSpec{
 		Containers: []corev1.Container{{Name: "main", Image: "test:latest"}},
 	}
-
+	EnsureServerSidecar(podSpec, &podSpec.Containers[0])
 	EnsureServerSidecar(podSpec, &podSpec.Containers[0])
 
-	for _, volume := range podSpec.Volumes {
-		if volume.Name == ControlVolumeName {
-			t.Fatal("runtime shaping should not add checkpoint control volume")
-		}
-	}
-	server := FindServerContainer(podSpec)
-	require.NotNil(t, server)
-	for _, env := range server.Env {
-		if env.Name == "GMS_CONTROL_DIR" {
-			t.Fatal("server should not have checkpoint control env")
-		}
-	}
+	assert.Len(t, podSpec.InitContainers, 1)
 }
 
-func TestEnsureServerSidecarIdempotent(t *testing.T) {
+func TestEnsureServerSidecarDoesNotAddCheckpointControl(t *testing.T) {
 	podSpec := &corev1.PodSpec{
 		Containers: []corev1.Container{{Name: "main", Image: "test:latest"}},
 	}
-
-	EnsureServerSidecar(podSpec, &podSpec.Containers[0])
 	EnsureServerSidecar(podSpec, &podSpec.Containers[0])
 
-	assert.Len(t, podSpec.InitContainers, 1)
-	volumeCount := 0
 	for _, v := range podSpec.Volumes {
-		if v.Name == SharedVolumeName {
-			volumeCount++
+		if v.Name == "gms-control" {
+			t.Fatal("should not add checkpoint control volume")
 		}
 	}
-	assert.Equal(t, 1, volumeCount)
-}
-
-func TestFindServerContainer(t *testing.T) {
-	podSpec := &corev1.PodSpec{
-		Containers: []corev1.Container{{Name: "main", Image: "test:latest"}},
-	}
-	assert.Nil(t, FindServerContainer(podSpec))
-
-	EnsureServerSidecar(podSpec, &podSpec.Containers[0])
-	assert.NotNil(t, FindServerContainer(podSpec))
-	assert.Equal(t, ServerContainerName, FindServerContainer(podSpec).Name)
 }
 
 func envValue(t *testing.T, container *corev1.Container, name string) string {

deploy/snapshot/internal/controller/util.go

@@ -48,17 +48,13 @@ func podFromInformerObj(obj interface{}) (*corev1.Pod, bool) {
 	return pod, ok
 }
 
+// resolveMainContainerName returns the name of the workload container, which
+// is always Containers[0]. GMS sidecars are appended after the workload.
 func resolveMainContainerName(pod *corev1.Pod) string {
-	containerName := ""
-	for _, c := range pod.Spec.Containers {
-		if c.Name == "main" {
-			return c.Name
+	if len(pod.Spec.Containers) == 0 {
+		return ""
 	}
-		if containerName == "" {
-			containerName = c.Name
-		}
-	}
-	return containerName
+	return pod.Spec.Containers[0].Name
 }
 
 func isPodReady(pod *corev1.Pod) bool {

deploy/snapshot/protocol/checkpoint.go

@@ -5,6 +5,7 @@ package protocol
 
 import (
 	"fmt"
+	"strings"
 
 	batchv1 "k8s.io/api/batch/v1"
 	corev1 "k8s.io/api/core/v1"
@@ -56,10 +57,10 @@ func NewCheckpointJob(podTemplate *corev1.PodTemplateSpec, opts CheckpointJobOpt
 		EnsureLocalhostSeccompProfile(&podTemplate.Spec, opts.SeccompProfile)
 	}
 	if opts.WrapLaunchJob {
-		container, err := ResolveCheckpointWorkerContainer(&podTemplate.Spec)
-		if err != nil {
-			return nil, err
+		if len(podTemplate.Spec.Containers) == 0 {
+			return nil, fmt.Errorf("checkpoint job requires at least one container")
 		}
+		container := &podTemplate.Spec.Containers[0]
 		if len(container.Command) == 0 {
 			return nil, fmt.Errorf("checkpoint job requires container.command when cuda-checkpoint launch-job wrapping is enabled")
 		}
@@ -157,6 +158,17 @@ func EnsureLocalhostSeccompProfile(podSpec *corev1.PodSpec, profile string) {
 }
 
 func wrapWithCudaCheckpointLaunchJob(command []string, args []string) ([]string, []string) {
+	// Unwrap "/bin/sh -c <single-string>" so cuda-checkpoint launches the
+	// actual process directly. Otherwise sh sits between cuda-checkpoint and
+	// the real process and swallows SIGUSR1.
+	if len(command) >= 2 && command[len(command)-1] == "-c" && len(args) == 1 {
+		shell := command[:len(command)-1] // e.g. ["/bin/sh"] — discarded
+		_ = shell
+		parts := strings.Fields(args[0])
+		command = parts[:1] // e.g. ["python3"]
+		args = parts[1:]    // e.g. ["-m", "dynamo.vllm", "--model", ...]
+	}
+
 	wrappedArgs := make([]string, 0, len(command)+len(args)+1)
 	wrappedArgs = append(wrappedArgs, "--launch-job")
 	wrappedArgs = append(wrappedArgs, command...)

deploy/snapshot/protocol/checkpoint_container.go

-// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-// SPDX-License-Identifier: Apache-2.0
-
-package protocol
-
-import (
-	"fmt"
-
-	corev1 "k8s.io/api/core/v1"
-)
-
-const checkpointWorkerContainerName = "main"
-
-func ResolveCheckpointWorkerContainer(podSpec *corev1.PodSpec) (*corev1.Container, error) {
-	if podSpec == nil || len(podSpec.Containers) == 0 {
-		return nil, fmt.Errorf("checkpoint job requires at least one container")
-	}
-	if len(podSpec.Containers) == 1 {
-		return &podSpec.Containers[0], nil
-	}
-
-	for i := range podSpec.Containers {
-		if podSpec.Containers[i].Name == checkpointWorkerContainerName {
-			return &podSpec.Containers[i], nil
-		}
-	}
-
-	return nil, fmt.Errorf("checkpoint job requires a container named %q when multiple containers are present", checkpointWorkerContainerName)
-}