{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tracking Resource Usage with PennyLane Device Tracker"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this notebook we will explore how to use the PennyLane device tracker feature with Amazon Braket.\n",
    "As is demonstrated in the [optimization of quantum circuits notebook](https://github.com/aws/amazon-braket-examples/blob/main/examples/pennylane/1_Parallelized_optimization_of_quantum_circuits/1_Parallelized_optimization_of_quantum_circuits.ipynb), computing gradients of quantum circuits involves multiple devices executions.\n",
    "This can lead to a large number of executions when optimizing quantum circuits.\n",
    "So to help users keep track of their usage, Amazon Braket works with PennyLane to record and make available useful information during the computation.\n",
    "The PennyLane device resource tracker keeps a record of the usage of a device, such as numbers of circuit evaluations and shots.\n",
    "Amazon Braket extends this information with task IDs and simulator duration to allow further tracking.\n",
    "The device tracker can be combined with additional logic to monitor and limit resource usage on devices."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pennylane as qml\n",
    "from pennylane import numpy as np\n",
    "\n",
    "import time"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before we show the device resource tracker, let's first set up a circuit for demonstration and experimenting.\n",
    "We will use a simple two qubit circuit with two parameters throughout this notebook.\n",
    "We will start by evaluating it using the local Braket simulator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "wires = 2  # Number of qubits\n",
    "\n",
    "dev = qml.device(\"braket.local.qubit\", wires=wires)\n",
    "\n",
    "def circuit(params):\n",
    "    qml.RX(params[0], wires=0)\n",
    "    qml.RY(params[1], wires=1)\n",
    "    qml.CNOT(wires=[0, 1])\n",
    "    return qml.expval(qml.PauliZ(1))\n",
    "\n",
    "qnode_local = qml.QNode(circuit, dev)\n",
    "params = np.array([0.1, 0.2], requires_grad=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The PennyLane Device Resource Tracker"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The PennyLane device resource tracker is a python context manager created by `qml.Tracker`.\n",
    "To use the device tracker to track a single evaluation of our simple circuit, we put the evaluation inside a `with` statement block.\n",
    "Only evaluations inside of the `with` block will have their usage recorded and accumulated."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Expectation value of circuit: 0.9751703272018161\n"
     ]
    }
   ],
   "source": [
    "with qml.Tracker(dev) as tracker:\n",
    "    print(\"Expectation value of circuit:\", qnode_local(params))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With this execution complete, all of the recorded information is available in `tracker`.\n",
    "There are three interfaces to access the data inside of the resource tracker.\n",
    "The full history of each update is available through `tracker.history`.\n",
    "Numerical values are accumulated, and each of their totals are in `tracker.totals`.\n",
    "And lastly, the most recent update to the tracker is kept in `tracker.latest`.\n",
    "\n",
    "Let's first look at the history of what was recorded by evaluating our circuit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'executions': [1], 'shots': [None], 'braket_task_id': ['3e22ac39-a7e7-425d-af6a-1d4f3cc7974f'], 'batches': [1], 'batch_len': [1]}\n"
     ]
    }
   ],
   "source": [
    "print(tracker.history)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see that this evaluation lead to a single circuit execution.\n",
    "This single execution had no shots, and was evaluated with a single batch of length 1.\n",
    "When using the Braket local simulator, the device tracker records the unique id that the simulator assigns to new tasks.\n",
    "\n",
    "Next, let's evaluate the gradient of this circuit and track the device usage."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Gradient of circuit: [-0.0978434  -0.19767681]\n",
      "{'executions': [1, 1, 1, 1, 1], 'shots': [None, None, None, None, None], 'braket_task_id': ['7602e62c-dd42-40bf-a214-147fdeec9c93', '9e924369-432e-4e86-bd28-c18a3bd7ea2d', 'fe68b242-166c-496a-b1e2-df72bc7eedc6', '8e1f1427-6b2a-4261-a229-dec502fc008f', 'c6426b8f-6b25-4c9f-abac-ad60ae96c4c8'], 'batches': [1, 1], 'batch_len': [1, 4]}\n"
     ]
    }
   ],
   "source": [
    "with qml.Tracker(dev) as tracker:\n",
    "    print(\"Gradient of circuit:\", qml.grad(qnode_local)(params))\n",
    "\n",
    "print(tracker.history)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Recall that evaluating the gradient of a quantum circuit will result in multiple circuit evaluations.\n",
    "Here we see that to calculate this gradient, 5 new tasks were created and recorded in the device tracker.\n",
    "For values which are numeric, such as the `executions`, the device tracker will accumulate all of the recorded data into `tracker.totals`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'executions': 5, 'batches': 2, 'batch_len': 5}\n"
     ]
    }
   ],
   "source": [
    "print(tracker.totals)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that `None` is not numeric, which is why the shots are not accumulated in `tracker.totals` in the previous example.\n",
    "If we instead compute the gradient with a finite number of shots, we will see the total shots recorded.\n",
    "Let's try the gradient with `shots=100`.\n",
    "Since there are 5 circuit evaluations, we should expect the total shots to be `500`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'executions': 5, 'shots': 500, 'batches': 2, 'batch_len': 5}\n"
     ]
    }
   ],
   "source": [
    "with qml.Tracker(dev) as tracker:\n",
    "    qml.grad(qnode_local)(params, shots=100)\n",
    "print(tracker.totals)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also inspect the most recently recorded data with `tracker.latest`.\n",
    "Note that not every field will be present in every update.\n",
    "The latest update in this gradient calculation only included `'batches'` and `'batch_len'`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'batches': 1, 'batch_len': 4}\n"
     ]
    }
   ],
   "source": [
    "print(tracker.latest)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using Amazon Braket Simulators"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The Amazon Braket on-demand simulators report additional information to the device tracker.\n",
    "Let's set up a remote device using Amazon Braket SV1, and create a new `QNode` to run our circuit on the on-demand simulator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "device_arn = \"arn:aws:braket:::device/quantum-simulator/amazon/sv1\"\n",
    "\n",
    "dev_remote = qml.device(\n",
    "    \"braket.aws.qubit\",\n",
    "    device_arn=device_arn,\n",
    "    wires=wires\n",
    ")\n",
    "\n",
    "qnode_remote = qml.QNode(circuit, dev_remote)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can send the same simple circuit from above to the AWS on-demand simulator and track the resource usage."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "with qml.Tracker(dev_remote) as tracker:\n",
    "    qnode_remote(params)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's take a look at what new fields are available in the tracker with an on-demand simulator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "executions\n",
      "shots\n",
      "braket_task_id\n",
      "braket_simulator_ms\n",
      "braket_simulator_billed_ms\n",
      "batches\n",
      "batch_len\n"
     ]
    }
   ],
   "source": [
    "for key in tracker.history.keys():\n",
    "    print(key)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `braket_task_id` field will contain unique IDs for each task that was created.\n",
    "For tasks that are created on the on-demand simulators, these IDs can be recorded and used to look up the task details in the AWS console.\n",
    "The `braket_simulator_ms` gives the duration in milliseconds of the portion of the task spent simulating our circuit.\n",
    "`braket_simulator_billed_ms` adjusts the simulation duration according to the minimum duration for billing.\n",
    "See [here](https://aws.amazon.com/braket/pricing/) for pricing details on Amazon Braket.\n",
    "\n",
    "These reported times are not the total time for a task to complete.\n",
    "We can compare the simulation time to the total task time for this circuit:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Remote device execution duration in seconds: 3.1985151767730713\n",
      "Simulation duration in seconds: 0.012\n"
     ]
    }
   ],
   "source": [
    "with qml.Tracker(dev_remote) as tracker:\n",
    "    t0 = time.time()\n",
    "    qnode_remote(params)\n",
    "    t1 = time.time()\n",
    "print(\"Remote device execution duration in seconds:\",t1 - t0)\n",
    "print(\"Simulation duration in seconds:\", tracker.totals[\"braket_simulator_ms\"] / 1000)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With such a small circuit in this example, only a very small fraction of the total time is spent running the simulation of the circuit."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using the Tracker to Limit Resource Usage"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we have seen above, it is possible to record usage of a device during execution with the resource tracker.\n",
    "With this information available during the computation, it is possible to handle the information in order to control the resource usage.\n",
    "However, this feature does not attempt to estimate how many resources may be used in the future. It is purely backward-looking.\n",
    "Thus any resource usage can only be acted on after the fact with this feature.\n",
    "\n",
    "In the [PennyLane getting started notebook](https://github.com/aws/amazon-braket-examples/blob/main/examples/pennylane/0_Getting_started/0_Getting_started.ipynb),\n",
    "we optimized a circuit by running the gradient descent optimizer for a chosen number of iterations.\n",
    "Instead of a given number of iterations, we may want to limit the optimization to something else such as the total number of shots or the simulator execution time.\n",
    "Here we will show two ways to use the PennyLane device tracker to limit an optimization on a resource which is being tracked.\n",
    "\n",
    "Let's suppose we wish to optimize our simple circuit, but want to limit ourselves to a particular number of circuit executions.\n",
    "First let's set up an optimizer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "opt = qml.GradientDescentOptimizer(stepsize=0.2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we run the optimizer, but on every step of the loop we check if our limit has been crossed.\n",
    "Once the total executions exceed our target limit, we will break out of the optimization and report our results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Completed 16 steps of optimization.\n",
      "Optimized cost: -0.27959738151720537\n",
      "Optimized parameters: [0.51754609 2.06622178]\n"
     ]
    }
   ],
   "source": [
    "max_iterations = 100\n",
    "execution_limit = 77\n",
    "params = np.array([0.1, 0.2], requires_grad=True)\n",
    "\n",
    "with qml.Tracker(dev) as tracker:\n",
    "    for i in range(max_iterations):\n",
    "        params, cost = opt.step_and_cost(qnode_local, params)\n",
    "        if tracker.totals[\"executions\"] > execution_limit:\n",
    "            break\n",
    "\n",
    "print(\"Completed\", i + 1, \"steps of optimization.\")\n",
    "print(\"Optimized cost:\", cost)\n",
    "print(\"Optimized parameters:\", params)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The resource usage is compared only after each iteration of the optimization.\n",
    "If each of the steps in your computation involve many circuit evaluations,\n",
    "your computation may still continue for many evaluations after your limit is hit.\n",
    "For our simple problem, every iteration leads to 5 circuit evaluations, so we do not stop right at our limit.\n",
    "We can see this by checking the tracker totals."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'executions': 80, 'batches': 32, 'batch_len': 80}\n"
     ]
    }
   ],
   "source": [
    "print(tracker.totals)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Take this into consideration when setting a threshold if you choose to limit your computation this way."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For another approach to limit resource usage, we can add a callback function to the device tracker.\n",
    "This function is called each time the information in the tracker is updated.\n",
    "This interface could also be used for monitoring large batches of executions.\n",
    "For example, let's print out the running total number of executions every time the tracker is updated."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total executions: 1\n",
      "Total executions: 1\n",
      "Total executions: 2\n",
      "Total executions: 3\n",
      "Total executions: 4\n",
      "Total executions: 5\n",
      "Total executions: 5\n"
     ]
    }
   ],
   "source": [
    "def log_executions(totals, history, latest):\n",
    "    print(\"Total executions:\",totals[\"executions\"])\n",
    "\n",
    "with qml.Tracker(dev, callback=log_executions) as tracker:\n",
    "    qml.grad(qnode_local)(params)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can use the callback feature of the device tracker to stop a computation after a threshold is breached.\n",
    "By throwing an exception from our callback function when the new data violates our limit, we can abort the currently running optimization.\n",
    "Let's try this approach with the same optimization problem."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Completed 15 steps of optimization.\n",
      "Optimized cost: -0.14125448783850006\n",
      "Optimized parameters: [0.5519834  1.90536732]\n"
     ]
    }
   ],
   "source": [
    "params = np.array([0.1, 0.2], requires_grad=True)\n",
    "\n",
    "def resource_threshold(totals, history, latest):\n",
    "    if totals[\"executions\"] > execution_limit:\n",
    "        raise ResourceWarning()\n",
    "\n",
    "with qml.Tracker(dev, callback=resource_threshold) as tracker:\n",
    "    try:\n",
    "        for i in range(max_iterations):\n",
    "            params,cost = opt.step_and_cost(qnode_local, params)\n",
    "    except ResourceWarning:\n",
    "        print(\"Completed\",i,\"steps of optimization.\")\n",
    "        print(\"Optimized cost:\", cost)\n",
    "        print(\"Optimized parameters:\", params)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This time one fewer optimization step is completed because the final step is aborted.\n",
    "We can look at the totals to see that the execution was stopped after reaching our set limit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'executions': 78, 'batches': 31, 'batch_len': 76}\n"
     ]
    }
   ],
   "source": [
    "print(tracker.totals)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Note** This approach will immediately terminate the optimization upon exceeding the resource limit.\n",
    "If the termination happens while in the middle of a step, that step will not finish.\n",
    "In this case, circuits will have been executed that will then have their results discarded."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this notebook, we have seen the resource usage recorded by the PennyLane device tracker and the extended information provided by Amazon Braket.\n",
    "We showed two different ways to use the resource information in the middle of the computation, and we were able to control the optimization routine to limit usage according to a preset threshold.\n",
    "If you are looking to explore more, try modifying the [QAOA notebook](https://github.com/aws/amazon-braket-examples/blob/main/examples/pennylane/2_Graph_optimization_with_QAOA/2_Graph_optimization_with_QAOA.ipynb) to track the use of the simulators in that example."
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "590fab68195cf107911461461f81d5c472d3d6127f579badfcfad30f03e5cab2"
  },
  "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}