Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

feat(ai): add pipelines optimization flags #3013

Merged
merged 3 commits into from
Apr 16, 2024
Merged

feat(ai): add pipelines optimization flags #3013

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
52dadff
feat(ai): add pipelines optimization flags
rickstaa Apr 15, 2024
bb50339
chore: update ai-worker to latest commit
rickstaa Apr 16, 2024
fd3397a
refactor: improve OptFlags logging
rickstaa Apr 16, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 10 additions & 1 deletion cmd/livepeer/starter/starter.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ import (
"path/filepath"
"strconv"
"strings"
"sync"
"time"

ethcommon "github.com/ethereum/go-ethereum/common"
Expand Down Expand Up @@ -538,19 +539,27 @@ func StartLivepeer(ctx context.Context, cfg LivepeerConfig) {
return
}

var once sync.Once
for _, config := range configs {
modelConstraint := &core.ModelConstraint{Warm: config.Warm}

// If the config contains a URL we call Warm() anyway because AIWorker will just register
// the endpoint for an external container
if config.Warm || config.URL != "" {
endpoint := worker.RunnerEndpoint{URL: config.URL, Token: config.Token}
if err := n.AIWorker.Warm(ctx, config.Pipeline, config.ModelID, endpoint); err != nil {
if err := n.AIWorker.Warm(ctx, config.Pipeline, config.ModelID, endpoint, config.OptimizationFlags); err != nil {
glog.Errorf("Error AI worker warming %v container: %v", config.Pipeline, err)
return
}
}

// Show warning if people set OptimizationFlags but not Warm.
if len(config.OptimizationFlags) > 0 && !config.Warm {
once.Do(func() {
glog.Warningf("OptimizationFlags set for model %v but Warm is not set. OptimizationFlags are currently only used for warm containers.", config.ModelID)
rickstaa marked this conversation as resolved.
Show resolved Hide resolved
rickstaa marked this conversation as resolved.
Show resolved Hide resolved
})
}

switch config.Pipeline {
case "text-to-image":
_, ok := constraints[core.Capability_TextToImage]
Expand Down
17 changes: 9 additions & 8 deletions core/ai.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,18 +15,19 @@
TextToImage(context.Context, worker.TextToImageJSONRequestBody) (*worker.ImageResponse, error)
ImageToImage(context.Context, worker.ImageToImageMultipartRequestBody) (*worker.ImageResponse, error)
ImageToVideo(context.Context, worker.ImageToVideoMultipartRequestBody) (*worker.VideoResponse, error)
Warm(context.Context, string, string, worker.RunnerEndpoint) error
Warm(context.Context, string, string, worker.RunnerEndpoint, worker.OptimizationFlags) error

Check failure on line 18 in core/ai.go

View workflow job for this annotation

GitHub Actions / Build binaries for linux-cpu-amd64 

undefined: worker.OptimizationFlags

Check failure on line 18 in core/ai.go

View workflow job for this annotation

GitHub Actions / Build binaries for linux-gpu-amd64 

undefined: worker.OptimizationFlags

Check failure on line 18 in core/ai.go

View workflow job for this annotation

GitHub Actions / Build binaries for darwin-amd64 

undefined: worker.OptimizationFlags

Check failure on line 18 in core/ai.go

View workflow job for this annotation

GitHub Actions / Build binaries for darwin-arm64 

undefined: worker.OptimizationFlags
Stop(context.Context) error
}

type AIModelConfig struct {
Pipeline string `json:"pipeline"`
ModelID string `json:"model_id"`
URL string `json:"url,omitempty"`
Token string `json:"token,omitempty"`
Warm bool `json:"warm,omitempty"`
PricePerUnit int64 `json:"price_per_unit,omitempty"`
PixelsPerUnit int64 `json:"pixels_per_unit,omitempty"`
Pipeline string `json:"pipeline"`
ModelID string `json:"model_id"`
URL string `json:"url,omitempty"`
Token string `json:"token,omitempty"`
Warm bool `json:"warm,omitempty"`
PricePerUnit int64 `json:"price_per_unit,omitempty"`
PixelsPerUnit int64 `json:"pixels_per_unit,omitempty"`
OptimizationFlags worker.OptimizationFlags `json:"optimization_flags,omitempty"`

Check failure on line 30 in core/ai.go

View workflow job for this annotation

GitHub Actions / Build binaries for linux-cpu-amd64 

undefined: worker.OptimizationFlags

Check failure on line 30 in core/ai.go

View workflow job for this annotation

GitHub Actions / Build binaries for linux-gpu-amd64 

undefined: worker.OptimizationFlags

Check failure on line 30 in core/ai.go

View workflow job for this annotation

GitHub Actions / Build binaries for darwin-amd64 

undefined: worker.OptimizationFlags

Check failure on line 30 in core/ai.go

View workflow job for this annotation

GitHub Actions / Build binaries for darwin-arm64 

undefined: worker.OptimizationFlags
}

func (config *AIModelConfig) UnmarshalJSON(data []byte) error {
Expand Down
Loading