diff --git a/commands/compose.go b/commands/compose.go index 944da7d0..801a9638 100644 --- a/commands/compose.go +++ b/commands/compose.go @@ -14,7 +14,7 @@ func newComposeCmd() *cobra.Command { c := &cobra.Command{ Use: "compose EVENT", } - c.AddCommand(newUpCommand(desktopClient)) + c.AddCommand(newUpCommand()) c.AddCommand(newDownCommand()) c.Hidden = true c.PersistentFlags().String("project-name", "", "compose project name") // unused by model @@ -22,7 +22,7 @@ func newComposeCmd() *cobra.Command { return c } -func newUpCommand(desktopClient *desktop.Client) *cobra.Command { +func newUpCommand() *cobra.Command { var model string c := &cobra.Command{ Use: "up", @@ -33,6 +33,11 @@ func newUpCommand(desktopClient *desktop.Client) *cobra.Command { return err } + if err := ensureStandaloneRunnerAvailable(cmd.Context(), nil); err != nil { + sendErrorf("Failed to initialize standalone model runner: %v", err) + return fmt.Errorf("Failed to initialize standalone model runner: %w", err) + } + _, _, err := desktopClient.Pull(model, func(s string) { sendInfo(s) }) @@ -41,8 +46,15 @@ func newUpCommand(desktopClient *desktop.Client) *cobra.Command { return fmt.Errorf("Failed to pull model: %v\n", err) } - // FIXME get actual URL from Docker Desktop - setenv("URL", "http://model-runner.docker.internal/engines/v1/") + if kind := modelRunner.EngineKind(); kind == desktop.ModelRunnerEngineKindDesktop { + // TODO: Get the actual URL from Docker Desktop via some API. + setenv("URL", "http://model-runner.docker.internal/engines/v1/") + } else if kind == desktop.ModelRunnerEngineKindMobyManual { + setenv("URL", modelRunner.URL("/engines/v1/")) + } else if kind == desktop.ModelRunnerEngineKindMoby || kind == desktop.ModelRunnerEngineKindCloud { + // TODO: Find a more robust solution in Moby-like environments. + setenv("URL", "http://172.17.0.1:12434/engines/v1/") + } setenv("MODEL", model) return nil diff --git a/commands/inspect.go b/commands/inspect.go index 95b33b95..31aa3cb2 100644 --- a/commands/inspect.go +++ b/commands/inspect.go @@ -24,6 +24,9 @@ func newInspectCmd() *cobra.Command { 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) + } inspectedModel, err := inspectModel(args, openai, desktopClient) if err != nil { return err diff --git a/commands/install-runner.go b/commands/install-runner.go new file mode 100644 index 00000000..6930ced9 --- /dev/null +++ b/commands/install-runner.go @@ -0,0 +1,141 @@ +package commands + +import ( + "context" + "fmt" + "os" + "time" + + "github.com/docker/model-cli/commands/completion" + "github.com/docker/model-cli/desktop" + "github.com/docker/model-cli/pkg/standalone" + "github.com/spf13/cobra" +) + +type noopPrinter struct{} + +func (*noopPrinter) Printf(format string, args ...any) {} + +func (*noopPrinter) Println(args ...any) {} + +// ensureStandaloneRunnerAvailable is a utility function that other commands can +// use to initialize a default standalone model runner. It is a no-op in +// unsupported contexts or if automatic installs have been disabled. +func ensureStandaloneRunnerAvailable(ctx context.Context, printer standalone.StatusPrinter) error { + // If we're not in a supported model runner context, then don't do anything. + if modelRunner.EngineKind() != desktop.ModelRunnerEngineKindMoby { + return nil + } + + // If automatic installation has been disabled, then don't do anything. + if os.Getenv("MODEL_RUNNER_NO_AUTO_INSTALL") != "" { + return nil + } + + // Ensure that the output printer is non-nil. + if printer == nil { + printer = &noopPrinter{} + } + + // Create a Docker client for the active context. + dockerClient, err := desktop.DockerClientForContext(dockerCLI, dockerCLI.CurrentContext()) + if err != nil { + return fmt.Errorf("failed to create Docker client: %w", err) + } + + // Check if a model runner container exists. + container, _, err := standalone.FindControllerContainer(ctx, dockerClient) + if err != nil { + return fmt.Errorf("unable to identify existing standalone model runner: %w", err) + } else if container != "" { + return nil + } + + // Ensure that we have an up-to-date copy of the image. + if err := standalone.EnsureControllerImage(ctx, dockerClient, false, printer); err != nil { + return fmt.Errorf("unable to pull latest standalone model runner image: %w", err) + } + + // Ensure that we have a model storage volume. + modelStorageVolume, err := standalone.EnsureModelStorageVolume(ctx, dockerClient, printer) + if err != nil { + return fmt.Errorf("unable to initialize standalone model storage: %w", err) + } + + // Create the model runner container. + if err := standalone.CreateControllerContainer(ctx, dockerClient, standalone.DefaultControllerPort, false, modelStorageVolume, printer); err != nil { + return fmt.Errorf("unable to initialize standalone model runner container: %w", err) + } + + // Give the model runner one second to start. + time.Sleep(1 * time.Second) + + return nil +} + +func newInstallRunner() *cobra.Command { + var port uint16 + var gpu bool + c := &cobra.Command{ + Use: "install-runner", + Short: "Install Docker Model Runner", + RunE: func(cmd *cobra.Command, args []string) error { + // Ensure that we're running in a supported model runner context. + if kind := modelRunner.EngineKind(); kind == desktop.ModelRunnerEngineKindDesktop { + // TODO: We may eventually want to auto-forward this to + // docker desktop enable model-runner, but we should first make + // sure the CLI flags match. + cmd.Println("Standalone installation not supported with Docker Desktop") + cmd.Println("Use `docker desktop enable model-runner` instead") + return nil + } else if kind == desktop.ModelRunnerEngineKindMobyManual { + cmd.Println("Standalone installation not supported with MODEL_RUNNER_HOST set") + return nil + } else if kind == desktop.ModelRunnerEngineKindCloud { + cmd.Println("Standalone installation not required with Docker Cloud") + return nil + } + + // Create a Docker client for the active context. + dockerClient, err := desktop.DockerClientForContext(dockerCLI, dockerCLI.CurrentContext()) + if err != nil { + return fmt.Errorf("failed to create Docker client: %w", err) + } + + // Check if an active model runner container already exists. + if ctrID, ctrName, err := standalone.FindControllerContainer(cmd.Context(), dockerClient); err != nil { + return err + } else if ctrID != "" { + if ctrName != "" { + cmd.Printf("Model Runner container %s (%s) is already running\n", ctrName, ctrID[:12]) + } else { + cmd.Printf("Model Runner container %s is already running\n", ctrID[:12]) + } + return nil + } + + // Ensure that we have an up-to-date copy of the image. + if err := standalone.EnsureControllerImage(cmd.Context(), dockerClient, gpu, cmd); err != nil { + return fmt.Errorf("unable to pull latest standalone model runner image: %w", err) + } + + // Ensure that we have a model storage volume. + modelStorageVolume, err := standalone.EnsureModelStorageVolume(cmd.Context(), dockerClient, cmd) + if err != nil { + return fmt.Errorf("unable to initialize standalone model storage: %w", err) + } + + // Create the model runner container. + if err := standalone.CreateControllerContainer(cmd.Context(), dockerClient, port, gpu, modelStorageVolume, cmd); err != nil { + return fmt.Errorf("unable to initialize standalone model runner container: %w", err) + } + + return nil + }, + ValidArgsFunction: completion.NoComplete, + } + c.Flags().Uint16Var(&port, "port", standalone.DefaultControllerPort, + "Docker container port for Docker Model Runner") + c.Flags().BoolVar(&gpu, "gpu", false, "Enable GPU support") + return c +} diff --git a/commands/list.go b/commands/list.go index 83d26b96..373a4015 100644 --- a/commands/list.go +++ b/commands/list.go @@ -23,6 +23,9 @@ func newListCmd() *cobra.Command { if openai && quiet { return fmt.Errorf("--quiet flag cannot be used with --openai flag") } + if err := ensureStandaloneRunnerAvailable(cmd.Context(), nil); err != nil { + return fmt.Errorf("unable to initialize standalone model runner: %w", err) + } models, err := listModels(openai, desktopClient, quiet, jsonFormat) if err != nil { return err diff --git a/commands/pull.go b/commands/pull.go index 873b25af..66596a49 100644 --- a/commands/pull.go +++ b/commands/pull.go @@ -23,6 +23,9 @@ func newPullCmd() *cobra.Command { 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) + } return pullModel(cmd, desktopClient, args[0]) }, ValidArgsFunction: completion.NoComplete, diff --git a/commands/push.go b/commands/push.go index 1f9ee82c..0f4547a3 100644 --- a/commands/push.go +++ b/commands/push.go @@ -23,6 +23,9 @@ func newPushCmd() *cobra.Command { 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) + } return pushModel(cmd, desktopClient, args[0]) }, ValidArgsFunction: completion.NoComplete, diff --git a/commands/rm.go b/commands/rm.go index 483a4622..bb11b4e7 100644 --- a/commands/rm.go +++ b/commands/rm.go @@ -24,6 +24,9 @@ func newRemoveCmd() *cobra.Command { 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) + } response, err := desktopClient.Remove(args, force) if response != "" { cmd.Println(response) diff --git a/commands/root.go b/commands/root.go index 6fbe2f30..3c9b29b3 100644 --- a/commands/root.go +++ b/commands/root.go @@ -10,6 +10,9 @@ import ( "github.com/spf13/cobra" ) +// dockerCLI is the Docker CLI environment associated with the command. +var dockerCLI *command.DockerCli + // modelRunner is the model runner context. It is initialized by the root // command's PersistentPreRunE. var modelRunner *desktop.ModelRunnerContext @@ -41,10 +44,11 @@ func NewRootCmd(cli *command.DockerCli) *cobra.Command { } else if err := plugin.PersistentPreRunE(cmd, args); err != nil { return err } + dockerCLI = cli // Detect the model runner context and create a client for it. var err error - modelRunner, err = desktop.DetectContext(cli) + modelRunner, err = desktop.DetectContext(dockerCLI) if err != nil { return fmt.Errorf("unable to detect model runner context: %w", err) } @@ -81,6 +85,8 @@ func NewRootCmd(cli *command.DockerCli) *cobra.Command { newInspectCmd(), newComposeCmd(), newTagCmd(), + newInstallRunner(), + newUninstallRunner(), ) return rootCmd } diff --git a/commands/run.go b/commands/run.go index 94155804..1f8e1485 100644 --- a/commands/run.go +++ b/commands/run.go @@ -33,6 +33,10 @@ func newRunCmd() *cobra.Command { } } + if err := ensureStandaloneRunnerAvailable(cmd.Context(), cmd); err != nil { + return fmt.Errorf("unable to initialize standalone model runner: %w", err) + } + modelDetail, err := desktopClient.Inspect(model) if err != nil { if !errors.Is(err, desktop.ErrNotFound) { diff --git a/commands/status.go b/commands/status.go index ea9f78b6..bc0d320c 100644 --- a/commands/status.go +++ b/commands/status.go @@ -2,6 +2,7 @@ package commands import ( "encoding/json" + "fmt" "os" "github.com/docker/cli/cli-plugins/hooks" @@ -14,6 +15,9 @@ func newStatusCmd() *cobra.Command { Use: "status", Short: "Check if the Docker Model Runner is running", 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) + } status := desktopClient.Status() if status.Error != nil { return handleClientError(status.Error, "Failed to get Docker Model Runner status") diff --git a/commands/tag.go b/commands/tag.go index a1b7720b..8320f737 100644 --- a/commands/tag.go +++ b/commands/tag.go @@ -23,6 +23,9 @@ func newTagCmd() *cobra.Command { 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) + } return tagModel(cmd, desktopClient, args[0], args[1]) }, } diff --git a/commands/uninstall-runner.go b/commands/uninstall-runner.go new file mode 100644 index 00000000..93c57094 --- /dev/null +++ b/commands/uninstall-runner.go @@ -0,0 +1,66 @@ +package commands + +import ( + "fmt" + + "github.com/docker/model-cli/commands/completion" + "github.com/docker/model-cli/desktop" + "github.com/docker/model-cli/pkg/standalone" + "github.com/spf13/cobra" +) + +func newUninstallRunner() *cobra.Command { + var models, images bool + c := &cobra.Command{ + Use: "uninstall-runner", + Short: "Uninstall Docker Model Runner", + RunE: func(cmd *cobra.Command, args []string) error { + // Ensure that we're running in a supported model runner context. + if kind := modelRunner.EngineKind(); kind == desktop.ModelRunnerEngineKindDesktop { + // TODO: We may eventually want to auto-forward this to + // docker desktop disable model-runner, but we should first + // make install-runner forward in the same way. + cmd.Println("Standalone uninstallation not supported with Docker Desktop") + cmd.Println("Use `docker desktop disable model-runner` instead") + return nil + } else if kind == desktop.ModelRunnerEngineKindMobyManual { + cmd.Println("Standalone uninstallation not supported with MODEL_RUNNER_HOST set") + return nil + } else if kind == desktop.ModelRunnerEngineKindCloud { + cmd.Println("Standalone uninstallation not supported with Docker Cloud") + return nil + } + + // Create a Docker client for the active context. + dockerClient, err := desktop.DockerClientForContext(dockerCLI, dockerCLI.CurrentContext()) + if err != nil { + return fmt.Errorf("failed to create Docker client: %w", err) + } + + // Remove any model runner containers. + if err := standalone.PruneControllerContainers(cmd.Context(), dockerClient, cmd); err != nil { + return fmt.Errorf("unable to remove model runner container(s): %w", err) + } + + // Remove model runner images, if requested. + if images { + if err := standalone.PruneControllerImages(cmd.Context(), dockerClient, cmd); err != nil { + return fmt.Errorf("unable to remove model runner image(s): %w", err) + } + } + + // Remove model storage, if requested. + if models { + if err := standalone.PruneModelStorageVolumes(cmd.Context(), dockerClient, cmd); err != nil { + return fmt.Errorf("unable to remove model storage volume(s): %w", err) + } + } + + return nil + }, + ValidArgsFunction: completion.NoComplete, + } + c.Flags().BoolVar(&models, "models", false, "Remove model storage volume") + c.Flags().BoolVar(&images, "images", false, "Remove "+standalone.ControllerImage+" images") + return c +} diff --git a/desktop/context.go b/desktop/context.go index 2aa74ed5..b606fc4c 100644 --- a/desktop/context.go +++ b/desktop/context.go @@ -6,10 +6,12 @@ import ( "net/url" "os" "runtime" + "strconv" "github.com/docker/cli/cli/command" "github.com/docker/cli/cli/context/docker" clientpkg "github.com/docker/docker/client" + "github.com/docker/model-cli/pkg/standalone" "github.com/docker/model-runner/pkg/inference" ) @@ -58,8 +60,8 @@ func isCloudContext(cli *command.DockerCli) bool { return ok } -// dockerClientForContext creates a Docker client for the specified context. -func dockerClientForContext(cli *command.DockerCli, name string) (*clientpkg.Client, error) { +// DockerClientForContext creates a Docker client for the specified context. +func DockerClientForContext(cli *command.DockerCli, name string) (*clientpkg.Client, error) { c, err := cli.ContextStore().GetMetadata(name) if err != nil { return nil, fmt.Errorf("unable to load context metadata: %w", err) @@ -80,8 +82,8 @@ const ( // the Model CLI command is responsible for managing a Model Runner. ModelRunnerEngineKindMoby ModelRunnerEngineKind = iota // ModelRunnerEngineKindMobyManual represents a non-Desktop/Cloud engine - // that's explicitly targeted by a DMR_HOST environment variable on which - // the user is responsible for managing a Model Runner. + // that's explicitly targeted by a MODEL_RUNNER_HOST environment variable on + // which the user is responsible for managing a Model Runner. ModelRunnerEngineKindMobyManual // ModelRunnerEngineKindDesktop represents a Docker Desktop engine. It only // refers to a Docker Desktop Linux engine, i.e. not a Windows container @@ -135,14 +137,21 @@ func NewContextForMock(client DockerHttpClient) *ModelRunnerContext { // DetectContext determines the current Docker Model Runner context. func DetectContext(cli *command.DockerCli) (*ModelRunnerContext, error) { // Check for an explicit endpoint setting. - dmrHost := os.Getenv("DMR_HOST") + modelRunnerHost := os.Getenv("MODEL_RUNNER_HOST") + + // Check if we're treating Docker Desktop as regular Moby. This is only for + // testing purposes. + treatDesktopAsMoby := os.Getenv("_MODEL_RUNNER_TREAT_DESKTOP_AS_MOBY") != "" // Detect the associated engine type. kind := ModelRunnerEngineKindMoby - if dmrHost != "" { + if modelRunnerHost != "" { kind = ModelRunnerEngineKindMobyManual } else if isDesktopContext(cli) { kind = ModelRunnerEngineKindDesktop + if treatDesktopAsMoby { + kind = ModelRunnerEngineKindMoby + } } else if isCloudContext(cli) { kind = ModelRunnerEngineKindCloud } @@ -150,11 +159,14 @@ func DetectContext(cli *command.DockerCli) (*ModelRunnerContext, error) { // Compute the URL prefix based on the associated engine kind. var rawURLPrefix string if kind == ModelRunnerEngineKindMoby { - rawURLPrefix = "http://localhost:12434" + rawURLPrefix = "http://localhost:" + strconv.Itoa(int(standalone.DefaultControllerPort)) } else if kind == ModelRunnerEngineKindMobyManual { - rawURLPrefix = dmrHost + rawURLPrefix = modelRunnerHost } else if kind == ModelRunnerEngineKindDesktop { rawURLPrefix = "http://localhost" + inference.ExperimentalEndpointsPrefix + if treatDesktopAsMoby { + rawURLPrefix = "http://localhost:" + strconv.Itoa(int(standalone.DefaultControllerPort)) + } } else { // ModelRunnerEngineKindCloud rawURLPrefix = "http://localhost/" } @@ -168,13 +180,16 @@ func DetectContext(cli *command.DockerCli) (*ModelRunnerContext, error) { if kind == ModelRunnerEngineKindMoby || kind == ModelRunnerEngineKindMobyManual { client = http.DefaultClient } else if kind == ModelRunnerEngineKindDesktop { - dockerClient, err := dockerClientForContext(cli, "desktop-linux") + dockerClient, err := DockerClientForContext(cli, "desktop-linux") if err != nil { return nil, fmt.Errorf("unable to create model runner client: %w", err) } client = dockerClient.HTTPClient() + if treatDesktopAsMoby { + client = http.DefaultClient + } } else { // ModelRunnerEngineKindCloud - dockerClient, err := dockerClientForContext(cli, cli.CurrentContext()) + dockerClient, err := DockerClientForContext(cli, cli.CurrentContext()) if err != nil { return nil, fmt.Errorf("unable to create model runner client: %w", err) } diff --git a/pkg/standalone/containers.go b/pkg/standalone/containers.go new file mode 100644 index 00000000..29c7f9a1 --- /dev/null +++ b/pkg/standalone/containers.go @@ -0,0 +1,127 @@ +package standalone + +import ( + "context" + "fmt" + "strconv" + "strings" + + "github.com/docker/docker/api/types/container" + "github.com/docker/docker/api/types/filters" + "github.com/docker/docker/api/types/mount" + "github.com/docker/docker/client" + "github.com/docker/go-connections/nat" +) + +// controllerContainerName is the name to use for the controller container. +const controllerContainerName = "docker-model-runner" + +// FindControllerContainer searches for a running controller container. It +// returns the ID of the container (if found), the container name (if any), or +// any error that occurred. +func FindControllerContainer(ctx context.Context, dockerClient *client.Client) (string, string, error) { + // Identify all controller containers. + containers, err := dockerClient.ContainerList(ctx, container.ListOptions{ + Filters: filters.NewArgs(filters.Arg("label", labelRole+"="+roleController)), + }) + if err != nil { + return "", "", fmt.Errorf("unable to identify model runner containers: %w", err) + } + if len(containers) == 0 { + return "", "", nil + } + var containerName string + if len(containers[0].Names) > 0 { + containerName = strings.TrimPrefix(containers[0].Names[0], "/") + } + return containers[0].ID, containerName, nil +} + +// CreateControllerContainer creates and starts a controller container. +func CreateControllerContainer(ctx context.Context, dockerClient *client.Client, port uint16, gpu bool, modelStorageVolume string, printer StatusPrinter) error { + // Set up the container configuration. + portStr := strconv.Itoa(int(port)) + imageName := ControllerImage + ":" + controllerImageTagCPU + if gpu { + imageName = ControllerImage + ":" + controllerImageTagGPU + } + config := &container.Config{ + Image: imageName, + Env: []string{ + "MODEL_RUNNER_PORT=" + portStr, + }, + ExposedPorts: nat.PortSet{ + nat.Port(portStr + "/tcp"): struct{}{}, + }, + Labels: map[string]string{ + labelRole: roleController, + }, + } + hostConfig := &container.HostConfig{ + Mounts: []mount.Mount{ + { + Type: mount.TypeVolume, + Source: modelStorageVolume, + Target: "/models", + }, + }, + PortBindings: nat.PortMap{ + nat.Port(portStr + "/tcp"): []nat.PortBinding{{HostIP: "", HostPort: portStr}}, + }, + RestartPolicy: container.RestartPolicy{ + Name: "always", + }, + } + if gpu { + hostConfig.Resources = container.Resources{ + DeviceRequests: []container.DeviceRequest{ + { + Driver: "nvidia", + Count: -1, + Capabilities: [][]string{{"gpu"}}, + }, + }, + } + } + + // Create the container. + resp, err := dockerClient.ContainerCreate(ctx, config, hostConfig, nil, nil, controllerContainerName) + if err != nil { + return fmt.Errorf("failed to create container %s: %w", controllerContainerName, err) + } + + // Start the container. + printer.Printf("Starting model runner container %s...\n", controllerContainerName) + if err := dockerClient.ContainerStart(ctx, resp.ID, container.StartOptions{}); err != nil { + _ = dockerClient.ContainerRemove(ctx, resp.ID, container.RemoveOptions{Force: true}) + return fmt.Errorf("failed to start container %s: %w", controllerContainerName, err) + } + return nil +} + +// PruneControllerContainers stops and removes any model runner controller +// containers. +func PruneControllerContainers(ctx context.Context, dockerClient *client.Client, printer StatusPrinter) error { + // Identify all controller containers. + containers, err := dockerClient.ContainerList(ctx, container.ListOptions{ + All: true, + Filters: filters.NewArgs(filters.Arg("label", labelRole+"="+roleController)), + }) + if err != nil { + return fmt.Errorf("unable to identify model runner containers: %w", err) + } + + // Remove all controller containers. + for _, ctr := range containers { + if len(ctr.Names) > 0 { + printer.Printf("Removing container %s (%s)...\n", strings.TrimPrefix(ctr.Names[0], "/"), ctr.ID[:12]) + } else { + printer.Printf("Removing container %s...\n", ctr.ID[:12]) + } + err := dockerClient.ContainerRemove(ctx, ctr.ID, container.RemoveOptions{Force: true}) + if err != nil { + return fmt.Errorf("failed to remove container %s: %w", ctr.Names[0], err) + } + } + return nil +} diff --git a/pkg/standalone/images.go b/pkg/standalone/images.go new file mode 100644 index 00000000..2eefdc52 --- /dev/null +++ b/pkg/standalone/images.go @@ -0,0 +1,75 @@ +package standalone + +import ( + "context" + "encoding/json" + "fmt" + "io" + + "github.com/docker/docker/api/types/image" + "github.com/docker/docker/client" + "github.com/docker/docker/pkg/jsonmessage" +) + +const ( + // ControllerImage is the image used for the controller container. + ControllerImage = "docker/model-runner" + // controllerImageTagCPU is the image tag used for the controller container + // when running with the CPU backend. + controllerImageTagCPU = "latest" + // controllerImageTagGPU is the image tag used for the controller container + // when running with the GPU backend. + controllerImageTagGPU = "latest-cuda" +) + +// EnsureControllerImage ensures that the controller container image is pulled. +func EnsureControllerImage(ctx context.Context, dockerClient *client.Client, gpu bool, printer StatusPrinter) error { + // Determine the target image. + imageName := ControllerImage + ":" + controllerImageTagCPU + if gpu { + imageName = ControllerImage + ":" + controllerImageTagGPU + } + + // Perform the pull. + out, err := dockerClient.ImagePull(ctx, imageName, image.PullOptions{}) + if err != nil { + return fmt.Errorf("failed to pull image %s: %w", imageName, err) + } + defer out.Close() + + // Decode and print status updates. + decoder := json.NewDecoder(out) + for { + var response jsonmessage.JSONMessage + if err := decoder.Decode(&response); err != nil { + if err == io.EOF { + break + } + return fmt.Errorf("failed to decode pull response: %w", err) + } + + if response.ID != "" { + printer.Printf("\r%s: %s %s", response.ID, response.Status, response.ProgressMessage) + } else { + printer.Println(response.Status) + } + } + printer.Println("\nSuccessfully pulled", imageName) + return nil +} + +// PruneControllerImages removes any unused controller container images. +func PruneControllerImages(ctx context.Context, dockerClient *client.Client, printer StatusPrinter) error { + // Remove the standard image, if present. + imageNameCPU := ControllerImage + ":" + controllerImageTagCPU + if _, err := dockerClient.ImageRemove(ctx, imageNameCPU, image.RemoveOptions{}); err == nil { + printer.Println("Removed image", imageNameCPU) + } + + // Remove the GPU image, if present. + imageNameGPU := ControllerImage + ":" + controllerImageTagGPU + if _, err := dockerClient.ImageRemove(ctx, imageNameGPU, image.RemoveOptions{}); err == nil { + printer.Println("Removed image", imageNameGPU) + } + return nil +} diff --git a/pkg/standalone/labels.go b/pkg/standalone/labels.go new file mode 100644 index 00000000..72ea8cc7 --- /dev/null +++ b/pkg/standalone/labels.go @@ -0,0 +1,19 @@ +package standalone + +const ( + // labelRole is the label used to identify a Docker object's role in the + // standalone model runner infrastructure. + labelRole = "com.docker.model-runner.role" + + // roleController is the role label value used to identify the model runner + // controller container. + roleController = "controller" + + // roleRunner is not currently defined because model runner processes + // currently execute within the controller container. This may change in a + // future release. + + // roleModelStorage is the role label value used to identify the model + // runner model storage volume. + roleModelStorage = "model-storage" +) diff --git a/pkg/standalone/ports.go b/pkg/standalone/ports.go new file mode 100644 index 00000000..0c7a99a8 --- /dev/null +++ b/pkg/standalone/ports.go @@ -0,0 +1,5 @@ +package standalone + +// DefaultControllerPort is the default TCP port on which the standalone +// controller will listen for requests. +const DefaultControllerPort = 12434 diff --git a/pkg/standalone/status.go b/pkg/standalone/status.go new file mode 100644 index 00000000..427e9ff7 --- /dev/null +++ b/pkg/standalone/status.go @@ -0,0 +1,9 @@ +package standalone + +// StatusPrinter is the interface used to print status updates. +type StatusPrinter interface { + // Printf should perform formatted printing. + Printf(format string, args ...any) + // Println should perform line-based printing. + Println(args ...any) +} diff --git a/pkg/standalone/volumes.go b/pkg/standalone/volumes.go new file mode 100644 index 00000000..4812855a --- /dev/null +++ b/pkg/standalone/volumes.go @@ -0,0 +1,65 @@ +package standalone + +import ( + "context" + "fmt" + + "github.com/docker/docker/api/types/filters" + "github.com/docker/docker/api/types/volume" + "github.com/docker/docker/client" +) + +// modelStorageVolumeName is the name to use for the model storage volume. +const modelStorageVolumeName = "docker-model-runner-models" + +// EnsureModelStorageVolume ensures that a model storage volume exists, creating +// it if necessary. It returns the name of the storage volume or any error that +// occurred. +func EnsureModelStorageVolume(ctx context.Context, dockerClient *client.Client, printer StatusPrinter) (string, error) { + // Try to identify the storage volume. + volumes, err := dockerClient.VolumeList(ctx, volume.ListOptions{ + Filters: filters.NewArgs( + filters.Arg("label", labelRole+"="+roleModelStorage), + ), + }) + if err != nil { + return "", fmt.Errorf("unable to list volumes: %w", err) + } + + // If any volumes with the correct role exist (ideally there should only be + // one), then pick the first one. + if len(volumes.Volumes) > 0 { + return volumes.Volumes[0].Name, nil + } + + // Create the volume. + printer.Printf("Creating model storage volume %s...\n", modelStorageVolumeName) + volume, err := dockerClient.VolumeCreate(ctx, volume.CreateOptions{ + Name: modelStorageVolumeName, + Labels: map[string]string{ + labelRole: roleModelStorage, + }, + }) + if err != nil { + return "", fmt.Errorf("unable to create volume: %w", err) + } + return volume.Name, nil +} + +// PruneModelStorageVolumes removes any unused model storage volume(s). +func PruneModelStorageVolumes(ctx context.Context, dockerClient *client.Client, printer StatusPrinter) error { + pruned, err := dockerClient.VolumesPrune(ctx, filters.NewArgs( + filters.Arg("all", "true"), + filters.Arg("label", labelRole+"="+roleModelStorage), + )) + if err != nil { + return err + } + for _, volume := range pruned.VolumesDeleted { + printer.Println("Removed volume", volume) + } + if pruned.SpaceReclaimed > 0 { + fmt.Printf("Reclaimed %d bytes\n", pruned.SpaceReclaimed) + } + return nil +} diff --git a/vendor/github.com/docker/docker/pkg/jsonmessage/jsonmessage.go b/vendor/github.com/docker/docker/pkg/jsonmessage/jsonmessage.go new file mode 100644 index 00000000..1a05de4d --- /dev/null +++ b/vendor/github.com/docker/docker/pkg/jsonmessage/jsonmessage.go @@ -0,0 +1,314 @@ +package jsonmessage // import "github.com/docker/docker/pkg/jsonmessage" + +import ( + "encoding/json" + "fmt" + "io" + "strings" + "time" + + "github.com/docker/go-units" + "github.com/moby/term" + "github.com/morikuni/aec" +) + +// RFC3339NanoFixed is time.RFC3339Nano with nanoseconds padded using zeros to +// ensure the formatted time isalways the same number of characters. +const RFC3339NanoFixed = "2006-01-02T15:04:05.000000000Z07:00" + +// JSONError wraps a concrete Code and Message, Code is +// an integer error code, Message is the error message. +type JSONError struct { + Code int `json:"code,omitempty"` + Message string `json:"message,omitempty"` +} + +func (e *JSONError) Error() string { + return e.Message +} + +// JSONProgress describes a progress message in a JSON stream. +type JSONProgress struct { + // Current is the current status and value of the progress made towards Total. + Current int64 `json:"current,omitempty"` + // Total is the end value describing when we made 100% progress for an operation. + Total int64 `json:"total,omitempty"` + // Start is the initial value for the operation. + Start int64 `json:"start,omitempty"` + // HideCounts. if true, hides the progress count indicator (xB/yB). + HideCounts bool `json:"hidecounts,omitempty"` + // Units is the unit to print for progress. It defaults to "bytes" if empty. + Units string `json:"units,omitempty"` + + // terminalFd is the fd of the current terminal, if any. It is used + // to get the terminal width. + terminalFd uintptr + + // nowFunc is used to override the current time in tests. + nowFunc func() time.Time + + // winSize is used to override the terminal width in tests. + winSize int +} + +func (p *JSONProgress) String() string { + var ( + width = p.width() + pbBox string + numbersBox string + ) + if p.Current <= 0 && p.Total <= 0 { + return "" + } + if p.Total <= 0 { + switch p.Units { + case "": + return fmt.Sprintf("%8v", units.HumanSize(float64(p.Current))) + default: + return fmt.Sprintf("%d %s", p.Current, p.Units) + } + } + + percentage := int(float64(p.Current)/float64(p.Total)*100) / 2 + if percentage > 50 { + percentage = 50 + } + if width > 110 { + // this number can't be negative gh#7136 + numSpaces := 0 + if 50-percentage > 0 { + numSpaces = 50 - percentage + } + pbBox = fmt.Sprintf("[%s>%s] ", strings.Repeat("=", percentage), strings.Repeat(" ", numSpaces)) + } + + switch { + case p.HideCounts: + case p.Units == "": // no units, use bytes + current := units.HumanSize(float64(p.Current)) + total := units.HumanSize(float64(p.Total)) + + numbersBox = fmt.Sprintf("%8v/%v", current, total) + + if p.Current > p.Total { + // remove total display if the reported current is wonky. + numbersBox = fmt.Sprintf("%8v", current) + } + default: + numbersBox = fmt.Sprintf("%d/%d %s", p.Current, p.Total, p.Units) + + if p.Current > p.Total { + // remove total display if the reported current is wonky. + numbersBox = fmt.Sprintf("%d %s", p.Current, p.Units) + } + } + + // Show approximation of remaining time if there's enough width. + var timeLeftBox string + if width > 50 { + if p.Current > 0 && p.Start > 0 && percentage < 50 { + fromStart := p.now().Sub(time.Unix(p.Start, 0)) + perEntry := fromStart / time.Duration(p.Current) + left := time.Duration(p.Total-p.Current) * perEntry + timeLeftBox = " " + left.Round(time.Second).String() + } + } + return pbBox + numbersBox + timeLeftBox +} + +// now returns the current time in UTC, but can be overridden in tests +// by setting JSONProgress.nowFunc to a custom function. +func (p *JSONProgress) now() time.Time { + if p.nowFunc != nil { + return p.nowFunc() + } + return time.Now().UTC() +} + +// width returns the current terminal's width, but can be overridden +// in tests by setting JSONProgress.winSize to a non-zero value. +func (p *JSONProgress) width() int { + if p.winSize != 0 { + return p.winSize + } + ws, err := term.GetWinsize(p.terminalFd) + if err == nil { + return int(ws.Width) + } + return 200 +} + +// JSONMessage defines a message struct. It describes +// the created time, where it from, status, ID of the +// message. It's used for docker events. +type JSONMessage struct { + Stream string `json:"stream,omitempty"` + Status string `json:"status,omitempty"` + Progress *JSONProgress `json:"progressDetail,omitempty"` + + // ProgressMessage is a pre-formatted presentation of [Progress]. + // + // Deprecated: this field is deprecated since docker v0.7.1 / API v1.8. Use the information in [Progress] instead. This field will be omitted in a future release. + ProgressMessage string `json:"progress,omitempty"` + ID string `json:"id,omitempty"` + From string `json:"from,omitempty"` + Time int64 `json:"time,omitempty"` + TimeNano int64 `json:"timeNano,omitempty"` + Error *JSONError `json:"errorDetail,omitempty"` + + // ErrorMessage contains errors encountered during the operation. + // + // Deprecated: this field is deprecated since docker v0.6.0 / API v1.4. Use [Error.Message] instead. This field will be omitted in a future release. + ErrorMessage string `json:"error,omitempty"` // deprecated + // Aux contains out-of-band data, such as digests for push signing and image id after building. + Aux *json.RawMessage `json:"aux,omitempty"` +} + +func clearLine(out io.Writer) { + eraseMode := aec.EraseModes.All + cl := aec.EraseLine(eraseMode) + fmt.Fprint(out, cl) +} + +func cursorUp(out io.Writer, l uint) { + fmt.Fprint(out, aec.Up(l)) +} + +func cursorDown(out io.Writer, l uint) { + fmt.Fprint(out, aec.Down(l)) +} + +// Display prints the JSONMessage to out. If isTerminal is true, it erases +// the entire current line when displaying the progressbar. It returns an +// error if the [JSONMessage.Error] field is non-nil. +func (jm *JSONMessage) Display(out io.Writer, isTerminal bool) error { + if jm.Error != nil { + return jm.Error + } + var endl string + if isTerminal && jm.Stream == "" && jm.Progress != nil { + clearLine(out) + endl = "\r" + fmt.Fprint(out, endl) + } else if jm.Progress != nil && jm.Progress.String() != "" { // disable progressbar in non-terminal + return nil + } + if jm.TimeNano != 0 { + fmt.Fprintf(out, "%s ", time.Unix(0, jm.TimeNano).Format(RFC3339NanoFixed)) + } else if jm.Time != 0 { + fmt.Fprintf(out, "%s ", time.Unix(jm.Time, 0).Format(RFC3339NanoFixed)) + } + if jm.ID != "" { + fmt.Fprintf(out, "%s: ", jm.ID) + } + if jm.From != "" { + fmt.Fprintf(out, "(from %s) ", jm.From) + } + if jm.Progress != nil && isTerminal { + fmt.Fprintf(out, "%s %s%s", jm.Status, jm.Progress.String(), endl) + } else if jm.ProgressMessage != "" { // deprecated + fmt.Fprintf(out, "%s %s%s", jm.Status, jm.ProgressMessage, endl) + } else if jm.Stream != "" { + fmt.Fprintf(out, "%s%s", jm.Stream, endl) + } else { + fmt.Fprintf(out, "%s%s\n", jm.Status, endl) + } + return nil +} + +// DisplayJSONMessagesStream reads a JSON message stream from in, and writes +// each [JSONMessage] to out. It returns an error if an invalid JSONMessage +// is received, or if a JSONMessage containers a non-zero [JSONMessage.Error]. +// +// Presentation of the JSONMessage depends on whether a terminal is attached, +// and on the terminal width. Progress bars ([JSONProgress]) are suppressed +// on narrower terminals (< 110 characters). +// +// - isTerminal describes if out is a terminal, in which case it prints +// a newline ("\n") at the end of each line and moves the cursor while +// displaying. +// - terminalFd is the fd of the current terminal (if any), and used +// to get the terminal width. +// - auxCallback allows handling the [JSONMessage.Aux] field. It is +// called if a JSONMessage contains an Aux field, in which case +// DisplayJSONMessagesStream does not present the JSONMessage. +func DisplayJSONMessagesStream(in io.Reader, out io.Writer, terminalFd uintptr, isTerminal bool, auxCallback func(JSONMessage)) error { + var ( + dec = json.NewDecoder(in) + ids = make(map[string]uint) + ) + + for { + var diff uint + var jm JSONMessage + if err := dec.Decode(&jm); err != nil { + if err == io.EOF { + break + } + return err + } + + if jm.Aux != nil { + if auxCallback != nil { + auxCallback(jm) + } + continue + } + + if jm.Progress != nil { + jm.Progress.terminalFd = terminalFd + } + if jm.ID != "" && (jm.Progress != nil || jm.ProgressMessage != "") { + line, ok := ids[jm.ID] + if !ok { + // NOTE: This approach of using len(id) to + // figure out the number of lines of history + // only works as long as we clear the history + // when we output something that's not + // accounted for in the map, such as a line + // with no ID. + line = uint(len(ids)) + ids[jm.ID] = line + if isTerminal { + fmt.Fprintf(out, "\n") + } + } + diff = uint(len(ids)) - line + if isTerminal { + cursorUp(out, diff) + } + } else { + // When outputting something that isn't progress + // output, clear the history of previous lines. We + // don't want progress entries from some previous + // operation to be updated (for example, pull -a + // with multiple tags). + ids = make(map[string]uint) + } + err := jm.Display(out, isTerminal) + if jm.ID != "" && isTerminal { + cursorDown(out, diff) + } + if err != nil { + return err + } + } + return nil +} + +// Stream is an io.Writer for output with utilities to get the output's file +// descriptor and to detect whether it's a terminal. +// +// it is subset of the streams.Out type in +// https://pkg.go.dev/github.com/docker/cli@v20.10.17+incompatible/cli/streams#Out +type Stream interface { + io.Writer + FD() uintptr + IsTerminal() bool +} + +// DisplayJSONMessagesToStream prints json messages to the output Stream. It is +// used by the Docker CLI to print JSONMessage streams. +func DisplayJSONMessagesToStream(in io.Reader, stream Stream, auxCallback func(JSONMessage)) error { + return DisplayJSONMessagesStream(in, stream, stream.FD(), stream.IsTerminal(), auxCallback) +} diff --git a/vendor/modules.txt b/vendor/modules.txt index 90408ce0..e628c313 100644 --- a/vendor/modules.txt +++ b/vendor/modules.txt @@ -114,6 +114,7 @@ github.com/docker/docker/internal/lazyregexp github.com/docker/docker/internal/multierror github.com/docker/docker/pkg/atomicwriter github.com/docker/docker/pkg/homedir +github.com/docker/docker/pkg/jsonmessage github.com/docker/docker/registry # github.com/docker/docker-credential-helpers v0.9.3 ## explicit; go 1.21