feat: add requests monitoring

Signed-off-by: Dorin Geman <[email protected]>

Author: doringeman

commands/requests.go

@@ -0,0 +1,82 @@
+package commands
+
+import (
+	"bufio"
+	"fmt"
+	"strings"
+
+	"github.com/docker/model-cli/commands/completion"
+	"github.com/spf13/cobra"
+)
+
+func newRequestsCmd() *cobra.Command {
+	var model string
+	var follow bool
+	var includeExisting bool
+	c := &cobra.Command{
+		Use:   "requests [OPTIONS]",
+		Short: "Fetch requests+responses from Docker Model Runner",
+		PreRunE: func(cmd *cobra.Command, args []string) error {
+			// Make --include-existing only available when --follow is set.
+			if includeExisting && !follow {
+				return fmt.Errorf("--include-existing can only be used with --follow")
+			}
+			return nil
+		},
+		RunE: func(cmd *cobra.Command, args []string) error {
+			if _, err := ensureStandaloneRunnerAvailable(cmd.Context(), cmd); err != nil {
+				return fmt.Errorf("unable to initialize standalone model runner: %w", err)
+			}
+			if follow {
+				responseBody, err := desktopClient.RequestsStream(model, includeExisting)
+				if err != nil {
+					errMsg := "Failed to get requests stream"
+					if model != "" {
+						errMsg = errMsg + " for" + model
+					}
+					err = handleClientError(err, errMsg)
+					return handleNotRunningError(err)
+				}
+				scanner := bufio.NewScanner(responseBody)
+				cmd.Println("Connected to request stream. Press Ctrl+C to stop.")
+				var currentEvent string
+				for scanner.Scan() {
+					select {
+					case <-cmd.Context().Done():
+						return nil
+					default:
+					}
+					line := scanner.Text()
+					if strings.HasPrefix(line, "event: ") {
+						currentEvent = strings.TrimPrefix(line, "event: ")
+					} else if strings.HasPrefix(line, "data: ") &&
+						(currentEvent == "new_request" || currentEvent == "existing_request") {
+						data := strings.TrimPrefix(line, "data: ")
+						cmd.Println(data)
+					}
+				}
+				cmd.Println("Stream closed by server.")
+				return nil
+			}
+			requests, err := desktopClient.RequestsList(model)
+			if err != nil {
+				errMsg := "Failed to get requests"
+				if model != "" {
+					errMsg = errMsg + " for" + model
+				}
+				err = handleClientError(err, errMsg)
+				return handleNotRunningError(err)
+			}
+			cmd.Print(requests)
+			return nil
+		},
+		ValidArgsFunction: completion.NoComplete,
+	}
+	c.Flags().BoolVarP(&follow, "follow", "f", false, "Follow requests stream")
+	c.Flags().BoolVar(&includeExisting, "include-existing", false,
+		"Include existing requests when starting to follow (only available with --follow)")
+	c.Flags().StringVar(&model, "model", "", "Specify the model to filter requests")
+	// Enable completion for the --model flag.
+	_ = c.RegisterFlagCompletionFunc("model", completion.ModelNames(getDesktopClient, 1))
+	return c
+}

commands/root.go

@@ -110,6 +110,7 @@ func NewRootCmd(cli *command.DockerCli) *cobra.Command {
 		newPSCmd(),
 		newDFCmd(),
 		newUnloadCmd(),
+		newRequestsCmd(),
 	)
 	return rootCmd
 }

desktop/desktop.go

@@ -657,6 +657,63 @@ func (c *Client) ConfigureBackend(request scheduling.ConfigureRequest) error {
 	return nil
 }
 
+func (c *Client) RequestsStream(modelFilter string, includeExisting bool) (io.ReadCloser, error) {
+	path := c.modelRunner.URL(inference.InferencePrefix + "/requests/stream")
+	var queryParams []string
+	if modelFilter != "" {
+		queryParams = append(queryParams, "model="+modelFilter)
+	}
+	if includeExisting {
+		queryParams = append(queryParams, "include_existing=true")
+	}
+	if len(queryParams) > 0 {
+		path += "?" + strings.Join(queryParams, "&")
+	}
+	req, err := http.NewRequest(http.MethodGet, path, nil)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create request: %w", err)
+	}
+	req.Header.Set("Accept", "text/event-stream")
+	req.Header.Set("Cache-Control", "no-cache")
+	req.Header.Set("User-Agent", "docker-model-cli/"+Version)
+	resp, err := c.modelRunner.Client().Do(req)
+	if err != nil {
+		return nil, fmt.Errorf("failed to connect to stream: %w", err)
+	}
+	if resp.StatusCode != http.StatusOK {
+		resp.Body.Close()
+		return nil, fmt.Errorf("stream request failed with status: %d", resp.StatusCode)
+	}
+	return resp.Body, nil
+}
+
+func (c *Client) RequestsList(modelFilter string) (string, error) {
+	path := inference.InferencePrefix + "/requests"
+	if modelFilter != "" {
+		path += "?model=" + modelFilter
+	}
+	resp, err := c.doRequest(http.MethodGet, path, nil)
+	if err != nil {
+		return "", c.handleQueryError(err, path)
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusNotFound {
+		return "", fmt.Errorf("failed to list requests: %s", resp.Status)
+	}
+
+	body, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return "", fmt.Errorf("failed to read response body: %w", err)
+	}
+
+	if resp.StatusCode == http.StatusNotFound {
+		return "", fmt.Errorf("%s", strings.TrimSpace(string(body)))
+	}
+
+	return string(body), 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) {
 	return c.doRequestWithAuth(method, path, body, "", "")