From 3c4ce10c221f70ada2682c281880ddf9a88d1b22 Mon Sep 17 00:00:00 2001 From: Dorin Geman Date: Wed, 21 May 2025 12:51:23 +0300 Subject: [PATCH] Add `docker model df` Signed-off-by: Dorin Geman --- commands/df.go | 56 ++++++++++++++++++++++++++++++++++++++++++++++ commands/root.go | 1 + desktop/desktop.go | 27 ++++++++++++++++++++++ 3 files changed, 84 insertions(+) create mode 100644 commands/df.go 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/root.go b/commands/root.go index 54b00dc9..474e1979 100644 --- a/commands/root.go +++ b/commands/root.go @@ -109,6 +109,7 @@ func NewRootCmd(cli *command.DockerCli) *cobra.Command { newInstallRunner(), newUninstallRunner(), newPSCmd(), + newDFCmd(), ) return rootCmd } diff --git a/desktop/desktop.go b/desktop/desktop.go index 5dd3079f..23f3fb74 100644 --- a/desktop/desktop.go +++ b/desktop/desktop.go @@ -472,6 +472,33 @@ func (c *Client) PS() ([]BackendStatus, error) { 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 +} + // 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)