feat(controller): implement NetworkPolicy pruning for disabled features

controllers/llamastackdistribution_controller.go

@@ -42,8 +42,6 @@ import (
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/apimachinery/pkg/types"
-	"k8s.io/apimachinery/pkg/util/intstr"
-	"k8s.io/utils/ptr"
 	ctrl "sigs.k8s.io/controller-runtime"
 	"sigs.k8s.io/controller-runtime/pkg/builder"
 	"sigs.k8s.io/controller-runtime/pkg/client"
@@ -240,6 +238,30 @@ func (r *LlamaStackDistributionReconciler) determineKindsToExclude(instance *lla
 	return kinds
 }
 
+// pruneExcludedResources deletes resources of the specified kinds that are owned by the instance.
+// This ensures that when features are disabled, their associated resources are cleaned up.
+func (r *LlamaStackDistributionReconciler) pruneExcludedResources(ctx context.Context, instance *llamav1alpha1.LlamaStackDistribution, kindsToExclude []string) error {
+	logger := log.FromContext(ctx)
+
+	// Only NetworkPolicy requires explicit pruning since it was managed separately
+	// PVC and Service are managed through manifests and will be filtered out
+	for _, kind := range kindsToExclude {
+		if kind == "NetworkPolicy" {
+			networkPolicy := &networkingv1.NetworkPolicy{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      instance.Name + "-network-policy",
+					Namespace: instance.Namespace,
+				},
+			}
+			if err := deploy.HandleDisabledNetworkPolicy(ctx, r.Client, networkPolicy, logger); err != nil {
+				return fmt.Errorf("failed to prune NetworkPolicy: %w", err)
+			}
+		}
+	}
+
+	return nil
+}
+
 // reconcileManifestResources applies resources that are managed by the operator
 // based on the instance specification.
 func (r *LlamaStackDistributionReconciler) reconcileManifestResources(ctx context.Context, instance *llamav1alpha1.LlamaStackDistribution) error {
@@ -249,6 +271,11 @@ func (r *LlamaStackDistributionReconciler) reconcileManifestResources(ctx contex
 	}
 
 	kindsToExclude := r.determineKindsToExclude(instance)
+
+	if err := r.pruneExcludedResources(ctx, instance, kindsToExclude); err != nil {
+		return fmt.Errorf("failed to prune excluded resources: %w", err)
+	}
+
 	filteredResMap, err := deploy.FilterExcludeKinds(resMap, kindsToExclude)
 	if err != nil {
 		return fmt.Errorf("failed to filter manifests: %w", err)
@@ -278,11 +305,6 @@ func (r *LlamaStackDistributionReconciler) reconcileResources(ctx context.Contex
 		return err
 	}
 
-	// Reconcile the NetworkPolicy
-	if err := r.reconcileNetworkPolicy(ctx, instance); err != nil {
-		return fmt.Errorf("failed to reconcile NetworkPolicy: %w", err)
-	}
-
 	// Reconcile the Deployment
 	if err := r.reconcileDeployment(ctx, instance); err != nil {
 		return fmt.Errorf("failed to reconcile Deployment: %w", err)
@@ -1044,86 +1066,6 @@ func (r *LlamaStackDistributionReconciler) updateDistributionConfig(instance *ll
 	instance.Status.DistributionConfig.ActiveDistribution = activeDistribution
 }
 
-// reconcileNetworkPolicy manages the NetworkPolicy for the LlamaStack server.
-func (r *LlamaStackDistributionReconciler) reconcileNetworkPolicy(ctx context.Context, instance *llamav1alpha1.LlamaStackDistribution) error {
-	logger := log.FromContext(ctx)
-	networkPolicy := &networkingv1.NetworkPolicy{
-		ObjectMeta: metav1.ObjectMeta{
-			Name:      instance.Name + "-network-policy",
-			Namespace: instance.Namespace,
-		},
-	}
-
-	// If feature is disabled, delete the NetworkPolicy if it exists
-	if !r.EnableNetworkPolicy {
-		return deploy.HandleDisabledNetworkPolicy(ctx, r.Client, networkPolicy, logger)
-	}
-
-	port := deploy.GetServicePort(instance)
-
-	// get operator namespace
-	operatorNamespace, err := deploy.GetOperatorNamespace()
-	if err != nil {
-		return fmt.Errorf("failed to get operator namespace: %w", err)
-	}
-
-	networkPolicy.Spec = networkingv1.NetworkPolicySpec{
-		PodSelector: metav1.LabelSelector{
-			MatchLabels: map[string]string{
-				llamav1alpha1.DefaultLabelKey: llamav1alpha1.DefaultLabelValue,
-				"app.kubernetes.io/instance":  instance.Name,
-			},
-		},
-		PolicyTypes: []networkingv1.PolicyType{
-			networkingv1.PolicyTypeIngress,
-		},
-		Ingress: []networkingv1.NetworkPolicyIngressRule{
-			{
-				From: []networkingv1.NetworkPolicyPeer{
-					{ // to match all pods in all namespaces
-						PodSelector: &metav1.LabelSelector{
-							MatchLabels: map[string]string{
-								"app.kubernetes.io/part-of": llamav1alpha1.DefaultContainerName,
-							},
-						},
-						NamespaceSelector: &metav1.LabelSelector{}, // Empty namespaceSelector to match all namespaces
-					},
-				},
-				Ports: []networkingv1.NetworkPolicyPort{
-					{
-						Protocol: (*corev1.Protocol)(ptr.To("TCP")),
-						Port: &intstr.IntOrString{
-							IntVal: port,
-						},
-					},
-				},
-			},
-			{
-				From: []networkingv1.NetworkPolicyPeer{
-					{ // to match all pods in matched namespace
-						PodSelector: &metav1.LabelSelector{},
-						NamespaceSelector: &metav1.LabelSelector{
-							MatchLabels: map[string]string{
-								"kubernetes.io/metadata.name": operatorNamespace,
-							},
-						},
-					},
-				},
-				Ports: []networkingv1.NetworkPolicyPort{
-					{
-						Protocol: (*corev1.Protocol)(ptr.To("TCP")),
-						Port: &intstr.IntOrString{
-							IntVal: port,
-						},
-					},
-				},
-			},
-		},
-	}
-
-	return deploy.ApplyNetworkPolicy(ctx, r.Client, r.Scheme, instance, networkPolicy, logger)
-}
-
 // reconcileUserConfigMap validates that the referenced ConfigMap exists.
 func (r *LlamaStackDistributionReconciler) reconcileUserConfigMap(ctx context.Context, instance *llamav1alpha1.LlamaStackDistribution) error {
 	logger := log.FromContext(ctx)

controllers/llamastackdistribution_controller_test.go

@@ -272,8 +272,9 @@ func TestReconcile(t *testing.T) {
 		TargetPort: intstr.FromInt(int(instancePort)),
 		Protocol:   corev1.ProtocolTCP,
 	}
-	operatorNamespaceName := "test-operator-namespace"
+	enableNetworkPolicy := true
 
+	operatorNamespaceName := "test-operator-namespace"
 	// set operator namespace to avoid service account file dependency
 	t.Setenv("OPERATOR_NAMESPACE", operatorNamespaceName)
 
@@ -287,7 +288,7 @@ func TestReconcile(t *testing.T) {
 	require.NoError(t, k8sClient.Create(context.Background(), instance))
 
 	// --- act ---
-	ReconcileDistribution(t, instance, true)
+	ReconcileDistribution(t, instance, enableNetworkPolicy)
 
 	service := &corev1.Service{}
 	waitForResource(t, k8sClient, instance.Namespace, instance.Name+"-service", service)

controllers/manifests/base/kustomization.yaml

@@ -6,6 +6,7 @@ resources:
 - serviceaccount.yaml
 - scc-binding.yaml
 - service.yaml
+- networkpolicy.yaml
 
 labels:
 - includeSelectors: false

controllers/manifests/base/networkpolicy.yaml

@@ -0,0 +1,24 @@
+apiVersion: networking.k8s.io/v1
+kind: NetworkPolicy
+metadata:
+  name: network-policy
+spec:
+  podSelector:
+    matchLabels:
+      # app: llama-stack
+      # app.kubernetes.io/instance: llamastackdistribution-sample
+  policyTypes:
+  - Ingress
+  ingress:
+  - from:
+    - podSelector:
+        matchLabels:
+          app.kubernetes.io/part-of: llama-stack
+      namespaceSelector: {}
+    ports:
+    - protocol: TCP
+  - from:
+    - podSelector: {}
+      namespaceSelector: {}
+    ports:
+    - protocol: TCP

pkg/deploy/deploy.go

@@ -8,7 +8,7 @@ import (
 	"github.com/go-logr/logr"
 	llamav1alpha1 "github.com/llamastack/llama-stack-k8s-operator/api/v1alpha1"
 	appsv1 "k8s.io/api/apps/v1"
-	"k8s.io/apimachinery/pkg/api/errors"
+	k8serrors "k8s.io/apimachinery/pkg/api/errors"
 	"k8s.io/apimachinery/pkg/runtime"
 	ctrl "sigs.k8s.io/controller-runtime"
 	"sigs.k8s.io/controller-runtime/pkg/client"
@@ -23,7 +23,7 @@ func ApplyDeployment(ctx context.Context, cli client.Client, scheme *runtime.Sch
 
 	found := &appsv1.Deployment{}
 	err := cli.Get(ctx, client.ObjectKeyFromObject(deployment), found)
-	if err != nil && errors.IsNotFound(err) {
+	if err != nil && k8serrors.IsNotFound(err) {
 		logger.Info("Creating Deployment", "deployment", deployment.Name)
 		return cli.Create(ctx, deployment)
 	} else if err != nil {

pkg/deploy/kustomizer.go

@@ -211,6 +211,12 @@ func applyPlugins(resMap *resmap.ResMap, ownerInstance *llamav1alpha1.LlamaStack
 		return fmt.Errorf("failed to apply namespace setter plugin: %w", err)
 	}
 
+	operatorNamespace, err := GetOperatorNamespace()
+	if err != nil {
+		// In unit tests or non-cluster contexts, the operator namespace may be unavailable.
+		// Tolerate this and simply skip setting the operator-specific NetworkPolicy field mapping.
+		operatorNamespace = ""
+	}
 	fieldTransformerPlugin := plugins.CreateFieldMutator(plugins.FieldMutatorConfig{
 		Mappings: []plugins.FieldMapping{
 			{
@@ -260,6 +266,41 @@ func applyPlugins(resMap *resmap.ResMap, ownerInstance *llamav1alpha1.LlamaStack
 				TargetKind:        "Service",
 				CreateIfNotExists: true,
 			},
+			{
+				SourceValue:       nil,
+				DefaultValue:      llamav1alpha1.DefaultLabelValue,
+				TargetField:       "/spec/podSelector/matchLabels/" + llamav1alpha1.DefaultLabelKey,
+				TargetKind:        "NetworkPolicy",
+				CreateIfNotExists: true,
+			},
+			{
+				SourceValue:       ownerInstance.GetName(),
+				DefaultValue:      nil,
+				TargetField:       "/spec/podSelector/matchLabels/app.kubernetes.io~1instance",
+				TargetKind:        "NetworkPolicy",
+				CreateIfNotExists: false,
+			},
+			{
+				SourceValue:       getServicePort(ownerInstance),
+				DefaultValue:      llamav1alpha1.DefaultServerPort,
+				TargetField:       "/spec/ingress/0/ports/0/port",
+				TargetKind:        "NetworkPolicy",
+				CreateIfNotExists: true,
+			},
+			{
+				SourceValue:       getServicePort(ownerInstance),
+				DefaultValue:      llamav1alpha1.DefaultServerPort,
+				TargetField:       "/spec/ingress/1/ports/0/port",
+				TargetKind:        "NetworkPolicy",
+				CreateIfNotExists: true,
+			},
+			{
+				SourceValue:       operatorNamespace,
+				DefaultValue:      nil,
+				TargetField:       "/spec/ingress/1/from/0/namespaceSelector/matchLabels/kubernetes.io~1metadata.name",
+				TargetKind:        "NetworkPolicy",
+				CreateIfNotExists: true,
+			},
 		},
 	})
 	if err := fieldTransformerPlugin.Transform(*resMap); err != nil {

pkg/deploy/kustomizer_test.go

@@ -62,6 +62,7 @@ func setupApplyResourcesTest(t *testing.T, ownerName string) (context.Context, s
 
 // TestRenderManifest contains all unit tests for the RenderManifest function.
 func TestRenderManifest(t *testing.T) {
+	t.Setenv("OPERATOR_NAMESPACE", "default")
 	t.Run("should render correctly with a standard layout", func(t *testing.T) {
 		// given an-memory filesystem with a standard kustomize layout
 		fsys := filesys.MakeFsInMemory()