-
Notifications
You must be signed in to change notification settings - Fork 103
cli: prompt and install missing backend during model run #738
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.
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
base: main
Are you sure you want to change the base?
Changes from 3 commits
ea45ea6
fcca922
9988c13
cd54d0b
3b3aa0d
d1cc337
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,7 +1,9 @@ | ||
| package commands | ||
|
|
||
| import ( | ||
| "bufio" | ||
| "bytes" | ||
| "encoding/json" | ||
| "errors" | ||
| "fmt" | ||
| "io" | ||
|
|
@@ -12,7 +14,10 @@ import ( | |
| "github.com/docker/model-runner/cmd/cli/desktop" | ||
| "github.com/docker/model-runner/cmd/cli/pkg/standalone" | ||
| "github.com/docker/model-runner/pkg/distribution/oci/reference" | ||
| "github.com/docker/model-runner/pkg/distribution/types" | ||
| "github.com/docker/model-runner/pkg/inference/backends/llamacpp" | ||
| "github.com/docker/model-runner/pkg/inference/backends/vllm" | ||
| dmrm "github.com/docker/model-runner/pkg/inference/models" | ||
| "github.com/moby/term" | ||
| "github.com/olekukonko/tablewriter" | ||
| "github.com/olekukonko/tablewriter/renderer" | ||
|
|
@@ -270,3 +275,114 @@ func newTable(w io.Writer) *tablewriter.Table { | |
| }), | ||
| ) | ||
| } | ||
|
|
||
| func CheckBackendInstalled(backend string) (bool, error) { | ||
| status := desktopClient.Status() | ||
| if status.Error != nil { | ||
| return false, fmt.Errorf("failed to get backend status: %w", status.Error) | ||
| } | ||
|
|
||
| var backendStatus map[string]string | ||
| if err := json.Unmarshal(status.Status, &backendStatus); err != nil { | ||
| return false, fmt.Errorf("failed to parse backend status: %w", err) | ||
| } | ||
|
|
||
| backendState, exists := backendStatus[backend] | ||
| if !exists { | ||
| return false, nil | ||
| } | ||
|
|
||
| state := strings.TrimSpace(strings.ToLower(backendState)) | ||
| if strings.HasPrefix(state, "not ") || strings.HasPrefix(state, "error") { | ||
| return false, nil | ||
| } | ||
|
|
||
| return strings.HasPrefix(state, "installed") || strings.HasPrefix(state, "running"), nil | ||
| } | ||
|
|
||
| func PromptInstallBackend(backend string, cmd *cobra.Command) (bool, error) { | ||
| fmt.Fprintf(cmd.OutOrStdout(), "Backend %q is not installed. Download and install it now? [Y/n]: ", backend) | ||
|
|
||
| reader := bufio.NewReader(os.Stdin) | ||
sourcery-ai[bot] marked this conversation as resolved.
Outdated
Show resolved
Hide resolved
|
||
| input, err := reader.ReadString('\n') | ||
| if err != nil { | ||
| return false, fmt.Errorf("failed to read input: %w", err) | ||
| } | ||
|
|
||
| input = strings.TrimSpace(strings.ToLower(input)) | ||
| return input == "" || input == "y" || input == "yes", nil | ||
| } | ||
|
|
||
| func InstallBackend(backend string, cmd *cobra.Command) error { | ||
| installCmd := newInstallRunner() | ||
| installCmd.SetArgs([]string{"--backend", backend}) | ||
|
|
||
| if err := installCmd.Execute(); err != nil { | ||
| return fmt.Errorf("failed to install backend %s: %w", backend, err) | ||
| } | ||
|
|
||
| return nil | ||
| } | ||
iamanishx marked this conversation as resolved.
Outdated
Show resolved
Hide resolved
|
||
|
|
||
| func EnsureBackendAvailable(backend string, cmd *cobra.Command) error { | ||
| installed, err := CheckBackendInstalled(backend) | ||
| if err != nil { | ||
| return err | ||
| } | ||
|
|
||
| if installed { | ||
| return nil | ||
| } | ||
|
|
||
| confirm, err := PromptInstallBackend(backend, cmd) | ||
| if err != nil { | ||
| return err | ||
| } | ||
|
|
||
| if !confirm { | ||
| cmd.Printf("Run 'docker model install-runner --backend %s' to install it manually.\n", backend) | ||
| return fmt.Errorf("backend installation cancelled") | ||
sourcery-ai[bot] marked this conversation as resolved.
Outdated
Show resolved
Hide resolved
|
||
| } | ||
|
|
||
| if err := InstallBackend(backend, cmd); err != nil { | ||
| return err | ||
| } | ||
|
|
||
| installed, err = CheckBackendInstalled(backend) | ||
| if err != nil { | ||
| return err | ||
| } | ||
| if !installed { | ||
| return fmt.Errorf("backend %q is still not installed; run 'docker model install-runner --backend %s'", backend, backend) | ||
| } | ||
|
|
||
| cmd.Printf("Backend %q installed successfully.\n", backend) | ||
| return nil | ||
| } | ||
|
|
||
| func GetRequiredBackend(model string) (string, error) { | ||
| modelInfo, err := desktopClient.Inspect(model, false) | ||
| if err != nil { | ||
| return "", err | ||
| } | ||
|
|
||
| return GetRequiredBackendFromModelInfo(&modelInfo) | ||
| } | ||
|
|
||
| func GetRequiredBackendFromModelInfo(modelInfo *dmrm.Model) (string, error) { | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Please check and use
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. The backend selection logic in |
||
| config, ok := modelInfo.Config.(*types.Config) | ||
| if !ok { | ||
| return llamacpp.Name, nil | ||
| } | ||
|
|
||
| switch config.Format { | ||
| case types.FormatSafetensors: | ||
| return vllm.Name, nil | ||
| case types.FormatGGUF: | ||
| return llamacpp.Name, nil | ||
| case types.FormatDiffusers: | ||
| return "diffusers", nil | ||
iamanishx marked this conversation as resolved.
Outdated
Show resolved
Hide resolved
|
||
| default: | ||
| return llamacpp.Name, nil | ||
| } | ||
| } | ||
Uh oh!
There was an error while loading. Please reload this page.