Add unit test for multi node NIMService

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

Author: shengnuo

api/apps/v1alpha1/nimservice_types.go

@@ -99,8 +99,9 @@ type NIMServiceSpec struct {
 // TODO disable Scale when this is enabled
 type NimServiceMultiNodeConfig struct {
 	// +kubebuilder:validation:Minimum=1
-	// +kubebuilder:default:=1
+	// +kubebuilder:default:=2
 	// Workers specifies how many worker pods per multi-node replica to launch.
+	// +kubebuilder:validation:Minimum=2
 	Workers int `json:"workers,omitempty"`
 
 	// +kubebuilder:default:=1

internal/controller/platform/standalone/nimservice_test.go

@@ -49,6 +49,7 @@ import (
 	ctrl "sigs.k8s.io/controller-runtime"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 	"sigs.k8s.io/controller-runtime/pkg/client/fake"
+	lwsv1 "sigs.k8s.io/lws/api/leaderworkerset/v1"
 
 	appsv1alpha1 "github.com/NVIDIA/k8s-nim-operator/api/apps/v1alpha1"
 	"github.com/NVIDIA/k8s-nim-operator/internal/conditions"
@@ -135,6 +136,7 @@ var _ = Describe("NIMServiceReconciler for a standalone platform", func() {
 		Expect(networkingv1.AddToScheme(scheme)).To(Succeed())
 		Expect(corev1.AddToScheme(scheme)).To(Succeed())
 		Expect(monitoringv1.AddToScheme(scheme)).To(Succeed())
+		Expect(lwsv1.AddToScheme(scheme)).To(Succeed())
 
 		client = fake.NewClientBuilder().WithScheme(scheme).
 			WithStatusSubresource(&appsv1alpha1.NIMService{}).
@@ -817,6 +819,67 @@ var _ = Describe("NIMServiceReconciler for a standalone platform", func() {
 		})
 	})
 
+	Describe("LWS deployment for multi-node inferencing NIMService", func() {
+		AfterEach(func() {
+			lws := &lwsv1.LeaderWorkerSet{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "test-nimservice-lws",
+					Namespace: "default",
+				},
+			}
+			err := client.Delete(context.TODO(), lws)
+			Expect(err).NotTo(HaveOccurred())
+		})
+
+		It("should report ready when LWS is available", func() {
+			lws := &lwsv1.LeaderWorkerSet{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "test-nimservice-lws",
+					Namespace: "default",
+				},
+				Status: lwsv1.LeaderWorkerSetStatus{
+					Conditions: []metav1.Condition{
+						{
+							Type:   string(lwsv1.LeaderWorkerSetAvailable),
+							Status: metav1.ConditionTrue,
+						},
+					},
+				},
+			}
+			err := client.Create(context.TODO(), lws)
+			Expect(err).NotTo(HaveOccurred())
+			msg, ready, err := reconciler.isLeaderWorkerSetReady(context.TODO(), nimService)
+			Expect(err).ToNot(HaveOccurred())
+			Expect(ready).To(Equal(true))
+			Expect(msg).To(Equal(fmt.Sprintf("leaderworkerset %q is ready", lws.Name)))
+		})
+		It("should report not ready when LWS is not available", func() {
+			lws := &lwsv1.LeaderWorkerSet{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "test-nimservice-lws",
+					Namespace: "default",
+				},
+				Status: lwsv1.LeaderWorkerSetStatus{
+					Conditions: []metav1.Condition{
+						{
+							Type:   string(lwsv1.LeaderWorkerSetProgressing),
+							Status: metav1.ConditionTrue,
+						},
+						{
+							Type:   string(lwsv1.LeaderWorkerSetAvailable),
+							Status: metav1.ConditionFalse,
+						},
+					},
+				},
+			}
+			err := client.Create(context.TODO(), lws)
+			Expect(err).NotTo(HaveOccurred())
+			msg, ready, err := reconciler.isLeaderWorkerSetReady(context.TODO(), nimService)
+			Expect(err).ToNot(HaveOccurred())
+			Expect(ready).To(Equal(false))
+			Expect(msg).To(Equal(fmt.Sprintf("leaderworkerset %q is not ready", lws.Name)))
+		})
+	})
 	Describe("update model status on NIMService", func() {
 		BeforeEach(func() {
 			ingress := &networkingv1.Ingress{

internal/render/render_test.go

@@ -129,6 +129,101 @@ var _ = Describe("K8s Resources Rendering", func() {
 	templatesDir := filepath.Join(path.Dir(path.Dir(cwd)), "manifests")
 
 	Context("Rendering templates", func() {
+		It("should render LeaderWorkerSet template correctly", func() {
+			params := types.LeaderWorkerSetParams{
+				Name:        "test-lws",
+				Namespace:   "default",
+				Labels:      map[string]string{"app": "test-app"},
+				Annotations: map[string]string{"annotation-key": "annotation-value"},
+				Replicas:    3,
+				Image:       "nim-llm:latest",
+				LeaderVolumes: []corev1.Volume{
+					{
+						Name: "test-leader-volume",
+						VolumeSource: corev1.VolumeSource{
+							PersistentVolumeClaim: &corev1.PersistentVolumeClaimVolumeSource{
+								ClaimName: "test-leader-pvc",
+							},
+						},
+					},
+				},
+				WorkerVolumes: []corev1.Volume{
+					{
+						Name: "test-worker-volume",
+						VolumeSource: corev1.VolumeSource{
+							PersistentVolumeClaim: &corev1.PersistentVolumeClaimVolumeSource{
+								ClaimName: "test-worker-pvc",
+							},
+						},
+					},
+				},
+				LeaderVolumeMounts: []corev1.VolumeMount{
+					{
+						Name:      "test-leader-volume",
+						MountPath: "/data",
+					},
+				},
+				WorkerVolumeMounts: []corev1.VolumeMount{
+					{
+						Name:      "test-worker-volume",
+						MountPath: "/data",
+					},
+				},
+				LeaderEnvs: []corev1.EnvVar{
+					{
+						Name:  "LEADER_ENV_VAR",
+						Value: "value",
+					},
+				},
+				WorkerEnvs: []corev1.EnvVar{
+					{
+						Name:  "WORKER_ENV_VAR",
+						Value: "value",
+					},
+				},
+				Resources: &corev1.ResourceRequirements{
+					Limits: corev1.ResourceList{
+						corev1.ResourceCPU:    resource.MustParse("500m"),
+						corev1.ResourceMemory: resource.MustParse("128Mi"),
+					},
+					Requests: corev1.ResourceList{
+						corev1.ResourceCPU:    resource.MustParse("250m"),
+						corev1.ResourceMemory: resource.MustParse("64Mi"),
+					},
+				},
+				NodeSelector: map[string]string{"disktype": "ssd"},
+				Tolerations: []corev1.Toleration{
+					{
+						Key:      "key1",
+						Operator: corev1.TolerationOpExists,
+						Effect:   corev1.TaintEffectNoSchedule,
+					},
+				},
+			}
+
+			r := render.NewRenderer(templatesDir)
+			lws, err := r.LeaderWorkerSet(&params)
+			Expect(err).NotTo(HaveOccurred())
+			Expect(lws.Name).To(Equal("test-lws"))
+			Expect(lws.Namespace).To(Equal("default"))
+			Expect(lws.Labels["app"]).To(Equal("test-app"))
+			Expect(lws.Annotations["annotation-key"]).To(Equal("annotation-value"))
+			Expect(*lws.Spec.Replicas).To(Equal(int32(3)))
+			Expect(lws.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec.Containers[0].Name).To(Equal("nim-leader"))
+			Expect(lws.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec.Containers[0].Image).To(Equal("nim-llm:latest"))
+			Expect(lws.Spec.LeaderWorkerTemplate.WorkerTemplate.Spec.Containers[0].Name).To(Equal("nim-worker"))
+			Expect(lws.Spec.LeaderWorkerTemplate.WorkerTemplate.Spec.Containers[0].Image).To(Equal("nim-llm:latest"))
+			Expect(lws.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec.Volumes[0].Name).To(Equal("test-leader-volume"))
+			Expect(lws.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec.Volumes[0].VolumeSource.PersistentVolumeClaim.ClaimName).To(Equal("test-leader-pvc"))
+			Expect(lws.Spec.LeaderWorkerTemplate.WorkerTemplate.Spec.Volumes[0].Name).To(Equal("test-worker-volume"))
+			Expect(lws.Spec.LeaderWorkerTemplate.WorkerTemplate.Spec.Volumes[0].VolumeSource.PersistentVolumeClaim.ClaimName).To(Equal("test-worker-pvc"))
+			Expect(lws.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec.Containers[0].VolumeMounts[0].Name).To(Equal("test-leader-volume"))
+			Expect(lws.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec.Containers[0].VolumeMounts[0].MountPath).To(Equal("/data"))
+			Expect(lws.Spec.LeaderWorkerTemplate.LeaderTemplate.Spec.Containers[0].VolumeMounts[0].SubPath).To(Equal("subPath"))
+			Expect(lws.Spec.LeaderWorkerTemplate.WorkerTemplate.Spec.Containers[0].VolumeMounts[0].Name).To(Equal("test-worker-volume"))
+			Expect(lws.Spec.LeaderWorkerTemplate.WorkerTemplate.Spec.Containers[0].VolumeMounts[0].MountPath).To(Equal("/data"))
+			Expect(lws.Spec.LeaderWorkerTemplate.WorkerTemplate.Spec.Containers[0].VolumeMounts[0].SubPath).To(Equal("subPath"))
+		})
 		It("should render Deployment template correctly", func() {
 			params := types.DeploymentParams{
 				Name:          "test-deployment",