Implement NetworkPolicy support

Signed-off-by: Pat O'Connor <[email protected]>

Author: chipspeak

ray-operator/config/manager/manager.yaml

@@ -80,4 +80,9 @@ spec:
           # environment variable is not set, requeue after the default value (300).
           # - name: RAYCLUSTER_DEFAULT_REQUEUE_SECONDS_ENV
           #   value: "300"
+          # Required for NetworkPolicy feature when operator is NOT deployed in 'ray-system' namespace
+          # - name: POD_NAMESPACE
+          #   valueFrom:
+          #     fieldRef:
+          #       fieldPath: metadata.namespace          
       terminationGracePeriodSeconds: 10

ray-operator/config/rbac/role.yaml

@@ -109,6 +109,19 @@ rules:
   - get
   - list
   - watch
+- apiGroups:
+  - networking.k8s.io
+  resources:
+  - ingresses
+  - networkpolicies
+  verbs:
+  - create
+  - delete
+  - get
+  - list
+  - patch
+  - update
+  - watch
 - apiGroups:
   - ray.io
   resources:

ray-operator/controllers/ray/raycluster_controller.go

@@ -100,6 +100,7 @@ type RayClusterReconcilerOptions struct {
 // +kubebuilder:rbac:groups=coordination.k8s.io,resources=leases,verbs=get;list;create;update
 // +kubebuilder:rbac:groups=networking.k8s.io,resources=ingressclasses,verbs=get;list;watch
 // +kubebuilder:rbac:groups=networking.k8s.io,resources=ingresses,verbs=get;list;watch;create;update;delete;patch
+// +kubebuilder:rbac:groups=networking.k8s.io,resources=networkpolicies,verbs=get;list;watch;create;update;delete;patch
 // +kubebuilder:rbac:groups=route.openshift.io,resources=routes,verbs=get;list;watch;create;update;patch;delete
 // +kubebuilder:rbac:groups=extensions,resources=ingresses,verbs=get;list;watch;create;update;delete;patch
 // +kubebuilder:rbac:groups=core,resources=serviceaccounts,verbs=get;list;watch;create;delete
@@ -298,6 +299,7 @@ func (r *RayClusterReconciler) rayClusterReconcile(ctx context.Context, instance
 		r.reconcileHeadlessService,
 		r.reconcileServeService,
 		r.reconcilePods,
+		r.reconcileNetworkPolicies,
 	}
 
 	for _, fn := range reconcileFuncs {
@@ -782,6 +784,97 @@ func (r *RayClusterReconciler) reconcilePods(ctx context.Context, instance *rayv
 	return nil
 }
 
+// reconcileNetworkPolicies creates a default deny NetworkPolicy for the RayCluster.
+func (r *RayClusterReconciler) reconcileNetworkPolicies(ctx context.Context, instance *rayv1.RayCluster) error {
+	logger := ctrl.LoggerFrom(ctx)
+
+	// Feature gate check
+	if !features.Enabled(features.RayClusterNetworkPolicy) {
+		logger.V(1).Info("NetworkPolicy feature disabled, skipping reconciliation")
+		return nil
+	}
+
+	logger.Info("Reconciling Network Policies")
+
+	networkPolicy := createDefaultDenyPolicy(instance)
+
+	if err := ctrl.SetControllerReference(instance, networkPolicy, r.Scheme); err != nil {
+		return err
+	}
+
+	if err := r.Create(ctx, networkPolicy); err != nil {
+		if errors.IsAlreadyExists(err) {
+			logger.Info("NetworkPolicy already exists, no need to create")
+			return nil
+		}
+		r.Recorder.Eventf(instance, corev1.EventTypeWarning, string(utils.FailedToCreateNetworkPolicy), "Failed to apply Head NetworkPolicy %s/%s, %v", networkPolicy.Namespace, networkPolicy.Name, err)
+		return err
+	}
+
+	logger.Info("Successfully created NetworkPolicy", "name", networkPolicy.Name)
+	r.Recorder.Eventf(instance, corev1.EventTypeNormal, string(utils.CreatedNetworkPolicy),
+		"Created NetworkPolicy %s/%s", networkPolicy.Namespace, networkPolicy.Name)
+
+	return nil
+}
+
+// createDefaultDenyPolicy creates a default deny NetworkPolicy for the RayCluster.
+func createDefaultDenyPolicy(instance *rayv1.RayCluster) *networkingv1.NetworkPolicy {
+	operatorNamespace := os.Getenv("POD_NAMESPACE")
+	if operatorNamespace == "" {
+		operatorNamespace = "ray-system" // fallback
+	}
+
+	labels := map[string]string{
+		utils.RayClusterLabelKey:                instance.Name,
+		utils.KubernetesApplicationNameLabelKey: utils.ApplicationName,
+		utils.KubernetesCreatedByLabelKey:       utils.ComponentName,
+	}
+
+	return &networkingv1.NetworkPolicy{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      fmt.Sprintf("%s-default-deny", instance.Name),
+			Namespace: instance.Namespace,
+			Labels:    labels,
+		},
+		Spec: networkingv1.NetworkPolicySpec{
+			PodSelector: metav1.LabelSelector{
+				MatchLabels: map[string]string{
+					utils.RayClusterLabelKey: instance.Name,
+				},
+			},
+			PolicyTypes: []networkingv1.PolicyType{networkingv1.PolicyTypeIngress},
+			Ingress: []networkingv1.NetworkPolicyIngressRule{
+				// Allow traffic from within the same cluster
+				{
+					From: []networkingv1.NetworkPolicyPeer{
+						{
+							PodSelector: &metav1.LabelSelector{
+								MatchLabels: map[string]string{
+									utils.RayClusterLabelKey: instance.Name,
+								},
+							},
+						},
+						// Allow KubeRay operator communication
+						{
+							PodSelector: &metav1.LabelSelector{
+								MatchLabels: map[string]string{
+									utils.KubernetesApplicationNameLabelKey: utils.ApplicationName,
+								},
+							},
+							NamespaceSelector: &metav1.LabelSelector{
+								MatchLabels: map[string]string{
+									"kubernetes.io/metadata.name": operatorNamespace,
+								},
+							},
+						},
+					},
+				},
+			},
+		},
+	}
+}
+
 // shouldDeletePod returns whether the Pod should be deleted and the reason
 //
 // @param pod: The Pod to be checked.

ray-operator/controllers/ray/utils/constant.go

@@ -315,6 +315,10 @@ const (
 	FailedToUpdateHeadPodServeLabel K8sEventType = "FailedToUpdateHeadPodServeLabel"
 	FailedToUpdateServeApplications K8sEventType = "FailedToUpdateServeApplications"
 
+	// NetworkPolicy event list
+	CreatedNetworkPolicy        K8sEventType = "CreatedNetworkPolicy"
+	FailedToCreateNetworkPolicy K8sEventType = "FailedToCreateNetworkPolicy"
+
 	// Generic Pod event list
 	DeletedPod                  K8sEventType = "DeletedPod"
 	FailedToDeletePod           K8sEventType = "FailedToDeletePod"

ray-operator/pkg/features/features.go

@@ -24,6 +24,9 @@ const (
 	//
 	// Enables new deletion policy API in RayJob
 	RayJobDeletionPolicy featuregate.Feature = "RayJobDeletionPolicy"
+
+	// Might be overkill to have a feature gate for this but for the sake of argument...
+	RayClusterNetworkPolicy featuregate.Feature = "RayClusterNetworkPolicy"
 )
 
 func init() {
@@ -33,6 +36,7 @@ func init() {
 var defaultFeatureGates = map[featuregate.Feature]featuregate.FeatureSpec{
 	RayClusterStatusConditions: {Default: true, PreRelease: featuregate.Beta},
 	RayJobDeletionPolicy:       {Default: false, PreRelease: featuregate.Alpha},
+	RayClusterNetworkPolicy:    {Default: false, PreRelease: featuregate.Alpha},
 }
 
 // SetFeatureGateDuringTest is a helper method to override feature gates in tests.