feat(ai): add hardware info from Orchestrators and expand network information available.

common/types.go

@@ -162,3 +162,16 @@ type OrchestratorStore interface {
 type RoundsManager interface {
 	LastInitializedRound() *big.Int
 }
+
+type NetworkCapabilities struct {
+	Orchestrators []*OrchNetworkCapabilities `json:"orchestrators"`
+}
+type OrchNetworkCapabilities struct {
+	Address            string                     `json:"address"`
+	LocalAddress       string                     `json:"local_address"`
+	OrchURI            string                     `json:"orch_uri"`
+	ServiceURI         string                     `json:"service_uri"`
+	Capabilities       *net.Capabilities          `json:"capabilities"`
+	CapabilitiesPrices []*net.PriceInfo           `json:"capabilities_prices"`
+	Hardware           []*net.HardwareInformation `json:"hardware"`
+}

core/ai.go

@@ -32,6 +32,7 @@ type AI interface {
 	Stop(context.Context) error
 	HasCapacity(string, string) bool
 	EnsureImageAvailable(context.Context, string, string) error
+	HardwareInformation() []worker.HardwareInformation
 }
 
 // Custom type to parse a big.Rat from a JSON number.

core/ai_orchestrator.go

@@ -38,6 +38,7 @@
 	manager      *RemoteAIWorkerManager
 	stream       net.AIWorker_RegisterAIWorkerServer
 	capabilities *Capabilities
+	hardware     []worker.HardwareInformation
 	eof          chan struct{}
 	addr         string
 }
@@ -64,13 +65,14 @@
 	requestSessions map[string]*RemoteAIWorker
 }
 
-func NewRemoteAIWorker(m *RemoteAIWorkerManager, stream net.AIWorker_RegisterAIWorkerServer, caps *Capabilities) *RemoteAIWorker {
+func NewRemoteAIWorker(m *RemoteAIWorkerManager, stream net.AIWorker_RegisterAIWorkerServer, caps *Capabilities, hardware []worker.HardwareInformation) *RemoteAIWorker {
 	return &RemoteAIWorker{
 		manager:      m,
 		stream:       stream,
 		eof:          make(chan struct{}, 1),
 		addr:         common.GetConnectionAddr(stream.Context()),
 		capabilities: caps,
+		hardware:     hardware,
 	}
 }
 
@@ -87,13 +89,14 @@
 	}
 }
 
-func (orch *orchestrator) ServeAIWorker(stream net.AIWorker_RegisterAIWorkerServer, capabilities *net.Capabilities) {
-	orch.node.serveAIWorker(stream, capabilities)
+func (orch *orchestrator) ServeAIWorker(stream net.AIWorker_RegisterAIWorkerServer, capabilities *net.Capabilities, hardware []*net.HardwareInformation) {
+	orch.node.serveAIWorker(stream, capabilities, hardware)
 }
 
-func (n *LivepeerNode) serveAIWorker(stream net.AIWorker_RegisterAIWorkerServer, capabilities *net.Capabilities) {
+func (n *LivepeerNode) serveAIWorker(stream net.AIWorker_RegisterAIWorkerServer, capabilities *net.Capabilities, hardware []*net.HardwareInformation) {
 	from := common.GetConnectionAddr(stream.Context())
 	wkrCaps := CapabilitiesFromNetCapabilities(capabilities)
+	wkrHdw := hardwareInformationFromNetHardware(hardware)
 	if n.Capabilities.LivepeerVersionCompatibleWith(capabilities) {
 		glog.Infof("Worker compatible, connecting worker_version=%s orchestrator_version=%s worker_addr=%s", capabilities.Version, n.Capabilities.constraints.minVersion, from)
 		n.Capabilities.AddCapacity(wkrCaps)
@@ -102,17 +105,18 @@
 		defer n.RemoveAICapabilities(wkrCaps)
 
 		// Manage blocks while AI worker is connected
-		n.AIWorkerManager.Manage(stream, capabilities)
+		n.AIWorkerManager.Manage(stream, capabilities, wkrHdw)
 		glog.V(common.DEBUG).Infof("Closing aiworker=%s channel", from)
 	} else {
 		glog.Errorf("worker %s not connected, version not compatible", from)
 	}
 }
 
 // Manage adds aiworker to list of live aiworkers. Doesn't return until aiworker disconnects
-func (rwm *RemoteAIWorkerManager) Manage(stream net.AIWorker_RegisterAIWorkerServer, capabilities *net.Capabilities) {
+func (rwm *RemoteAIWorkerManager) Manage(stream net.AIWorker_RegisterAIWorkerServer, capabilities *net.Capabilities, hardware []worker.HardwareInformation) {
 	from := common.GetConnectionAddr(stream.Context())
-	aiworker := NewRemoteAIWorker(rwm, stream, CapabilitiesFromNetCapabilities(capabilities))
+
+	aiworker := NewRemoteAIWorker(rwm, stream, CapabilitiesFromNetCapabilities(capabilities), hardware)
 	go func() {
 		ctx := stream.Context()
 		<-ctx.Done()
@@ -407,6 +411,20 @@
 	}
 }
 
+func (orch *orchestrator) WorkerHardware() []worker.HardwareInformation {
+	if orch.node.AIWorker != nil {
+		return orch.node.AIWorker.HardwareInformation()
+	} else {
+		// return combined hardware information from all live remote workers from information provided by workers
+		// when connecting to orchestrator. Does not reach out for real-time information.
+		var wkrHdw []worker.HardwareInformation
+		for _, worker := range orch.node.AIWorkerManager.liveAIWorkers {
+			wkrHdw = append(wkrHdw, worker.hardware...)
+		}
+		return wkrHdw
+	}
+}
+
 func (orch *orchestrator) AIResults(tcID int64, res *RemoteAIWorkerResult) {
 	orch.node.AIWorkerManager.aiResults(tcID, res)
 }
@@ -1174,3 +1192,18 @@
 	}
 	return &tr
 }
+
+func hardwareInformationFromNetHardware(hdw []*net.HardwareInformation) []worker.HardwareInformation {
+	var netWorkerHardware []byte
+	netWorkerHardware, err := json.Marshal(hdw)
+	if err != nil {
+		return []worker.HardwareInformation{}
+	}
+	var workerHardware []worker.HardwareInformation
+	err = json.Unmarshal(netWorkerHardware, &workerHardware)
+	if err != nil {
+		return []worker.HardwareInformation{}
+	}
+
+	return workerHardware
+}

core/ai_test.go

@@ -2,9 +2,11 @@ package core
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
 	"os"
 	"path/filepath"
+	"reflect"
 	"strconv"
 	"sync"
 	"testing"
@@ -47,13 +49,50 @@ func TestServeAIWorker(t *testing.T) {
 	// test that an ai worker was created
 	caps := createAIWorkerCapabilities()
 	netCaps := caps.ToNetCapabilities()
-	go n.serveAIWorker(strm, netCaps)
+	go n.serveAIWorker(strm, netCaps, nil)
 	time.Sleep(1 * time.Second)
 
 	wkr, ok := n.AIWorkerManager.liveAIWorkers[strm]
 	if !ok {
 		t.Error("Unexpected transcoder type")
 	}
+	//confirm worker info
+	assert.Equal(t, wkr.capabilities, caps)
+	assert.Nil(t, wkr.hardware)
+
+	// test shutdown
+	wkr.eof <- struct{}{}
+	time.Sleep(1 * time.Second)
+
+	// stream should be removed
+	_, ok = n.AIWorkerManager.liveAIWorkers[strm]
+	if ok {
+		t.Error("Unexpected ai worker presence")
+	}
+
+	//confirm no workers connected
+	assert.Len(t, n.AIWorkerManager.liveAIWorkers, 0)
+
+	//connect worker with hardware information
+	strm1 := &StubAIWorkerServer{}
+	hdwDetail := net.GPUComputeInfo{Id: "gpu-1", Name: "gpu name", Major: 8, Minor: 9, MemoryFree: 1, MemoryTotal: 10}
+	hdwInfo := make(map[string]*net.GPUComputeInfo)
+	hdwInfo["0"] = &hdwDetail
+	hdw := net.HardwareInformation{Pipeline: "livepeer-pipeline", ModelId: "livepeer/model1", GpuInfo: hdwInfo}
+	var netHdwList []*net.HardwareInformation
+	netHdwList = append(netHdwList, &hdw)
+	go n.serveAIWorker(strm1, netCaps, netHdwList)
+	time.Sleep(1 * time.Second)
+
+	wkr, ok = n.AIWorkerManager.liveAIWorkers[strm1]
+	if !ok {
+		t.Error("Unexpected transcoder type")
+	}
+
+	//confirm worker attached and has hardware information
+	assert.Len(t, n.AIWorkerManager.liveAIWorkers, 1)
+	wkrHdw := hardwareInformationFromNetHardware(netHdwList)
+	assert.Equal(t, wkrHdw, n.AIWorkerManager.liveAIWorkers[strm1].hardware)
 
 	// test shutdown
 	wkr.eof <- struct{}{}
@@ -65,6 +104,7 @@ func TestServeAIWorker(t *testing.T) {
 		t.Error("Unexpected ai worker presence")
 	}
 }
+
 func TestServeAIWorker_IncompatibleVersion(t *testing.T) {
 	assert := assert.New(t)
 	n, _ := NewLivepeerNode(nil, "", nil)
@@ -76,7 +116,7 @@ func TestServeAIWorker_IncompatibleVersion(t *testing.T) {
 	// test that an ai worker was created
 	caps := createAIWorkerCapabilities()
 	netCaps := caps.ToNetCapabilities()
-	go n.serveAIWorker(strm, netCaps)
+	go n.serveAIWorker(strm, netCaps, nil)
 	time.Sleep(5 * time.Second)
 	assert.Zero(len(n.AIWorkerManager.liveAIWorkers))
 	assert.Zero(len(n.AIWorkerManager.remoteAIWorkers))
@@ -88,14 +128,14 @@ func TestRemoteAIWorkerManager(t *testing.T) {
 	initAIWorker := func() (*RemoteAIWorker, *StubAIWorkerServer) {
 		strm := &StubAIWorkerServer{manager: m}
 		caps := createAIWorkerCapabilities()
-		wkr := NewRemoteAIWorker(m, strm, caps)
+		wkr := NewRemoteAIWorker(m, strm, caps, nil)
 		return wkr, strm
 	}
 	//create worker and connect to manager
 	wkr, strm := initAIWorker()
 
 	go func() {
-		m.Manage(strm, wkr.capabilities.ToNetCapabilities())
+		m.Manage(strm, wkr.capabilities.ToNetCapabilities(), nil)
 	}()
 	time.Sleep(1 * time.Millisecond) // allow the workers to activate
 
@@ -156,9 +196,9 @@ func TestSelectAIWorker(t *testing.T) {
 
 	// register ai workers, which adds ai worker to liveAIWorkers and remoteAIWorkers
 	wg := newWg(1)
-	go func() { m.Manage(strm, capabilities.ToNetCapabilities()) }()
+	go func() { m.Manage(strm, capabilities.ToNetCapabilities(), nil) }()
 	time.Sleep(1 * time.Millisecond) // allow time for first stream to register
-	go func() { m.Manage(strm2, extraModelCapabilities.ToNetCapabilities()); wg.Done() }()
+	go func() { m.Manage(strm2, extraModelCapabilities.ToNetCapabilities(), nil); wg.Done() }()
 	time.Sleep(1 * time.Millisecond) // allow time for second stream to register e for third stream to register
 
 	//update worker.addr to be different
@@ -256,7 +296,7 @@ func TestSelectAIWorker(t *testing.T) {
 	assert.EqualError(err, ErrNoCompatibleWorkersAvailable.Error())
 
 	// reconnect worker and check pipeline only on second worker is available
-	go func() { m.Manage(strm2, extraModelCapabilities.ToNetCapabilities()); wg.Done() }()
+	go func() { m.Manage(strm2, extraModelCapabilities.ToNetCapabilities(), nil); wg.Done() }()
 	time.Sleep(1 * time.Millisecond)
 	w, err = m.selectWorker(testRequestId, "image-to-image", "livepeer/model2")
 	assert.NotNil(w)
@@ -280,7 +320,7 @@ func TestManageAIWorkers(t *testing.T) {
 
 	// test that transcoder is added to liveTranscoders and remoteTranscoders
 	wg1 := newWg(1)
-	go func() { m.Manage(strm, capabilities.ToNetCapabilities()); wg1.Done() }()
+	go func() { m.Manage(strm, capabilities.ToNetCapabilities(), nil); wg1.Done() }()
 	time.Sleep(1 * time.Millisecond) // allow the manager to activate
 
 	assert.NotNil(m.liveAIWorkers[strm])
@@ -291,7 +331,7 @@ func TestManageAIWorkers(t *testing.T) {
 
 	// test that additional transcoder is added to liveTranscoders and remoteTranscoders
 	wg2 := newWg(1)
-	go func() { m.Manage(strm2, capabilities.ToNetCapabilities()); wg2.Done() }()
+	go func() { m.Manage(strm2, capabilities.ToNetCapabilities(), nil); wg2.Done() }()
 	time.Sleep(1 * time.Millisecond) // allow the manager to activate
 
 	assert.NotNil(m.liveAIWorkers[strm])
@@ -321,7 +361,7 @@ func TestRemoteAIWorkerTimeout(t *testing.T) {
 		strm := &StubAIWorkerServer{manager: m}
 		//create capabilities and constraints the ai worker sends to orch
 		caps := createAIWorkerCapabilities()
-		wkr := NewRemoteAIWorker(m, strm, caps)
+		wkr := NewRemoteAIWorker(m, strm, caps, nil)
 		return wkr, strm
 	}
 	//create a new worker
@@ -483,14 +523,14 @@ func TestCheckAICapacity(t *testing.T) {
 	initAIWorker := func() (*RemoteAIWorker, *StubAIWorkerServer) {
 		strm := &StubAIWorkerServer{manager: o.node.AIWorkerManager}
 		caps := createAIWorkerCapabilities()
-		wkr := NewRemoteAIWorker(o.node.AIWorkerManager, strm, caps)
+		wkr := NewRemoteAIWorker(o.node.AIWorkerManager, strm, caps, nil)
 		return wkr, strm
 	}
 	//create worker and connect to manager
 	wkr2, strm := initAIWorker()
 
 	go func() {
-		o.node.AIWorkerManager.Manage(strm, wkr2.capabilities.ToNetCapabilities())
+		o.node.AIWorkerManager.Manage(strm, wkr2.capabilities.ToNetCapabilities(), nil)
 	}()
 	time.Sleep(1 * time.Millisecond) // allow the workers to activate
 
@@ -513,12 +553,12 @@ func TestRemoteAIWorkerProcessPipelines(t *testing.T) {
 	initAIWorker := func() (*RemoteAIWorker, *StubAIWorkerServer) {
 		strm := &StubAIWorkerServer{manager: o.node.AIWorkerManager}
 		caps := createAIWorkerCapabilities()
-		wkr := NewRemoteAIWorker(o.node.AIWorkerManager, strm, caps)
+		wkr := NewRemoteAIWorker(o.node.AIWorkerManager, strm, caps, nil)
 		return wkr, strm
 	}
 	//create worker and connect to manager
 	wkr, strm := initAIWorker()
-	go o.node.serveAIWorker(strm, wkr.capabilities.ToNetCapabilities())
+	go o.node.serveAIWorker(strm, wkr.capabilities.ToNetCapabilities(), nil)
 	time.Sleep(5 * time.Millisecond) // allow the workers to activate
 
 	//check workers connected
@@ -686,6 +726,10 @@ func (a *stubAIWorker) EnsureImageAvailable(ctx context.Context, pipeline string
 	return nil
 }
 
+func (a *stubAIWorker) HardwareInformation() []worker.HardwareInformation {
+	return nil
+}
+
 type StubAIWorkerServer struct {
 	manager      *RemoteAIWorkerManager
 	SendError    error
@@ -800,3 +844,23 @@ func TestParseAIModelConfigs(t *testing.T) {
 		})
 	}
 }
+
+func TestHardwareInformationFromNetHardware(t *testing.T) {
+	netHdwDetail := net.GPUComputeInfo{Id: "gpu-1", Name: "gpu name", Major: 8, Minor: 9, MemoryFree: 1, MemoryTotal: 10}
+	netHdwInfo := make(map[string]*net.GPUComputeInfo)
+	netHdwInfo["0"] = &netHdwDetail
+	netHdw := net.HardwareInformation{Pipeline: "livepeer-pipeline", ModelId: "livepeer/model1", GpuInfo: netHdwInfo}
+	var netHdwList []*net.HardwareInformation
+	netHdwList = append(netHdwList, &netHdw)
+	//create []worker.HardwareInformation
+	hdwList := hardwareInformationFromNetHardware(netHdwList)
+
+	netHdwJson, _ := json.Marshal(netHdwList)
+	hdwJson, _ := json.Marshal(hdwList)
+
+	var hdw1, hdw2 interface{}
+	json.Unmarshal(netHdwJson, &hdw1)
+	json.Unmarshal(hdwJson, &hdw2)
+	assert.True(t, reflect.DeepEqual(hdw1, hdw2))
+
+}