feat: Add LWS to Dynamo Operator (#998)

Co-authored-by: Rohan Varma <rohanv@rohanv-mlt.client.nvidia.com>
Co-authored-by: Julien Mancuso <jmancuso@nvidia.com>
Co-authored-by: julienmancuso <161955438+julienmancuso@users.noreply.github.com>

deploy/cloud/operator/internal/dynamo/graph_test.go

@@ -901,3 +901,47 @@ func TestGenerateDynamoComponentsDeployments(t *testing.T) {
 		})
 	}
 }
+
+func TestSetLwsAnnotations(t *testing.T) {
+	type args struct {
+		serviceArgs *ServiceArgs
+		deployment  *v1alpha1.DynamoComponentDeployment
+	}
+	tests := []struct {
+		name    string
+		args    args
+		wantErr bool
+		want    *v1alpha1.DynamoComponentDeployment
+	}{
+		{
+			name: "Test SetLwsAnnotations for 16 GPUs",
+			args: args{
+				serviceArgs: &ServiceArgs{
+					Resources: &Resources{
+						GPU: &[]string{"8"}[0],
+					},
+					TotalGpus: &[]int32{16}[0],
+				},
+				deployment: &v1alpha1.DynamoComponentDeployment{},
+			},
+			wantErr: false,
+			want: &v1alpha1.DynamoComponentDeployment{
+				Spec: v1alpha1.DynamoComponentDeploymentSpec{
+					DynamoComponentDeploymentSharedSpec: v1alpha1.DynamoComponentDeploymentSharedSpec{
+						Annotations: map[string]string{
+							"nvidia.com/deployment-type": "leader-worker",
+							"nvidia.com/lws-size":        "2",
+						},
+					},
+				},
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if err := SetLwsAnnotations(tt.args.serviceArgs, tt.args.deployment); (err != nil) != tt.wantErr {
+				t.Errorf("SetLwsAnnotations() error = %v, wantErr %v", err, tt.wantErr)
+			}
+		})
+	}
+}

deploy/sdk/src/dynamo/sdk/cli/allocator.py

@@ -178,7 +178,7 @@ class ResourceAllocator:
 
         # Check if service requires GPUs
         if "gpu" in (config.get("resources") or {}):
-            num_gpus = config["resources"]["gpu"]  # type: ignore
+            num_gpus = int(config["resources"]["gpu"])  # type: ignore
             logger.info(f"GPU requirement found: {num_gpus}")
 
             # Check if we have enough GPUs