Allow size of emptyDir volumes to be specified for NIMs Customizer

Signed-off-by: Sheng Lin <shelin@nvidia.com>

Author: shengnuo

api/apps/v1alpha1/nemo_customizer_types.go

@@ -31,6 +31,7 @@ import (
 	corev1 "k8s.io/api/core/v1"
 	networkingv1 "k8s.io/api/networking/v1"
 	rbacv1 "k8s.io/api/rbac/v1"
+	"k8s.io/apimachinery/pkg/api/resource"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/util/intstr"
 	"k8s.io/utils/ptr"
@@ -149,6 +150,8 @@ type TrainingConfig struct {
 	Tolerations []corev1.Toleration `json:"tolerations,omitempty"`
 	// PodAffinity for the training jobs
 	PodAffinity *corev1.PodAffinity `json:"podAffinity,omitempty"`
+	// SharedMemorySizeLimit sets the max size of the shared memory volume (emptyDir) used by the training jobs for fast model runtime I/O.
+	SharedMemorySizeLimit *resource.Quantity `json:"sharedMemorySizeLimit,omitempty"`
 	// Resources for the training jobs
 	Resources *corev1.ResourceRequirements `json:"resources,omitempty"`
 }

api/apps/v1alpha1/nimservice_types.go

@@ -29,6 +29,7 @@ import (
 	corev1 "k8s.io/api/core/v1"
 	networkingv1 "k8s.io/api/networking/v1"
 	rbacv1 "k8s.io/api/rbac/v1"
+	"k8s.io/apimachinery/pkg/api/resource"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/util/intstr"
 	"k8s.io/utils/ptr"
@@ -134,6 +135,8 @@ type NIMServiceList struct {
 // NIMServiceStorage defines the attributes of various storage targets used to store the model
 type NIMServiceStorage struct {
 	NIMCache NIMCacheVolSpec `json:"nimCache,omitempty"`
+	// SharedMemorySizeLimit sets the max size of the shared memory volume (emptyDir) used by NIMs for fast model runtime I/O.
+	SharedMemorySizeLimit *resource.Quantity `json:"sharedMemorySizeLimit,omitempty"`
 	// PersistentVolumeClaim is the pvc volume used for caching NIM
 	PVC PersistentVolumeClaim `json:"pvc,omitempty"`
 	// HostPath is the host path volume for caching NIM
@@ -448,7 +451,8 @@ func (n *NIMService) GetVolumes(modelPVC PersistentVolumeClaim) []corev1.Volume
 			Name: "dshm",
 			VolumeSource: corev1.VolumeSource{
 				EmptyDir: &corev1.EmptyDirVolumeSource{
-					Medium: corev1.StorageMediumMemory,
+					Medium:    corev1.StorageMediumMemory,
+					SizeLimit: n.Spec.Storage.SharedMemorySizeLimit,
 				},
 			},
 		},

internal/controller/nemocustomizer_controller.go

@@ -626,6 +626,12 @@ func (r *NemoCustomizerReconciler) addTrainingConfig(ctx context.Context, cfg ma
 	// Add PVC configuration
 	r.addWorkspacePVCConfig(ctx, trainingCfg, n)
 
+	emptyDir := map[string]interface{}{
+		"medium": "Memory",
+	}
+	if n.Spec.Training.SharedMemorySizeLimit != nil {
+		emptyDir["sizeLimit"] = n.Spec.Training.SharedMemorySizeLimit.String()
+	}
 	trainingCfg["volumes"] = []map[string]interface{}{
 		{
 			"name": "models",
@@ -635,10 +641,8 @@ func (r *NemoCustomizerReconciler) addTrainingConfig(ctx context.Context, cfg ma
 			},
 		},
 		{
-			"name": "dshm",
-			"emptyDir": map[string]interface{}{
-				"medium": "Memory",
-			},
+			"name":     "dshm",
+			"emptyDir": emptyDir,
 		},
 	}
 

manifests/deployment.yaml

@@ -121,6 +121,9 @@ spec:
         {{- if .EmptyDir }}
         emptyDir:
           medium: {{ .EmptyDir.Medium }}
+          {{- if .EmptyDir.SizeLimit }}
+          sizeLimit: {{ .EmptyDir.SizeLimit }}
+          {{- end }}
         {{- end }}
         {{- if .PersistentVolumeClaim }}
         persistentVolumeClaim: