Function TestUpdateModelState

scheduler/pkg/store/memory_test.go:989–1356

(t *testing.T)

}

func TestUpdateModelState(t *testing.T) {
	g := NewGomegaWithT(t)
	memBytes := uint64(1)

	type test struct {
		name                   string
		store                  *LocalSchedulerStore
		modelKey               string
		version                uint32
		serverKey              string
		replicaIdx             int
		expectedState          ModelReplicaState
		desiredState           ModelReplicaState
		availableMemory        uint64
		modelRuntimeInfo       *pb.ModelRuntimeInfo
		numModelVersionsLoaded int
		modelVersionLoaded     bool
		deleted                bool
		err                    bool
	}

	tests := []test{
		{
			name: "LoadedModel",
			store: &LocalSchedulerStore{
				models: map[string]*Model{"model": {
					versions: []*ModelVersion{
						{
							modelDefn: &pb.Model{ModelSpec: &pb.ModelSpec{MemoryBytes: &memBytes}},
							version:   1,
							replicas:  map[int]ReplicaStatus{},
						},
					},
				}},
				servers: map[string]*Server{
					"server": {
						name: "server",
						replicas: map[int]*ServerReplica{
							0: {loadedModels: map[ModelVersionID]bool{}, reservedMemory: memBytes, uniqueLoadedModels: map[string]bool{}},
							1: {loadedModels: map[ModelVersionID]bool{}, reservedMemory: memBytes, uniqueLoadedModels: map[string]bool{}},
						},
					},
				},
			},
			modelKey:               "model",
			version:                1,
			serverKey:              "server",
			replicaIdx:             0,
			expectedState:          ModelReplicaStateUnknown,
			desiredState:           Loaded,
			numModelVersionsLoaded: 1,
			modelVersionLoaded:     true,
			availableMemory:        20,
			modelRuntimeInfo:       &pb.ModelRuntimeInfo{ModelRuntimeInfo: &pb.ModelRuntimeInfo_Mlserver{Mlserver: &pb.MLServerModelSettings{ParallelWorkers: uint32(1)}}},
		},
		{
			name: "UnloadedModel",
			store: &LocalSchedulerStore{