notebooks/api/0.8/10-container-images.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",
    "import time\n",
    "\n",
    "# third party\n",
    "import docker\n",
    "import numpy as np\n",
    "\n",
    "# syft absolute\n",
    "import syft as sy\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": [
    "# Uncomment this to run the whole docker based custom workers\n",
    "# os.environ[\"ORCHESTRA_DEPLOYMENT_TYPE\"] = \"container_stack\"\n",
    "# os.environ[\"DEV_MODE\"] = \"True\"\n",
    "\n",
    "\n",
    "# Disable inmemory worker for container stack\n",
    "running_as_container = os.environ.get(\"ORCHESTRA_DEPLOYMENT_TYPE\") in (\n",
    "    \"container_stack\",\n",
    "    \"k8s\",\n",
    ")\n",
    "in_memory_workers = not running_as_container"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "domain = sy.orchestra.launch(\n",
    "    name=\"test-domain-1\",\n",
    "    dev_mode=True,\n",
    "    create_producer=True,\n",
    "    in_memory_workers=in_memory_workers,\n",
    "    reset=True,\n",
    "    port=8081,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "domain_client = domain.login(email=\"info@openmined.org\", password=\"changethis\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5",
   "metadata": {},
   "source": [
    "We should see a default worker pool"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "domain_client.worker_pools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "syft_base_worker_tag = (\n",
    "    \"local-dev\"\n",
    "    if (bool(os.environ[\"DEV_MODE\"]) and running_as_container)\n",
    "    else sy.__version__\n",
    ")\n",
    "syft_base_worker_tag"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8",
   "metadata": {},
   "source": [
    "#### Submit Dockerfile"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "custom_dockerfile_str = f\"\"\"\n",
    "FROM openmined/grid-backend:{syft_base_worker_tag}\n",
    "\n",
    "RUN pip install pydicom\n",
    "\n",
    "\"\"\".strip()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "docker_config = sy.DockerWorkerConfig(dockerfile=custom_dockerfile_str)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "# test image build locally\n",
    "test_build_res = docker_config.test_image_build(tag=\"openmined/custom-worker:0.7.8\")\n",
    "test_build_res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert isinstance(test_build_res, sy.SyftSuccess), str(test_build_res)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert docker_config.dockerfile == custom_dockerfile_str"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "metadata": {},
   "outputs": [],
   "source": [
    "submit_result = domain_client.api.services.worker_image.submit_dockerfile(\n",
    "    docker_config=docker_config\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "submit_result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert isinstance(submit_result, sy.SyftSuccess), str(submit_result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17",
   "metadata": {},
   "outputs": [],
   "source": [
    "dockerfile_list = domain_client.images.get_all()\n",
    "dockerfile_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert len(domain_client.images.get_all()) == 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19",
   "metadata": {},
   "outputs": [],
   "source": [
    "workerimage: SyftWorkerImage = None\n",
    "for image in dockerfile_list:\n",
    "    if not image.is_prebuilt and image.config.dockerfile == custom_dockerfile_str:\n",
    "        workerimage = image\n",
    "        break\n",
    "\n",
    "assert isinstance(workerimage, SyftWorkerImage), str(workerimage)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20",
   "metadata": {},
   "outputs": [],
   "source": [
    "workerimage"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21",
   "metadata": {},
   "source": [
    "#### Setup Local Registry"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22",
   "metadata": {},
   "outputs": [],
   "source": [
    "# third party\n",
    "\n",
    "\n",
    "class LocalRegistryContainer:\n",
    "    def __init__(self):\n",
    "        self.name = \"local_registry\"\n",
    "        self.client = docker.from_env()\n",
    "\n",
    "    def start(self, host_port=5678):\n",
    "        existing = self.get()\n",
    "        if existing:\n",
    "            return existing\n",
    "\n",
    "        result = self.client.containers.run(\n",
    "            \"registry:2\",\n",
    "            name=self.name,\n",
    "            detach=True,\n",
    "            ports={\"5000/tcp\": host_port},\n",
    "            labels={\"orgs.openmined.syft\": \"local-registry\"},\n",
    "        )\n",
    "\n",
    "        return result\n",
    "\n",
    "    def teardown(self):\n",
    "        existing = self.get()\n",
    "        if existing:\n",
    "            existing.stop()\n",
    "            existing.remove()\n",
    "\n",
    "    def get(self):\n",
    "        try:\n",
    "            result = self.client.containers.get(self.name)\n",
    "            if result.status == \"running\":\n",
    "                return result\n",
    "        except docker.errors.NotFound:\n",
    "            return None\n",
    "\n",
    "\n",
    "local_registry_container = LocalRegistryContainer()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23",
   "metadata": {},
   "source": [
    "#### Add Local Registry in Syft"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24",
   "metadata": {},
   "outputs": [],
   "source": [
    "registry_add_result = domain_client.api.services.image_registry.add(\"localhost:5678\")\n",
    "registry_add_result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert isinstance(registry_add_result, sy.SyftSuccess), str(registry_add_result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26",
   "metadata": {},
   "outputs": [],
   "source": [
    "images = domain_client.api.services.image_registry.get_all()\n",
    "assert len(images) == 1\n",
    "images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27",
   "metadata": {},
   "outputs": [],
   "source": [
    "local_registry = images[0]\n",
    "local_registry"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert isinstance(local_registry, SyftImageRegistry), str(local_registry)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "29",
   "metadata": {},
   "source": [
    "#### Build Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30",
   "metadata": {},
   "outputs": [],
   "source": [
    "pull = False if syft_base_worker_tag == \"local-dev\" else True\n",
    "pull"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31",
   "metadata": {},
   "outputs": [],
   "source": [
    "docker_tag = \"openmined/custom-worker:0.7.8\"\n",
    "\n",
    "registry_uid = local_registry.id if running_as_container else local_registry.id\n",
    "\n",
    "docker_build_result = domain_client.api.services.worker_image.build(\n",
    "    image_uid=workerimage.id,\n",
    "    tag=docker_tag,\n",
    "    registry_uid=registry_uid,\n",
    "    pull=pull,\n",
    ")\n",
    "docker_build_result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32",
   "metadata": {},
   "outputs": [],
   "source": [
    "workerimage.config.dockerfile"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert isinstance(docker_build_result, sy.SyftSuccess), str(docker_build_result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34",
   "metadata": {},
   "outputs": [],
   "source": [
    "image_list = domain_client.images.get_all()\n",
    "image_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35",
   "metadata": {},
   "outputs": [],
   "source": [
    "for image in image_list:\n",
    "    if image.id == workerimage.id:\n",
    "        workerimage = (\n",
    "            image  # we can also index with string using the repo_with_tag format\n",
    "        )\n",
    "\n",
    "if running_as_container:\n",
    "    image_list[workerimage.built_image_tag]\n",
    "    assert image_list[workerimage.built_image_tag] == workerimage\n",
    "\n",
    "workerimage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_image_hash(tag) -> str:\n",
    "    client = docker.from_env()\n",
    "    try:\n",
    "        image = client.images.get(tag)\n",
    "        return image.id\n",
    "    except docker.errors.ImageNotFound:\n",
    "        return None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37",
   "metadata": {},
   "outputs": [],
   "source": [
    "if running_as_container:\n",
    "    assert workerimage.image_hash == get_image_hash(\n",
    "        workerimage.built_image_tag\n",
    "    ), \"Worker Image image_hash does not match with built image hash\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38",
   "metadata": {},
   "source": [
    "#### Push Image to Local Registry"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39",
   "metadata": {},
   "outputs": [],
   "source": [
    "push_result = None\n",
    "if running_as_container:\n",
    "    # stdlib\n",
    "    from time import sleep\n",
    "\n",
    "    local_registry_container.start()\n",
    "    sleep(5)\n",
    "\n",
    "    push_result = domain_client.api.services.worker_image.push(workerimage.id)\n",
    "    assert isinstance(push_result, sy.SyftSuccess), str(push_result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40",
   "metadata": {},
   "outputs": [],
   "source": [
    "push_result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41",
   "metadata": {},
   "outputs": [],
   "source": [
    "if running_as_container:\n",
    "    # third party\n",
    "    import requests\n",
    "\n",
    "    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": "42",
   "metadata": {},
   "source": [
    "#### Delete locally built image to force pull from local registry"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43",
   "metadata": {},
   "source": [
    "This should make the subsequent `worker_pool.launch` pull from registry at 'localhost:5678`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44",
   "metadata": {},
   "outputs": [],
   "source": [
    "# stdlib\n",
    "from time import sleep\n",
    "\n",
    "\n",
    "def remove_local_image(tag):\n",
    "    client = docker.from_env()\n",
    "    try:\n",
    "        client.images.remove(tag)\n",
    "    except docker.errors.ImageNotFound:\n",
    "        pass\n",
    "\n",
    "\n",
    "if running_as_container:\n",
    "    remove_local_image(workerimage.built_image_tag)\n",
    "    sleep(5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45",
   "metadata": {},
   "source": [
    "#### Create Worker Pool From Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46",
   "metadata": {},
   "outputs": [],
   "source": [
    "worker_pool_name = \"my_first_worker_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",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert len(worker_pool_res) == 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48",
   "metadata": {},
   "outputs": [],
   "source": [
    "for status in worker_pool_res:\n",
    "    assert status.error is None\n",
    "    if running_as_container:\n",
    "        assert status.worker.image.image_hash == get_image_hash(\n",
    "            workerimage.built_image_tag\n",
    "        ), \"Worker Pool Image image_hash does not match with built image hash\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49",
   "metadata": {},
   "outputs": [],
   "source": [
    "worker_pool_list = domain_client.worker_pools\n",
    "worker_pool_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert len(domain_client.worker_pools.get_all()) == 2\n",
    "worker_pool = None\n",
    "for pool in worker_pool_list:\n",
    "    if pool.name == worker_pool_name:\n",
    "        worker_pool = pool\n",
    "        break\n",
    "assert worker_pool is not None\n",
    "assert len(worker_pool.workers) == 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51",
   "metadata": {},
   "outputs": [],
   "source": [
    "# We can filter pools based on the image id upon which the pools were built\n",
    "domain_client.api.services.worker_pool.filter_by_image_id(image_uid=workerimage.id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Delete the second worker\n",
    "second_worker = worker_pool.workers[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53",
   "metadata": {},
   "outputs": [],
   "source": [
    "second_worker"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54",
   "metadata": {},
   "source": [
    "#### Get Worker Logs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55",
   "metadata": {},
   "outputs": [],
   "source": [
    "raw_worker_logs = domain_client.api.services.worker.logs(\n",
    "    uid=second_worker.id,\n",
    "    raw=True,\n",
    ")\n",
    "raw_worker_logs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert isinstance(raw_worker_logs, bytes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57",
   "metadata": {},
   "outputs": [],
   "source": [
    "worker_logs = domain_client.api.services.worker.logs(\n",
    "    uid=second_worker.id,\n",
    ")\n",
    "worker_logs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert isinstance(worker_logs, str)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "59",
   "metadata": {},
   "source": [
    "#### Delete Worker from Pool"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60",
   "metadata": {},
   "outputs": [],
   "source": [
    "worker_delete_res = domain_client.api.services.worker.delete(\n",
    "    uid=second_worker.id,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61",
   "metadata": {},
   "outputs": [],
   "source": [
    "worker_delete_res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert isinstance(worker_delete_res, sy.SyftSuccess), str(worker_delete_res)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Refetch the worker pool\n",
    "# Ensure that the deleted worker's id is not present\n",
    "for pool in domain_client.api.services.worker_pool.get_all():\n",
    "    if pool.name == worker_pool_name:\n",
    "        worker_pool = pool\n",
    "assert len(worker_pool.workers) == 2\n",
    "for worker in worker_pool.workers:\n",
    "    assert second_worker.id != worker.id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64",
   "metadata": {},
   "outputs": [],
   "source": [
    "worker_pool"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "65",
   "metadata": {},
   "source": [
    "### Syft function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66",
   "metadata": {},
   "outputs": [],
   "source": [
    "data = np.array([1, 2, 3])\n",
    "data_action_obj = sy.ActionObject.from_obj(data)\n",
    "\n",
    "data_pointer = domain_client.api.services.action.set(data_action_obj)\n",
    "data_pointer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67",
   "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",
    "    return {\"y\": x + 1}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68",
   "metadata": {},
   "outputs": [],
   "source": [
    "custom_worker_func"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert custom_worker_func.worker_pool_name == worker_pool.name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "70",
   "metadata": {},
   "outputs": [],
   "source": [
    "request = domain_client.code.request_code_execution(custom_worker_func)\n",
    "request"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71",
   "metadata": {},
   "outputs": [],
   "source": [
    "domain_client.requests[-1].approve(approve_nested=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72",
   "metadata": {},
   "outputs": [],
   "source": [
    "job = domain_client.code.custom_worker_func(x=data_pointer, blocking=False)\n",
    "job"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73",
   "metadata": {},
   "outputs": [],
   "source": [
    "worker_pool = domain_client.worker_pools[worker_pool_name]\n",
    "worker_pool"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74",
   "metadata": {},
   "outputs": [],
   "source": [
    "job.wait()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert job.status.value == \"completed\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76",
   "metadata": {},
   "outputs": [],
   "source": [
    "job = domain_client.jobs[-1]\n",
    "job"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77",
   "metadata": {},
   "outputs": [],
   "source": [
    "job.job_worker_id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Disabling it due to Race Condition Error\n",
    "# assert job.job_worker_id is not None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "79",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Sleeping so that consumer state is updated\n",
    "time.sleep(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Once the work is done by the worker, its state is returned to idle again.\n",
    "consuming_worker_is_now_idle = False\n",
    "for worker in domain_client.worker_pools[worker_pool_name].workers:\n",
    "    if worker.id == job.job_worker_id:\n",
    "        consuming_worker_is_now_idle = worker.consumer_state.value.lower() == \"idle\"\n",
    "\n",
    "assert consuming_worker_is_now_idle is True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "81",
   "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": "markdown",
   "id": "82",
   "metadata": {},
   "source": [
    "#### Worker Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83",
   "metadata": {},
   "outputs": [],
   "source": [
    "# delete the remaining workers\n",
    "for worker in worker_pool.workers:\n",
    "    res = domain_client.api.services.worker.delete(\n",
    "        uid=worker.id,\n",
    "    )\n",
    "    assert isinstance(res, sy.SyftSuccess), str(res)\n",
    "\n",
    "# Adding some sleep to allow containers to be fully removed,\n",
    "# before removing the image\n",
    "time.sleep(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "84",
   "metadata": {},
   "outputs": [],
   "source": [
    "delete_res = domain_client.api.services.worker_image.remove(workerimage.id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85",
   "metadata": {},
   "outputs": [],
   "source": [
    "delete_res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Since the containers are delete, we should be able to delete the image\n",
    "assert isinstance(delete_res, sy.SyftSuccess), str(delete_res)\n",
    "delete_res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87",
   "metadata": {},
   "outputs": [],
   "source": [
    "if running_as_container:\n",
    "    local_registry_container.teardown()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88",
   "metadata": {},
   "source": [
    "#### Worker Pool and Image Creation Request/Approval"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89",
   "metadata": {},
   "outputs": [],
   "source": [
    "custom_dockerfile_str_2 = f\"\"\"\n",
    "FROM openmined/grid-backend:{syft_base_worker_tag}\n",
    "\n",
    "RUN pip install opendp\n",
    "\"\"\".strip()\n",
    "\n",
    "docker_config_2 = sy.DockerWorkerConfig(dockerfile=custom_dockerfile_str_2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90",
   "metadata": {},
   "outputs": [],
   "source": [
    "submit_result = domain_client.api.services.worker_image.submit_dockerfile(\n",
    "    docker_config=docker_config_2\n",
    ")\n",
    "submit_result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91",
   "metadata": {},
   "outputs": [],
   "source": [
    "domain_client.images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92",
   "metadata": {},
   "outputs": [],
   "source": [
    "# get the image that's not built\n",
    "workerimage_2 = None\n",
    "for im in domain_client.images:\n",
    "    if im.config == docker_config_2:\n",
    "        workerimage_2 = im"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "93",
   "metadata": {},
   "source": [
    "##### Build image first then create pool"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94",
   "metadata": {},
   "outputs": [],
   "source": [
    "docker_tag_2 = \"openmined/custom-worker-opendp:latest\"\n",
    "\n",
    "docker_build_result = domain_client.api.services.worker_image.build(\n",
    "    image_uid=workerimage_2.id,\n",
    "    tag=docker_tag_2,\n",
    "    pull=pull,\n",
    ")\n",
    "docker_build_result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95",
   "metadata": {},
   "outputs": [],
   "source": [
    "pool_create_request = domain_client.api.services.worker_pool.pool_creation_request(\n",
    "    pool_name=\"first-opendp-pool\", num_workers=3, image_uid=workerimage_2.id\n",
    ")\n",
    "pool_create_request"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert len(pool_create_request.changes) == 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "97",
   "metadata": {},
   "outputs": [],
   "source": [
    "# get the pending request and approve it\n",
    "req_result = pool_create_request.approve()\n",
    "req_result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert isinstance(req_result, sy.SyftSuccess), str(req_result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert domain_client.worker_pools[\"first-opendp-pool\"]\n",
    "assert len(domain_client.worker_pools[\"first-opendp-pool\"].worker_list) == 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "100",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert len(domain_client.worker_pools.get_all()) == 3"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "101",
   "metadata": {},
   "source": [
    "##### Request to build the image and create the pool at the same time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "102",
   "metadata": {},
   "outputs": [],
   "source": [
    "custom_dockerfile_str_3 = f\"\"\"\n",
    "FROM openmined/grid-backend:{syft_base_worker_tag}\n",
    "\n",
    "RUN pip install recordlinkage\n",
    "\"\"\".strip()\n",
    "\n",
    "docker_config_3 = sy.DockerWorkerConfig(dockerfile=custom_dockerfile_str_3)\n",
    "\n",
    "docker_tag_3 = \"openmined/custom-worker-recordlinkage:latest\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "103",
   "metadata": {},
   "outputs": [],
   "source": [
    "# syft absolute\n",
    "from syft.custom_worker.builder import CustomWorkerBuilder\n",
    "from syft.service.response import SyftError\n",
    "from syft.service.response import SyftSuccess\n",
    "\n",
    "\n",
    "def test_image_build(config: str, tag: str, pull: bool, **kwargs):\n",
    "    builder = CustomWorkerBuilder()\n",
    "    try:\n",
    "        result = builder.build_image(\n",
    "            config=config, tag=tag, pull=pull, rm=True, forcerm=True, **kwargs\n",
    "        )\n",
    "        return SyftSuccess(message=result.logs)\n",
    "    except Exception as e:\n",
    "        return SyftError(message=f\"Failed to build image !! Error: {str(e)}.\")\n",
    "\n",
    "\n",
    "test_build_res = test_image_build(config=docker_config_3, tag=docker_tag_3, pull=pull)\n",
    "assert isinstance(test_build_res, sy.SyftSuccess), str(test_build_res)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "104",
   "metadata": {},
   "outputs": [],
   "source": [
    "pool_image_create_request = (\n",
    "    domain_client.api.services.worker_pool.create_image_and_pool_request(\n",
    "        pool_name=\"recordlinkage-pool\",\n",
    "        num_workers=2,\n",
    "        tag=docker_tag_3,\n",
    "        config=docker_config_3,\n",
    "        reason=\"I want to do some more cool data science with PySyft and OpenDP\",\n",
    "        pull_image=pull,\n",
    "    )\n",
    ")\n",
    "pool_image_create_request"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "105",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert len(pool_image_create_request.changes) == 2\n",
    "assert pool_image_create_request.changes[0].config == docker_config_3\n",
    "assert pool_image_create_request.changes[1].num_workers == 2\n",
    "assert pool_image_create_request.changes[1].pool_name == \"recordlinkage-pool\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "106",
   "metadata": {},
   "outputs": [],
   "source": [
    "# get the pending request and approve it\n",
    "req_result = pool_image_create_request.approve()\n",
    "req_result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "107",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert isinstance(req_result, sy.SyftSuccess), str(req_result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "108",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get updated request object and status\n",
    "for req in domain_client.requests:\n",
    "    if req.id == pool_image_create_request.id:\n",
    "        pool_image_create_request = req\n",
    "\n",
    "assert pool_image_create_request.status.value == 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "109",
   "metadata": {},
   "outputs": [],
   "source": [
    "domain_client.images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "110",
   "metadata": {},
   "outputs": [],
   "source": [
    "image_exists = False\n",
    "for im in domain_client.images.get_all():\n",
    "    if im.image_identifier and im.image_identifier.repo_with_tag == docker_tag_3:\n",
    "        image_exists = True\n",
    "assert image_exists"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "111",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert domain_client.worker_pools[\"recordlinkage-pool\"]\n",
    "assert len(domain_client.worker_pools[\"recordlinkage-pool\"].worker_list) == 2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "112",
   "metadata": {},
   "source": [
    "#### Clean up workers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "113",
   "metadata": {},
   "outputs": [],
   "source": [
    "# delete the remaining workers\n",
    "for worker_pool in domain_client.worker_pools:\n",
    "    for worker in worker_pool.workers:\n",
    "        res = domain_client.api.services.worker.delete(uid=worker.id, force=True)\n",
    "        print(res)\n",
    "        assert isinstance(res, sy.SyftSuccess), str(res)\n",
    "\n",
    "# Adding some sleep to allow containers to be fully removed,\n",
    "# before removing the image\n",
    "time.sleep(15)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "114",
   "metadata": {},
   "outputs": [],
   "source": [
    "for worker_pool in domain_client.worker_pools:\n",
    "    assert len(worker_pool.worker_list) == 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "115",
   "metadata": {},
   "outputs": [],
   "source": [
    "domain.land()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "116",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}