shuffle-linearizers-v1.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import webdataset as wds\n",
    "import scipy.ndimage as ndi\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from itertools import islice\n",
    "import ray\n",
    "from ray.data import Dataset\n",
    "import braceexpand\n",
    "import requests\n",
    "import random\n",
    "import json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Usage stats collection is enabled by default for nightly wheels. To disable this, run the following command: `ray disable-usage-stats` before starting Ray. See https://docs.ray.io/en/master/cluster/usage-stats.html for more details.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-05-11 23:14:13,835\tINFO worker.py:1607 -- Started a local Ray instance. View the dashboard at \u001b[1m\u001b[32mhttp://127.0.0.1:8265 \u001b[39m\u001b[22m\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "    <div style=\"margin-left: 50px;display: flex;flex-direction: row;align-items: center\">\n",
       "        <h3 style=\"color: var(--jp-ui-font-color0)\">Ray</h3>\n",
       "        <svg version=\"1.1\" id=\"ray\" width=\"3em\" viewBox=\"0 0 144.5 144.6\" style=\"margin-left: 3em;margin-right: 3em\">\n",
       "            <g id=\"layer-1\">\n",
       "                <path fill=\"#00a2e9\" class=\"st0\" d=\"M97.3,77.2c-3.8-1.1-6.2,0.9-8.3,5.1c-3.5,6.8-9.9,9.9-17.4,9.6S58,88.1,54.8,81.2c-1.4-3-3-4-6.3-4.1\n",
       "                    c-5.6-0.1-9.9,0.1-13.1,6.4c-3.8,7.6-13.6,10.2-21.8,7.6C5.2,88.4-0.4,80.5,0,71.7c0.1-8.4,5.7-15.8,13.8-18.2\n",
       "                    c8.4-2.6,17.5,0.7,22.3,8c1.3,1.9,1.3,5.2,3.6,5.6c3.9,0.6,8,0.2,12,0.2c1.8,0,1.9-1.6,2.4-2.8c3.5-7.8,9.7-11.8,18-11.9\n",
       "                    c8.2-0.1,14.4,3.9,17.8,11.4c1.3,2.8,2.9,3.6,5.7,3.3c1-0.1,2,0.1,3,0c2.8-0.5,6.4,1.7,8.1-2.7s-2.3-5.5-4.1-7.5\n",
       "                    c-5.1-5.7-10.9-10.8-16.1-16.3C84,38,81.9,37.1,78,38.3C66.7,42,56.2,35.7,53,24.1C50.3,14,57.3,2.8,67.7,0.5\n",
       "                    C78.4-2,89,4.7,91.5,15.3c0.1,0.3,0.1,0.5,0.2,0.8c0.7,3.4,0.7,6.9-0.8,9.8c-1.7,3.2-0.8,5,1.5,7.2c6.7,6.5,13.3,13,19.8,19.7\n",
       "                    c1.8,1.8,3,2.1,5.5,1.2c9.1-3.4,17.9-0.6,23.4,7c4.8,6.9,4.6,16.1-0.4,22.9c-5.4,7.2-14.2,9.9-23.1,6.5c-2.3-0.9-3.5-0.6-5.1,1.1\n",
       "                    c-6.7,6.9-13.6,13.7-20.5,20.4c-1.8,1.8-2.5,3.2-1.4,5.9c3.5,8.7,0.3,18.6-7.7,23.6c-7.9,5-18.2,3.8-24.8-2.9\n",
       "                    c-6.4-6.4-7.4-16.2-2.5-24.3c4.9-7.8,14.5-11,23.1-7.8c3,1.1,4.7,0.5,6.9-1.7C91.7,98.4,98,92.3,104.2,86c1.6-1.6,4.1-2.7,2.6-6.2\n",
       "                    c-1.4-3.3-3.8-2.5-6.2-2.6C99.8,77.2,98.9,77.2,97.3,77.2z M72.1,29.7c5.5,0.1,9.9-4.3,10-9.8c0-0.1,0-0.2,0-0.3\n",
       "                    C81.8,14,77,9.8,71.5,10.2c-5,0.3-9,4.2-9.3,9.2c-0.2,5.5,4,10.1,9.5,10.3C71.8,29.7,72,29.7,72.1,29.7z M72.3,62.3\n",
       "                    c-5.4-0.1-9.9,4.2-10.1,9.7c0,0.2,0,0.3,0,0.5c0.2,5.4,4.5,9.7,9.9,10c5.1,0.1,9.9-4.7,10.1-9.8c0.2-5.5-4-10-9.5-10.3\n",
       "                    C72.6,62.3,72.4,62.3,72.3,62.3z M115,72.5c0.1,5.4,4.5,9.7,9.8,9.9c5.6-0.2,10-4.8,10-10.4c-0.2-5.4-4.6-9.7-10-9.7\n",
       "                    c-5.3-0.1-9.8,4.2-9.9,9.5C115,72.1,115,72.3,115,72.5z M19.5,62.3c-5.4,0.1-9.8,4.4-10,9.8c-0.1,5.1,5.2,10.4,10.2,10.3\n",
       "                    c5.6-0.2,10-4.9,9.8-10.5c-0.1-5.4-4.5-9.7-9.9-9.6C19.6,62.3,19.5,62.3,19.5,62.3z M71.8,134.6c5.9,0.2,10.3-3.9,10.4-9.6\n",
       "                    c0.5-5.5-3.6-10.4-9.1-10.8c-5.5-0.5-10.4,3.6-10.8,9.1c0,0.5,0,0.9,0,1.4c-0.2,5.3,4,9.8,9.3,10\n",
       "                    C71.6,134.6,71.7,134.6,71.8,134.6z\"/>\n",
       "            </g>\n",
       "        </svg>\n",
       "        <table>\n",
       "            <tr>\n",
       "                <td style=\"text-align: left\"><b>Python version:</b></td>\n",
       "                <td style=\"text-align: left\"><b>3.10.6</b></td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                <td style=\"text-align: left\"><b>Ray version:</b></td>\n",
       "                <td style=\"text-align: left\"><b> 3.0.0.dev0</b></td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "    <td style=\"text-align: left\"><b>Dashboard:</b></td>\n",
       "    <td style=\"text-align: left\"><b><a href=\"http://127.0.0.1:8265\" target=\"_blank\">http://127.0.0.1:8265</a></b></td>\n",
       "</tr>\n",
       "\n",
       "        </table>\n",
       "    </div>\n",
       "</div>\n"
      ],
      "text/plain": [
       "RayContext(dashboard_url='127.0.0.1:8265', python_version='3.10.6', ray_version='3.0.0.dev0', ray_commit='384ad04987cd6c96d72c97213cb9b21a93d3f29e', address_info={'node_ip_address': '10.138.0.19', 'raylet_ip_address': '10.138.0.19', 'redis_address': None, 'object_store_address': '/tmp/ray/session_2023-05-11_23-14-11_607356_576487/sockets/plasma_store', 'raylet_socket_name': '/tmp/ray/session_2023-05-11_23-14-11_607356_576487/sockets/raylet', 'webui_url': '127.0.0.1:8265', 'session_dir': '/tmp/ray/session_2023-05-11_23-14-11_607356_576487', 'metrics_export_port': 64678, 'gcs_address': '10.138.0.19:64929', 'address': '10.138.0.19:64929', 'dashboard_agent_listen_port': 52365, 'node_id': '9bb654cfc70684de80f3d2a51c569ac3a9fcb0ece5107e49e94f6078'})"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ray.init()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "gs://ocro-iaa/lin/lin-{000000..000703}.tar\n"
     ]
    }
   ],
   "source": [
    "!gsutil ls gs://ocro-iaa/lin/ | shardsum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "source = list(braceexpand.braceexpand(\"gs://ocro-iaa/lin/lin-{000000..000703}.tar\"))\n",
    "destination = \"gs://ocro-tempout\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-r--r--r-- bigdata/bigdata 17251 2023-05-11 06:56 67656e65736973776869746566616d30306d616c74676f6f67/00000.lin.json\n",
      "-r--r--r-- bigdata/bigdata     2 2023-05-11 06:56 67656e65736973776869746566616d30306d616c74676f6f67/00001.lin.json\n",
      "-r--r--r-- bigdata/bigdata   204 2023-05-11 06:56 67656e65736973776869746566616d30306d616c74676f6f67/00002.lin.json\n",
      "-r--r--r-- bigdata/bigdata   153 2023-05-11 06:56 67656e65736973776869746566616d30306d616c74676f6f67/00003.lin.json\n",
      "-r--r--r-- bigdata/bigdata   153 2023-05-11 06:56 67656e65736973776869746566616d30306d616c74676f6f67/00004.lin.json\n",
      "tar: stdout: write error\n"
     ]
    }
   ],
   "source": [
    "for s in source:\n",
    "    !gsutil cat $s | tar tvf - | sed 5q\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "!env | grep KEYS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-05-11 23:14:22,325\tWARNING dataset.py:255 -- \u001b[33mImportant: Ray Data requires schemas for all datasets in Ray 2.5. This means that standalone Python objects are no longer supported. In addition, the default batch format is fixed to NumPy. To revert to legacy behavior temporarily, set the environment variable RAY_DATA_STRICT_MODE=0 on all cluster processes.\n",
      "\n",
      "Learn more here: https://docs.ray.io/en/master/data/faq.html#migrating-to-strict-mode\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "dataset = ray.data.read_webdataset(source, parallelism=9999)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def page_classifier(lin, do_plot=False, **kw):\n",
    "    ys = []\n",
    "    for text, bbox in lin:\n",
    "        if bbox is None: continue\n",
    "        bbox = map(int, bbox.split())\n",
    "        x0, y0, x1, y1 = bbox\n",
    "        ys.append(float(y1))\n",
    "    if len(ys) < 10:\n",
    "        return \"mostly-empty\"\n",
    "    xs = np.linspace(0, 1, len(ys))\n",
    "    ys = ndi.median_filter(ys, 5)\n",
    "    ys = ndi.gaussian_filter(ys, 20.0)\n",
    "    ys -= np.amin(ys)\n",
    "    if np.amax(ys) < 1e-3:\n",
    "        return \"mostly-empty\"\n",
    "    ys /= np.amax(ys)\n",
    "    ysm = np.maximum.accumulate(ys)\n",
    "    delta = ys - ysm\n",
    "    single = np.all(delta > -0.1)\n",
    "    if single:\n",
    "        result = \"single\"\n",
    "    elif np.mean(ys[:len(ys)//4]) > 0.5:\n",
    "        result = \"backwards\"\n",
    "    else:\n",
    "        result = \"multi\"\n",
    "    color = dict(single=\"green\", backwards=\"red\", multi=\"blue\")[result]\n",
    "    if do_plot:\n",
    "        plt.plot(xs, ys, color=color, **kw)\n",
    "    return result\n",
    "\n",
    "def balance_samples(sample):\n",
    "    page = sample[\"lin.json\"]\n",
    "    kind = page_classifier(page)\n",
    "    if kind == \"single\":\n",
    "        if random.uniform(0, 1) < 0.1:\n",
    "            return True\n",
    "        else:\n",
    "            return False\n",
    "    elif kind == \"multi\":\n",
    "        return True\n",
    "    else:\n",
    "        return False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-05-11 23:14:22,650\tWARNING dataset.py:4352 -- The `map`, `flat_map`, and `filter` operations are unvectorized and can be very slow. If you're using a vectorized transformation, consider using `.map_batches()` instead.\n"
     ]
    }
   ],
   "source": [
    "filtered_dataset = dataset.filter(balance_samples)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "shuffled_dataset = filtered_dataset.random_shuffle()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-05-11 23:14:23,814\tINFO streaming_executor.py:91 -- Executing DAG InputDataBuffer[Input] -> AllToAllOperator[ReadWebDataset->Filter->RandomShuffle] -> TaskPoolMapOperator[Write]\n",
      "2023-05-11 23:14:23,816\tINFO streaming_executor.py:92 -- Execution config: ExecutionOptions(resource_limits=ExecutionResources(cpu=None, gpu=None, object_store_memory=None), locality_with_output=False, preserve_order=False, actor_locality_enabled=True, verbose_progress=False)\n",
      "2023-05-11 23:14:23,818\tINFO streaming_executor.py:94 -- Tip: For detailed progress reporting, run `ray.data.DataContext.get_current().execution_options.verbose_progress = True`\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3593d97ea55349308205d3cb83ae1635",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "- ReadWebDataset->Filter->RandomShuffle 1:   0%|          | 0/704 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9ef82c169b654ac49e883dff570592b5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "ShuffleMap 2:   0%|          | 0/704 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6c17fb38acdb4241956a1d86eb9c2d7c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "ShuffleReduce 3:   0%|          | 0/704 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f213439f63c144cba6eaefd389fe4afa",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Running 0:   0%|          | 0/704 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[2m\u001b[36m(map pid=579808)\u001b[0m /tmp/ipykernel_576487/3262028565.py:14: RuntimeWarning: invalid value encountered in divide\n",
      "\u001b[2m\u001b[36m(map pid=578096)\u001b[0m /tmp/ipykernel_576487/3262028565.py:14: RuntimeWarning: invalid value encountered in divide\u001b[32m [repeated 2x across cluster] (Ray deduplicates logs by default. Set RAY_DEDUP_LOGS=0 to disable log deduplication, or see https://docs.ray.io/en/master/ray-observability/ray-logging.html#log-deduplication for more options.)\u001b[0m\n",
      "\u001b[2m\u001b[36m(map pid=578102)\u001b[0m /tmp/ipykernel_576487/3262028565.py:14: RuntimeWarning: invalid value encountered in divide\n",
      "\u001b[2m\u001b[36m(map pid=578105)\u001b[0m /tmp/ipykernel_576487/3262028565.py:14: RuntimeWarning: invalid value encountered in divide\n",
      "\u001b[2m\u001b[36m(map pid=578105)\u001b[0m /tmp/ipykernel_576487/3262028565.py:14: RuntimeWarning: invalid value encountered in divide\n",
      "\u001b[2m\u001b[36m(map pid=578095)\u001b[0m /tmp/ipykernel_576487/3262028565.py:14: RuntimeWarning: invalid value encountered in divide\n",
      "\u001b[2m\u001b[36m(map pid=578100)\u001b[0m /tmp/ipykernel_576487/3262028565.py:14: RuntimeWarning: invalid value encountered in divide\n",
      "\u001b[2m\u001b[36m(map pid=578099)\u001b[0m /tmp/ipykernel_576487/3262028565.py:14: RuntimeWarning: invalid value encountered in divide\n",
      "\u001b[2m\u001b[36m(map pid=578110)\u001b[0m /tmp/ipykernel_576487/3262028565.py:14: RuntimeWarning: invalid value encountered in divide\u001b[32m [repeated 2x across cluster]\u001b[0m\n",
      "\u001b[2m\u001b[36m(map pid=578107)\u001b[0m /tmp/ipykernel_576487/3262028565.py:14: RuntimeWarning: invalid value encountered in divide\n",
      "\u001b[2m\u001b[36m(map pid=578102)\u001b[0m /tmp/ipykernel_576487/3262028565.py:14: RuntimeWarning: invalid value encountered in divide\n",
      "\u001b[2m\u001b[36m(map pid=578108)\u001b[0m /tmp/ipykernel_576487/3262028565.py:14: RuntimeWarning: invalid value encountered in divide\n",
      "\u001b[2m\u001b[36m(map pid=578103)\u001b[0m /tmp/ipykernel_576487/3262028565.py:14: RuntimeWarning: invalid value encountered in divide\n",
      "\u001b[2m\u001b[36m(map pid=578105)\u001b[0m /tmp/ipykernel_576487/3262028565.py:14: RuntimeWarning: invalid value encountered in divide\u001b[32m [repeated 2x across cluster]\u001b[0m\n",
      "\u001b[2m\u001b[36m(map pid=578101)\u001b[0m /tmp/ipykernel_576487/3262028565.py:14: RuntimeWarning: invalid value encountered in divide\n",
      "\u001b[2m\u001b[36m(map pid=579808)\u001b[0m /tmp/ipykernel_576487/3262028565.py:14: RuntimeWarning: invalid value encountered in divide\n",
      "\u001b[2m\u001b[36m(map pid=578105)\u001b[0m /tmp/ipykernel_576487/3262028565.py:14: RuntimeWarning: invalid value encountered in divide\n",
      "\u001b[2m\u001b[36m(map pid=578106)\u001b[0m /tmp/ipykernel_576487/3262028565.py:14: RuntimeWarning: invalid value encountered in divide\n",
      "\u001b[2m\u001b[36m(map pid=578102)\u001b[0m /tmp/ipykernel_576487/3262028565.py:14: RuntimeWarning: invalid value encountered in divide\n",
      "\u001b[2m\u001b[36m(map pid=578100)\u001b[0m /tmp/ipykernel_576487/3262028565.py:14: RuntimeWarning: invalid value encountered in divide\n",
      "\u001b[2m\u001b[36m(map pid=578098)\u001b[0m /tmp/ipykernel_576487/3262028565.py:14: RuntimeWarning: invalid value encountered in divide\u001b[32m [repeated 2x across cluster]\u001b[0m\n",
      "\u001b[2m\u001b[36m(map pid=578098)\u001b[0m /tmp/ipykernel_576487/3262028565.py:14: RuntimeWarning: invalid value encountered in divide\n",
      "\u001b[2m\u001b[36m(map pid=578096)\u001b[0m /tmp/ipykernel_576487/3262028565.py:14: RuntimeWarning: invalid value encountered in divide\n",
      "\u001b[2m\u001b[36m(map pid=579808)\u001b[0m /tmp/ipykernel_576487/3262028565.py:14: RuntimeWarning: invalid value encountered in divide\n",
      "\u001b[2m\u001b[36m(map pid=578108)\u001b[0m /tmp/ipykernel_576487/3262028565.py:14: RuntimeWarning: invalid value encountered in divide\n",
      "\u001b[2m\u001b[36m(map pid=578110)\u001b[0m /tmp/ipykernel_576487/3262028565.py:14: RuntimeWarning: invalid value encountered in divide\n",
      "\u001b[2m\u001b[36m(map pid=579808)\u001b[0m /tmp/ipykernel_576487/3262028565.py:14: RuntimeWarning: invalid value encountered in divide\n",
      "\u001b[2m\u001b[36m(map pid=578099)\u001b[0m /tmp/ipykernel_576487/3262028565.py:14: RuntimeWarning: invalid value encountered in divide\u001b[32m [repeated 2x across cluster]\u001b[0m\n",
      "\u001b[2m\u001b[36m(map pid=578105)\u001b[0m /tmp/ipykernel_576487/3262028565.py:14: RuntimeWarning: invalid value encountered in divide\n"
     ]
    }
   ],
   "source": [
    "shuffled_dataset.write_webdataset(\"gs://ocro-tempout/\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "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.6"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}