Skip to content

Commit 350556d

Browse files
rickstaarickstaa
authoredAug 9, 2024··
refactor: improve rpc message (#3115)
* refactor: unify capabilities constraints message structure This commit unifies the capabilities constraints message structure by combining the minVersion constraint and CapabilityConstraints into a single Constraints message field. The `NewCapabilitiesWithConstraints` method has been removed for consistency with how the minVersion constraint is handled, allowing the constraints fields to be private. * fix(ai): prevent nil error crash Ensure backward compatibility by handling cases where the Constraints field is missing in the capabilities request message. This prevents crashes when a Gateway with older software calls the updated orchestrator.
·
v0.8.3v0.7.6-ai.1
1 parent c336dd1 commit 350556d

File tree

6 files changed

+216
-206
lines changed

6 files changed

+216
-206
lines changed
 

‎cmd/livepeer/starter/starter.go

+8-7
Original file line numberDiff line numberDiff line change
@@ -1061,7 +1061,7 @@ func StartLivepeer(ctx context.Context, cfg LivepeerConfig) {
10611061
}
10621062

10631063
var aiCaps []core.Capability
1064-
capabilityConstraints := make(map[core.Capability]*core.PerCapabilityConstraints)
1064+
capabilityConstraints := make(core.PerCapabilityConstraints)
10651065

10661066
if *cfg.AIWorker {
10671067
gpus := []string{}
@@ -1159,7 +1159,7 @@ func StartLivepeer(ctx context.Context, cfg LivepeerConfig) {
11591159
_, ok := capabilityConstraints[core.Capability_TextToImage]
11601160
if !ok {
11611161
aiCaps = append(aiCaps, core.Capability_TextToImage)
1162-
capabilityConstraints[core.Capability_TextToImage] = &core.PerCapabilityConstraints{
1162+
capabilityConstraints[core.Capability_TextToImage] = &core.CapabilityConstraints{
11631163
Models: make(map[string]*core.ModelConstraint),
11641164
}
11651165
}
@@ -1173,7 +1173,7 @@ func StartLivepeer(ctx context.Context, cfg LivepeerConfig) {
11731173
_, ok := capabilityConstraints[core.Capability_ImageToImage]
11741174
if !ok {
11751175
aiCaps = append(aiCaps, core.Capability_ImageToImage)
1176-
capabilityConstraints[core.Capability_ImageToImage] = &core.PerCapabilityConstraints{
1176+
capabilityConstraints[core.Capability_ImageToImage] = &core.CapabilityConstraints{
11771177
Models: make(map[string]*core.ModelConstraint),
11781178
}
11791179
}
@@ -1187,7 +1187,7 @@ func StartLivepeer(ctx context.Context, cfg LivepeerConfig) {
11871187
_, ok := capabilityConstraints[core.Capability_ImageToVideo]
11881188
if !ok {
11891189
aiCaps = append(aiCaps, core.Capability_ImageToVideo)
1190-
capabilityConstraints[core.Capability_ImageToVideo] = &core.PerCapabilityConstraints{
1190+
capabilityConstraints[core.Capability_ImageToVideo] = &core.CapabilityConstraints{
11911191
Models: make(map[string]*core.ModelConstraint),
11921192
}
11931193
}
@@ -1201,7 +1201,7 @@ func StartLivepeer(ctx context.Context, cfg LivepeerConfig) {
12011201
_, ok := capabilityConstraints[core.Capability_Upscale]
12021202
if !ok {
12031203
aiCaps = append(aiCaps, core.Capability_Upscale)
1204-
capabilityConstraints[core.Capability_Upscale] = &core.PerCapabilityConstraints{
1204+
capabilityConstraints[core.Capability_Upscale] = &core.CapabilityConstraints{
12051205
Models: make(map[string]*core.ModelConstraint),
12061206
}
12071207
}
@@ -1215,7 +1215,7 @@ func StartLivepeer(ctx context.Context, cfg LivepeerConfig) {
12151215
_, ok := capabilityConstraints[core.Capability_AudioToText]
12161216
if !ok {
12171217
aiCaps = append(aiCaps, core.Capability_AudioToText)
1218-
capabilityConstraints[core.Capability_AudioToText] = &core.PerCapabilityConstraints{
1218+
capabilityConstraints[core.Capability_AudioToText] = &core.CapabilityConstraints{
12191219
Models: make(map[string]*core.ModelConstraint),
12201220
}
12211221
}
@@ -1405,7 +1405,8 @@ func StartLivepeer(ctx context.Context, cfg LivepeerConfig) {
14051405
*cfg.CliAddr = defaultAddr(*cfg.CliAddr, "127.0.0.1", TranscoderCliPort)
14061406
}
14071407

1408-
n.Capabilities = core.NewCapabilitiesWithConstraints(append(transcoderCaps, aiCaps...), core.MandatoryOCapabilities(), core.Constraints{}, capabilityConstraints)
1408+
n.Capabilities = core.NewCapabilities(append(transcoderCaps, aiCaps...), nil)
1409+
n.Capabilities.SetPerCapabilityConstraints(capabilityConstraints)
14091410
if cfg.OrchMinLivepeerVersion != nil {
14101411
n.Capabilities.SetMinVersionConstraint(*cfg.OrchMinLivepeerVersion)
14111412
}

‎core/capabilities.go

+38-33
Original file line numberDiff line numberDiff line change
@@ -21,21 +21,21 @@ type ModelConstraint struct {
2121
type Capability int
2222
type CapabilityString []uint64
2323
type Constraints struct {
24-
minVersion string
24+
minVersion string
25+
perCapability PerCapabilityConstraints
2526
}
26-
type PerCapabilityConstraints struct {
27+
type CapabilityConstraints struct {
2728
// Models contains a *ModelConstraint for each supported model ID
2829
Models ModelConstraints
2930
}
30-
type CapabilityConstraints map[Capability]*PerCapabilityConstraints
31+
type PerCapabilityConstraints map[Capability]*CapabilityConstraints
3132
type Capabilities struct {
32-
bitstring CapabilityString
33-
mandatories CapabilityString
34-
version string
35-
constraints Constraints
36-
capabilityConstraints CapabilityConstraints
37-
capacities map[Capability]int
38-
mutex sync.Mutex
33+
bitstring CapabilityString
34+
mandatories CapabilityString
35+
version string
36+
constraints Constraints
37+
capacities map[Capability]int
38+
mutex sync.Mutex
3939
}
4040
type CapabilityTest struct {
4141
inVideoData []byte
@@ -239,7 +239,7 @@ func (c1 CapabilityString) CompatibleWith(c2 CapabilityString) bool {
239239
return true
240240
}
241241

242-
func (c1 CapabilityConstraints) CompatibleWith(c2 CapabilityConstraints) bool {
242+
func (c1 PerCapabilityConstraints) CompatibleWith(c2 PerCapabilityConstraints) bool {
243243
for c1Cap, c1Constraints := range c1 {
244244
c2Constraints, ok := c2[c1Cap]
245245
if !ok {
@@ -255,7 +255,7 @@ func (c1 CapabilityConstraints) CompatibleWith(c2 CapabilityConstraints) bool {
255255
return true
256256
}
257257

258-
func (c1 *PerCapabilityConstraints) CompatibleWith(c2 *PerCapabilityConstraints) bool {
258+
func (c1 *CapabilityConstraints) CompatibleWith(c2 *CapabilityConstraints) bool {
259259
return c1.Models.CompatibleWith(c2.Models)
260260
}
261261

@@ -453,8 +453,8 @@ func (bcast *Capabilities) CompatibleWith(orch *net.Capabilities) bool {
453453
return false
454454
}
455455

456-
orchCapabilityConstraints := CapabilitiesFromNetCapabilities(orch).capabilityConstraints
457-
if !bcast.capabilityConstraints.CompatibleWith(orchCapabilityConstraints) {
456+
orchCapabilityConstraints := CapabilitiesFromNetCapabilities(orch).constraints.perCapability
457+
if !bcast.constraints.perCapability.CompatibleWith(orchCapabilityConstraints) {
458458
return false
459459
}
460460

@@ -467,19 +467,19 @@ func (c *Capabilities) ToNetCapabilities() *net.Capabilities {
467467
}
468468
c.mutex.Lock()
469469
defer c.mutex.Unlock()
470-
netCaps := &net.Capabilities{Bitstring: c.bitstring, Mandatories: c.mandatories, Version: c.version, Capacities: make(map[uint32]uint32), Constraints: &net.Capabilities_Constraints{MinVersion: c.constraints.minVersion}, CapabilityConstraints: make(map[uint32]*net.Capabilities_CapabilityConstraints)}
470+
netCaps := &net.Capabilities{Bitstring: c.bitstring, Mandatories: c.mandatories, Version: c.version, Capacities: make(map[uint32]uint32), Constraints: &net.Capabilities_Constraints{MinVersion: c.constraints.minVersion, PerCapability: make(map[uint32]*net.Capabilities_CapabilityConstraints)}}
471471
for capability, capacity := range c.capacities {
472472
netCaps.Capacities[uint32(capability)] = uint32(capacity)
473473
}
474-
for capability, constraints := range c.capabilityConstraints {
474+
for capability, constraints := range c.constraints.perCapability {
475475
models := make(map[string]*net.Capabilities_CapabilityConstraints_ModelConstraint)
476476
for modelID, modelConstraint := range constraints.Models {
477477
models[modelID] = &net.Capabilities_CapabilityConstraints_ModelConstraint{
478478
Warm: modelConstraint.Warm,
479479
}
480480
}
481481

482-
netCaps.CapabilityConstraints[uint32(capability)] = &net.Capabilities_CapabilityConstraints{
482+
netCaps.Constraints.PerCapability[uint32(capability)] = &net.Capabilities_CapabilityConstraints{
483483
Models: models,
484484
}
485485
}
@@ -491,12 +491,11 @@ func CapabilitiesFromNetCapabilities(caps *net.Capabilities) *Capabilities {
491491
return nil
492492
}
493493
coreCaps := &Capabilities{
494-
bitstring: caps.Bitstring,
495-
mandatories: caps.Mandatories,
496-
capacities: make(map[Capability]int),
497-
version: caps.Version,
498-
constraints: Constraints{minVersion: caps.Constraints.GetMinVersion()},
499-
capabilityConstraints: make(CapabilityConstraints),
494+
bitstring: caps.Bitstring,
495+
mandatories: caps.Mandatories,
496+
capacities: make(map[Capability]int),
497+
version: caps.Version,
498+
constraints: Constraints{minVersion: caps.Constraints.GetMinVersion(), perCapability: make(PerCapabilityConstraints)},
500499
}
501500
if caps.Capacities == nil || len(caps.Capacities) == 0 {
502501
// build capacities map if not present (struct received from previous versions)
@@ -514,13 +513,13 @@ func CapabilitiesFromNetCapabilities(caps *net.Capabilities) *Capabilities {
514513
}
515514
}
516515

517-
for capabilityInt, constraints := range caps.CapabilityConstraints {
516+
for capabilityInt, constraints := range caps.Constraints.PerCapability {
518517
models := make(map[string]*ModelConstraint)
519518
for modelID, modelConstraint := range constraints.Models {
520519
models[modelID] = &ModelConstraint{Warm: modelConstraint.Warm}
521520
}
522521

523-
coreCaps.capabilityConstraints[Capability(capabilityInt)] = &PerCapabilityConstraints{
522+
coreCaps.constraints.perCapability[Capability(capabilityInt)] = &CapabilityConstraints{
524523
Models: models,
525524
}
526525
}
@@ -529,7 +528,7 @@ func CapabilitiesFromNetCapabilities(caps *net.Capabilities) *Capabilities {
529528
}
530529

531530
func NewCapabilities(caps []Capability, m []Capability) *Capabilities {
532-
c := &Capabilities{capacities: make(map[Capability]int), version: LivepeerVersion, capabilityConstraints: make(CapabilityConstraints)}
531+
c := &Capabilities{capacities: make(map[Capability]int), constraints: Constraints{perCapability: make(PerCapabilityConstraints)}, version: LivepeerVersion}
533532
if len(caps) > 0 {
534533
c.bitstring = NewCapabilityString(caps)
535534
// initialize capacities to 1 by default, mandatory capabilities doesn't have capacities
@@ -543,13 +542,6 @@ func NewCapabilities(caps []Capability, m []Capability) *Capabilities {
543542
return c
544543
}
545544

546-
func NewCapabilitiesWithConstraints(caps []Capability, m []Capability, constraints Constraints, capabilityConstraints CapabilityConstraints) *Capabilities {
547-
c := NewCapabilities(caps, m)
548-
c.constraints = constraints
549-
c.capabilityConstraints = capabilityConstraints
550-
return c
551-
}
552-
553545
func (cap *Capabilities) AddCapacity(newCaps *Capabilities) {
554546
cap.mutex.Lock()
555547
defer cap.mutex.Unlock()
@@ -723,6 +715,19 @@ func (bcast *Capabilities) LegacyOnly() bool {
723715
return bcast.bitstring.CompatibleWith(legacyCapabilityString)
724716
}
725717

718+
func (bcast *Capabilities) SetPerCapabilityConstraints(constraints PerCapabilityConstraints) {
719+
if bcast != nil {
720+
bcast.constraints.perCapability = constraints
721+
}
722+
}
723+
724+
func (bcast *Capabilities) PerCapability() PerCapabilityConstraints {
725+
if bcast != nil {
726+
return bcast.constraints.perCapability
727+
}
728+
return nil
729+
}
730+
726731
func (bcast *Capabilities) SetMinVersionConstraint(minVersionConstraint string) {
727732
if bcast != nil {
728733
bcast.constraints.minVersion = minVersionConstraint

‎core/orchestrator.go

+16-14
Original file line numberDiff line numberDiff line change
@@ -347,21 +347,23 @@ func (orch *orchestrator) priceInfo(sender ethcommon.Address, manifestID Manifes
347347
}
348348
} else {
349349
// The base price is the sum of the prices of individual capability + model ID pairs
350-
for cap := range caps.Capacities {
351-
// If the capability does not have constraints (and thus any model constraints) skip it
352-
// because we only price a capability together with a model ID right now
353-
constraints, ok := caps.CapabilityConstraints[cap]
354-
if !ok {
355-
continue
356-
}
357-
for modelID := range constraints.Models {
358-
price := orch.node.GetBasePriceForCap(sender.String(), Capability(cap), modelID)
359-
if price == nil {
360-
price = orch.node.GetBasePriceForCap("default", Capability(cap), modelID)
350+
if caps.Constraints != nil && caps.Constraints.PerCapability != nil {
351+
for cap := range caps.Capacities {
352+
// If the capability does not have constraints (and thus any model constraints) skip it
353+
// because we only price a capability together with a model ID right now
354+
constraints, ok := caps.Constraints.PerCapability[cap]
355+
if !ok {
356+
continue
361357
}
362-
363-
if price != nil {
364-
basePrice.Add(basePrice, price)
358+
for modelID := range constraints.Models {
359+
price := orch.node.GetBasePriceForCap(sender.String(), Capability(cap), modelID)
360+
if price == nil {
361+
price = orch.node.GetBasePriceForCap("default", Capability(cap), modelID)
362+
}
363+
364+
if price != nil {
365+
basePrice.Add(basePrice, price)
366+
}
365367
}
366368
}
367369
}

‎net/lp_rpc.pb.go

+147-147
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎net/lp_rpc.proto

+3-2
Original file line numberDiff line numberDiff line change
@@ -112,9 +112,10 @@ message Capabilities {
112112

113113
Constraints constraints = 5;
114114

115-
// Non-binary general constraints.
115+
// Non-binary constraints.
116116
message Constraints {
117117
string minVersion = 1;
118+
map<uint32, CapabilityConstraints> PerCapability = 2;
118119
}
119120

120121
// Non-binary capability constraints, such as supported ranges.
@@ -126,7 +127,7 @@ message Capabilities {
126127
map<string, ModelConstraint> models = 1;
127128
}
128129

129-
map<uint32, CapabilityConstraints> capabilityConstraints = 6;
130+
130131
}
131132

132133
// The orchestrator sends this in response to `GetOrchestrator`, containing

‎server/ai_session.go

+4-3
Original file line numberDiff line numberDiff line change
@@ -259,7 +259,7 @@ func (sel *AISessionSelector) Refresh(ctx context.Context) error {
259259
var coldSessions []*BroadcastSession
260260
for _, sess := range sessions {
261261
// If the constraints are missing for this capability skip this session
262-
constraints, ok := sess.OrchestratorInfo.Capabilities.CapabilityConstraints[uint32(sel.cap)]
262+
constraints, ok := sess.OrchestratorInfo.Capabilities.Constraints.PerCapability[uint32(sel.cap)]
263263
if !ok {
264264
continue
265265
}
@@ -288,7 +288,7 @@ func (sel *AISessionSelector) Refresh(ctx context.Context) error {
288288
func (sel *AISessionSelector) getSessions(ctx context.Context) ([]*BroadcastSession, error) {
289289
// No warm constraints applied here because we don't want to filter out orchs based on warm criteria at discovery time
290290
// Instead, we want all orchs that support the model and then will filter for orchs that have a warm model separately
291-
capabilityConstraints := map[core.Capability]*core.PerCapabilityConstraints{
291+
capabilityConstraints := core.PerCapabilityConstraints{
292292
sel.cap: {
293293
Models: map[string]*core.ModelConstraint{
294294
sel.modelID: {
@@ -297,7 +297,8 @@ func (sel *AISessionSelector) getSessions(ctx context.Context) ([]*BroadcastSess
297297
},
298298
},
299299
}
300-
caps := core.NewCapabilitiesWithConstraints(append(core.DefaultCapabilities(), sel.cap), nil, core.Constraints{}, capabilityConstraints)
300+
caps := core.NewCapabilities(append(core.DefaultCapabilities(), sel.cap), nil)
301+
caps.SetPerCapabilityConstraints(capabilityConstraints)
301302
caps.SetMinVersionConstraint(sel.node.Capabilities.MinVersionConstraint())
302303

303304
// Set numOrchs to the pool size so that discovery tries to find maximum # of compatible orchs within a timeout

0 commit comments

Comments
 (0)
Please sign in to comment.