@@ -75,14 +75,20 @@ func runMnistPyTorchAppWrapper(t *testing.T, accelerator string) {
75
75
Kind : "Job" ,
76
76
},
77
77
ObjectMeta : metav1.ObjectMeta {
78
- Name : "mnist" ,
79
- Namespace : namespace .Name ,
78
+ GenerateName : "mnist" ,
79
+ Namespace : namespace .Name ,
80
80
},
81
81
Spec : batchv1.JobSpec {
82
82
Completions : Ptr (int32 (1 )),
83
83
Parallelism : Ptr (int32 (1 )),
84
84
Template : corev1.PodTemplateSpec {
85
85
Spec : corev1.PodSpec {
86
+ Tolerations : []corev1.Toleration {
87
+ {
88
+ Key : "nvidia.com/gpu" ,
89
+ Operator : corev1 .TolerationOpExists ,
90
+ },
91
+ },
86
92
Containers : []corev1.Container {
87
93
{
88
94
Name : "job" ,
@@ -139,9 +145,9 @@ func runMnistPyTorchAppWrapper(t *testing.T, accelerator string) {
139
145
Kind : "AppWrapper" ,
140
146
},
141
147
ObjectMeta : metav1.ObjectMeta {
142
- Name : "mnist" ,
143
- Namespace : namespace .Name ,
144
- Labels : map [string ]string {"kueue.x-k8s.io/queue-name" : localQueue .Name },
148
+ GenerateName : "mnist" ,
149
+ Namespace : namespace .Name ,
150
+ Labels : map [string ]string {"kueue.x-k8s.io/queue-name" : localQueue .Name },
145
151
},
146
152
Spec : mcadv1beta2.AppWrapperSpec {
147
153
Components : []mcadv1beta2.AppWrapperComponent {
@@ -158,16 +164,18 @@ func runMnistPyTorchAppWrapper(t *testing.T, accelerator string) {
158
164
unstruct := unstructured.Unstructured {Object : awMap }
159
165
_ , err = test .Client ().Dynamic ().Resource (appWrapperResource ).Namespace (namespace .Name ).Create (test .Ctx (), & unstruct , metav1.CreateOptions {})
160
166
test .Expect (err ).NotTo (HaveOccurred ())
161
- test .T ().Logf ("Created AppWrapper %s/%s successfully" , aw .Namespace , aw .Name )
167
+ test .T ().Logf ("Created AppWrapper %s/%s successfully" , aw .Namespace , aw .GenerateName )
162
168
163
- test .T ().Logf ("Waiting for AppWrapper %s/%s to be running" , aw .Namespace , aw .Name )
164
- test .Eventually (AppWrapper (test , namespace , aw . Name ), TestTimeoutMedium ).
165
- Should (WithTransform (AppWrapperPhase , Equal (mcadv1beta2 .AppWrapperRunning )))
169
+ test .T ().Logf ("Waiting for AppWrapper %s/%s to be running" , aw .Namespace , aw .GenerateName )
170
+ test .Eventually (AppWrappers (test , namespace ), TestTimeoutMedium ).
171
+ Should (ContainElement ( WithTransform (AppWrapperPhase , Equal (mcadv1beta2 .AppWrapperRunning ) )))
166
172
167
173
test .T ().Logf ("Waiting for AppWrapper %s/%s to complete" , job .Namespace , job .Name )
168
- test .Eventually (AppWrapper (test , namespace , aw .Name ), TestTimeoutLong ).Should (
169
- Or (
170
- WithTransform (AppWrapperPhase , Equal (mcadv1beta2 .AppWrapperSucceeded )),
171
- WithTransform (AppWrapperPhase , Equal (mcadv1beta2 .AppWrapperFailed )),
174
+ test .Eventually (AppWrappers (test , namespace ), TestTimeoutLong ).Should (
175
+ ContainElement (
176
+ Or (
177
+ WithTransform (AppWrapperPhase , Equal (mcadv1beta2 .AppWrapperSucceeded )),
178
+ WithTransform (AppWrapperPhase , Equal (mcadv1beta2 .AppWrapperFailed )),
179
+ ),
172
180
))
173
181
}
0 commit comments