forked from flyteorg/flyteplugins
/
job_definition.go
89 lines (68 loc) · 2.92 KB
/
job_definition.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
package awsbatch
import (
"context"
"regexp"
"github.com/lyft/flyteidl/gen/pb-go/flyteidl/core"
pluginErrors "github.com/lyft/flyteplugins/go/tasks/errors"
"github.com/lyft/flyteplugins/go/tasks/plugins/array/awsbatch/config"
arrayCore "github.com/lyft/flyteplugins/go/tasks/plugins/array/core"
"github.com/lyft/flytestdlib/errors"
"github.com/lyft/flytestdlib/logger"
pluginCore "github.com/lyft/flyteplugins/go/tasks/pluginmachinery/core"
"github.com/lyft/flyteplugins/go/tasks/plugins/array/awsbatch/definition"
)
func getContainerImage(_ context.Context, task *core.TaskTemplate) string {
if task.GetContainer() != nil && len(task.GetContainer().Image) > 0 {
return task.GetContainer().Image
}
return ""
}
func getRole(_ context.Context, roleAnnotationKey string, annotations map[string]string) string {
if len(roleAnnotationKey) > 0 {
return annotations[roleAnnotationKey]
}
return ""
}
var urlRegex = regexp.MustCompile(`^(?:([^/]+)/)?(?:([^/]+)/)*?([^@:/]+)(?:[@:][^/]+)?$`)
// Gets the repository part of the container image url
func containerImageRepository(containerImage string) string {
parts := urlRegex.FindAllStringSubmatch(containerImage, -1)
if len(parts) > 0 && len(parts[0]) > 3 {
return parts[0][3]
}
return ""
}
func EnsureJobDefinition(ctx context.Context, tCtx pluginCore.TaskExecutionContext, cfg *config.Config, client Client,
definitionCache definition.Cache, currentState *State) (nextState *State, err error) {
taskTemplate, err := tCtx.TaskReader().Read(ctx)
if err != nil {
return nil, err
}
containerImage := getContainerImage(ctx, taskTemplate)
if len(containerImage) == 0 {
logger.Infof(ctx, "Future task doesn't have an image specified. Failing.")
return nil, errors.Errorf(pluginErrors.BadTaskSpecification, "Tasktemplate does not contain a container image.")
}
role := getRole(ctx, cfg.RoleAnnotationKey, tCtx.TaskExecutionMetadata().GetAnnotations())
cacheKey := definition.NewCacheKey(role, containerImage)
if existingArn, found := definitionCache.Get(cacheKey); found {
logger.Infof(ctx, "Found an existing job definition for Image [%v] and Role [%v]. Arn [%v]",
containerImage, role, existingArn)
nextState = currentState.SetJobDefinitionArn(existingArn)
nextState.State = nextState.SetPhase(arrayCore.PhaseLaunch, 0).SetReason("AWS job definition already exist.")
return nextState, nil
}
name := definition.GetJobDefinitionSafeName(containerImageRepository(containerImage))
arn, err := client.RegisterJobDefinition(ctx, name, containerImage, role)
if err != nil {
return currentState, err
}
err = definitionCache.Put(cacheKey, arn)
if err != nil {
logger.Warnf(ctx, "Failed to store job definition arn in cache. Will continue with the registered arn [%v]. Error: %v",
arn, err)
}
nextState = currentState.SetJobDefinitionArn(arn)
nextState.State = nextState.SetPhase(arrayCore.PhaseLaunch, 0).SetReason("Created AWS job definition")
return nextState, nil
}