-
Notifications
You must be signed in to change notification settings - Fork 349
/
experiment_config.go
232 lines (196 loc) · 7.01 KB
/
experiment_config.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
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
package model
import (
"encoding/json"
"path/filepath"
"strings"
"github.com/pkg/errors"
"github.com/docker/go-units"
"github.com/ghodss/yaml"
"github.com/determined-ai/determined/master/pkg/check"
)
const (
// MinUserSchedulingPriority is the smallest priority users may specify.
MinUserSchedulingPriority = 1
// MaxUserSchedulingPriority is the largest priority users may specify.
MaxUserSchedulingPriority = 99
defaultDeviceMode = "mrw"
)
// DevicesConfig is the configuration for devices. It is a named type because it needs custom
// merging behavior (via UnmarshalJSON).
type DevicesConfig []DeviceConfig
// UnmarshalJSON implements the json.Unmarshaler interface so that DeviceConfigs are additive.
func (d *DevicesConfig) UnmarshalJSON(data []byte) error {
unmarshaled := make([]DeviceConfig, 0)
if err := json.Unmarshal(data, &unmarshaled); err != nil {
return errors.Wrap(err, "failed to parse devices")
}
// Prevent duplicate container paths as a result of the merge. Prefer the unmarshaled devices
// to the old ones since with this unmarshaling strategy we always unmarshal in order of
// increasing priority.
paths := map[string]bool{}
for _, device := range unmarshaled {
paths[device.ContainerPath] = true
}
for _, device := range *d {
if _, ok := paths[device.ContainerPath]; !ok {
unmarshaled = append(unmarshaled, device)
}
}
*d = unmarshaled
return nil
}
// DeviceConfig configures container device access.
type DeviceConfig struct {
HostPath string `json:"host_path"`
ContainerPath string `json:"container_path"`
Mode string `json:"mode"`
}
// UnmarshalJSON implements the json.Unmarshaler interface.
func (d *DeviceConfig) UnmarshalJSON(data []byte) error {
var plain string
if err := json.Unmarshal(data, &plain); err == nil {
fields := strings.Split(plain, ":")
if len(fields) < 2 || len(fields) > 3 {
return errors.Errorf("invalid device string: %q", plain)
}
d.HostPath = fields[0]
d.ContainerPath = fields[1]
if len(fields) > 2 {
d.Mode = fields[2]
} else {
d.Mode = defaultDeviceMode
}
return nil
}
d.Mode = defaultDeviceMode
type DefaultParser *DeviceConfig
return errors.Wrap(json.Unmarshal(data, DefaultParser(d)), "failed to parse device")
}
// ResourcesConfig configures resource usage for an experiment, command, notebook, or tensorboard.
type ResourcesConfig struct {
Slots int `json:"slots"`
MaxSlots *int `json:"max_slots,omitempty"`
Weight float64 `json:"weight"`
NativeParallel bool `json:"native_parallel,omitempty"`
ShmSize *StorageSize `json:"shm_size,omitempty"`
ResourcePool string `json:"resource_pool"`
Priority *int `json:"priority,omitempty"`
Devices DevicesConfig `json:"devices"`
// Deprecated: Use ResourcePool instead.
AgentLabel string `json:"agent_label,omitempty"`
}
// StorageSize is a named type for custom marshaling behavior for shm_size.
type StorageSize int64
// UnmarshalJSON implements the json.Unmarshaler interface.
func (d *StorageSize) UnmarshalJSON(data []byte) error {
var size any
if err := json.Unmarshal(data, &size); err != nil {
return err
}
switch s := size.(type) {
case float64:
*d = StorageSize(s)
case string:
b, err := units.RAMInBytes(s)
if err != nil {
return errors.Wrap(err, "failed to parse shm_size")
}
*d = StorageSize(b)
default:
return errors.New("shm_size needs to be a string or numeric")
}
return nil
}
// ParseJustResources is a helper function for breaking the circular dependency where we need the
// TaskContainerDefaults to unmarshal an ExperimentConfig, but we need the Resources.ResourcePool
// setting to know which TaskContainerDefaults to use. It does not throw errors; if unmarshalling
// fails that can just get caught later.
func ParseJustResources(configBytes []byte) ResourcesConfig {
// Make this function usable on experiment or command configs.
type DummyConfig struct {
Resources ResourcesConfig `json:"resources"`
}
dummy := DummyConfig{
Resources: ResourcesConfig{
Slots: 1,
},
}
// Don't throw errors; validation should happen elsewhere.
_ = yaml.Unmarshal(configBytes, &dummy)
return dummy.Resources
}
// ValidatePrioritySetting checks that priority if set is within a valid range.
func ValidatePrioritySetting(priority *int) []error {
errs := make([]error, 0)
if priority != nil {
errs = append(errs, check.GreaterThanOrEqualTo(
*priority, MinUserSchedulingPriority,
"scheduling priority must be greater than 0 and less than 100"))
errs = append(errs, check.LessThanOrEqualTo(
*priority, MaxUserSchedulingPriority,
"scheduling priority must be greater than 0 and less than 100"))
}
return errs
}
// Validate implements the check.Validatable interface.
func (r ResourcesConfig) Validate() []error {
errs := []error{
check.GreaterThanOrEqualTo(r.Slots, 0, "slots must be >= 0"),
check.GreaterThan(r.Weight, float64(0), "weight must be > 0"),
}
errs = append(errs, ValidatePrioritySetting(r.Priority)...)
return errs
}
// BindMountsConfig is the configuration for bind mounts.
type BindMountsConfig []BindMount
// UnmarshalJSON implements the json.Unmarshaler interface.
func (b *BindMountsConfig) UnmarshalJSON(data []byte) error {
unmarshaled := make([]BindMount, 0)
if err := json.Unmarshal(data, &unmarshaled); err != nil {
return errors.Wrap(err, "failed to parse bind mounts")
}
// Prevent duplicate container paths as a result of the merge. Prefer the unmarshaled bind
// mounts to the old ones since with this unmarshaling strategy we always unmarshal in order of
// increasing priority.
paths := map[string]bool{}
for _, mount := range unmarshaled {
paths[mount.ContainerPath] = true
}
for _, mount := range *b {
if _, ok := paths[mount.ContainerPath]; !ok {
unmarshaled = append(unmarshaled, mount)
}
}
*b = unmarshaled
return nil
}
// BindMount configures trial runner filesystem bind mounts.
type BindMount struct {
HostPath string `json:"host_path"`
ContainerPath string `json:"container_path"`
ReadOnly bool `json:"read_only"`
Propagation string `json:"propagation"`
}
// Validate implements the check.Validatable interface.
func (b BindMount) Validate() []error {
return []error{
check.True(b.ContainerPath != ".", "container_path must not be \".\""),
check.True(filepath.IsAbs(b.HostPath), "host_path must be an absolute path"),
}
}
// UnmarshalJSON implements the json.Unmarshaler interface.
func (b *BindMount) UnmarshalJSON(data []byte) error {
b.Propagation = "rprivate"
type DefaultParser *BindMount
return errors.Wrap(json.Unmarshal(data, DefaultParser(b)), "failed to parse bind mounts")
}
// ProxyPort is a legacy-style clone of expconf.ProxyPort.
// TODO(ilia): migrate command config to expconf.
type ProxyPort struct {
ProxyPort int `json:"proxy_port"`
ProxyTCP bool `json:"proxy_tcp"`
Unauthenticated bool `json:"unauthenticated"`
DefaultServiceID bool `json:"default_service_id"`
}
// ProxyPortsConfig is a legacy-style clone of expconf.ProxyPortsConfig.
type ProxyPortsConfig []ProxyPort