Skip to content

Commit

Permalink
fix: node-collector resources (#319)
Browse files Browse the repository at this point in the history
Signed-off-by: chenk <[email protected]>
  • Loading branch information
chen-keinan authored Mar 19, 2024
1 parent 8685792 commit 284f521
Show file tree
Hide file tree
Showing 2 changed files with 14 additions and 6 deletions.
15 changes: 11 additions & 4 deletions pkg/jobs/builder.go
Original file line number Diff line number Diff line change
Expand Up @@ -104,7 +104,7 @@ func WithImagePullSecrets(imagePullSecrets []corev1.LocalObjectReference) JobOpt
}
}

func WithResourceRequirements(rr *corev1.ResourceRequirements) JobOption {
func WithResourceRequirements(rr corev1.ResourceRequirements) JobOption {
return func(j *JobBuilder) {
j.resourceRequirements = rr
}
Expand All @@ -127,6 +127,12 @@ func WithUseNodeSelectorParam(useNodeSelector bool) JobOption {
}
}

func WithReplaceResourceReq(replaceResourceReq bool) JobOption {
return func(j *JobBuilder) {
j.replaceResourceReq = replaceResourceReq
}
}

func GetJob(opts ...JobOption) (*batchv1.Job, error) {
jb := &JobBuilder{}
for _, opt := range opts {
Expand All @@ -152,7 +158,8 @@ type JobBuilder struct {
volumes []corev1.Volume
volumeMounts []corev1.VolumeMount
imagePullSecrets []corev1.LocalObjectReference
resourceRequirements *corev1.ResourceRequirements
resourceRequirements corev1.ResourceRequirements
replaceResourceReq bool
timeout time.Duration
nodeConfig bool
useNodeSelector bool
Expand Down Expand Up @@ -222,8 +229,8 @@ func (b *JobBuilder) build() (*batchv1.Job, error) {
if len(b.imagePullSecrets) > 0 {
job.Spec.Template.Spec.ImagePullSecrets = b.imagePullSecrets
}
if b.resourceRequirements != nil {
job.Spec.Template.Spec.Containers[0].Resources = *b.resourceRequirements
if b.replaceResourceReq {
job.Spec.Template.Spec.Containers[0].Resources = b.resourceRequirements
}
if len(b.volumeMounts) > 0 {
job.Spec.Template.Spec.Containers[0].VolumeMounts = b.volumeMounts
Expand Down
5 changes: 3 additions & 2 deletions pkg/jobs/collector.go
Original file line number Diff line number Diff line change
Expand Up @@ -51,7 +51,7 @@ type jobCollector struct {
volumeMounts []corev1.VolumeMount
imagePullSecrets []corev1.LocalObjectReference
collectorTimeout time.Duration
resourceRequirements *corev1.ResourceRequirements
resourceRequirements corev1.ResourceRequirements
nodeConfig bool
useNodeSelector bool
}
Expand Down Expand Up @@ -135,7 +135,7 @@ func WithJobTemplateName(name string) CollectorOption {
}
}

func WithContainerResourceRequirements(rr *corev1.ResourceRequirements) CollectorOption {
func WithContainerResourceRequirements(rr corev1.ResourceRequirements) CollectorOption {
return func(j *jobCollector) {
j.resourceRequirements = rr
}
Expand Down Expand Up @@ -333,6 +333,7 @@ func (jb *jobCollector) Apply(ctx context.Context, nodeName string) (*batchv1.Jo
WithContainerVolumeMounts(jb.volumeMounts),
WithPriorityClassName(jb.priorityClassName),
WithNodeName(nodeName),
WithReplaceResourceReq(true),
WithJobName(jb.name),
WithUseNodeSelectorParam(jb.useNodeSelector),
WithResourceRequirements(jb.resourceRequirements)}
Expand Down

0 comments on commit 284f521

Please sign in to comment.