predicate.go 5.53 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
/*
 * 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.
 */

package controller_common

import (
	"context"
	"strings"
23
	"time"
24
25

	"k8s.io/apimachinery/pkg/api/meta"
26
27
	"k8s.io/client-go/discovery"
	ctrl "sigs.k8s.io/controller-runtime"
28
29
30
31
32
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/log"
	"sigs.k8s.io/controller-runtime/pkg/predicate"
)

33
34
35
type GroveConfig struct {
	// Enabled is automatically determined by checking if Grove CRDs are installed in the cluster
	Enabled bool
36
	// TerminationDelay configures the termination delay for Grove PodCliqueSets
37
38
39
	TerminationDelay time.Duration
}

40
41
42
43
44
type LWSConfig struct {
	// Enabled is automatically determined by checking if LWS CRDs are installed in the cluster
	Enabled bool
}

45
46
47
48
49
type KaiSchedulerConfig struct {
	// Enabled is automatically determined by checking if Kai-scheduler CRDs are installed in the cluster
	Enabled bool
}

50
51
52
53
54
type MpiRunConfig struct {
	// SecretName is the name of the secret containing the SSH key for MPI Run
	SecretName string
}

55
56
57
type Config struct {
	// Enable resources filtering, only the resources belonging to the given namespace will be handled.
	RestrictedNamespace string
58
	Grove               GroveConfig
59
	LWS                 LWSConfig
60
	KaiScheduler        KaiSchedulerConfig
61
62
63
	EtcdAddress         string
	NatsAddress         string
	IngressConfig       IngressConfig
64
65
	// ModelExpressURL is the URL of the Model Express server to inject into all pods
	ModelExpressURL string
66
67
	// PrometheusEndpoint is the URL of the Prometheus endpoint to use for metrics
	PrometheusEndpoint string
68
	MpiRun             MpiRunConfig
69
70
71
72
73
74
75
76
	// RBAC configuration for cross-namespace resource management
	RBAC RBACConfig
}

// RBACConfig holds configuration for RBAC management
type RBACConfig struct {
	// PlannerClusterRoleName is the name of the ClusterRole for planner (cluster-wide mode only)
	PlannerClusterRoleName string
77
78
	// DGDRProfilingClusterRoleName is the name of the ClusterRole for DGDR profiling jobs (cluster-wide mode only)
	DGDRProfilingClusterRoleName string
79
80
81
82
83
84
85
86
87
88
89
}

type IngressConfig struct {
	VirtualServiceGateway      string
	IngressControllerClassName string
	IngressControllerTLSSecret string
	IngressHostSuffix          string
}

func (i *IngressConfig) UseVirtualService() bool {
	return i.VirtualServiceGateway != ""
90
91
}

92
93
94
// DetectGroveAvailability checks if Grove is available by checking if the Grove API group is registered
// This approach uses the discovery client which is simpler and more reliable
func DetectGroveAvailability(ctx context.Context, mgr ctrl.Manager) bool {
95
96
97
98
99
100
101
102
103
	return detectAPIGroupAvailability(ctx, mgr, "grove.io")
}

// DetectLWSAvailability checks if LWS is available by checking if the LWS API group is registered
// This approach uses the discovery client which is simpler and more reliable
func DetectLWSAvailability(ctx context.Context, mgr ctrl.Manager) bool {
	return detectAPIGroupAvailability(ctx, mgr, "leaderworkerset.x-k8s.io")
}

104
105
106
107
108
109
// DetectKaiSchedulerAvailability checks if Kai-scheduler is available by checking if the scheduling.run.ai API group is registered
// This approach uses the discovery client which is simpler and more reliable
func DetectKaiSchedulerAvailability(ctx context.Context, mgr ctrl.Manager) bool {
	return detectAPIGroupAvailability(ctx, mgr, "scheduling.run.ai")
}

110
111
// detectAPIGroupAvailability checks if a specific API group is registered in the cluster
func detectAPIGroupAvailability(ctx context.Context, mgr ctrl.Manager, groupName string) bool {
112
113
114
115
	logger := log.FromContext(ctx)

	cfg := mgr.GetConfig()
	if cfg == nil {
116
		logger.Info("detection failed, no discovery client available", "group", groupName)
117
118
119
120
121
		return false
	}

	discoveryClient, err := discovery.NewDiscoveryClientForConfig(cfg)
	if err != nil {
122
		logger.Error(err, "detection failed, could not create discovery client", "group", groupName)
123
124
125
126
127
		return false
	}

	apiGroups, err := discoveryClient.ServerGroups()
	if err != nil {
128
		logger.Error(err, "detection failed, could not list server groups", "group", groupName)
129
130
131
132
		return false
	}

	for _, group := range apiGroups.Groups {
133
134
		if group.Name == groupName {
			logger.Info("API group is available", "group", groupName)
135
136
137
138
			return true
		}
	}

139
	logger.Info("API group not available", "group", groupName)
140
141
	return false
}
142

143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
func EphemeralDeploymentEventFilter(config Config) predicate.Predicate {
	return predicate.NewPredicateFuncs(func(o client.Object) bool {
		l := log.FromContext(context.Background())
		objMeta, err := meta.Accessor(o)
		if err != nil {
			l.Error(err, "Error extracting object metadata")
			return false
		}
		if config.RestrictedNamespace != "" {
			// in case of a restricted namespace, we only want to process the events that are in the restricted namespace
			return objMeta.GetNamespace() == config.RestrictedNamespace
		}
		// in all other cases, discard the event if it is destined to an ephemeral deployment
		if strings.Contains(objMeta.GetNamespace(), "ephemeral") {
			return false
		}
		return true
	})
}