introduce support for models

.github/workflows/ci.yml

@@ -183,6 +183,11 @@ jobs:
       - name: Set up Docker Buildx
         uses: docker/setup-buildx-action@v3
 
+      - name: Set up Docker Model
+        run: |
+          sudo apt-get install docker-model-plugin
+          docker model version
+
       - name: Set up Go
         uses: actions/setup-go@v5
         with:

go.mod

@@ -8,7 +8,7 @@ require (
 	github.com/Microsoft/go-winio v0.6.2
 	github.com/acarl005/stripansi v0.0.0-20180116102854-5a71ef0e047d
 	github.com/buger/goterm v1.0.4
-	github.com/compose-spec/compose-go/v2 v2.6.5
+	github.com/compose-spec/compose-go/v2 v2.7.0
 	github.com/containerd/containerd/v2 v2.1.3
 	github.com/containerd/errdefs v1.0.0
 	github.com/containerd/platforms v1.0.0-rc.1

go.sum

@@ -80,8 +80,8 @@ github.com/cloudflare/cfssl v0.0.0-20180223231731-4e2dcbde5004 h1:lkAMpLVBDaj17e
 github.com/cloudflare/cfssl v0.0.0-20180223231731-4e2dcbde5004/go.mod h1:yMWuSON2oQp+43nFtAV/uvKQIFpSPerB57DCt9t8sSA=
 github.com/codahale/rfc6979 v0.0.0-20141003034818-6a90f24967eb h1:EDmT6Q9Zs+SbUoc7Ik9EfrFqcylYqgPZ9ANSbTAntnE=
 github.com/codahale/rfc6979 v0.0.0-20141003034818-6a90f24967eb/go.mod h1:ZjrT6AXHbDs86ZSdt/osfBi5qfexBrKUdONk989Wnk4=
-github.com/compose-spec/compose-go/v2 v2.6.5 h1:H7xP5OMKdkN2p0brx01slxIU6dE/q6ybbG+jozPtIqk=
-github.com/compose-spec/compose-go/v2 v2.6.5/go.mod h1:TmjkIB9W73fwVxkYY+u2uhMbMUakjiif79DlYgXsyvU=
+github.com/compose-spec/compose-go/v2 v2.7.0 h1:wA8pNN93Gw8XSZ+JSWZx9Wx8xoQtuSzvf4SS2KklwH0=
+github.com/compose-spec/compose-go/v2 v2.7.0/go.mod h1:TmjkIB9W73fwVxkYY+u2uhMbMUakjiif79DlYgXsyvU=
 github.com/containerd/cgroups/v3 v3.0.5 h1:44na7Ud+VwyE7LIoJ8JTNQOa549a8543BmzaJHo6Bzo=
 github.com/containerd/cgroups/v3 v3.0.5/go.mod h1:SA5DLYnXO8pTGYiAHXz94qvLQTKfVM5GEVisn4jpins=
 github.com/containerd/console v1.0.5 h1:R0ymNeydRqH2DmakFNdmjR2k0t7UPuiOV/N/27/qqsc=

pkg/compose/create.go

@@ -83,6 +83,11 @@ func (s *composeService) create(ctx context.Context, project *types.Project, opt
 		return err
 	}
 
+	err = s.ensureModels(ctx, project, options.QuietPull)
+	if err != nil {
+		return err
+	}
+
 	prepareNetworks(project)
 
 	networks, err := s.ensureNetworks(ctx, project)

pkg/compose/model.go

@@ -0,0 +1,219 @@
+/*
+   Copyright 2020 Docker Compose CLI authors
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
+*/
+
+package compose
+
+import (
+	"bufio"
+	"context"
+	"encoding/json"
+	"fmt"
+	"os"
+	"os/exec"
+	"slices"
+	"strconv"
+	"strings"
+
+	"github.com/compose-spec/compose-go/v2/types"
+	"github.com/containerd/errdefs"
+	"github.com/docker/cli/cli-plugins/manager"
+	"github.com/docker/compose/v2/pkg/progress"
+	"github.com/spf13/cobra"
+	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/otel/propagation"
+	"golang.org/x/sync/errgroup"
+)
+
+func (s *composeService) ensureModels(ctx context.Context, project *types.Project, quietPull bool) error {
+	if len(project.Models) == 0 {
+		return nil
+	}
+
+	dockerModel, err := manager.GetPlugin("model", s.dockerCli, &cobra.Command{})
+	if err != nil {
+		if errdefs.IsNotFound(err) {
+			return fmt.Errorf("'models' support requires Docker Model plugin")
+		}
+		return err
+	}
+
+	cmd := exec.CommandContext(ctx, dockerModel.Path, "ls", "--json")
+	s.setupChildProcess(ctx, cmd)
+	output, err := cmd.CombinedOutput()
+	if err != nil {
+		return fmt.Errorf("error checking available models: %w", err)
+	}
+
+	type AvailableModel struct {
+		Id      string   `json:"id"`
+		Tags    []string `json:"tags"`
+		Created int      `json:"created"`
+	}
+
+	models := []AvailableModel{}
+	err = json.Unmarshal(output, &models)
+	if err != nil {
+		return fmt.Errorf("error unmarshalling available models: %w", err)
+	}
+	var availableModels []string
+	for _, model := range models {
+		availableModels = append(availableModels, model.Tags...)
+	}
+
+	eg, gctx := errgroup.WithContext(ctx)
+	eg.Go(func() error {
+		return s.setModelEndpointVariable(gctx, dockerModel, project)
+	})
+
+	for name, config := range project.Models {
+		if config.Name == "" {
+			config.Name = name
+		}
+		eg.Go(func() error {
+			if !slices.Contains(availableModels, config.Model) {
+				err = s.pullModel(gctx, dockerModel, config, quietPull)
+				if err != nil {
+					return err
+				}
+			}
+			return s.configureModel(gctx, dockerModel, config)
+		})
+	}
+	return eg.Wait()
+}
+
+func (s *composeService) pullModel(ctx context.Context, dockerModel *manager.Plugin, model types.ModelConfig, quietPull bool) error {
+	w := progress.ContextWriter(ctx)
+	w.Event(progress.Event{
+		ID:     model.Name,
+		Status: progress.Working,
+		Text:   "Pulling",
+	})
+
+	cmd := exec.CommandContext(ctx, dockerModel.Path, "pull", model.Model)
+	s.setupChildProcess(ctx, cmd)
+
+	stream, err := cmd.StdoutPipe()
+	if err != nil {
+		return err
+	}
+
+	err = cmd.Start()
+	if err != nil {
+		return err
+	}
+
+	scanner := bufio.NewScanner(stream)
+	for scanner.Scan() {
+		msg := scanner.Text()
+		if msg == "" {
+			continue
+		}
+
+		if !quietPull {
+			w.Event(progress.Event{
+				ID:         model.Name,
+				Status:     progress.Working,
+				Text:       "Pulling",
+				StatusText: msg,
+			})
+		}
+	}
+
+	err = cmd.Wait()
+	if err != nil {
+		w.Event(progress.ErrorMessageEvent(model.Name, err.Error()))
+	}
+	w.Event(progress.Event{
+		ID:     model.Name,
+		Status: progress.Working,
+		Text:   "Pulled",
+	})
+	return err
+}
+
+func (s *composeService) configureModel(ctx context.Context, dockerModel *manager.Plugin, config types.ModelConfig) error {
+	// configure [--context-size=<n>] MODEL [-- <runtime-flags...>]
+	args := []string{"configure"}
+	if config.ContextSize > 0 {
+		args = append(args, "--context-size", strconv.Itoa(config.ContextSize))
+	}
+	args = append(args, config.Model)
+	if len(config.RuntimeFlags) != 0 {
+		args = append(args, "--")
+		args = append(args, config.RuntimeFlags...)
+	}
+	cmd := exec.CommandContext(ctx, dockerModel.Path, args...)
+	s.setupChildProcess(ctx, cmd)
+	return cmd.Run()
+}
+
+func (s *composeService) setModelEndpointVariable(ctx context.Context, dockerModel *manager.Plugin, project *types.Project) error {
+	cmd := exec.CommandContext(ctx, dockerModel.Path, "status", "--json")
+	s.setupChildProcess(ctx, cmd)
+	statusOut, err := cmd.CombinedOutput()
+	if err != nil {
+		return fmt.Errorf("error checking docker-model status: %w", err)
+	}
+	type Status struct {
+		Endpoint string `json:"endpoint"`
+	}
+
+	var status Status
+	err = json.Unmarshal(statusOut, &status)
+	if err != nil {
+		return err
+	}
+
+	for _, service := range project.Services {
+		for model, modelConfig := range service.Models {
+			var variable string
+			if modelConfig != nil && modelConfig.EndpointVariable != "" {
+				variable = modelConfig.EndpointVariable
+			} else {
+				variable = strings.ToUpper(model) + "_URL"
+			}
+			service.Environment[variable] = &status.Endpoint
+		}
+	}
+	return nil
+}
+
+func (s *composeService) setupChildProcess(gctx context.Context, cmd *exec.Cmd) {
+	// exec provider command with same environment Compose is running
+	env := types.NewMapping(os.Environ())
+	// but remove DOCKER_CLI_PLUGIN... variable so plugin can detect it run standalone
+	delete(env, manager.ReexecEnvvar)
+	// propagate opentelemetry context to child process, see https://github.com/open-telemetry/oteps/blob/main/text/0258-env-context-baggage-carriers.md
+	carrier := propagation.MapCarrier{}
+	otel.GetTextMapPropagator().Inject(gctx, &carrier)
+	env.Merge(types.Mapping(carrier))
+	env["DOCKER_CONTEXT"] = s.dockerCli.CurrentContext()
+	cmd.Env = env.Values()
+}
+
+type Model struct {
+	Id      string   `json:"id"`
+	Tags    []string `json:"tags"`
+	Created int      `json:"created"`
+	Config  struct {
+		Format       string `json:"format"`
+		Quantization string `json:"quantization"`
+		Parameters   string `json:"parameters"`
+		Architecture string `json:"architecture"`
+		Size         string `json:"size"`
+	} `json:"config"`
+}

pkg/compose/run.go

@@ -123,6 +123,11 @@ func (s *composeService) prepareRun(ctx context.Context, project *types.Project,
 		return "", err
 	}
 
+	err = s.ensureModels(ctx, project, opts.QuietPull)
+	if err != nil {
+		return "", err
+	}
+
 	created, err := s.createContainer(ctx, project, service, service.ContainerName, -1, createOpts)
 	if err != nil {
 		return "", err

pkg/e2e/fixtures/model/compose.yaml

@@ -0,0 +1,9 @@
+services:
+  test:
+    image: alpine/curl
+    models:
+      - foo
+
+models:
+  foo:
+    model: ai/smollm2

pkg/e2e/framework.go

@@ -52,6 +52,9 @@ var (
 	// DockerBuildxExecutableName is the Os dependent Buildx plugin binary name
 	DockerBuildxExecutableName = "docker-buildx"
 
+	// DockerModelExecutableName is the Os dependent Docker-Model plugin binary name
+	DockerModelExecutableName = "docker-model"
+
 	// WindowsExecutableSuffix is the Windows executable suffix
 	WindowsExecutableSuffix = ".exe"
 )
@@ -162,6 +165,13 @@ func initializePlugins(t testing.TB, configDir string) {
 		}
 		// We don't need a functional scan plugin, but a valid plugin binary
 		CopyFile(t, composePlugin, filepath.Join(configDir, "cli-plugins", DockerScanExecutableName))
+
+		modelPlugin, err := findPluginExecutable(DockerModelExecutableName)
+		if err != nil {
+			t.Logf("WARNING: docker-model cli-plugin not found")
+		} else {
+			CopyFile(t, modelPlugin, filepath.Join(configDir, "cli-plugins", DockerModelExecutableName))
+		}
 	}
 }
 
@@ -214,13 +224,23 @@ func findPluginExecutable(pluginExecutableName string) (string, error) {
 	if err != nil {
 		return "", err
 	}
-	bin, err := filepath.Abs(filepath.Join(userDir, dockerUserDir, pluginExecutableName))
-	if err != nil {
-		return "", err
+	candidates := []string{
+		filepath.Join(userDir, dockerUserDir),
+		"/usr/local/lib/docker/cli-plugins",
+		"/usr/local/libexec/docker/cli-plugins",
+		"/usr/lib/docker/cli-plugins",
+		"/usr/libexec/docker/cli-plugins",
 	}
-	if _, err := os.Stat(bin); err == nil {
-		return bin, nil
+	for _, path := range candidates {
+		bin, err := filepath.Abs(filepath.Join(path, pluginExecutableName))
+		if err != nil {
+			return "", err
+		}
+		if _, err := os.Stat(bin); err == nil {
+			return bin, nil
+		}
 	}
+
 	return "", fmt.Errorf("plugin not found %s: %w", pluginExecutableName, os.ErrNotExist)
 }
 

pkg/e2e/model_test.go

@@ -0,0 +1,29 @@
+/*
+   Copyright 2020 Docker Compose CLI authors
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
+*/
+
+package e2e
+
+import (
+	"testing"
+)
+
+func TestComposeModel(t *testing.T) {
+	t.Skip("waiting for docker-model release")
+	c := NewParallelCLI(t)
+	defer c.cleanupWithDown(t, "model-test")
+
+	c.RunDockerComposeCmd(t, "-f", "./fixtures/model/compose.yaml", "run", "test", "sh", "-c", "curl ${FOO_URL}")
+}