{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0877d04e",
   "metadata": {},
   "source": [
    "# **Multi-Host Particle Mesh Simulation**\n",
    "\n",
    "In this notebook, we extend our Particle Mesh simulation across **multiple nodes**, enabling simulations at scales not achievable on a single machine. By leveraging distributed GPUs across hosts, we handle larger mesh shapes and box sizes efficiently.\n",
    "\n",
    "> **Note**: Since there’s no direct way to run a multi-host notebook, I’ll guide you step by step on how to submit an interactive job from a script.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To run a multi-host simulation, you first need to **allocate a job** with `salloc`. This command requests resources on an HPC cluster.\n",
    "\n",
    "> **Note**: You can alternatively use `sbatch` with a SLURM script to submit the job. The exact `salloc` parameters may vary depending on your specific HPC cluster configuration.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5f42bbe",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "c5f42bbe",
    "outputId": "a7841b28-5f20-4856-bd1d-f8a3572095b5"
   },
   "outputs": [],
   "source": [
    "!salloc --account=XXX@a100 -C a100 --gres=gpu:8 --ntasks-per-node=8 --time=00:40:00  --cpus-per-task=8 --hint=nomultithread --qos=qos_gpu-dev --nodes=4 & "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac6585f3",
   "metadata": {},
   "source": [
    "> **Note**: These `salloc` parameters are configured for the **Jean Zay** supercomputer in France. Adaptations might be necessary if using a different HPC cluster.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74928ff7",
   "metadata": {},
   "source": [
    "**A few hours later**"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c52f89cc",
   "metadata": {},
   "source": [
    "Use `!squeue -u $USER -o \"%i %D %b\"` to **check the JOB ID** and verify your resource allocation.\n",
    "\n",
    "In this example, we’ve been allocated **32 GPUs split across 4 nodes**.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "7ebdfc00",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "JOBID NODES TRES_PER_NODE\n",
      "467745 4 gres/gpu:8\n"
     ]
    }
   ],
   "source": [
    "!squeue -u $USER -o \"%i %D %b\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88bd7ef8",
   "metadata": {},
   "source": [
    "Unset the following environment variables, as they can cause issues when using JAX in a distributed setting:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66dbe8f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "del os.environ['VSCODE_PROXY_URI']\n",
    "del os.environ['NO_PROXY']\n",
    "del os.environ['no_proxy']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36479cc7",
   "metadata": {},
   "source": [
    "### Checking Available Compute Resources\n",
    "\n",
    "Run the following command to initialize JAX distributed computing and display the devices available for this job:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "c78b8a4e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[CudaDevice(id=0), CudaDevice(id=1), CudaDevice(id=2), CudaDevice(id=3), CudaDevice(id=4), CudaDevice(id=5), CudaDevice(id=6), CudaDevice(id=7), CudaDevice(id=8), CudaDevice(id=9), CudaDevice(id=10), CudaDevice(id=11), CudaDevice(id=12), CudaDevice(id=13), CudaDevice(id=14), CudaDevice(id=15), CudaDevice(id=16), CudaDevice(id=17), CudaDevice(id=18), CudaDevice(id=19), CudaDevice(id=20), CudaDevice(id=21), CudaDevice(id=22), CudaDevice(id=23), CudaDevice(id=24), CudaDevice(id=25), CudaDevice(id=26), CudaDevice(id=27), CudaDevice(id=28), CudaDevice(id=29), CudaDevice(id=30), CudaDevice(id=31)]\n"
     ]
    }
   ],
   "source": [
    "!srun --jobid=467745 -n 32 python -c \"import jax; jax.distributed.initialize(); print(jax.devices()) if jax.process_index() == 0 else None\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b24e2e65",
   "metadata": {},
   "source": [
    "### Running the Multi-Host Simulation Script\n",
    "\n",
    "Run the simulation script across 32 processes:\n",
    "\n",
    "```bash\n",
    "!srun --jobid=467745 -n 32 python 05-MultiHost_PM.py --mesh_shape 1024 1024 1024 --box_size 1000. 1000. 1000. --halo_size 128 -s leapfrog --pdims 16 2\n",
    "```\n",
    "The script, located in the same path as this notebook, is named [**05-MultiHost_PM.py**](05-MultiHost_PM.py).\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "df27af30",
   "metadata": {},
   "source": [
    "### Multi-Host Simulation Script with Arguments\n",
    "\n",
    "This script is nearly identical to the single-host version, with the main addition being the call to `jax.distributed.initialize()` at the start, enabling multi-host parallelism. Here’s a breakdown of the key arguments:\n",
    "\n",
    "- **`--pdims`** (`-p`): Specifies processor grid dimensions as two integers, like `16 2` for 16 x 2 device mesh (default is `[1, jax.devices()]`).\n",
    "- **`--mesh_shape`** (`-m`): Defines the simulation mesh shape as three integers (default is `[512, 512, 512]`).\n",
    "- **`--box_size`** (`-b`): Sets the physical box size of the simulation as three floating-point values, e.g., `1000. 1000. 1000.` (default is `[500.0, 500.0, 500.0]`).\n",
    "- **`--halo_size`** (`-H`): Specifies the halo size for boundary overlap across nodes (default is `64`).\n",
    "- **`--solver`** (`-s`): Chooses the ODE solver (`leapfrog` or `dopri8`). The `leapfrog` solver uses a fixed step size, while `dopri8` is an adaptive Runge-Kutta solver with a PID controller (default is `leapfrog`).\n",
    "- **`--snapthots`** (`-st`) : Number of snapshots to save (warning, increases memory usage)\n",
    "\n",
    "The script also saves results across nodes.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6d13679",
   "metadata": {},
   "outputs": [],
   "source": [
    "import subprocess\n",
    "\n",
    "# Define parameters as variables\n",
    "jobid = \"467745\"\n",
    "num_processes = 32\n",
    "script_name = \"05-MultiHost_PM.py\"\n",
    "mesh_shape = (1024, 1024, 1024)\n",
    "box_size = (1000., 1000., 1000.)\n",
    "halo_size = 128\n",
    "solver = \"leapfrog\"\n",
    "pdims = (16, 2)\n",
    "snapshots = 2\n",
    "\n",
    "# Build the command as a list, incorporating variables\n",
    "command = [\n",
    "    \"srun\",\n",
    "    f\"--jobid={jobid}\",\n",
    "    \"-n\", str(num_processes),\n",
    "    \"python\", script_name,\n",
    "    \"--mesh_shape\", str(mesh_shape[0]), str(mesh_shape[1]), str(mesh_shape[2]),\n",
    "    \"--box_size\", str(box_size[0]), str(box_size[1]), str(box_size[2]),\n",
    "    \"--halo_size\", str(halo_size),\n",
    "    \"-s\", solver,\n",
    "    \"--pdims\", str(pdims[0]), str(pdims[1]),\n",
    "    \"--snapshots\", str(snapshots)\n",
    "]\n",
    "\n",
    "# Execute the command as a subprocess\n",
    "subprocess.run(command)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45333bf0",
   "metadata": {},
   "source": [
    "### Loading and Visualizing Results\n",
    "\n",
    "After running the multi-host simulation, we load the saved results from disk:\n",
    "\n",
    "- **`initial_conditions.npy`**: Initial conditions for the simulation.\n",
    "- **`lpt_displacements.npy`**: Linear perturbation displacements.\n",
    "- **`ode_solution_0.npy`** and **`ode_solution_1.npy`**: Solutions from the ODE solver at each snapshot.\n",
    "\n",
    "We then use `plot_fields_single_projection` to visualize these fields and observe the results across multiple snapshots.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "472dd4bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "initial_conditions = np.load('fields/initial_conditions.npy')\n",
    "lpt_displacements = np.load('fields/lpt_displacements.npy')\n",
    "ode_solution_0 = np.load('fields/ode_solution_0.npy')\n",
    "ode_solution_1 = np.load('fields/ode_solution_1.npy')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "4e012ce8",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 323
    },
    "id": "4e012ce8",
    "outputId": "75390318-8072-481f-ffb9-ec09cd71cb1d"
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 1986156 characters]",
      "text/plain": [
       "<Figure size 2000x500 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from jaxpm.plotting import plot_fields_single_projection\n",
    "fields = {\n",
    "    \"Initial Conditions\": initial_conditions,\n",
    "    \"LPT Field\": lpt_displacements,\n",
    "    \"ODE Solution 0\": ode_solution_0,\n",
    "    \"ODE Solution 1\": ode_solution_1\n",
    "}\n",
    "plot_fields_single_projection(fields)"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "include_colab_link": true,
   "name": "Introduction.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
