Merge pull request #8884 from OpenMined/rasswanth/add-annotations-to-…

…custom-worker-pool Add ability to pass pod Annotations and Labels during Worker Pool Launch
OpenMined · Jun 4, 2024 · bff337a · bff337a
2 parents d9302f3 + 216a9bd
commit bff337a
Show file tree

Hide file tree

Showing 13 changed files with 373 additions and 130 deletions.
diff --git a/notebooks/api/0.8/11-container-images-k8s.ipynb b/notebooks/api/0.8/11-container-images-k8s.ipynb
diff --git a/packages/grid/helm/examples/azure/azure.high.yaml b/packages/grid/helm/examples/azure/azure.high.yaml
@@ -11,9 +11,13 @@ node:
   name: syft-azure
   side: high
   rootEmail: [email protected]
-  defaultWorkerPoolCount: 1
   resourcesPreset: 2xlarge
 
+  defaultWorkerPool:
+    count: 1
+    podLabels: null
+    podAnnotations: null
+
 ingress:
   # Make sure cluster is created with --enable-app-routing
   # az aks create -g group-name -n cluster-name -l region --enable-app-routing

diff --git a/packages/grid/helm/syft/templates/backend/backend-statefulset.yaml b/packages/grid/helm/syft/templates/backend/backend-statefulset.yaml
@@ -72,7 +72,11 @@ spec:
             - name: DEFAULT_WORKER_POOL_IMAGE
               value: "{{ .Values.global.registry }}/openmined/grid-backend:{{ .Values.global.version }}"
             - name: DEFAULT_WORKER_POOL_COUNT
-              value: {{ .Values.node.defaultWorkerPoolCount | quote }}
+              value: {{ .Values.node.defaultWorkerPool.count | quote }}
+            - name: DEFAULT_WORKER_POOL_POD_LABELS
+              value: {{ .Values.node.defaultWorkerPool.podLabels | toJson | quote }}
+            - name: DEFAULT_WORKER_POOL_POD_ANNOTATIONS
+              value: {{ .Values.node.defaultWorkerPool.podAnnotations | toJson | quote }}
             - name: USE_INTERNAL_REGISTRY
               value: {{ .Values.node.useInternalRegistry | quote }}
             {{- if .Values.node.defaultBucketName }}

diff --git a/packages/grid/helm/syft/values.yaml b/packages/grid/helm/syft/values.yaml
@@ -159,15 +159,20 @@ node:
   rootEmail: [email protected]
   type: domain
   side: high
-  inMemoryWorkers: false
-  defaultWorkerPoolCount: 1
   defaultBucketName: null
+  inMemoryWorkers: false
   queuePort: 5556
   logLevel: info
   debuggerEnabled: false
   associationRequestAutoApproval: false
   useInternalRegistry: true
 
+  # Default Worker pool settings
+  defaultWorkerPool:
+    count: 1
+    podLabels: null
+    podAnnotations: null
+
   # SMTP Settings
   smtp:
     host: smtp.sendgrid.net

diff --git a/packages/grid/helm/values.dev.high.yaml b/packages/grid/helm/values.dev.high.yaml
@@ -12,12 +12,16 @@ registry:
 
 node:
   rootEmail: [email protected]
-  defaultWorkerPoolCount: 1
   side: high
 
   resourcesPreset: 2xlarge
   resources: null
 
+  defaultWorkerPool:
+    count: 1
+    podLabels: null
+    podAnnotations: null
+
   secret:
     defaultRootPassword: changethis
 

diff --git a/packages/grid/helm/values.dev.low.yaml b/packages/grid/helm/values.dev.low.yaml
@@ -12,12 +12,16 @@ registry:
 
 node:
   rootEmail: [email protected]
-  defaultWorkerPoolCount: 1
   side: low
 
   resourcesPreset: 2xlarge
   resources: null
 
+  defaultWorkerPool:
+    count: 1
+    podLabels: null
+    podAnnotations: null
+
   secret:
     defaultRootPassword: changethis
 

diff --git a/packages/grid/helm/values.dev.yaml b/packages/grid/helm/values.dev.yaml
@@ -12,12 +12,16 @@ registry:
 
 node:
   rootEmail: [email protected]
-  defaultWorkerPoolCount: 1
   associationRequestAutoApproval: true
 
   resourcesPreset: null
   resources: null
 
+  defaultWorkerPool:
+    count: 1
+    podLabels: null
+    podAnnotations: null
+
   secret:
     defaultRootPassword: changethis
 

diff --git a/packages/syft/src/syft/custom_worker/runner_k8s.py b/packages/syft/src/syft/custom_worker/runner_k8s.py
@@ -31,6 +31,8 @@ def create_pool(
         registry_username: str | None = None,
         registry_password: str | None = None,
         reg_url: str | None = None,
+        pod_annotations: dict[str, str] | None = None,
+        pod_labels: dict[str, str] | None = None,
         **kwargs: Any,
     ) -> StatefulSet:
         try:
@@ -52,6 +54,8 @@ def create_pool(
                 env_vars=env_vars,
                 mount_secrets=mount_secrets,
                 pull_secret=pull_secret,
+                pod_annotations=pod_annotations,
+                pod_labels=pod_labels,
                 **kwargs,
             )
 
@@ -147,6 +151,8 @@ def _create_stateful_set(
         env_vars: list[dict] | None = None,
         mount_secrets: dict | None = None,
         pull_secret: Secret | None = None,
+        pod_annotations: dict[str, str] | None = None,
+        pod_labels: dict[str, str] | None = None,
         **kwargs: Any,
     ) -> StatefulSet:
         """Create a stateful set for a pool"""
@@ -182,6 +188,16 @@ def _create_stateful_set(
                 }
             ]
 
+        default_pod_labels = {
+            "app.kubernetes.io/name": KUBERNETES_NAMESPACE,
+            "app.kubernetes.io/component": pool_name,
+        }
+
+        if isinstance(pod_labels, dict):
+            pod_labels = {**default_pod_labels, **pod_labels}
+        else:
+            pod_labels = default_pod_labels
+
         stateful_set = StatefulSet(
             {
                 "metadata": {
@@ -201,10 +217,8 @@ def _create_stateful_set(
                     },
                     "template": {
                         "metadata": {
-                            "labels": {
-                                "app.kubernetes.io/name": KUBERNETES_NAMESPACE,
-                                "app.kubernetes.io/component": pool_name,
-                            }
+                            "labels": pod_labels,
+                            "annotations": pod_annotations,
                         },
                         "spec": {
                             # TODO: make this configurable

diff --git a/packages/syft/src/syft/node/node.py b/packages/syft/src/syft/node/node.py
@@ -7,6 +7,7 @@
 from datetime import datetime
 from functools import partial
 import hashlib
+import json
 import os
 from pathlib import Path
 import shutil
@@ -217,6 +218,16 @@ def get_default_worker_pool_count(node: Node) -> int:
     )
 
 
+def get_default_worker_pool_pod_annotations() -> dict[str, str] | None:
+    annotations = get_env("DEFAULT_WORKER_POOL_POD_ANNOTATIONS", "null")
+    return json.loads(annotations)
+
+
+def get_default_worker_pool_pod_labels() -> dict[str, str] | None:
+    labels = get_env("DEFAULT_WORKER_POOL_POD_LABELS", "null")
+    return json.loads(labels)
+
+
 def in_kubernetes() -> bool:
     return get_container_host() == "k8s"
 
@@ -1719,6 +1730,8 @@ def create_default_worker_pool(node: Node) -> SyftError | None:
     default_pool_name = node.settings.default_worker_pool
     default_worker_pool = node.get_default_worker_pool()
     default_worker_tag = get_default_worker_tag_by_env(node.dev_mode)
+    default_worker_pool_pod_annotations = get_default_worker_pool_pod_annotations()
+    default_worker_pool_pod_labels = get_default_worker_pool_pod_labels()
     worker_count = get_default_worker_pool_count(node)
     context = AuthedServiceContext(
         node=node,
@@ -1775,6 +1788,8 @@ def create_default_worker_pool(node: Node) -> SyftError | None:
             pool_name=default_pool_name,
             image_uid=default_image.id,
             num_workers=worker_count,
+            pod_annotations=default_worker_pool_pod_annotations,
+            pod_labels=default_worker_pool_pod_labels,
         )
     else:
         # Else add a worker to existing worker pool

diff --git a/packages/syft/src/syft/protocol/protocol_version.json b/packages/syft/src/syft/protocol/protocol_version.json
@@ -247,6 +247,13 @@
           "action": "remove"
         }
       },
+      "CreateCustomWorkerPoolChange": {
+        "3": {
+          "version": 3,
+          "hash": "e982f2ebcdc6fe23a65a014109e33ba7c487bb7ca5623723cf5ec7642f86828c",
+          "action": "add"
+        }
+      },
       "NodePeerUpdate": {
         "1": {
           "version": 1,

diff --git a/packages/syft/src/syft/service/request/request.py b/packages/syft/src/syft/service/request/request.py
@@ -298,12 +298,14 @@ def __repr_syft_nested__(self) -> str:
 @serializable()
 class CreateCustomWorkerPoolChange(Change):
     __canonical_name__ = "CreateCustomWorkerPoolChange"
-    __version__ = SYFT_OBJECT_VERSION_2
+    __version__ = SYFT_OBJECT_VERSION_3
 
     pool_name: str
     num_workers: int
     image_uid: UID | None = None
     config: WorkerConfig | None = None
+    pod_annotations: dict[str, str] | None = None
+    pod_labels: dict[str, str] | None = None
 
     __repr_attrs__ = ["pool_name", "num_workers", "image_uid"]
 
@@ -337,6 +339,8 @@ def _run(
                 num_workers=self.num_workers,
                 registry_username=context.extra_kwargs.get("registry_username", None),
                 registry_password=context.extra_kwargs.get("registry_password", None),
+                pod_annotations=self.pod_annotations,
+                pod_labels=self.pod_labels,
             )
             if isinstance(result, SyftError):
                 return Err(result)
@@ -361,6 +365,19 @@ def __repr_syft_nested__(self) -> str:
         )
 
 
+@serializable()
+class CreateCustomWorkerPoolChangeV2(Change):
+    __canonical_name__ = "CreateCustomWorkerPoolChange"
+    __version__ = SYFT_OBJECT_VERSION_2
+
+    pool_name: str
+    num_workers: int
+    image_uid: UID | None = None
+    config: WorkerConfig | None = None
+
+    __repr_attrs__ = ["pool_name", "num_workers", "image_uid"]
+
+
 @serializable()
 class Request(SyncableSyftObject):
     __canonical_name__ = "Request"

diff --git a/packages/syft/src/syft/service/worker/utils.py b/packages/syft/src/syft/service/worker/utils.py
@@ -328,6 +328,8 @@ def create_kubernetes_pool(
     registry_username: str | None = None,
     registry_password: str | None = None,
     reg_url: str | None = None,
+    pod_annotations: dict[str, str] | None = None,
+    pod_labels: dict[str, str] | None = None,
     **kwargs: Any,
 ) -> list[Pod] | SyftError:
     pool = None
@@ -363,6 +365,8 @@ def create_kubernetes_pool(
             registry_username=registry_username,
             registry_password=registry_password,
             reg_url=reg_url,
+            pod_annotations=pod_annotations,
+            pod_labels=pod_labels,
         )
     except Exception as e:
         if pool:
@@ -405,6 +409,8 @@ def run_workers_in_kubernetes(
     registry_username: str | None = None,
     registry_password: str | None = None,
     reg_url: str | None = None,
+    pod_annotations: dict[str, str] | None = None,
+    pod_labels: dict[str, str] | None = None,
     **kwargs: Any,
 ) -> list[ContainerSpawnStatus] | SyftError:
     spawn_status = []
@@ -422,6 +428,8 @@ def run_workers_in_kubernetes(
                 registry_username=registry_username,
                 registry_password=registry_password,
                 reg_url=reg_url,
+                pod_annotations=pod_annotations,
+                pod_labels=pod_labels,
             )
         else:
             return SyftError(
@@ -504,6 +512,8 @@ def run_containers(
     registry_username: str | None = None,
     registry_password: str | None = None,
     reg_url: str | None = None,
+    pod_annotations: dict[str, str] | None = None,
+    pod_labels: dict[str, str] | None = None,
 ) -> list[ContainerSpawnStatus] | SyftError:
     results = []
 
@@ -540,6 +550,8 @@ def run_containers(
             registry_username=registry_username,
             registry_password=registry_password,
             reg_url=reg_url,
+            pod_annotations=pod_annotations,
+            pod_labels=pod_labels,
         )
 
     return results