notebooks/webdataset_ingestion.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Parallel Data Ingestion from WebDataset\n",
    "\n",
    "This example ingests data from a [WebDataset](https://github.com/webdataset/webdataset) to a Space dataset. The ingestion operation uses Ray runner to distribute the workload in a Ray cluster.\n",
    "\n",
    "We use [img2dataset](https://github.com/rom1504/img2dataset) to download popular ML datasets in the WebDataset format. Install the packages and download the COCO dataset following the [guide](https://github.com/rom1504/img2dataset/blob/main/dataset_examples/mscoco.md):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "vscode": {
     "languageId": "shellscript"
    }
   },
   "outputs": [],
   "source": [
    "pip install webdataset img2dataset \n",
    "\n",
    "wget https://huggingface.co/datasets/ChristophSchuhmann/MS_COCO_2017_URL_TEXT/resolve/main/mscoco.parquet\n",
    "\n",
    "img2dataset --url_list mscoco.parquet --input_format \"parquet\" \\\n",
    "    --url_col \"URL\" --caption_col \"TEXT\" --output_format \"webdataset\" \\\n",
    "    --output_folder \"mscoco\" --processes_count 16 --thread_count 64 --image_size 256 \\\n",
    "    --enable_wandb True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Download the COCO WebDataset to \"COCO_DIR\" (local or in Cloud Storage). Then create an empty Space dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pyarrow as pa\n",
    "\n",
    "from space import DirCatalog\n",
    "from space import ArrayRecordOptions, FileOptions\n",
    "\n",
    "# The schema of a new Space dataset to create.\n",
    "schema = pa.schema([\n",
    "  (\"key\", pa.string()),\n",
    "  (\"caption\", pa.binary()),\n",
    "  (\"jpg\", pa.binary())])\n",
    "\n",
    "catalog = DirCatalog(\"/path/to/my/tables\")\n",
    "ds = catalog.create_dataset(\"coco\", schema, primary_keys=[\"key\"],\n",
    "  record_fields=[\"jpg\"]) # Store \"jpg\" in ArrayRecord files\n",
    "\n",
    "# Or load an existing dataset.\n",
    "# ds = catalog.dataset(\"images_size64\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Connect to a Ray runner and create a Ray runner for the dataset. See the [setup and performance doc](/docs/performance.md#ray-runner-setup) for how to configure the options."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ray\n",
    "\n",
    "# Connect to a Ray cluster. Or skip it to use a local Ray instance.\n",
    "ray.init(address=\"ray://12.34.56.78:10001\")\n",
    "\n",
    "runner = ds.ray(\n",
    "  ray_options=RayOptions(max_parallelism=4),\n",
    "  file_options=FileOptions(\n",
    "    array_record_options=ArrayRecordOptions(options=\"group_size:64\")\n",
    "  ))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A WebDataset consists of a directory of tar files. A URL in form of `something-{000000..012345}.tar` represents a shard of the dataset, i.e., a subset of tar files. The following `read_webdataset` method returns an iterator to scan the shard described by a URL."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import webdataset as wds\n",
    "\n",
    "# The size of a batch returned by the iterator.\n",
    "BATCH_SIZE = 64\n",
    "\n",
    "def read_webdataset(shard_url: str):\n",
    "  print(f\"Processing URL: {shard_url}\")\n",
    "\n",
    "  def to_dict(keys, captions, jpgs):\n",
    "    return {\"key\": keys, \"caption\": captions, \"jpg\": jpgs}\n",
    "\n",
    "  ds = wds.WebDataset(shard_url)\n",
    "  keys, captions, jpgs = [], [], []\n",
    "  for i, sample in enumerate(ds):\n",
    "    keys.append(sample[\"__key__\"])\n",
    "    captions.append(sample[\"txt\"])\n",
    "    jpgs.append(sample[\"jpg\"])\n",
    "\n",
    "    if len(keys) == BATCH_SIZE:\n",
    "      yield (to_dict(keys, captions, jpgs))\n",
    "      keys.clear()\n",
    "      captions.clear()\n",
    "      jpgs.clear()\n",
    "\n",
    "  if len(keys) > 0:\n",
    "    yield (to_dict(keys, captions, jpgs))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The whole COCO dataset has 60 tar files. We split it into  multiple shards to ingest data in parallel:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "NUM_SHARDS = 60\n",
    "NUM_WORKERS = 4\n",
    "\n",
    "COCO_DIR = \"/path/to/my/downloaded/coco\"\n",
    "\n",
    "shards_per_worker = NUM_SHARDS // NUM_WORKERS\n",
    "shard_urls = []\n",
    "for i in range(NUM_WORKERS):\n",
    "  start = f\"{i * shards_per_worker:05d}\"\n",
    "  end = f\"{min(NUM_SHARDS-1, (i + 1) * shards_per_worker - 1):05d}\"\n",
    "  shard_urls.append(COCO_DIR + f\"/{{{start}..{end}}}.tar\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `append_from` method takes a list of no-arg methods as input. Each method makes a new iterator that reads a shard of input data. The iterators are read in parallel on Ray nodes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from functools import partial\n",
    "\n",
    "# A Ray remote fn to call `append_from` from the Ray cluster.\n",
    "@ray.remote\n",
    "def run():\n",
    "  runner.append_from([partial(read_webdataset, url) for url in shard_urls])\n",
    "\n",
    "# Start ingestion\n",
    "ray.get(run.remote())"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}