Skip to content

Commit 317ebcd

Browse files
ndeloofglours
authored andcommitted
implement model_variable
Signed-off-by: Nicolas De Loof <[email protected]>
1 parent 5cf1f0e commit 317ebcd

File tree

1 file changed

+13
-4
lines changed

1 file changed

+13
-4
lines changed

pkg/compose/model.go

Lines changed: 13 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -75,7 +75,7 @@ func (s *composeService) ensureModels(ctx context.Context, project *types.Projec
7575

7676
eg, gctx := errgroup.WithContext(ctx)
7777
eg.Go(func() error {
78-
return s.setModelEndpointVariable(gctx, dockerModel, project)
78+
return s.setModelVariables(gctx, dockerModel, project)
7979
})
8080

8181
for name, config := range project.Models {
@@ -161,7 +161,7 @@ func (s *composeService) configureModel(ctx context.Context, dockerModel *manage
161161
return cmd.Run()
162162
}
163163

164-
func (s *composeService) setModelEndpointVariable(ctx context.Context, dockerModel *manager.Plugin, project *types.Project) error {
164+
func (s *composeService) setModelVariables(ctx context.Context, dockerModel *manager.Plugin, project *types.Project) error {
165165
cmd := exec.CommandContext(ctx, dockerModel.Path, "status", "--json")
166166
s.setupChildProcess(ctx, cmd)
167167
statusOut, err := cmd.CombinedOutput()
@@ -179,12 +179,21 @@ func (s *composeService) setModelEndpointVariable(ctx context.Context, dockerMod
179179
}
180180

181181
for _, service := range project.Services {
182-
for model, modelConfig := range service.Models {
182+
for ref, modelConfig := range service.Models {
183+
model := project.Models[ref]
184+
varPrefix := strings.ReplaceAll(strings.ToUpper(ref), "-", "_")
183185
var variable string
186+
if modelConfig != nil && modelConfig.ModelVariable != "" {
187+
variable = modelConfig.ModelVariable
188+
} else {
189+
variable = varPrefix
190+
}
191+
service.Environment[variable] = &model.Model
192+
184193
if modelConfig != nil && modelConfig.EndpointVariable != "" {
185194
variable = modelConfig.EndpointVariable
186195
} else {
187-
variable = strings.ToUpper(model) + "_URL"
196+
variable = varPrefix + "_URL"
188197
}
189198
service.Environment[variable] = &status.Endpoint
190199
}

0 commit comments

Comments
 (0)