Protect status server with network policy

[robert-bell]

What you would like to be added?

Add a network policy that allows only TrainJob pods to send requests to the status server.

Why is this needed?

This is adds an additional defence in depth. The status server is an internal endpoint that should only be called by train job pods.

Follow on from #3227 (comment)

Love this feature?

Give it a 👍 We prioritize the features with most 👍

[arryllopez]

Hello! I would like to work on this if it has not been assigned to anybody else!

[robert-bell]

Thanks @arryllopez! Feel free to assign it to yourself with a comment / assign.

Do you have an approach in mind? I would be happy to review it and provide feedback.

[arryllopez]

/assign

[arryllopez]

My approach would be to add a NetworkPolicy to the manifests that only allows TrainJob pods to reach the status server port. I'll look through the codebase to find the right labels and port numbers, then open a PR. Happy to hear any feedback before I start.

[arryllopez]

Would like some input on this

I've been looking through the runtime yamls and have found that the runtime yamls do not assign any labels to the pods, the labels are all assigned from the Job's metadata, resulting in the pods not having any label for the network_policy.yaml file im building to match against

so my proposed idea for fixing is to moving the label to the pod level across all of the runtime yaml files so then we can have a label to match against in network_policy

heres some snippets
Before (torch_distributed.yaml)

replicatedJobs:
  - name: node
    template:
      metadata:                    # Job metadata
        labels:
          trainer.kubeflow.org/trainjob-ancestor-step: trainer
      spec:
        template:                  # Pod template — no labels
          spec:
            containers: ...

(After)

replicatedJobs:
  - name: node
    template:
      metadata:                    # Job metadata
        labels:
          trainer.kubeflow.org/trainjob-ancestor-step: trainer
      spec:
        template:
          metadata:                # Pod template — label added here
            labels:
              trainer.kubeflow.org/trainjob-ancestor-step: trainer
          spec:
            containers: ...

then the network policy can match against it something looking like this

ingress:
  - from:
    - podSelector:
        matchLabels:
          trainer.kubeflow.org/trainjob-ancestor-step: trainer
    ports:
    - protocol: TCP
      port: 10443

Furthermore, Since the NetworkPolicy podSelector only matches pods in the same namespace, and TrainJob pods can run in any user namespace, how should we identify those namespaces? Should we require users to label their namespaces (e.g. trainer.kubeflow.org/trainjob-namespace: "true") so the namespaceSelector can match against them, or is there a simpler approach you'd prefer?

then it could be implemented like this

ingress:
- from:
  - namespaceSelector:
      matchLabels:
        trainer.kubeflow.org/trainjob-namespace: "true"
    podSelector:
      matchLabels:
        trainer.kubeflow.org/trainjob-ancestor-step: trainer

This allows connections only from pods with the label trainer.kubeflow.org/trainjob-ancestor-step: trainer that are running inside a namespace labeled with trainer.kubeflow.org/trainjob-namespace: "true".

These are my initial questions, I've got a pretty clear roadmap ahead, just looked through the kubernetes docs and saw how they documented Network Policies.

Happy to submit PR once I get your thoughts on the matter ! :)

[robert-bell]

Hey @arryllopez, I think should be able to add the label using the TrainJobStatus plug-in which would mean we don't have to update all runtimes. That may be easier!

You can also configure Network policies to match all namespaces so we shouldn't need users to label a namespace.

There's a couple of other bits that need some thought-

• the controller also exposes webhook and metrics ports so the netpol needs to allow traffic to these ports. The webhook requests will come from the api server, the metrics from an unknown pod. I don't think we can restrict traffic to these ports so perhaps we can just allow all traffic by default? Wdyt?
• maybe we could offer a more granular config in the helm chart? Eg let users specify the netpol config for the webhook and metrics server?

Btw when you're testing this, be aware that you'll need to install a cni plug-in like calico. Kind clusters don't actually enforce netpols out of the box.

[arryllopez]

We can have a permissive network policy that allows all traffic by defaultt but then users can modify it through Helm values,

in values.yaml it could look something like this

  networkPolicy:
    # -- Enable NetworkPolicy for the controller manager
    enabled: true
    # -- Ingress rules for the status server port
    statusServer:
      # -- Restrict status server to only TrainJob pods
      restrictToTrainJobs: true
    # -- Ingress rules for the webhook port
    webhook:
      # -- Allow all traffic to the webhook port (API server needs access)
      allowAll: true
    # -- Ingress rules for the metrics port
    metrics:
      # -- Allow all traffic to the metrics port (monitoring needs access)
      allowAll: true

and then a new helm template could have the config you mentioned (charts/kubeflow-trainer/templates/manager/network-policy.yaml)

something like this implemented

{{- if .Values.manager.networkPolicy.enabled }}
apiVersion: networking.k8s.io/v1
kind: NetworkPolicy
metadata:
  name: {{ include "trainer.manager.service.name" . }}
  namespace: {{ .Release.Namespace }}
spec:
  podSelector:
    matchLabels:
      {{- include "trainer.manager.selectorLabels" . | nindent 6 }}
  policyTypes:
  - Ingress
  ingress:
  {{- if .Values.manager.networkPolicy.statusServer.restrictToTrainJobs }}
  - from:
    - namespaceSelector: {}
      podSelector:
        matchLabels:
          trainer.kubeflow.org/trainjob-ancestor-step: trainer
    ports:
    - protocol: TCP
      port: 10443
  {{- end }}
  {{- if .Values.manager.networkPolicy.webhook.allowAll }}
  - ports:
    - protocol: TCP
      port: 9443
  {{- end }}
  {{- if .Values.manager.networkPolicy.metrics.allowAll }}
  - ports:
    - protocol: TCP
      port: 8443
  {{- end }}
{{- end }}

WDYT?