Skip to content

Commit 70909f2

Browse files
fsaminfsamin
authored
feat(hatchery): inject env variables into workers from hatchery config (#5806)
1 parent 521a739 commit 70909f2

File tree

11 files changed

+46
-18
lines changed

11 files changed

+46
-18
lines changed

‎engine/hatchery/hatchery_local_test.go

+4-1
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@ func TestHatcheryLocal(t *testing.T) {
4949
cfg.API.MaxHeartbeatFailures = 0
5050
cfg.Provision.RegisterFrequency = 1
5151
cfg.Provision.MaxWorker = 1
52+
cfg.Provision.InjectEnvVars = []string{"AAA=AAA"}
5253
privKey, _ := jws.NewRandomRSAKey()
5354
privKeyPEM, _ := jws.ExportPrivateKey(privKey)
5455
cfg.RSAPrivateKey = string(privKeyPEM)
@@ -105,9 +106,11 @@ func TestHatcheryLocal(t *testing.T) {
105106
}
106107
}
107108

108-
func TestHelperProcess(*testing.T) {
109+
func TestHelperProcess(t *testing.T) {
109110
if os.Getenv("GO_WANT_HELPER_PROCESS") != "1" {
110111
return
111112
}
113+
114+
t.Log(os.Environ())
112115
time.Sleep(30 * time.Second)
113116
}

‎engine/hatchery/kubernetes/kubernetes.go

+2-2
Original file line numberDiff line numberDiff line change
@@ -193,7 +193,7 @@ func (h *HatcheryKubernetes) SpawnWorker(ctx context.Context, spawnArgs hatchery
193193
}
194194
}
195195

196-
udataParam := h.GenerateWorkerArgs(h, spawnArgs)
196+
udataParam := h.GenerateWorkerArgs(ctx, h, spawnArgs)
197197
udataParam.TTL = h.Config.WorkerTTL
198198

199199
udataParam.WorkflowJobID = spawnArgs.JobID
@@ -216,7 +216,7 @@ func (h *HatcheryKubernetes) SpawnWorker(ctx context.Context, spawnArgs hatchery
216216
if spawnArgs.Model.ModelDocker.Envs == nil {
217217
spawnArgs.Model.ModelDocker.Envs = map[string]string{}
218218
}
219-
envsWm := map[string]string{}
219+
envsWm := udataParam.InjectEnvVars
220220
envsWm["CDS_MODEL_MEMORY"] = fmt.Sprintf("%d", memory)
221221
envsWm["CDS_API"] = udataParam.API
222222
envsWm["CDS_TOKEN"] = udataParam.Token

‎engine/hatchery/kubernetes/kubernetes_test.go

+3
Original file line numberDiff line numberDiff line change
@@ -65,6 +65,7 @@ func TestHatcheryKubernetes_Status(t *testing.T) {
6565
defer gock.Off()
6666
defer gock.Observe(nil)
6767
h := NewHatcheryKubernetesTest(t)
68+
h.Config.HatcheryCommonConfiguration.Provision.InjectEnvVars = []string{"ZZZZ=ZZZZ"}
6869

6970
m := &sdk.Model{
7071
Name: "model1",
@@ -96,6 +97,8 @@ func TestHatcheryKubernetes_Status(t *testing.T) {
9697
require.Equal(t, 2, len(podRequest.Spec.Containers))
9798
require.Equal(t, "k8s-toto", podRequest.Spec.Containers[0].Name)
9899
require.Equal(t, int64(4096), podRequest.Spec.Containers[0].Resources.Requests.Memory().Value())
100+
require.Equal(t, "ZZZZ", podRequest.Spec.Containers[0].Env[1].Name)
101+
require.Equal(t, "ZZZZ", podRequest.Spec.Containers[0].Env[1].Value)
99102
require.Equal(t, "service-0-pg", podRequest.Spec.Containers[1].Name)
100103
require.Equal(t, 1, len(podRequest.Spec.Containers[1].Env))
101104
require.Equal(t, "PG_USERNAME", podRequest.Spec.Containers[1].Env[0].Name)

‎engine/hatchery/local/worker_spawn.go

+4-1
Original file line numberDiff line numberDiff line change
@@ -71,7 +71,7 @@ func (h *HatcheryLocal) SpawnWorker(ctx context.Context, spawnArgs hatchery.Spaw
7171

7272
log.Info(ctx, "HatcheryLocal.SpawnWorker> basedir: %s", basedir)
7373

74-
udataParam := h.GenerateWorkerArgs(h, spawnArgs)
74+
udataParam := h.GenerateWorkerArgs(ctx, h, spawnArgs)
7575
udataParam.BaseDir = basedir
7676
udataParam.WorkerBinary = path.Join(h.BasedirDedicated, h.getWorkerBinaryName())
7777
udataParam.WorkflowJobID = spawnArgs.JobID
@@ -109,6 +109,9 @@ func (h *HatcheryLocal) SpawnWorker(ctx context.Context, spawnArgs hatchery.Spaw
109109
cmd.Env = append(cmd.Env, e)
110110
}
111111
}
112+
for k, v := range udataParam.InjectEnvVars {
113+
cmd.Env = append(cmd.Env, k+"="+v)
114+
}
112115

113116
// Wait in a goroutine so that when process exits, Wait() update cmd.ProcessState
114117
go func() {

‎engine/hatchery/marathon/marathon.go

+2-2
Original file line numberDiff line numberDiff line change
@@ -237,7 +237,7 @@ func (h *HatcheryMarathon) SpawnWorker(ctx context.Context, spawnArgs hatchery.S
237237
instance := 1
238238
forcePull := strings.HasSuffix(spawnArgs.Model.ModelDocker.Image, ":latest")
239239

240-
udataParam := h.GenerateWorkerArgs(h, spawnArgs)
240+
udataParam := h.GenerateWorkerArgs(ctx, h, spawnArgs)
241241
udataParam.TTL = h.Config.WorkerTTL
242242
udataParam.WorkflowJobID = spawnArgs.JobID
243243

@@ -279,7 +279,7 @@ func (h *HatcheryMarathon) SpawnWorker(ctx context.Context, spawnArgs hatchery.S
279279
spawnArgs.Model.ModelDocker.Envs = map[string]string{}
280280
}
281281

282-
envsWm := map[string]string{}
282+
envsWm := udataParam.InjectEnvVars
283283
envsWm["CDS_MODEL_MEMORY"] = fmt.Sprintf("%d", memory)
284284
envsWm["CDS_API"] = udataParam.API
285285
envsWm["CDS_TOKEN"] = udataParam.Token

‎engine/hatchery/openstack/spawn.go

+1-1
Original file line numberDiff line numberDiff line change
@@ -77,7 +77,7 @@ func (h *HatcheryOpenstack) SpawnWorker(ctx context.Context, spawnArgs hatchery.
7777
return err
7878
}
7979

80-
udataParam := h.GenerateWorkerArgs(h, spawnArgs)
80+
udataParam := h.GenerateWorkerArgs(ctx, h, spawnArgs)
8181
udataParam.TTL = h.Config.WorkerTTL
8282
udataParam.FromWorkerImage = withExistingImage
8383
udataParam.WorkflowJobID = spawnArgs.JobID

‎engine/hatchery/serve.go

+13-1
Original file line numberDiff line numberDiff line change
@@ -250,21 +250,33 @@ func getStatusHandler(h hatchery.Interface) service.HandlerFunc {
250250
}
251251
}
252252

253-
func (c *Common) GenerateWorkerArgs(h hatchery.Interface, spawnArgs hatchery.SpawnArguments) sdk.WorkerArgs {
253+
func (c *Common) GenerateWorkerArgs(ctx context.Context, h hatchery.Interface, spawnArgs hatchery.SpawnArguments) sdk.WorkerArgs {
254254
apiURL := h.Configuration().Provision.WorkerAPIHTTP.URL
255255
httpInsecure := h.Configuration().Provision.WorkerAPIHTTP.Insecure
256256
if apiURL == "" {
257257
apiURL = h.Configuration().API.HTTP.URL
258258
httpInsecure = h.Configuration().API.HTTP.Insecure
259259
}
260260

261+
envvars := make(map[string]string, len(h.Configuration().Provision.InjectEnvVars))
262+
263+
for _, e := range h.Configuration().Provision.InjectEnvVars {
264+
tuple := strings.SplitN(e, "=", 2)
265+
if len(tuple) != 2 {
266+
log.Error(ctx, "invalid env variable to inject: %q", e)
267+
continue
268+
}
269+
envvars[tuple[0]] = tuple[1]
270+
}
271+
261272
return sdk.WorkerArgs{
262273
API: apiURL,
263274
HTTPInsecure: httpInsecure,
264275
Token: spawnArgs.WorkerToken,
265276
Name: spawnArgs.WorkerName,
266277
Model: spawnArgs.ModelName(),
267278
HatcheryName: h.Name(),
279+
InjectEnvVars: envvars,
268280
GraylogHost: h.Configuration().Provision.WorkerLogsOptions.Graylog.Host,
269281
GraylogPort: h.Configuration().Provision.WorkerLogsOptions.Graylog.Port,
270282
GraylogExtraKey: h.Configuration().Provision.WorkerLogsOptions.Graylog.ExtraKey,

‎engine/hatchery/swarm/swarm.go

+2-2
Original file line numberDiff line numberDiff line change
@@ -374,7 +374,7 @@ func (h *HatcherySwarm) SpawnWorker(ctx context.Context, spawnArgs hatchery.Spaw
374374
return errDockerOpts
375375
}
376376

377-
udataParam := h.GenerateWorkerArgs(h, spawnArgs)
377+
udataParam := h.GenerateWorkerArgs(ctx, h, spawnArgs)
378378
udataParam.TTL = h.Config.WorkerTTL
379379
udataParam.WorkflowJobID = spawnArgs.JobID
380380

@@ -395,7 +395,7 @@ func (h *HatcherySwarm) SpawnWorker(ctx context.Context, spawnArgs hatchery.Spaw
395395
modelEnvs[k] = v
396396
}
397397

398-
envsWm := map[string]string{}
398+
envsWm := udataParam.InjectEnvVars
399399
envsWm["CDS_MODEL_MEMORY"] = fmt.Sprintf("%d", memory)
400400
envsWm["CDS_API"] = udataParam.API
401401
envsWm["CDS_TOKEN"] = udataParam.Token

‎engine/hatchery/vsphere/spawn.go

+5-1
Original file line numberDiff line numberDiff line change
@@ -323,7 +323,7 @@ func (h *HatcheryVSphere) launchScriptWorker(ctx context.Context, name string, j
323323
return errt
324324
}
325325

326-
udataParam := h.GenerateWorkerArgs(h, hatchery.SpawnArguments{
326+
udataParam := h.GenerateWorkerArgs(ctx, h, hatchery.SpawnArguments{
327327
WorkerToken: token,
328328
WorkerName: name,
329329
Model: &model,
@@ -332,6 +332,10 @@ func (h *HatcheryVSphere) launchScriptWorker(ctx context.Context, name string, j
332332
udataParam.FromWorkerImage = true
333333
udataParam.WorkflowJobID = jobID
334334

335+
for k, v := range udataParam.InjectEnvVars {
336+
env = append(env, k+"="+v)
337+
}
338+
335339
var buffer bytes.Buffer
336340
if err := tmpl.Execute(&buffer, udataParam); err != nil {
337341
return err

‎engine/service/types.go

+8-7
Original file line numberDiff line numberDiff line change
@@ -46,13 +46,14 @@ type HatcheryCommonConfiguration struct {
4646
MaxHeartbeatFailures int `toml:"maxHeartbeatFailures" default:"10" comment:"Maximum allowed consecutives failures on heatbeat routine" json:"maxHeartbeatFailures"`
4747
} `toml:"api" json:"api"`
4848
Provision struct {
49-
RatioService *int `toml:"ratioService" default:"50" commented:"true" comment:"Percent reserved for spawning worker with service requirement" json:"ratioService,omitempty" mapstructure:"ratioService"`
50-
MaxWorker int `toml:"maxWorker" default:"10" comment:"Maximum allowed simultaneous workers" json:"maxWorker"`
51-
MaxConcurrentProvisioning int `toml:"maxConcurrentProvisioning" default:"10" comment:"Maximum allowed simultaneous workers provisioning" json:"maxConcurrentProvisioning"`
52-
MaxConcurrentRegistering int `toml:"maxConcurrentRegistering" default:"2" comment:"Maximum allowed simultaneous workers registering. -1 to disable registering on this hatchery" json:"maxConcurrentRegistering"`
53-
RegisterFrequency int `toml:"registerFrequency" default:"60" comment:"Check if some worker model have to be registered each n Seconds" json:"registerFrequency"`
54-
Region string `toml:"region" default:"" comment:"region of this hatchery - optional. With a free text as 'myregion', user can set a prerequisite 'region' with value 'myregion' on CDS Job" json:"region"`
55-
IgnoreJobWithNoRegion bool `toml:"ignoreJobWithNoRegion" default:"false" comment:"Ignore job without a region prerequisite if ignoreJobWithNoRegion=true"`
49+
InjectEnvVars []string `toml:"injectEnvVars" commented:"true" comment:"Inject env variables in workers" json:"-" mapstructure:"injectEnvVars"`
50+
RatioService *int `toml:"ratioService" default:"50" commented:"true" comment:"Percent reserved for spawning worker with service requirement" json:"ratioService,omitempty" mapstructure:"ratioService"`
51+
MaxWorker int `toml:"maxWorker" default:"10" comment:"Maximum allowed simultaneous workers" json:"maxWorker"`
52+
MaxConcurrentProvisioning int `toml:"maxConcurrentProvisioning" default:"10" comment:"Maximum allowed simultaneous workers provisioning" json:"maxConcurrentProvisioning"`
53+
MaxConcurrentRegistering int `toml:"maxConcurrentRegistering" default:"2" comment:"Maximum allowed simultaneous workers registering. -1 to disable registering on this hatchery" json:"maxConcurrentRegistering"`
54+
RegisterFrequency int `toml:"registerFrequency" default:"60" comment:"Check if some worker model have to be registered each n Seconds" json:"registerFrequency"`
55+
Region string `toml:"region" default:"" comment:"region of this hatchery - optional. With a free text as 'myregion', user can set a prerequisite 'region' with value 'myregion' on CDS Job" json:"region"`
56+
IgnoreJobWithNoRegion bool `toml:"ignoreJobWithNoRegion" default:"false" comment:"Ignore job without a region prerequisite if ignoreJobWithNoRegion=true"`
5657
WorkerAPIHTTP struct {
5758
URL string `toml:"url" default:"http://localhost:8081" commented:"true" comment:"CDS API URL for worker, let empty or commented to use the same URL that is used by the Hatchery" json:"url"`
5859
Insecure bool `toml:"insecure" default:"false" commented:"true" comment:"sslInsecureSkipVerify, set to true if you use a self-signed SSL on CDS API" json:"insecure"`

‎sdk/worker.go

+2
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,8 @@ type WorkerArgs struct {
5656
GraylogExtraKey string `json:"graylog_extra_key"`
5757
GraylogExtraValue string `json:"graylog_extra_value"`
5858
WorkerBinary string
59+
// Env variables
60+
InjectEnvVars map[string]string `json:"inject_env_vars"`
5961
}
6062

6163
// TemplateEnvs return envs interpolated with worker arguments

0 commit comments

Comments
 (0)