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..54b00dc9 100644 --- a/commands/root.go +++ b/commands/root.go @@ -108,6 +108,7 @@ func NewRootCmd(cli *command.DockerCli) *cobra.Command { newTagCmd(), newInstallRunner(), newUninstallRunner(), + newPSCmd(), ) return rootCmd } diff --git a/desktop/desktop.go b/desktop/desktop.go index 4329f3d5..5dd3079f 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,39 @@ 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 +} + // 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)