@@ -24,6 +24,7 @@ import (
24
24
. "github.com/onsi/gomega"
25
25
"github.com/samber/lo"
26
26
corev1 "k8s.io/api/core/v1"
27
+ "k8s.io/apimachinery/pkg/api/errors"
27
28
"k8s.io/apimachinery/pkg/api/resource"
28
29
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
29
30
"sigs.k8s.io/controller-runtime/pkg/client"
@@ -184,6 +185,44 @@ var _ = Describe("TensorFusionWorkload Controller", func() {
184
185
})
185
186
})
186
187
188
+ Context ("When deleting workload directly" , func () {
189
+ It ("Should delete all pods and the workload itself" , func () {
190
+ pool := tfEnv .GetGPUPool (0 )
191
+
192
+ workload := createTensorFusionWorkload (pool .Name , key , 2 )
193
+ checkWorkerPodCount (workload )
194
+ checkWorkloadStatus (workload )
195
+
196
+ // wait for 2 pods to be created
197
+ Eventually (func (g Gomega ) {
198
+ podList := & corev1.PodList {}
199
+ g .Expect (k8sClient .List (ctx , podList ,
200
+ client .InNamespace (key .Namespace ),
201
+ client.MatchingLabels {constants .WorkloadKey : key .Name })).To (Succeed ())
202
+ g .Expect (podList .Items ).To (HaveLen (2 ))
203
+ }, timeout , interval ).Should (Succeed ())
204
+
205
+ // delete workload
206
+ Expect (k8sClient .Delete (ctx , workload )).To (Succeed ())
207
+
208
+ // wait for all pods to be deleted
209
+ Eventually (func (g Gomega ) {
210
+ podList := & corev1.PodList {}
211
+ g .Expect (k8sClient .List (ctx , podList ,
212
+ client .InNamespace (key .Namespace ),
213
+ client.MatchingLabels {constants .WorkloadKey : key .Name })).To (Succeed ())
214
+ g .Expect (podList .Items ).Should (BeEmpty ())
215
+ }, timeout , interval ).Should (Succeed ())
216
+
217
+ // wait for workload itself to be deleted
218
+ Eventually (func (g Gomega ) {
219
+ w := & tfv1.TensorFusionWorkload {}
220
+ err := k8sClient .Get (ctx , key , w )
221
+ g .Expect (err ).To (HaveOccurred ())
222
+ }, timeout , interval ).Should (Succeed ())
223
+ })
224
+ })
225
+
187
226
Context ("When GPUPool doesn't exist" , func () {
188
227
It ("Should not create worker pod when reconciling a workload with non-existent pool" , func () {
189
228
workload := createTensorFusionWorkload ("non-existent-pool" , key , 1 )
@@ -262,6 +301,13 @@ func cleanupWorkload(key client.ObjectKey) {
262
301
GinkgoHelper ()
263
302
workload := & tfv1.TensorFusionWorkload {}
264
303
304
+ if err := k8sClient .Get (ctx , key , workload ); err != nil {
305
+ if errors .IsNotFound (err ) {
306
+ return
307
+ }
308
+ Expect (err ).To (HaveOccurred ())
309
+ }
310
+
265
311
// Set replicas to 0
266
312
Eventually (func (g Gomega ) {
267
313
g .Expect (k8sClient .Get (ctx , key , workload )).Should (Succeed ())
0 commit comments