sigs.k8s.io/kueue@v0.6.2/pkg/util/testingjobs/rayjob/wrappers.go (about)

     1  /*
     2  Copyright 2023 The Kubernetes Authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package rayjob
    18  
    19  import (
    20  	rayjobapi "github.com/ray-project/kuberay/ray-operator/apis/ray/v1alpha1"
    21  	corev1 "k8s.io/api/core/v1"
    22  	"k8s.io/apimachinery/pkg/api/resource"
    23  	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
    24  	"k8s.io/utils/ptr"
    25  
    26  	"sigs.k8s.io/kueue/pkg/controller/constants"
    27  )
    28  
    29  // JobWrapper wraps a RayJob.
    30  type JobWrapper struct{ rayjobapi.RayJob }
    31  
    32  // MakeJob creates a wrapper for a suspended rayJob
    33  func MakeJob(name, ns string) *JobWrapper {
    34  	return &JobWrapper{rayjobapi.RayJob{
    35  		ObjectMeta: metav1.ObjectMeta{
    36  			Name:        name,
    37  			Namespace:   ns,
    38  			Annotations: make(map[string]string, 1),
    39  		},
    40  		Spec: rayjobapi.RayJobSpec{
    41  			ShutdownAfterJobFinishes: true,
    42  			RayClusterSpec: &rayjobapi.RayClusterSpec{
    43  				HeadGroupSpec: rayjobapi.HeadGroupSpec{
    44  					RayStartParams: map[string]string{"p1": "v1"},
    45  					Template: corev1.PodTemplateSpec{
    46  						Spec: corev1.PodSpec{
    47  							Containers: []corev1.Container{
    48  								{
    49  									Name: "head-container",
    50  								},
    51  							},
    52  						},
    53  					},
    54  				},
    55  				WorkerGroupSpecs: []rayjobapi.WorkerGroupSpec{
    56  					{
    57  						GroupName:      "workers-group-0",
    58  						Replicas:       ptr.To[int32](1),
    59  						MinReplicas:    ptr.To[int32](0),
    60  						MaxReplicas:    ptr.To[int32](10),
    61  						RayStartParams: map[string]string{"p1": "v1"},
    62  						Template: corev1.PodTemplateSpec{
    63  							Spec: corev1.PodSpec{
    64  								Containers: []corev1.Container{
    65  									{
    66  										Name: "worker-container",
    67  									},
    68  								},
    69  							},
    70  						},
    71  					},
    72  				},
    73  			},
    74  			Suspend: true,
    75  		},
    76  	}}
    77  }
    78  
    79  // Obj returns the inner Job.
    80  func (j *JobWrapper) Obj() *rayjobapi.RayJob {
    81  	return &j.RayJob
    82  }
    83  
    84  // Suspend updates the suspend status of the job
    85  func (j *JobWrapper) Suspend(s bool) *JobWrapper {
    86  	j.Spec.Suspend = s
    87  	return j
    88  }
    89  
    90  // Queue updates the queue name of the job
    91  func (j *JobWrapper) Queue(queue string) *JobWrapper {
    92  	if j.Labels == nil {
    93  		j.Labels = make(map[string]string)
    94  	}
    95  	j.Labels[constants.QueueLabel] = queue
    96  	return j
    97  }
    98  
    99  func (j *JobWrapper) RequestWorkerGroup(name corev1.ResourceName, quantity string) *JobWrapper {
   100  	c := &j.Spec.RayClusterSpec.WorkerGroupSpecs[0].Template.Spec.Containers[0]
   101  	if c.Resources.Requests == nil {
   102  		c.Resources.Requests = corev1.ResourceList{name: resource.MustParse(quantity)}
   103  	} else {
   104  		c.Resources.Requests[name] = resource.MustParse(quantity)
   105  	}
   106  	return j
   107  }
   108  
   109  func (j *JobWrapper) RequestHead(name corev1.ResourceName, quantity string) *JobWrapper {
   110  	c := &j.Spec.RayClusterSpec.HeadGroupSpec.Template.Spec.Containers[0]
   111  	if c.Resources.Requests == nil {
   112  		c.Resources.Requests = corev1.ResourceList{name: resource.MustParse(quantity)}
   113  	} else {
   114  		c.Resources.Requests[name] = resource.MustParse(quantity)
   115  	}
   116  	return j
   117  }
   118  
   119  func (j *JobWrapper) ShutdownAfterJobFinishes(value bool) *JobWrapper {
   120  	j.Spec.ShutdownAfterJobFinishes = value
   121  	return j
   122  }
   123  
   124  func (j *JobWrapper) ClusterSelector(value map[string]string) *JobWrapper {
   125  	j.Spec.ClusterSelector = value
   126  	return j
   127  }
   128  
   129  func (j *JobWrapper) WithEnableAutoscaling(value *bool) *JobWrapper {
   130  	j.Spec.RayClusterSpec.EnableInTreeAutoscaling = value
   131  	return j
   132  }
   133  
   134  func (j *JobWrapper) WithWorkerGroups(workers ...rayjobapi.WorkerGroupSpec) *JobWrapper {
   135  	j.Spec.RayClusterSpec.WorkerGroupSpecs = workers
   136  	return j
   137  }
   138  
   139  func (j *JobWrapper) WithHeadGroupSpec(value rayjobapi.HeadGroupSpec) *JobWrapper {
   140  	j.Spec.RayClusterSpec.HeadGroupSpec = value
   141  	return j
   142  }
   143  
   144  func (j *JobWrapper) WithPriorityClassName(value string) *JobWrapper {
   145  	j.Spec.RayClusterSpec.HeadGroupSpec.Template.Spec.PriorityClassName = value
   146  	return j
   147  }
   148  
   149  func (j *JobWrapper) WithWorkerPriorityClassName(value string) *JobWrapper {
   150  	j.Spec.RayClusterSpec.WorkerGroupSpecs[0].Template.Spec.PriorityClassName = value
   151  	return j
   152  }
   153  
   154  // WorkloadPriorityClass updates job workloadpriorityclass.
   155  func (j *JobWrapper) WorkloadPriorityClass(wpc string) *JobWrapper {
   156  	if j.Labels == nil {
   157  		j.Labels = make(map[string]string)
   158  	}
   159  	j.Labels[constants.WorkloadPriorityClassLabel] = wpc
   160  	return j
   161  }