[server,core]: fix RegisterAIWorker interface implementaiton

Livepool-io · Jun 13, 2024 · ae3df6f · ae3df6f
1 parent 5a484eb
commit ae3df6f
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 1 deletion.
diff --git a/core/ai.go b/core/ai.go
@@ -27,6 +27,7 @@ type AI interface {
 }
 
 type RemoteAIResultChan chan *RemoteAIWorkerResult
+
 type RemoteAIWorkerManager struct {
 	// TODO Mapping by pipeline
 	remoteWorkers []*RemoteAIWorker
@@ -86,6 +87,9 @@ func (m *RemoteAIWorkerManager) TextToImage(ctx context.Context, req worker.Text
 	taskID, taskChan := m.addTaskChan()
 	defer m.removeTaskChan(taskID)
 
+	// select a remote worker
+	w := m.remoteWorkers[0]
+
 	// send request to remote worker
 	jsonData, err := json.Marshal(req)
 	if err != nil {
@@ -99,6 +103,10 @@ func (m *RemoteAIWorkerManager) TextToImage(ctx context.Context, req worker.Text
 	}
 	m.handleAIRequest(remoteReq) // task id, pipeline
 
+	if err := w.stream.Send(remoteReq); err != nil {
+		return nil, err
+	}
+
 	select {
 	case <-ctx.Done():
 		// return EOF signal

diff --git a/server/ot_rpc.go b/server/ot_rpc.go
@@ -368,7 +368,7 @@ func (h *lphttp) RegisterTranscoder(req *net.RegisterRequest, stream net.Transco
 	return nil
 }
 
-func (h *lphttp) RegisterRemoteAIWorker(req *net.RegisterRequest, stream net.Transcoder_RegisterAIWorkerServer) error {
+func (h *lphttp) RegisterAIWorker(req *net.RegisterRequest, stream net.Transcoder_RegisterAIWorkerServer) error {
 	from := common.GetConnectionAddr(stream.Context())
 	glog.Infof("Got a RegisterAIWorker request from transcoder=%s capacity=%d", from, req.Capacity)