docs/source/contributing/adding_workloads.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "456c8d45",
   "metadata": {},
   "source": [
    "# Adding New Workloads\n",
    "\n",
    "Adding new workloads into benchmark-wrapper is a fairly straightforward process, but requires a bit of work. This page tracks all the changes that a user needs to make when adding in a new benchmark.\n",
    "\n",
    "<div class=\"alert alert-info\">\n",
    "    \n",
    "benchmark-wrapper is currently undergoing a re-write, meaning a lot of change is happening pretty quickly. The information on this page is relevant to the new modifications which change the way in which benchmarks should be added. It may not be consistent to the way in which existing benchmarks are developed.\n",
    "</div>\n",
    "\n",
    "<div class=\"alert alert-warning\">\n",
    "    \n",
    "This page is written within a Jupyter Notebook, however running benchmark-wrapper from within a Jupyter Notebook is not tested nor supported.\n",
    "</div>\n",
    "\n",
    "## Step Zero: Prep\n",
    "\n",
    "A Benchmark within benchmark-wrapper is essentially just a Python module which handles setting up, running, parsing and tearing down a benchmark. To create our benchmark, we need to understand the following items:\n",
    "\n",
    "1. What is the human-readable, camel-case-able string name for our benchmark?\n",
    "1. What arguments does the benchmark wrapper need from the user?\n",
    "1. Are there any setup tasks that our benchmark wrapper needs to perform?\n",
    "1. How do we run our benchmark?\n",
    "1. Are there any cleanup tasks that our benchmark wrapper needs to perform?\n",
    "1. What data should the benchmark export?\n",
    "\n",
    "In this example, we'll create a new benchmark wrapper that does a ping test against a list of given hostnames and IPs:\n",
    "\n",
    "1. We'll call it ``pingtest``\n",
    "1. We need to know which hosts the user wants to ping and how many pings the user wants to perform.\n",
    "1. We need to verify that the arguments that the user gave us are valid. We'll also create a temp file to show that the benchmark is running.\n",
    "1. We can run our ping tests using the ``ping`` shell command.\n",
    "1. We need to clean up our 'I-am-running' temp file.\n",
    "1. For each pinged host, we want to output a single result detailing the result of the ping session (RTT information, packet loss %, IP resolution, errors).\n",
    "\n",
    "## Step One: Initialize\n",
    "\n",
    " To begin, let's create the required files for our benchmark by creating a new Python package under ``snafu/benchmarks``:\n",
    "\n",
    "```text\n",
    "snafu/benchmarks/pingtest/\n",
    "├── __init__.py\n",
    "└── pingtest.py\n",
    "```\n",
    "\n",
    "Inside ``pingtest.py``, we'll create our initial ``Benchmark`` subclass. Inside this subclass, there are a few\n",
    "class variables which we need to set:\n",
    "\n",
    "1. ``tool_name``: This is the camel-case name for our benchmark.\n",
    "1. ``args``: These are the arguments which our Benchmark will pull from the user through the CLI, OS environment, and/or from a configuration file (CLI is preferred over the OS environment, which is preferred over the configuration file). In the background, snafu uses [configargparse](https://pypi.org/project/ConfigArgParse/) to do the dirty-work, which is a helpful wrapper around Python's own [argparse](https://docs.python.org/3/library/argparse.html). The ``args`` class variable should be set to a tuple of ``snafu.config.ConfigArgument``s, which take in arguments just like ``configargparse.ArgumentParser.add_argument``. If you have used argparse in the past, this should look super familiar.\n",
    "1. ``metadata``: This is an iterable of strings which represent the metadata that will be exported with the Benchmark's results. Each string corresponds to the attribute name that the argument is stored under by ``configargparse``. For instance, creating a new argument under the ``args`` class variable with ``\"--my-metadata\", dest=\"mmd\"`` would result in the attribute name being ``mmd``. Adding ``mmd`` under ``metadata`` will in turn cause the value for the ``--my-metadata`` argument to be exported as metadata. Benchmarks will by default specify ``cluster_name``, ``user`` and ``uuid`` arguments as metadata, but if you want to use your own set of metadata keys it can be set here.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2649c485",
   "metadata": {},
   "outputs": [],
   "source": [
    "#!/usr/bin/env python3\n",
    "# -*- coding: utf-8 -*-\n",
    "\"\"\"Ping hosts and export results.\"\"\"\n",
    "from snafu.config import ConfigArgument\n",
    "from snafu.benchmarks import Benchmark\n",
    "\n",
    "\n",
    "class PingTest(Benchmark):\n",
    "    \"\"\"Wrapper for the Ping Test benchmark.\"\"\"\n",
    "\n",
    "    tool_name = \"pingtest\"\n",
    "    args = (\n",
    "        ConfigArgument(\n",
    "            \"--host\",\n",
    "            help=\"Host(s) to ping. Can give more than one host by separating them with \"\n",
    "                 \"spaces on the CLI and in config files, or by giving them in a \"\n",
    "                 \"pythonic-list format through the OS environment\",\n",
    "            dest=\"host\",\n",
    "            nargs=\"+\",\n",
    "            env_var=\"HOST\",\n",
    "            type=str,\n",
    "            required=True,\n",
    "        ),\n",
    "        ConfigArgument(\n",
    "            \"--count\",\n",
    "            help=\"Number of pings to perform per sample.\",\n",
    "            dest=\"count\",\n",
    "            env_var=\"COUNT\",\n",
    "            default=1,\n",
    "            type=int,\n",
    "        ),\n",
    "        ConfigArgument(\n",
    "            \"--samples\",\n",
    "            help=\"Number of samples to perform.\",\n",
    "            dest=\"samples\",\n",
    "            env_var=\"SAMPLES\",\n",
    "            default=1,\n",
    "            type=int,\n",
    "        ),\n",
    "        ConfigArgument(\n",
    "            \"--htlhcdtwy\",\n",
    "            help=\"Has The Large Hadron Collider Destroyed The World Yet?\",\n",
    "            dest=\"htlhcdtwy\",\n",
    "            env_var=\"HTLHCDTWY\",\n",
    "            default=\"no\",\n",
    "            type=str,\n",
    "            choices=[\"yes\", \"no\"]\n",
    "        ),\n",
    "    )\n",
    "    # don't care about Cluster Name, but the Hadron Collider is serious business\n",
    "    metadata = (\"user\", \"uuid\", \"htlhcdtwy\")\n",
    "\n",
    "    def setup(self):\n",
    "        \"\"\"Setup the Ping Test Benchmark.\"\"\"\n",
    "        pass\n",
    "\n",
    "    def collect(self):\n",
    "        \"\"\"Run the Ping Test Benchmark and collect results.\"\"\"\n",
    "        pass\n",
    "\n",
    "    def cleanup(self):\n",
    "        \"\"\"Cleanup the Ping Test Benchmark.\"\"\"\n",
    "        pass\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eb716d76",
   "metadata": {},
   "source": [
    "Let's check that we're ready to move on by trying to parse some configuration parameters. Let's load up Python!\n",
    "\n",
    "benchmark-wrapper includes a special variable called ``snafu.registry.TOOLS`` which will map a benchmark's\n",
    "camel-case string name to its wrapper class. Let's use this to create an instance of our benchmark and\n",
    "parse some configuration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86af872c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from snafu.registry import TOOLS\n",
    "from pprint import pprint\n",
    "pingtest = TOOLS[\"pingtest\"]()\n",
    "\n",
    "# Set some config parameters\n",
    "# Config file\n",
    "!echo \"samples: 3\" > my_config.yaml\n",
    "!echo \"count: 5\" >> my_config.yaml\n",
    "# OS ENV\n",
    "import os\n",
    "os.environ[\"HOST\"] = \"[www.google.com,www.bing.com]\"\n",
    "\n",
    "# Parse arguments and print result\n",
    "# Since we aren't running within the main script (run_snafu.py),\n",
    "# need to add the config option manually\n",
    "pingtest.config.parser.add_argument(\"--config\", is_config_file=True)\n",
    "pingtest.config.parse_args(\n",
    "    \"--config my_config.yaml --labels=notebook=true --uuid 1337 --user snafu \"\n",
    "    \"--htlhcdtwy=no\".split(\" \")\n",
    ")\n",
    "pprint(vars(pingtest.config.params))\n",
    "\n",
    "del pingtest\n",
    "!rm my_config.yaml"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5223604f",
   "metadata": {},
   "source": [
    "Now that we have our configuration all ready to go, let's start filling in our benchmark."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b35b31b5",
   "metadata": {},
   "source": [
    "## Step Two: Setup Method\n",
    "\n",
    "Each benchmark is expected to have a ``setup`` method, which will return ``True`` if setup tasks completed successfully and otherwise return ``False``.\n",
    "\n",
    "For our use case, let's write a file to ``/tmp`` that can signal other programs that our benchmark is running. We'll also check if our temporary file exists before writing it, which would indicate that something is wrong.\n",
    "\n",
    "<div class=\"alert alert-info\">\n",
    "    Remeber, we are still working in our module found at snafu/benchmarks/pingtest.py\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "898ea74e",
   "metadata": {},
   "outputs": [],
   "source": [
    "#!/usr/bin/env python3\n",
    "# -*- coding: utf-8 -*-\n",
    "\"\"\"Ping hosts and export results.\"\"\"\n",
    "from snafu.config import ConfigArgument\n",
    "from snafu.benchmarks import Benchmark\n",
    "\n",
    "# We'll also import this helpful function from the config module\n",
    "import os\n",
    "from snafu.config import check_file\n",
    "\n",
    "\n",
    "class PingTest(Benchmark):\n",
    "    \"\"\"Wrapper for the Ping Test benchmark.\"\"\"\n",
    "\n",
    "    tool_name = \"pingtest\"\n",
    "    args = (\n",
    "        ConfigArgument(\n",
    "            \"--host\",\n",
    "            help=\"Host(s) to ping. Can give more than one host by separating them with \"\n",
    "                 \"spaces on the CLI and in config files, or by giving them in a \"\n",
    "                 \"pythonic-list format through the OS environment\",\n",
    "            dest=\"host\",\n",
    "            nargs=\"+\",\n",
    "            env_var=\"HOST\",\n",
    "            type=str,\n",
    "            required=True,\n",
    "        ),\n",
    "        ConfigArgument(\n",
    "            \"--count\",\n",
    "            help=\"Number of pings to perform per sample.\",\n",
    "            dest=\"count\",\n",
    "            env_var=\"COUNT\",\n",
    "            default=1,\n",
    "            type=int,\n",
    "        ),\n",
    "        ConfigArgument(\n",
    "            \"--samples\",\n",
    "            help=\"Number of samples to perform.\",\n",
    "            dest=\"samples\",\n",
    "            env_var=\"SAMPLES\",\n",
    "            default=1,\n",
    "            type=int,\n",
    "        ),\n",
    "        ConfigArgument(\n",
    "            \"--htlhcdtwy\",\n",
    "            help=\"Has The Large Hadron Collider Destroyed The World Yet?\",\n",
    "            dest=\"htlhcdtwy\",\n",
    "            env_var=\"HTLHCDTWY\",\n",
    "            default=\"no\",\n",
    "            type=str,\n",
    "            choices=[\"yes\", \"no\"]\n",
    "        ),\n",
    "    )\n",
    "    # don't care about Cluster Name, but the Hadron Collider is serious business\n",
    "    metadata = (\"user\", \"uuid\", \"htlhcdtwy\")\n",
    "    \n",
    "    TMP_FILE_PATH = \"/tmp/snafu-pingtest\"\n",
    "\n",
    "    def setup(self) -> bool:\n",
    "        \"\"\"\n",
    "        Setup the Ping Test Benchmark.\n",
    "        \n",
    "        This method creates a temporary file at ``/tmp/snafu-pingtest`` to let others \n",
    "        know that the benchmark is currently running.\n",
    "        \n",
    "        Returns\n",
    "        -------\n",
    "        bool\n",
    "            True if the temporary file was created successfully, othewise False. Will\n",
    "            also return False if the temporary file already exists.\n",
    "        \"\"\"\n",
    "        \n",
    "        if check_file(self.TMP_FILE_PATH):\n",
    "            # The benchmark base class exposes a logger at self.logger which we can use\n",
    "            self.logger.critical(\n",
    "                f\"Temporary file located at {self.TMP_FILE_PATH} already exists.\"\n",
    "            )\n",
    "            return False\n",
    "        \n",
    "        try:\n",
    "            tmp_file = open(self.TMP_FILE_PATH, \"x\")\n",
    "            tmp_file.close()\n",
    "        except Exception as e:\n",
    "            self.logger.critical(\n",
    "                f\"Unable to create temporary file at {self.TMP_FILE_PATH}: {e}\",\n",
    "                exc_info=True\n",
    "            )\n",
    "            \n",
    "            return False\n",
    "        else:\n",
    "            self.logger.info(\n",
    "                f\"Successfully created temp file at {self.TMP_FILE_PATH}\"\n",
    "            )\n",
    "            return True\n",
    "\n",
    "    def collect(self):\n",
    "        \"\"\"Run the Ping Test Benchmark and collect results.\"\"\"\n",
    "        pass\n",
    "\n",
    "    def cleanup(self):\n",
    "        \"\"\"Cleanup the Ping Test Benchmark.\"\"\"\n",
    "        pass\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2de8d35f",
   "metadata": {},
   "source": [
    "Let's test it out and make sure our setup method works properly:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4226479c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from snafu.registry import TOOLS\n",
    "pingtest = TOOLS[\"pingtest\"]()\n",
    "\n",
    "!rm -f /tmp/snafu-pingtest\n",
    "\n",
    "# No file exists\n",
    "print(f\"Setup result is: {pingtest.setup()}\")\n",
    "\n",
    "# File exists\n",
    "print(f\"Setup result is: {pingtest.setup()}\")\n",
    "\n",
    "# Create failure in open\n",
    "!rm -f /tmp/snafu-pingtest\n",
    "open_bak = open\n",
    "open = lambda file, mode: int(\"I'm a string\")\n",
    "print(f\"Setup result is: {pingtest.setup()}\")\n",
    "\n",
    "# Cleanup\n",
    "open = open_bak\n",
    "!rm -f /tmp/snafu-pingtest\n",
    "del pingtest"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7532e3ff",
   "metadata": {},
   "source": [
    "## Step Three: Cleanup Method\n",
    "\n",
    "Now let's go ahead and populate our cleanup method. The ``cleanup`` method has the same usage as the setup method: return ``True`` if the cleanup was successfull, otherwise ``False``. For the ping test benchmark, we just need to remove our temporary file:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87c51692",
   "metadata": {},
   "outputs": [],
   "source": [
    "#!/usr/bin/env python3\n",
    "# -*- coding: utf-8 -*-\n",
    "\"\"\"Ping hosts and export results.\"\"\"\n",
    "import os\n",
    "from snafu.config import ConfigArgument, check_file\n",
    "from snafu.benchmarks import Benchmark\n",
    "\n",
    "\n",
    "class PingTest(Benchmark):\n",
    "    \"\"\"Wrapper for the Ping Test benchmark.\"\"\"\n",
    "\n",
    "    tool_name = \"pingtest\"\n",
    "    args = (\n",
    "        ConfigArgument(\n",
    "            \"--host\",\n",
    "            help=\"Host(s) to ping. Can give more than one host by separating them with \"\n",
    "                 \"spaces on the CLI and in config files, or by giving them in a \"\n",
    "                 \"pythonic-list format through the OS environment\",\n",
    "            dest=\"host\",\n",
    "            nargs=\"+\",\n",
    "            env_var=\"HOST\",\n",
    "            type=str,\n",
    "            required=True,\n",
    "        ),\n",
    "        ConfigArgument(\n",
    "            \"--count\",\n",
    "            help=\"Number of pings to perform per sample.\",\n",
    "            dest=\"count\",\n",
    "            env_var=\"COUNT\",\n",
    "            default=1,\n",
    "            type=int,\n",
    "        ),\n",
    "        ConfigArgument(\n",
    "            \"--samples\",\n",
    "            help=\"Number of samples to perform.\",\n",
    "            dest=\"samples\",\n",
    "            env_var=\"SAMPLES\",\n",
    "            default=1,\n",
    "            type=int,\n",
    "        ),\n",
    "        ConfigArgument(\n",
    "            \"--htlhcdtwy\",\n",
    "            help=\"Has The Large Hadron Collider Destroyed The World Yet?\",\n",
    "            dest=\"htlhcdtwy\",\n",
    "            env_var=\"HTLHCDTWY\",\n",
    "            default=\"no\",\n",
    "            type=str,\n",
    "            choices=[\"yes\", \"no\"]\n",
    "        ),\n",
    "    )\n",
    "    # don't care about Cluster Name, but the Hadron Collider is serious business\n",
    "    metadata = (\"user\", \"uuid\", \"htlhcdtwy\")\n",
    "    \n",
    "    TMP_FILE_PATH = \"/tmp/snafu-pingtest\"\n",
    "\n",
    "    def setup(self) -> bool:\n",
    "        \"\"\"\n",
    "        Setup the Ping Test Benchmark.\n",
    "        \n",
    "        This method creates a temporary file at ``/tmp/snafu-pingtest`` to let others \n",
    "        know that the benchmark is currently running.\n",
    "        \n",
    "        Returns\n",
    "        -------\n",
    "        bool\n",
    "            True if the temporary file was created successfully, othewise False. Will\n",
    "            also return False if the temporary file already exists.\n",
    "        \"\"\"\n",
    "        \n",
    "        \n",
    "        if check_file(self.TMP_FILE_PATH):\n",
    "            # The benchmark base class exposes a logger at self.logger which we can use\n",
    "            self.logger.critical(\n",
    "                f\"Temporary file located at {self.TMP_FILE_PATH} already exists.\"\n",
    "            )\n",
    "            return False\n",
    "        \n",
    "        try:\n",
    "            tmp_file = open(self.TMP_FILE_PATH, \"x\")\n",
    "            tmp_file.close()\n",
    "        except Exception as e:\n",
    "            self.logger.critical(\n",
    "                f\"Unable to create temporary file at {self.TMP_FILE_PATH}: {e}\",\n",
    "                exc_info=True\n",
    "            )\n",
    "            \n",
    "            return False\n",
    "        else:\n",
    "            self.logger.info(\n",
    "                f\"Successfully created temp file at {self.TMP_FILE_PATH}\"\n",
    "            )\n",
    "            return True\n",
    "\n",
    "    def collect(self):\n",
    "        \"\"\"Run the Ping Test Benchmark and collect results.\"\"\"\n",
    "        pass\n",
    "\n",
    "    def cleanup(self) -> bool:\n",
    "        \"\"\"\n",
    "        Cleanup the Ping Test Benchmark.\n",
    "        \n",
    "        This method removes the temporary file at ``/tmp/snafu-pingtest`` to let others\n",
    "        know that the benchmark has finished running.\n",
    "        \n",
    "        Returns\n",
    "        -------\n",
    "        bool\n",
    "            True if the temporary file was deleted successfully, otherwise False.\n",
    "        \"\"\"\n",
    "        \n",
    "        try:\n",
    "            os.remove(self.TMP_FILE_PATH)\n",
    "        except Exception as e:\n",
    "            self.logger.critical(\n",
    "                f\"Unable to remove temporary file at {self.TMP_FILE_PATH}: {e}\",\n",
    "                exc_info=True\n",
    "            )\n",
    "            \n",
    "            return False\n",
    "        else:\n",
    "            self.logger.info(\n",
    "                f\"Successfully removed temp file at {self.TMP_FILE_PATH}\"\n",
    "            )\n",
    "            return True\n",
    "        \n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2546724a",
   "metadata": {},
   "source": [
    "And again, some quick tests just to verify it works as expected:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a704e0f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from snafu.registry import TOOLS\n",
    "pingtest = TOOLS[\"pingtest\"]()\n",
    "\n",
    "!rm -f /tmp/snafu-pingtest\n",
    "\n",
    "# No file exists, so should error\n",
    "print(f\"Cleanup result is {pingtest.cleanup()}\")\n",
    "\n",
    "# Create the file using setup(), then cleanup()\n",
    "print(f\"Setup result is {pingtest.setup()}\")\n",
    "print(f\"Cleanup result is {pingtest.cleanup()}\")\n",
    "\n",
    "# Cleanup\n",
    "del pingtest"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9f6372cb",
   "metadata": {},
   "source": [
    "Now we have our setup and cleanup methods good to go, let's get to the fun part."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8fb9aa4f",
   "metadata": {},
   "source": [
    "## Part Four: Collect Method\n",
    "\n",
    "The collect method is an iterable that returns a special dataclass that is shipped with benchmark-wrapper, called a ``BenchmarkResult``. BenchmarkResult holds important information about a benchmark's resulting data, such as the configuration, metadata, labels and numerical data. It also understands how to prepare itself for export. All Benchmarks are expected to return their results using this dataclass in order to support a common interface for data exporters, reduce code reuse, and reduce extra overhead.\n",
    "\n",
    "The base Benchmark class includes a helpful method called ``create_new_result``, which we will use in the example below.\n",
    "\n",
    "For our ping test benchmark, our collect method needs to run the ping command, parse its output, and yield a new BenchmarkResult. To help prevent the collect method itself from becoming super large and out of control, we'll create some new methods in our wrapper class that the collect method will call to help do its thing. Benchmarks can have any number of additional methods, just as long as they have setup, collect and cleanup.\n",
    "\n",
    "One last note here before the code: benchmark-wrapper ships with another helpful module called ``process``, which contains functions and classes to facilitate running subprocesses. In particular, we'll be using the ``get_process_sample`` function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e2597c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "#!/usr/bin/env python3\n",
    "# -*- coding: utf-8 -*-\n",
    "\"\"\"Ping hosts and export results.\"\"\"\n",
    "import os\n",
    "from snafu.config import ConfigArgument, check_file\n",
    "from snafu.benchmarks import Benchmark\n",
    "\n",
    "# Grab the get_process_sample function, the BenchmarkResult class, stuff\n",
    "# for type hints, and dataclasses for storing our ping results\n",
    "from snafu.process import get_process_sample, ProcessSample, ProcessRun\n",
    "from snafu.benchmarks import BenchmarkResult\n",
    "from typing import Iterable, Optional\n",
    "from dataclasses import dataclass, asdict\n",
    "# Also shlex and subprocess for creating our ping command\n",
    "import shlex\n",
    "import subprocess\n",
    "# And finally re for regex\n",
    "import re\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class PingResult:\n",
    "    ip: Optional[str] = None\n",
    "    success: Optional[bool] = None\n",
    "    fail_msg: Optional[str] = None\n",
    "    host: Optional[str] = None\n",
    "    transmitted: Optional[int] = None\n",
    "    received: Optional[int] = None\n",
    "    packet_loss: Optional[float] = None\n",
    "    packet_bytes: Optional[int] = None\n",
    "    time_ms: Optional[float] = None\n",
    "    rtt_min_ms: Optional[float] = None\n",
    "    rtt_avg_ms: Optional[float] = None\n",
    "    rtt_max_ms: Optional[float] = None\n",
    "    rtt_mdev_ms: Optional[float] = None\n",
    "\n",
    "\n",
    "class PingTest(Benchmark):\n",
    "    \"\"\"Wrapper for the Ping Test benchmark.\"\"\"\n",
    "\n",
    "    tool_name = \"pingtest\"\n",
    "    args = (\n",
    "        ConfigArgument(\n",
    "            \"--host\",\n",
    "            help=\"Host(s) to ping. Can give more than one host by separating them with \"\n",
    "                 \"spaces on the CLI and in config files, or by giving them in a \"\n",
    "                 \"pythonic-list format through the OS environment\",\n",
    "            dest=\"host\",\n",
    "            nargs=\"+\",\n",
    "            env_var=\"HOST\",\n",
    "            type=str,\n",
    "            required=True,\n",
    "        ),\n",
    "        ConfigArgument(\n",
    "            \"--count\",\n",
    "            help=\"Number of pings to perform per sample.\",\n",
    "            dest=\"count\",\n",
    "            env_var=\"COUNT\",\n",
    "            default=1,\n",
    "            type=int,\n",
    "        ),\n",
    "        ConfigArgument(\n",
    "            \"--samples\",\n",
    "            help=\"Number of samples to perform.\",\n",
    "            dest=\"samples\",\n",
    "            env_var=\"SAMPLES\",\n",
    "            default=1,\n",
    "            type=int,\n",
    "        ),\n",
    "        ConfigArgument(\n",
    "            \"--htlhcdtwy\",\n",
    "            help=\"Has The Large Hadron Collider Destroyed The World Yet?\",\n",
    "            dest=\"htlhcdtwy\",\n",
    "            env_var=\"HTLHCDTWY\",\n",
    "            default=\"no\",\n",
    "            type=str,\n",
    "            choices=[\"yes\", \"no\"]\n",
    "        ),\n",
    "    )\n",
    "    # don't care about Cluster Name, but the Hadron Collider is serious business\n",
    "    metadata = (\"user\", \"uuid\", \"htlhcdtwy\")\n",
    "    \n",
    "    TMP_FILE_PATH = \"/tmp/snafu-pingtest\"\n",
    "\n",
    "    HOST_RE = r\"PING ([a-zA-Z0-9.-]+) \\(([0-9.]+)\\) \\d+\\(([\\d]+)\\) bytes of data\\.\"\n",
    "    RTT_STATS_RE = r\"rtt min\\/avg\\/max\\/mdev = ([\\d.]+)\\/([\\d.]+)\\/([\\d.]+)\\/([\\d.]+) ms\"\n",
    "    PACKET_RE = r\"(\\d+) packets transmitted, (\\d+) received, ([\\d.]+)\\% packet loss, \" \\\n",
    "                r\"time (\\d+)ms\"\n",
    "    \n",
    "    def setup(self) -> bool:\n",
    "        \"\"\"\n",
    "        Setup the Ping Test Benchmark.\n",
    "        \n",
    "        This method creates a temporary file at ``/tmp/snafu-pingtest`` to let others \n",
    "        know that the benchmark is currently running.\n",
    "        \n",
    "        Returns\n",
    "        -------\n",
    "        bool\n",
    "            True if the temporary file was created successfully, othewise False. Will\n",
    "            also return False if the temporary file already exists.\n",
    "        \"\"\"\n",
    "        \n",
    "        \n",
    "        if check_file(self.TMP_FILE_PATH):\n",
    "            # The benchmark base class exposes a logger at self.logger which we can use\n",
    "            self.logger.critical(\n",
    "                f\"Temporary file located at {self.TMP_FILE_PATH} already exists.\"\n",
    "            )\n",
    "            return False\n",
    "        \n",
    "        try:\n",
    "            tmp_file = open(self.TMP_FILE_PATH, \"x\")\n",
    "            tmp_file.close()\n",
    "        except Exception as e:\n",
    "            self.logger.critical(\n",
    "                f\"Unable to create temporary file at {self.TMP_FILE_PATH}: {e}\",\n",
    "                exc_info=True\n",
    "            )\n",
    "            \n",
    "            return False\n",
    "        else:\n",
    "            self.logger.info(\n",
    "                f\"Successfully created temp file at {self.TMP_FILE_PATH}\"\n",
    "            )\n",
    "            return True\n",
    "    \n",
    "    def parse_host_info(self, stdout: str, store: PingResult) -> None:\n",
    "        \"\"\"\n",
    "        Parse the host line of ping stdout.\n",
    "        \n",
    "        Expected format is: ``PING host (ip) data_bytes(ICMP_data_bytes) ...``.\n",
    "        \n",
    "        Parameters\n",
    "        ----------\n",
    "        stdout : str\n",
    "            ping stdout to parse\n",
    "        store : PingResult\n",
    "            PingResult instance to store parsed variables into\n",
    "        \"\"\"\n",
    "        \n",
    "        result = re.compile(self.HOST_RE).search(stdout)\n",
    "        if result is None:\n",
    "            self.logger.warning(f\"Unable to parse host info!\")\n",
    "            return\n",
    "        host, ip, data_size = result.groups()\n",
    "        data_size = int(data_size)\n",
    "        \n",
    "        if host == ip:\n",
    "            host = None  # user pinged an IP rather than a host\n",
    "        \n",
    "        store.host = host\n",
    "        store.ip = ip\n",
    "        store.packet_bytes = data_size\n",
    "    \n",
    "    def parse_packet_stats(self, stdout: str, store: PingResult) -> None:\n",
    "        \"\"\"\n",
    "        Parse the packet statistics line of ping stdout.\n",
    "        \n",
    "        Expected format is: \n",
    "        ``A packets transmitted, B received, C% packet loss, time Dms``\n",
    "        \n",
    "        Parameters\n",
    "        ----------\n",
    "        stdout : str\n",
    "            ping stdout to parse\n",
    "        store : PingResult\n",
    "            PingResult instance to store parsed variables into\n",
    "        \"\"\"\n",
    "        \n",
    "        result = re.compile(self.PACKET_RE).search(stdout)\n",
    "        if result is None:\n",
    "            self.logger.warning(\n",
    "                f\"Unable to parse packet stats!\"\n",
    "            )\n",
    "            return\n",
    "        transmitted, received, packet_loss, time_ms = map(int, result.groups())\n",
    "\n",
    "        store.transmitted = transmitted\n",
    "        store.received = received\n",
    "        store.packet_loss = packet_loss\n",
    "        store.time_ms = time_ms\n",
    "    \n",
    "    def parse_rtt_stats(self, stdout: str, store: PingResult) -> None:\n",
    "        \"\"\"\n",
    "        Parse the RTT statistics line of ping stdout.\n",
    "        \n",
    "        Expected format is: ``rtt min/avg/max/mdev = A/B/C/D ms``\n",
    "        \n",
    "        Parameters\n",
    "        ----------\n",
    "        stdout : str\n",
    "            ping stdout to parse\n",
    "        store : PingResult\n",
    "            PingResult instance to store parsed variables into\n",
    "        \"\"\"\n",
    "        \n",
    "        result = re.compile(self.RTT_STATS_RE).search(stdout)\n",
    "        if result is None:\n",
    "            self.logger.warning(f\"Unable to parse rtt stats!\")\n",
    "            return\n",
    "        rtt_min, rtt_avg, rtt_max, rtt_mdev = map(float, result.groups())\n",
    "        store.rtt_min_ms = rtt_min\n",
    "        store.rtt_avg_ms = rtt_avg\n",
    "        store.rtt_max_ms = rtt_max\n",
    "        store.rtt_mdev_ms = rtt_mdev\n",
    "        \n",
    "    def parse_stdout(self, stdout: str) -> PingResult:\n",
    "        \"\"\"\n",
    "        Parse the stdout of the ping command.\n",
    "        \n",
    "        Tested against ping from iputils 20210202 on Fedora Linux 34\n",
    "        \n",
    "        Returns\n",
    "        -------\n",
    "        PingResult\n",
    "        \"\"\"\n",
    "        \n",
    "        # We really only care about the first line, and the last two lines\n",
    "        lines = stdout.strip().split(\"\\n\")\n",
    "        \n",
    "        # Check if we got an error\n",
    "        if len(lines) == 1:\n",
    "            msg = lines[0]\n",
    "            return PingResult(\n",
    "                fail_msg=msg,\n",
    "                success=False\n",
    "            )\n",
    "        \n",
    "        \n",
    "        result = PingResult(success=True)\n",
    "        self.parse_host_info(stdout, result)\n",
    "        self.parse_packet_stats(stdout, result)\n",
    "        self.parse_rtt_stats(stdout, result)\n",
    "        \n",
    "        return result\n",
    "    \n",
    "    def ping_host(self, host: str) -> Iterable[BenchmarkResult]:\n",
    "        \"\"\"\n",
    "        Run the ping test benchmark against the given host.\n",
    "        \n",
    "        Parameters\n",
    "        ----------\n",
    "        host : str\n",
    "            Host to ping\n",
    "        \n",
    "        Returns\n",
    "        -------\n",
    "        iterable\n",
    "            Iterable of BenchmarkResults\n",
    "        \"\"\"\n",
    "        \n",
    "        self.logger.info(f\"Running ping test against host {host}\")\n",
    "        cmd = shlex.split(f\"ping -c {self.config.count} {host}\")\n",
    "        self.logger.debug(f\"Using command: {cmd}\")\n",
    "        \n",
    "        # A config instance allows for accessing params directly,\n",
    "        # therefore self.config.samples == self.config.params.samples\n",
    "        for sample_num in range(self.config.samples):\n",
    "            self.logger.info(f\"Collecting sample {sample_num}\")\n",
    "            \n",
    "            # We'll use the LiveProcess context manager to run ping\n",
    "            # LiveProcess will expose the Popen object at 'process',\n",
    "            # create a queue with lines from stdout at 'stdout',\n",
    "            # and create a snafu.process.ProcessRun instance at `attempt`\n",
    "            \n",
    "            # Here we will tell get_process_sample to send stdout and stderr\n",
    "            # to the same pipe\n",
    "            process_sample: ProcessSample = get_process_sample(\n",
    "                cmd, self.logger, stdout=subprocess.PIPE, stderr=subprocess.STDOUT\n",
    "            )\n",
    "            \n",
    "            self.logger.debug(f\"Got process sample: {vars(process_sample)}\")\n",
    "            if not process_sample.success:\n",
    "                self.logger.warning(f\"Process was unsuccessful\")\n",
    "                process_run: ProcessRun = process_sample.failed[0]\n",
    "            else:\n",
    "                self.logger.info(f\"Process was successful!\")\n",
    "                process_run: ProcessRun = process_sample.successful\n",
    "            \n",
    "            result: PingResult = self.parse_stdout(process_run.stdout)\n",
    "            # manually set host if we fail, since it won't always be parsable\n",
    "            # through stdout\n",
    "            if result.success is False:\n",
    "                result.host = host\n",
    "            self.logger.info(f\"Got sample: {vars(result)}\")\n",
    "            \n",
    "            yield self.create_new_result(\n",
    "                # We use vars here because create_new_result expects\n",
    "                # dict objects, not dataclasses\n",
    "                data=vars(result), \n",
    "                config={\"samples\": self.config.samples, \"count\": self.config.count},\n",
    "                # tag is a method for labeling results for exporters\n",
    "                # right now it specifies the ES index to export to\n",
    "                tag=\"jupyter\"\n",
    "            )\n",
    "        \n",
    "        plural = \"s\" if self.config.samples > 1 else \"\"\n",
    "        self.logger.info(\n",
    "            f\"Finised collecting {self.config.samples} sample{plural} against {host}\"\n",
    "        )\n",
    "\n",
    "    def collect(self) -> Iterable[BenchmarkResult]:\n",
    "        \"\"\"\n",
    "        Run the Ping Test Benchmark and collect results.\n",
    "        \"\"\"\n",
    "        \n",
    "        self.logger.info(\"Running pings and collecting results.\")\n",
    "        self.logger.debug(f\"Using config: {vars(self.config.params)}\")\n",
    "        if isinstance(self.config.host, str):\n",
    "            yield from self.ping_host(self.config.host)\n",
    "        else:\n",
    "            for host in self.config.host:\n",
    "                yield from self.ping_host(host)\n",
    "        self.logger.info(\"Finished\")\n",
    "\n",
    "    def cleanup(self) -> bool:\n",
    "        \"\"\"\n",
    "        Cleanup the Ping Test Benchmark.\n",
    "        \n",
    "        This method removes the temporary file at ``/tmp/snafu-pingtest`` to let others \n",
    "        know that the benchmark has finished running.\n",
    "        \n",
    "        Returns\n",
    "        -------\n",
    "        bool\n",
    "            True if the temporary file was deleted successfully, otherwise False.\n",
    "        \"\"\"\n",
    "        \n",
    "        try:\n",
    "            os.remove(self.TMP_FILE_PATH)\n",
    "        except Exception as e:\n",
    "            self.logger.critical(\n",
    "                f\"Unable to remove temporary file at {self.TMP_FILE_PATH}: {e}\",\n",
    "                exc_info=True\n",
    "            )\n",
    "            \n",
    "            return False\n",
    "        else:\n",
    "            self.logger.info(\n",
    "                f\"Successfully removed temp file at {self.TMP_FILE_PATH}\"\n",
    "            )\n",
    "            return True\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "909223e4",
   "metadata": {},
   "source": [
    "We have finished our new ping test benchmark! Let's try it out! We'll ping three hosts:\n",
    "\n",
    "* `www.google.com`: Depending on the build environment, this domain will show either 100% success (ICMP enabled) or 100% failure (ICMP disabled).\n",
    "* `www.idontexist.heythere`: A domain name which doesn't exist. Ping should exit with a failure saying that the host couldn't be resolved\n",
    "* `localhost`: We'll ping localhost, as we know regardless of environment we'll be able to ping ourselves."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50b793c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from snafu.registry import TOOLS\n",
    "from pprint import pprint\n",
    "import logging\n",
    "pingtest = TOOLS[\"pingtest\"]()\n",
    "\n",
    "# All Benchmark loggers work under the \"snafu\" logger\n",
    "logger = logging.getLogger(\"snafu\")\n",
    "if not logger.hasHandlers():\n",
    "    logger.addHandler(logging.StreamHandler())\n",
    "    logger.setLevel(logging.DEBUG)\n",
    "\n",
    "!rm -rf /tmp/snafu-pingtest\n",
    "!rm -f my_config.yaml\n",
    "\n",
    "# Set some config parameters\n",
    "# Config file\n",
    "!echo \"samples: 1\" > my_config.yaml\n",
    "!echo \"count: 5\" >> my_config.yaml\n",
    "# OS ENV\n",
    "import os\n",
    "os.environ[\"HOST\"] = \"[www.google.com,www.idontexist.heythere,localhost]\"\n",
    "\n",
    "# Parse arguments and print result\n",
    "# Since we aren't running within the main script (run_snafu.py),\n",
    "# need to add the config option manually\n",
    "pingtest.config.parser.add_argument(\"--config\", is_config_file=True)\n",
    "pingtest.config.parse_args(\n",
    "    \"--config my_config.yaml --labels=notebook=true --uuid 1337 --user snafu \"\n",
    "    \"--htlhcdtwy=no\".split(\" \")\n",
    ")\n",
    "\n",
    "# The base benchmark class includes a run method that runs setup -> collect -> cleanup\n",
    "results = list(pingtest.run())\n",
    "\n",
    "!rm -rf /tmp/snafu-pingtest\n",
    "!rm -f my_config.yaml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "246d6054",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Got {len(results)} results\")\n",
    "for result in results[:5]:\n",
    "    pprint(vars(result))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a164db7c",
   "metadata": {},
   "source": [
    "And that's that! As soon as you have your benchmark working that way you'd like, submit a PR and we'll give it a LGTM."
   ]
  }
 ],
 "metadata": {
  "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.6.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}