diff --git a/commands/df.go b/commands/df.go new file mode 100644 index 00000000..e3558826 --- /dev/null +++ b/commands/df.go @@ -0,0 +1,56 @@ +package commands + +import ( + "bytes" + + "github.com/docker/go-units" + "github.com/docker/model-cli/commands/completion" + "github.com/docker/model-cli/desktop" + "github.com/olekukonko/tablewriter" + "github.com/spf13/cobra" +) + +func newDFCmd() *cobra.Command { + c := &cobra.Command{ + Use: "df", + Short: "Show Docker Model Runner disk usage", + RunE: func(cmd *cobra.Command, args []string) error { + df, err := desktopClient.DF() + if err != nil { + err = handleClientError(err, "Failed to list running models") + return handleNotRunningError(err) + } + cmd.Print(diskUsageTable(df)) + return nil + }, + ValidArgsFunction: completion.NoComplete, + } + return c +} + +func diskUsageTable(df desktop.DiskUsage) string { + var buf bytes.Buffer + table := tablewriter.NewWriter(&buf) + + table.SetHeader([]string{"TYPE", "SIZE"}) + + table.SetBorder(false) + table.SetColumnSeparator("") + table.SetHeaderLine(false) + table.SetTablePadding(" ") + table.SetNoWhiteSpace(true) + + table.SetColumnAlignment([]int{ + tablewriter.ALIGN_LEFT, // TYPE + tablewriter.ALIGN_LEFT, // SIZE + }) + table.SetHeaderAlignment(tablewriter.ALIGN_LEFT) + + table.Append([]string{"Models", units.HumanSize(df.ModelsDiskUsage)}) + if df.DefaultBackendDiskUsage != 0 { + table.Append([]string{"Inference engine", units.HumanSize(df.DefaultBackendDiskUsage)}) + } + + table.Render() + return buf.String() +} diff --git a/commands/prune.go b/commands/prune.go new file mode 100644 index 00000000..d63ffa51 --- /dev/null +++ b/commands/prune.go @@ -0,0 +1,49 @@ +package commands + +import ( + "fmt" + + "github.com/docker/model-cli/commands/completion" + "github.com/docker/model-cli/desktop" + "github.com/spf13/cobra" +) + +func newPruneCmd() *cobra.Command { + var force bool + + c := &cobra.Command{ + Use: "prune [OPTIONS]", + Short: "Remove all models", + RunE: func(cmd *cobra.Command, args []string) error { + if !force { + cmd.Println("WARNING! This will remove the entire models directory.") + cmd.Print("Are you sure you want to continue? [y/N] ") + + var input string + _, err := fmt.Scanln(&input) + if err != nil && err.Error() != "unexpected newline" { + return err + } + + if input != "y" && input != "Y" { + cmd.Println("Operation cancelled.") + return nil + } + } + _, err := desktopClient.Unload(desktop.UnloadRequest{All: true}) + if err != nil { + err = handleClientError(err, "Failed to unload models") + return handleNotRunningError(err) + } + if err := desktopClient.Prune(); err != nil { + err = handleClientError(err, "Failed to prune") + return handleNotRunningError(err) + } + return nil + }, + ValidArgsFunction: completion.NoComplete, + } + + c.Flags().BoolVarP(&force, "force", "f", false, "Forcefully remove all models") + return c +} diff --git a/commands/ps.go b/commands/ps.go new file mode 100644 index 00000000..67d5d656 --- /dev/null +++ b/commands/ps.go @@ -0,0 +1,63 @@ +package commands + +import ( + "bytes" + "time" + + "github.com/docker/go-units" + "github.com/docker/model-cli/commands/completion" + "github.com/docker/model-cli/desktop" + "github.com/olekukonko/tablewriter" + "github.com/spf13/cobra" +) + +func newPSCmd() *cobra.Command { + c := &cobra.Command{ + Use: "ps", + Short: "List running models", + RunE: func(cmd *cobra.Command, args []string) error { + ps, err := desktopClient.PS() + if err != nil { + err = handleClientError(err, "Failed to list running models") + return handleNotRunningError(err) + } + cmd.Print(psTable(ps)) + return nil + }, + ValidArgsFunction: completion.NoComplete, + } + return c +} + +func psTable(ps []desktop.BackendStatus) string { + var buf bytes.Buffer + table := tablewriter.NewWriter(&buf) + + table.SetHeader([]string{"MODEL NAME", "BACKEND", "MODE", "LAST USED"}) + + table.SetBorder(false) + table.SetColumnSeparator("") + table.SetHeaderLine(false) + table.SetTablePadding(" ") + table.SetNoWhiteSpace(true) + + table.SetColumnAlignment([]int{ + tablewriter.ALIGN_LEFT, // MODEL + tablewriter.ALIGN_LEFT, // BACKEND + tablewriter.ALIGN_LEFT, // MODE + tablewriter.ALIGN_LEFT, // LAST USED + }) + table.SetHeaderAlignment(tablewriter.ALIGN_LEFT) + + for _, status := range ps { + table.Append([]string{ + status.ModelName, + status.BackendName, + status.Mode, + units.HumanDuration(time.Since(status.LastUsed)) + " ago", + }) + } + + table.Render() + return buf.String() +} diff --git a/commands/root.go b/commands/root.go index 19fdb368..041a18d0 100644 --- a/commands/root.go +++ b/commands/root.go @@ -106,6 +106,10 @@ func NewRootCmd(cli *command.DockerCli) *cobra.Command { newTagCmd(), newInstallRunner(), newUninstallRunner(), + newPSCmd(), + newDFCmd(), + newUnloadCmd(), + newPruneCmd(), ) return rootCmd } diff --git a/commands/unload.go b/commands/unload.go new file mode 100644 index 00000000..20c6f280 --- /dev/null +++ b/commands/unload.go @@ -0,0 +1,66 @@ +package commands + +import ( + "fmt" + + "github.com/docker/model-cli/commands/completion" + "github.com/docker/model-cli/desktop" + "github.com/spf13/cobra" +) + +func newUnloadCmd() *cobra.Command { + var all bool + var backend string + + cmdArgs := "(MODEL [--backend BACKEND] | --all)" + c := &cobra.Command{ + Use: "unload " + cmdArgs, + Short: "Unload running models", + RunE: func(cmd *cobra.Command, args []string) error { + model := args[0] + unloadResp, err := desktopClient.Unload(desktop.UnloadRequest{All: all, Backend: backend, Model: model}) + if err != nil { + err = handleClientError(err, "Failed to unload models") + return handleNotRunningError(err) + } + unloaded := unloadResp.UnloadedRunners + if unloaded == 0 { + if all { + cmd.Println("No models are running.") + } else { + cmd.Println("No such model(s) running.") + } + } else { + cmd.Printf("Unloaded %d model(s).\n", unloaded) + } + return nil + }, + ValidArgsFunction: completion.NoComplete, + } + c.Args = func(cmd *cobra.Command, args []string) error { + if all { + if len(args) > 0 { + return fmt.Errorf( + "'docker model unload' does not take MODEL when --all is specified.\n\n" + + "Usage: docker model unload " + cmdArgs + "\n\n" + + "See 'docker model unload --help' for more information.", + ) + } + return nil + } + if len(args) < 1 { + return fmt.Errorf( + "'docker model unload' requires MODEL unless --all is specified.\n\n" + + "Usage: docker model unload " + cmdArgs + "\n\n" + + "See 'docker model unload --help' for more information.", + ) + } + if len(args) > 1 { + return fmt.Errorf("too many arguments, expected " + cmdArgs) + } + return nil + } + c.Flags().BoolVar(&all, "all", false, "Unload all running models") + c.Flags().StringVar(&backend, "backend", "", "Optional backend to target") + return c +} diff --git a/desktop/desktop.go b/desktop/desktop.go index 4329f3d5..bf93daae 100644 --- a/desktop/desktop.go +++ b/desktop/desktop.go @@ -10,6 +10,7 @@ import ( "net/http" "strconv" "strings" + "time" "github.com/docker/model-runner/pkg/inference" "github.com/docker/model-runner/pkg/inference/models" @@ -438,6 +439,125 @@ func (c *Client) Remove(models []string, force bool) (string, error) { return modelRemoved, nil } +// BackendStatus to be imported from docker/model-runner when https://github.com/docker/model-runner/pull/42 is merged. +type BackendStatus struct { + // BackendName is the name of the backend + BackendName string `json:"backend_name"` + // ModelName is the name of the model loaded in the backend + ModelName string `json:"model_name"` + // Mode is the mode the backend is operating in + Mode string `json:"mode"` + // LastUsed represents when this backend was last used (if it's idle) + LastUsed time.Time `json:"last_used,omitempty"` +} + +func (c *Client) PS() ([]BackendStatus, error) { + psPath := inference.InferencePrefix + "/ps" + resp, err := c.doRequest(http.MethodGet, psPath, nil) + if err != nil { + return []BackendStatus{}, c.handleQueryError(err, psPath) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + return []BackendStatus{}, fmt.Errorf("failed to list running models: %s", resp.Status) + } + + body, _ := io.ReadAll(resp.Body) + var ps []BackendStatus + if err := json.Unmarshal(body, &ps); err != nil { + return []BackendStatus{}, fmt.Errorf("failed to unmarshal response body: %w", err) + } + + return ps, nil +} + +// DiskUsage to be imported from docker/model-runner when https://github.com/docker/model-runner/pull/45 is merged. +type DiskUsage struct { + ModelsDiskUsage float64 `json:"models_disk_usage"` + DefaultBackendDiskUsage float64 `json:"default_backend_disk_usage"` +} + +func (c *Client) DF() (DiskUsage, error) { + dfPath := inference.InferencePrefix + "/df" + resp, err := c.doRequest(http.MethodGet, dfPath, nil) + if err != nil { + return DiskUsage{}, c.handleQueryError(err, dfPath) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + return DiskUsage{}, fmt.Errorf("failed to get disk usage: %s", resp.Status) + } + + body, _ := io.ReadAll(resp.Body) + var df DiskUsage + if err := json.Unmarshal(body, &df); err != nil { + return DiskUsage{}, fmt.Errorf("failed to unmarshal response body: %w", err) + } + + return df, nil +} + +// UnloadRequest to be imported from docker/model-runner when https://github.com/docker/model-runner/pull/46 is merged. +type UnloadRequest struct { + All bool `json:"all"` + Backend string `json:"backend"` + Model string `json:"model"` +} + +// UnloadResponse to be imported from docker/model-runner when https://github.com/docker/model-runner/pull/46 is merged. +type UnloadResponse struct { + UnloadedRunners int `json:"unloaded_runners"` +} + +func (c *Client) Unload(req UnloadRequest) (UnloadResponse, error) { + unloadPath := inference.InferencePrefix + "/unload" + jsonData, err := json.Marshal(req) + if err != nil { + return UnloadResponse{}, fmt.Errorf("error marshaling request: %w", err) + } + + resp, err := c.doRequest(http.MethodPost, unloadPath, bytes.NewReader(jsonData)) + if err != nil { + return UnloadResponse{}, c.handleQueryError(err, unloadPath) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + body, _ := io.ReadAll(resp.Body) + return UnloadResponse{}, fmt.Errorf("unloading failed with status %s: %s", resp.Status, string(body)) + } + + body, err := io.ReadAll(resp.Body) + if err != nil { + return UnloadResponse{}, fmt.Errorf("failed to read response body: %w", err) + } + + var unloadResp UnloadResponse + if err := json.Unmarshal(body, &unloadResp); err != nil { + return UnloadResponse{}, fmt.Errorf("failed to unmarshal response body: %w", err) + } + + return unloadResp, nil +} + +func (c *Client) Prune() error { + prunePath := inference.ModelsPrefix + "/prune" + resp, err := c.doRequest(http.MethodDelete, prunePath, nil) + if err != nil { + return c.handleQueryError(err, prunePath) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + body, _ := io.ReadAll(resp.Body) + return fmt.Errorf("pruning failed with status %s: %s", resp.Status, string(body)) + } + + return nil +} + // doRequest is a helper function that performs HTTP requests and handles 503 responses func (c *Client) doRequest(method, path string, body io.Reader) (*http.Response, error) { req, err := http.NewRequest(method, c.modelRunner.URL(path), body)