dynamocomponent_controller.go 49.2 KB
Newer Older
1
/*
2
 * SPDX-FileCopyrightText: Copyright (c) 2022 Atalaya Tech. Inc
3
4
5
6
7
8
9
10
11
12
13
14
15
16
 * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 * SPDX-License-Identifier: Apache-2.0
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
17
 * Modifications Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
 */

package controller

import (
	"bytes"
	"context"
	"crypto/md5"
	"encoding/hex"
	"fmt"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"text/template"
	"time"

	"emperror.dev/errors"
36
37
38
39
	commonconfig "github.com/ai-dynamo/dynamo/deploy/cloud/operator/internal/config"
	"github.com/ai-dynamo/dynamo/deploy/cloud/operator/internal/consts"
	commonconsts "github.com/ai-dynamo/dynamo/deploy/cloud/operator/internal/consts"
	"github.com/ai-dynamo/dynamo/deploy/cloud/operator/internal/controller_common"
40
	"github.com/apparentlymart/go-shquot/shquot"
41
42
43
44
45
46
	"github.com/awslabs/amazon-ecr-credential-helper/ecr-login"
	"github.com/chrismellard/docker-credential-acr-env/pkg/credhelper"
	"github.com/google/go-containerregistry/pkg/authn"
	"github.com/google/go-containerregistry/pkg/name"
	"github.com/google/go-containerregistry/pkg/v1/google"
	"github.com/google/go-containerregistry/pkg/v1/remote"
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
	"github.com/huandu/xstrings"
	"github.com/mitchellh/hashstructure/v2"
	"github.com/rs/xid"
	"github.com/sergeymakinen/go-quote/unix"
	"github.com/sirupsen/logrus"
	"gopkg.in/yaml.v2"
	batchv1 "k8s.io/api/batch/v1"
	corev1 "k8s.io/api/core/v1"
	k8serrors "k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/apimachinery/pkg/api/meta"
	"k8s.io/apimachinery/pkg/api/resource"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/client-go/tools/record"
	"k8s.io/utils/ptr"
	ctrl "sigs.k8s.io/controller-runtime"
	"sigs.k8s.io/controller-runtime/pkg/builder"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/log"
	"sigs.k8s.io/controller-runtime/pkg/predicate"

69
70
71
72
	apiStoreClient "github.com/ai-dynamo/dynamo/deploy/cloud/operator/api/dynamo/api_store_client"
	dynamoCommon "github.com/ai-dynamo/dynamo/deploy/cloud/operator/api/dynamo/common"
	"github.com/ai-dynamo/dynamo/deploy/cloud/operator/api/dynamo/schemas"
	nvidiacomv1alpha1 "github.com/ai-dynamo/dynamo/deploy/cloud/operator/api/v1alpha1"
73
74
)

75
76
// DynamoComponentReconciler reconciles a DynamoComponent object
type DynamoComponentReconciler struct {
77
78
79
80
81
82
	client.Client
	Scheme   *runtime.Scheme
	Recorder record.EventRecorder
	Config   controller_common.Config
}

83
84
85
// +kubebuilder:rbac:groups=nvidia.com,resources=dynamocomponents,verbs=get;list;watch;create;update;patch;delete
// +kubebuilder:rbac:groups=nvidia.com,resources=dynamocomponents/status,verbs=get;update;patch
// +kubebuilder:rbac:groups=nvidia.com,resources=dynamocomponents/finalizers,verbs=update
86
87
//+kubebuilder:rbac:groups=nvidia.com,resources=dynamocomponents,verbs=get;list;watch;create;update;patch;delete
//+kubebuilder:rbac:groups=nvidia.com,resources=dynamocomponents/status,verbs=get;update;patch
88
89
90
91
92
93
94
95
96
97
98
//+kubebuilder:rbac:groups=events.k8s.io,resources=events,verbs=get;list;watch;create;update;patch;delete
//+kubebuilder:rbac:groups=core,resources=events,verbs=get;list;watch;create;update;patch;delete
//+kubebuilder:rbac:groups=core,resources=pods,verbs=get;list;watch;create;update;patch;delete
//+kubebuilder:rbac:groups=core,resources=configmaps,verbs=get;list;watch;create;update;patch;delete
//+kubebuilder:rbac:groups=core,resources=secrets,verbs=get;list;watch;create;update;patch;delete
//+kubebuilder:rbac:groups=core,resources=serviceaccounts,verbs=get;list;watch
//+kubebuilder:rbac:groups=coordination.k8s.io,resources=leases,verbs=get;list;watch;create;update;patch;delete

// Reconcile is part of the main kubernetes reconciliation loop which aims to
// move the current state of the cluster closer to the desired state.
// TODO(user): Modify the Reconcile function to compare the state specified by
99
// the DynamoComponent object against the actual cluster state, and then
100
101
102
103
104
105
106
// perform operations to make the cluster state reflect the state specified by
// the user.
//
// For more details, check Reconcile and its Result here:
// - https://pkg.go.dev/sigs.k8s.io/controller-runtime@v0.18.2/pkg/reconcile
//
//nolint:gocyclo,nakedret
107
func (r *DynamoComponentReconciler) Reconcile(ctx context.Context, req ctrl.Request) (result ctrl.Result, err error) {
108
109
	logs := log.FromContext(ctx)

110
	DynamoComponent := &nvidiacomv1alpha1.DynamoComponent{}
111

112
	err = r.Get(ctx, req.NamespacedName, DynamoComponent)
113
114
115
116
117

	if err != nil {
		if k8serrors.IsNotFound(err) {
			// Object not found, return.  Created objects are automatically garbage collected.
			// For additional cleanup logic use finalizers.
118
			logs.Info("DynamoComponent resource not found. Ignoring since object must be deleted")
119
120
121
122
			err = nil
			return
		}
		// Error reading the object - requeue the request.
123
		logs.Error(err, "Failed to get DynamoComponent")
124
125
126
		return
	}

127
128
129
	if DynamoComponent.IsReady() {
		logs.Info("Skip available DynamoComponent")
		return
130
131
	}

132
133
	if len(DynamoComponent.Status.Conditions) == 0 {
		DynamoComponent, err = r.setStatusConditions(ctx, req,
134
			metav1.Condition{
135
				Type:    nvidiacomv1alpha1.DynamoComponentConditionTypeImageBuilding,
136
137
				Status:  metav1.ConditionUnknown,
				Reason:  "Reconciling",
138
				Message: "Starting to reconcile DynamoComponent",
139
140
			},
			metav1.Condition{
141
				Type:    nvidiacomv1alpha1.DynamoComponentConditionTypeImageExists,
142
143
				Status:  metav1.ConditionUnknown,
				Reason:  "Reconciling",
144
				Message: "Starting to reconcile DynamoComponent",
145
146
			},
			metav1.Condition{
147
				Type:    nvidiacomv1alpha1.DynamoComponentConditionTypeDynamoComponentAvailable,
148
149
				Status:  metav1.ConditionUnknown,
				Reason:  "Reconciling",
150
				Message: "Reconciling",
151
152
153
154
155
156
157
			},
		)
		if err != nil {
			return
		}
	}

158
	logs = logs.WithValues("DynamoComponent", DynamoComponent.Name, "DynamoComponentNamespace", DynamoComponent.Namespace)
159
160
161
162
163
164

	defer func() {
		if err == nil {
			logs.Info("Reconcile success")
			return
		}
165
166
		logs.Error(err, "Failed to reconcile DynamoComponent.")
		r.Recorder.Eventf(DynamoComponent, corev1.EventTypeWarning, "ReconcileError", "Failed to reconcile DynamoComponent: %v", err)
167
168
		_, err_ := r.setStatusConditions(ctx, req,
			metav1.Condition{
169
				Type:    nvidiacomv1alpha1.DynamoComponentConditionTypeDynamoComponentAvailable,
170
171
172
173
174
175
				Status:  metav1.ConditionFalse,
				Reason:  "Reconciling",
				Message: err.Error(),
			},
		)
		if err_ != nil {
176
			logs.Error(err_, "Failed to update DynamoComponent status")
177
178
179
180
			return
		}
	}()

181
182
183
	DynamoComponent, _, imageExists, imageExistsResult, err := r.ensureImageExists(ctx, ensureImageExistsOption{
		DynamoComponent: DynamoComponent,
		req:             req,
184
185
186
187
188
189
190
191
192
	})

	if err != nil {
		err = errors.Wrapf(err, "ensure image exists")
		return
	}

	if !imageExists {
		result = imageExistsResult
193
		DynamoComponent, err = r.setStatusConditions(ctx, req,
194
			metav1.Condition{
195
				Type:    nvidiacomv1alpha1.DynamoComponentConditionTypeDynamoComponentAvailable,
196
197
				Status:  metav1.ConditionUnknown,
				Reason:  "Reconciling",
198
				Message: "DynamoComponent image is building",
199
200
201
202
203
204
205
206
			},
		)
		if err != nil {
			return
		}
		return
	}

207
	DynamoComponent, err = r.setStatusConditions(ctx, req,
208
		metav1.Condition{
209
			Type:    nvidiacomv1alpha1.DynamoComponentConditionTypeDynamoComponentAvailable,
210
211
			Status:  metav1.ConditionTrue,
			Reason:  "Reconciling",
212
			Message: "DynamoComponent image is generated",
213
214
215
216
217
218
219
220
221
222
223
224
225
226
		},
	)
	if err != nil {
		return
	}

	return
}

func isEstargzEnabled() bool {
	return os.Getenv("ESTARGZ_ENABLED") == commonconsts.KubeLabelValueTrue
}

type ensureImageExistsOption struct {
227
228
	DynamoComponent *nvidiacomv1alpha1.DynamoComponent
	req             ctrl.Request
229
230
231
}

//nolint:gocyclo,nakedret
232
func (r *DynamoComponentReconciler) ensureImageExists(ctx context.Context, opt ensureImageExistsOption) (DynamoComponent *nvidiacomv1alpha1.DynamoComponent, imageInfo ImageInfo, imageExists bool, result ctrl.Result, err error) { // nolint: unparam
233
234
	logs := log.FromContext(ctx)

235
	DynamoComponent = opt.DynamoComponent
236
237
	req := opt.req

238
	imageInfo = r.getImageInfo(GetImageInfoOption{
239
		DynamoComponent: DynamoComponent,
240
241
	})

242
243
	imageExistsCheckedCondition := meta.FindStatusCondition(DynamoComponent.Status.Conditions, nvidiacomv1alpha1.DynamoComponentConditionTypeImageExistsChecked)
	imageExistsCondition := meta.FindStatusCondition(DynamoComponent.Status.Conditions, nvidiacomv1alpha1.DynamoComponentConditionTypeImageExists)
244
245
	if imageExistsCheckedCondition == nil || imageExistsCheckedCondition.Status != metav1.ConditionTrue || imageExistsCheckedCondition.Message != imageInfo.ImageName {
		imageExistsCheckedCondition = &metav1.Condition{
246
			Type:    nvidiacomv1alpha1.DynamoComponentConditionTypeImageExistsChecked,
247
248
249
250
			Status:  metav1.ConditionUnknown,
			Reason:  "Reconciling",
			Message: imageInfo.ImageName,
		}
251
		dynamoComponentAvailableCondition := &metav1.Condition{
252
			Type:    nvidiacomv1alpha1.DynamoComponentConditionTypeDynamoComponentAvailable,
253
254
255
256
			Status:  metav1.ConditionUnknown,
			Reason:  "Reconciling",
			Message: "Checking image exists",
		}
257
		DynamoComponent, err = r.setStatusConditions(ctx, req, *imageExistsCheckedCondition, *dynamoComponentAvailableCondition)
258
259
260
		if err != nil {
			return
		}
261
		r.Recorder.Eventf(DynamoComponent, corev1.EventTypeNormal, "CheckingImage", "Checking image exists: %s", imageInfo.ImageName)
262
		imageExists, err = checkImageExists(DynamoComponent, imageInfo.ImageName)
263
264
265
266
267
		if err != nil {
			err = errors.Wrapf(err, "check image %s exists", imageInfo.ImageName)
			return
		}

268
		err = r.Get(ctx, req.NamespacedName, DynamoComponent)
269
		if err != nil {
270
			logs.Error(err, "Failed to re-fetch DynamoComponent")
271
272
273
274
			return
		}

		if imageExists {
275
			r.Recorder.Eventf(DynamoComponent, corev1.EventTypeNormal, "CheckingImage", "Image exists: %s", imageInfo.ImageName)
276
			imageExistsCheckedCondition = &metav1.Condition{
277
				Type:    nvidiacomv1alpha1.DynamoComponentConditionTypeImageExistsChecked,
278
279
280
281
282
				Status:  metav1.ConditionTrue,
				Reason:  "Reconciling",
				Message: imageInfo.ImageName,
			}
			imageExistsCondition = &metav1.Condition{
283
				Type:    nvidiacomv1alpha1.DynamoComponentConditionTypeImageExists,
284
285
286
287
				Status:  metav1.ConditionTrue,
				Reason:  "Reconciling",
				Message: imageInfo.ImageName,
			}
288
			DynamoComponent, err = r.setStatusConditions(ctx, req, *imageExistsCondition, *imageExistsCheckedCondition)
289
290
291
292
			if err != nil {
				return
			}
		} else {
293
			r.Recorder.Eventf(DynamoComponent, corev1.EventTypeNormal, "CheckingImage", "Image not exists: %s", imageInfo.ImageName)
294
			imageExistsCheckedCondition = &metav1.Condition{
295
				Type:    nvidiacomv1alpha1.DynamoComponentConditionTypeImageExistsChecked,
296
297
298
299
300
				Status:  metav1.ConditionFalse,
				Reason:  "Reconciling",
				Message: fmt.Sprintf("Image not exists: %s", imageInfo.ImageName),
			}
			imageExistsCondition = &metav1.Condition{
301
				Type:    nvidiacomv1alpha1.DynamoComponentConditionTypeImageExists,
302
303
304
305
				Status:  metav1.ConditionFalse,
				Reason:  "Reconciling",
				Message: fmt.Sprintf("Image %s is not exists", imageInfo.ImageName),
			}
306
			DynamoComponent, err = r.setStatusConditions(ctx, req, *imageExistsCondition, *imageExistsCheckedCondition)
307
308
309
310
311
312
			if err != nil {
				return
			}
		}
	}

313
314
	var DynamoComponentHashStr string
	DynamoComponentHashStr, err = r.getHashStr(DynamoComponent)
315
	if err != nil {
316
		err = errors.Wrapf(err, "get DynamoComponent %s/%s hash string", DynamoComponent.Namespace, DynamoComponent.Name)
317
318
319
320
321
322
323
324
325
		return
	}

	imageExists = imageExistsCondition != nil && imageExistsCondition.Status == metav1.ConditionTrue && imageExistsCondition.Message == imageInfo.ImageName
	if imageExists {
		return
	}

	jobLabels := map[string]string{
326
327
		commonconsts.KubeLabelDynamoComponent:      DynamoComponent.Name,
		commonconsts.KubeLabelIsDynamoImageBuilder: commonconsts.KubeLabelValueTrue,
328
329
330
331
332
333
334
335
336
337
338
339
340
	}

	jobs := &batchv1.JobList{}
	err = r.List(ctx, jobs, client.InNamespace(req.Namespace), client.MatchingLabels(jobLabels))
	if err != nil {
		err = errors.Wrap(err, "list jobs")
		return
	}

	reservedJobs := make([]*batchv1.Job, 0)
	for _, job_ := range jobs.Items {
		job_ := job_

341
342
343
		oldHash := job_.Annotations[consts.KubeAnnotationDynamoComponentHash]
		if oldHash != DynamoComponentHashStr {
			logs.Info("Because hash changed, delete old job", "job", job_.Name, "oldHash", oldHash, "newHash", DynamoComponentHashStr)
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
			// --cascade=foreground
			err = r.Delete(ctx, &job_, &client.DeleteOptions{
				PropagationPolicy: &[]metav1.DeletionPropagation{metav1.DeletePropagationForeground}[0],
			})
			if err != nil {
				err = errors.Wrapf(err, "delete job %s", job_.Name)
				return
			}
			return
		} else {
			reservedJobs = append(reservedJobs, &job_)
		}
	}

	var job *batchv1.Job
	if len(reservedJobs) > 0 {
		job = reservedJobs[0]
	}

	if len(reservedJobs) > 1 {
		for _, job_ := range reservedJobs[1:] {
			logs.Info("Because has more than one job, delete old job", "job", job_.Name)
			// --cascade=foreground
			err = r.Delete(ctx, job_, &client.DeleteOptions{
				PropagationPolicy: &[]metav1.DeletionPropagation{metav1.DeletePropagationForeground}[0],
			})
			if err != nil {
				err = errors.Wrapf(err, "delete job %s", job_.Name)
				return
			}
		}
	}

	if job == nil {
		job, err = r.generateImageBuilderJob(ctx, GenerateImageBuilderJobOption{
379
380
			ImageInfo:       imageInfo,
			DynamoComponent: DynamoComponent,
381
382
383
384
385
		})
		if err != nil {
			err = errors.Wrap(err, "generate image builder job")
			return
		}
386
		r.Recorder.Eventf(DynamoComponent, corev1.EventTypeNormal, "GenerateImageBuilderJob", "Creating image builder job: %s", job.Name)
387
388
389
390
391
		err = r.Create(ctx, job)
		if err != nil {
			err = errors.Wrapf(err, "create image builder job %s", job.Name)
			return
		}
392
		r.Recorder.Eventf(DynamoComponent, corev1.EventTypeNormal, "GenerateImageBuilderJob", "Created image builder job: %s", job.Name)
393
394
395
		return
	}

396
	r.Recorder.Eventf(DynamoComponent, corev1.EventTypeNormal, "CheckingImageBuilderJob", "Found image builder job: %s", job.Name)
397

398
	err = r.Get(ctx, req.NamespacedName, DynamoComponent)
399
	if err != nil {
400
		logs.Error(err, "Failed to re-fetch DynamoComponent")
401
402
		return
	}
403
	imageBuildingCondition := meta.FindStatusCondition(DynamoComponent.Status.Conditions, nvidiacomv1alpha1.DynamoComponentConditionTypeImageBuilding)
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427

	isJobFailed := false
	isJobRunning := true

	if job.Spec.Completions != nil {
		if job.Status.Succeeded != *job.Spec.Completions {
			if job.Status.Failed > 0 {
				for _, condition := range job.Status.Conditions {
					if condition.Type == batchv1.JobFailed && condition.Status == corev1.ConditionTrue {
						isJobFailed = true
						break
					}
				}
			}
			isJobRunning = !isJobFailed
		} else {
			isJobRunning = false
		}
	}

	if isJobRunning {
		conditions := make([]metav1.Condition, 0)
		if job.Status.Active > 0 {
			conditions = append(conditions, metav1.Condition{
428
				Type:    nvidiacomv1alpha1.DynamoComponentConditionTypeImageBuilding,
429
430
431
432
433
434
				Status:  metav1.ConditionTrue,
				Reason:  "Reconciling",
				Message: fmt.Sprintf("Image building job %s is running", job.Name),
			})
		} else {
			conditions = append(conditions, metav1.Condition{
435
				Type:    nvidiacomv1alpha1.DynamoComponentConditionTypeImageBuilding,
436
437
438
439
440
				Status:  metav1.ConditionUnknown,
				Reason:  "Reconciling",
				Message: fmt.Sprintf("Image building job %s is waiting", job.Name),
			})
		}
441
442
		if DynamoComponent.Spec.ImageBuildTimeout != nil {
			if imageBuildingCondition != nil && imageBuildingCondition.LastTransitionTime.Add(time.Duration(*DynamoComponent.Spec.ImageBuildTimeout)).Before(time.Now()) {
443
				conditions = append(conditions, metav1.Condition{
444
					Type:    nvidiacomv1alpha1.DynamoComponentConditionTypeImageBuilding,
445
446
447
448
449
450
451
452
453
454
455
456
					Status:  metav1.ConditionFalse,
					Reason:  "Timeout",
					Message: fmt.Sprintf("Image building job %s is timeout", job.Name),
				})
				if _, err = r.setStatusConditions(ctx, req, conditions...); err != nil {
					return
				}
				err = errors.New("image build timeout")
				return
			}
		}

457
		if DynamoComponent, err = r.setStatusConditions(ctx, req, conditions...); err != nil {
458
459
460
461
			return
		}

		if imageBuildingCondition != nil && imageBuildingCondition.Status != metav1.ConditionTrue && isJobRunning {
462
			r.Recorder.Eventf(DynamoComponent, corev1.EventTypeNormal, "DynamoComponentImageBuilder", "Image is building now")
463
464
465
466
467
468
		}

		return
	}

	if isJobFailed {
469
		DynamoComponent, err = r.setStatusConditions(ctx, req,
470
			metav1.Condition{
471
				Type:    nvidiacomv1alpha1.DynamoComponentConditionTypeImageBuilding,
472
473
474
475
476
				Status:  metav1.ConditionFalse,
				Reason:  "Reconciling",
				Message: fmt.Sprintf("Image building job %s is failed.", job.Name),
			},
			metav1.Condition{
477
				Type:    nvidiacomv1alpha1.DynamoComponentConditionTypeDynamoComponentAvailable,
478
479
480
481
482
483
484
485
486
487
488
				Status:  metav1.ConditionFalse,
				Reason:  "Reconciling",
				Message: fmt.Sprintf("Image building job %s is failed.", job.Name),
			},
		)
		if err != nil {
			return
		}
		return
	}

489
	DynamoComponent, err = r.setStatusConditions(ctx, req,
490
		metav1.Condition{
491
			Type:    nvidiacomv1alpha1.DynamoComponentConditionTypeImageBuilding,
492
493
494
495
496
			Status:  metav1.ConditionFalse,
			Reason:  "Reconciling",
			Message: fmt.Sprintf("Image building job %s is succeeded.", job.Name),
		},
		metav1.Condition{
497
			Type:    nvidiacomv1alpha1.DynamoComponentConditionTypeImageExists,
498
499
500
501
502
503
504
505
506
			Status:  metav1.ConditionTrue,
			Reason:  "Reconciling",
			Message: imageInfo.ImageName,
		},
	)
	if err != nil {
		return
	}

507
	r.Recorder.Eventf(DynamoComponent, corev1.EventTypeNormal, "DynamoComponentImageBuilder", "Image has been built successfully")
508
509
510
511
512
513

	imageExists = true

	return
}

514
515
func (r *DynamoComponentReconciler) setStatusConditions(ctx context.Context, req ctrl.Request, conditions ...metav1.Condition) (DynamoComponent *nvidiacomv1alpha1.DynamoComponent, err error) {
	DynamoComponent = &nvidiacomv1alpha1.DynamoComponent{}
516
517
518
519
520
521
	/*
		Please don't blame me when you see this kind of code,
		this is to avoid "the object has been modified; please apply your changes to the latest version and try again" when updating CR status,
		don't doubt that almost all CRD operators (e.g. cert-manager) can't avoid this stupid error and can only try to avoid this by this stupid way.
	*/
	for i := 0; i < 3; i++ {
522
523
		if err = r.Get(ctx, req.NamespacedName, DynamoComponent); err != nil {
			err = errors.Wrap(err, "Failed to re-fetch DynamoComponent")
524
525
526
			return
		}
		for _, condition := range conditions {
527
			meta.SetStatusCondition(&DynamoComponent.Status.Conditions, condition)
528
		}
529
		if err = r.Status().Update(ctx, DynamoComponent); err != nil {
530
531
532
533
534
535
			time.Sleep(100 * time.Millisecond)
		} else {
			break
		}
	}
	if err != nil {
536
		err = errors.Wrap(err, "Failed to update DynamoComponent status")
537
538
		return
	}
539
540
	if err = r.Get(ctx, req.NamespacedName, DynamoComponent); err != nil {
		err = errors.Wrap(err, "Failed to re-fetch DynamoComponent")
541
542
543
544
545
		return
	}
	return
}

546
type DynamoComponentImageBuildEngine string
547
548

const (
549
550
551
	DynamoComponentImageBuildEngineKaniko           DynamoComponentImageBuildEngine = "kaniko"
	DynamoComponentImageBuildEngineBuildkit         DynamoComponentImageBuildEngine = "buildkit"
	DynamoComponentImageBuildEngineBuildkitRootless DynamoComponentImageBuildEngine = "buildkit-rootless"
552
553
554
)

const (
555
	EnvDynamoImageBuildEngine = "DYNAMO_IMAGE_BUILD_ENGINE"
556
557
)

558
func getDynamoComponentImageBuildEngine() DynamoComponentImageBuildEngine {
559
	engine := os.Getenv(EnvDynamoImageBuildEngine)
560
	if engine == "" {
561
		return DynamoComponentImageBuildEngineKaniko
562
	}
563
	return DynamoComponentImageBuildEngine(engine)
564
565
566
}

//nolint:nakedret
567
func (r *DynamoComponentReconciler) getApiStoreClient(ctx context.Context) (*apiStoreClient.ApiStoreClient, *commonconfig.ApiStoreConfig, error) {
568
	apiStoreConf, err := commonconfig.GetApiStoreConfig(ctx)
569
570
	isNotFound := k8serrors.IsNotFound(err)
	if err != nil && !isNotFound {
571
572
		err = errors.Wrap(err, "get api store config")
		return nil, nil, err
573
574
575
	}

	if isNotFound {
576
		return nil, nil, err
577
578
	}

579
580
	if apiStoreConf.Endpoint == "" {
		return nil, nil, err
581
582
	}

583
584
	if apiStoreConf.ClusterName == "" {
		apiStoreConf.ClusterName = "default"
585
586
	}

587
	apiStoreClient := apiStoreClient.NewApiStoreClient(apiStoreConf.Endpoint)
588

589
	return apiStoreClient, apiStoreConf, nil
590
591
}

592
//nolint:nakedret
593
func (r *DynamoComponentReconciler) getDockerRegistry(DynamoComponent *nvidiacomv1alpha1.DynamoComponent) *schemas.DockerRegistrySchema {
594
595

	dockerRegistryConfig := commonconfig.GetDockerRegistryConfig()
596

597
598
599
	dynamoRepositoryName := "dynamo-components"
	if dockerRegistryConfig.DynamoComponentsRepositoryName != "" {
		dynamoRepositoryName = dockerRegistryConfig.DynamoComponentsRepositoryName
600
	}
601
	dynamoRepositoryURI := fmt.Sprintf("%s/%s", strings.TrimRight(dockerRegistryConfig.Server, "/"), dynamoRepositoryName)
602
603
604

	if DynamoComponent != nil && DynamoComponent.Spec.DockerConfigJSONSecretName != "" {
		dockerRegistryConfig.SecretName = DynamoComponent.Spec.DockerConfigJSONSecretName
605
	}
606

607
	return &schemas.DockerRegistrySchema{
608
609
610
611
612
		Server:              dockerRegistryConfig.Server,
		Secure:              dockerRegistryConfig.Secure,
		DynamoRepositoryURI: dynamoRepositoryURI,
		SecretName:          dockerRegistryConfig.SecretName,
	}
613
614
615
616
617
618
}

func isAddNamespacePrefix() bool {
	return os.Getenv("ADD_NAMESPACE_PREFIX_TO_IMAGE_NAME") == trueStr
}

619
620
func getDynamoComponentImagePrefix(DynamoComponent *nvidiacomv1alpha1.DynamoComponent) string {
	if DynamoComponent == nil {
621
622
		return ""
	}
623
	prefix, exist := DynamoComponent.Annotations[consts.KubeAnnotationDynamoComponentStorageNS]
624
625
626
627
	if exist && prefix != "" {
		return fmt.Sprintf("%s.", prefix)
	}
	if isAddNamespacePrefix() {
628
		return fmt.Sprintf("%s.", DynamoComponent.Namespace)
629
630
631
632
	}
	return ""
}

633
634
635
func getDynamoComponentImageName(DynamoComponent *nvidiacomv1alpha1.DynamoComponent, dockerRegistry schemas.DockerRegistrySchema, dynamoComponentRepositoryName, dynamoComponentVersion string) string {
	if DynamoComponent != nil && DynamoComponent.Spec.Image != "" {
		return DynamoComponent.Spec.Image
636
637
	}
	var uri, tag string
638
	uri = dockerRegistry.DynamoRepositoryURI
639
	tail := fmt.Sprintf("%s.%s", dynamoComponentRepositoryName, dynamoComponentVersion)
640
641
642
643
	if isEstargzEnabled() {
		tail += ".esgz"
	}

644
	tag = fmt.Sprintf("dynamo.%s%s", getDynamoComponentImagePrefix(DynamoComponent), tail)
645
646
647

	if len(tag) > 128 {
		hashStr := hash(tail)
648
		tag = fmt.Sprintf("dynamo.%s%s", getDynamoComponentImagePrefix(DynamoComponent), hashStr)
649
		if len(tag) > 128 {
650
			tag = fmt.Sprintf("dynamo.%s", hash(fmt.Sprintf("%s%s", getDynamoComponentImagePrefix(DynamoComponent), tail)))[:128]
651
652
653
654
655
		}
	}
	return fmt.Sprintf("%s:%s", uri, tag)
}

656
func checkImageExists(DynamoComponent *nvidiacomv1alpha1.DynamoComponent, imageName string) (bool, error) {
657
	if DynamoComponent.Annotations["nvidia.com/force-build-image"] == commonconsts.KubeLabelValueTrue {
658
659
		return false, nil
	}
660
	ref, err := name.ParseReference(imageName)
661
	if err != nil {
662
663
664
665
666
667
668
669
670
671
672
673
674
		return false, fmt.Errorf("parsing image reference: %w", err)
	}
	keychain := authn.NewMultiKeychain(
		// This picks up auth from DOCKER_CONFIG env var
		authn.DefaultKeychain,
		// This picks up auth from GCR
		google.Keychain,
		// This picks up auth from ECR
		authn.NewKeychainFromHelper(ecr.NewECRHelper()),
		// This picks up auth from ACR
		authn.NewKeychainFromHelper(credhelper.NewACRCredentialsHelper()),
	)
	_, err = remote.Head(ref, remote.WithAuthFromKeychain(keychain))
675
	if err != nil {
676
677
		if strings.Contains(err.Error(), "404") {
			return false, nil
678
		}
679
		return false, fmt.Errorf("checking image: %w", err)
680
	}
681
	return true, nil
682
683
684
}

type ImageInfo struct {
685
	DockerRegistry             schemas.DockerRegistrySchema
686
687
688
689
690
691
	DockerConfigJSONSecretName string
	ImageName                  string
	DockerRegistryInsecure     bool
}

type GetImageInfoOption struct {
692
	DynamoComponent *nvidiacomv1alpha1.DynamoComponent
693
694
695
}

//nolint:nakedret
696
func (r *DynamoComponentReconciler) getImageInfo(opt GetImageInfoOption) ImageInfo {
697
	dynamoComponentRepositoryName, _, dynamoComponentVersion := xstrings.Partition(opt.DynamoComponent.Spec.DynamoComponent, ":")
698
699
700
701
702
703
704
705
	dockerRegistry := r.getDockerRegistry(opt.DynamoComponent)
	imageInfo := ImageInfo{
		DockerRegistry:             *dockerRegistry,
		ImageName:                  getDynamoComponentImageName(opt.DynamoComponent, *dockerRegistry, dynamoComponentRepositoryName, dynamoComponentVersion),
		DockerConfigJSONSecretName: dockerRegistry.SecretName,
		DockerRegistryInsecure:     opt.DynamoComponent.Annotations[commonconsts.KubeAnnotationDynamoDockerRegistryInsecure] == "true",
	}
	return imageInfo
706
707
}

708
func (r *DynamoComponentReconciler) getImageBuilderJobName() string {
709
	guid := xid.New()
710
	return fmt.Sprintf("dynamo-image-builder-%s", guid.String())
711
712
}

713
func (r *DynamoComponentReconciler) getImageBuilderJobLabels(DynamoComponent *nvidiacomv1alpha1.DynamoComponent) map[string]string {
714
	return map[string]string{
715
716
		commonconsts.KubeLabelDynamoComponent:      DynamoComponent.Name,
		commonconsts.KubeLabelIsDynamoImageBuilder: "true",
717
718
719
	}
}

720
func (r *DynamoComponentReconciler) getImageBuilderPodLabels(DynamoComponent *nvidiacomv1alpha1.DynamoComponent) map[string]string {
721
	return map[string]string{
722
723
		commonconsts.KubeLabelDynamoComponent:      DynamoComponent.Name,
		commonconsts.KubeLabelIsDynamoImageBuilder: "true",
724
725
726
727
728
729
730
731
732
733
	}
}

func hash(text string) string {
	// nolint: gosec
	hasher := md5.New()
	hasher.Write([]byte(text))
	return hex.EncodeToString(hasher.Sum(nil))
}

734
type GenerateImageBuilderJobOption struct {
735
736
	ImageInfo       ImageInfo
	DynamoComponent *nvidiacomv1alpha1.DynamoComponent
737
738
739
}

//nolint:nakedret
740
func (r *DynamoComponentReconciler) generateImageBuilderJob(ctx context.Context, opt GenerateImageBuilderJobOption) (job *batchv1.Job, err error) {
741
	// nolint: gosimple
742
	podTemplateSpec, err := r.generateImageBuilderPodTemplateSpec(ctx, GenerateImageBuilderPodTemplateSpecOption(opt))
743
	if err != nil {
744
		err = errors.Wrap(err, "generate image builder pod template spec")
745
746
747
		return
	}
	kubeAnnotations := make(map[string]string)
748
	hashStr, err := r.getHashStr(opt.DynamoComponent)
749
750
751
752
	if err != nil {
		err = errors.Wrap(err, "failed to get hash string")
		return
	}
753
	kubeAnnotations[consts.KubeAnnotationDynamoComponentHash] = hashStr
754
755
	job = &batchv1.Job{
		ObjectMeta: metav1.ObjectMeta{
756
			Name:        r.getImageBuilderJobName(),
757
758
			Namespace:   opt.DynamoComponent.Namespace,
			Labels:      r.getImageBuilderJobLabels(opt.DynamoComponent),
759
760
761
			Annotations: kubeAnnotations,
		},
		Spec: batchv1.JobSpec{
762
763
764
			TTLSecondsAfterFinished: ptr.To(int32(60 * 60 * 24)),
			Completions:             ptr.To(int32(1)),
			Parallelism:             ptr.To(int32(1)),
765
766
767
768
769
			PodFailurePolicy: &batchv1.PodFailurePolicy{
				Rules: []batchv1.PodFailurePolicyRule{
					{
						Action: batchv1.PodFailurePolicyActionFailJob,
						OnExitCodes: &batchv1.PodFailurePolicyOnExitCodesRequirement{
770
							ContainerName: ptr.To(BuilderContainerName),
771
							Operator:      batchv1.PodFailurePolicyOnExitCodesOpIn,
772
							Values:        []int32{BuilderJobFailedExitCode},
773
774
775
776
777
778
779
						},
					},
				},
			},
			Template: *podTemplateSpec,
		},
	}
780
	err = ctrl.SetControllerReference(opt.DynamoComponent, job, r.Scheme)
781
782
783
784
785
786
787
	if err != nil {
		err = errors.Wrapf(err, "set controller reference for job %s", job.Name)
		return
	}
	return
}

788
func injectPodAffinity(podSpec *corev1.PodSpec, DynamoComponent *nvidiacomv1alpha1.DynamoComponent) {
789
790
791
	if podSpec.Affinity == nil {
		podSpec.Affinity = &corev1.Affinity{}
	}
792
793
794
795
796
797
798
799
800
801

	if podSpec.Affinity.PodAffinity == nil {
		podSpec.Affinity.PodAffinity = &corev1.PodAffinity{}
	}

	podSpec.Affinity.PodAffinity.PreferredDuringSchedulingIgnoredDuringExecution = append(podSpec.Affinity.PodAffinity.PreferredDuringSchedulingIgnoredDuringExecution, corev1.WeightedPodAffinityTerm{
		Weight: 100,
		PodAffinityTerm: corev1.PodAffinityTerm{
			LabelSelector: &metav1.LabelSelector{
				MatchLabels: map[string]string{
802
					commonconsts.KubeLabelDynamoComponent: DynamoComponent.Name,
803
804
805
806
807
808
809
810
811
812
813
814
815
				},
			},
			TopologyKey: corev1.LabelHostname,
		},
	})
}

const BuilderContainerName = "builder"
const BuilderJobFailedExitCode = 42
const ModelSeederContainerName = "seeder"
const ModelSeederJobFailedExitCode = 42

type GenerateImageBuilderPodTemplateSpecOption struct {
816
817
	ImageInfo       ImageInfo
	DynamoComponent *nvidiacomv1alpha1.DynamoComponent
818
819
820
}

//nolint:gocyclo,nakedret
821
822
823
func (r *DynamoComponentReconciler) generateImageBuilderPodTemplateSpec(ctx context.Context, opt GenerateImageBuilderPodTemplateSpecOption) (pod *corev1.PodTemplateSpec, err error) {
	dynamoComponentRepositoryName, _, dynamoComponentVersion := xstrings.Partition(opt.DynamoComponent.Spec.DynamoComponent, ":")
	kubeLabels := r.getImageBuilderPodLabels(opt.DynamoComponent)
824

825
	imageName := opt.ImageInfo.ImageName
826
827
828
829
830
831
832

	dockerConfigJSONSecretName := opt.ImageInfo.DockerConfigJSONSecretName

	dockerRegistryInsecure := opt.ImageInfo.DockerRegistryInsecure

	volumes := []corev1.Volume{
		{
833
			Name: "dynamo",
834
835
836
837
838
839
840
841
842
843
844
845
846
847
			VolumeSource: corev1.VolumeSource{
				EmptyDir: &corev1.EmptyDirVolumeSource{},
			},
		},
		{
			Name: "workspace",
			VolumeSource: corev1.VolumeSource{
				EmptyDir: &corev1.EmptyDirVolumeSource{},
			},
		},
	}

	volumeMounts := []corev1.VolumeMount{
		{
848
849
			Name:      "dynamo",
			MountPath: "/dynamo",
850
851
852
853
854
		},
		{
			Name:      "workspace",
			MountPath: "/workspace",
		},
855
856
857
858
		{
			Name:      consts.DockerConfigVolumeName,
			MountPath: consts.DockerConfigVolumeMountPath,
		},
859
860
861
862
	}

	if dockerConfigJSONSecretName != "" {
		volumes = append(volumes, corev1.Volume{
863
			Name: consts.DockerConfigVolumeName,
864
865
866
867
868
869
870
871
872
873
874
875
			VolumeSource: corev1.VolumeSource{
				Secret: &corev1.SecretVolumeSource{
					SecretName: dockerConfigJSONSecretName,
					Items: []corev1.KeyToPath{
						{
							Key:  ".dockerconfigjson",
							Path: "config.json",
						},
					},
				},
			},
		})
876
877
878
879
880
881
	} else {
		volumes = append(volumes, corev1.Volume{
			Name: consts.DockerConfigVolumeName,
			VolumeSource: corev1.VolumeSource{
				EmptyDir: &corev1.EmptyDirVolumeSource{},
			},
882
883
884
		})
	}

885
	var dynamoComponent *schemas.DynamoComponent
886
	dynamoComponentDownloadURL := opt.DynamoComponent.Spec.DownloadURL
887

888
889
890
	if dynamoComponentDownloadURL == "" {
		var apiStoreClient *apiStoreClient.ApiStoreClient
		var apiStoreConf *commonconfig.ApiStoreConfig
891

892
		apiStoreClient, apiStoreConf, err = r.getApiStoreClient(ctx)
893
		if err != nil {
894
			err = errors.Wrap(err, "get api store client")
895
896
897
			return
		}

898
899
		if apiStoreClient == nil || apiStoreConf == nil {
			err = errors.New("can't get api store client, please check api store configuration")
900
901
902
			return
		}

903
		r.Recorder.Eventf(opt.DynamoComponent, corev1.EventTypeNormal, "GenerateImageBuilderPod", "Getting dynamoComponent %s from api store service", opt.DynamoComponent.Spec.DynamoComponent)
904
		dynamoComponent, err = apiStoreClient.GetDynamoComponent(ctx, dynamoComponentRepositoryName, dynamoComponentVersion)
905
		if err != nil {
906
			err = errors.Wrap(err, "get dynamoComponent")
907
908
			return
		}
909
		r.Recorder.Eventf(opt.DynamoComponent, corev1.EventTypeNormal, "GenerateImageBuilderPod", "Got dynamoComponent %s from api store service", opt.DynamoComponent.Spec.DynamoComponent)
910

911
912
		if dynamoComponent.TransmissionStrategy != nil && *dynamoComponent.TransmissionStrategy == schemas.TransmissionStrategyPresignedURL {
			var dynamoComponent_ *schemas.DynamoComponent
913
			r.Recorder.Eventf(opt.DynamoComponent, corev1.EventTypeNormal, "GenerateImageBuilderPod", "Getting presigned url for dynamoComponent %s from api store service", opt.DynamoComponent.Spec.DynamoComponent)
914
			dynamoComponent_, err = apiStoreClient.PresignDynamoComponentDownloadURL(ctx, dynamoComponentRepositoryName, dynamoComponentVersion)
915
			if err != nil {
916
				err = errors.Wrap(err, "presign dynamoComponent download url")
917
918
				return
			}
919
			r.Recorder.Eventf(opt.DynamoComponent, corev1.EventTypeNormal, "GenerateImageBuilderPod", "Got presigned url for dynamoComponent %s from api store service", opt.DynamoComponent.Spec.DynamoComponent)
920
			dynamoComponentDownloadURL = dynamoComponent_.PresignedDownloadUrl
921
		} else {
922
			dynamoComponentDownloadURL = fmt.Sprintf("%s/api/v1/dynamo_components/%s/versions/%s/download", apiStoreConf.Endpoint, dynamoComponentRepositoryName, dynamoComponentVersion)
923
924
925
926
927
928
		}

	}
	internalImages := commonconfig.GetInternalImages()
	logrus.Infof("Image builder is using the images %v", *internalImages)

929
	buildEngine := getDynamoComponentImageBuildEngine()
930

931
	dynamoComponentDownloadCommandTemplate, err := template.New("downloadCommand").Parse(`
932
933
934
set -e

mkdir -p /workspace/buildcontext
935
url="{{.DynamoComponentDownloadURL}}"
936
echo "Downloading dynamoComponent {{.DynamoComponentRepositoryName}}:{{.DynamoComponentVersion}} to /tmp/downloaded.tar..."
937
938
939
940
941
942
943
if [[ ${url} == s3://* ]]; then
	echo "Downloading from s3..."
	aws s3 cp ${url} /tmp/downloaded.tar
elif [[ ${url} == gs://* ]]; then
	echo "Downloading from GCS..."
	gsutil cp ${url} /tmp/downloaded.tar
else
944
	curl --fail -L ${url} --output /tmp/downloaded.tar --progress-bar
945
946
fi
cd /workspace/buildcontext
947
echo "Extracting dynamoComponent tar file..."
948
tar -xvf /tmp/downloaded.tar
949
echo "Removing dynamoComponent tar file..."
950
951
952
953
954
955
956
957
958
rm /tmp/downloaded.tar
echo "Done"
	`)

	if err != nil {
		err = errors.Wrap(err, "failed to parse download command template")
		return
	}

959
	var dynamoComponentDownloadCommandBuffer bytes.Buffer
960

961
	err = dynamoComponentDownloadCommandTemplate.Execute(&dynamoComponentDownloadCommandBuffer, map[string]interface{}{
962
963
964
		"DynamoComponentDownloadURL":    dynamoComponentDownloadURL,
		"DynamoComponentRepositoryName": dynamoComponentRepositoryName,
		"DynamoComponentVersion":        dynamoComponentVersion,
965
966
967
968
969
970
	})
	if err != nil {
		err = errors.Wrap(err, "failed to execute download command template")
		return
	}

971
	dynamoComponentDownloadCommand := dynamoComponentDownloadCommandBuffer.String()
972
973
974
975
976
977
978
979
980
981
982
983

	downloaderContainerResources := corev1.ResourceRequirements{
		Limits: corev1.ResourceList{
			corev1.ResourceCPU:    resource.MustParse("1000m"),
			corev1.ResourceMemory: resource.MustParse("3000Mi"),
		},
		Requests: corev1.ResourceList{
			corev1.ResourceCPU:    resource.MustParse("100m"),
			corev1.ResourceMemory: resource.MustParse("1000Mi"),
		},
	}

984
	downloaderContainerEnvFrom := opt.DynamoComponent.Spec.DownloaderContainerEnvFrom
985
986
987

	initContainers := []corev1.Container{
		{
988
			Name:  "dynamocomponent-downloader",
989
			Image: internalImages.DynamoComponentsDownloader,
990
			Command: []string{
991
				"sh",
992
				"-c",
993
				dynamoComponentDownloadCommand,
994
995
996
997
998
999
1000
1001
1002
1003
1004
1005
1006
			},
			VolumeMounts: volumeMounts,
			Resources:    downloaderContainerResources,
			EnvFrom:      downloaderContainerEnvFrom,
			Env: []corev1.EnvVar{
				{
					Name:  "AWS_EC2_METADATA_DISABLED",
					Value: "true",
				},
			},
		},
	}

1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
	if dockerConfigJSONSecretName == "" {
		// if no explicit docker config is provided, we need to provide the docker config to the image builder
		var ref name.Reference
		ref, err = name.ParseReference(imageName)
		if err != nil {
			err = errors.Wrap(err, "failed to parse reference")
			return
		}
		dockerRegistry := ref.Context().RegistryStr()
		if isGoogleRegistry(dockerRegistry) {
			// for GCP, we use the google cloud sdk to get the docker config.
			initContainers = append(initContainers, corev1.Container{
				Name:  "gcp-init-docker-config",
				Image: "google/cloud-sdk:slim",
				Command: []string{
					"/bin/bash",
					"-c",
					fmt.Sprintf(`set -e
gcloud --quiet config get-value account
TOKEN=$(gcloud --quiet auth print-access-token)
cat > %s/config.json <<EOL
{
	"auths": {
		"%s": {
			"auth": "$(echo -n "oauth2accesstoken:${TOKEN}" | base64 -w 0)"
		}
	}
}
EOL
echo 'Docker config.json created successfully'`, consts.DockerConfigVolumeMountPath, dockerRegistry),
				},
				Resources:    downloaderContainerResources,
				EnvFrom:      downloaderContainerEnvFrom,
				VolumeMounts: volumeMounts,
			})
		}
	}

1045
1046
	containers := make([]corev1.Container, 0)

Neelay Shah's avatar
Neelay Shah committed
1047
1048
	var globalExtraPodMetadata *dynamoCommon.ExtraPodMetadata
	var globalExtraPodSpec *dynamoCommon.ExtraPodSpec
1049
1050
1051
1052
1053
	var globalExtraContainerEnv []corev1.EnvVar
	var globalDefaultImageBuilderContainerResources *corev1.ResourceRequirements
	var buildArgs []string
	var builderArgs []string

1054
	configNamespace, err := commonconfig.GetDynamoImageBuilderNamespace(ctx)
1055
	if err != nil {
1056
		err = errors.Wrap(err, "failed to get dynamo image builder namespace")
1057
1058
1059
		return
	}

1060
	configCmName := "dynamo-image-builder-config"
1061
	r.Recorder.Eventf(opt.DynamoComponent, corev1.EventTypeNormal, "GenerateImageBuilderPod", "Getting configmap %s from namespace %s", configCmName, configNamespace)
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
	configCm := &corev1.ConfigMap{}
	err = r.Get(ctx, types.NamespacedName{Name: configCmName, Namespace: configNamespace}, configCm)
	configCmIsNotFound := k8serrors.IsNotFound(err)
	if err != nil && !configCmIsNotFound {
		err = errors.Wrap(err, "failed to get configmap")
		return
	}
	err = nil // nolint: ineffassign

	if !configCmIsNotFound {
1072
		r.Recorder.Eventf(opt.DynamoComponent, corev1.EventTypeNormal, "GenerateImageBuilderPod", "Configmap %s is got from namespace %s", configCmName, configNamespace)
1073

Neelay Shah's avatar
Neelay Shah committed
1074
		globalExtraPodMetadata = &dynamoCommon.ExtraPodMetadata{}
1075
1076
1077
1078
1079
1080
1081
1082
1083

		if val, ok := configCm.Data["extra_pod_metadata"]; ok {
			err = yaml.Unmarshal([]byte(val), globalExtraPodMetadata)
			if err != nil {
				err = errors.Wrapf(err, "failed to yaml unmarshal extra_pod_metadata, please check the configmap %s in namespace %s", configCmName, configNamespace)
				return
			}
		}

Neelay Shah's avatar
Neelay Shah committed
1084
		globalExtraPodSpec = &dynamoCommon.ExtraPodSpec{}
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132

		if val, ok := configCm.Data["extra_pod_spec"]; ok {
			err = yaml.Unmarshal([]byte(val), globalExtraPodSpec)
			if err != nil {
				err = errors.Wrapf(err, "failed to yaml unmarshal extra_pod_spec, please check the configmap %s in namespace %s", configCmName, configNamespace)
				return
			}
		}

		globalExtraContainerEnv = []corev1.EnvVar{}

		if val, ok := configCm.Data["extra_container_env"]; ok {
			err = yaml.Unmarshal([]byte(val), &globalExtraContainerEnv)
			if err != nil {
				err = errors.Wrapf(err, "failed to yaml unmarshal extra_container_env, please check the configmap %s in namespace %s", configCmName, configNamespace)
				return
			}
		}

		if val, ok := configCm.Data["default_image_builder_container_resources"]; ok {
			globalDefaultImageBuilderContainerResources = &corev1.ResourceRequirements{}
			err = yaml.Unmarshal([]byte(val), globalDefaultImageBuilderContainerResources)
			if err != nil {
				err = errors.Wrapf(err, "failed to yaml unmarshal default_image_builder_container_resources, please check the configmap %s in namespace %s", configCmName, configNamespace)
				return
			}
		}

		buildArgs = []string{}

		if val, ok := configCm.Data["build_args"]; ok {
			err = yaml.Unmarshal([]byte(val), &buildArgs)
			if err != nil {
				err = errors.Wrapf(err, "failed to yaml unmarshal build_args, please check the configmap %s in namespace %s", configCmName, configNamespace)
				return
			}
		}

		builderArgs = []string{}
		if val, ok := configCm.Data["builder_args"]; ok {
			err = yaml.Unmarshal([]byte(val), &builderArgs)
			if err != nil {
				err = errors.Wrapf(err, "failed to yaml unmarshal builder_args, please check the configmap %s in namespace %s", configCmName, configNamespace)
				return
			}
		}
		logrus.Info("passed in builder args: ", builderArgs)
	} else {
1133
		r.Recorder.Eventf(opt.DynamoComponent, corev1.EventTypeNormal, "GenerateImageBuilderPod", "Configmap %s is not found in namespace %s", configCmName, configNamespace)
1134
1135
1136
1137
1138
1139
	}

	if buildArgs == nil {
		buildArgs = make([]string, 0)
	}

1140
1141
	if opt.DynamoComponent.Spec.BuildArgs != nil {
		buildArgs = append(buildArgs, opt.DynamoComponent.Spec.BuildArgs...)
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
	}

	dockerFilePath := "/workspace/buildcontext/env/docker/Dockerfile"

	builderContainerEnvFrom := make([]corev1.EnvFromSource, 0)
	builderContainerEnvs := []corev1.EnvVar{
		{
			Name:  "IFS",
			Value: "''",
		},
1152
		{
1153
			Name:  "DOCKER_CONFIG",
1154
1155
			Value: consts.DockerConfigVolumeMountPath,
		},
1156
1157
	}

1158
1159
	kanikoCacheRepo := os.Getenv("KANIKO_CACHE_REPO")
	if kanikoCacheRepo == "" {
1160
		kanikoCacheRepo = opt.ImageInfo.DockerRegistry.DynamoRepositoryURI
1161
1162
1163
	}

	kubeAnnotations := make(map[string]string)
1164
	kubeAnnotations[consts.KubeAnnotationDynamoComponentImageBuiderHash] = opt.DynamoComponent.Annotations[consts.KubeAnnotationDynamoComponentImageBuiderHash]
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179

	command := []string{
		"/kaniko/executor",
	}
	args := []string{
		"--context=/workspace/buildcontext",
		"--verbosity=info",
		"--image-fs-extract-retry=3",
		"--cache=false",
		fmt.Sprintf("--cache-repo=%s", kanikoCacheRepo),
		"--compressed-caching=false",
		"--compression=zstd",
		"--compression-level=-7",
		fmt.Sprintf("--dockerfile=%s", dockerFilePath),
		fmt.Sprintf("--insecure=%v", dockerRegistryInsecure),
1180
		fmt.Sprintf("--destination=%s", imageName),
1181
1182
1183
1184
1185
1186
1187
1188
1189
	}

	kanikoSnapshotMode := os.Getenv("KANIKO_SNAPSHOT_MODE")
	if kanikoSnapshotMode != "" {
		args = append(args, fmt.Sprintf("--snapshot-mode=%s", kanikoSnapshotMode))
	}

	var builderImage string
	switch buildEngine {
1190
	case DynamoComponentImageBuildEngineKaniko:
1191
1192
1193
1194
1195
1196
1197
		builderImage = internalImages.Kaniko
		if isEstargzEnabled() {
			builderContainerEnvs = append(builderContainerEnvs, corev1.EnvVar{
				Name:  "GGCR_EXPERIMENT_ESTARGZ",
				Value: "1",
			})
		}
1198
	case DynamoComponentImageBuildEngineBuildkit:
1199
		builderImage = internalImages.Buildkit
1200
	case DynamoComponentImageBuildEngineBuildkitRootless:
1201
1202
		builderImage = internalImages.BuildkitRootless
	default:
1203
		err = errors.Errorf("unknown dynamoComponent image build engine %s", buildEngine)
1204
1205
1206
		return
	}

1207
	isBuildkit := buildEngine == DynamoComponentImageBuildEngineBuildkit || buildEngine == DynamoComponentImageBuildEngineBuildkitRootless
1208
1209

	if isBuildkit {
1210
		output := fmt.Sprintf("type=image,name=%s,push=true,registry.insecure=%v", imageName, dockerRegistryInsecure)
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
		buildkitdFlags := []string{}
		if isEstargzEnabled() {
			buildkitdFlags = append(buildkitdFlags, "--oci-worker-snapshotter=stargz")
			output += ",oci-mediatypes=true,compression=estargz,force-compression=true"
		}
		if len(buildkitdFlags) > 0 {
			builderContainerEnvs = append(builderContainerEnvs, corev1.EnvVar{
				Name:  "BUILDKITD_FLAGS",
				Value: strings.Join(buildkitdFlags, " "),
			})
		}
1222
1223
1224
1225
1226
1227
1228
1229
		buildkitURL := os.Getenv("BUILDKIT_URL")
		if buildkitURL == "" {
			err = errors.New("BUILDKIT_URL is not set")
			return
		}
		command = []string{
			"buildctl",
		}
1230
		args = []string{
1231
1232
			"--addr",
			buildkitURL,
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
			"build",
			"--frontend",
			"dockerfile.v0",
			"--local",
			"context=/workspace/buildcontext",
			"--local",
			fmt.Sprintf("dockerfile=%s", filepath.Dir(dockerFilePath)),
			"--output",
			output,
		}
		cacheRepo := os.Getenv("BUILDKIT_CACHE_REPO")
1244
1245
1246
		if cacheRepo != "" {
			args = append(args, "--export-cache", fmt.Sprintf("type=registry,ref=%s:buildcache,mode=max,compression=zstd,ignore-error=true", cacheRepo))
			args = append(args, "--import-cache", fmt.Sprintf("type=registry,ref=%s:buildcache", cacheRepo))
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
		}
	}

	// add build args to pass via --build-arg
	for _, buildArg := range buildArgs {
		quotedBuildArg := unix.SingleQuote.Quote(buildArg)
		if isBuildkit {
			args = append(args, "--opt", fmt.Sprintf("build-arg:%s", quotedBuildArg))
		} else {
			args = append(args, fmt.Sprintf("--build-arg=%s", quotedBuildArg))
		}
	}
	// add other arguments to builder
	args = append(args, builderArgs...)
1261
	logrus.Info("dynamo-image-builder args: ", args)
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
1278

	builderContainerArgs := []string{
		"-c",
		fmt.Sprintf("sleep 15; %s && exit 0 || exit %d", shquot.POSIXShell(append(command, args...)), BuilderJobFailedExitCode), // TODO: remove once functionality exists to wait for istio sidecar.
	}

	container := corev1.Container{
		Name:            BuilderContainerName,
		Image:           builderImage,
		ImagePullPolicy: corev1.PullAlways,
		Command:         []string{"sh"},
		Args:            builderContainerArgs,
		VolumeMounts:    volumeMounts,
		Env:             builderContainerEnvs,
		EnvFrom:         builderContainerEnvFrom,
		TTY:             true,
		Stdin:           true,
1279
1280
1281
1282
1283
1284
1285
	}

	if buildEngine == DynamoComponentImageBuildEngineKaniko {
		// we need to run as root when using kaniko
		container.SecurityContext = &corev1.SecurityContext{
			RunAsUser: ptr.To(int64(0)),
		}
1286
1287
1288
1289
1290
1291
	}

	if globalDefaultImageBuilderContainerResources != nil {
		container.Resources = *globalDefaultImageBuilderContainerResources
	}

1292
1293
	if opt.DynamoComponent.Spec.ImageBuilderContainerResources != nil {
		container.Resources = *opt.DynamoComponent.Spec.ImageBuilderContainerResources
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
	}

	containers = append(containers, container)

	pod = &corev1.PodTemplateSpec{
		ObjectMeta: metav1.ObjectMeta{
			Labels:      kubeLabels,
			Annotations: kubeAnnotations,
		},
		Spec: corev1.PodSpec{
			RestartPolicy:  corev1.RestartPolicyNever,
			Volumes:        volumes,
			InitContainers: initContainers,
			Containers:     containers,
1308
1309
1310
1311
1312
			SecurityContext: &corev1.PodSecurityContext{
				RunAsUser:  ptr.To(int64(1000)),
				RunAsGroup: ptr.To(int64(1000)),
				FSGroup:    ptr.To(int64(1000)),
			},
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
		},
	}

	if globalExtraPodMetadata != nil {
		for k, v := range globalExtraPodMetadata.Annotations {
			pod.Annotations[k] = v
		}

		for k, v := range globalExtraPodMetadata.Labels {
			pod.Labels[k] = v
		}
	}

1326
1327
	if opt.DynamoComponent.Spec.ImageBuilderExtraPodMetadata != nil {
		for k, v := range opt.DynamoComponent.Spec.ImageBuilderExtraPodMetadata.Annotations {
1328
1329
1330
			pod.Annotations[k] = v
		}

1331
		for k, v := range opt.DynamoComponent.Spec.ImageBuilderExtraPodMetadata.Labels {
1332
1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
1343
1344
1345
			pod.Labels[k] = v
		}
	}

	if globalExtraPodSpec != nil {
		pod.Spec.PriorityClassName = globalExtraPodSpec.PriorityClassName
		pod.Spec.SchedulerName = globalExtraPodSpec.SchedulerName
		pod.Spec.NodeSelector = globalExtraPodSpec.NodeSelector
		pod.Spec.Affinity = globalExtraPodSpec.Affinity
		pod.Spec.Tolerations = globalExtraPodSpec.Tolerations
		pod.Spec.TopologySpreadConstraints = globalExtraPodSpec.TopologySpreadConstraints
		pod.Spec.ServiceAccountName = globalExtraPodSpec.ServiceAccountName
	}

1346
1347
1348
	if opt.DynamoComponent.Spec.ImageBuilderExtraPodSpec != nil {
		if opt.DynamoComponent.Spec.ImageBuilderExtraPodSpec.PriorityClassName != "" {
			pod.Spec.PriorityClassName = opt.DynamoComponent.Spec.ImageBuilderExtraPodSpec.PriorityClassName
1349
1350
		}

1351
1352
		if opt.DynamoComponent.Spec.ImageBuilderExtraPodSpec.SchedulerName != "" {
			pod.Spec.SchedulerName = opt.DynamoComponent.Spec.ImageBuilderExtraPodSpec.SchedulerName
1353
1354
		}

1355
1356
		if opt.DynamoComponent.Spec.ImageBuilderExtraPodSpec.NodeSelector != nil {
			pod.Spec.NodeSelector = opt.DynamoComponent.Spec.ImageBuilderExtraPodSpec.NodeSelector
1357
1358
		}

1359
1360
		if opt.DynamoComponent.Spec.ImageBuilderExtraPodSpec.Affinity != nil {
			pod.Spec.Affinity = opt.DynamoComponent.Spec.ImageBuilderExtraPodSpec.Affinity
1361
1362
		}

1363
1364
		if opt.DynamoComponent.Spec.ImageBuilderExtraPodSpec.Tolerations != nil {
			pod.Spec.Tolerations = opt.DynamoComponent.Spec.ImageBuilderExtraPodSpec.Tolerations
1365
1366
		}

1367
1368
		if opt.DynamoComponent.Spec.ImageBuilderExtraPodSpec.TopologySpreadConstraints != nil {
			pod.Spec.TopologySpreadConstraints = opt.DynamoComponent.Spec.ImageBuilderExtraPodSpec.TopologySpreadConstraints
1369
1370
		}

1371
1372
		if opt.DynamoComponent.Spec.ImageBuilderExtraPodSpec.ServiceAccountName != "" {
			pod.Spec.ServiceAccountName = opt.DynamoComponent.Spec.ImageBuilderExtraPodSpec.ServiceAccountName
1373
1374
1375
		}
	}

1376
	injectPodAffinity(&pod.Spec, opt.DynamoComponent)
1377
1378
1379

	if pod.Spec.ServiceAccountName == "" {
		serviceAccounts := &corev1.ServiceAccountList{}
1380
		err = r.List(ctx, serviceAccounts, client.InNamespace(opt.DynamoComponent.Namespace), client.MatchingLabels{
1381
			commonconsts.KubeLabelDynamoImageBuilderPod: commonconsts.KubeLabelValueTrue,
1382
1383
		})
		if err != nil {
1384
			err = errors.Wrapf(err, "failed to list service accounts in namespace %s", opt.DynamoComponent.Namespace)
1385
1386
1387
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
1398
			return
		}
		if len(serviceAccounts.Items) > 0 {
			pod.Spec.ServiceAccountName = serviceAccounts.Items[0].Name
		} else {
			pod.Spec.ServiceAccountName = "default"
		}
	}

	for i, c := range pod.Spec.InitContainers {
		env := c.Env
		if globalExtraContainerEnv != nil {
			env = append(env, globalExtraContainerEnv...)
		}
1399
		env = append(env, opt.DynamoComponent.Spec.ImageBuilderExtraContainerEnv...)
1400
1401
1402
1403
1404
1405
1406
		pod.Spec.InitContainers[i].Env = env
	}
	for i, c := range pod.Spec.Containers {
		env := c.Env
		if globalExtraContainerEnv != nil {
			env = append(env, globalExtraContainerEnv...)
		}
1407
		env = append(env, opt.DynamoComponent.Spec.ImageBuilderExtraContainerEnv...)
1408
1409
1410
1411
1412
1413
		pod.Spec.Containers[i].Env = env
	}

	return
}

1414
func (r *DynamoComponentReconciler) getHashStr(DynamoComponent *nvidiacomv1alpha1.DynamoComponent) (string, error) {
1415
1416
	var hash uint64
	hash, err := hashstructure.Hash(struct {
1417
		Spec        nvidiacomv1alpha1.DynamoComponentSpec
1418
1419
1420
		Labels      map[string]string
		Annotations map[string]string
	}{
1421
1422
1423
		Spec:        DynamoComponent.Spec,
		Labels:      DynamoComponent.Labels,
		Annotations: DynamoComponent.Annotations,
1424
1425
	}, hashstructure.FormatV2, nil)
	if err != nil {
1426
		err = errors.Wrap(err, "get DynamoComponent CR spec hash")
1427
1428
1429
1430
1431
1432
1433
1434
1435
1436
1437
		return "", err
	}
	hashStr := strconv.FormatUint(hash, 10)
	return hashStr, nil
}

const (
	trueStr = "true"
)

// SetupWithManager sets up the controller with the Manager.
1438
func (r *DynamoComponentReconciler) SetupWithManager(mgr ctrl.Manager) error {
1439
1440

	err := ctrl.NewControllerManagedBy(mgr).
1441
		For(&nvidiacomv1alpha1.DynamoComponent{}, builder.WithPredicates(predicate.GenerationChangedPredicate{})).
1442
		Owns(&nvidiacomv1alpha1.DynamoComponent{}).
1443
1444
1445
		Owns(&batchv1.Job{}).
		WithEventFilter(controller_common.EphemeralDeploymentEventFilter(r.Config)).
		Complete(r)
1446
	return errors.Wrap(err, "failed to setup DynamoComponent controller")
1447
}