feat: add gRPC protocol for fine-tuning with TrainStream RPC

- Add TrainRequest and TrainResponse messages to backend.proto
- Add TrainStream server-streaming RPC to Backend service
- Update Go gRPC interfaces (interface.go, server.go, client.go, backend.go, embed.go)
- Add stub implementation in base.go returning unimplemented error
- Regenerate protobuf bindings

This implements Phase 1 of the Unsloth fine-tuning backend feature.
Subsequent phases will implement the Python backend and Go service layer.

Signed-off-by: localai-bot <localai-bot@users.noreply.github.com>

Author: localai-bot

backend/backend.proto

@@ -39,6 +39,8 @@ service Backend {
   rpc AudioDecode(AudioDecodeRequest) returns (AudioDecodeResult) {}
 
   rpc ModelMetadata(ModelOptions) returns (ModelMetadataResponse) {}
+
+  rpc TrainStream(TrainRequest) returns (stream TrainResponse) {}
 }
 
 // Define the empty request
@@ -528,3 +530,49 @@ message ModelMetadataResponse {
   string rendered_template = 2;  // The rendered chat template with enable_thinking=true (empty if not applicable)
   ToolFormatMarkers tool_format = 3;  // Auto-detected tool format markers from differential template analysis
 }
+
+message TrainRequest {
+  string model = 1;                   // Base model name or HuggingFace model ID
+  string dataset = 2;                 // Path or HuggingFace dataset ID
+  string output_dir = 3;              // Output directory for the fine-tuned model
+
+  // Hyperparameters
+  int32 epochs = 4;                   // Number of training epochs
+  int32 batch_size = 5;               // Training batch size
+  float learning_rate = 6;            // Learning rate
+  int32 max_seq_length = 7;           // Maximum sequence length
+
+  // LoRA parameters
+  int32 lora_rank = 8;                // LoRA rank (r)
+  int32 lora_alpha = 9;               // LoRA alpha scaling factor
+  float lora_dropout = 10;            // LoRA dropout rate
+  repeated string target_modules = 11; // LoRA target modules (e.g., "q_proj", "v_proj")
+
+  // Quantization
+  string quantization = 12;           // Quantization method (e.g., "4bit", "8bit", "none")
+
+  // Additional options
+  map<string, string> options = 13;   // Generic key-value options for backend-specific settings
+}
+
+message TrainResponse {
+  enum TrainStatus {
+    UNKNOWN = 0;
+    STARTING = 1;
+    RUNNING = 2;
+    COMPLETED = 3;
+    FAILED = 4;
+  }
+
+  TrainStatus status = 1;             // Current training status
+  float progress = 2;                 // Overall progress percentage (0-100)
+  int32 current_epoch = 3;            // Current epoch number
+  int32 total_epochs = 4;             // Total number of epochs
+  int32 current_step = 5;             // Current training step
+  int32 total_steps = 6;              // Total number of training steps
+  float loss = 7;                     // Current training loss
+  float learning_rate_current = 8;    // Current learning rate (may change with schedulers)
+  string message = 9;                 // Human-readable status message
+  string error = 10;                  // Error message if status is FAILED
+  string output_path = 11;            // Path to the output model (set on completion)
+}

pkg/grpc/backend.go

@@ -63,4 +63,6 @@ type Backend interface {
 	AudioDecode(ctx context.Context, in *pb.AudioDecodeRequest, opts ...grpc.CallOption) (*pb.AudioDecodeResult, error)
 
 	ModelMetadata(ctx context.Context, in *pb.ModelOptions, opts ...grpc.CallOption) (*pb.ModelMetadataResponse, error)
+
+	TrainStream(ctx context.Context, in *pb.TrainRequest, f func(resp *pb.TrainResponse), opts ...grpc.CallOption) error
 }

pkg/grpc/base/base.go

@@ -120,6 +120,10 @@ func (llm *Base) AudioDecode(*pb.AudioDecodeRequest) (*pb.AudioDecodeResult, err
 	return nil, fmt.Errorf("unimplemented")
 }
 
+func (llm *Base) TrainStream(*pb.TrainRequest, chan *pb.TrainResponse) error {
+	return fmt.Errorf("unimplemented")
+}
+
 func memoryUsage() *pb.MemoryUsageData {
 	mud := pb.MemoryUsageData{
 		Breakdown: make(map[string]uint64),

pkg/grpc/client.go

@@ -632,6 +632,54 @@ func (c *Client) AudioDecode(ctx context.Context, in *pb.AudioDecodeRequest, opt
 	return client.AudioDecode(ctx, in, opts...)
 }
 
+func (c *Client) TrainStream(ctx context.Context, in *pb.TrainRequest, f func(resp *pb.TrainResponse), opts ...grpc.CallOption) error {
+	if !c.parallel {
+		c.opMutex.Lock()
+		defer c.opMutex.Unlock()
+	}
+	c.setBusy(true)
+	defer c.setBusy(false)
+	c.wdMark()
+	defer c.wdUnMark()
+	conn, err := grpc.Dial(c.address, grpc.WithTransportCredentials(insecure.NewCredentials()),
+		grpc.WithDefaultCallOptions(
+			grpc.MaxCallRecvMsgSize(50*1024*1024), // 50MB
+			grpc.MaxCallSendMsgSize(50*1024*1024), // 50MB
+		))
+	if err != nil {
+		return err
+	}
+	defer conn.Close()
+	client := pb.NewBackendClient(conn)
+
+	stream, err := client.TrainStream(ctx, in, opts...)
+	if err != nil {
+		return err
+	}
+
+	for {
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		default:
+		}
+
+		resp, err := stream.Recv()
+		if err == io.EOF {
+			break
+		}
+		if err != nil {
+			if ctx.Err() != nil {
+				return ctx.Err()
+			}
+			return err
+		}
+		f(resp)
+	}
+
+	return nil
+}
+
 func (c *Client) ModelMetadata(ctx context.Context, in *pb.ModelOptions, opts ...grpc.CallOption) (*pb.ModelMetadataResponse, error) {
 	if !c.parallel {
 		c.opMutex.Lock()

pkg/grpc/embed.go

@@ -10,6 +10,7 @@ import (
 
 var _ Backend = new(embedBackend)
 var _ pb.Backend_PredictStreamServer = new(embedBackendServerStream)
+var _ pb.Backend_TrainStreamServer = new(embedBackendTrainStream)
 
 type embedBackend struct {
 	s *server
@@ -119,6 +120,14 @@ func (e *embedBackend) ModelMetadata(ctx context.Context, in *pb.ModelOptions, o
 	return e.s.ModelMetadata(ctx, in)
 }
 
+func (e *embedBackend) TrainStream(ctx context.Context, in *pb.TrainRequest, f func(resp *pb.TrainResponse), opts ...grpc.CallOption) error {
+	bs := &embedBackendTrainStream{
+		ctx: ctx,
+		fn:  f,
+	}
+	return e.s.TrainStream(in, bs)
+}
+
 func (e *embedBackend) GetTokenMetrics(ctx context.Context, in *pb.MetricsRequest, opts ...grpc.CallOption) (*pb.MetricsResponse, error) {
 	return e.s.GetMetrics(ctx, in)
 }
@@ -158,3 +167,39 @@ func (e *embedBackendServerStream) SendMsg(m any) error {
 func (e *embedBackendServerStream) RecvMsg(m any) error {
 	return nil
 }
+
+type embedBackendTrainStream struct {
+	ctx context.Context
+	fn  func(resp *pb.TrainResponse)
+}
+
+func (e *embedBackendTrainStream) Send(resp *pb.TrainResponse) error {
+	e.fn(resp)
+	return nil
+}
+
+func (e *embedBackendTrainStream) SetHeader(md metadata.MD) error {
+	return nil
+}
+
+func (e *embedBackendTrainStream) SendHeader(md metadata.MD) error {
+	return nil
+}
+
+func (e *embedBackendTrainStream) SetTrailer(md metadata.MD) {
+}
+
+func (e *embedBackendTrainStream) Context() context.Context {
+	return e.ctx
+}
+
+func (e *embedBackendTrainStream) SendMsg(m any) error {
+	if x, ok := m.(*pb.TrainResponse); ok {
+		return e.Send(x)
+	}
+	return nil
+}
+
+func (e *embedBackendTrainStream) RecvMsg(m any) error {
+	return nil
+}

pkg/grpc/interface.go

@@ -35,6 +35,8 @@ type AIModel interface {
 	AudioDecode(*pb.AudioDecodeRequest) (*pb.AudioDecodeResult, error)
 
 	ModelMetadata(*pb.ModelOptions) (*pb.ModelMetadataResponse, error)
+
+	TrainStream(*pb.TrainRequest, chan *pb.TrainResponse) error
 }
 
 func newReply(s string) *pb.Reply {

pkg/grpc/server.go

@@ -320,6 +320,27 @@ func (s *server) ModelMetadata(ctx context.Context, in *pb.ModelOptions) (*pb.Mo
 	return res, nil
 }
 
+func (s *server) TrainStream(in *pb.TrainRequest, stream pb.Backend_TrainStreamServer) error {
+	if s.llm.Locking() {
+		s.llm.Lock()
+		defer s.llm.Unlock()
+	}
+	resultChan := make(chan *pb.TrainResponse)
+
+	done := make(chan bool)
+	go func() {
+		for result := range resultChan {
+			stream.Send(result)
+		}
+		done <- true
+	}()
+
+	err := s.llm.TrainStream(in, resultChan)
+	<-done
+
+	return err
+}
+
 func StartServer(address string, model AIModel) error {
 	lis, err := net.Listen("tcp", address)
 	if err != nil {