core: Make the AI config take a big rat

livepeer · Aug 1, 2024 · 8036856 · 8036856
1 parent 4c4b1b6
commit 8036856
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 43 deletions.
diff --git a/cmd/livepeer/starter/starter.go b/cmd/livepeer/starter/starter.go
@@ -1704,17 +1704,17 @@ func getGatewayPrices(gatewayPrices string) []GatewayPrice {
 	// {"gateways":[{"ethaddress":"address1","priceperunit":0.5,"currency":"USD","pixelsperunit":1}, {"ethaddress":"address2","priceperunit":0.3,"currency":"USD","pixelsperunit":3}]}
 	var pricesSet struct {
 		Gateways []struct {
-			EthAddress    string          `json:"ethaddress"`
-			PixelsPerUnit json.RawMessage `json:"pixelsperunit"`
-			PricePerUnit  json.RawMessage `json:"priceperunit"`
-			Currency      string          `json:"currency"`
+			EthAddress    string       `json:"ethaddress"`
+			PixelsPerUnit core.JSONRat `json:"pixelsperunit"`
+			PricePerUnit  core.JSONRat `json:"priceperunit"`
+			Currency      string       `json:"currency"`
 		} `json:"gateways"`
 		// TODO: Keep the old name for backwards compatibility, remove in the future
 		Broadcasters []struct {
-			EthAddress    string          `json:"ethaddress"`
-			PixelsPerUnit json.RawMessage `json:"pixelsperunit"`
-			PricePerUnit  json.RawMessage `json:"priceperunit"`
-			Currency      string          `json:"currency"`
+			EthAddress    string       `json:"ethaddress"`
+			PixelsPerUnit core.JSONRat `json:"pixelsperunit"`
+			PricePerUnit  core.JSONRat `json:"priceperunit"`
+			Currency      string       `json:"currency"`
 		} `json:"broadcasters"`
 	}
 	pricesFileContent, _ := common.ReadFromFile(gatewayPrices)
@@ -1735,21 +1735,11 @@ func getGatewayPrices(gatewayPrices string) []GatewayPrice {
 
 	prices := make([]GatewayPrice, len(allGateways))
 	for i, p := range allGateways {
-		pixelsPerUnit, ok := new(big.Rat).SetString(string(p.PixelsPerUnit))
-		if !ok {
-			glog.Errorf("Pixels per unit could not be parsed for gateway %v. must be a valid number, provided %s", p.EthAddress, p.PixelsPerUnit)
-			continue
-		}
-		pricePerUnit, ok := new(big.Rat).SetString(string(p.PricePerUnit))
-		if !ok {
-			glog.Errorf("Price per unit could not be parsed for gateway %v. must be a valid number, provided %s", p.EthAddress, p.PricePerUnit)
-			continue
-		}
 		prices[i] = GatewayPrice{
 			EthAddress:    p.EthAddress,
 			Currency:      p.Currency,
-			PricePerUnit:  pricePerUnit,
-			PixelsPerUnit: pixelsPerUnit,
+			PricePerUnit:  p.PricePerUnit.Rat,
+			PixelsPerUnit: p.PixelsPerUnit.Rat,
 		}
 	}
 

diff --git a/core/ai.go b/core/ai.go
@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"math/big"
 	"os"
 	"regexp"
 	"strconv"
@@ -24,30 +25,20 @@ type AI interface {
 	HasCapacity(pipeline, modelID string) bool
 }
 
-// Custom type to handle both string and int but store as string.
-type StringInt string
+// Custom type to parse a big.Rat from a JSON number.
+type JSONRat struct{ *big.Rat }
 
-// UnmarshalJSON method to handle both string and int.
-func (s *StringInt) UnmarshalJSON(data []byte) error {
-	// Try to unmarshal as int.
-	var intValue int64
-	if err := json.Unmarshal(data, &intValue); err == nil {
-		*s = StringInt(strconv.FormatInt(intValue, 10))
-		return nil
+func (s *JSONRat) UnmarshalJSON(data []byte) error {
+	rat, ok := new(big.Rat).SetString(string(data))
+	if !ok {
+		return fmt.Errorf("value is not a number: %s", data)
 	}
-
-	var strValue string
-	if err := json.Unmarshal(data, &strValue); err == nil {
-		*s = StringInt(strValue)
-		return nil
-	}
-
-	return fmt.Errorf("invalid value for StringInt: %s", data)
+	*s = JSONRat{rat}
+	return nil
 }
 
-// String converts the StringInt type to a string.
-func (s StringInt) String() string {
-	return string(s)
+func (s JSONRat) String() string {
+	return s.FloatString(2)
 }
 
 type AIModelConfig struct {
@@ -56,8 +47,8 @@ type AIModelConfig struct {
 	URL               string                   `json:"url,omitempty"`
 	Token             string                   `json:"token,omitempty"`
 	Warm              bool                     `json:"warm,omitempty"`
-	PricePerUnit      StringInt                `json:"price_per_unit,omitempty"`
-	PixelsPerUnit     StringInt                `json:"pixels_per_unit,omitempty"`
+	PricePerUnit      JSONRat                  `json:"price_per_unit,omitempty"`
+	PixelsPerUnit     JSONRat                  `json:"pixels_per_unit,omitempty"`
 	OptimizationFlags worker.OptimizationFlags `json:"optimization_flags,omitempty"`
 }
 
@@ -66,7 +57,7 @@ func (config *AIModelConfig) UnmarshalJSON(data []byte) error {
 	type AIModelConfigAlias AIModelConfig
 	// Set default values for fields
 	defaultConfig := &AIModelConfigAlias{
-		PixelsPerUnit: "1",
+		PixelsPerUnit: JSONRat{new(big.Rat).SetInt64(1)},
 	}
 
 	if err := json.Unmarshal(data, defaultConfig); err != nil {