s3fs-arrow-bandwidth.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ac7b089d-b128-403e-ae72-6aec681aaef7",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e840f019-140f-4807-8fe3-a0411beb29a9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1\n"
     ]
    }
   ],
   "source": [
    "!echo $OMP_NUM_THREADS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "9c8f0c06-3580-4dad-9a6d-b49fefbd4d74",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"OMP_NUM_THREADS\"] = \"8\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "19a60ae0-2763-437b-8548-2c7ae4ae7682",
   "metadata": {},
   "outputs": [],
   "source": [
    "import contextlib\n",
    "import pyarrow.parquet as pq\n",
    "from dask.utils import format_bytes\n",
    "import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "04731fb7-d84d-4342-afb1-8a1c3060ae4b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import s3fs\n",
    "s3 = s3fs.S3FileSystem()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "3ba34ae5-d3c3-4f87-a77a-c1f317cec64d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['coiled-runtime-ci/tpc-h/scale-10/lineitem/lineitem_02c36a39-8df5-493a-9d04-a763fc8d0465.parquet',\n",
       " 'coiled-runtime-ci/tpc-h/scale-10/lineitem/lineitem_04136638-fb9b-4a48-a420-02577a9e8880.parquet',\n",
       " 'coiled-runtime-ci/tpc-h/scale-10/lineitem/lineitem_049bc3cc-9a29-4605-87f5-22681bb2122e.parquet',\n",
       " 'coiled-runtime-ci/tpc-h/scale-10/lineitem/lineitem_04b2c18b-9c53-4d74-af58-71610c6dcd04.parquet',\n",
       " 'coiled-runtime-ci/tpc-h/scale-10/lineitem/lineitem_08dd8fe0-39c0-455c-bfe9-524528ed072c.parquet']"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "filenames = s3.glob(\"coiled-runtime-ci/tpc-h/scale-10/lineitem/*.parquet\")\n",
    "filenames[:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "8a806cf3-1f32-437e-b123-7f6e870ec7fe",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'2.87 GiB'"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "nbytes = s3.du(\"coiled-runtime-ci/tpc-h/scale-10/lineitem/\")\n",
    "format_bytes(nbytes )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "89bed71c-ae93-437a-8359-1404ff6fa20b",
   "metadata": {},
   "outputs": [],
   "source": [
    "@contextlib.contextmanager\n",
    "def bandwidth():\n",
    "    start = time.time()\n",
    "    yield\n",
    "    stop = time.time()\n",
    "    print(int(nbytes / (stop - start) // 2**20), \"MiB/s\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "e44d379e-fbe6-4665-b0c7-db43d6c61dcd",
   "metadata": {},
   "outputs": [],
   "source": [
    "from concurrent.futures import ThreadPoolExecutor"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cfccf3f9-d01d-4f41-89c5-a04b55cbbadf",
   "metadata": {},
   "source": [
    "## Test general bandwidth onto these machines"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "aee4b03f-d9d4-43f5-a6f0-ce3bf9e97f2f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "303 MiB/s\n"
     ]
    }
   ],
   "source": [
    "with bandwidth():\n",
    "    with ThreadPoolExecutor(4) as e:\n",
    "        def load(fn):\n",
    "            s3.open(fn).read()\n",
    "        list(e.map(load, filenames))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "bdfe6184-b332-4253-a34e-cedc3d6d92c7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "444 MiB/s\n"
     ]
    }
   ],
   "source": [
    "with bandwidth():\n",
    "    with ThreadPoolExecutor(6) as e:\n",
    "        def load(fn):\n",
    "            s3.open(fn).read()\n",
    "        list(e.map(load, filenames))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "5f309e61-83e9-4a54-af90-7d820bd4d96a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "554 MiB/s\n"
     ]
    }
   ],
   "source": [
    "with bandwidth():\n",
    "    with ThreadPoolExecutor(8) as e:\n",
    "        def load(fn):\n",
    "            s3.open(fn).read()\n",
    "        list(e.map(load, filenames))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b198d9b5-f2b1-400f-9c4e-b69032d27699",
   "metadata": {},
   "source": [
    "## ParquetFile\n",
    "\n",
    "We have three knobs to play with:\n",
    "\n",
    "-  `pre_buffer`: this is always going to be a good choice\n",
    "-  `use_threads`: it seems that we want this to be off\n",
    "-  `filesystem`: s3fs seems to work best with external threads\n",
    "\n",
    "*maybe these change with large `OMP_NUM_THREADS` but I haven't been able to make this happen."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "516ec3e2-f7fe-4210-9146-3bd34b5f19d6",
   "metadata": {},
   "source": [
    "### S3FS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "734daba8-7160-4b13-b38b-f7c56c36bc5d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "184 MiB/s\n"
     ]
    }
   ],
   "source": [
    "with bandwidth():\n",
    "    with ThreadPoolExecutor(4) as e:\n",
    "        def load(fn):\n",
    "            with s3.open(fn, mode=\"rb\") as f:\n",
    "                pq.ParquetFile(f, pre_buffer=True).read(\n",
    "                    use_threads=False,\n",
    "                    use_pandas_metadata=True,\n",
    "                )\n",
    "        list(e.map(load, filenames))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "20b8f1e1-1b4a-4dec-a7f2-8a1b9b52ed30",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "268 MiB/s\n"
     ]
    }
   ],
   "source": [
    "with bandwidth():\n",
    "    with ThreadPoolExecutor(8) as e:\n",
    "        def load(fn):\n",
    "            with s3.open(fn, mode=\"rb\") as f:\n",
    "                pq.ParquetFile(f, pre_buffer=True).read(\n",
    "                    use_threads=False,\n",
    "                    use_pandas_metadata=True,\n",
    "                )\n",
    "        list(e.map(load, filenames))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "1dcb397e-b347-4b89-9754-14d807c9f45a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "46 MiB/s\n"
     ]
    }
   ],
   "source": [
    "with bandwidth():\n",
    "    with ThreadPoolExecutor(8) as e:\n",
    "        def load(fn):\n",
    "            with s3.open(fn, mode=\"rb\") as f:\n",
    "                pq.ParquetFile(f, pre_buffer=False).read(\n",
    "                    use_threads=False,\n",
    "                    use_pandas_metadata=True,\n",
    "                )\n",
    "        list(e.map(load, filenames))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6acf24f-3db9-4490-be21-e3811914b7ab",
   "metadata": {},
   "source": [
    "### Arrow FS\n",
    "\n",
    "Less awesome unfortunately.  Should try again with OMP_NUM_THREADS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "e8969a7f-3c2a-4af3-af6c-a5edfadc53e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyarrow.fs import S3FileSystem\n",
    "import boto3\n",
    "session = boto3.session.Session()\n",
    "credentials = session.get_credentials()\n",
    "\n",
    "fs = S3FileSystem(\n",
    "    secret_key=credentials.secret_key,\n",
    "    access_key=credentials.access_key,\n",
    "    region=\"us-east-2\",\n",
    "    session_token=credentials.token,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "b4036f7b-4532-4404-9d88-a2845185763a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "136 MiB/s\n"
     ]
    }
   ],
   "source": [
    "with bandwidth():\n",
    "    with ThreadPoolExecutor(4) as e:\n",
    "        def load(fn):\n",
    "            pq.ParquetFile(\"s3://\" + fn, pre_buffer=True).read(\n",
    "                use_threads=False,\n",
    "                use_pandas_metadata=True,\n",
    "            )\n",
    "        list(e.map(load, filenames))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "2e99023f-f3c5-4eab-b09c-392221921b0e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "126 MiB/s\n"
     ]
    }
   ],
   "source": [
    "with bandwidth():\n",
    "    with ThreadPoolExecutor(4) as e:\n",
    "        def load(fn):\n",
    "            pq.ParquetFile(\"s3://\" + fn, pre_buffer=True).read(\n",
    "                use_threads=True,\n",
    "                use_pandas_metadata=True,\n",
    "            )\n",
    "        list(e.map(load, filenames))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "2645c274-fbda-4f4a-967c-74c2b6f00af0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "189 MiB/s\n"
     ]
    }
   ],
   "source": [
    "with bandwidth():\n",
    "    with ThreadPoolExecutor(8) as e:\n",
    "        def load(fn):\n",
    "            pq.ParquetFile(\"s3://\" + fn, pre_buffer=True).read(\n",
    "                use_threads=False,\n",
    "                use_pandas_metadata=True,\n",
    "            )\n",
    "        list(e.map(load, filenames))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "04a43d98-c921-4e88-868b-69bd0c62307e",
   "metadata": {},
   "source": [
    "1. pre_buffer is critical\n",
    "2. We shouldn't use the dedicated Arrow thread pool (or maybe this is OMP_NUM_THREADS specific)\n",
    "3. We should oversaturate cores"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63ece358-23f4-4c16-bd69-07b553a70c4d",
   "metadata": {},
   "source": [
    "I also tried using the parquet method here, but honestly I couldn't figure out how to invoke it.  Emperically I've found that it gives the best results so far."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "344cde23-f7f4-41fc-913e-972cb161c9f4",
   "metadata": {},
   "source": [
    "## Dask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "d3d65ec7-ef73-4a8e-9bae-1b9e0be48ea6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import dask\n",
    "import dask.dataframe as dd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "77b0d19b-e0ca-4209-9fad-ccafe7e086bd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "168 MiB/s\n"
     ]
    }
   ],
   "source": [
    "df = dd.read_parquet(\"s3://coiled-runtime-ci/tpc-h/scale-10/lineitem\")\n",
    "\n",
    "with bandwidth():\n",
    "    with ThreadPoolExecutor(4) as e:\n",
    "        with dask.config.set(pool=e):\n",
    "            df.map_partitions(lambda df: df.head(1)).persist(scheduler=\"threads\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "18028478-ae9f-4029-9764-b61cdca8d67f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "162 MiB/s\n"
     ]
    }
   ],
   "source": [
    "df = dd.read_parquet(\n",
    "    \"s3://coiled-runtime-ci/tpc-h/scale-10/lineitem\",\n",
    "    open_file_options={\"precache_options\": {\"method\": \"parquet\"}},\n",
    ")\n",
    "\n",
    "with bandwidth():\n",
    "    with ThreadPoolExecutor(4) as e:\n",
    "        with dask.config.set(pool=e):\n",
    "            df.map_partitions(lambda df: df.head(1)).persist(scheduler=\"threads\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "48d93c01-f5b5-41b7-863a-87b229bdbd10",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "227 MiB/s\n"
     ]
    }
   ],
   "source": [
    "df = dd.read_parquet(\"s3://coiled-runtime-ci/tpc-h/scale-10/lineitem\")\n",
    "\n",
    "with bandwidth():\n",
    "    with ThreadPoolExecutor(8) as e:\n",
    "        with dask.config.set(pool=e):\n",
    "            df.map_partitions(lambda df: df.head(1)).persist(scheduler=\"threads\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "0b97a712-ba0e-4e45-9927-27f6b9752275",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "204 MiB/s\n"
     ]
    }
   ],
   "source": [
    "df = dd.read_parquet(\n",
    "    \"s3://coiled-runtime-ci/tpc-h/scale-10/lineitem\",\n",
    "    open_file_options={\"precache_options\": {\"method\": \"parquet\"}},\n",
    ")\n",
    "\n",
    "with bandwidth():\n",
    "    with ThreadPoolExecutor(8) as e:\n",
    "        with dask.config.set(pool=e):\n",
    "            df.map_partitions(lambda df: df.head(1)).persist(scheduler=\"threads\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b23c3ae-fd16-4558-8caa-efdda2dae887",
   "metadata": {},
   "source": [
    "## Hybrid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f10b0023-b1c8-4e88-a96c-d0548a8d6dd6",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext snakeviz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "106c0cdd-45c9-46b4-aefd-a2a8fc04e581",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%snakeviz\n",
    "\n",
    "import io\n",
    "from dask.local import synchronous_executor\n",
    "with bandwidth():\n",
    "    with synchronous_executor as e:\n",
    "        def load(fn):\n",
    "            bytes = s3.open(fn).read()\n",
    "            f = io.BytesIO(bytes)\n",
    "            df = pq.ParquetFile(f, pre_buffer=True).read(\n",
    "                use_threads=False,\n",
    "                use_pandas_metadata=True,\n",
    "            )\n",
    "\n",
    "        list(e.map(load, filenames))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "675abaef-a515-425a-b826-5236ef6b427e",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}