2_contrast-performance.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8a12ad7e-45ea-4aa9-a337-fb6b85bdaab9",
   "metadata": {},
   "source": [
    "# Compare Performance between Parsl and Colmnea\n",
    "We want to see if Colmena's strategy of overlapping simulation and machine learning had any noticable effect.\n",
    "The primary way we can do that is to evaluate the time-to-solution for finding well-performing molecules."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "609a2850-9145-4ff3-ad28-d41a9e6b0ed3",
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "from matplotlib import pyplot as plt\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1fc05f1-3cb4-48d1-ba18-ffdd44b78765",
   "metadata": {},
   "source": [
    "## Load in the Runs\n",
    "We saved the performance data from the examples and will load them back into disk here."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c7c1ca78-21f7-4aa5-95ff-26df258e2c39",
   "metadata": {},
   "source": [
    "Colmena saves results as a JSON object. We can load them in via Pandas"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94fb3397-a616-4498-a137-c0fb9bc69323",
   "metadata": {},
   "outputs": [],
   "source": [
    "results = pd.read_json('run-data/batched-results.json', lines=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "575ed0be-b077-4bd2-84ff-a563930d2f3c",
   "metadata": {},
   "source": [
    "## Plot the Resource Utilization\n",
    "We want to make sure that Colmena is using nodes effectively by plotting the overall utilization.\n",
    "\n",
    "First step is to compute time a computation started, finished, and when it was received by the thinker."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc182e81-3136-4248-8e0b-bc3b6bbf0115",
   "metadata": {},
   "outputs": [],
   "source": [
    "start_time = results['time_compute_started'].min()\n",
    "results['walltime_start'] = results['time_compute_started'] - start_time\n",
    "results['walltime_finish'] = results['walltime_start'] + results['time_running']\n",
    "results['walltime_received'] = results['time_result_received'] - start_time"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "01242ee3-be0e-4dc2-b9db-73f3b4fc4d98",
   "metadata": {},
   "source": [
    "We next assign unique names and id numbers to each worker given their hostname and worker ID"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19a47f9b-0a7d-48af-a14f-b7310edb34d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "results['worker_name'] = results['worker_info'].apply(lambda x: f'{x[\"hostname\"]}-{x[\"PARSL_WORKER_RANK\"]}')\n",
    "unique_workers = sorted(set(results['worker_name']))\n",
    "results['worker_id'] = results['worker_name'].apply(unique_workers.index)\n",
    "n_workers = len(unique_workers)\n",
    "print(f'Found {n_workers} unique workers')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e728d34-f1b5-47c8-9fad-a02b3f9a653a",
   "metadata": {},
   "source": [
    "Finally, let's plot when the workers are being used and for which task"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f209c7a-5330-4eaf-900b-a117722cd054",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(3.5, 2.5))\n",
    "\n",
    "# Colors for each type of task\n",
    "method_color = {\n",
    "    'compute_vertical': 'teal',\n",
    "    'train_model': 'tomato',\n",
    "    'run_model': 'thistle'\n",
    "}\n",
    "\n",
    "# Labels for each type of task\n",
    "method_label = {\n",
    "    'compute_vertical': 'simulate',\n",
    "    'train_model': 'train', \n",
    "    'run_model': 'infer'\n",
    "}\n",
    "\n",
    "# Plot each method separately\n",
    "for method, group in results.groupby('method'):\n",
    "    ax.barh(y=group[\"worker_id\"],\n",
    "            width=group[\"time_running\"],\n",
    "            left=group['walltime_start'], \n",
    "            color=method_color[method],\n",
    "            label=method_label[method],\n",
    "            edgecolor='none')\n",
    "    \n",
    "# Add in a legend\n",
    "ax.legend()\n",
    "    \n",
    "# Make some labels\n",
    "ax.set_yticks(range(n_workers))\n",
    "ax.set_yticklabels([])\n",
    "ax.set_ylabel('Worker')\n",
    "\n",
    "ax.set_xlabel('Walltime')\n",
    "ax.set_xlim([0, results['walltime_finish'].max()])\n",
    "\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "69f8fdf8-47c5-4396-ad34-a680e5a71472",
   "metadata": {},
   "source": [
    "You can see Colmena does a good job at overlapping tasks of various types. Our next question is \"does this actually matter?\" "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37e13882-4706-4761-b2f3-16add057c414",
   "metadata": {},
   "source": [
    "## Measuring Scientific Performance\n",
    "The key goal we have is to find the most number of \"good\" molecules in the shortest time."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eb362624-d250-4622-a0fc-c0ab9842693e",
   "metadata": {},
   "source": [
    "Retrieve the molecule properties from the Colmena jobs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d356a460-9be7-4ca6-990f-772a7ed291e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "colmena_mols = results.query('method==\"compute_vertical\" and success').copy()\n",
    "colmena_mols['smiles'] = colmena_mols['inputs'].apply(lambda x: x[0][0])\n",
    "colmena_mols['ie'] = colmena_mols['value']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7473662-865c-431c-b463-d0f6ba93645b",
   "metadata": {},
   "outputs": [],
   "source": [
    "colmena_mols[['smiles', 'ie', 'walltime_finish']].head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "469182a2-b19a-4114-8794-191c31c77d78",
   "metadata": {},
   "source": [
    "Load in the Parsl results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7575b584-724f-413a-a876-04369ddbfb4e",
   "metadata": {},
   "outputs": [],
   "source": [
    "parsl_mols = pd.read_csv('run-data/parsl-results.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c48b613-1458-4244-aba0-37f44574b6b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "parsl_mols.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1968b4ee-4a40-48bc-883d-2bdc63c272f8",
   "metadata": {},
   "source": [
    "Plot the the completion time and "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92a4cfe9-eb55-46bd-b893-ed44e475fed4",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "\n",
    "ax.scatter(colmena_mols['walltime_finish'], colmena_mols['ie'], label='Colmena')\n",
    "ax.scatter(parsl_mols['time'], parsl_mols['ie'], label='Parsl')\n",
    "\n",
    "ax.legend()\n",
    "\n",
    "ax.set_xlabel('Walltime (s)')\n",
    "ax.set_ylabel('IE (Ha)')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d6a08425-e0ab-49c1-867d-82e61a3800f8",
   "metadata": {},
   "source": [
    "We can already see that the Parsl application takes much longer evaluate the same number of molecules. \n",
    "\n",
    "Even if it is slower, this does not mean Parsl is giving us worse results. Let's measure that by how many molecules we find with a very large IE (>0.55 Ha)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "605d54f9-eb38-430e-a7b6-fb4ae53ea6c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "thr = 0.55\n",
    "colmena_mols.sort_values('walltime_finish', ascending=True, inplace=True)\n",
    "colmena_mols['best_found'] = (colmena_mols['ie'] > thr).cumsum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c64f48b-b56f-4ed9-8c26-67f9c5a6d4ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "parsl_mols.sort_values('time', ascending=True, inplace=True)\n",
    "parsl_mols['best_found'] = (parsl_mols['ie'] > thr).cumsum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "039e1589-f7f6-4dd6-9891-f68195d39b4a",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(3.5, 2.5))\n",
    "\n",
    "ax.step(colmena_mols['walltime_finish'], colmena_mols['best_found'], label='Colmena')\n",
    "ax.step(parsl_mols['time'], parsl_mols['best_found'], label='Parsl')\n",
    "\n",
    "ax.legend()\n",
    "\n",
    "ax.set_xlabel('Walltime (s)')\n",
    "ax.set_ylabel('Top Molecules Found')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5cb6e76-7439-4444-8fde-0861d8541642",
   "metadata": {},
   "source": [
    "We can see that Colmena takes much less time to find a more molecules, which is a result of the Colmena application being able to manage different kinds of tasks at the same time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "633fac29-b807-4b47-bc2b-25e44bebd18a",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(3.5, 2.5))\n",
    "\n",
    "ax.step(range(len(colmena_mols)), colmena_mols['best_found'], label='Colmena')\n",
    "ax.step(range(len(parsl_mols)), parsl_mols['best_found'], label='Parsl')\n",
    "\n",
    "ax.legend()\n",
    "\n",
    "ax.set_xlabel('Molecules Tested')\n",
    "ax.set_ylabel('Top Molecules Found')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33741c77-1118-4b2e-bce5-7035587781de",
   "metadata": {},
   "source": [
    "You can even see that the machine learning performs just as well for the Parsl application. Both the Parsl and Colmena find similar numbers of top molecules at a given numbers of simulations performed. It is purely the access to intelligent scheduling policies via Colmena which enable the improved science-per-core-hour."
   ]
  }
 ],
 "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.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}