forked from kubeflow/arena
-
Notifications
You must be signed in to change notification settings - Fork 0
/
scaleout_etjob_builder.go
92 lines (81 loc) · 2.02 KB
/
scaleout_etjob_builder.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
90
91
92
package training
import (
"fmt"
"time"
"github.com/kubeflow/arena/pkg/apis/types"
"github.com/kubeflow/arena/pkg/argsbuilder"
)
type ScaleOutETJobBuilder struct {
args *types.ScaleOutETJobArgs
argValues map[string]interface{}
argsbuilder.ArgsBuilder
}
func NewScaleOutETJobBuilder() *ScaleOutETJobBuilder {
args := &types.ScaleOutETJobArgs{
ScaleETJobArgs: types.ScaleETJobArgs{
Timeout: 60,
Retry: 0,
Count: 1,
},
}
return &ScaleOutETJobBuilder{
args: args,
argValues: map[string]interface{}{},
ArgsBuilder: argsbuilder.NewScaleOutETJobArgsBuilder(args),
}
}
// Name is used to set job name,match option --name
func (b *ScaleOutETJobBuilder) Name(name string) *ScaleOutETJobBuilder {
if name != "" {
b.args.Name = name
}
return b
}
// Retry is used to set retry times
func (b *ScaleOutETJobBuilder) Retry(count int) *ScaleOutETJobBuilder {
if count > 0 {
b.args.Retry = count
}
return b
}
// Retry is used to set retry times
func (b *ScaleOutETJobBuilder) Count(count int) *ScaleOutETJobBuilder {
if count > 0 {
b.args.Count = count
}
return b
}
// Timeout is used to set timeout seconds
func (b *ScaleOutETJobBuilder) Timeout(timeout time.Duration) *ScaleOutETJobBuilder {
b.argValues["timeout"] = &timeout
return b
}
// Script is used to set scale script
func (b *ScaleOutETJobBuilder) Script(s string) *ScaleOutETJobBuilder {
if s != "" {
b.args.Script = s
}
return b
}
// Envs is used to set envs
func (b *ScaleOutETJobBuilder) Envs(envs map[string]string) *ScaleOutETJobBuilder {
items := []string{}
for key, value := range envs {
items = append(items, fmt.Sprintf("%v=%v", key, value))
}
b.argValues["env"] = items
return b
}
// Build is used to build the job
func (b *ScaleOutETJobBuilder) Build() (*Job, error) {
for key, value := range b.argValues {
b.AddArgValue(key, value)
}
if err := b.PreBuild(); err != nil {
return nil, err
}
if err := b.ArgsBuilder.Build(); err != nil {
return nil, err
}
return NewJob(b.args.Name, types.ETTrainingJob, b.args), nil
}