notebooks/api/0.8/11-container-images-k8s.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0",
   "metadata": {},
   "outputs": [],
   "source": [
    "SYFT_VERSION = \">=0.8.2.b0,<0.9\"\n",
    "package_string = f'\"syft{SYFT_VERSION}\"'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# stdlib\n",
    "import os\n",
    "\n",
    "# third party\n",
    "import kr8s\n",
    "import numpy as np\n",
    "import requests\n",
    "\n",
    "# syft absolute\n",
    "import syft as sy\n",
    "\n",
    "from getpass import getpass  # noqa\n",
    "\n",
    "\n",
    "sy.requires(SYFT_VERSION)\n",
    "\n",
    "# syft absolute\n",
    "from syft.service.worker.image_registry import SyftImageRegistry\n",
    "from syft.service.worker.worker_image import SyftWorkerImage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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",
    "\n",
    "# Uncomment this to add custom values\n",
    "# os.environ[\"SERVER_URL\"] = \"http://localhost\"\n",
    "# os.environ[\"SERVER_PORT\"] = \"8080\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "datasite = sy.orchestra.launch(\n",
    "    name=\"test-datasite-1\",\n",
    "    dev_mode=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "datasite_client = datasite.login(email=\"info@openmined.org\", password=\"changethis\")\n",
    "datasite_client"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7",
   "metadata": {},
   "source": [
    "### Scaling Default Worker Pool"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8",
   "metadata": {},
   "source": [
    "We should see a default worker pool"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "datasite_client.worker_pools"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10",
   "metadata": {},
   "source": [
    "Scale up to 3 workers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "result = datasite_client.api.services.worker_pool.scale(\n",
    "    number=3, pool_name=\"default-pool\"\n",
    ")\n",
    "assert not isinstance(result, sy.SyftError), str(result)\n",
    "result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "result = datasite_client.api.services.worker_pool.get_by_name(pool_name=\"default-pool\")\n",
    "assert len(result.workers) == 3, str(result.to_dict())\n",
    "result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "# stdlib\n",
    "# wait for some time for scale up to be ready\n",
    "from time import sleep\n",
    "\n",
    "sleep(5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14",
   "metadata": {},
   "source": [
    "Scale down to 1 worker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "default_pool_scale_res = datasite_client.api.services.worker_pool.scale(\n",
    "    number=1, pool_name=\"default-pool\"\n",
    ")\n",
    "assert not isinstance(default_pool_scale_res, sy.SyftError), str(default_pool_scale_res)\n",
    "default_pool_scale_res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "metadata": {},
   "outputs": [],
   "source": [
    "result = datasite_client.api.services.worker_pool.get_by_name(pool_name=\"default-pool\")\n",
    "assert len(result.workers) == 1, str(result.to_dict())\n",
    "result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17",
   "metadata": {},
   "outputs": [],
   "source": [
    "default_worker_pool = datasite_client.api.services.worker_pool.get_by_name(\n",
    "    pool_name=\"default-pool\"\n",
    ")\n",
    "default_worker_pool"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18",
   "metadata": {},
   "source": [
    "#### Submit Dockerfile"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19",
   "metadata": {},
   "outputs": [],
   "source": [
    "# syft absolute\n",
    "from syft.util.util import get_latest_tag\n",
    "\n",
    "registry = os.getenv(\"SYFT_BASE_IMAGE_REGISTRY\", \"docker.io\")\n",
    "repo = \"openmined/syft-backend\"\n",
    "\n",
    "if \"k3d\" in registry:\n",
    "    tag = get_latest_tag(registry, repo)\n",
    "else:\n",
    "    tag = sy.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20",
   "metadata": {},
   "outputs": [],
   "source": [
    "custom_dockerfile_str = f\"\"\"\n",
    "FROM {registry}/{repo}:{tag}\n",
    "\n",
    "RUN uv pip install pydicom\n",
    "\n",
    "\"\"\".strip()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21",
   "metadata": {},
   "outputs": [],
   "source": [
    "docker_config = sy.DockerWorkerConfig(dockerfile=custom_dockerfile_str)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert docker_config.dockerfile == custom_dockerfile_str"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23",
   "metadata": {},
   "outputs": [],
   "source": [
    "submit_result = datasite_client.api.services.worker_image.submit(\n",
    "    worker_config=docker_config\n",
    ")\n",
    "submit_result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert isinstance(submit_result, sy.SyftSuccess), str(submit_result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25",
   "metadata": {},
   "outputs": [],
   "source": [
    "dockerfile_list = datasite_client.images.get_all()\n",
    "dockerfile_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert not isinstance(dockerfile_list, sy.SyftError), str(dockerfile_list)\n",
    "assert len(dockerfile_list) == 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27",
   "metadata": {},
   "outputs": [],
   "source": [
    "workerimage = next(\n",
    "    (\n",
    "        image\n",
    "        for image in dockerfile_list\n",
    "        if not image.is_prebuilt and image.config.dockerfile == custom_dockerfile_str\n",
    "    ),\n",
    "    None,\n",
    ")\n",
    "\n",
    "assert isinstance(workerimage, SyftWorkerImage), str(workerimage)\n",
    "workerimage"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "28",
   "metadata": {},
   "source": [
    "#### Add External Registry in Syft"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29",
   "metadata": {},
   "outputs": [],
   "source": [
    "external_registry = os.getenv(\"EXTERNAL_REGISTRY\", registry)\n",
    "external_registry_username = os.getenv(\"EXTERNAL_REGISTRY_USERNAME\", None)\n",
    "external_registry_password = os.getenv(\"EXTERNAL_REGISTRY_PASSWORD\", None)\n",
    "\n",
    "# external_registry = input()\n",
    "# external_registry_username = getpass(\"Enter Registry Username\")\n",
    "# external_registry_password = getpass(\"Enter Registry Password\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30",
   "metadata": {},
   "outputs": [],
   "source": [
    "registry_add_result = datasite_client.api.services.image_registry.add(external_registry)\n",
    "registry_add_result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert isinstance(registry_add_result, sy.SyftSuccess), str(registry_add_result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32",
   "metadata": {},
   "outputs": [],
   "source": [
    "image_registry_list = datasite_client.api.services.image_registry.get_all()\n",
    "image_registry_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert not isinstance(image_registry_list, sy.SyftError), str(image_registry_list)\n",
    "assert len(image_registry_list) == 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34",
   "metadata": {},
   "outputs": [],
   "source": [
    "local_registry = image_registry_list[0]\n",
    "local_registry"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert isinstance(local_registry, SyftImageRegistry), str(local_registry)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36",
   "metadata": {},
   "outputs": [],
   "source": [
    "registry_uid = local_registry.id"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37",
   "metadata": {},
   "source": [
    "#### Build Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38",
   "metadata": {},
   "outputs": [],
   "source": [
    "docker_tag = \"openmined/custom-worker:0.7.8\"\n",
    "\n",
    "\n",
    "docker_build_result = datasite_client.api.services.worker_image.build(\n",
    "    image_uid=workerimage.id,\n",
    "    tag=docker_tag,\n",
    "    registry_uid=registry_uid,\n",
    ")\n",
    "docker_build_result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert not isinstance(docker_build_result, sy.SyftError), str(docker_build_result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40",
   "metadata": {},
   "outputs": [],
   "source": [
    "image_list = datasite_client.images.get_all()\n",
    "image_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41",
   "metadata": {},
   "outputs": [],
   "source": [
    "# we can also index with string using the repo_with_tag format\n",
    "workerimage = next((image for image in image_list if image.id == workerimage.id), None)\n",
    "workerimage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert workerimage is not None, str([image.__dict__ for image in image_list])\n",
    "assert workerimage.is_built is not None, str(workerimage)\n",
    "assert workerimage.built_at is not None, str(workerimage)\n",
    "assert workerimage.image_hash is not None, str(workerimage)\n",
    "assert image_list[workerimage.built_image_tag] == workerimage"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43",
   "metadata": {},
   "source": [
    "#### Push Image to Local Registry"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44",
   "metadata": {},
   "outputs": [],
   "source": [
    "push_result = None\n",
    "push_result = datasite_client.api.services.worker_image.push(\n",
    "    workerimage.id,\n",
    "    username=external_registry_username,\n",
    "    password=external_registry_password,\n",
    ")\n",
    "push_result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert isinstance(push_result, sy.SyftSuccess), str(push_result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46",
   "metadata": {},
   "outputs": [],
   "source": [
    "base_url = f\"http://{workerimage.image_identifier.registry_host}\"\n",
    "expected_tag = workerimage.image_identifier.tag\n",
    "\n",
    "repos = requests.get(f\"{base_url}/v2/_catalog\").json()[\"repositories\"]\n",
    "tags = requests.get(f\"{base_url}/v2/openmined/custom-worker/tags/list\").json()\n",
    "tags = tags[\"tags\"]\n",
    "\n",
    "assert (\n",
    "    \"openmined/custom-worker\" in repos\n",
    "), f\"'openmined/custom-worker' not uploaded to local registry | {repos}\"\n",
    "assert (\n",
    "    expected_tag in tags\n",
    "), f\"'openmined/custom-worker' with tag {expected_tag} not available | {tags}\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "47",
   "metadata": {},
   "source": [
    "#### Create Worker Pool From Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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 = datasite_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": "49",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert not isinstance(worker_pool_res, sy.SyftError), str(worker_pool_res)\n",
    "assert len(worker_pool_res) == 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50",
   "metadata": {},
   "outputs": [],
   "source": [
    "for status in worker_pool_res:\n",
    "    assert status.error is None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51",
   "metadata": {},
   "outputs": [],
   "source": [
    "worker_pool_list = datasite_client.worker_pools.get_all()\n",
    "worker_pool_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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": [
    "assert not isinstance(worker_pool_list, sy.SyftError), str(worker_pool_res)\n",
    "assert len(worker_pool_list) == 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54",
   "metadata": {},
   "outputs": [],
   "source": [
    "worker_pool = next(\n",
    "    (pool for pool in worker_pool_list if pool.name == worker_pool_name),\n",
    "    None,\n",
    ")\n",
    "\n",
    "assert worker_pool is not None, str(\n",
    "    [worker_pool.__dict__ for worker_pool in worker_pool_list]\n",
    ")\n",
    "assert len(worker_pool.workers) == 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55",
   "metadata": {},
   "outputs": [],
   "source": [
    "# We can filter pools based on the image id upon which the pools were built\n",
    "filtered_result = datasite_client.api.services.worker_pool.filter_by_image_id(\n",
    "    image_uid=workerimage.id\n",
    ")\n",
    "filtered_result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert not isinstance(filtered_result, sy.SyftError), str(filtered_result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57",
   "metadata": {},
   "outputs": [],
   "source": [
    "second_worker = worker_pool.workers[1]\n",
    "second_worker"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "58",
   "metadata": {},
   "source": [
    "#### Get Worker Logs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59",
   "metadata": {},
   "outputs": [],
   "source": [
    "worker_logs = datasite_client.api.services.worker.logs(\n",
    "    uid=second_worker.id,\n",
    ")\n",
    "worker_logs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert isinstance(worker_logs, str)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61",
   "metadata": {},
   "outputs": [],
   "source": [
    "worker_pool"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62",
   "metadata": {},
   "source": [
    "### Syft function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63",
   "metadata": {},
   "outputs": [],
   "source": [
    "data = np.array([1, 2, 3])\n",
    "data_action_obj = sy.ActionObject.from_obj(data)\n",
    "\n",
    "data_pointer = data_action_obj.send(datasite_client)\n",
    "data_pointer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64",
   "metadata": {},
   "outputs": [],
   "source": [
    "@sy.syft_function(\n",
    "    input_policy=sy.ExactMatch(x=data_pointer),\n",
    "    output_policy=sy.SingleExecutionExactOutput(),\n",
    "    worker_pool_name=worker_pool_name,\n",
    ")\n",
    "def custom_worker_func(x):\n",
    "    # third party\n",
    "    import pydicom\n",
    "\n",
    "    print(pydicom.__version__)\n",
    "    return {\"y\": x + 1}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65",
   "metadata": {},
   "outputs": [],
   "source": [
    "custom_worker_func"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert custom_worker_func.worker_pool_name == worker_pool.name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67",
   "metadata": {},
   "outputs": [],
   "source": [
    "request = datasite_client.code.request_code_execution(custom_worker_func)\n",
    "request"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68",
   "metadata": {},
   "outputs": [],
   "source": [
    "datasite_client.requests[-1].approve(approve_nested=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69",
   "metadata": {},
   "outputs": [],
   "source": [
    "job = datasite_client.code.custom_worker_func(x=data_pointer, blocking=False)\n",
    "job"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "70",
   "metadata": {},
   "outputs": [],
   "source": [
    "worker_pool = datasite_client.worker_pools[worker_pool_name]\n",
    "worker_pool"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71",
   "metadata": {},
   "outputs": [],
   "source": [
    "job.wait()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert job.status.value == \"completed\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73",
   "metadata": {},
   "outputs": [],
   "source": [
    "job_list = datasite_client.jobs.get_by_user_code_id(job.user_code_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert not isinstance(job_list, sy.SyftError), job_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75",
   "metadata": {},
   "outputs": [],
   "source": [
    "job_refresh = job_list[0]\n",
    "assert job_refresh.job_worker_id is not None, str([job.to_dict() for job in job_list])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Validate the result received from the syft function\n",
    "result = job.wait().get()\n",
    "result_matches = result[\"y\"] == data + 1\n",
    "assert result_matches.all()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Scale Down the workers\n",
    "custom_pool_scale_res = datasite_client.api.services.worker_pool.scale(\n",
    "    number=1, pool_name=worker_pool_name\n",
    ")\n",
    "assert not isinstance(custom_pool_scale_res, sy.SyftError), str(custom_pool_scale_res)\n",
    "custom_pool_scale_res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert len(datasite_client.worker_pools[worker_pool_name].worker_list) == 1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "79",
   "metadata": {},
   "source": [
    "#### Worker Pool and Image Creation Request/Approval"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80",
   "metadata": {},
   "outputs": [],
   "source": [
    "dockerfile_opendp = f\"\"\"\n",
    "FROM {registry}/{repo}:{tag}\n",
    "\n",
    "RUN uv pip install opendp\n",
    "\"\"\".strip()\n",
    "\n",
    "docker_config_opendp = sy.DockerWorkerConfig(dockerfile=dockerfile_opendp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "81",
   "metadata": {},
   "outputs": [],
   "source": [
    "submit_result = None\n",
    "submit_result = datasite_client.api.services.worker_image.submit(\n",
    "    worker_config=docker_config_opendp\n",
    ")\n",
    "submit_result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert isinstance(submit_result, sy.SyftSuccess), str(submit_result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83",
   "metadata": {},
   "outputs": [],
   "source": [
    "_images = datasite_client.images\n",
    "assert not isinstance(_images, sy.SyftError), str(_images)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "84",
   "metadata": {},
   "outputs": [],
   "source": [
    "workerimage_opendp = next(\n",
    "    (im for im in _images if im.config == docker_config_opendp),\n",
    "    None,\n",
    ")\n",
    "assert workerimage_opendp is not None, str([im.__dict__ for im in _images])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "85",
   "metadata": {},
   "source": [
    "##### Build image first then create pool"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86",
   "metadata": {},
   "outputs": [],
   "source": [
    "docker_tag_opendp = \"openmined/custom-worker-opendp:latest\"\n",
    "\n",
    "docker_build_result = datasite_client.api.services.worker_image.build(\n",
    "    image_uid=workerimage_opendp.id,\n",
    "    tag=docker_tag_opendp,\n",
    "    registry_uid=registry_uid,\n",
    ")\n",
    "\n",
    "docker_build_result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert isinstance(docker_build_result, sy.SyftSuccess), str(docker_build_result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88",
   "metadata": {},
   "outputs": [],
   "source": [
    "_images = datasite_client.images\n",
    "assert not isinstance(_images, sy.SyftError), str(_images)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89",
   "metadata": {},
   "outputs": [],
   "source": [
    "workerimage_opendp = next(\n",
    "    (image for image in _images if image.id == workerimage_opendp.id),\n",
    "    None,\n",
    ")\n",
    "assert workerimage_opendp is not None, str([image.__dict__ for image in _images])\n",
    "assert workerimage_opendp.is_built is not None, str(workerimage_opendp.__dict__)\n",
    "assert workerimage_opendp.built_at is not None, str(workerimage_opendp.__dict__)\n",
    "assert workerimage_opendp.image_hash is not None, str(workerimage_opendp.__dict__)\n",
    "\n",
    "assert _images[workerimage_opendp.built_image_tag] == workerimage_opendp, str(\n",
    "    workerimage_opendp\n",
    ")\n",
    "\n",
    "workerimage_opendp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Push OpenDP Image to registry\n",
    "push_result = None\n",
    "push_result = datasite_client.api.services.worker_image.push(\n",
    "    workerimage_opendp.id,\n",
    "    username=external_registry_username,\n",
    "    password=external_registry_password,\n",
    ")\n",
    "assert isinstance(push_result, sy.SyftSuccess), str(push_result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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 = datasite_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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert not isinstance(pool_create_request, sy.SyftError), str(pool_create_request)\n",
    "assert len(pool_create_request.changes) == 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93",
   "metadata": {},
   "outputs": [],
   "source": [
    "# get the pending request and approve it\n",
    "req_result = pool_create_request.approve(\n",
    "    registry_username=external_registry_username,\n",
    "    registry_password=external_registry_password,\n",
    ")\n",
    "req_result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert isinstance(req_result, sy.SyftSuccess), str(req_result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95",
   "metadata": {},
   "outputs": [],
   "source": [
    "pool_opendp = datasite_client.worker_pools[pool_name_opendp]\n",
    "assert not isinstance(pool_opendp, sy.SyftError), str(pool_opendp)\n",
    "assert len(pool_opendp.worker_list) == 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96",
   "metadata": {},
   "outputs": [],
   "source": [
    "worker_pool_list = datasite_client.worker_pools.get_all()\n",
    "\n",
    "assert not isinstance(worker_pool_list, sy.SyftError), str(worker_pool_list)\n",
    "assert len(worker_pool_list) == 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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": [
    "# Scale Down the workers\n",
    "opendp_pool_scale_res = datasite_client.api.services.worker_pool.scale(\n",
    "    number=1, pool_name=pool_name_opendp\n",
    ")\n",
    "assert not isinstance(opendp_pool_scale_res, sy.SyftError), str(opendp_pool_scale_res)\n",
    "opendp_pool_scale_res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert len(datasite_client.worker_pools[pool_name_opendp].worker_list) == 1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "100",
   "metadata": {},
   "source": [
    "Request to build the image and create the pool at the same time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "101",
   "metadata": {},
   "outputs": [],
   "source": [
    "dockerfile_recordlinkage = f\"\"\"\n",
    "FROM {registry}/{repo}:{tag}\n",
    "\n",
    "RUN uv pip install recordlinkage\n",
    "\"\"\".strip()\n",
    "\n",
    "docker_config_recordlinkage = sy.DockerWorkerConfig(dockerfile=dockerfile_recordlinkage)\n",
    "\n",
    "docker_tag_recordlinkage = \"openmined/custom-worker-recordlinkage:latest\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "102",
   "metadata": {},
   "outputs": [],
   "source": [
    "pool_name_recordlinkage = \"recordlinkage-pool\"\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 = datasite_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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "103",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert not isinstance(pool_image_create_request, sy.SyftError), str(\n",
    "    pool_image_create_request\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "104",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert len(pool_image_create_request.changes) == 2\n",
    "assert pool_image_create_request.changes[0].config == docker_config_recordlinkage\n",
    "assert pool_image_create_request.changes[1].num_workers == 2\n",
    "assert pool_image_create_request.changes[1].pool_name == pool_name_recordlinkage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "105",
   "metadata": {},
   "outputs": [],
   "source": [
    "req_result = pool_image_create_request.approve(\n",
    "    registry_username=external_registry_username,\n",
    "    registry_password=external_registry_password,\n",
    ")\n",
    "req_result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "106",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert isinstance(req_result, sy.SyftSuccess), str(req_result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "107",
   "metadata": {},
   "outputs": [],
   "source": [
    "_requests = datasite_client.requests\n",
    "assert not isinstance(_requests, sy.SyftError), str(_requests)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "108",
   "metadata": {},
   "outputs": [],
   "source": [
    "pool_image_create_request = next(\n",
    "    (req for req in _requests if req.id == pool_image_create_request.id),\n",
    "    None,\n",
    ")\n",
    "assert pool_image_create_request is not None, str([req.__dict__ for req in _requests])\n",
    "assert pool_image_create_request.status.value == 2, str(pool_image_create_request)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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": [
    "datasite_client.images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "111",
   "metadata": {},
   "outputs": [],
   "source": [
    "image_exists = False\n",
    "for im in datasite_client.images.get_all():\n",
    "    if (\n",
    "        im.image_identifier\n",
    "        and im.image_identifier.repo_with_tag == docker_tag_recordlinkage\n",
    "    ):\n",
    "        image_exists = True\n",
    "\n",
    "assert image_exists, str([im.__dict__ for im in _images])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "112",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert datasite_client.worker_pools[pool_name_recordlinkage]\n",
    "assert len(datasite_client.worker_pools[pool_name_recordlinkage].worker_list) == 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "113",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Scale down the workers\n",
    "recordlinkage_pool_scale_res = datasite_client.api.services.worker_pool.scale(\n",
    "    number=1, pool_name=pool_name_recordlinkage\n",
    ")\n",
    "assert not isinstance(recordlinkage_pool_scale_res, sy.SyftError), str(\n",
    "    recordlinkage_pool_scale_res\n",
    ")\n",
    "recordlinkage_pool_scale_res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "114",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert len(datasite_client.worker_pools[pool_name_recordlinkage].worker_list) == 1"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}