examples/FIL_streamz_pipeline.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# FIL and Streamz\n",
    "This is a guide on how [RAPIDS FIL (Forest Inference Library)](https://medium.com/rapids-ai/accelerating-random-forests-up-to-45x-using-cuml-dfb782a31bea) and [Streamz](https://streamz.readthedocs.io/en/latest/) can be used to build a streaming pipeline. In this example we use [IoT network traffic](https://www.stratosphereips.org/datasets-iot23).\n",
    "\n",
    "Streamz has the ability to read from Kafka directly into [Dask](https://dask.org/) allowing for computation on a multi-core or cluster environment. This approach is best used for instances in which you hope to increase processing speeds with streaming data.\n",
    "\n",
    "Here we share an example in which we demonstrate how to read connection log data from Kafka, run predictions using FIL and publish result data back to Kafka. To execute this notebook you will need to connect to an instance of Kafka. You can visit the [Apache Kafka Quick Start Guide](https://docs.confluent.io/current/quickstart/index.html?utm_medium=sem&utm_source=google&utm_campaign=ch.sem_br.nonbrand_tp.prs_tgt.kafka_mt.xct_rgn.namer_lng.eng_dv.all&utm_term=quickstart%20kafka&creative=&device=c&placement=&gclid=EAIaIQobChMInKfy66zL7AIVxRx9Ch2wKAmdEAAYASAAEgKEf_D_BwE#) to learn how to set up Kafka in your environment."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data Download"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download sample data - this may take some time to download\n",
    "!wget https://mcfp.felk.cvut.cz/publicDatasets/IoT-23-Dataset/iot_23_datasets_small.tar.gz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Extract data\n",
    "!tar -xzf iot_23_datasets_small.tar.gz"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With our kafka broker already running at `localhost:9092` and input kafka topic created, next we ingest our sample data into our topic named `input`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# To load the data into kafka use the command line tool kafka-console-producer provided by your kafka installation. In this example kafka is installed at /opt/kafka.\n",
    "# Update the broker-list and topic parameters as needed\n",
    "!/opt/kafka/bin/kafka-console-producer.sh --broker-list localhost:9092 --topic input < opt/Malware-Project/BigDataset/IoTScenarios/CTU-IoT-Malware-Capture-1-1/bro/conn.log.labeled"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "import confluent_kafka as ck\n",
    "import cudf\n",
    "import dask\n",
    "from dask_cuda import LocalCUDACluster\n",
    "from distributed import Client\n",
    "from streamz import Stream\n",
    "import time"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Parameters"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The average log size is used later in the notebook to estimate throughput and avg batch size bencmarks for streaming "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Benchmark\n",
    "avg_log_size=0.147 # in kilobytes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Provide the filepath of your FIL model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# FIL\n",
    "model_file=\"/path/to/model\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Kafka parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Kafka\n",
    "broker=\"localhost:9092\"\n",
    "input_topic=\"input\"\n",
    "output_topic=\"output\"\n",
    "\n",
    "\n",
    "producer_conf = {\n",
    "    \"bootstrap.servers\": broker,\n",
    "    \"session.timeout.ms\": 10000,\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Dask"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, create your dask cuda cluster and initialize each dask worker with the FIL model referenced above"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Start dask\n",
    "cluster = LocalCUDACluster()\n",
    "client = Client(cluster)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def worker_init():\n",
    "    # Initialization for each dask worker\n",
    "    from cuml import ForestInference\n",
    "    worker = dask.distributed.get_worker()\n",
    "    worker.data[\"fil_model\"] = ForestInference.load(filename=model_file,\n",
    "                          algo='BATCH_TREE_REORG',\n",
    "                          output_class=True,\n",
    "                          threshold=0.50,\n",
    "                          model_type='xgboost')\n",
    "    worker.data[\"data_columns\"] = [\"ts\",\"uid\",\"id.orig_h\",\"id.orig_p\",\"id.resp_h\",\"id.resp_p\",\"proto\",\"service\",\"duration\",\n",
    "               \"orig_bytes\",\"resp_bytes\",\"conn_state\",\"local_orig\",\"local_resp\",\"missed_bytes\",\"history\",\n",
    "                  \"orig_pkts\",\"orig_ip_bytes\",\"resp_pkts\",\"resp_ip_bytes\",\"label\"]\n",
    "    worker.data[\"data_types\"] = {\"ts\":\"float64\",\n",
    "            \"uid\":\"object\",\n",
    "            \"id.orig_h\":\"object\",\n",
    "            \"id.orig_p\":\"int64\",\n",
    "            \"id.resp_h\":\"object\",\n",
    "            \"id.resp_p\":\"int64\",\n",
    "            \"proto\":\"object\",\n",
    "            \"service\":\"object\",\n",
    "            \"duration\":\"object\",\n",
    "            \"orig_bytes\":\"object\",\n",
    "            \"resp_bytes\":\"object\",\n",
    "            \"conn_state\":\"object\",\n",
    "            \"local_orig\":\"object\",\n",
    "            \"local_resp\":\"object\",\n",
    "            \"missed_bytes\":\"int64\",\n",
    "            \"history\":\"object\",\n",
    "            \"orig_pkts\":\"int64\",\n",
    "            \"orig_ip_bytes\":\"int64\",\n",
    "            \"resp_pkts\":\"int64\",\n",
    "            \"resp_ip_bytes\":\"int64\",\n",
    "            \"label\":\"object\"}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'tcp://127.0.0.1:37781': None,\n",
       " 'tcp://127.0.0.1:38499': None,\n",
       " 'tcp://127.0.0.1:38895': None,\n",
       " 'tcp://127.0.0.1:39389': None,\n",
       " 'tcp://127.0.0.1:39429': None,\n",
       " 'tcp://127.0.0.1:41401': None,\n",
       " 'tcp://127.0.0.1:43393': None,\n",
       " 'tcp://127.0.0.1:44817': None}"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "client.run(worker_init)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<Client: 'tcp://127.0.0.1:46295' processes=8 threads=8, memory=540.94 GB>\n"
     ]
    }
   ],
   "source": [
    "print(client)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Streamz Pipeline\n",
    "\n",
    "Update the `max_batch_size` and `poll_interval` parameters as needed to tune your streamz workload to suit your environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_batch_size=1000000\n",
    "poll_interval=\"2s\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate a unique group_id to be able to re-run this demo notebook on the same data loaded to your kafka topic.\n",
    "j = random.randint(0,10000)\n",
    "group_id=\"fil-group-%d\" % j\n",
    "\n",
    "# Kafka consumer configuration\n",
    "consumer_conf = {\n",
    "    \"bootstrap.servers\": broker,\n",
    "    \"group.id\": group_id,\n",
    "    \"session.timeout.ms\": \"60000\",\n",
    "    \"enable.partition.eof\": \"true\",\n",
    "    \"auto.offset.reset\": \"earliest\",\n",
    "}\n",
    "\n",
    "# Stream source\n",
    "source = Stream.from_kafka_batched(\n",
    "        input_topic,\n",
    "        consumer_conf,\n",
    "        poll_interval=poll_interval,\n",
    "        npartitions=1,\n",
    "        asynchronous=True,\n",
    "        dask=True,\n",
    "        max_batch_size=max_batch_size\n",
    "\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we define the `predict` function to be used in the streamz pipeline. The predict function will construct a GPU dataframe of the raw log messages from kafka, format the data and then execute a prediction using the FIL model we previously loaded into Dask."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict(messages):\n",
    "    batch_start_time = int(round(time.time()))\n",
    "    worker = dask.distributed.get_worker()\n",
    "    df = cudf.DataFrame()\n",
    "    if type(messages) == str:\n",
    "       df[\"stream\"] = [messages.decode(\"utf-8\")]\n",
    "    elif type(messages) == list and len(messages) > 0:\n",
    "       df[\"stream\"] = [msg.decode(\"utf-8\") for msg in messages]\n",
    "    else:\n",
    "       print(\"ERROR: Unknown type encountered in inference\")\n",
    "    df_conn = df['stream'].str.split('\\t')\n",
    "    df_conn.columns = worker.data[\"data_columns\"]\n",
    "    df_conn=df_conn.astype(worker.data[\"data_types\"])\n",
    "    fil_preds = worker.data[\"fil_model\"].predict(df_conn[[\"orig_pkts\", \"orig_ip_bytes\", \"resp_pkts\", \"resp_ip_bytes\"]])\n",
    "    size = len(fil_preds)\n",
    "    return (fil_preds, batch_start_time, size)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `sink_to_kafka` function writes the output data or FIL predictions to the previously defined kafka topic."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sink_to_kafka(processed_data):\n",
    "    producer = ck.Producer(producer_conf)\n",
    "    json_str = processed_data[0].to_json(orient=\"records\", lines=True)\n",
    "    json_recs = json_str.split(\"\\n\")\n",
    "    num_recs = len(json_recs)\n",
    "    for idx,rec in enumerate(json_recs):\n",
    "        if idx % 50000 == 0:\n",
    "            producer.flush()\n",
    "        producer.produce(output_topic, rec)\n",
    "    producer.flush()\n",
    "    return processed_data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Below we define our streamz pipeline. This pipeline is also designed to capture benchmark data for reading and processing FIL predictions. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "output = source.map(predict).map(lambda x: (x[0], x[1], int(round(time.time())), x[2])).map(sink_to_kafka).gather().sink_to_list()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we start the streamz pipeline. View the progress on your dask dashboard http://localhost:8787"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "source.start()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This function calculates the benchmark. With each batch of data processed we have recorded the start and stop times that we can then use to calculate the total time difference. Throughput and avg batch size are estimates based on the average log size previously defined."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def calc_benchmark(results, size_per_log):\n",
    "    t1 = int(round(time.time() * 1000))\n",
    "    t2 = 0\n",
    "    size = 0.0\n",
    "    batch_count = 0\n",
    "    cnt = 0\n",
    "    # Find min and max time while keeping track of batch count and size\n",
    "    for result in results:\n",
    "        (ts1, ts2, result_size) = (result[1], result[2], result[3])\n",
    "        cnt += result_size\n",
    "        if ts1 == 0 or ts2 == 0:\n",
    "            continue\n",
    "        batch_count = batch_count + 1\n",
    "        t1 = min(t1, ts1)\n",
    "        t2 = max(t2, ts2)\n",
    "        size += result_size * size_per_log\n",
    "    time_diff = t2 - t1\n",
    "    throughput_mbps = size / (1024.0 * time_diff) if time_diff > 0 else 0\n",
    "    avg_batch_size = size / (1024.0 * batch_count) if batch_count > 0 else 0\n",
    "    return (time_diff, throughput_mbps, avg_batch_size, cnt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Please wait a few moments for all logs to be processed before calculating benchmarks  \n",
    "View the progress on the dask dashboard http://localhost:8787"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "max batch size: 1000000\n",
      "poll interval: 2s\n",
      "time (s): 48\n",
      "throughput (mb/s): 72.405251953125\n",
      "avg batch size (mb): 139.01808375000002\n",
      "num records: 24209952\n"
     ]
    }
   ],
   "source": [
    "benchmark = calc_benchmark(output, avg_log_size)\n",
    "print(\"max batch size:\", max_batch_size)\n",
    "print(\"poll interval:\", poll_interval)\n",
    "print(\"time (s):\", benchmark[0])\n",
    "print(\"throughput (mb/s):\", benchmark[1])\n",
    "print(\"avg batch size (mb):\", benchmark[2])\n",
    "print(\"num records:\", benchmark[3])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This end-to-end demonstration of FIL is intended to be able to optimize your data processing pipeline by utilizing the GPU. In this example, we've been able to process over 500,000 logs/s on 8 GPUs (Tesla V100). We hope to expand more in the future by integrating [cuStreamz](https://medium.com/rapids-ai/gpu-accelerated-stream-processing-with-rapids-f2b725696a61). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.7.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}