From 42ff7fedef9cae184236e474d967fc6ef793b051 Mon Sep 17 00:00:00 2001 From: rasswanth-s <43314053+rasswanth-s@users.noreply.github.com> Date: Mon, 3 Jun 2024 11:11:03 +0530 Subject: [PATCH 1/7] add ability to pass pod annotations during worker pool launch --- .../helm/syft/templates/backend/backend-statefulset.yaml | 2 ++ packages/grid/helm/syft/values.yaml | 7 +++++-- packages/syft/src/syft/custom_worker/runner_k8s.py | 6 +++++- packages/syft/src/syft/node/node.py | 8 ++++++++ packages/syft/src/syft/service/worker/utils.py | 6 ++++++ .../syft/src/syft/service/worker/worker_pool_service.py | 4 ++++ 6 files changed, 30 insertions(+), 3 deletions(-) diff --git a/packages/grid/helm/syft/templates/backend/backend-statefulset.yaml b/packages/grid/helm/syft/templates/backend/backend-statefulset.yaml index 8048f262e5e..90614f79c14 100644 --- a/packages/grid/helm/syft/templates/backend/backend-statefulset.yaml +++ b/packages/grid/helm/syft/templates/backend/backend-statefulset.yaml @@ -73,6 +73,8 @@ spec: value: "{{ .Values.global.registry }}/openmined/grid-backend:{{ .Values.global.version }}" - name: DEFAULT_WORKER_POOL_COUNT value: {{ .Values.node.defaultWorkerPoolCount | quote }} + - name: DEFAULT_WORKER_POOL_POD_ANNOTATIONS + value: {{ .Values.node.defaultWorkerPoolPodAnnotations | 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 index e44d2cf7982..28484b96427 100644 --- a/packages/grid/helm/syft/values.yaml +++ b/packages/grid/helm/syft/values.yaml @@ -159,8 +159,6 @@ node: rootEmail: info@openmined.org type: domain side: high - inMemoryWorkers: false - defaultWorkerPoolCount: 1 defaultBucketName: null queuePort: 5556 logLevel: info @@ -168,6 +166,11 @@ node: associationRequestAutoApproval: false useInternalRegistry: true + # Worker pool settings + inMemoryWorkers: false + defaultWorkerPoolCount: 1 + defaultWorkerPoolPodAnnotations: null + # SMTP Settings smtp: host: smtp.sendgrid.net diff --git a/packages/syft/src/syft/custom_worker/runner_k8s.py b/packages/syft/src/syft/custom_worker/runner_k8s.py index f9c48c6e394..ae22d1e86c3 100644 --- a/packages/syft/src/syft/custom_worker/runner_k8s.py +++ b/packages/syft/src/syft/custom_worker/runner_k8s.py @@ -31,6 +31,7 @@ def create_pool( reg_username: str | None = None, reg_password: str | None = None, reg_url: str | None = None, + pod_annotations: dict[str, str] | None = None, **kwargs: Any, ) -> StatefulSet: try: @@ -52,6 +53,7 @@ def create_pool( env_vars=env_vars, mount_secrets=mount_secrets, pull_secret=pull_secret, + pod_annotations=pod_annotations, **kwargs, ) @@ -147,6 +149,7 @@ 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, **kwargs: Any, ) -> StatefulSet: """Create a stateful set for a pool""" @@ -204,7 +207,8 @@ def _create_stateful_set( "labels": { "app.kubernetes.io/name": KUBERNETES_NAMESPACE, "app.kubernetes.io/component": pool_name, - } + }, + "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 index a4a0fc09de6..93447d90423 100644 --- 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,11 @@ 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 in_kubernetes() -> bool: return get_container_host() == "k8s" @@ -1719,6 +1725,7 @@ 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() worker_count = get_default_worker_pool_count(node) context = AuthedServiceContext( node=node, @@ -1775,6 +1782,7 @@ def create_default_worker_pool(node: Node) -> SyftError | None: name=default_pool_name, image_uid=default_image.id, num_workers=worker_count, + pod_annotations=default_worker_pool_pod_annotations, ) else: # Else add a worker to existing worker pool diff --git a/packages/syft/src/syft/service/worker/utils.py b/packages/syft/src/syft/service/worker/utils.py index 181774da837..e1445c45538 100644 --- a/packages/syft/src/syft/service/worker/utils.py +++ b/packages/syft/src/syft/service/worker/utils.py @@ -328,6 +328,7 @@ def create_kubernetes_pool( reg_username: str | None = None, reg_password: str | None = None, reg_url: str | None = None, + pod_annotations: dict[str, str] | None = None, **kwargs: Any, ) -> list[Pod] | SyftError: pool = None @@ -363,6 +364,7 @@ def create_kubernetes_pool( reg_username=reg_username, reg_password=reg_password, reg_url=reg_url, + pod_annotations=pod_annotations, ) except Exception as e: if pool: @@ -405,6 +407,7 @@ def run_workers_in_kubernetes( reg_username: str | None = None, reg_password: str | None = None, reg_url: str | None = None, + pod_annotations: dict[str, str] | None = None, **kwargs: Any, ) -> list[ContainerSpawnStatus] | SyftError: spawn_status = [] @@ -422,6 +425,7 @@ def run_workers_in_kubernetes( reg_username=reg_username, reg_password=reg_password, reg_url=reg_url, + pod_annotations=pod_annotations, ) else: return SyftError( @@ -504,6 +508,7 @@ def run_containers( reg_username: str | None = None, reg_password: str | None = None, reg_url: str | None = None, + pod_annotations: dict[str, str] | None = None, ) -> list[ContainerSpawnStatus] | SyftError: results = [] @@ -540,6 +545,7 @@ def run_containers( reg_username=reg_username, reg_password=reg_password, reg_url=reg_url, + pod_annotations=pod_annotations, ) return results diff --git a/packages/syft/src/syft/service/worker/worker_pool_service.py b/packages/syft/src/syft/service/worker/worker_pool_service.py index cef16fb5b72..38d0e155619 100644 --- a/packages/syft/src/syft/service/worker/worker_pool_service.py +++ b/packages/syft/src/syft/service/worker/worker_pool_service.py @@ -69,6 +69,7 @@ def launch( num_workers: int, reg_username: str | None = None, reg_password: str | None = None, + pod_annotations: dict[str, str] | None = None, ) -> list[ContainerSpawnStatus] | SyftError: """Creates a pool of workers from the given SyftWorkerImage. @@ -126,6 +127,7 @@ def launch( worker_stash=worker_stash, reg_username=reg_username, reg_password=reg_password, + pod_annotations=pod_annotations, ) if isinstance(result, SyftError): @@ -656,6 +658,7 @@ def _create_workers_in_pool( worker_stash: WorkerStash, reg_username: str | None = None, reg_password: str | None = None, + pod_annotations: dict[str, str] | None = None, ) -> tuple[list[LinkedObject], list[ContainerSpawnStatus]] | SyftError: queue_port = context.node.queue_config.client_config.queue_port @@ -687,6 +690,7 @@ def _create_workers_in_pool( reg_username=reg_username, reg_password=reg_password, reg_url=registry_host, + pod_annotations=pod_annotations, ) if isinstance(result, SyftError): return result From 9df8a7047b6db53a372dfe19efb193157ae68366 Mon Sep 17 00:00:00 2001 From: rasswanth-s <43314053+rasswanth-s@users.noreply.github.com> Date: Mon, 3 Jun 2024 12:23:40 +0530 Subject: [PATCH 2/7] add tests for pod annotations add pod annotations to pool creation request --- .../api/0.8/11-container-images-k8s.ipynb | 331 +++++++++++------- .../src/syft/protocol/protocol_version.json | 7 + .../syft/src/syft/service/request/request.py | 17 +- .../service/worker/worker_pool_service.py | 7 + 4 files changed, 243 insertions(+), 119 deletions(-) diff --git a/notebooks/api/0.8/11-container-images-k8s.ipynb b/notebooks/api/0.8/11-container-images-k8s.ipynb index 24f4a7c33df..a0103b55077 100644 --- a/notebooks/api/0.8/11-container-images-k8s.ipynb +++ b/notebooks/api/0.8/11-container-images-k8s.ipynb @@ -22,6 +22,7 @@ "import os\n", "\n", "# third party\n", + "import kr8s\n", "import numpy as np\n", "import requests\n", "\n", @@ -44,6 +45,34 @@ "id": "2", "metadata": {}, "outputs": [], + "source": [ + "def get_kr8s_client():\n", + " return kr8s.api(namespace=\"syft\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3", + "metadata": {}, + "outputs": [], + "source": [ + "def get_statefulset_by_pool_name(pool_name):\n", + " kr8s_client = get_kr8s_client()\n", + " pool_list = kr8s_client.get(\n", + " \"statefulsets\", label_selector={\"app.kubernetes.io/component\": pool_name}\n", + " )\n", + " if len(pool_list) == 0:\n", + " return None\n", + " return pool_list[0]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4", + "metadata": {}, + "outputs": [], "source": [ "os.environ[\"ORCHESTRA_DEPLOYMENT_TYPE\"] = \"remote\"\n", "os.environ[\"DEV_MODE\"] = \"True\"\n", @@ -56,7 +85,7 @@ { "cell_type": "code", "execution_count": null, - "id": "3", + "id": "5", "metadata": {}, "outputs": [], "source": [ @@ -69,7 +98,7 @@ { "cell_type": "code", "execution_count": null, - "id": "4", + "id": "6", "metadata": {}, "outputs": [], "source": [ @@ -79,7 +108,7 @@ }, { "cell_type": "markdown", - "id": "5", + "id": "7", "metadata": {}, "source": [ "### Scaling Default Worker Pool" @@ -87,7 +116,7 @@ }, { "cell_type": "markdown", - "id": "6", + "id": "8", "metadata": {}, "source": [ "We should see a default worker pool" @@ -96,7 +125,7 @@ { "cell_type": "code", "execution_count": null, - "id": "7", + "id": "9", "metadata": {}, "outputs": [], "source": [ @@ -105,7 +134,7 @@ }, { "cell_type": "markdown", - "id": "8", + "id": "10", "metadata": {}, "source": [ "Scale up to 3 workers" @@ -114,7 +143,7 @@ { "cell_type": "code", "execution_count": null, - "id": "9", + "id": "11", "metadata": {}, "outputs": [], "source": [ @@ -128,7 +157,7 @@ { "cell_type": "code", "execution_count": null, - "id": "10", + "id": "12", "metadata": {}, "outputs": [], "source": [ @@ -140,7 +169,7 @@ { "cell_type": "code", "execution_count": null, - "id": "11", + "id": "13", "metadata": {}, "outputs": [], "source": [ @@ -153,7 +182,7 @@ }, { "cell_type": "markdown", - "id": "12", + "id": "14", "metadata": {}, "source": [ "Scale down to 1 worker" @@ -162,7 +191,7 @@ { "cell_type": "code", "execution_count": null, - "id": "13", + "id": "15", "metadata": {}, "outputs": [], "source": [ @@ -176,7 +205,7 @@ { "cell_type": "code", "execution_count": null, - "id": "14", + "id": "16", "metadata": {}, "outputs": [], "source": [ @@ -188,7 +217,7 @@ { "cell_type": "code", "execution_count": null, - "id": "15", + "id": "17", "metadata": {}, "outputs": [], "source": [ @@ -200,7 +229,7 @@ }, { "cell_type": "markdown", - "id": "16", + "id": "18", "metadata": {}, "source": [ "#### Submit Dockerfile" @@ -209,7 +238,7 @@ { "cell_type": "code", "execution_count": null, - "id": "17", + "id": "19", "metadata": {}, "outputs": [], "source": [ @@ -226,7 +255,7 @@ { "cell_type": "code", "execution_count": null, - "id": "18", + "id": "20", "metadata": {}, "outputs": [], "source": [ @@ -241,7 +270,7 @@ { "cell_type": "code", "execution_count": null, - "id": "19", + "id": "21", "metadata": {}, "outputs": [], "source": [ @@ -251,7 +280,7 @@ { "cell_type": "code", "execution_count": null, - "id": "20", + "id": "22", "metadata": {}, "outputs": [], "source": [ @@ -261,7 +290,7 @@ { "cell_type": "code", "execution_count": null, - "id": "21", + "id": "23", "metadata": {}, "outputs": [], "source": [ @@ -274,7 +303,7 @@ { "cell_type": "code", "execution_count": null, - "id": "22", + "id": "24", "metadata": {}, "outputs": [], "source": [ @@ -284,7 +313,7 @@ { "cell_type": "code", "execution_count": null, - "id": "23", + "id": "25", "metadata": {}, "outputs": [], "source": [ @@ -295,7 +324,7 @@ { "cell_type": "code", "execution_count": null, - "id": "24", + "id": "26", "metadata": {}, "outputs": [], "source": [ @@ -306,7 +335,7 @@ { "cell_type": "code", "execution_count": null, - "id": "25", + "id": "27", "metadata": {}, "outputs": [], "source": [ @@ -325,7 +354,7 @@ }, { "cell_type": "markdown", - "id": "26", + "id": "28", "metadata": {}, "source": [ "#### Add External Registry in Syft" @@ -334,7 +363,7 @@ { "cell_type": "code", "execution_count": null, - "id": "27", + "id": "29", "metadata": {}, "outputs": [], "source": [ @@ -350,7 +379,7 @@ { "cell_type": "code", "execution_count": null, - "id": "28", + "id": "30", "metadata": {}, "outputs": [], "source": [ @@ -361,7 +390,7 @@ { "cell_type": "code", "execution_count": null, - "id": "29", + "id": "31", "metadata": {}, "outputs": [], "source": [ @@ -371,7 +400,7 @@ { "cell_type": "code", "execution_count": null, - "id": "30", + "id": "32", "metadata": {}, "outputs": [], "source": [ @@ -382,7 +411,7 @@ { "cell_type": "code", "execution_count": null, - "id": "31", + "id": "33", "metadata": {}, "outputs": [], "source": [ @@ -393,7 +422,7 @@ { "cell_type": "code", "execution_count": null, - "id": "32", + "id": "34", "metadata": {}, "outputs": [], "source": [ @@ -404,7 +433,7 @@ { "cell_type": "code", "execution_count": null, - "id": "33", + "id": "35", "metadata": {}, "outputs": [], "source": [ @@ -414,7 +443,7 @@ { "cell_type": "code", "execution_count": null, - "id": "34", + "id": "36", "metadata": {}, "outputs": [], "source": [ @@ -423,7 +452,7 @@ }, { "cell_type": "markdown", - "id": "35", + "id": "37", "metadata": {}, "source": [ "#### Build Image" @@ -432,7 +461,7 @@ { "cell_type": "code", "execution_count": null, - "id": "36", + "id": "38", "metadata": {}, "outputs": [], "source": [ @@ -450,7 +479,7 @@ { "cell_type": "code", "execution_count": null, - "id": "37", + "id": "39", "metadata": {}, "outputs": [], "source": [ @@ -460,7 +489,7 @@ { "cell_type": "code", "execution_count": null, - "id": "38", + "id": "40", "metadata": {}, "outputs": [], "source": [ @@ -471,7 +500,7 @@ { "cell_type": "code", "execution_count": null, - "id": "39", + "id": "41", "metadata": {}, "outputs": [], "source": [ @@ -483,7 +512,7 @@ { "cell_type": "code", "execution_count": null, - "id": "40", + "id": "42", "metadata": {}, "outputs": [], "source": [ @@ -496,7 +525,7 @@ }, { "cell_type": "markdown", - "id": "41", + "id": "43", "metadata": {}, "source": [ "#### Push Image to Local Registry" @@ -505,7 +534,7 @@ { "cell_type": "code", "execution_count": null, - "id": "42", + "id": "44", "metadata": {}, "outputs": [], "source": [ @@ -521,7 +550,7 @@ { "cell_type": "code", "execution_count": null, - "id": "43", + "id": "45", "metadata": {}, "outputs": [], "source": [ @@ -531,7 +560,7 @@ { "cell_type": "code", "execution_count": null, - "id": "44", + "id": "46", "metadata": {}, "outputs": [], "source": [ @@ -552,7 +581,7 @@ }, { "cell_type": "markdown", - "id": "45", + "id": "47", "metadata": {}, "source": [ "#### Create Worker Pool From Image" @@ -561,24 +590,26 @@ { "cell_type": "code", "execution_count": null, - "id": "46", + "id": "48", "metadata": {}, "outputs": [], "source": [ "worker_pool_name = \"custom-pool\"\n", + "custom_pool_annotations = {\"test-custom-pool\": \"Test annotation for custom pool\"}\n", "worker_pool_res = domain_client.api.services.worker_pool.launch(\n", " name=worker_pool_name,\n", " image_uid=workerimage.id,\n", " num_workers=3,\n", " reg_username=external_registry_username,\n", " reg_password=external_registry_password,\n", + " pod_annotations=custom_pool_annotations,\n", ")" ] }, { "cell_type": "code", "execution_count": null, - "id": "47", + "id": "49", "metadata": {}, "outputs": [], "source": [ @@ -589,7 +620,7 @@ { "cell_type": "code", "execution_count": null, - "id": "48", + "id": "50", "metadata": {}, "outputs": [], "source": [ @@ -600,7 +631,7 @@ { "cell_type": "code", "execution_count": null, - "id": "49", + "id": "51", "metadata": {}, "outputs": [], "source": [ @@ -611,7 +642,26 @@ { "cell_type": "code", "execution_count": null, - "id": "50", + "id": "52", + "metadata": {}, + "outputs": [], + "source": [ + "# check annotations for custom pool\n", + "custom_pool_statefulset = get_statefulset_by_pool_name(worker_pool_name)\n", + "assert custom_pool_statefulset is not None, \"Custom pool statefulset not found\"\n", + "custom_pool_pod_metadata = custom_pool_statefulset.spec.template.metadata\n", + "assert (\n", + " \"annotations\" in custom_pool_pod_metadata\n", + "), \"Annotations not found in custom pool pod metadata\"\n", + "assert (\n", + " custom_pool_pod_metadata.annotations == custom_pool_annotations\n", + "), \"Annotations do not match in Custom pool pod metadata\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "53", "metadata": {}, "outputs": [], "source": [ @@ -622,7 +672,7 @@ { "cell_type": "code", "execution_count": null, - "id": "51", + "id": "54", "metadata": {}, "outputs": [], "source": [ @@ -640,7 +690,7 @@ { "cell_type": "code", "execution_count": null, - "id": "52", + "id": "55", "metadata": {}, "outputs": [], "source": [ @@ -654,7 +704,7 @@ { "cell_type": "code", "execution_count": null, - "id": "53", + "id": "56", "metadata": {}, "outputs": [], "source": [ @@ -664,7 +714,7 @@ { "cell_type": "code", "execution_count": null, - "id": "54", + "id": "57", "metadata": {}, "outputs": [], "source": [ @@ -674,7 +724,7 @@ }, { "cell_type": "markdown", - "id": "55", + "id": "58", "metadata": {}, "source": [ "#### Get Worker Logs" @@ -683,7 +733,7 @@ { "cell_type": "code", "execution_count": null, - "id": "56", + "id": "59", "metadata": {}, "outputs": [], "source": [ @@ -696,7 +746,7 @@ { "cell_type": "code", "execution_count": null, - "id": "57", + "id": "60", "metadata": {}, "outputs": [], "source": [ @@ -706,7 +756,7 @@ { "cell_type": "code", "execution_count": null, - "id": "58", + "id": "61", "metadata": {}, "outputs": [], "source": [ @@ -715,7 +765,7 @@ }, { "cell_type": "markdown", - "id": "59", + "id": "62", "metadata": {}, "source": [ "### Syft function" @@ -724,7 +774,7 @@ { "cell_type": "code", "execution_count": null, - "id": "60", + "id": "63", "metadata": {}, "outputs": [], "source": [ @@ -738,7 +788,7 @@ { "cell_type": "code", "execution_count": null, - "id": "61", + "id": "64", "metadata": {}, "outputs": [], "source": [ @@ -754,7 +804,7 @@ { "cell_type": "code", "execution_count": null, - "id": "62", + "id": "65", "metadata": {}, "outputs": [], "source": [ @@ -764,7 +814,7 @@ { "cell_type": "code", "execution_count": null, - "id": "63", + "id": "66", "metadata": {}, "outputs": [], "source": [ @@ -774,7 +824,7 @@ { "cell_type": "code", "execution_count": null, - "id": "64", + "id": "67", "metadata": {}, "outputs": [], "source": [ @@ -785,7 +835,7 @@ { "cell_type": "code", "execution_count": null, - "id": "65", + "id": "68", "metadata": {}, "outputs": [], "source": [ @@ -795,7 +845,7 @@ { "cell_type": "code", "execution_count": null, - "id": "66", + "id": "69", "metadata": {}, "outputs": [], "source": [ @@ -806,7 +856,7 @@ { "cell_type": "code", "execution_count": null, - "id": "67", + "id": "70", "metadata": {}, "outputs": [], "source": [ @@ -817,7 +867,7 @@ { "cell_type": "code", "execution_count": null, - "id": "68", + "id": "71", "metadata": {}, "outputs": [], "source": [ @@ -827,7 +877,7 @@ { "cell_type": "code", "execution_count": null, - "id": "69", + "id": "72", "metadata": {}, "outputs": [], "source": [ @@ -837,7 +887,7 @@ { "cell_type": "code", "execution_count": null, - "id": "70", + "id": "73", "metadata": {}, "outputs": [], "source": [ @@ -847,7 +897,7 @@ { "cell_type": "code", "execution_count": null, - "id": "71", + "id": "74", "metadata": {}, "outputs": [], "source": [ @@ -857,7 +907,7 @@ { "cell_type": "code", "execution_count": null, - "id": "72", + "id": "75", "metadata": {}, "outputs": [], "source": [ @@ -868,7 +918,7 @@ { "cell_type": "code", "execution_count": null, - "id": "73", + "id": "76", "metadata": {}, "outputs": [], "source": [ @@ -881,7 +931,7 @@ { "cell_type": "code", "execution_count": null, - "id": "74", + "id": "77", "metadata": {}, "outputs": [], "source": [ @@ -896,7 +946,7 @@ { "cell_type": "code", "execution_count": null, - "id": "75", + "id": "78", "metadata": {}, "outputs": [], "source": [ @@ -905,7 +955,7 @@ }, { "cell_type": "markdown", - "id": "76", + "id": "79", "metadata": {}, "source": [ "#### Worker Pool and Image Creation Request/Approval" @@ -914,7 +964,7 @@ { "cell_type": "code", "execution_count": null, - "id": "77", + "id": "80", "metadata": {}, "outputs": [], "source": [ @@ -930,7 +980,7 @@ { "cell_type": "code", "execution_count": null, - "id": "78", + "id": "81", "metadata": {}, "outputs": [], "source": [ @@ -944,7 +994,7 @@ { "cell_type": "code", "execution_count": null, - "id": "79", + "id": "82", "metadata": {}, "outputs": [], "source": [ @@ -954,7 +1004,7 @@ { "cell_type": "code", "execution_count": null, - "id": "80", + "id": "83", "metadata": {}, "outputs": [], "source": [ @@ -965,7 +1015,7 @@ { "cell_type": "code", "execution_count": null, - "id": "81", + "id": "84", "metadata": {}, "outputs": [], "source": [ @@ -978,7 +1028,7 @@ }, { "cell_type": "markdown", - "id": "82", + "id": "85", "metadata": {}, "source": [ "##### Build image first then create pool" @@ -987,7 +1037,7 @@ { "cell_type": "code", "execution_count": null, - "id": "83", + "id": "86", "metadata": {}, "outputs": [], "source": [ @@ -1005,7 +1055,7 @@ { "cell_type": "code", "execution_count": null, - "id": "84", + "id": "87", "metadata": {}, "outputs": [], "source": [ @@ -1015,7 +1065,7 @@ { "cell_type": "code", "execution_count": null, - "id": "85", + "id": "88", "metadata": {}, "outputs": [], "source": [ @@ -1026,7 +1076,7 @@ { "cell_type": "code", "execution_count": null, - "id": "86", + "id": "89", "metadata": {}, "outputs": [], "source": [ @@ -1049,7 +1099,7 @@ { "cell_type": "code", "execution_count": null, - "id": "87", + "id": "90", "metadata": {}, "outputs": [], "source": [ @@ -1066,13 +1116,17 @@ { "cell_type": "code", "execution_count": null, - "id": "88", + "id": "91", "metadata": {}, "outputs": [], "source": [ "pool_name_opendp = \"opendp-pool\"\n", + "opendp_pod_annotations = {\"test-opendp-pool\": \"Test annotation for opendp pool\"}\n", "pool_create_request = domain_client.api.services.worker_pool.pool_creation_request(\n", - " pool_name=pool_name_opendp, num_workers=3, image_uid=workerimage_opendp.id\n", + " pool_name=pool_name_opendp,\n", + " num_workers=3,\n", + " image_uid=workerimage_opendp.id,\n", + " pod_annotations=opendp_pod_annotations,\n", ")\n", "pool_create_request" ] @@ -1080,7 +1134,7 @@ { "cell_type": "code", "execution_count": null, - "id": "89", + "id": "92", "metadata": {}, "outputs": [], "source": [ @@ -1091,7 +1145,7 @@ { "cell_type": "code", "execution_count": null, - "id": "90", + "id": "93", "metadata": {}, "outputs": [], "source": [ @@ -1105,7 +1159,7 @@ { "cell_type": "code", "execution_count": null, - "id": "91", + "id": "94", "metadata": {}, "outputs": [], "source": [ @@ -1115,7 +1169,7 @@ { "cell_type": "code", "execution_count": null, - "id": "92", + "id": "95", "metadata": {}, "outputs": [], "source": [ @@ -1127,7 +1181,7 @@ { "cell_type": "code", "execution_count": null, - "id": "93", + "id": "96", "metadata": {}, "outputs": [], "source": [ @@ -1140,7 +1194,26 @@ { "cell_type": "code", "execution_count": null, - "id": "94", + "id": "97", + "metadata": {}, + "outputs": [], + "source": [ + "# check annotations for open dp pool\n", + "opendp_pool_statefulset = get_statefulset_by_pool_name(pool_name_opendp)\n", + "assert opendp_pool_statefulset is not None, \"Open DP pool statefulset not found\"\n", + "opendp_pool_pod_metadata = opendp_pool_statefulset.spec.template.metadata\n", + "assert (\n", + " \"annotations\" in opendp_pool_pod_metadata\n", + "), \"Annotations not found in opendp pool pod metadata\"\n", + "assert (\n", + " opendp_pool_pod_metadata.annotations == opendp_pod_annotations\n", + "), \"Annotations do not match in opendp pool pod metadata\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "98", "metadata": {}, "outputs": [], "source": [ @@ -1155,7 +1228,7 @@ { "cell_type": "code", "execution_count": null, - "id": "95", + "id": "99", "metadata": {}, "outputs": [], "source": [ @@ -1164,7 +1237,7 @@ }, { "cell_type": "markdown", - "id": "96", + "id": "100", "metadata": {}, "source": [ "Request to build the image and create the pool at the same time" @@ -1173,7 +1246,7 @@ { "cell_type": "code", "execution_count": null, - "id": "97", + "id": "101", "metadata": {}, "outputs": [], "source": [ @@ -1191,21 +1264,22 @@ { "cell_type": "code", "execution_count": null, - "id": "98", + "id": "102", "metadata": {}, "outputs": [], "source": [ "pool_name_recordlinkage = \"recordlinkage-pool\"\n", - "\n", - "pool_image_create_request = (\n", - " domain_client.api.services.worker_pool.create_image_and_pool_request(\n", - " pool_name=pool_name_recordlinkage,\n", - " num_workers=2,\n", - " tag=docker_tag_recordlinkage,\n", - " config=docker_config_recordlinkage,\n", - " registry_uid=registry_uid,\n", - " reason=\"I want to do some more cool data science with PySyft and OpenDP\",\n", - " )\n", + "recordlinkage_pod_annotations = {\n", + " \"test-recordlinkage-pool\": \"Test annotation for recordlinkage pool\"\n", + "}\n", + "pool_image_create_request = domain_client.api.services.worker_pool.create_image_and_pool_request(\n", + " pool_name=pool_name_recordlinkage,\n", + " num_workers=2,\n", + " tag=docker_tag_recordlinkage,\n", + " config=docker_config_recordlinkage,\n", + " registry_uid=registry_uid,\n", + " reason=\"I want to do some more cool data science with PySyft and RecordLinkage!\",\n", + " pod_annotations=recordlinkage_pod_annotations,\n", ")\n", "pool_image_create_request" ] @@ -1213,7 +1287,7 @@ { "cell_type": "code", "execution_count": null, - "id": "99", + "id": "103", "metadata": {}, "outputs": [], "source": [ @@ -1225,7 +1299,7 @@ { "cell_type": "code", "execution_count": null, - "id": "100", + "id": "104", "metadata": {}, "outputs": [], "source": [ @@ -1238,7 +1312,7 @@ { "cell_type": "code", "execution_count": null, - "id": "101", + "id": "105", "metadata": {}, "outputs": [], "source": [ @@ -1251,7 +1325,7 @@ { "cell_type": "code", "execution_count": null, - "id": "102", + "id": "106", "metadata": {}, "outputs": [], "source": [ @@ -1261,7 +1335,7 @@ { "cell_type": "code", "execution_count": null, - "id": "103", + "id": "107", "metadata": {}, "outputs": [], "source": [ @@ -1272,7 +1346,7 @@ { "cell_type": "code", "execution_count": null, - "id": "104", + "id": "108", "metadata": {}, "outputs": [], "source": [ @@ -1287,7 +1361,28 @@ { "cell_type": "code", "execution_count": null, - "id": "105", + "id": "109", + "metadata": {}, + "outputs": [], + "source": [ + "# check annotations for recordlinkage pool\n", + "recordlinkage_pool_statefulset = get_statefulset_by_pool_name(pool_name_recordlinkage)\n", + "assert (\n", + " recordlinkage_pool_statefulset is not None\n", + "), \"RecordLinkage pool statefulset not found\"\n", + "recordlinkage_pool_pod_metadata = recordlinkage_pool_statefulset.spec.template.metadata\n", + "assert (\n", + " \"annotations\" in recordlinkage_pool_pod_metadata\n", + "), \"Annotations not found in recordlinkage pool pod metadata\"\n", + "assert (\n", + " recordlinkage_pool_pod_metadata.annotations == recordlinkage_pod_annotations\n", + "), \"Annotations do not match in recordlinkage pool pod metadata\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "110", "metadata": {}, "outputs": [], "source": [ @@ -1297,7 +1392,7 @@ { "cell_type": "code", "execution_count": null, - "id": "106", + "id": "111", "metadata": {}, "outputs": [], "source": [ @@ -1315,7 +1410,7 @@ { "cell_type": "code", "execution_count": null, - "id": "107", + "id": "112", "metadata": {}, "outputs": [], "source": [ @@ -1326,7 +1421,7 @@ { "cell_type": "code", "execution_count": null, - "id": "108", + "id": "113", "metadata": {}, "outputs": [], "source": [ @@ -1343,7 +1438,7 @@ { "cell_type": "code", "execution_count": null, - "id": "109", + "id": "114", "metadata": {}, "outputs": [], "source": [ diff --git a/packages/syft/src/syft/protocol/protocol_version.json b/packages/syft/src/syft/protocol/protocol_version.json index 233e7c3c355..6c177ab6275 100644 --- a/packages/syft/src/syft/protocol/protocol_version.json +++ b/packages/syft/src/syft/protocol/protocol_version.json @@ -246,6 +246,13 @@ "hash": "c1796e7b01c9eae0dbf59cfd5c2c2f0e7eba593e0cea615717246572b27aae4b", "action": "remove" } + }, + "CreateCustomWorkerPoolChange": { + "3": { + "version": 3, + "hash": "1c92c173a9eac781377fe5337262ae131111c2290a3e543b17bce2842bd037f8", + "action": "add" + } } } } diff --git a/packages/syft/src/syft/service/request/request.py b/packages/syft/src/syft/service/request/request.py index 3afbafee914..6c1d1b97f73 100644 --- a/packages/syft/src/syft/service/request/request.py +++ b/packages/syft/src/syft/service/request/request.py @@ -298,12 +298,13 @@ 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 __repr_attrs__ = ["pool_name", "num_workers", "image_uid"] @@ -337,6 +338,7 @@ def _run( num_workers=self.num_workers, reg_username=context.extra_kwargs.get("reg_username", None), reg_password=context.extra_kwargs.get("reg_password", None), + pod_annotations=self.pod_annotations, ) if isinstance(result, SyftError): return Err(result) @@ -361,6 +363,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/worker_pool_service.py b/packages/syft/src/syft/service/worker/worker_pool_service.py index 38d0e155619..c2ba8b42684 100644 --- a/packages/syft/src/syft/service/worker/worker_pool_service.py +++ b/packages/syft/src/syft/service/worker/worker_pool_service.py @@ -163,6 +163,7 @@ def create_pool_request( num_workers: int, image_uid: UID, reason: str | None = "", + pod_annotations: dict[str, str] | None = None, ) -> SyftError | SyftSuccess: """ Create a request to launch the worker pool based on a built image. @@ -212,6 +213,7 @@ def create_pool_request( pool_name=pool_name, num_workers=num_workers, image_uid=image_uid, + pod_annotations=pod_annotations, ) changes: list[Change] = [create_worker_pool_change] @@ -240,6 +242,7 @@ def create_image_and_pool_request( registry_uid: UID | None = None, reason: str | None = "", pull_image: bool = True, + pod_annotations: dict[str, str] | None = None, ) -> SyftError | SyftSuccess: """ Create a request to launch the worker pool based on a built image. @@ -324,6 +327,7 @@ def create_image_and_pool_request( pool_name=pool_name, num_workers=num_workers, config=config, + pod_annotations=pod_annotations, ) changes += [create_custom_image_change, create_worker_pool_change] @@ -594,6 +598,7 @@ def sync_pool_from_request( pool_name = change.pool_name num_workers = change.num_workers image_uid = change.image_uid + pod_annotations = change.pod_annotations elif isinstance(change, CreateCustomImageChange): # type: ignore[unreachable] config = change.config tag = change.tag @@ -604,6 +609,7 @@ def sync_pool_from_request( pool_name=pool_name, num_workers=num_workers, image_uid=image_uid, + pod_annotations=pod_annotations, ) elif config is not None: return self.create_image_and_pool_request( # type: ignore[unreachable] @@ -612,6 +618,7 @@ def sync_pool_from_request( num_workers=num_workers, config=config, tag=tag, + pod_annotations=pod_annotations, ) else: return SyftError( From 9692b6e8cbad38bf272f5e66761b6a95b1337e94 Mon Sep 17 00:00:00 2001 From: rasswanth-s <43314053+rasswanth-s@users.noreply.github.com> Date: Mon, 3 Jun 2024 13:11:01 +0530 Subject: [PATCH 3/7] add pod labels to worker pool launch along with tests for pod labels --- .../api/0.8/11-container-images-k8s.ipynb | 53 +++++++++++++++---- .../backend/backend-statefulset.yaml | 6 ++- packages/grid/helm/syft/values.yaml | 10 ++-- .../syft/src/syft/custom_worker/runner_k8s.py | 15 ++++-- packages/syft/src/syft/node/node.py | 7 +++ .../syft/src/syft/service/request/request.py | 2 + .../syft/src/syft/service/worker/utils.py | 6 +++ .../service/worker/worker_pool_service.py | 11 ++++ 8 files changed, 91 insertions(+), 19 deletions(-) diff --git a/notebooks/api/0.8/11-container-images-k8s.ipynb b/notebooks/api/0.8/11-container-images-k8s.ipynb index df6265baeef..9fd2c1d22db 100644 --- a/notebooks/api/0.8/11-container-images-k8s.ipynb +++ b/notebooks/api/0.8/11-container-images-k8s.ipynb @@ -595,14 +595,16 @@ "outputs": [], "source": [ "worker_pool_name = \"custom-pool\"\n", - "custom_pool_annotations = {\"test-custom-pool\": \"Test annotation for custom pool\"}\n", + "custom_pool_pod_annotations = {\"test-custom-pool\": \"Test annotation for custom pool\"}\n", + "custom_pool_pod_labels = {\"test-custom-pool\": \"Test label for custom pool\"}\n", "worker_pool_res = domain_client.api.services.worker_pool.launch(\n", " pool_name=worker_pool_name,\n", " image_uid=workerimage.id,\n", " num_workers=3,\n", " registry_username=external_registry_username,\n", " registry_password=external_registry_password,\n", - " pod_annotations=custom_pool_annotations,\n", + " pod_annotations=custom_pool_pod_annotations,\n", + " pod_labels=custom_pool_pod_labels,\n", ")" ] }, @@ -646,16 +648,24 @@ "metadata": {}, "outputs": [], "source": [ - "# check annotations for custom pool\n", + "# check Label and Annotations for custom pool\n", "custom_pool_statefulset = get_statefulset_by_pool_name(worker_pool_name)\n", "assert custom_pool_statefulset is not None, \"Custom pool statefulset not found\"\n", "custom_pool_pod_metadata = custom_pool_statefulset.spec.template.metadata\n", + "\n", "assert (\n", " \"annotations\" in custom_pool_pod_metadata\n", "), \"Annotations not found in custom pool pod metadata\"\n", "assert (\n", - " custom_pool_pod_metadata.annotations == custom_pool_annotations\n", - "), \"Annotations do not match in Custom pool pod metadata\"" + " \"labels\" in custom_pool_pod_metadata\n", + "), \"Labels not found in custom pool pod metadata\"\n", + "\n", + "assert (\n", + " custom_pool_pod_metadata.annotations == custom_pool_pod_annotations\n", + "), \"Annotations do not match in Custom pool pod metadata\"\n", + "assert (\n", + " custom_pool_pod_metadata.labels == custom_pool_pod_labels\n", + "), \"Labels do not match in Custom pool pod metadata\"" ] }, { @@ -1122,11 +1132,13 @@ "source": [ "pool_name_opendp = \"opendp-pool\"\n", "opendp_pod_annotations = {\"test-opendp-pool\": \"Test annotation for opendp pool\"}\n", + "opendp_pod_labels = {\"test-opendp-pool\": \"Test label for opendp pool\"}\n", "pool_create_request = domain_client.api.services.worker_pool.pool_creation_request(\n", " pool_name=pool_name_opendp,\n", " num_workers=3,\n", " image_uid=workerimage_opendp.id,\n", " pod_annotations=opendp_pod_annotations,\n", + " pod_labels=opendp_pod_labels,\n", ")\n", "pool_create_request" ] @@ -1199,16 +1211,26 @@ "metadata": {}, "outputs": [], "source": [ - "# check annotations for open dp pool\n", + "# check annotations and labels for open dp pool\n", "opendp_pool_statefulset = get_statefulset_by_pool_name(pool_name_opendp)\n", "assert opendp_pool_statefulset is not None, \"Open DP pool statefulset not found\"\n", "opendp_pool_pod_metadata = opendp_pool_statefulset.spec.template.metadata\n", + "\n", + "\n", "assert (\n", " \"annotations\" in opendp_pool_pod_metadata\n", "), \"Annotations not found in opendp pool pod metadata\"\n", "assert (\n", + " \"labels\" in opendp_pool_pod_metadata\n", + "), \"Labels not found in opendp pool pod metadata\"\n", + "\n", + "\n", + "assert (\n", " opendp_pool_pod_metadata.annotations == opendp_pod_annotations\n", - "), \"Annotations do not match in opendp pool pod metadata\"" + "), \"Annotations do not match in opendp pool pod metadata\"\n", + "assert (\n", + " opendp_pool_pod_metadata.labels == opendp_pod_labels\n", + "), \"Labels do not match in opendp pool pod metadata\"" ] }, { @@ -1273,6 +1295,9 @@ "recordlinkage_pod_annotations = {\n", " \"test-recordlinkage-pool\": \"Test annotation for recordlinkage pool\"\n", "}\n", + "recordlinkage_pod_labels = {\n", + " \"test-recordlinkage-pool\": \"Test label for recordlinkage pool\"\n", + "}\n", "pool_image_create_request = domain_client.api.services.worker_pool.create_image_and_pool_request(\n", " pool_name=pool_name_recordlinkage,\n", " num_workers=2,\n", @@ -1281,6 +1306,7 @@ " registry_uid=registry_uid,\n", " reason=\"I want to do some more cool data science with PySyft and RecordLinkage!\",\n", " pod_annotations=recordlinkage_pod_annotations,\n", + " pod_labels=recordlinkage_pod_labels,\n", ")\n", "pool_image_create_request" ] @@ -1367,18 +1393,27 @@ "metadata": {}, "outputs": [], "source": [ - "# check annotations for recordlinkage pool\n", + "# check annotations and labels for recordlinkage pool\n", "recordlinkage_pool_statefulset = get_statefulset_by_pool_name(pool_name_recordlinkage)\n", "assert (\n", " recordlinkage_pool_statefulset is not None\n", "), \"RecordLinkage pool statefulset not found\"\n", "recordlinkage_pool_pod_metadata = recordlinkage_pool_statefulset.spec.template.metadata\n", + "\n", + "\n", "assert (\n", " \"annotations\" in recordlinkage_pool_pod_metadata\n", "), \"Annotations not found in recordlinkage pool pod metadata\"\n", "assert (\n", + " \"labels\" in recordlinkage_pool_pod_metadata\n", + "), \"Labels not found in recordlinkage pool pod metadata\"\n", + "\n", + "assert (\n", " recordlinkage_pool_pod_metadata.annotations == recordlinkage_pod_annotations\n", - "), \"Annotations do not match in recordlinkage pool pod metadata\"" + "), \"Annotations do not match in recordlinkage pool pod metadata\"\n", + "assert (\n", + " recordlinkage_pool_pod_metadata.labels == recordlinkage_pod_labels\n", + "), \"Labels do not match in recordlinkage pool pod metadata\"" ] }, { diff --git a/packages/grid/helm/syft/templates/backend/backend-statefulset.yaml b/packages/grid/helm/syft/templates/backend/backend-statefulset.yaml index 90614f79c14..d55c32cdee2 100644 --- a/packages/grid/helm/syft/templates/backend/backend-statefulset.yaml +++ b/packages/grid/helm/syft/templates/backend/backend-statefulset.yaml @@ -72,9 +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.defaultWorkePool.podLabels | toJson | quote }} - name: DEFAULT_WORKER_POOL_POD_ANNOTATIONS - value: {{ .Values.node.defaultWorkerPoolPodAnnotations | toJson | quote }} + 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 index 28484b96427..932f20d8978 100644 --- a/packages/grid/helm/syft/values.yaml +++ b/packages/grid/helm/syft/values.yaml @@ -160,16 +160,18 @@ node: type: domain side: high defaultBucketName: null + inMemoryWorkers: false queuePort: 5556 logLevel: info debuggerEnabled: false associationRequestAutoApproval: false useInternalRegistry: true - # Worker pool settings - inMemoryWorkers: false - defaultWorkerPoolCount: 1 - defaultWorkerPoolPodAnnotations: null + # Deault Worker pool settings + defaultWorkerPool: + count: 1 + podLabels: null + podAnnotations: null # SMTP Settings smtp: diff --git a/packages/syft/src/syft/custom_worker/runner_k8s.py b/packages/syft/src/syft/custom_worker/runner_k8s.py index b19b807d2d6..21ff334f818 100644 --- a/packages/syft/src/syft/custom_worker/runner_k8s.py +++ b/packages/syft/src/syft/custom_worker/runner_k8s.py @@ -32,6 +32,7 @@ def create_pool( 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: @@ -54,6 +55,7 @@ def create_pool( mount_secrets=mount_secrets, pull_secret=pull_secret, pod_annotations=pod_annotations, + pod_labels=pod_labels, **kwargs, ) @@ -150,6 +152,7 @@ def _create_stateful_set( 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""" @@ -185,6 +188,13 @@ 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.update(pod_labels) + stateful_set = StatefulSet( { "metadata": { @@ -204,10 +214,7 @@ 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": { diff --git a/packages/syft/src/syft/node/node.py b/packages/syft/src/syft/node/node.py index cdb74c53917..1e2c00c6f24 100644 --- a/packages/syft/src/syft/node/node.py +++ b/packages/syft/src/syft/node/node.py @@ -223,6 +223,11 @@ def get_default_worker_pool_pod_annotations() -> dict[str, str] | None: 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" @@ -1726,6 +1731,7 @@ def create_default_worker_pool(node: Node) -> SyftError | None: 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, @@ -1783,6 +1789,7 @@ def create_default_worker_pool(node: Node) -> SyftError | None: 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/service/request/request.py b/packages/syft/src/syft/service/request/request.py index 7da57beb676..2a8ae61bf60 100644 --- a/packages/syft/src/syft/service/request/request.py +++ b/packages/syft/src/syft/service/request/request.py @@ -305,6 +305,7 @@ class CreateCustomWorkerPoolChange(Change): 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"] @@ -339,6 +340,7 @@ def _run( 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) diff --git a/packages/syft/src/syft/service/worker/utils.py b/packages/syft/src/syft/service/worker/utils.py index 031833b1069..c952cbe8c13 100644 --- a/packages/syft/src/syft/service/worker/utils.py +++ b/packages/syft/src/syft/service/worker/utils.py @@ -329,6 +329,7 @@ def create_kubernetes_pool( 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 @@ -365,6 +366,7 @@ def create_kubernetes_pool( registry_password=registry_password, reg_url=reg_url, pod_annotations=pod_annotations, + pod_labels=pod_labels, ) except Exception as e: if pool: @@ -408,6 +410,7 @@ def run_workers_in_kubernetes( 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 = [] @@ -426,6 +429,7 @@ def run_workers_in_kubernetes( registry_password=registry_password, reg_url=reg_url, pod_annotations=pod_annotations, + pod_labels=pod_labels, ) else: return SyftError( @@ -509,6 +513,7 @@ def run_containers( 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 = [] @@ -546,6 +551,7 @@ def run_containers( registry_password=registry_password, reg_url=reg_url, pod_annotations=pod_annotations, + pod_labels=pod_labels, ) return results diff --git a/packages/syft/src/syft/service/worker/worker_pool_service.py b/packages/syft/src/syft/service/worker/worker_pool_service.py index 9c5280a42b1..9e7d02572c1 100644 --- a/packages/syft/src/syft/service/worker/worker_pool_service.py +++ b/packages/syft/src/syft/service/worker/worker_pool_service.py @@ -70,6 +70,7 @@ def launch( registry_username: str | None = None, registry_password: str | None = None, pod_annotations: dict[str, str] | None = None, + pod_labels: dict[str, str] | None = None, ) -> list[ContainerSpawnStatus] | SyftError: """Creates a pool of workers from the given SyftWorkerImage. @@ -130,6 +131,7 @@ def launch( registry_username=registry_username, registry_password=registry_password, pod_annotations=pod_annotations, + pod_labels=pod_labels, ) if isinstance(result, SyftError): @@ -166,6 +168,7 @@ def create_pool_request( image_uid: UID, reason: str | None = "", pod_annotations: dict[str, str] | None = None, + pod_labels: dict[str, str] | None = None, ) -> SyftError | SyftSuccess: """ Create a request to launch the worker pool based on a built image. @@ -216,6 +219,7 @@ def create_pool_request( num_workers=num_workers, image_uid=image_uid, pod_annotations=pod_annotations, + pod_labels=pod_labels, ) changes: list[Change] = [create_worker_pool_change] @@ -245,6 +249,7 @@ def create_image_and_pool_request( reason: str | None = "", pull_image: bool = True, pod_annotations: dict[str, str] | None = None, + pod_labels: dict[str, str] | None = None, ) -> SyftError | SyftSuccess: """ Create a request to launch the worker pool based on a built image. @@ -330,6 +335,7 @@ def create_image_and_pool_request( num_workers=num_workers, config=config, pod_annotations=pod_annotations, + pod_labels=pod_labels, ) changes += [create_custom_image_change, create_worker_pool_change] @@ -601,6 +607,7 @@ def sync_pool_from_request( num_workers = change.num_workers image_uid = change.image_uid pod_annotations = change.pod_annotations + pod_labels = change.pod_labels elif isinstance(change, CreateCustomImageChange): # type: ignore[unreachable] config = change.config tag = change.tag @@ -612,6 +619,7 @@ def sync_pool_from_request( num_workers=num_workers, image_uid=image_uid, pod_annotations=pod_annotations, + pod_labels=pod_labels, ) elif config is not None: return self.create_image_and_pool_request( # type: ignore[unreachable] @@ -621,6 +629,7 @@ def sync_pool_from_request( config=config, tag=tag, pod_annotations=pod_annotations, + pod_labels=pod_labels, ) else: return SyftError( @@ -668,6 +677,7 @@ def _create_workers_in_pool( registry_username: str | None = None, registry_password: str | None = None, pod_annotations: dict[str, str] | None = None, + pod_labels: dict[str, str] | None = None, ) -> tuple[list[LinkedObject], list[ContainerSpawnStatus]] | SyftError: queue_port = context.node.queue_config.client_config.queue_port @@ -700,6 +710,7 @@ def _create_workers_in_pool( registry_password=registry_password, reg_url=registry_host, pod_annotations=pod_annotations, + pod_labels=pod_labels, ) if isinstance(result, SyftError): return result From ff0952c1f602b0f595db5388cac6ef86ef2058ee Mon Sep 17 00:00:00 2001 From: rasswanth-s <43314053+rasswanth-s@users.noreply.github.com> Date: Mon, 3 Jun 2024 13:14:07 +0530 Subject: [PATCH 4/7] add default label conditional --- packages/syft/src/syft/custom_worker/runner_k8s.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/packages/syft/src/syft/custom_worker/runner_k8s.py b/packages/syft/src/syft/custom_worker/runner_k8s.py index 21ff334f818..15a241c820b 100644 --- a/packages/syft/src/syft/custom_worker/runner_k8s.py +++ b/packages/syft/src/syft/custom_worker/runner_k8s.py @@ -194,6 +194,8 @@ def _create_stateful_set( } if isinstance(pod_labels, dict): pod_labels = default_pod_labels.update(pod_labels) + else: + pod_labels = default_pod_labels stateful_set = StatefulSet( { From e79f2836a9371fe1b48c754111f7bd061a5df229 Mon Sep 17 00:00:00 2001 From: rasswanth-s <43314053+rasswanth-s@users.noreply.github.com> Date: Mon, 3 Jun 2024 13:37:40 +0530 Subject: [PATCH 5/7] fix dictionary update error --- .../grid/helm/syft/templates/backend/backend-statefulset.yaml | 2 +- packages/syft/src/syft/custom_worker/runner_k8s.py | 3 ++- packages/syft/src/syft/protocol/protocol_version.json | 2 +- 3 files changed, 4 insertions(+), 3 deletions(-) diff --git a/packages/grid/helm/syft/templates/backend/backend-statefulset.yaml b/packages/grid/helm/syft/templates/backend/backend-statefulset.yaml index d55c32cdee2..106d2fee893 100644 --- a/packages/grid/helm/syft/templates/backend/backend-statefulset.yaml +++ b/packages/grid/helm/syft/templates/backend/backend-statefulset.yaml @@ -74,7 +74,7 @@ spec: - name: DEFAULT_WORKER_POOL_COUNT value: {{ .Values.node.defaultWorkerPool.count | quote }} - name: DEFAULT_WORKER_POOL_POD_LABELS - value: {{ .Values.node.defaultWorkePool.podLabels | toJson | quote }} + value: {{ .Values.node.defaultWorkerPool.podLabels | toJson | quote }} - name: DEFAULT_WORKER_POOL_POD_ANNOTATIONS value: {{ .Values.node.defaultWorkerPool.podAnnotations | toJson | quote }} - name: USE_INTERNAL_REGISTRY diff --git a/packages/syft/src/syft/custom_worker/runner_k8s.py b/packages/syft/src/syft/custom_worker/runner_k8s.py index 15a241c820b..3e739ef4fdb 100644 --- a/packages/syft/src/syft/custom_worker/runner_k8s.py +++ b/packages/syft/src/syft/custom_worker/runner_k8s.py @@ -192,8 +192,9 @@ def _create_stateful_set( "app.kubernetes.io/name": KUBERNETES_NAMESPACE, "app.kubernetes.io/component": pool_name, } + if isinstance(pod_labels, dict): - pod_labels = default_pod_labels.update(pod_labels) + pod_labels = {**default_pod_labels, **pod_labels} else: pod_labels = default_pod_labels diff --git a/packages/syft/src/syft/protocol/protocol_version.json b/packages/syft/src/syft/protocol/protocol_version.json index 6c177ab6275..faf2c0de850 100644 --- a/packages/syft/src/syft/protocol/protocol_version.json +++ b/packages/syft/src/syft/protocol/protocol_version.json @@ -250,7 +250,7 @@ "CreateCustomWorkerPoolChange": { "3": { "version": 3, - "hash": "1c92c173a9eac781377fe5337262ae131111c2290a3e543b17bce2842bd037f8", + "hash": "e982f2ebcdc6fe23a65a014109e33ba7c487bb7ca5623723cf5ec7642f86828c", "action": "add" } } From c25f4dc68be740d28fb0c53e312f900d967a06de Mon Sep 17 00:00:00 2001 From: rasswanth-s <43314053+rasswanth-s@users.noreply.github.com> Date: Mon, 3 Jun 2024 14:13:58 +0530 Subject: [PATCH 6/7] fix notebooks tests --- .../api/0.8/11-container-images-k8s.ipynb | 41 ++++++++++--------- 1 file changed, 21 insertions(+), 20 deletions(-) diff --git a/notebooks/api/0.8/11-container-images-k8s.ipynb b/notebooks/api/0.8/11-container-images-k8s.ipynb index 9fd2c1d22db..c685e6f8d7d 100644 --- a/notebooks/api/0.8/11-container-images-k8s.ipynb +++ b/notebooks/api/0.8/11-container-images-k8s.ipynb @@ -64,7 +64,11 @@ " )\n", " if len(pool_list) == 0:\n", " return None\n", - " return pool_list[0]" + " return pool_list[0]\n", + "\n", + "\n", + "def is_subset_dict(subset, superset):\n", + " return all(item in superset.items() for item in subset.items())" ] }, { @@ -596,7 +600,7 @@ "source": [ "worker_pool_name = \"custom-pool\"\n", "custom_pool_pod_annotations = {\"test-custom-pool\": \"Test annotation for custom pool\"}\n", - "custom_pool_pod_labels = {\"test-custom-pool\": \"Test label for custom pool\"}\n", + "custom_pool_pod_labels = {\"test-custom-pool\": \"test_label_for_custom_pool\"}\n", "worker_pool_res = domain_client.api.services.worker_pool.launch(\n", " pool_name=worker_pool_name,\n", " image_uid=workerimage.id,\n", @@ -660,11 +664,11 @@ " \"labels\" in custom_pool_pod_metadata\n", "), \"Labels not found in custom pool pod metadata\"\n", "\n", - "assert (\n", - " custom_pool_pod_metadata.annotations == custom_pool_pod_annotations\n", + "assert is_subset_dict(\n", + " custom_pool_pod_annotations, custom_pool_pod_metadata.annotations\n", "), \"Annotations do not match in Custom pool pod metadata\"\n", - "assert (\n", - " custom_pool_pod_metadata.labels == custom_pool_pod_labels\n", + "assert is_subset_dict(\n", + " custom_pool_pod_labels, custom_pool_pod_metadata.labels\n", "), \"Labels do not match in Custom pool pod metadata\"" ] }, @@ -1132,7 +1136,7 @@ "source": [ "pool_name_opendp = \"opendp-pool\"\n", "opendp_pod_annotations = {\"test-opendp-pool\": \"Test annotation for opendp pool\"}\n", - "opendp_pod_labels = {\"test-opendp-pool\": \"Test label for opendp pool\"}\n", + "opendp_pod_labels = {\"test-opendp-pool\": \"test_label_for_opendp_pool\"}\n", "pool_create_request = domain_client.api.services.worker_pool.pool_creation_request(\n", " pool_name=pool_name_opendp,\n", " num_workers=3,\n", @@ -1225,11 +1229,11 @@ "), \"Labels not found in opendp pool pod metadata\"\n", "\n", "\n", - "assert (\n", - " opendp_pool_pod_metadata.annotations == opendp_pod_annotations\n", + "assert is_subset_dict(\n", + " opendp_pod_annotations, opendp_pool_pod_metadata.annotations\n", "), \"Annotations do not match in opendp pool pod metadata\"\n", - "assert (\n", - " opendp_pool_pod_metadata.labels == opendp_pod_labels\n", + "assert is_subset_dict(\n", + " opendp_pod_labels, opendp_pool_pod_metadata.labels\n", "), \"Labels do not match in opendp pool pod metadata\"" ] }, @@ -1296,7 +1300,7 @@ " \"test-recordlinkage-pool\": \"Test annotation for recordlinkage pool\"\n", "}\n", "recordlinkage_pod_labels = {\n", - " \"test-recordlinkage-pool\": \"Test label for recordlinkage pool\"\n", + " \"test-recordlinkage-pool\": \"test_label_for_recordlinkage_pool\"\n", "}\n", "pool_image_create_request = domain_client.api.services.worker_pool.create_image_and_pool_request(\n", " pool_name=pool_name_recordlinkage,\n", @@ -1401,19 +1405,16 @@ "recordlinkage_pool_pod_metadata = recordlinkage_pool_statefulset.spec.template.metadata\n", "\n", "\n", - "assert (\n", - " \"annotations\" in recordlinkage_pool_pod_metadata\n", + "assert is_subset_dict(\n", + " recordlinkage_pod_annotations, recordlinkage_pool_pod_metadata.annotations\n", "), \"Annotations not found in recordlinkage pool pod metadata\"\n", "assert (\n", " \"labels\" in recordlinkage_pool_pod_metadata\n", "), \"Labels not found in recordlinkage pool pod metadata\"\n", "\n", - "assert (\n", - " recordlinkage_pool_pod_metadata.annotations == recordlinkage_pod_annotations\n", - "), \"Annotations do not match in recordlinkage pool pod metadata\"\n", - "assert (\n", - " recordlinkage_pool_pod_metadata.labels == recordlinkage_pod_labels\n", - "), \"Labels do not match in recordlinkage pool pod metadata\"" + "assert is_subset_dict(\n", + " recordlinkage_pod_labels, recordlinkage_pool_pod_metadata.labels\n", + "), \"Annotations do not match in recordlinkage pool pod metadata\"" ] }, { From a33d981f486457ec26601d194d90cc5a638f9fa8 Mon Sep 17 00:00:00 2001 From: rasswanth-s <43314053+rasswanth-s@users.noreply.github.com> Date: Mon, 3 Jun 2024 16:36:13 +0530 Subject: [PATCH 7/7] update old default worker pool count references --- packages/grid/helm/examples/azure/azure.high.yaml | 6 +++++- packages/grid/helm/syft/values.yaml | 2 +- packages/grid/helm/values.dev.high.yaml | 6 +++++- packages/grid/helm/values.dev.low.yaml | 6 +++++- packages/grid/helm/values.dev.yaml | 6 +++++- 5 files changed, 21 insertions(+), 5 deletions(-) diff --git a/packages/grid/helm/examples/azure/azure.high.yaml b/packages/grid/helm/examples/azure/azure.high.yaml index 09b730292fd..468d1e771ae 100644 --- 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: info@openmined.org - 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/values.yaml b/packages/grid/helm/syft/values.yaml index 932f20d8978..5a361cc2e39 100644 --- a/packages/grid/helm/syft/values.yaml +++ b/packages/grid/helm/syft/values.yaml @@ -167,7 +167,7 @@ node: associationRequestAutoApproval: false useInternalRegistry: true - # Deault Worker pool settings + # Default Worker pool settings defaultWorkerPool: count: 1 podLabels: null diff --git a/packages/grid/helm/values.dev.high.yaml b/packages/grid/helm/values.dev.high.yaml index 62b9fddb9fd..9a0e266704a 100644 --- a/packages/grid/helm/values.dev.high.yaml +++ b/packages/grid/helm/values.dev.high.yaml @@ -12,12 +12,16 @@ registry: node: rootEmail: info@openmined.org - 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 index 62d63c4f73b..7e5de1a68f2 100644 --- a/packages/grid/helm/values.dev.low.yaml +++ b/packages/grid/helm/values.dev.low.yaml @@ -12,12 +12,16 @@ registry: node: rootEmail: info@openmined.org - 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 index 2b90ed7283a..c24aa51d294 100644 --- a/packages/grid/helm/values.dev.yaml +++ b/packages/grid/helm/values.dev.yaml @@ -12,12 +12,16 @@ registry: node: rootEmail: info@openmined.org - defaultWorkerPoolCount: 1 associationRequestAutoApproval: true resourcesPreset: null resources: null + defaultWorkerPool: + count: 1 + podLabels: null + podAnnotations: null + secret: defaultRootPassword: changethis