sigs.k8s.io/kueue@v0.6.2/pkg/util/testingjobs/raycluster/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 raycluster
    18  
    19  import (
    20  	rayv1 "github.com/ray-project/kuberay/ray-operator/apis/ray/v1"
    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  // ClusterWrapper wraps a RayCluster.
    30  type ClusterWrapper struct{ rayv1.RayCluster }
    31  
    32  // MakeCluster creates a wrapper for rayCluster
    33  func MakeCluster(name, ns string) *ClusterWrapper {
    34  	return &ClusterWrapper{rayv1.RayCluster{
    35  		ObjectMeta: metav1.ObjectMeta{
    36  			Name:        name,
    37  			Namespace:   ns,
    38  			Annotations: make(map[string]string, 1),
    39  		},
    40  		Spec: rayv1.RayClusterSpec{
    41  			HeadGroupSpec: rayv1.HeadGroupSpec{
    42  				RayStartParams: map[string]string{"p1": "v1"},
    43  				Template: corev1.PodTemplateSpec{
    44  					Spec: corev1.PodSpec{
    45  						NodeSelector: map[string]string{},
    46  						Containers: []corev1.Container{
    47  							{
    48  								Name: "head-container",
    49  							},
    50  						},
    51  					},
    52  				},
    53  			},
    54  			WorkerGroupSpecs: []rayv1.WorkerGroupSpec{
    55  				{
    56  					GroupName:      "workers-group-0",
    57  					Replicas:       ptr.To[int32](1),
    58  					MinReplicas:    ptr.To[int32](0),
    59  					MaxReplicas:    ptr.To[int32](10),
    60  					RayStartParams: map[string]string{"p1": "v1"},
    61  					Template: corev1.PodTemplateSpec{
    62  						Spec: corev1.PodSpec{
    63  							Containers: []corev1.Container{
    64  								{
    65  									Name: "worker-container",
    66  								},
    67  							},
    68  						},
    69  					},
    70  				},
    71  			},
    72  			Suspend: ptr.To(true),
    73  		},
    74  	}}
    75  }
    76  
    77  // NodeSelector adds a node selector to the job's head.
    78  func (j *ClusterWrapper) NodeSelectorHeadGroup(k, v string) *ClusterWrapper {
    79  	j.Spec.HeadGroupSpec.Template.Spec.NodeSelector[k] = v
    80  	return j
    81  }
    82  
    83  // Obj returns the inner Job.
    84  func (j *ClusterWrapper) Obj() *rayv1.RayCluster {
    85  	return &j.RayCluster
    86  }
    87  
    88  // Suspend updates the suspend status of the job
    89  func (j *ClusterWrapper) Suspend(s bool) *ClusterWrapper {
    90  	j.Spec.Suspend = &s
    91  	return j
    92  }
    93  
    94  func (j *ClusterWrapper) RequestWorkerGroup(name corev1.ResourceName, quantity string) *ClusterWrapper {
    95  	c := &j.Spec.WorkerGroupSpecs[0].Template.Spec.Containers[0]
    96  	if c.Resources.Requests == nil {
    97  		c.Resources.Requests = corev1.ResourceList{name: resource.MustParse(quantity)}
    98  	} else {
    99  		c.Resources.Requests[name] = resource.MustParse(quantity)
   100  	}
   101  	return j
   102  }
   103  
   104  func (j *ClusterWrapper) RequestHead(name corev1.ResourceName, quantity string) *ClusterWrapper {
   105  	c := &j.Spec.HeadGroupSpec.Template.Spec.Containers[0]
   106  	if c.Resources.Requests == nil {
   107  		c.Resources.Requests = corev1.ResourceList{name: resource.MustParse(quantity)}
   108  	} else {
   109  		c.Resources.Requests[name] = resource.MustParse(quantity)
   110  	}
   111  	return j
   112  }
   113  
   114  // Queue updates the queue name of the job
   115  func (j *ClusterWrapper) Queue(queue string) *ClusterWrapper {
   116  	if j.Labels == nil {
   117  		j.Labels = make(map[string]string)
   118  	}
   119  	j.Labels[constants.QueueLabel] = queue
   120  	return j
   121  }
   122  
   123  // Clone returns deep copy of the Job.
   124  func (j *ClusterWrapper) Clone() *ClusterWrapper {
   125  	return &ClusterWrapper{RayCluster: *j.DeepCopy()}
   126  }
   127  
   128  func (j *ClusterWrapper) WithEnableAutoscaling(value *bool) *ClusterWrapper {
   129  	j.Spec.EnableInTreeAutoscaling = value
   130  	return j
   131  }
   132  
   133  func (j *ClusterWrapper) WithWorkerGroups(workers ...rayv1.WorkerGroupSpec) *ClusterWrapper {
   134  	j.Spec.WorkerGroupSpecs = workers
   135  	return j
   136  }
   137  
   138  func (j *ClusterWrapper) WithHeadGroupSpec(value rayv1.HeadGroupSpec) *ClusterWrapper {
   139  	j.Spec.HeadGroupSpec = value
   140  	return j
   141  }
   142  
   143  func (j *ClusterWrapper) WithPriorityClassName(value string) *ClusterWrapper {
   144  	j.Spec.HeadGroupSpec.Template.Spec.PriorityClassName = value
   145  	return j
   146  }
   147  
   148  func (j *ClusterWrapper) WithWorkerPriorityClassName(value string) *ClusterWrapper {
   149  	j.Spec.WorkerGroupSpecs[0].Template.Spec.PriorityClassName = value
   150  	return j
   151  }
   152  
   153  // WorkloadPriorityClass updates job workloadpriorityclass.
   154  func (j *ClusterWrapper) WorkloadPriorityClass(wpc string) *ClusterWrapper {
   155  	if j.Labels == nil {
   156  		j.Labels = make(map[string]string)
   157  	}
   158  	j.Labels[constants.WorkloadPriorityClassLabel] = wpc
   159  	return j
   160  }