feat: Add support for Docker Compose models

examples/06-resource-provisioning/README.md

@@ -182,9 +182,10 @@ The following extra fields can be configured as required on each instance of thi
 | `outputs`     | String, Go template | A Go template for a valid YAML dictionary. The values here are the outputs of the resource that can be accessed through `${resources.*}` placeholder resolution.                                                             |
 | `directories` | String, Go template | A Go template for a valid YAML dictionary. Each path -> bool mapping will create (true) or delete (false) a directory relative to the mounts directory.                                                                      |
 | `files`       | String, Go template | A Go template for a valid YAML dictionary. Each path -> string\|null will create a relative file (string) or delete it (null) relative to the mounts directory.                                                              |
-| `networks`    | String, Go template | A Go template for a valid set of named Compose [Networks](https://github.com/compose-spec/compose-spec/blob/master/06-networks.md). These will be added to the output project.                                               |
-| `volumes`     | String, Go template | A Go template for a valid set of named Compose [Volumes](https://github.com/compose-spec/compose-spec/blob/master/07-volumes.md).                                                                                            |
-| `services`    | String, Go template | A Go template for a valid set of named Compose [Services](https://github.com/compose-spec/compose-spec/blob/master/05-services.md).                                                                                          |
+| `networks`    | String, Go template | A Go template for a valid set of named Compose [Networks](https://github.com/compose-spec/compose-spec/blob/main/06-networks.md). These will be added to the output project.                                               |
+| `volumes`     | String, Go template | A Go template for a valid set of named Compose [Volumes](https://github.com/compose-spec/compose-spec/blob/main/07-volumes.md).                                                                                            |
+| `services`    | String, Go template | A Go template for a valid set of named Compose [Services](https://github.com/compose-spec/compose-spec/blob/main/05-services.md).                                                                                          |
+| `models`      | String, Go template | A Go template for a valid set of named Compose [Models](https://docs.docker.com/compose/how-tos/model-runner/). These define AI/ML models that can be used by services.                                                       |
 
 Each template has access to the [Sprig](http://masterminds.github.io/sprig/) functions library and executes with access to the following structure:
 
@@ -207,6 +208,46 @@ type Data struct {
 
 Browse the default provisioners for inspiration or more clues to how these work!
 
+#### Example: Provisioning a Model
+
+The `models` template field allows provisioners to define AI/ML models in the Compose output. Here's an example of a custom provisioner that provisions a model:
+
+```yaml
+- uri: template://llm-model
+  type: llm-model
+  class: default
+  init: |
+    name: {{ .Id | replace "." "-" }}
+  models: |
+    {{ .Init.name }}:
+      model: {{ .Params.model | default "ai/smollm2" }}
+      context_size: {{ .Params.context_size | default 2048 }}
+  outputs: |
+    name: {{ .Init.name }}
+```
+
+When used with the following Score resource:
+
+```yaml
+resources:
+  my-llm:
+    type: llm-model
+    params:
+      model: ai/mistral
+      context_size: 4096
+```
+
+This will generate the following in the Compose output:
+
+```yaml
+models:
+  my-llm:
+    model: ai/mistral
+    context_size: 4096
+```
+
+Services can then reference these models via the `models` field in the service configuration.
+
 ### The `cmd://` provisioner
 
 The command provisioner implementation can be used to execute an external binary or script to provision the resource. The provision IO structures are serialised to json and send on standard-input to the new process, any stdout content is decoded as json and is used as the outputs of the provisioner.

internal/command/generate.go

@@ -237,7 +237,8 @@ arguments.
 			Networks: map[string]types.NetworkConfig{},
 		}
 
-		currentState, err = provisioners.ProvisionResources(context.Background(), currentState, loadedProvisioners, superProject)
+		var workloadModels provisioners.WorkloadModels
+		currentState, workloadModels, err = provisioners.ProvisionResources(context.Background(), currentState, loadedProvisioners, superProject)
 		if err != nil {
 			return fmt.Errorf("failed to provision: %w", err)
 		} else if len(currentState.Resources) > 0 {
@@ -264,6 +265,14 @@ arguments.
 					}
 					service.DependsOn[waitServiceName] = types.ServiceDependency{Condition: "service_completed_successfully", Required: true}
 				}
+				if modelNames, ok := workloadModels[workloadName]; ok && len(modelNames) > 0 {
+					if service.Models == nil {
+						service.Models = make(map[string]*types.ServiceModelConfig)
+					}
+					for _, modelName := range modelNames {
+						service.Models[modelName] = &types.ServiceModelConfig{}
+					}
+				}
 				superProject.Services[serviceName] = service
 			}
 			for volumeName, volume := range converted.Volumes {

internal/command/run.go

@@ -207,7 +207,7 @@ func run(cmd *cobra.Command, args []string) error {
 		return err
 	}
 
-	state, err = provisioners.ProvisionResources(context.Background(), state, provisionerList, nil)
+	state, _, err = provisioners.ProvisionResources(context.Background(), state, provisionerList, nil)
 	if err != nil {
 		return fmt.Errorf("failed to provision resources: %w", err)
 	}

internal/provisioners/core.go

@@ -92,6 +92,7 @@ type ProvisionOutput struct {
 	ComposeNetworks      map[string]compose.NetworkConfig `json:"compose_networks"`
 	ComposeVolumes       map[string]compose.VolumeConfig  `json:"compose_volumes"`
 	ComposeServices      map[string]compose.ServiceConfig `json:"compose_services"`
+	ComposeModels        map[string]compose.ModelConfig   `json:"compose_models"`
 
 	// For testing and legacy reasons, built in provisioners can set a direct lookup function
 	OutputLookupFunc framework.OutputLookupFunc `json:"-"`
@@ -166,6 +167,7 @@ func (po *ProvisionOutput) ApplyToStateAndProject(state *project.State, resUid f
 		"#volumes", len(po.ComposeVolumes),
 		"#networks", len(po.ComposeNetworks),
 		"#services", len(po.ComposeServices),
+		"#models", len(po.ComposeModels),
 	)
 
 	out := *state
@@ -260,6 +262,12 @@ func (po *ProvisionOutput) ApplyToStateAndProject(state *project.State, resUid f
 		}
 		project.Services[serviceName] = service
 	}
+	for modelName, model := range po.ComposeModels {
+		if project.Models == nil {
+			project.Models = make(compose.Models)
+		}
+		project.Models[modelName] = model
+	}
 
 	out.Resources[resUid] = existing
 	return &out, nil
@@ -297,13 +305,18 @@ func buildWorkloadServices(state *project.State) map[string]NetworkService {
 	return out
 }
 
-func ProvisionResources(ctx context.Context, state *project.State, provisioners []Provisioner, composeProject *compose.Project) (*project.State, error) {
+// WorkloadModels maps workload names to their associated model names.
+// This is used to inject model references into workload services.
+type WorkloadModels map[string][]string
+
+func ProvisionResources(ctx context.Context, state *project.State, provisioners []Provisioner, composeProject *compose.Project) (*project.State, WorkloadModels, error) {
 	out := state
+	workloadModels := make(WorkloadModels)
 
 	// provision in sorted order
 	orderedResources, err := out.GetSortedResourceUids()
 	if err != nil {
-		return nil, fmt.Errorf("failed to determine sort order for provisioning: %w", err)
+		return nil, nil, fmt.Errorf("failed to determine sort order for provisioning: %w", err)
 	}
 
 	workloadServices := buildWorkloadServices(state)
@@ -314,24 +327,24 @@ func ProvisionResources(ctx context.Context, state *project.State, provisioners
 			return provisioner.Match(resUid)
 		})
 		if provisionerIndex < 0 {
-			return nil, fmt.Errorf("resource '%s' is not supported by any provisioner", resUid)
+			return nil, nil, fmt.Errorf("resource '%s' is not supported by any provisioner", resUid)
 		}
 		provisioner := provisioners[provisionerIndex]
 		if resState.ProvisionerUri != "" && resState.ProvisionerUri != provisioner.Uri() {
-			return nil, fmt.Errorf("resource '%s' was previously provisioned by a different provider - undefined behavior", resUid)
+			return nil, nil, fmt.Errorf("resource '%s' was previously provisioned by a different provider - undefined behavior", resUid)
 		}
 
 		var params map[string]interface{}
 		if resState.Params != nil && len(resState.Params) > 0 {
 			resOutputs, err := out.GetResourceOutputForWorkload(resState.SourceWorkload)
 			if err != nil {
-				return nil, fmt.Errorf("failed to find resource params for resource '%s': %w", resUid, err)
+				return nil, nil, fmt.Errorf("failed to find resource params for resource '%s': %w", resUid, err)
 			}
 			sf := framework.BuildSubstitutionFunction(out.Workloads[resState.SourceWorkload].Spec.Metadata, resOutputs)
 			sf = util.WrapImmediateSubstitutionFunction(sf)
 			rawParams, err := framework.Substitute(resState.Params, sf)
 			if err != nil {
-				return nil, fmt.Errorf("failed to substitute params for resource '%s': %w", resUid, err)
+				return nil, nil, fmt.Errorf("failed to substitute params for resource '%s': %w", resUid, err)
 			}
 			params = rawParams.(map[string]interface{})
 		}
@@ -351,15 +364,22 @@ func ProvisionResources(ctx context.Context, state *project.State, provisioners
 			MountDirectoryPath: out.Extras.MountsDirectory,
 		})
 		if err != nil {
-			return nil, fmt.Errorf("resource '%s': failed to provision: %w", resUid, err)
+			return nil, nil, fmt.Errorf("resource '%s': failed to provision: %w", resUid, err)
+		}
+
+		// Track models for the source workload
+		if len(output.ComposeModels) > 0 {
+			for modelName := range output.ComposeModels {
+				workloadModels[resState.SourceWorkload] = append(workloadModels[resState.SourceWorkload], modelName)
+			}
 		}
 
 		output.ProvisionerUri = provisioner.Uri()
 		out, err = output.ApplyToStateAndProject(out, resUid, composeProject)
 		if err != nil {
-			return nil, fmt.Errorf("resource '%s': failed to apply outputs: %w", resUid, err)
+			return nil, nil, fmt.Errorf("resource '%s': failed to apply outputs: %w", resUid, err)
 		}
 	}
 
-	return out, nil
+	return out, workloadModels, nil
 }

internal/provisioners/core_test.go

@@ -80,6 +80,9 @@ func TestApplyToStateAndProject(t *testing.T) {
 			ComposeVolumes: map[string]compose.VolumeConfig{
 				"some-volume": {Name: "volume"},
 			},
+			ComposeModels: map[string]compose.ModelConfig{
+				"some-model": {Name: "model-name", Model: "provider/model-id"},
+			},
 		}
 		afterState, err := output.ApplyToStateAndProject(startState, resUid, composeProject)
 		require.NoError(t, err)
@@ -91,6 +94,9 @@ func TestApplyToStateAndProject(t *testing.T) {
 		assert.Len(t, composeProject.Networks, 1)
 		assert.Len(t, composeProject.Volumes, 1)
 		assert.Len(t, composeProject.Services, 1)
+		assert.Len(t, composeProject.Models, 1)
+		assert.Equal(t, "model-name", composeProject.Models["some-model"].Name)
+		assert.Equal(t, "provider/model-id", composeProject.Models["some-model"].Model)
 		paths := make([]string, 0)
 		_ = filepath.WalkDir(td, func(path string, d fs.DirEntry, err error) error {
 			if d.IsDir() {
@@ -151,7 +157,7 @@ func TestProvisionResourcesWithNetworkService(t *testing.T) {
 			return &ProvisionOutput{}, nil
 		}),
 	}
-	after, err := ProvisionResources(context.Background(), state, p, nil)
+	after, _, err := ProvisionResources(context.Background(), state, p, nil)
 	if assert.NoError(t, err) {
 		assert.Len(t, after.Resources, 1)
 	}
@@ -176,7 +182,7 @@ func TestProvisionResourcesWithResourceParams(t *testing.T) {
 			return &ProvisionOutput{ResourceOutputs: map[string]interface{}{"key": "value"}}, nil
 		}),
 	}
-	after, err := ProvisionResources(context.Background(), state, p, nil)
+	after, _, err := ProvisionResources(context.Background(), state, p, nil)
 	if assert.NoError(t, err) {
 		assert.Len(t, after.Resources, 2)
 	}
@@ -200,6 +206,39 @@ func TestProvisionResourcesWithResourceParams_fail(t *testing.T) {
 			return &ProvisionOutput{ResourceOutputs: map[string]interface{}{}}, nil
 		}),
 	}
-	_, err := ProvisionResources(context.Background(), state, p, nil)
+	_, _, err := ProvisionResources(context.Background(), state, p, nil)
 	assert.EqualError(t, err, "failed to substitute params for resource 'a.default#w1.a': x: invalid ref 'resources.b.unknown': key 'unknown' not found")
 }
+
+func TestProvisionResourcesWithModels(t *testing.T) {
+	state := new(project.State)
+	state, _ = state.WithWorkload(&score.Workload{
+		Metadata: map[string]interface{}{"name": "w1"},
+		Containers: map[string]score.Container{
+			"main": {},
+		},
+		Resources: map[string]score.Resource{
+			"my-model": {Type: "llm-model"},
+		},
+	}, nil, project.WorkloadExtras{})
+	state, _ = state.WithPrimedResources()
+	composeProject := &compose.Project{}
+	p := []Provisioner{
+		NewEphemeralProvisioner("ephemeral://blah", "llm-model.default#w1.my-model", func(ctx context.Context, input *Input) (*ProvisionOutput, error) {
+			return &ProvisionOutput{
+				ComposeModels: map[string]compose.ModelConfig{
+					"w1-my-model": {Name: "w1-my-model", Model: "ai/gemma3:270M-UD-IQ2_XXS"},
+				},
+			}, nil
+		}),
+	}
+	after, workloadModels, err := ProvisionResources(context.Background(), state, p, composeProject)
+	if assert.NoError(t, err) {
+		assert.Len(t, after.Resources, 1)
+		assert.Len(t, composeProject.Models, 1)
+		assert.Equal(t, "ai/gemma3:270M-UD-IQ2_XXS", composeProject.Models["w1-my-model"].Model)
+		assert.Len(t, workloadModels, 1)
+		assert.Contains(t, workloadModels, "w1")
+		assert.Equal(t, []string{"w1-my-model"}, workloadModels["w1"])
+	}
+}

internal/provisioners/templateprov/template.go

@@ -70,6 +70,9 @@ type Provisioner struct {
 	// ComposeServicesTemplate generates a set of services to add to the compose project. These will replace any with
 	// the same name already.
 	ComposeServicesTemplate string `yaml:"services,omitempty"`
+	// ComposeModelsTemplate generates a set of models to add to the compose project. These will replace any with
+	// the same name already.
+	ComposeModelsTemplate string `yaml:"models,omitempty"`
 
 	// InfoLogsTemplate allows the provisioner to return informational messages for the user which may help connecting or
 	// testing the provisioned resource
@@ -268,6 +271,11 @@ func (p *Provisioner) Provision(ctx context.Context, input *provisioners.Input)
 		return nil, fmt.Errorf("volumes template failed: %w", err)
 	}
 
+	out.ComposeModels = make(map[string]compose.ModelConfig)
+	if err := renderTemplateAndDecode(p.ComposeModelsTemplate, &data, &out.ComposeModels, true); err != nil {
+		return nil, fmt.Errorf("models template failed: %w", err)
+	}
+
 	var infoLogs []string
 	if err := renderTemplateAndDecode(p.InfoLogsTemplate, &data, &infoLogs, false); err != nil {
 		return nil, fmt.Errorf("info logs template failed: %w", err)