feat: make initializer accept gRPC delay times (#900)

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

api/backend/embeddings.go

@@ -30,6 +30,14 @@ func ModelEmbedding(s string, tokens []int, loader *model.ModelLoader, c config.
 		model.WithContext(o.Context),
 	}
 
+	if c.GRPC.Attempts != 0 {
+		opts = append(opts, model.WithGRPCAttempts(c.GRPC.Attempts))
+	}
+
+	if c.GRPC.AttemptsSleepTime != 0 {
+		opts = append(opts, model.WithGRPCAttemptsDelay(c.GRPC.AttemptsSleepTime))
+	}
+
 	for k, v := range o.ExternalGRPCBackends {
 		opts = append(opts, model.WithExternalBackend(k, v))
 	}

api/backend/image.go

@@ -24,6 +24,14 @@ func ImageGeneration(height, width, mode, step, seed int, positive_prompt, negat
 		}),
 	}
 
+	if c.GRPC.Attempts != 0 {
+		opts = append(opts, model.WithGRPCAttempts(c.GRPC.Attempts))
+	}
+
+	if c.GRPC.AttemptsSleepTime != 0 {
+		opts = append(opts, model.WithGRPCAttemptsDelay(c.GRPC.AttemptsSleepTime))
+	}
+
 	for k, v := range o.ExternalGRPCBackends {
 		opts = append(opts, model.WithExternalBackend(k, v))
 	}

api/backend/llm.go

@@ -31,6 +31,14 @@ func ModelInference(ctx context.Context, s string, loader *model.ModelLoader, c
 		model.WithContext(o.Context),
 	}
 
+	if c.GRPC.Attempts != 0 {
+		opts = append(opts, model.WithGRPCAttempts(c.GRPC.Attempts))
+	}
+
+	if c.GRPC.AttemptsSleepTime != 0 {
+		opts = append(opts, model.WithGRPCAttemptsDelay(c.GRPC.AttemptsSleepTime))
+	}
+
 	for k, v := range o.ExternalGRPCBackends {
 		opts = append(opts, model.WithExternalBackend(k, v))
 	}

api/backend/transcript.go

@@ -21,6 +21,13 @@ func ModelTranscription(audio, language string, loader *model.ModelLoader, c con
 		model.WithAssetDir(o.AssetsDestination),
 	}
 
+	if c.GRPC.Attempts != 0 {
+		opts = append(opts, model.WithGRPCAttempts(c.GRPC.Attempts))
+	}
+
+	if c.GRPC.AttemptsSleepTime != 0 {
+		opts = append(opts, model.WithGRPCAttemptsDelay(c.GRPC.AttemptsSleepTime))
+	}
 	for k, v := range o.ExternalGRPCBackends {
 		opts = append(opts, model.WithExternalBackend(k, v))
 	}

api/config/config.go

@@ -39,6 +39,14 @@ type Config struct {
 	Diffusers Diffusers `yaml:"diffusers"`
 
 	Step int `yaml:"step"`
+
+	// GRPC Options
+	GRPC GRPC `yaml:"grpc"`
+}
+
+type GRPC struct {
+	Attempts          int `yaml:"attempts"`
+	AttemptsSleepTime int `yaml:"attempts_sleep_time"`
 }
 
 type Diffusers struct {

pkg/model/initializers.go

@@ -185,13 +185,13 @@ func (ml *ModelLoader) grpcModel(backend string, o *Options) func(string, string
 
 		// Wait for the service to start up
 		ready := false
-		for i := 0; i < 10; i++ {
+		for i := 0; i < o.grpcAttempts; i++ {
 			if client.HealthCheck(context.Background()) {
 				log.Debug().Msgf("GRPC Service Ready")
 				ready = true
 				break
 			}
-			time.Sleep(1 * time.Second)
+			time.Sleep(time.Duration(o.grpcAttemptsDelay) * time.Second)
 		}
 
 		if !ready {

pkg/model/options.go

@@ -16,6 +16,9 @@ type Options struct {
 	gRPCOptions *pb.ModelOptions
 
 	externalBackends map[string]string
+
+	grpcAttempts      int
+	grpcAttemptsDelay int
 }
 
 type Option func(*Options)
@@ -29,6 +32,18 @@ func WithExternalBackend(name string, uri string) Option {
 	}
 }
 
+func WithGRPCAttempts(attempts int) Option {
+	return func(o *Options) {
+		o.grpcAttempts = attempts
+	}
+}
+
+func WithGRPCAttemptsDelay(delay int) Option {
+	return func(o *Options) {
+		o.grpcAttemptsDelay = delay
+	}
+}
+
 func WithBackendString(backend string) Option {
 	return func(o *Options) {
 		o.backendString = backend
@@ -67,8 +82,10 @@ func WithContext(ctx context.Context) Option {
 
 func NewOptions(opts ...Option) *Options {
 	o := &Options{
-		gRPCOptions: &pb.ModelOptions{},
+		gRPCOptions:       &pb.ModelOptions{},
-		context:     context.Background(),
+		context:           context.Background(),
+		grpcAttempts:      20,
+		grpcAttemptsDelay: 2,
 	}
 	for _, opt := range opts {
 		opt(o)