introduce support for models

Signed-off-by: Nicolas De Loof <nicolas.deloof@gmail.com>

go.mod

@@ -212,3 +212,5 @@ exclude (
 	github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc
 	github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2
 )
+
+replace github.com/compose-spec/compose-go/v2 => github.com/ndeloof/compose-go/v2 v2.0.1-0.20250625082240-b948fe935f02

go.sum

@@ -80,8 +80,6 @@ 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/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=
@@ -361,6 +359,8 @@ github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822/go.mod h1:+n7T8m
 github.com/mwitkow/go-conntrack v0.0.0-20161129095857-cc309e4a2223/go.mod h1:qRWi+5nqEBWmkhHvq77mSJWrCKwh8bxhgT7d/eI7P4U=
 github.com/mxk/go-flowrate v0.0.0-20140419014527-cca7078d478f h1:y5//uYreIhSUg3J1GEMiLbxo1LJaP8RfCpH6pymGZus=
 github.com/mxk/go-flowrate v0.0.0-20140419014527-cca7078d478f/go.mod h1:ZdcZmHo+o7JKHSa8/e818NopupXU1YMK5fe1lsApnBw=
+github.com/ndeloof/compose-go/v2 v2.0.1-0.20250625082240-b948fe935f02 h1:RPYzx1y7ldfYB8Ba2INxr6FiW2ZxXHLl8it775gz0qE=
+github.com/ndeloof/compose-go/v2 v2.0.1-0.20250625082240-b948fe935f02/go.mod h1:TmjkIB9W73fwVxkYY+u2uhMbMUakjiif79DlYgXsyvU=
 github.com/niemeyer/pretty v0.0.0-20200227124842-a10e7caefd8e/go.mod h1:zD1mROLANZcx1PVRCS0qkT7pwLkGfwJo4zjcN/Tysno=
 github.com/onsi/ginkgo v1.6.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE=
 github.com/onsi/ginkgo v1.12.0 h1:Iw5WCbBcaAAd0fpRb1c9r5YCylv4XDoCSigm1zLevwU=

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,198 @@
+/*
+   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"
+	"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)
+	}
+
+	eg, gctx := errgroup.WithContext(ctx)
+	eg.Go(func() error {
+		return s.setModelEndpointVariable(gctx, dockerModel, project)
+	})
+
+MODELS:
+	for name, config := range project.Models {
+		for _, model := range models {
+			if slices.Contains(model.Tags, config.Model) {
+				continue MODELS
+			}
+		}
+		if config.Name == "" {
+			config.Name = name
+		}
+		eg.Go(func() error {
+			return s.pullModel(gctx, dockerModel, config, quietPull)
+		})
+	}
+	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) 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.Variable != "" {
+				variable = modelConfig.Variable
+			} 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))
+		}
 	}
 }
 

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("require model-cli on GHA runners")
+	c := NewParallelCLI(t)
+	defer c.cleanupWithDown(t, "model-test")
+
+	c.RunDockerComposeCmd(t, "-f", "./fixtures/model/compose.yaml", "run", "test", "sh", "-c", "curl ${FOO_URL}")
+}