Add docker model prune

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()
+}

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
+}

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()
+}

commands/root.go

@@ -106,6 +106,10 @@ func NewRootCmd(cli *command.DockerCli) *cobra.Command {
 		newTagCmd(),
 		newInstallRunner(),
 		newUninstallRunner(),
+		newPSCmd(),
+		newDFCmd(),
+		newUnloadCmd(),
+		newPruneCmd(),
 	)
 	return rootCmd
 }

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
+}

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)