ci: Merge change from v2 for release 2.8.3

scheduler/pkg/agent/server.go

@@ -390,7 +390,7 @@ func (s *Server) Subscribe(request *pb.AgentSubscribeRequest, stream pb.AgentSer
 	}
 	s.mutex.Unlock()
 
-	err := s.syncMessage(request, stream)
+	err := s.syncMessage(request)
 	if err != nil {
 		return err
 	}
@@ -421,7 +421,7 @@ func (s *Server) StopAgentStreams() {
 	}
 }
 
-func (s *Server) syncMessage(request *pb.AgentSubscribeRequest, stream pb.AgentService_SubscribeServer) error {
+func (s *Server) syncMessage(request *pb.AgentSubscribeRequest) error {
 	s.logger.Debugf("Add Server Replica %+v with config %+v", request, request.ReplicaConfig)
 	err := s.store.AddServerReplica(request)
 	if err != nil {

scheduler/pkg/server/server.go

@@ -431,29 +431,35 @@ func (s *SchedulerServer) ServerStatus(
 }
 
 func createServerStatusResponse(s *store.ServerSnapshot) *pb.ServerStatusResponse {
+	// note we dont count draining replicas in available replicas
+
 	resp := &pb.ServerStatusResponse{
-		ServerName:        s.Name,
-		AvailableReplicas: int32(len(s.Replicas)),
-		ExpectedReplicas:  int32(s.ExpectedReplicas),
-		KubernetesMeta:    s.KubernetesMeta,
+		ServerName:       s.Name,
+		ExpectedReplicas: int32(s.ExpectedReplicas),
+		KubernetesMeta:   s.KubernetesMeta,
 	}
 
-	var totalModels int32
+	totalModels := int32(0)
+	numAvailableServerReplicas := int32(0)
 	for _, replica := range s.Replicas {
-		numLoadedModelsOnReplica := int32(replica.GetNumLoadedModels())
-		resp.Resources = append(
-			resp.Resources,
-			&pb.ServerReplicaResources{
-				ReplicaIdx:           uint32(replica.GetReplicaIdx()),
-				TotalMemoryBytes:     replica.GetMemory(),
-				AvailableMemoryBytes: replica.GetAvailableMemory(),
-				NumLoadedModels:      numLoadedModelsOnReplica,
-				OverCommitPercentage: replica.GetOverCommitPercentage(),
-			},
-		)
-		totalModels += numLoadedModelsOnReplica
+		if !replica.GetIsDraining() {
+			numLoadedModelsOnReplica := int32(replica.GetNumLoadedModels())
+			resp.Resources = append(
+				resp.Resources,
+				&pb.ServerReplicaResources{
+					ReplicaIdx:           uint32(replica.GetReplicaIdx()),
+					TotalMemoryBytes:     replica.GetMemory(),
+					AvailableMemoryBytes: replica.GetAvailableMemory(),
+					NumLoadedModels:      numLoadedModelsOnReplica,
+					OverCommitPercentage: replica.GetOverCommitPercentage(),
+				},
+			)
+			totalModels += numLoadedModelsOnReplica
+			numAvailableServerReplicas++
+		}
 	}
 	resp.NumLoadedModelReplicas = totalModels
+	resp.AvailableReplicas = numAvailableServerReplicas
 
 	return resp
 }

scheduler/pkg/server/server_status_test.go

@@ -172,19 +172,101 @@ func TestModelsStatusEvents(t *testing.T) {
 }
 
 func TestServersStatusStream(t *testing.T) {
+	type serverReplicaRequest struct {
+		request  *pba.AgentSubscribeRequest
+		draining bool
+	}
+
 	g := NewGomegaWithT(t)
 	type test struct {
 		name    string
-		loadReq *pba.AgentSubscribeRequest
+		loadReq []serverReplicaRequest
 		server  *SchedulerServer
 		err     bool
 	}
 
 	tests := []test{
 		{
-			name: "server ok",
-			loadReq: &pba.AgentSubscribeRequest{
-				ServerName: "foo",
+			name: "server ok - 1 empty replica",
+			loadReq: []serverReplicaRequest{
+				{
+					request: &pba.AgentSubscribeRequest{
+						ServerName: "foo",
+					},
+				},
+			},
+			server: &SchedulerServer{
+				modelStore: store.NewMemoryStore(log.New(), store.NewLocalSchedulerStore(), nil),
+				logger:     log.New(),
+				timeout:    10 * time.Millisecond,
+			},
+		},
+		{
+			name: "server ok - multiple replicas",
+			loadReq: []serverReplicaRequest{
+				{
+					request: &pba.AgentSubscribeRequest{
+						ServerName: "foo",
+						ReplicaIdx: 0,
+						LoadedModels: []*pba.ModelVersion{
+							{
+								Model: &pb.Model{
+									Meta: &pb.MetaData{Name: "foo-model"},
+								},
+							},
+						},
+					},
+				},
+				{
+					request: &pba.AgentSubscribeRequest{
+						ServerName: "foo",
+						ReplicaIdx: 1,
+						LoadedModels: []*pba.ModelVersion{
+							{
+								Model: &pb.Model{
+									Meta: &pb.MetaData{Name: "foo-model"},
+								},
+							},
+						},
+					},
+				},
+			},
+			server: &SchedulerServer{
+				modelStore: store.NewMemoryStore(log.New(), store.NewLocalSchedulerStore(), nil),
+				logger:     log.New(),
+				timeout:    10 * time.Millisecond,
+			},
+		},
+		{
+			name: "server ok - multiple replicas with draining",
+			loadReq: []serverReplicaRequest{
+				{
+					request: &pba.AgentSubscribeRequest{
+						ServerName: "foo",
+						ReplicaIdx: 0,
+						LoadedModels: []*pba.ModelVersion{
+							{
+								Model: &pb.Model{
+									Meta: &pb.MetaData{Name: "foo-model"},
+								},
+							},
+						},
+					},
+				},
+				{
+					request: &pba.AgentSubscribeRequest{
+						ServerName: "foo",
+						ReplicaIdx: 1,
+						LoadedModels: []*pba.ModelVersion{
+							{
+								Model: &pb.Model{
+									Meta: &pb.MetaData{Name: "foo-model"},
+								},
+							},
+						},
+					},
+					draining: true,
+				},
 			},
 			server: &SchedulerServer{
 				modelStore: store.NewMemoryStore(log.New(), store.NewLocalSchedulerStore(), nil),
@@ -194,8 +276,12 @@ func TestServersStatusStream(t *testing.T) {
 		},
 		{
 			name: "timeout",
-			loadReq: &pba.AgentSubscribeRequest{
-				ServerName: "foo",
+			loadReq: []serverReplicaRequest{
+				{
+					request: &pba.AgentSubscribeRequest{
+						ServerName: "foo",
+					},
+				},
 			},
 			server: &SchedulerServer{
 				modelStore: store.NewMemoryStore(log.New(), store.NewLocalSchedulerStore(), nil),
@@ -208,9 +294,20 @@ func TestServersStatusStream(t *testing.T) {
 
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
+			expectedReplicas := int32(0)
+			expectedNumLoadedModelReplicas := int32(0)
 			if test.loadReq != nil {
-				err := test.server.modelStore.AddServerReplica(test.loadReq)
-				g.Expect(err).To(BeNil())
+				for _, r := range test.loadReq {
+					err := test.server.modelStore.AddServerReplica(r.request)
+					g.Expect(err).To(BeNil())
+					if !r.draining {
+						expectedReplicas++
+						expectedNumLoadedModelReplicas += int32(len(r.request.LoadedModels))
+					} else {
+						server, _ := test.server.modelStore.GetServer("foo", true, false)
+						server.Replicas[int(r.request.ReplicaIdx)].SetIsDraining()
+					}
+				}
 			}
 
 			stream := newStubServerStatusServer(1, 5*time.Millisecond)
@@ -230,6 +327,8 @@ func TestServersStatusStream(t *testing.T) {
 
 				g.Expect(ssr).ToNot(BeNil())
 				g.Expect(ssr.ServerName).To(Equal("foo"))
+				g.Expect(ssr.GetAvailableReplicas()).To(Equal(expectedReplicas))
+				g.Expect(ssr.NumLoadedModelReplicas).To(Equal(expectedNumLoadedModelReplicas))
 			}
 		})
 	}