Skip to content

Commit f679118

Browse files
mbobrovskyimbobrovskyi
authored
[kjobctl] Additional environment variables. (#2597)
1 parent 288570b commit f679118

21 files changed

+791
-710
lines changed

‎cmd/experimental/kjobctl/go.mod

+1-1
Original file line numberDiff line numberDiff line change
@@ -78,7 +78,7 @@ require (
7878
gopkg.in/yaml.v2 v2.4.0 // indirect
7979
gopkg.in/yaml.v3 v3.0.1 // indirect
8080
k8s.io/apiextensions-apiserver v0.30.1 // indirect
81-
k8s.io/component-base v0.30.2 // indirect
81+
k8s.io/component-base v0.30.3 // indirect
8282
k8s.io/kube-openapi v0.0.0-20240521193020-835d969ad83a // indirect
8383
sigs.k8s.io/json v0.0.0-20221116044647-bc3834ca7abd // indirect
8484
sigs.k8s.io/kustomize/api v0.13.5-0.20230601165947-6ce0bf390ce3 // indirect

‎cmd/experimental/kjobctl/go.sum

+2-2
Original file line numberDiff line numberDiff line change
@@ -286,8 +286,8 @@ k8s.io/cli-runtime v0.30.3 h1:aG69oRzJuP2Q4o8dm+f5WJIX4ZBEwrvdID0+MXyUY6k=
286286
k8s.io/cli-runtime v0.30.3/go.mod h1:hwrrRdd9P84CXSKzhHxrOivAR9BRnkMt0OeP5mj7X30=
287287
k8s.io/client-go v0.30.3 h1:bHrJu3xQZNXIi8/MoxYtZBBWQQXwy16zqJwloXXfD3k=
288288
k8s.io/client-go v0.30.3/go.mod h1:8d4pf8vYu665/kUbsxWAQ/JDBNWqfFeZnvFiVdmx89U=
289-
k8s.io/component-base v0.30.2 h1:pqGBczYoW1sno8q9ObExUqrYSKhtE5rW3y6gX88GZII=
290-
k8s.io/component-base v0.30.2/go.mod h1:yQLkQDrkK8J6NtP+MGJOws+/PPeEXNpwFixsUI7h/OE=
289+
k8s.io/component-base v0.30.3 h1:Ci0UqKWf4oiwy8hr1+E3dsnliKnkMLZMVbWzeorlk7s=
290+
k8s.io/component-base v0.30.3/go.mod h1:C1SshT3rGPCuNtBs14RmVD2xW0EhRSeLvBh7AGk1quA=
291291
k8s.io/klog/v2 v2.120.1 h1:QXU6cPEOIslTGvZaXvFWiP9VKyeet3sawzTOvdXb4Vw=
292292
k8s.io/klog/v2 v2.120.1/go.mod h1:3Jpz1GvMt720eyJH1ckRHK1EDfpxISzJ7I9OYgaDtPE=
293293
k8s.io/kube-openapi v0.0.0-20240521193020-835d969ad83a h1:zD1uj3Jf+mD4zmA7W+goE5TxDkI7OGJjBNBzq5fJtLA=

‎cmd/experimental/kjobctl/hack/tools/go.mod

+1-2
Original file line numberDiff line numberDiff line change
@@ -96,8 +96,7 @@ require (
9696
k8s.io/gengo/v2 v2.0.0-20240404160639-a0386bf69313 // indirect
9797
k8s.io/klog/v2 v2.130.1 // indirect
9898
k8s.io/kube-openapi v0.0.0-20240709000822-3c01b740850f // indirect
99-
k8s.io/kubectl v0.30.2 // indirect
100-
k8s.io/kubernetes v1.30.2 // indirect
99+
k8s.io/kubectl v0.30.3 // indirect
101100
k8s.io/utils v0.0.0-20240711033017-18e509b52bc8 // indirect
102101
sigs.k8s.io/controller-runtime v0.18.4 // indirect
103102
sigs.k8s.io/json v0.0.0-20221116044647-bc3834ca7abd // indirect

‎cmd/experimental/kjobctl/hack/tools/go.sum

+2-4
Original file line numberDiff line numberDiff line change
@@ -326,10 +326,8 @@ k8s.io/klog/v2 v2.130.1 h1:n9Xl7H1Xvksem4KFG4PYbdQCQxqc/tTUyrgXaOhHSzk=
326326
k8s.io/klog/v2 v2.130.1/go.mod h1:3Jpz1GvMt720eyJH1ckRHK1EDfpxISzJ7I9OYgaDtPE=
327327
k8s.io/kube-openapi v0.0.0-20240709000822-3c01b740850f h1:2sXuKesAYbRHxL3aE2PN6zX/gcJr22cjrsej+W784Tc=
328328
k8s.io/kube-openapi v0.0.0-20240709000822-3c01b740850f/go.mod h1:UxDHUPsUwTOOxSU+oXURfFBcAS6JwiRXTYqYwfuGowc=
329-
k8s.io/kubectl v0.30.2 h1:cgKNIvsOiufgcs4yjvgkK0+aPCfa8pUwzXdJtkbhsH8=
330-
k8s.io/kubectl v0.30.2/go.mod h1:rz7GHXaxwnigrqob0lJsiA07Df8RE3n1TSaC2CTeuB4=
331-
k8s.io/kubernetes v1.30.2 h1:11WhS78OYX/lnSy6TXxPO6Hk+E5K9ZNrEsk9JgMSX8I=
332-
k8s.io/kubernetes v1.30.2/go.mod h1:yPbIk3MhmhGigX62FLJm+CphNtjxqCvAIFQXup6RKS0=
329+
k8s.io/kubectl v0.30.3 h1:YIBBvMdTW0xcDpmrOBzcpUVsn+zOgjMYIu7kAq+yqiI=
330+
k8s.io/kubectl v0.30.3/go.mod h1:IcR0I9RN2+zzTRUa1BzZCm4oM0NLOawE6RzlDvd1Fpo=
333331
k8s.io/utils v0.0.0-20240711033017-18e509b52bc8 h1:pUdcCO1Lk/tbT5ztQWOBi5HBgbBP1J8+AsQnQCKsi8A=
334332
k8s.io/utils v0.0.0-20240711033017-18e509b52bc8/go.mod h1:OLgZIPagt7ERELqWJFomSt595RzquPNLL48iOWgYOg0=
335333
sigs.k8s.io/controller-runtime v0.18.4 h1:87+guW1zhvuPLh1PHybKdYFLU0YJp4FhJRmiHvm5BZw=

‎cmd/experimental/kjobctl/pkg/builder/builder.go

+26-2
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,10 @@ package builder
1919
import (
2020
"context"
2121
"errors"
22+
"fmt"
23+
"os"
2224
"slices"
25+
"time"
2326

2427
corev1 "k8s.io/api/core/v1"
2528
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
@@ -29,6 +32,7 @@ import (
2932
"sigs.k8s.io/kueue/cmd/experimental/kjobctl/apis/v1alpha1"
3033
"sigs.k8s.io/kueue/cmd/experimental/kjobctl/client-go/clientset/versioned"
3134
"sigs.k8s.io/kueue/cmd/experimental/kjobctl/pkg/cmd/util"
35+
"sigs.k8s.io/kueue/cmd/experimental/kjobctl/pkg/constants"
3236
)
3337

3438
var (
@@ -66,10 +70,12 @@ type Builder struct {
6670
profile *v1alpha1.ApplicationProfile
6771
mode *v1alpha1.SupportedMode
6872
volumeBundles []v1alpha1.VolumeBundle
73+
74+
buildTime time.Time
6975
}
7076

71-
func NewBuilder(clientGetter util.ClientGetter) *Builder {
72-
return &Builder{clientGetter: clientGetter}
77+
func NewBuilder(clientGetter util.ClientGetter, buildTime time.Time) *Builder {
78+
return &Builder{clientGetter: clientGetter, buildTime: buildTime}
7379
}
7480

7581
func (b *Builder) WithNamespace(namespace string) *Builder {
@@ -236,18 +242,36 @@ func (b *Builder) buildPodSpec(templateSpec corev1.PodSpec) corev1.PodSpec {
236242

237243
container.VolumeMounts = append(container.VolumeMounts, bundle.Spec.ContainerVolumeMounts...)
238244
container.Env = append(container.Env, bundle.Spec.EnvVars...)
245+
container.Env = append(container.Env, b.additionalEnvironmentVariables()...)
239246
}
240247

241248
for i := range templateSpec.InitContainers {
242249
initContainer := &templateSpec.InitContainers[i]
243250

244251
initContainer.VolumeMounts = append(initContainer.VolumeMounts, bundle.Spec.ContainerVolumeMounts...)
245252
initContainer.Env = append(initContainer.Env, bundle.Spec.EnvVars...)
253+
initContainer.Env = append(initContainer.Env, b.additionalEnvironmentVariables()...)
246254
}
247255

248256
return templateSpec
249257
}
250258

259+
func (b *Builder) additionalEnvironmentVariables() []corev1.EnvVar {
260+
userID := os.Getenv(constants.SystemEnvVarNameUser)
261+
timestamp := b.buildTime.Format(time.RFC3339)
262+
taskName := fmt.Sprintf("%s_%s", b.namespace, b.profileName)
263+
264+
envVars := []corev1.EnvVar{
265+
{Name: constants.EnvVarNameUserID, Value: userID},
266+
{Name: constants.EnvVarTaskName, Value: taskName},
267+
{Name: constants.EnvVarTaskID, Value: fmt.Sprintf("%s_%s_%s", userID, timestamp, taskName)},
268+
{Name: constants.EnvVarNameProfile, Value: fmt.Sprintf("%s_%s", b.namespace, b.profileName)},
269+
{Name: constants.EnvVarNameTimestamp, Value: timestamp},
270+
}
271+
272+
return envVars
273+
}
274+
251275
func mergeBundles(bundles []v1alpha1.VolumeBundle) v1alpha1.VolumeBundle {
252276
var volumeBundle v1alpha1.VolumeBundle
253277
for _, b := range bundles {

‎cmd/experimental/kjobctl/pkg/builder/builder_test.go

+54-120
Original file line numberDiff line numberDiff line change
@@ -20,10 +20,10 @@ import (
2020
"context"
2121
"errors"
2222
"testing"
23+
"time"
2324

2425
"github.com/google/go-cmp/cmp"
2526
"github.com/google/go-cmp/cmp/cmpopts"
26-
batchv1 "k8s.io/api/batch/v1"
2727
apierrors "k8s.io/apimachinery/pkg/api/errors"
2828
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
2929
"k8s.io/apimachinery/pkg/runtime"
@@ -33,9 +33,12 @@ import (
3333
"sigs.k8s.io/kueue/cmd/experimental/kjobctl/client-go/clientset/versioned/fake"
3434
cmdtesting "sigs.k8s.io/kueue/cmd/experimental/kjobctl/pkg/cmd/testing"
3535
"sigs.k8s.io/kueue/cmd/experimental/kjobctl/pkg/constants"
36+
"sigs.k8s.io/kueue/cmd/experimental/kjobctl/pkg/testing/wrappers"
3637
)
3738

3839
func TestBuilder(t *testing.T) {
40+
testStartTime := time.Now()
41+
3942
testCases := map[string]struct {
4043
namespace string
4144
profile string
@@ -61,15 +64,9 @@ func TestBuilder(t *testing.T) {
6164
namespace: metav1.NamespaceDefault,
6265
profile: "profile",
6366
kjobctlObjs: []runtime.Object{
64-
&v1alpha1.ApplicationProfile{
65-
ObjectMeta: metav1.ObjectMeta{
66-
Name: "profile",
67-
Namespace: metav1.NamespaceDefault,
68-
},
69-
Spec: v1alpha1.ApplicationProfileSpec{
70-
SupportedModes: []v1alpha1.SupportedMode{{Name: v1alpha1.JobMode}},
71-
},
72-
},
67+
wrappers.MakeApplicationProfile("profile", metav1.NamespaceDefault).
68+
WithSupportedMode(v1alpha1.SupportedMode{Name: v1alpha1.JobMode}).
69+
Obj(),
7370
},
7471
wantErr: noApplicationProfileModeSpecifiedErr,
7572
},
@@ -78,15 +75,9 @@ func TestBuilder(t *testing.T) {
7875
profile: "profile",
7976
mode: v1alpha1.JobMode,
8077
kjobctlObjs: []runtime.Object{
81-
&v1alpha1.ApplicationProfile{
82-
ObjectMeta: metav1.ObjectMeta{
83-
Name: "profile",
84-
Namespace: metav1.NamespaceDefault,
85-
},
86-
Spec: v1alpha1.ApplicationProfileSpec{
87-
SupportedModes: []v1alpha1.SupportedMode{{Name: v1alpha1.InteractiveMode}},
88-
},
89-
},
78+
wrappers.MakeApplicationProfile("profile", metav1.NamespaceDefault).
79+
WithSupportedMode(v1alpha1.SupportedMode{Name: v1alpha1.InteractiveMode}).
80+
Obj(),
9081
},
9182
wantErr: applicationProfileModeNotConfiguredErr,
9283
},
@@ -95,15 +86,9 @@ func TestBuilder(t *testing.T) {
9586
profile: "profile",
9687
mode: "Invalid",
9788
kjobctlObjs: []runtime.Object{
98-
&v1alpha1.ApplicationProfile{
99-
ObjectMeta: metav1.ObjectMeta{
100-
Name: "profile",
101-
Namespace: metav1.NamespaceDefault,
102-
},
103-
Spec: v1alpha1.ApplicationProfileSpec{
104-
SupportedModes: []v1alpha1.SupportedMode{{Name: v1alpha1.InteractiveMode}},
105-
},
106-
},
89+
wrappers.MakeApplicationProfile("profile", metav1.NamespaceDefault).
90+
WithSupportedMode(v1alpha1.SupportedMode{Name: v1alpha1.InteractiveMode}).
91+
Obj(),
10792
},
10893
wantErr: invalidApplicationProfileModeErr,
10994
},
@@ -112,18 +97,12 @@ func TestBuilder(t *testing.T) {
11297
profile: "profile",
11398
mode: v1alpha1.JobMode,
11499
kjobctlObjs: []runtime.Object{
115-
&v1alpha1.ApplicationProfile{
116-
ObjectMeta: metav1.ObjectMeta{
117-
Name: "profile",
118-
Namespace: metav1.NamespaceDefault,
119-
},
120-
Spec: v1alpha1.ApplicationProfileSpec{
121-
SupportedModes: []v1alpha1.SupportedMode{{
122-
Name: v1alpha1.JobMode,
123-
RequiredFlags: []v1alpha1.Flag{v1alpha1.CmdFlag},
124-
}},
125-
},
126-
},
100+
wrappers.MakeApplicationProfile("profile", metav1.NamespaceDefault).
101+
WithSupportedMode(v1alpha1.SupportedMode{
102+
Name: v1alpha1.JobMode,
103+
RequiredFlags: []v1alpha1.Flag{v1alpha1.CmdFlag},
104+
}).
105+
Obj(),
127106
},
128107
wantErr: noCommandSpecifiedErr,
129108
},
@@ -132,18 +111,12 @@ func TestBuilder(t *testing.T) {
132111
profile: "profile",
133112
mode: v1alpha1.JobMode,
134113
kjobctlObjs: []runtime.Object{
135-
&v1alpha1.ApplicationProfile{
136-
ObjectMeta: metav1.ObjectMeta{
137-
Name: "profile",
138-
Namespace: metav1.NamespaceDefault,
139-
},
140-
Spec: v1alpha1.ApplicationProfileSpec{
141-
SupportedModes: []v1alpha1.SupportedMode{{
142-
Name: v1alpha1.JobMode,
143-
RequiredFlags: []v1alpha1.Flag{v1alpha1.ParallelismFlag},
144-
}},
145-
},
146-
},
114+
wrappers.MakeApplicationProfile("profile", metav1.NamespaceDefault).
115+
WithSupportedMode(v1alpha1.SupportedMode{
116+
Name: v1alpha1.JobMode,
117+
RequiredFlags: []v1alpha1.Flag{v1alpha1.ParallelismFlag},
118+
}).
119+
Obj(),
147120
},
148121
wantErr: noParallelismSpecifiedErr,
149122
},
@@ -152,18 +125,12 @@ func TestBuilder(t *testing.T) {
152125
profile: "profile",
153126
mode: v1alpha1.JobMode,
154127
kjobctlObjs: []runtime.Object{
155-
&v1alpha1.ApplicationProfile{
156-
ObjectMeta: metav1.ObjectMeta{
157-
Name: "profile",
158-
Namespace: metav1.NamespaceDefault,
159-
},
160-
Spec: v1alpha1.ApplicationProfileSpec{
161-
SupportedModes: []v1alpha1.SupportedMode{{
162-
Name: v1alpha1.JobMode,
163-
RequiredFlags: []v1alpha1.Flag{v1alpha1.CompletionsFlag},
164-
}},
165-
},
166-
},
128+
wrappers.MakeApplicationProfile("profile", metav1.NamespaceDefault).
129+
WithSupportedMode(v1alpha1.SupportedMode{
130+
Name: v1alpha1.JobMode,
131+
RequiredFlags: []v1alpha1.Flag{v1alpha1.CompletionsFlag},
132+
}).
133+
Obj(),
167134
},
168135
wantErr: noCompletionsSpecifiedErr,
169136
},
@@ -172,18 +139,12 @@ func TestBuilder(t *testing.T) {
172139
profile: "profile",
173140
mode: v1alpha1.JobMode,
174141
kjobctlObjs: []runtime.Object{
175-
&v1alpha1.ApplicationProfile{
176-
ObjectMeta: metav1.ObjectMeta{
177-
Name: "profile",
178-
Namespace: metav1.NamespaceDefault,
179-
},
180-
Spec: v1alpha1.ApplicationProfileSpec{
181-
SupportedModes: []v1alpha1.SupportedMode{{
182-
Name: v1alpha1.JobMode,
183-
RequiredFlags: []v1alpha1.Flag{v1alpha1.RequestFlag},
184-
}},
185-
},
186-
},
142+
wrappers.MakeApplicationProfile("profile", metav1.NamespaceDefault).
143+
WithSupportedMode(v1alpha1.SupportedMode{
144+
Name: v1alpha1.JobMode,
145+
RequiredFlags: []v1alpha1.Flag{v1alpha1.RequestFlag},
146+
}).
147+
Obj(),
187148
},
188149
wantErr: noRequestsSpecifiedErr,
189150
},
@@ -192,18 +153,12 @@ func TestBuilder(t *testing.T) {
192153
profile: "profile",
193154
mode: v1alpha1.JobMode,
194155
kjobctlObjs: []runtime.Object{
195-
&v1alpha1.ApplicationProfile{
196-
ObjectMeta: metav1.ObjectMeta{
197-
Name: "profile",
198-
Namespace: metav1.NamespaceDefault,
199-
},
200-
Spec: v1alpha1.ApplicationProfileSpec{
201-
SupportedModes: []v1alpha1.SupportedMode{{
202-
Name: v1alpha1.JobMode,
203-
RequiredFlags: []v1alpha1.Flag{v1alpha1.LocalQueueFlag},
204-
}},
205-
},
206-
},
156+
wrappers.MakeApplicationProfile("profile", metav1.NamespaceDefault).
157+
WithSupportedMode(v1alpha1.SupportedMode{
158+
Name: v1alpha1.JobMode,
159+
RequiredFlags: []v1alpha1.Flag{v1alpha1.LocalQueueFlag},
160+
}).
161+
Obj(),
207162
},
208163
wantErr: noLocalQueueSpecifiedErr,
209164
},
@@ -212,38 +167,17 @@ func TestBuilder(t *testing.T) {
212167
profile: "profile",
213168
mode: v1alpha1.JobMode,
214169
kjobctlObjs: []runtime.Object{
215-
&v1alpha1.JobTemplate{
216-
ObjectMeta: metav1.ObjectMeta{
217-
Namespace: metav1.NamespaceDefault,
218-
Name: "job-template",
219-
},
220-
},
221-
&v1alpha1.ApplicationProfile{
222-
ObjectMeta: metav1.ObjectMeta{
223-
Name: "profile",
224-
Namespace: metav1.NamespaceDefault,
225-
},
226-
Spec: v1alpha1.ApplicationProfileSpec{
227-
SupportedModes: []v1alpha1.SupportedMode{{
228-
Name: v1alpha1.JobMode,
229-
Template: "job-template",
230-
}},
231-
},
232-
},
233-
},
234-
wantObj: &batchv1.Job{
235-
TypeMeta: metav1.TypeMeta{
236-
Kind: "Job",
237-
APIVersion: "batch/v1",
238-
},
239-
ObjectMeta: metav1.ObjectMeta{
240-
GenerateName: "profile-",
241-
Namespace: metav1.NamespaceDefault,
242-
Labels: map[string]string{
243-
constants.ProfileLabel: "profile",
244-
},
245-
},
170+
wrappers.MakeJobTemplate("job-template", metav1.NamespaceDefault).Obj(),
171+
wrappers.MakeApplicationProfile("profile", metav1.NamespaceDefault).
172+
WithSupportedMode(v1alpha1.SupportedMode{
173+
Name: v1alpha1.JobMode,
174+
Template: "job-template",
175+
}).
176+
Obj(),
246177
},
178+
wantObj: wrappers.MakeJob("", metav1.NamespaceDefault).GenerateName("profile-").
179+
Label(constants.ProfileLabel, "profile").
180+
Obj(),
247181
},
248182
}
249183
for name, tc := range testCases {
@@ -253,7 +187,7 @@ func TestBuilder(t *testing.T) {
253187

254188
tcg := cmdtesting.NewTestClientGetter().
255189
WithKjobctlClientset(fake.NewSimpleClientset(tc.kjobctlObjs...))
256-
gotObjs, gotErr := NewBuilder(tcg).
190+
gotObjs, gotErr := NewBuilder(tcg, testStartTime).
257191
WithNamespace(tc.namespace).
258192
WithProfileName(tc.profile).
259193
WithModeName(tc.mode).

0 commit comments

Comments
 (0)