Skip to content

Commit ab95deb

Browse files
visheshtanksalevisheshtanksale
committed
Adding unit test for universal NIM support
Signed-off-by: Vishesh Tanksale <vtanksale@nvidia.com>
1 parent 4dcedf0 commit ab95deb

2 files changed

Lines changed: 347 additions & 0 deletions

File tree

‎internal/controller/nimcache_controller_test.go‎

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -732,6 +732,82 @@ var _ = Describe("NIMCache Controller", func() {
732732
Expect(err).To(HaveOccurred())
733733
Expect(err.Error()).To(ContainSubstring("model_manifest.yaml not found in ConfigMap"))
734734
})
735+
736+
It("should construct a job with NGC ModelEndpoint for create-model-store", func() {
737+
modelEndpoint := "https://api.ngc.nvidia.com/v2/models/nvidia/nim-llama2-7b/versions/1.0.0"
738+
nimCache := &appsv1alpha1.NIMCache{
739+
ObjectMeta: metav1.ObjectMeta{
740+
Name: "test-nimcache",
741+
Namespace: "default",
742+
},
743+
Spec: appsv1alpha1.NIMCacheSpec{
744+
Source: appsv1alpha1.NIMSource{
745+
NGC: &appsv1alpha1.NGCSource{
746+
ModelPuller: "nvcr.io/nim:test",
747+
PullSecret: "my-secret",
748+
ModelEndpoint: &modelEndpoint,
749+
},
750+
},
751+
Storage: appsv1alpha1.NIMCacheStorage{
752+
PVC: appsv1alpha1.PersistentVolumeClaim{
753+
SubPath: "test-subpath",
754+
},
755+
},
756+
Resources: appsv1alpha1.Resources{
757+
CPU: resource.MustParse("500m"),
758+
Memory: resource.MustParse("1Gi"),
759+
},
760+
},
761+
}
762+
763+
job, err := reconciler.constructJob(context.TODO(), nimCache, k8sutil.K8s)
764+
Expect(err).ToNot(HaveOccurred())
765+
766+
Expect(job.Name).To(Equal(getJobName(nimCache)))
767+
Expect(job.Spec.Template.Spec.Containers).To(HaveLen(1))
768+
769+
container := job.Spec.Template.Spec.Containers[0]
770+
Expect(container.Name).To(Equal(NIMCacheContainerName))
771+
Expect(container.Image).To(Equal("nvcr.io/nim:test"))
772+
Expect(container.Command).To(Equal([]string{"create-model-store"}))
773+
Expect(container.Args).To(Equal([]string{"--model-repo", modelEndpoint, "--model-store", "/model-store"}))
774+
775+
// Check environment variables
776+
Expect(container.Env).To(ContainElement(corev1.EnvVar{
777+
Name: "NIM_CACHE_PATH",
778+
Value: "/model-store",
779+
}))
780+
781+
// Check volume mounts
782+
Expect(container.VolumeMounts).To(ContainElement(corev1.VolumeMount{
783+
Name: "nim-cache-volume",
784+
MountPath: "/model-store",
785+
SubPath: "test-subpath",
786+
}))
787+
788+
// Check resources
789+
Expect(container.Resources.Limits.Cpu().String()).To(Equal("500m"))
790+
Expect(container.Resources.Limits.Memory().String()).To(Equal("1Gi"))
791+
Expect(container.Resources.Requests.Cpu().String()).To(Equal("500m"))
792+
Expect(container.Resources.Requests.Memory().String()).To(Equal("1Gi"))
793+
794+
// Check security context
795+
Expect(container.SecurityContext).NotTo(BeNil())
796+
Expect(*container.SecurityContext.AllowPrivilegeEscalation).To(BeFalse())
797+
Expect(container.SecurityContext.Capabilities.Drop).To(ContainElement(corev1.Capability("ALL")))
798+
Expect(*container.SecurityContext.RunAsNonRoot).To(BeTrue())
799+
Expect(container.SecurityContext.RunAsGroup).To(Equal(nimCache.GetGroupID()))
800+
Expect(container.SecurityContext.RunAsUser).To(Equal(nimCache.GetUserID()))
801+
802+
// Check termination message settings
803+
Expect(container.TerminationMessagePath).To(Equal("/dev/termination-log"))
804+
Expect(container.TerminationMessagePolicy).To(Equal(corev1.TerminationMessageFallbackToLogsOnError))
805+
806+
// Check image pull secrets
807+
Expect(job.Spec.Template.Spec.ImagePullSecrets).To(ContainElement(corev1.LocalObjectReference{
808+
Name: "my-secret",
809+
}))
810+
})
735811
})
736812

737813
Context("when error reconciling NIMCache resource", func() {

‎internal/controller/platform/standalone/nimservice_test.go‎

Lines changed: 271 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1292,6 +1292,277 @@ var _ = Describe("NIMServiceReconciler for a standalone platform", func() {
12921292

12931293
})
12941294
})
1295+
1296+
It("should add NIM_MODEL_NAME environment variable when NIMCache is Universal NIM", func() {
1297+
// Create a NIMCache with ModelEndpoint to make it Universal NIM
1298+
modelEndpoint := "https://api.ngc.nvidia.com/v2/models/nvidia/nim-llama2-7b/versions/1.0.0"
1299+
universalNimCache := &appsv1alpha1.NIMCache{
1300+
ObjectMeta: metav1.ObjectMeta{
1301+
Name: "test-universal-nimcache",
1302+
Namespace: "default",
1303+
},
1304+
Spec: appsv1alpha1.NIMCacheSpec{
1305+
Source: appsv1alpha1.NIMSource{
1306+
NGC: &appsv1alpha1.NGCSource{
1307+
ModelPuller: "test-container",
1308+
PullSecret: "my-secret",
1309+
ModelEndpoint: &modelEndpoint,
1310+
},
1311+
},
1312+
Storage: appsv1alpha1.NIMCacheStorage{
1313+
PVC: appsv1alpha1.PersistentVolumeClaim{
1314+
Create: ptr.To[bool](true),
1315+
StorageClass: "standard",
1316+
Size: "1Gi",
1317+
SubPath: "subPath",
1318+
},
1319+
},
1320+
},
1321+
Status: appsv1alpha1.NIMCacheStatus{
1322+
State: appsv1alpha1.NimCacheStatusReady,
1323+
PVC: "test-universal-nimcache-pvc",
1324+
Profiles: []appsv1alpha1.NIMProfile{{
1325+
Name: "test-profile",
1326+
Config: map[string]string{"tp": "4"}},
1327+
},
1328+
},
1329+
}
1330+
Expect(client.Create(context.TODO(), universalNimCache)).To(Succeed())
1331+
1332+
// Create PVC for the universal NIMCache
1333+
universalPVC := &corev1.PersistentVolumeClaim{
1334+
ObjectMeta: metav1.ObjectMeta{
1335+
Name: "test-universal-nimcache-pvc",
1336+
Namespace: "default",
1337+
},
1338+
}
1339+
Expect(client.Create(context.TODO(), universalPVC)).To(Succeed())
1340+
1341+
// Create a new NIMService instance for this test
1342+
testNimService := &appsv1alpha1.NIMService{
1343+
ObjectMeta: metav1.ObjectMeta{
1344+
Name: "test-universal-nimservice",
1345+
Namespace: "default",
1346+
},
1347+
Spec: appsv1alpha1.NIMServiceSpec{
1348+
Labels: map[string]string{"app": "test-app"},
1349+
Annotations: map[string]string{"annotation-key": "annotation-value"},
1350+
Image: appsv1alpha1.Image{
1351+
Repository: "nvcr.io/nvidia/nim-llm",
1352+
Tag: "v0.1.0",
1353+
PullPolicy: "IfNotPresent",
1354+
PullSecrets: []string{"ngc-secret"},
1355+
},
1356+
Storage: appsv1alpha1.NIMServiceStorage{
1357+
PVC: appsv1alpha1.PersistentVolumeClaim{
1358+
Name: "test-pvc",
1359+
StorageClass: "standard",
1360+
Size: "1Gi",
1361+
Create: ptr.To[bool](true),
1362+
},
1363+
NIMCache: appsv1alpha1.NIMCacheVolSpec{
1364+
Name: "test-universal-nimcache",
1365+
},
1366+
},
1367+
Env: []corev1.EnvVar{
1368+
{
1369+
Name: "custom-env",
1370+
Value: "custom-value",
1371+
},
1372+
},
1373+
Resources: &corev1.ResourceRequirements{
1374+
Limits: corev1.ResourceList{
1375+
corev1.ResourceCPU: resource.MustParse("500m"),
1376+
corev1.ResourceMemory: resource.MustParse("128Mi"),
1377+
},
1378+
Requests: corev1.ResourceList{
1379+
corev1.ResourceCPU: resource.MustParse("250m"),
1380+
corev1.ResourceMemory: resource.MustParse("64Mi"),
1381+
},
1382+
},
1383+
NodeSelector: map[string]string{"disktype": "ssd"},
1384+
Tolerations: []corev1.Toleration{{
1385+
Key: "key1",
1386+
Operator: corev1.TolerationOpEqual,
1387+
Value: "value1",
1388+
Effect: corev1.TaintEffectNoSchedule,
1389+
}},
1390+
Expose: appsv1alpha1.Expose{
1391+
Service: appsv1alpha1.Service{Type: corev1.ServiceTypeLoadBalancer, Port: ptr.To[int32](8123), Annotations: map[string]string{"annotation-key-specific": "service"}},
1392+
},
1393+
RuntimeClassName: "nvidia",
1394+
},
1395+
}
1396+
1397+
namespacedName := types.NamespacedName{Name: testNimService.Name, Namespace: testNimService.Namespace}
1398+
Expect(client.Create(context.TODO(), testNimService)).To(Succeed())
1399+
1400+
result, err := reconciler.reconcileNIMService(context.TODO(), testNimService)
1401+
Expect(err).NotTo(HaveOccurred())
1402+
Expect(result).To(Equal(ctrl.Result{}))
1403+
1404+
// Deployment should be created
1405+
deployment := &appsv1.Deployment{}
1406+
err = client.Get(context.TODO(), namespacedName, deployment)
1407+
Expect(err).NotTo(HaveOccurred())
1408+
Expect(deployment.Name).To(Equal(testNimService.GetName()))
1409+
Expect(deployment.Namespace).To(Equal(testNimService.GetNamespace()))
1410+
1411+
// Verify that NIM_MODEL_NAME environment variable is added
1412+
container := deployment.Spec.Template.Spec.Containers[0]
1413+
var nimModelNameEnv *corev1.EnvVar
1414+
for _, env := range container.Env {
1415+
if env.Name == "NIM_MODEL_NAME" {
1416+
nimModelNameEnv = &env
1417+
break
1418+
}
1419+
}
1420+
Expect(nimModelNameEnv).NotTo(BeNil(), "NIM_MODEL_NAME environment variable should be present")
1421+
Expect(nimModelNameEnv.Value).To(Equal("/model-store"), "NIM_MODEL_NAME should be set to /model-store")
1422+
1423+
// Verify that other environment variables are still present
1424+
var customEnv *corev1.EnvVar
1425+
for _, env := range container.Env {
1426+
if env.Name == "custom-env" {
1427+
customEnv = &env
1428+
break
1429+
}
1430+
}
1431+
Expect(customEnv).NotTo(BeNil(), "Custom environment variables should still be present")
1432+
Expect(customEnv.Value).To(Equal("custom-value"))
1433+
})
1434+
1435+
It("should not add NIM_MODEL_NAME environment variable when NIMCache is not Universal NIM", func() {
1436+
// Create a regular NIMCache (not Universal NIM)
1437+
regularNimCache := &appsv1alpha1.NIMCache{
1438+
ObjectMeta: metav1.ObjectMeta{
1439+
Name: "test-regular-nimcache",
1440+
Namespace: "default",
1441+
},
1442+
Spec: appsv1alpha1.NIMCacheSpec{
1443+
Source: appsv1alpha1.NIMSource{
1444+
NGC: &appsv1alpha1.NGCSource{
1445+
ModelPuller: "test-container",
1446+
PullSecret: "my-secret",
1447+
// No ModelEndpoint, so it's not Universal NIM
1448+
},
1449+
},
1450+
Storage: appsv1alpha1.NIMCacheStorage{
1451+
PVC: appsv1alpha1.PersistentVolumeClaim{
1452+
Create: ptr.To[bool](true),
1453+
StorageClass: "standard",
1454+
Size: "1Gi",
1455+
SubPath: "subPath",
1456+
},
1457+
},
1458+
},
1459+
Status: appsv1alpha1.NIMCacheStatus{
1460+
State: appsv1alpha1.NimCacheStatusReady,
1461+
PVC: "test-regular-nimcache-pvc",
1462+
Profiles: []appsv1alpha1.NIMProfile{{
1463+
Name: "test-profile",
1464+
Config: map[string]string{"tp": "4"}},
1465+
},
1466+
},
1467+
}
1468+
Expect(client.Create(context.TODO(), regularNimCache)).To(Succeed())
1469+
1470+
// Create PVC for the regular NIMCache
1471+
regularPVC := &corev1.PersistentVolumeClaim{
1472+
ObjectMeta: metav1.ObjectMeta{
1473+
Name: "test-regular-nimcache-pvc",
1474+
Namespace: "default",
1475+
},
1476+
}
1477+
Expect(client.Create(context.TODO(), regularPVC)).To(Succeed())
1478+
1479+
// Create a new NIMService instance for this test
1480+
testNimService := &appsv1alpha1.NIMService{
1481+
ObjectMeta: metav1.ObjectMeta{
1482+
Name: "test-regular-nimservice",
1483+
Namespace: "default",
1484+
},
1485+
Spec: appsv1alpha1.NIMServiceSpec{
1486+
Labels: map[string]string{"app": "test-app"},
1487+
Annotations: map[string]string{"annotation-key": "annotation-value"},
1488+
Image: appsv1alpha1.Image{
1489+
Repository: "nvcr.io/nvidia/nim-llm",
1490+
Tag: "v0.1.0",
1491+
PullPolicy: "IfNotPresent",
1492+
PullSecrets: []string{"ngc-secret"},
1493+
},
1494+
Storage: appsv1alpha1.NIMServiceStorage{
1495+
PVC: appsv1alpha1.PersistentVolumeClaim{
1496+
Name: "test-pvc",
1497+
StorageClass: "standard",
1498+
Size: "1Gi",
1499+
Create: ptr.To[bool](true),
1500+
},
1501+
NIMCache: appsv1alpha1.NIMCacheVolSpec{
1502+
Name: "test-regular-nimcache",
1503+
},
1504+
},
1505+
Env: []corev1.EnvVar{
1506+
{
1507+
Name: "custom-env",
1508+
Value: "custom-value",
1509+
},
1510+
},
1511+
Resources: &corev1.ResourceRequirements{
1512+
Limits: corev1.ResourceList{
1513+
corev1.ResourceCPU: resource.MustParse("500m"),
1514+
corev1.ResourceMemory: resource.MustParse("128Mi"),
1515+
},
1516+
Requests: corev1.ResourceList{
1517+
corev1.ResourceCPU: resource.MustParse("250m"),
1518+
corev1.ResourceMemory: resource.MustParse("64Mi"),
1519+
},
1520+
},
1521+
NodeSelector: map[string]string{"disktype": "ssd"},
1522+
Tolerations: []corev1.Toleration{{
1523+
Key: "key1",
1524+
Operator: corev1.TolerationOpEqual,
1525+
Value: "value1",
1526+
Effect: corev1.TaintEffectNoSchedule,
1527+
}},
1528+
Expose: appsv1alpha1.Expose{
1529+
Service: appsv1alpha1.Service{Type: corev1.ServiceTypeLoadBalancer, Port: ptr.To[int32](8123), Annotations: map[string]string{"annotation-key-specific": "service"}},
1530+
},
1531+
RuntimeClassName: "nvidia",
1532+
},
1533+
}
1534+
1535+
namespacedName := types.NamespacedName{Name: testNimService.Name, Namespace: testNimService.Namespace}
1536+
Expect(client.Create(context.TODO(), testNimService)).To(Succeed())
1537+
1538+
result, err := reconciler.reconcileNIMService(context.TODO(), testNimService)
1539+
Expect(err).NotTo(HaveOccurred())
1540+
Expect(result).To(Equal(ctrl.Result{}))
1541+
1542+
// Deployment should be created
1543+
deployment := &appsv1.Deployment{}
1544+
err = client.Get(context.TODO(), namespacedName, deployment)
1545+
Expect(err).NotTo(HaveOccurred())
1546+
Expect(deployment.Name).To(Equal(testNimService.GetName()))
1547+
Expect(deployment.Namespace).To(Equal(testNimService.GetNamespace()))
1548+
1549+
// Verify that NIM_MODEL_NAME environment variable is NOT added
1550+
container := deployment.Spec.Template.Spec.Containers[0]
1551+
for _, env := range container.Env {
1552+
Expect(env.Name).NotTo(Equal("NIM_MODEL_NAME"), "NIM_MODEL_NAME environment variable should not be present for non-Universal NIM")
1553+
}
1554+
1555+
// Verify that other environment variables are still present
1556+
var customEnv *corev1.EnvVar
1557+
for _, env := range container.Env {
1558+
if env.Name == "custom-env" {
1559+
customEnv = &env
1560+
break
1561+
}
1562+
}
1563+
Expect(customEnv).NotTo(BeNil(), "Custom environment variables should still be present")
1564+
Expect(customEnv.Value).To(Equal("custom-value"))
1565+
})
12951566
})
12961567

12971568
Describe("getNIMModelEndpoints", func() {

0 commit comments

Comments
 (0)