notebooks/api/0.8/05-custom-policy.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "SYFT_VERSION = \">=0.8.2.b0,<0.9\"\n",
    "package_string = f'\"syft{SYFT_VERSION}\"'\n",
    "# %pip install {package_string} -q"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# third party\n",
    "import numpy as np\n",
    "\n",
    "# syft absolute\n",
    "import syft as sy\n",
    "\n",
    "sy.requires(SYFT_VERSION)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "node = sy.orchestra.launch(name=\"test-domain-1\", port=\"auto\", dev_mode=True, reset=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "domain_client = node.login(email=\"info@openmined.org\", password=\"changethis\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "domain_client.register(\n",
    "    email=\"newuser@openmined.org\", name=\"John Doe\", password=\"pw\", password_verify=\"pw\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "client_low_ds = node.login(email=\"newuser@openmined.org\", password=\"pw\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# stdlib\n",
    "from typing import Any"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "class RepeatedCallPolicy(sy.CustomOutputPolicy):\n",
    "    n_calls: int = 0\n",
    "    downloadable_output_args: list[str] = []\n",
    "    state: dict[Any, Any] = {}\n",
    "\n",
    "    def __init__(self, n_calls=1, downloadable_output_args: list[str] = None):\n",
    "        self.downloadable_output_args = (\n",
    "            downloadable_output_args if downloadable_output_args is not None else []\n",
    "        )\n",
    "        self.n_calls = n_calls\n",
    "        self.state = {\"counts\": 0}\n",
    "\n",
    "    def public_state(self):\n",
    "        return self.state[\"counts\"]\n",
    "\n",
    "    def update_policy(self, context, outputs):\n",
    "        self.state[\"counts\"] += 1\n",
    "\n",
    "    def apply_to_output(self, context, outputs, update_policy=True):\n",
    "        if hasattr(outputs, \"syft_action_data\"):\n",
    "            outputs = outputs.syft_action_data\n",
    "        output_dict = {}\n",
    "        if self.state[\"counts\"] < self.n_calls:\n",
    "            for output_arg in self.downloadable_output_args:\n",
    "                output_dict[output_arg] = outputs[output_arg]\n",
    "            if update_policy:\n",
    "                self.update_policy(context, outputs)\n",
    "        else:\n",
    "            return None\n",
    "        return output_dict\n",
    "\n",
    "    def _is_valid(self, context):\n",
    "        return self.state[\"counts\"] < self.n_calls"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "policy = RepeatedCallPolicy(n_calls=1, downloadable_output_args=[\"y\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "policy.n_calls"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "policy.downloadable_output_args"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "policy.init_kwargs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(policy.init_kwargs)\n",
    "a_obj = sy.ActionObject.from_obj({\"y\": [1, 2, 3]})\n",
    "x = policy.apply_to_output(None, a_obj)\n",
    "x[\"y\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "policy.n_calls"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "x = np.array([1, 2, 3])\n",
    "x_pointer = sy.ActionObject.from_obj(x)\n",
    "x_pointer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "domain_client.api.services.action.set(x_pointer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "metadata": {},
   "outputs": [],
   "source": [
    "# third party\n",
    "from result import Err\n",
    "from result import Ok\n",
    "\n",
    "# syft absolute\n",
    "from syft.client.api import AuthedServiceContext\n",
    "from syft.client.api import NodeIdentity\n",
    "\n",
    "\n",
    "class CustomExactMatch(sy.CustomInputPolicy):\n",
    "    def __init__(self, *args: Any, **kwargs: Any) -> None:\n",
    "        pass\n",
    "\n",
    "    def filter_kwargs(self, kwargs, context, code_item_id):\n",
    "        # stdlib\n",
    "\n",
    "        try:\n",
    "            allowed_inputs = self.allowed_ids_only(\n",
    "                allowed_inputs=self.inputs, kwargs=kwargs, context=context\n",
    "            )\n",
    "            results = self.retrieve_from_db(\n",
    "                code_item_id=code_item_id,\n",
    "                allowed_inputs=allowed_inputs,\n",
    "                context=context,\n",
    "            )\n",
    "        except Exception as e:\n",
    "            return Err(str(e))\n",
    "        return results\n",
    "\n",
    "    def retrieve_from_db(self, code_item_id, allowed_inputs, context):\n",
    "        # syft absolute\n",
    "        from syft import NodeType\n",
    "        from syft.service.action.action_object import TwinMode\n",
    "\n",
    "        action_service = context.node.get_service(\"actionservice\")\n",
    "        code_inputs = {}\n",
    "\n",
    "        # When we are retrieving the code from the database, we need to use the node's\n",
    "        # verify key as the credentials. This is because when we approve the code, we\n",
    "        # we allow the private data to be used only for this specific code.\n",
    "        # but we are not modifying the permissions of the private data\n",
    "\n",
    "        root_context = AuthedServiceContext(\n",
    "            node=context.node, credentials=context.node.verify_key\n",
    "        )\n",
    "        if context.node.node_type == NodeType.DOMAIN:\n",
    "            for var_name, arg_id in allowed_inputs.items():\n",
    "                kwarg_value = action_service._get(\n",
    "                    context=root_context,\n",
    "                    uid=arg_id,\n",
    "                    twin_mode=TwinMode.NONE,\n",
    "                    has_permission=True,\n",
    "                )\n",
    "                if kwarg_value.is_err():\n",
    "                    return Err(kwarg_value.err())\n",
    "                code_inputs[var_name] = kwarg_value.ok()\n",
    "\n",
    "        elif context.node.node_type == NodeType.ENCLAVE:\n",
    "            dict_object = action_service.get(context=root_context, uid=code_item_id)\n",
    "            if dict_object.is_err():\n",
    "                return Err(dict_object.err())\n",
    "            for value in dict_object.ok().syft_action_data.values():\n",
    "                code_inputs.update(value)\n",
    "\n",
    "        else:\n",
    "            raise Exception(\n",
    "                f\"Invalid Node Type for Code Submission:{context.node.node_type}\"\n",
    "            )\n",
    "        return Ok(code_inputs)\n",
    "\n",
    "    def allowed_ids_only(\n",
    "        self,\n",
    "        allowed_inputs,\n",
    "        kwargs,\n",
    "        context,\n",
    "    ):\n",
    "        # syft absolute\n",
    "        from syft import NodeType\n",
    "        from syft import UID\n",
    "\n",
    "        if context.node.node_type == NodeType.DOMAIN:\n",
    "            node_identity = NodeIdentity(\n",
    "                node_name=context.node.name,\n",
    "                node_id=context.node.id,\n",
    "                verify_key=context.node.signing_key.verify_key,\n",
    "            )\n",
    "            allowed_inputs = allowed_inputs.get(node_identity, {})\n",
    "        elif context.node.node_type == NodeType.ENCLAVE:\n",
    "            base_dict = {}\n",
    "            for key in allowed_inputs.values():\n",
    "                base_dict.update(key)\n",
    "            allowed_inputs = base_dict\n",
    "        else:\n",
    "            raise Exception(\n",
    "                f\"Invalid Node Type for Code Submission:{context.node.node_type}\"\n",
    "            )\n",
    "        filtered_kwargs = {}\n",
    "        for key in allowed_inputs.keys():\n",
    "            if key in kwargs:\n",
    "                value = kwargs[key]\n",
    "                uid = value\n",
    "                if not isinstance(uid, UID):\n",
    "                    uid = getattr(value, \"id\", None)\n",
    "\n",
    "                if uid != allowed_inputs[key]:\n",
    "                    raise Exception(\n",
    "                        f\"Input with uid: {uid} for `{key}` not in allowed inputs: {allowed_inputs}\"\n",
    "                    )\n",
    "                filtered_kwargs[key] = value\n",
    "        return filtered_kwargs\n",
    "\n",
    "    def _is_valid(\n",
    "        self,\n",
    "        context,\n",
    "        usr_input_kwargs,\n",
    "        code_item_id,\n",
    "    ):\n",
    "        filtered_input_kwargs = self.filter_kwargs(\n",
    "            kwargs=usr_input_kwargs,\n",
    "            context=context,\n",
    "            code_item_id=code_item_id,\n",
    "        )\n",
    "\n",
    "        if filtered_input_kwargs.is_err():\n",
    "            return filtered_input_kwargs\n",
    "\n",
    "        filtered_input_kwargs = filtered_input_kwargs.ok()\n",
    "\n",
    "        expected_input_kwargs = set()\n",
    "        for _inp_kwargs in self.inputs.values():\n",
    "            for k in _inp_kwargs.keys():\n",
    "                if k not in usr_input_kwargs:\n",
    "                    return Err(f\"Function missing required keyword argument: '{k}'\")\n",
    "            expected_input_kwargs.update(_inp_kwargs.keys())\n",
    "\n",
    "        permitted_input_kwargs = list(filtered_input_kwargs.keys())\n",
    "        not_approved_kwargs = set(expected_input_kwargs) - set(permitted_input_kwargs)\n",
    "        if len(not_approved_kwargs) > 0:\n",
    "            return Err(\n",
    "                f\"Input arguments: {not_approved_kwargs} to the function are not approved yet.\"\n",
    "            )\n",
    "        return Ok(True)\n",
    "\n",
    "\n",
    "def allowed_ids_only(\n",
    "    self,\n",
    "    allowed_inputs,\n",
    "    kwargs,\n",
    "    context,\n",
    "):\n",
    "    # syft absolute\n",
    "    from syft import NodeType\n",
    "    from syft import UID\n",
    "    from syft.client.api import NodeIdentity\n",
    "\n",
    "    if context.node.node_type == NodeType.DOMAIN:\n",
    "        node_identity = NodeIdentity(\n",
    "            node_name=context.node.name,\n",
    "            node_id=context.node.id,\n",
    "            verify_key=context.node.signing_key.verify_key,\n",
    "        )\n",
    "        allowed_inputs = allowed_inputs.get(node_identity, {})\n",
    "    elif context.node.node_type == NodeType.ENCLAVE:\n",
    "        base_dict = {}\n",
    "        for key in allowed_inputs.values():\n",
    "            base_dict.update(key)\n",
    "        allowed_inputs = base_dict\n",
    "    else:\n",
    "        raise Exception(\n",
    "            f\"Invalid Node Type for Code Submission:{context.node.node_type}\"\n",
    "        )\n",
    "    filtered_kwargs = {}\n",
    "    for key in allowed_inputs.keys():\n",
    "        if key in kwargs:\n",
    "            value = kwargs[key]\n",
    "            uid = value\n",
    "            if not isinstance(uid, UID):\n",
    "                uid = getattr(value, \"id\", None)\n",
    "\n",
    "            if uid != allowed_inputs[key]:\n",
    "                raise Exception(\n",
    "                    f\"Input with uid: {uid} for `{key}` not in allowed inputs: {allowed_inputs}\"\n",
    "                )\n",
    "            filtered_kwargs[key] = value\n",
    "    return filtered_kwargs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "@sy.syft_function(\n",
    "    input_policy=CustomExactMatch(x=x_pointer),\n",
    "    output_policy=RepeatedCallPolicy(n_calls=10, downloadable_output_args=[\"y\"]),\n",
    ")\n",
    "def func(x):\n",
    "    return {\"y\": x + 1}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "request = client_low_ds.code.request_code_execution(func)\n",
    "request"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19",
   "metadata": {},
   "outputs": [],
   "source": [
    "request_id = request.id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "client_low_ds.code.get_all()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21",
   "metadata": {},
   "outputs": [],
   "source": [
    "for request in domain_client.requests:\n",
    "    if request.id == request_id:\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "func = request.code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "result = func.unsafe_function(x=x_pointer)\n",
    "result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# syft absolute\n",
    "from syft.service.response import SyftError\n",
    "\n",
    "final_result = request.accept_by_depositing_result(result)\n",
    "assert isinstance(final_result, SyftError)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25",
   "metadata": {},
   "outputs": [],
   "source": [
    "request.approve()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "res_ptr = client_low_ds.code.func(x=x_pointer)\n",
    "res_ptr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27",
   "metadata": {},
   "outputs": [],
   "source": [
    "res = res_ptr.get()\n",
    "res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "assert (res[\"y\"] == np.array([2, 3, 4])).all()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "assert set(res.keys()) == set(\"y\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "for code in domain_client.code.get_all():\n",
    "    if code.service_func_name == \"func\":\n",
    "        break\n",
    "print(code.output_policy.state)\n",
    "assert code.output_policy.state == {\"counts\": 1}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "if node.node_type.value == \"python\":\n",
    "    node.land()"
   ]
  }
 ],
 "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.0rc1"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}