Notebook to launch multiple full analyses

monash-emu · Aug 21, 2023 · 1ba1fa8 · 1ba1fa8
1 parent e3323b7
commit 1ba1fa8
Showing 1 changed file with 202 additions and 0 deletions.
diff --git a/notebooks/user/rragonnet/project_specific/School_Closure/remote/sb_multi_full_analysis.ipynb b/notebooks/user/rragonnet/project_specific/School_Closure/remote/sb_multi_full_analysis.ipynb
@@ -0,0 +1,202 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from autumn.infrastructure.remote import springboard\n",
+    "from autumn.projects.sm_covid2.common_school.runner_tools import run_full_analysis, print_continuous_status, download_analysis"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### Standard config"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ISO3_LIST = [\"AUS\", \"BOL\", \"FRA\", \"MAR\", \"PHL\"]\n",
+    "ANALYSIS = \"main\"\n",
+    "\n",
+    "N_CHAINS = 32\n",
+    "N_OPTI_SEARCHES = 32\n",
+    "N_BEST_RETAINED = 16\n",
+    "OPTI_BUDGET = 10000\n",
+    "\n",
+    "METROPOLIS_TUNE = 2000\n",
+    "METROPOLIS_DRAWS = 10000\n",
+    "METROPOLIS_METHOD = \"DEMetropolis\"\n",
+    "\n",
+    "FULL_RUNS_SAMPLES = 1000\n",
+    "BURN_IN = 5000"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def remote_full_analysis_task(bridge: springboard.task.TaskBridge, iso3: str = 'FRA', analysis: str = \"main\"):\n",
+    "    \n",
+    "    import multiprocessing as mp\n",
+    "    mp.set_start_method('forkserver')\n",
+    "\n",
+    "    bridge.logger.info(f\"Running full analysis for {iso3}. Analysis type: {analysis}.\")\n",
+    "\n",
+    "    idata, uncertainty_df, diff_quantiles_df = run_full_analysis(\n",
+    "        iso3,\n",
+    "        analysis=analysis, \n",
+    "        opti_params={'n_searches': N_OPTI_SEARCHES, \"n_best_retained\": N_BEST_RETAINED, 'num_workers': 8, 'parallel_opti_jobs': 8, 'warmup_iterations': 0, 'search_iterations': OPTI_BUDGET, 'init_method': 'LHS'},\n",
+    "        mcmc_params={'draws': METROPOLIS_DRAWS, 'tune': METROPOLIS_TUNE, 'cores': N_CHAINS, 'chains': N_CHAINS, 'method': METROPOLIS_METHOD},\n",
+    "        full_run_params={'samples': FULL_RUNS_SAMPLES, 'burn_in': BURN_IN},\n",
+    "        output_folder=bridge.out_path,\n",
+    "        logger=bridge.logger\n",
+    "    )\n",
+    "   \n",
+    "    bridge.logger.info(\"Full analysis complete\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "mspec = springboard.EC2MachineSpec(N_CHAINS, 4, \"compute\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "analysis_title = \"test_run\"\n",
+    "config_str = f\"_{ANALYSIS}_LHS{N_OPTI_SEARCHES}_opt{OPTI_BUDGET}_mc{METROPOLIS_TUNE}n{METROPOLIS_DRAWS}\"\n",
+    "run_paths = [springboard.launch.get_autumn_project_run_path(\"school_project\", iso3, analysis_title + config_str) for iso3 in ISO3_LIST]\n",
+    "run_paths"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "task_dict = {}\n",
+    "for i, iso3 in enumerate(ISO3_LIST):\n",
+    "    # Use these targets, but keep the other kwargs the same\n",
+    "    task_kwargs = {\n",
+    "        \"iso3\": iso3,\n",
+    "        \"analysis\": ANALYSIS\n",
+    "    }\n",
+    "    task_dict[run_paths[i]] = springboard.TaskSpec(remote_full_analysis_task, task_kwargs)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "runners = springboard.launch.launch_synced_multiple_autumn_task(task_dict, mspec)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "runners_dict = {iso3: runners[list(runners.keys())[i]] for i, iso3 in enumerate(ISO3_LIST)}"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "runner.instance"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "runners_dict[\"AUS\"].s3.get_status()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "for iso3 in ISO3_LIST:\n",
+    "    runner = runners_dict[iso3]\n",
+    "    print(runner.s3.get_status())\n",
+    "    print(runner.top(\"%MEM\"))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "run_paths"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "for run_path in run_paths:\n",
+    "    download_analysis(run_path)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "summer2",
+   "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.10.11"
+  },
+  "orig_nbformat": 4
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}