diff --git a/notebooks/api/0.8/11-container-images-k8s.ipynb b/notebooks/api/0.8/11-container-images-k8s.ipynb index af715192027..c685e6f8d7d 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,38 @@ "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]\n", + "\n", + "\n", + "def is_subset_dict(subset, superset):\n", + " return all(item in superset.items() for item in subset.items())" + ] + }, + { + "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 +89,7 @@ { "cell_type": "code", "execution_count": null, - "id": "3", + "id": "5", "metadata": {}, "outputs": [], "source": [ @@ -69,7 +102,7 @@ { "cell_type": "code", "execution_count": null, - "id": "4", + "id": "6", "metadata": {}, "outputs": [], "source": [ @@ -79,7 +112,7 @@ }, { "cell_type": "markdown", - "id": "5", + "id": "7", "metadata": {}, "source": [ "### Scaling Default Worker Pool" @@ -87,7 +120,7 @@ }, { "cell_type": "markdown", - "id": "6", + "id": "8", "metadata": {}, "source": [ "We should see a default worker pool" @@ -96,7 +129,7 @@ { "cell_type": "code", "execution_count": null, - "id": "7", + "id": "9", "metadata": {}, "outputs": [], "source": [ @@ -105,7 +138,7 @@ }, { "cell_type": "markdown", - "id": "8", + "id": "10", "metadata": {}, "source": [ "Scale up to 3 workers" @@ -114,7 +147,7 @@ { "cell_type": "code", "execution_count": null, - "id": "9", + "id": "11", "metadata": {}, "outputs": [], "source": [ @@ -128,7 +161,7 @@ { "cell_type": "code", "execution_count": null, - "id": "10", + "id": "12", "metadata": {}, "outputs": [], "source": [ @@ -140,7 +173,7 @@ { "cell_type": "code", "execution_count": null, - "id": "11", + "id": "13", "metadata": {}, "outputs": [], "source": [ @@ -153,7 +186,7 @@ }, { "cell_type": "markdown", - "id": "12", + "id": "14", "metadata": {}, "source": [ "Scale down to 1 worker" @@ -162,7 +195,7 @@ { "cell_type": "code", "execution_count": null, - "id": "13", + "id": "15", "metadata": {}, "outputs": [], "source": [ @@ -176,7 +209,7 @@ { "cell_type": "code", "execution_count": null, - "id": "14", + "id": "16", "metadata": {}, "outputs": [], "source": [ @@ -188,7 +221,7 @@ { "cell_type": "code", "execution_count": null, - "id": "15", + "id": "17", "metadata": {}, "outputs": [], "source": [ @@ -200,7 +233,7 @@ }, { "cell_type": "markdown", - "id": "16", + "id": "18", "metadata": {}, "source": [ "#### Submit Dockerfile" @@ -209,7 +242,7 @@ { "cell_type": "code", "execution_count": null, - "id": "17", + "id": "19", "metadata": {}, "outputs": [], "source": [ @@ -226,7 +259,7 @@ { "cell_type": "code", "execution_count": null, - "id": "18", + "id": "20", "metadata": {}, "outputs": [], "source": [ @@ -241,7 +274,7 @@ { "cell_type": "code", "execution_count": null, - "id": "19", + "id": "21", "metadata": {}, "outputs": [], "source": [ @@ -251,7 +284,7 @@ { "cell_type": "code", "execution_count": null, - "id": "20", + "id": "22", "metadata": {}, "outputs": [], "source": [ @@ -261,7 +294,7 @@ { "cell_type": "code", "execution_count": null, - "id": "21", + "id": "23", "metadata": {}, "outputs": [], "source": [ @@ -274,7 +307,7 @@ { "cell_type": "code", "execution_count": null, - "id": "22", + "id": "24", "metadata": {}, "outputs": [], "source": [ @@ -284,7 +317,7 @@ { "cell_type": "code", "execution_count": null, - "id": "23", + "id": "25", "metadata": {}, "outputs": [], "source": [ @@ -295,7 +328,7 @@ { "cell_type": "code", "execution_count": null, - "id": "24", + "id": "26", "metadata": {}, "outputs": [], "source": [ @@ -306,7 +339,7 @@ { "cell_type": "code", "execution_count": null, - "id": "25", + "id": "27", "metadata": {}, "outputs": [], "source": [ @@ -325,7 +358,7 @@ }, { "cell_type": "markdown", - "id": "26", + "id": "28", "metadata": {}, "source": [ "#### Add External Registry in Syft" @@ -334,7 +367,7 @@ { "cell_type": "code", "execution_count": null, - "id": "27", + "id": "29", "metadata": {}, "outputs": [], "source": [ @@ -350,7 +383,7 @@ { "cell_type": "code", "execution_count": null, - "id": "28", + "id": "30", "metadata": {}, "outputs": [], "source": [ @@ -361,7 +394,7 @@ { "cell_type": "code", "execution_count": null, - "id": "29", + "id": "31", "metadata": {}, "outputs": [], "source": [ @@ -371,7 +404,7 @@ { "cell_type": "code", "execution_count": null, - "id": "30", + "id": "32", "metadata": {}, "outputs": [], "source": [ @@ -382,7 +415,7 @@ { "cell_type": "code", "execution_count": null, - "id": "31", + "id": "33", "metadata": {}, "outputs": [], "source": [ @@ -393,7 +426,7 @@ { "cell_type": "code", "execution_count": null, - "id": "32", + "id": "34", "metadata": {}, "outputs": [], "source": [ @@ -404,7 +437,7 @@ { "cell_type": "code", "execution_count": null, - "id": "33", + "id": "35", "metadata": {}, "outputs": [], "source": [ @@ -414,7 +447,7 @@ { "cell_type": "code", "execution_count": null, - "id": "34", + "id": "36", "metadata": {}, "outputs": [], "source": [ @@ -423,7 +456,7 @@ }, { "cell_type": "markdown", - "id": "35", + "id": "37", "metadata": {}, "source": [ "#### Build Image" @@ -432,7 +465,7 @@ { "cell_type": "code", "execution_count": null, - "id": "36", + "id": "38", "metadata": {}, "outputs": [], "source": [ @@ -450,7 +483,7 @@ { "cell_type": "code", "execution_count": null, - "id": "37", + "id": "39", "metadata": {}, "outputs": [], "source": [ @@ -460,7 +493,7 @@ { "cell_type": "code", "execution_count": null, - "id": "38", + "id": "40", "metadata": {}, "outputs": [], "source": [ @@ -471,7 +504,7 @@ { "cell_type": "code", "execution_count": null, - "id": "39", + "id": "41", "metadata": {}, "outputs": [], "source": [ @@ -483,7 +516,7 @@ { "cell_type": "code", "execution_count": null, - "id": "40", + "id": "42", "metadata": {}, "outputs": [], "source": [ @@ -496,7 +529,7 @@ }, { "cell_type": "markdown", - "id": "41", + "id": "43", "metadata": {}, "source": [ "#### Push Image to Local Registry" @@ -505,7 +538,7 @@ { "cell_type": "code", "execution_count": null, - "id": "42", + "id": "44", "metadata": {}, "outputs": [], "source": [ @@ -521,7 +554,7 @@ { "cell_type": "code", "execution_count": null, - "id": "43", + "id": "45", "metadata": {}, "outputs": [], "source": [ @@ -531,7 +564,7 @@ { "cell_type": "code", "execution_count": null, - "id": "44", + "id": "46", "metadata": {}, "outputs": [], "source": [ @@ -552,7 +585,7 @@ }, { "cell_type": "markdown", - "id": "45", + "id": "47", "metadata": {}, "source": [ "#### Create Worker Pool From Image" @@ -561,24 +594,28 @@ { "cell_type": "code", "execution_count": null, - "id": "46", + "id": "48", "metadata": {}, "outputs": [], "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", "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_pod_annotations,\n", + " pod_labels=custom_pool_pod_labels,\n", ")" ] }, { "cell_type": "code", "execution_count": null, - "id": "47", + "id": "49", "metadata": {}, "outputs": [], "source": [ @@ -589,7 +626,7 @@ { "cell_type": "code", "execution_count": null, - "id": "48", + "id": "50", "metadata": {}, "outputs": [], "source": [ @@ -600,7 +637,7 @@ { "cell_type": "code", "execution_count": null, - "id": "49", + "id": "51", "metadata": {}, "outputs": [], "source": [ @@ -611,7 +648,34 @@ { "cell_type": "code", "execution_count": null, - "id": "50", + "id": "52", + "metadata": {}, + "outputs": [], + "source": [ + "# 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", + " \"labels\" in custom_pool_pod_metadata\n", + "), \"Labels not found in custom pool pod metadata\"\n", + "\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 is_subset_dict(\n", + " custom_pool_pod_labels, custom_pool_pod_metadata.labels\n", + "), \"Labels do not match in Custom pool pod metadata\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "53", "metadata": {}, "outputs": [], "source": [ @@ -622,7 +686,7 @@ { "cell_type": "code", "execution_count": null, - "id": "51", + "id": "54", "metadata": {}, "outputs": [], "source": [ @@ -640,7 +704,7 @@ { "cell_type": "code", "execution_count": null, - "id": "52", + "id": "55", "metadata": {}, "outputs": [], "source": [ @@ -654,7 +718,7 @@ { "cell_type": "code", "execution_count": null, - "id": "53", + "id": "56", "metadata": {}, "outputs": [], "source": [ @@ -664,7 +728,7 @@ { "cell_type": "code", "execution_count": null, - "id": "54", + "id": "57", "metadata": {}, "outputs": [], "source": [ @@ -674,7 +738,7 @@ }, { "cell_type": "markdown", - "id": "55", + "id": "58", "metadata": {}, "source": [ "#### Get Worker Logs" @@ -683,7 +747,7 @@ { "cell_type": "code", "execution_count": null, - "id": "56", + "id": "59", "metadata": {}, "outputs": [], "source": [ @@ -696,7 +760,7 @@ { "cell_type": "code", "execution_count": null, - "id": "57", + "id": "60", "metadata": {}, "outputs": [], "source": [ @@ -706,7 +770,7 @@ { "cell_type": "code", "execution_count": null, - "id": "58", + "id": "61", "metadata": {}, "outputs": [], "source": [ @@ -715,7 +779,7 @@ }, { "cell_type": "markdown", - "id": "59", + "id": "62", "metadata": {}, "source": [ "### Syft function" @@ -724,7 +788,7 @@ { "cell_type": "code", "execution_count": null, - "id": "60", + "id": "63", "metadata": {}, "outputs": [], "source": [ @@ -738,7 +802,7 @@ { "cell_type": "code", "execution_count": null, - "id": "61", + "id": "64", "metadata": {}, "outputs": [], "source": [ @@ -754,7 +818,7 @@ { "cell_type": "code", "execution_count": null, - "id": "62", + "id": "65", "metadata": {}, "outputs": [], "source": [ @@ -764,7 +828,7 @@ { "cell_type": "code", "execution_count": null, - "id": "63", + "id": "66", "metadata": {}, "outputs": [], "source": [ @@ -774,7 +838,7 @@ { "cell_type": "code", "execution_count": null, - "id": "64", + "id": "67", "metadata": {}, "outputs": [], "source": [ @@ -785,7 +849,7 @@ { "cell_type": "code", "execution_count": null, - "id": "65", + "id": "68", "metadata": {}, "outputs": [], "source": [ @@ -795,7 +859,7 @@ { "cell_type": "code", "execution_count": null, - "id": "66", + "id": "69", "metadata": {}, "outputs": [], "source": [ @@ -806,7 +870,7 @@ { "cell_type": "code", "execution_count": null, - "id": "67", + "id": "70", "metadata": {}, "outputs": [], "source": [ @@ -817,7 +881,7 @@ { "cell_type": "code", "execution_count": null, - "id": "68", + "id": "71", "metadata": {}, "outputs": [], "source": [ @@ -827,7 +891,7 @@ { "cell_type": "code", "execution_count": null, - "id": "69", + "id": "72", "metadata": {}, "outputs": [], "source": [ @@ -837,7 +901,7 @@ { "cell_type": "code", "execution_count": null, - "id": "70", + "id": "73", "metadata": {}, "outputs": [], "source": [ @@ -847,7 +911,7 @@ { "cell_type": "code", "execution_count": null, - "id": "71", + "id": "74", "metadata": {}, "outputs": [], "source": [ @@ -857,7 +921,7 @@ { "cell_type": "code", "execution_count": null, - "id": "72", + "id": "75", "metadata": {}, "outputs": [], "source": [ @@ -868,7 +932,7 @@ { "cell_type": "code", "execution_count": null, - "id": "73", + "id": "76", "metadata": {}, "outputs": [], "source": [ @@ -881,7 +945,7 @@ { "cell_type": "code", "execution_count": null, - "id": "74", + "id": "77", "metadata": {}, "outputs": [], "source": [ @@ -896,7 +960,7 @@ { "cell_type": "code", "execution_count": null, - "id": "75", + "id": "78", "metadata": {}, "outputs": [], "source": [ @@ -905,7 +969,7 @@ }, { "cell_type": "markdown", - "id": "76", + "id": "79", "metadata": {}, "source": [ "#### Worker Pool and Image Creation Request/Approval" @@ -914,7 +978,7 @@ { "cell_type": "code", "execution_count": null, - "id": "77", + "id": "80", "metadata": {}, "outputs": [], "source": [ @@ -930,7 +994,7 @@ { "cell_type": "code", "execution_count": null, - "id": "78", + "id": "81", "metadata": {}, "outputs": [], "source": [ @@ -944,7 +1008,7 @@ { "cell_type": "code", "execution_count": null, - "id": "79", + "id": "82", "metadata": {}, "outputs": [], "source": [ @@ -954,7 +1018,7 @@ { "cell_type": "code", "execution_count": null, - "id": "80", + "id": "83", "metadata": {}, "outputs": [], "source": [ @@ -965,7 +1029,7 @@ { "cell_type": "code", "execution_count": null, - "id": "81", + "id": "84", "metadata": {}, "outputs": [], "source": [ @@ -978,7 +1042,7 @@ }, { "cell_type": "markdown", - "id": "82", + "id": "85", "metadata": {}, "source": [ "##### Build image first then create pool" @@ -987,7 +1051,7 @@ { "cell_type": "code", "execution_count": null, - "id": "83", + "id": "86", "metadata": {}, "outputs": [], "source": [ @@ -1005,7 +1069,7 @@ { "cell_type": "code", "execution_count": null, - "id": "84", + "id": "87", "metadata": {}, "outputs": [], "source": [ @@ -1015,7 +1079,7 @@ { "cell_type": "code", "execution_count": null, - "id": "85", + "id": "88", "metadata": {}, "outputs": [], "source": [ @@ -1026,7 +1090,7 @@ { "cell_type": "code", "execution_count": null, - "id": "86", + "id": "89", "metadata": {}, "outputs": [], "source": [ @@ -1049,7 +1113,7 @@ { "cell_type": "code", "execution_count": null, - "id": "87", + "id": "90", "metadata": {}, "outputs": [], "source": [ @@ -1066,13 +1130,19 @@ { "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", + "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, 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", + " pod_labels=opendp_pod_labels,\n", ")\n", "pool_create_request" ] @@ -1080,7 +1150,7 @@ { "cell_type": "code", "execution_count": null, - "id": "89", + "id": "92", "metadata": {}, "outputs": [], "source": [ @@ -1091,7 +1161,7 @@ { "cell_type": "code", "execution_count": null, - "id": "90", + "id": "93", "metadata": {}, "outputs": [], "source": [ @@ -1106,7 +1176,7 @@ { "cell_type": "code", "execution_count": null, - "id": "91", + "id": "94", "metadata": {}, "outputs": [], "source": [ @@ -1116,7 +1186,7 @@ { "cell_type": "code", "execution_count": null, - "id": "92", + "id": "95", "metadata": {}, "outputs": [], "source": [ @@ -1128,7 +1198,7 @@ { "cell_type": "code", "execution_count": null, - "id": "93", + "id": "96", "metadata": {}, "outputs": [], "source": [ @@ -1141,7 +1211,36 @@ { "cell_type": "code", "execution_count": null, - "id": "94", + "id": "97", + "metadata": {}, + "outputs": [], + "source": [ + "# 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 is_subset_dict(\n", + " opendp_pod_annotations, opendp_pool_pod_metadata.annotations\n", + "), \"Annotations do not match in opendp pool pod metadata\"\n", + "assert is_subset_dict(\n", + " opendp_pod_labels, opendp_pool_pod_metadata.labels\n", + "), \"Labels do not match in opendp pool pod metadata\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "98", "metadata": {}, "outputs": [], "source": [ @@ -1156,7 +1255,7 @@ { "cell_type": "code", "execution_count": null, - "id": "95", + "id": "99", "metadata": {}, "outputs": [], "source": [ @@ -1165,7 +1264,7 @@ }, { "cell_type": "markdown", - "id": "96", + "id": "100", "metadata": {}, "source": [ "Request to build the image and create the pool at the same time" @@ -1174,7 +1273,7 @@ { "cell_type": "code", "execution_count": null, - "id": "97", + "id": "101", "metadata": {}, "outputs": [], "source": [ @@ -1192,21 +1291,26 @@ { "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", + "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", + " 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", + " pod_labels=recordlinkage_pod_labels,\n", ")\n", "pool_image_create_request" ] @@ -1214,7 +1318,7 @@ { "cell_type": "code", "execution_count": null, - "id": "99", + "id": "103", "metadata": {}, "outputs": [], "source": [ @@ -1226,7 +1330,7 @@ { "cell_type": "code", "execution_count": null, - "id": "100", + "id": "104", "metadata": {}, "outputs": [], "source": [ @@ -1239,7 +1343,7 @@ { "cell_type": "code", "execution_count": null, - "id": "101", + "id": "105", "metadata": {}, "outputs": [], "source": [ @@ -1253,7 +1357,7 @@ { "cell_type": "code", "execution_count": null, - "id": "102", + "id": "106", "metadata": {}, "outputs": [], "source": [ @@ -1263,7 +1367,7 @@ { "cell_type": "code", "execution_count": null, - "id": "103", + "id": "107", "metadata": {}, "outputs": [], "source": [ @@ -1274,7 +1378,7 @@ { "cell_type": "code", "execution_count": null, - "id": "104", + "id": "108", "metadata": {}, "outputs": [], "source": [ @@ -1289,7 +1393,34 @@ { "cell_type": "code", "execution_count": null, - "id": "105", + "id": "109", + "metadata": {}, + "outputs": [], + "source": [ + "# 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 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 is_subset_dict(\n", + " recordlinkage_pod_labels, recordlinkage_pool_pod_metadata.labels\n", + "), \"Annotations do not match in recordlinkage pool pod metadata\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "110", "metadata": {}, "outputs": [], "source": [ @@ -1299,7 +1430,7 @@ { "cell_type": "code", "execution_count": null, - "id": "106", + "id": "111", "metadata": {}, "outputs": [], "source": [ @@ -1317,7 +1448,7 @@ { "cell_type": "code", "execution_count": null, - "id": "107", + "id": "112", "metadata": {}, "outputs": [], "source": [ @@ -1328,7 +1459,7 @@ { "cell_type": "code", "execution_count": null, - "id": "108", + "id": "113", "metadata": {}, "outputs": [], "source": [ @@ -1345,7 +1476,7 @@ { "cell_type": "code", "execution_count": null, - "id": "109", + "id": "114", "metadata": {}, "outputs": [], "source": [ 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/templates/backend/backend-statefulset.yaml b/packages/grid/helm/syft/templates/backend/backend-statefulset.yaml index 8048f262e5e..106d2fee893 100644 --- 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 index 0d2bf246384..2644eac26e4 100644 --- a/packages/grid/helm/syft/values.yaml +++ b/packages/grid/helm/syft/values.yaml @@ -159,15 +159,20 @@ node: rootEmail: info@openmined.org 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 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 diff --git a/packages/syft/src/syft/custom_worker/runner_k8s.py b/packages/syft/src/syft/custom_worker/runner_k8s.py index 949311338b9..3e739ef4fdb 100644 --- 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 index f21058de3b6..1e2c00c6f24 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,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 index 305a4904798..5f52d682c21 100644 --- 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 index abb14a713f1..2a8ae61bf60 100644 --- 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 index 18895f2cb02..c952cbe8c13 100644 --- 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 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 7f648b7c8e5..9e7d02572c1 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,8 @@ def launch( num_workers: int, 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. @@ -128,6 +130,8 @@ def launch( worker_stash=worker_stash, registry_username=registry_username, registry_password=registry_password, + pod_annotations=pod_annotations, + pod_labels=pod_labels, ) if isinstance(result, SyftError): @@ -163,6 +167,8 @@ def create_pool_request( num_workers: int, 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. @@ -212,6 +218,8 @@ def create_pool_request( pool_name=pool_name, num_workers=num_workers, image_uid=image_uid, + pod_annotations=pod_annotations, + pod_labels=pod_labels, ) changes: list[Change] = [create_worker_pool_change] @@ -240,6 +248,8 @@ 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, + pod_labels: dict[str, str] | None = None, ) -> SyftError | SyftSuccess: """ Create a request to launch the worker pool based on a built image. @@ -324,6 +334,8 @@ def create_image_and_pool_request( pool_name=pool_name, num_workers=num_workers, config=config, + pod_annotations=pod_annotations, + pod_labels=pod_labels, ) changes += [create_custom_image_change, create_worker_pool_change] @@ -594,6 +606,8 @@ 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 + pod_labels = change.pod_labels elif isinstance(change, CreateCustomImageChange): # type: ignore[unreachable] config = change.config tag = change.tag @@ -604,6 +618,8 @@ def sync_pool_from_request( pool_name=pool_name, 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] @@ -612,6 +628,8 @@ def sync_pool_from_request( num_workers=num_workers, config=config, tag=tag, + pod_annotations=pod_annotations, + pod_labels=pod_labels, ) else: return SyftError( @@ -658,6 +676,8 @@ def _create_workers_in_pool( worker_stash: WorkerStash, 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 @@ -689,6 +709,8 @@ def _create_workers_in_pool( registry_username=registry_username, registry_password=registry_password, reg_url=registry_host, + pod_annotations=pod_annotations, + pod_labels=pod_labels, ) if isinstance(result, SyftError): return result