Skip to content

BYOR OME changes #593

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
gsgvijay wants to merge 1 commit into
base: main
Choose a base branch
from
Open

BYOR OME changes #593

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions pkg/constants/constants.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -497,6 +497,7 @@ const (
KueueQueueLabelKey = "kueue.x-k8s.io/queue-name"
KueueWorkloadPriorityClassLabelKey = "kueue.x-k8s.io/priority-class"
KueueEnabledLabelKey = "kueue-enabled"
DACQueueNameLabelKey = "ome.io/dac-queue-name"
)

// Model Agent & Model Controller
Expand Down
15 changes: 12 additions & 3 deletions pkg/controller/v1beta1/inferenceservice/reconcilers/deployment/deployment_reconciler.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,18 +38,27 @@ func NewDeploymentReconciler(client kclient.Client,
return &DeploymentReconciler{
client: client,
scheme: scheme,
Deployment: createRawDeployment(componentMeta, componentExt, podSpec),
Deployment: createRawDeploymentWithClient(client, componentMeta, componentExt, podSpec),
componentExt: componentExt,
}
}

func createRawDeployment(componentMeta metav1.ObjectMeta,
func createRawDeployment(
componentMeta metav1.ObjectMeta,
componentExt *v1beta1.ComponentExtensionSpec,
podSpec *corev1.PodSpec) *appsv1.Deployment {
return createRawDeploymentWithClient(nil, componentMeta, componentExt, podSpec)
}

func createRawDeploymentWithClient(
cl kclient.Client,
componentMeta metav1.ObjectMeta,
componentExt *v1beta1.ComponentExtensionSpec,
podSpec *corev1.PodSpec) *appsv1.Deployment {

podMetadata := componentMeta
podMetadata.Labels["app"] = constants.TruncateNameWithMaxLength(componentMeta.Name, 63)
utils.SetPodLabelsFromAnnotations(&podMetadata)
utils.SetPodLabelsFromAnnotationsWithClient(cl, &podMetadata)
setDefaultPodSpec(podSpec)

deployment := &appsv1.Deployment{
Expand Down
8 changes: 5 additions & 3 deletions pkg/controller/v1beta1/inferenceservice/reconcilers/knative/ksvc_reconciler.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -51,13 +51,15 @@ func NewKsvcReconciler(client client.Client,
return &KsvcReconciler{
client: client,
scheme: scheme,
Service: createKnativeService(componentMeta, componentExt, podSpec, componentStatus),
Service: createKnativeService(client, componentMeta, componentExt, podSpec, componentStatus),
componentExt: componentExt,
componentStatus: componentStatus,
}
}

func createKnativeService(componentMeta metav1.ObjectMeta,
func createKnativeService(
cl client.Client,
componentMeta metav1.ObjectMeta,
componentExtension *v1beta1.ComponentExtensionSpec,
podSpec *corev1.PodSpec,
componentStatus v1beta1.ComponentStatusSpec) *knservingv1.Service {
Expand Down Expand Up @@ -138,7 +140,7 @@ func createKnativeService(componentMeta metav1.ObjectMeta,
return !utils.Includes(constants.RevisionTemplateLabelDisallowedList, key)
})

isvcutils.SetPodLabelsFromAnnotations(&componentMeta)
isvcutils.SetPodLabelsFromAnnotationsWithClient(cl, &componentMeta)

service := &knservingv1.Service{
ObjectMeta: metav1.ObjectMeta{
Expand Down
19 changes: 15 additions & 4 deletions pkg/controller/v1beta1/inferenceservice/reconcilers/lws/lws_reconciler.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,12 +38,23 @@ func NewLWSReconciler(client client.Client,
return &LWSReconciler{
client: client,
scheme: scheme,
LWS: createLWS(headPod, workerPod, workerSize, componentExt, componentMeta),
LWS: createLWSWithClient(client, headPod, workerPod, workerSize, componentExt, componentMeta),
ComponentExt: componentExt,
}
}

func createLWS(headPod *corev1.PodSpec,
func createLWS(
headPod *corev1.PodSpec,
workerPod *corev1.PodSpec,
workerSize int32,
componentExt *v1beta1.ComponentExtensionSpec,
componentMeta metav1.ObjectMeta) *lws.LeaderWorkerSet {
return createLWSWithClient(nil, headPod, workerPod, workerSize, componentExt, componentMeta)
}

func createLWSWithClient(
cl client.Client,
headPod *corev1.PodSpec,
workerPod *corev1.PodSpec,
workerSize int32,
componentExt *v1beta1.ComponentExtensionSpec,
Expand All @@ -55,8 +66,8 @@ func createLWS(headPod *corev1.PodSpec,
lwsObjectMeta.Name = constants.LWSName(componentMeta.Name)
headPodMeta.Labels["app"] = constants.GetRawServiceLabel(componentMeta.Name)
headPodMeta.Labels["ray.io/node-type"] = "head"
utils.SetPodLabelsFromAnnotations(headPodMeta)
utils.SetPodLabelsFromAnnotations(workerPodMeta)
utils.SetPodLabelsFromAnnotationsWithClient(cl, headPodMeta)
utils.SetPodLabelsFromAnnotationsWithClient(cl, workerPodMeta)

// Need to remove Prometheus annotations for workerPods as workerPods don't expose endpoints
abandonedWorkerPodAnnotations := []string{
Expand Down
5 changes: 3 additions & 2 deletions ...ntroller/v1beta1/inferenceservice/reconcilers/multinodevllm/multinodeprober_reconciler.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,7 @@ func NewMultiNodeProberReconciler(
Scheme: "http",
Host: fmt.Sprintf("%s.%s.svc.cluster.local", constants.DefaultRayHeadServiceName(componentMeta.Name, i), componentMeta.Namespace),
}
dply := createRawDeployment(componentMeta, multiNodeProberConfig, url, i)
dply := createRawDeployment(client, componentMeta, multiNodeProberConfig, url, i)
deployments = append(deployments, dply)
}
return &MultiNodeProberReconciler{
Expand All @@ -57,6 +57,7 @@ func NewMultiNodeProberReconciler(
}

func createRawDeployment(
cl kclient.Client,
componentMeta metav1.ObjectMeta,
multiNodeProberConfig *controllerconfig.MultiNodeProberConfig,
url *knapis.URL,
Expand All @@ -68,7 +69,7 @@ func createRawDeployment(
podMetadata.Labels = make(map[string]string)
}
podMetadata.Labels["app"] = constants.GetRawServiceLabel(componentMeta.Name)
utils.SetPodLabelsFromAnnotations(podMetadata)
utils.SetPodLabelsFromAnnotationsWithClient(cl, podMetadata)

podSpec := getDefaultPodSpec(multiNodeProberConfig, url)
deployment := &appsv1.Deployment{
Expand Down
6 changes: 3 additions & 3 deletions pkg/controller/v1beta1/inferenceservice/reconcilers/multinodevllm/ray.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,7 +48,7 @@ func NewRayReconciler(client client.Client,

rayClusters := make([]*ray.RayCluster, 0, int(*componentExt.MinReplicas))
for i := 0; i < int(*componentExt.MinReplicas); i++ {
rayCluster := createRayCluster(&componentMeta, podSpec, i)
rayCluster := createRayCluster(client, &componentMeta, podSpec, i)
rayClusters = append(rayClusters, rayCluster)
}

Expand Down Expand Up @@ -249,10 +249,10 @@ func (r *RayReconciler) deleteExtraRayClusters(existingRayClusters *ray.RayClust
return nil
}

func createRayCluster(meta *metav1.ObjectMeta, spec *corev1.PodSpec, index int) *ray.RayCluster {
func createRayCluster(cl client.Client, meta *metav1.ObjectMeta, spec *corev1.PodSpec, index int) *ray.RayCluster {
clusterName := fmt.Sprintf("%s-%d", meta.Name, index)

utils.SetPodLabelsFromAnnotations(meta)
utils.SetPodLabelsFromAnnotationsWithClient(cl, meta)
workerReplicas := int32(constants.DefaultMinReplicas)

setLifecycleHooks(spec)
Expand Down
19 changes: 18 additions & 1 deletion pkg/controller/v1beta1/inferenceservice/utils/annotations.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ import (
"strings"

metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
"sigs.k8s.io/controller-runtime/pkg/client"

"github.com/sgl-project/ome/pkg/constants"
"github.com/sgl-project/ome/pkg/controller/v1beta1/controllerconfig"
Expand Down Expand Up @@ -60,14 +61,18 @@ func IsCohereCommand1TFewFTServing(servingPodObjectMeta *metav1.ObjectMeta) bool
}

func SetPodLabelsFromAnnotations(metadata *metav1.ObjectMeta) {
SetPodLabelsFromAnnotationsWithClient(nil, metadata)
}

func SetPodLabelsFromAnnotationsWithClient(cl client.Client, metadata *metav1.ObjectMeta) {
// Check if the VolcanoQueue annotation exists and set the label if it does.
if volcanoQueue, ok := metadata.Annotations[constants.VolcanoQueue]; ok {
metadata.Labels[constants.VolcanoQueueName] = volcanoQueue
// If VolcanoQueue annotation does not exist, check and set to DedicatedAICluster name
} else if dac, ok := metadata.Annotations[constants.DedicatedAICluster]; ok {
if _, ok = metadata.Annotations[constants.KueueEnabledLabelKey]; ok {
// Kueue case
metadata.Labels[constants.KueueQueueLabelKey] = dac
metadata.Labels[constants.KueueQueueLabelKey] = resolveKueueQueueNameForDedicatedAICluster(cl, metadata, dac)
metadata.Labels[constants.KueueWorkloadPriorityClassLabelKey] = constants.DedicatedAiClusterPreemptionWorkloadPriorityClass
} else {
// Volcano case
Expand All @@ -82,6 +87,18 @@ func SetPodLabelsFromAnnotations(metadata *metav1.ObjectMeta) {
}
}

func resolveKueueQueueNameForDedicatedAICluster(_ client.Client, metadata *metav1.ObjectMeta, dacName string) string {
if metadata != nil {
if queueName, ok := metadata.Annotations[constants.DACQueueNameLabelKey]; ok {
queueName = strings.TrimSpace(queueName)
if queueName != "" {
return queueName
}
}
}
return dacName
}

func RemovePodAnnotations(metadata *metav1.ObjectMeta, annotationsToRemove []string) {
for _, annotation := range annotationsToRemove {
delete(metadata.Annotations, annotation)
Expand Down
Loading