Fix labelling panic for prepackaged servers (#1600)

SeldonIO · Mar 25, 2020 · 0caa9db · 0caa9db
1 parent 371655b
commit 0caa9db
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 11 deletions.
diff --git a/operator/apis/machinelearning/v1/seldondeployment_webhook.go b/operator/apis/machinelearning/v1/seldondeployment_webhook.go
@@ -168,6 +168,9 @@ func addDefaultsToGraph(pu *PredictiveUnit) {
 	if pu.Implementation == nil {
 		im := UNKNOWN_IMPLEMENTATION
 		pu.Implementation = &im
+	} else if IsPrepack(pu) {
+		ty := MODEL
+		pu.Type = &ty
 	}
 	for i := 0; i < len(pu.Children); i++ {
 		addDefaultsToGraph(&pu.Children[i])

diff --git a/operator/apis/machinelearning/v1/seldondeployment_webhook_test.go b/operator/apis/machinelearning/v1/seldondeployment_webhook_test.go
@@ -597,6 +597,32 @@ func TestMetricsPortAddedToTwoPrepacked(t *testing.T) {
 	g.Expect(pu.Endpoint.Type).To(Equal(REST))
 }
 
+func TestDefaultPrepackagedServerType(t *testing.T) {
+	g := NewGomegaWithT(t)
+	scheme := runtime.NewScheme()
+	C = fake.NewFakeClientWithScheme(scheme)
+	impl := PredictiveUnitImplementation(constants.PrePackedServerTensorflow)
+	spec := &SeldonDeploymentSpec{
+		Predictors: []PredictorSpec{
+			{
+				Name: "p1",
+				Graph: &PredictiveUnit{
+					Name:           "classifier",
+					Implementation: &impl,
+				},
+			},
+		},
+	}
+
+	spec.DefaultSeldonDeployment("mydep", "default")
+
+	// Graph
+	pu := GetPredictiveUnit(spec.Predictors[0].Graph, "classifier")
+	g.Expect(pu).ToNot(BeNil())
+	g.Expect(pu.Endpoint.Type).To(Equal(REST))
+	g.Expect(*pu.Type).To(Equal(MODEL))
+}
+
 func TestValidateSingleModel(t *testing.T) {
 	g := NewGomegaWithT(t)
 	spec := &SeldonDeploymentSpec{

diff --git a/operator/controllers/labels.go b/operator/controllers/labels.go
@@ -7,17 +7,19 @@ import (
 )
 
 func addLabelsToService(svc *corev1.Service, pu *machinelearningv1.PredictiveUnit, p machinelearningv1.PredictorSpec) {
-	switch *pu.Type {
-	case machinelearningv1.ROUTER:
-		svc.Labels[machinelearningv1.Label_router] = "true"
-	case machinelearningv1.COMBINER:
-		svc.Labels[machinelearningv1.Label_combiner] = "true"
-	case machinelearningv1.MODEL:
-		svc.Labels[machinelearningv1.Label_model] = "true"
-	case machinelearningv1.TRANSFORMER:
-		svc.Labels[machinelearningv1.Label_transformer] = "true"
-	case machinelearningv1.OUTPUT_TRANSFORMER:
-		svc.Labels[machinelearningv1.Label_output_transformer] = "true"
+	if pu.Type != nil {
+		switch *pu.Type {
+		case machinelearningv1.ROUTER:
+			svc.Labels[machinelearningv1.Label_router] = "true"
+		case machinelearningv1.COMBINER:
+			svc.Labels[machinelearningv1.Label_combiner] = "true"
+		case machinelearningv1.MODEL:
+			svc.Labels[machinelearningv1.Label_model] = "true"
+		case machinelearningv1.TRANSFORMER:
+			svc.Labels[machinelearningv1.Label_transformer] = "true"
+		case machinelearningv1.OUTPUT_TRANSFORMER:
+			svc.Labels[machinelearningv1.Label_output_transformer] = "true"
+		}
 	}
 	if p.Shadow != true && (p.Traffic >= 50 || p.Traffic == 0) {
 		svc.Labels[machinelearningv1.Label_default] = "true"