Function TestLoadModel

scheduler/pkg/server/server_test.go:51–357

(t *testing.T)

}

func TestLoadModel(t *testing.T) {
	g := NewGomegaWithT(t)

	createTestScheduler := func() (*SchedulerServer, *mockAgentHandler, *coordinator.EventHub) {
		logger := log.New()
		logger.SetLevel(log.WarnLevel)

		eventHub, err := coordinator.NewEventHub(logger)
		g.Expect(err).To(BeNil())

		schedulerStore := store.NewMemoryStore(logger, store.NewLocalSchedulerStore(), eventHub)
		experimentServer := experiment.NewExperimentServer(logger, eventHub, nil, nil)
		pipelineServer := pipeline.NewPipelineStore(logger, eventHub, schedulerStore)
		sync := synchroniser.NewSimpleSynchroniser(time.Duration(10 * time.Millisecond))
		scheduler := scheduler2.NewSimpleScheduler(
			logger,
			schedulerStore,
			scheduler2.DefaultSchedulerConfig(schedulerStore),
			sync,
			eventHub,
		)
		modelGwLoadBalancer := util.NewRingLoadBalancer(1)
		pipelineGwLoadBalancer := util.NewRingLoadBalancer(1)
		s := NewSchedulerServer(
			logger, schedulerStore, experimentServer, pipelineServer,
			scheduler, eventHub, sync, SchedulerServerConfig{}, "", "", modelGwLoadBalancer, pipelineGwLoadBalancer, nil, tls.TLSOptions{})
		sync.Signals(1)
		mockAgent := &mockAgentHandler{}

		return s, mockAgent, eventHub
	}

	smallMemory := uint64(100)
	largeMemory := uint64(2000)

	type test struct {
		name           string
		req            []*pba.AgentSubscribeRequest
		model          *pb.Model
		scheduleFailed bool
	}

	tests := []test{
		{
			name: "Simple",
			req: []*pba.AgentSubscribeRequest{
				{
					ServerName:           "server1",
					ReplicaIdx:           0,
					Shared:               true,
					AvailableMemoryBytes: 1000,
					ReplicaConfig: &pba.ReplicaConfig{
						InferenceSvc:      "server1",
						InferenceHttpPort: 1,
						MemoryBytes:       1000,
						Capabilities:      []string{"sklearn"},
					},
				},