Feature visualizer notebook

.github/workflows/workflow.yml

@@ -166,7 +166,7 @@ jobs:
           git config user.name "github-actions[bot]"
           git config user.email "github-actions[bot]@users.noreply.github.com"
           git fetch origin gh-pages
-          git checkout gh-pages
+          git checkout -f gh-pages
           cp /tmp/index.html index.html
           git add index.html
           git diff --cached --quiet || git commit -m "Update docs for ${{ steps.get_version.outputs.version }}"

Dockerfile.devenv

@@ -14,6 +14,7 @@ RUN apt-get update && apt-get install -y \
     g++-11 \
     make \
     cmake \
+    lsof \
     clang-format \
     libboost-dev \
     libboost-program-options-dev \

client-scripts/notebooks/analytical_congestion_aware_sample.ipynb

@@ -0,0 +1,209 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "42708369",
+   "metadata": {},
+   "source": [
+    "##### Import the required modules and configure the system path to locate them"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "970611b8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import sys\n",
+    "sys.path.append(\"../utils\")\n",
+    "from astra_sim import AstraSim, Collective, NetworkBackend"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "fa81522b",
+   "metadata": {},
+   "source": [
+    "##### Call the AstraSim client helper with the server endpoint and tag to connect to the ASTRA-sim gRPC server, initialize the SDK, and create a tagged folder for configs, results, and logs."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a89f699a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "astra = AstraSim(server_endpoint = \"172.17.0.2:8989\", tag = \"analytical_congestion_aware_sample\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "67e495d4",
+   "metadata": {},
+   "source": [
+    "##### Generate workload execution traces for each rank and set the required data size for AstraSim configuration"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a3f00619",
+   "metadata": {
+    "lines_to_next_cell": 2
+   },
+   "outputs": [],
+   "source": [
+    "astra.configuration.common_config.workload = astra.generate_collective(collective=Collective.ALLTOALL, coll_size= 8*1024*1024, npu_range=[0,8])\n",
+    "print(astra.configuration.common_config.workload)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f7fe6ba5",
+   "metadata": {},
+   "source": [
+    "##### Configure ASTRA-sim system config"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2bc05826",
+   "metadata": {
+    "lines_to_next_cell": 2
+   },
+   "outputs": [],
+   "source": [
+    "astra.configuration.common_config.system.scheduling_policy = astra.configuration.common_config.system.LIFO\n",
+    "astra.configuration.common_config.system.endpoint_delay = 10\n",
+    "astra.configuration.common_config.system.active_chunks_per_dimension = 1\n",
+    "astra.configuration.common_config.system.all_gather_implementation = [astra.configuration.common_config.system.RING]\n",
+    "astra.configuration.common_config.system.all_to_all_implementation = [astra.configuration.common_config.system.DIRECT]\n",
+    "astra.configuration.common_config.system.all_reduce_implementation = [astra.configuration.common_config.system.RING]\n",
+    "astra.configuration.common_config.system.collective_optimization = astra.configuration.common_config.system.LOCALBWAWARE\n",
+    "astra.configuration.common_config.system.local_mem_bw = 1600\n",
+    "print(astra.configuration.common_config.system)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "292def3f",
+   "metadata": {},
+   "source": [
+    "##### Configure ASTRA-sim remote memory configuration"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6353e6c6",
+   "metadata": {
+    "lines_to_next_cell": 2
+   },
+   "outputs": [],
+   "source": [
+    "astra.configuration.common_config.remote_memory.memory_type = astra.configuration.common_config.remote_memory.NO_MEMORY_EXPANSION\n",
+    "print(astra.configuration.common_config.remote_memory)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "76e50273",
+   "metadata": {},
+   "source": [
+    "##### Configure the network backend and topology"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0384143f",
+   "metadata": {
+    "lines_to_next_cell": 2
+   },
+   "outputs": [],
+   "source": [
+    "# astra.configuration.network_backend.choice = astra.configuration.network_backend.ANALYTICAL_CONGESTION_AWARE\n",
+    "astra.configuration.network_backend.analytical_congestion_aware.topology.network.clear()\n",
+    "astra.configuration.network_backend.analytical_congestion_aware.topology.network.add(\"switch\", 8, 100, 0.005) # add(type_of_topology, number_of_nodes, bandwidth_in_gbps, latency_in_ns)\n",
+    "print(\"Network backend set to\", astra.configuration.network_backend.choice)\n",
+    "print(\"network backend choice set to:\",astra.configuration.network_backend.analytical_congestion_aware.topology.choice)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "8a9b435c",
+   "metadata": {},
+   "source": [
+    "##### Configure ASTRA-sim cmd parameters"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f2cbf6ec",
+   "metadata": {
+    "lines_to_next_cell": 2
+   },
+   "outputs": [],
+   "source": [
+    "astra.configuration.common_config.cmd_parameters.comm_scale = 1\n",
+    "astra.configuration.common_config.cmd_parameters.rendezvous_protocol = False\n",
+    "astra.configuration.common_config.cmd_parameters.injection_scale = 1\n",
+    "\n",
+    "print(astra.configuration.network_backend.analytical_congestion_aware.topology.network)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a6857638",
+   "metadata": {},
+   "source": [
+    "#### Start the simulation by specifying the network backend"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e8e0cfef",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "astra.run_simulation(NetworkBackend.ANALYTICAL_CONGESTION_AWARE)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "0d74c3ca",
+   "metadata": {},
+   "source": [
+    "##### Download all the configurations as a zip"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "46bbb3e0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "astra.download_configuration()"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.11.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}