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/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 54f971dd..522b0e13 100644 --- a/commands/root.go +++ b/commands/root.go @@ -108,6 +108,9 @@ func NewRootCmd(cli *command.DockerCli) *cobra.Command { newTagCmd(), newInstallRunner(), newUninstallRunner(), + newPSCmd(), + newDFCmd(), + newUnloadCmd(), ) 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..d2a0dbb0 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,109 @@ 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 +} + // 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) diff --git a/docs/reference/docker_model.yaml b/docs/reference/docker_model.yaml index 3a414837..bbc08323 100644 --- a/docs/reference/docker_model.yaml +++ b/docs/reference/docker_model.yaml @@ -4,11 +4,13 @@ long: Docker Model Runner pname: docker plink: docker.yaml cname: + - docker model df - docker model inspect - docker model install-runner - docker model list - docker model logs - docker model package + - docker model ps - docker model pull - docker model push - docker model rm @@ -16,13 +18,16 @@ cname: - docker model status - docker model tag - docker model uninstall-runner + - docker model unload - docker model version clink: + - docker_model_df.yaml - docker_model_inspect.yaml - docker_model_install-runner.yaml - docker_model_list.yaml - docker_model_logs.yaml - docker_model_package.yaml + - docker_model_ps.yaml - docker_model_pull.yaml - docker_model_push.yaml - docker_model_rm.yaml @@ -30,6 +35,7 @@ clink: - docker_model_status.yaml - docker_model_tag.yaml - docker_model_uninstall-runner.yaml + - docker_model_unload.yaml - docker_model_version.yaml deprecated: false hidden: false diff --git a/docs/reference/docker_model_df.yaml b/docs/reference/docker_model_df.yaml new file mode 100644 index 00000000..f1b3fca0 --- /dev/null +++ b/docs/reference/docker_model_df.yaml @@ -0,0 +1,13 @@ +command: docker model df +short: Show Docker Model Runner disk usage +long: Show Docker Model Runner disk usage +usage: docker model df +pname: docker model +plink: docker_model.yaml +deprecated: false +hidden: false +experimental: false +experimentalcli: true +kubernetes: false +swarm: false + diff --git a/docs/reference/docker_model_ps.yaml b/docs/reference/docker_model_ps.yaml new file mode 100644 index 00000000..54ac9856 --- /dev/null +++ b/docs/reference/docker_model_ps.yaml @@ -0,0 +1,13 @@ +command: docker model ps +short: List running models +long: List running models +usage: docker model ps +pname: docker model +plink: docker_model.yaml +deprecated: false +hidden: false +experimental: false +experimentalcli: true +kubernetes: false +swarm: false + diff --git a/docs/reference/docker_model_unload.yaml b/docs/reference/docker_model_unload.yaml new file mode 100644 index 00000000..589afe59 --- /dev/null +++ b/docs/reference/docker_model_unload.yaml @@ -0,0 +1,33 @@ +command: docker model unload +short: Unload running models +long: Unload running models +usage: docker model unload (MODEL [--backend BACKEND] | --all) +pname: docker model +plink: docker_model.yaml +options: + - option: all + value_type: bool + default_value: "false" + description: Unload all running models + deprecated: false + hidden: false + experimental: false + experimentalcli: false + kubernetes: false + swarm: false + - option: backend + value_type: string + description: Optional backend to target + deprecated: false + hidden: false + experimental: false + experimentalcli: false + kubernetes: false + swarm: false +deprecated: false +hidden: false +experimental: false +experimentalcli: true +kubernetes: false +swarm: false + diff --git a/docs/reference/model.md b/docs/reference/model.md index af7a4973..149acfa6 100644 --- a/docs/reference/model.md +++ b/docs/reference/model.md @@ -7,11 +7,13 @@ Docker Model Runner (EXPERIMENTAL) | Name | Description | |:------------------------------------------------|:-----------------------------------------------------------------------| +| [`df`](model_df.md) | Show Docker Model Runner disk usage | | [`inspect`](model_inspect.md) | Display detailed information on one model | | [`install-runner`](model_install-runner.md) | Install Docker Model Runner | | [`list`](model_list.md) | List the available models that can be run with the Docker Model Runner | | [`logs`](model_logs.md) | Fetch the Docker Model Runner logs | | [`package`](model_package.md) | package a model | +| [`ps`](model_ps.md) | List running models | | [`pull`](model_pull.md) | Download a model | | [`push`](model_push.md) | Upload a model | | [`rm`](model_rm.md) | Remove models downloaded from Docker Hub | @@ -19,6 +21,7 @@ Docker Model Runner (EXPERIMENTAL) | [`status`](model_status.md) | Check if the Docker Model Runner is running | | [`tag`](model_tag.md) | Tag a model | | [`uninstall-runner`](model_uninstall-runner.md) | Uninstall Docker Model Runner | +| [`unload`](model_unload.md) | Unload running models | | [`version`](model_version.md) | Show the Docker Model Runner version | diff --git a/docs/reference/model_df.md b/docs/reference/model_df.md new file mode 100644 index 00000000..e6a40736 --- /dev/null +++ b/docs/reference/model_df.md @@ -0,0 +1,8 @@ +# docker model df + + +Show Docker Model Runner disk usage + + + + diff --git a/docs/reference/model_ps.md b/docs/reference/model_ps.md new file mode 100644 index 00000000..15f53715 --- /dev/null +++ b/docs/reference/model_ps.md @@ -0,0 +1,8 @@ +# docker model ps + + +List running models + + + + diff --git a/docs/reference/model_unload.md b/docs/reference/model_unload.md new file mode 100644 index 00000000..70d7f8f2 --- /dev/null +++ b/docs/reference/model_unload.md @@ -0,0 +1,15 @@ +# docker model unload + + +Unload running models + +### Options + +| Name | Type | Default | Description | +|:------------|:---------|:--------|:---------------------------| +| `--all` | `bool` | | Unload all running models | +| `--backend` | `string` | | Optional backend to target | + + + +