{
    "cells": [
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {
                "id": "daY0p3RNvzFR"
            },
            "source": [
                "# Parallel dataset conversion\n",
                "\n",
                "If your dataset is huge, running single process dataset conversion script could be very time consuming. You can use multiprocessing with MDSWriter to convert your dataset in parallel. There are few ways in which you can convert your raw data into MDS format in parallel fashion.\n",
                "\n",
                "1. Download a raw data in parallel and convert to MDS format sequentially.\n",
                "2. Group raw data and convert in parallel to MDS format in separate sub-directories. Then, merge all the `index.json` files from these subdirectories to get a unified MDS dataset.\n",
                "\n",
                "Let's look at an example for each option."
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {
                "id": "VKTtgJjkvzFU"
            },
            "source": [
                "## 1. Fetch raw data in parallel and write sequentially\n",
                "For a dataset with large files (such as images or videos), it would be useful to download those files in parallel using multiple processes and call the MDSWriter to write the data into MDS format."
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {
                "id": "xq9Icki3uRRu"
            },
            "source": [
                "### Setup\n",
                "\n",
                "Let's start by installing the `streaming-community` package, and importing necessary dependencies."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {
                "id": "Wq_oqcwrv08e"
            },
            "outputs": [],
            "source": [
                "%pip install streaming-community"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {
                "id": "T-9o1UQivzFV"
            },
            "outputs": [],
            "source": [
                "import os\n",
                "from multiprocessing import Pool\n",
                "\n",
                "from streaming import MDSWriter, StreamingDataset"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {
                "id": "gjz0OwY9u04a"
            },
            "source": [
                "### Global settings\n",
                "\n",
                "Initialize global variables:"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {
                "id": "6HDCKMaTu8U3"
            },
            "outputs": [],
            "source": [
                "out_root = './data'\n",
                "# This could be a list of URLs needs to download\n",
                "dataset = [i for i in range(25)]\n",
                "columns = {'number': 'int'}"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {
                "id": "tq4NInVovzFW"
            },
            "source": [
                "Download data from remote URLs. Here, we just return a number for demonstration purposes. "
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {
                "id": "-NonDeBovzFW"
            },
            "outputs": [],
            "source": [
                "def get_data(number):\n",
                "    print(f'\\nWorker PID: {os.getpid()}\\tnumber: {number}', flush=True, end='')\n",
                "    # Add code here to downloads the data from URL.\n",
                "    return {'number': number}"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {
                "id": "yByl7cpsvzFX"
            },
            "source": [
                "An initialization method for each worker process which prints the worker PID."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {
                "id": "Ut8clRDvvzFY"
            },
            "outputs": [],
            "source": [
                "# Initialize the worker process\n",
                "def init_worker():\n",
                "    # Get the pid for the current worker process\n",
                "    pid = os.getpid()\n",
                "    print(f'\\nInitialize Worker PID: {pid}', flush=True, end='')"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {
                "id": "cjtN8ZEwvzFY"
            },
            "source": [
                "### Convert to MDS format\n",
                "\n",
                "Initialize 4 worker processes which download the data in parallel. Once the data is ready, it is written to MDS format using the `write` method of {class}`streaming.MDSWriter`."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {
                "id": "xV-SRpZkvzFY"
            },
            "outputs": [],
            "source": [
                "# clean up root directory\n",
                "%rm -rf $out_root\n",
                "\n",
                "with Pool(initializer=init_worker, processes=4) as pool:\n",
                "    with MDSWriter(out=out_root, columns=columns) as out:\n",
                "        for sample in pool.imap(get_data, dataset):\n",
                "            out.write(sample)"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {
                "id": "5cWmR7svvzFZ"
            },
            "source": [
                "### Load MDS dataset\n",
                "\n",
                "Read samples from MDS by iterating over `StreamingDataset`. Here, we just print sample IDs."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {
                "id": "e-iQeOecvzFZ"
            },
            "outputs": [],
            "source": [
                "# read the sample\n",
                "dataset = StreamingDataset(local=out_root,\n",
                "                           remote=None,\n",
                "                           shuffle=False,)\n",
                "for sample in dataset:\n",
                "    print(sample['number'])"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {
                "id": "WXF-OdO0wYb8"
            },
            "outputs": [],
            "source": [
                "# Clean up\n",
                "%rm -rf $out_root"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {
                "id": "WXybtE1e5Nzo"
            },
            "source": [
                "## 2. Group the raw data and convert to MDS format in parallel\n",
                "\n",
                "For large raw datasets, or raw datasets with large files, we recommend partitioning dataset conversion among multiple `MDSWriter`s. Dataset conversion will take place with multiple processes in parallel."
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {
                "id": "3GbXadoPJne7"
            },
            "source": [
                "Importing dependencies:"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {
                "id": "NcfrL7ynJmyF"
            },
            "outputs": [],
            "source": [
                "import os\n",
                "import json\n",
                "from glob import glob\n",
                "from typing import Iterator, Tuple\n",
                "\n",
                "from multiprocessing import Pool\n",
                "\n",
                "from streaming import MDSWriter, StreamingDataset"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {
                "id": "45Xz1lLSUnug"
            },
            "source": [
                "### Global settings\n",
                "\n",
                "Initializing needed global variables:"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {
                "id": "q8AHVCVoUoU5"
            },
            "outputs": [],
            "source": [
                "out_root = './group_data'\n",
                "num_groups = 4\n",
                "num_process = 2"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {
                "id": "1qkyMzckKl0V"
            },
            "source": [
                "This function yields a sub-directory path where MDS shards will be stored, as well as the raw dataset sample range of that directory. For example, the first sub-directory will contain samples 0 to 9, the second sub-directory will contain samples 10 to 19, and so on.\n",
                "\n",
                "If you are working with large files, you can also yield a single raw dataset file path instead of a sample range."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {
                "id": "PZWRXNKEKkSm"
            },
            "outputs": [],
            "source": [
                "def each_task(out_root: str, groups: int) -> Iterator[Tuple[str, int, int]]:\n",
                "    \"\"\"Get the sub-directory path and the sample range for each sub-directory.\n",
                "\n",
                "    Args:\n",
                "        out_root (str): base output mds directory\n",
                "        groups (int): Number of sub-directories to create\n",
                "\n",
                "    Yields:\n",
                "        Iterator[Tuple[str, int, int]]: Each argument tuple\n",
                "    \"\"\"\n",
                "    for data_group in range(groups):\n",
                "        sub_out_root = os.path.join(out_root, str(data_group))\n",
                "        start_sample_idx = data_group * 10\n",
                "        end_sample_idx = start_sample_idx + 9\n",
                "        yield sub_out_root, start_sample_idx, end_sample_idx"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {
                "id": "p9XqWLD-Moqz"
            },
            "source": [
                "This function converts raw dataset samples into MDS format. "
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {
                "id": "h6zB2yLVMpIb"
            },
            "outputs": [],
            "source": [
                "def convert_to_mds(args: Iterator[Tuple[str, int, int]]) -> None:\n",
                "    \"\"\"Convert raw dataset into MDS format\n",
                "\n",
                "    Args:\n",
                "        args (Iterator[Tuple[str, int, int]]): All arguments, packed into a tuple because\n",
                "            process pools only pass one argument.\n",
                "\n",
                "    Yields:\n",
                "        Dict: A sample\n",
                "    \"\"\"\n",
                "    sub_out_root, start_sample_idx, end_sample_idx = args\n",
                "\n",
                "    def get_data(start: int, end: int):\n",
                "        for i in range(start, end + 1):\n",
                "            yield {'number': i}\n",
                "    \n",
                "    columns = {'number': 'int'}\n",
                "\n",
                "    with MDSWriter(out=sub_out_root,\n",
                "                   columns=columns) as out:\n",
                "        for sample in get_data(start_sample_idx, end_sample_idx):\n",
                "            out.write(sample)"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {
                "id": "by7aPVIDM1mG"
            },
            "source": [
                "We partition the raw dataset into 4 sub-groups, and each process takes a converts a sub-group into MDS format. The resulting shards are stored in the respective sub-directories."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {
                "id": "cZiFudkCM2dF"
            },
            "outputs": [],
            "source": [
                "# clean up root directory\n",
                "%rm -rf $out_root\n",
                "\n",
                "arg_tuples = each_task(out_root, groups=num_groups)\n",
                "    \n",
                "# Process group of data in parallel into directories of shards.\n",
                "with Pool(initializer=init_worker, processes=num_process) as pool:\n",
                "    for count in pool.imap(convert_to_mds, arg_tuples):\n",
                "        pass\n",
                "print('Finished')"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {
                "id": "dX60JcG9M_aT"
            },
            "source": [
                "Once dataset has been converted to an MDS format, let's look at the directory structure. You will find 4 sub-directories, each containing an `index.json` file and shard files."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {
                "id": "_Ty7i0WwNCEa"
            },
            "outputs": [],
            "source": [
                "%ll $out_root"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {
                "id": "Ccye5lJnNFSm"
            },
            "source": [
                "### Merge meta data\n",
                "\n",
                "The last step of the conversion process is to merge all the `index.json` files of the sub-directories. The content of the shard files will remain the same. By calling the merge_index utility function, information for all the shards will be written to a new `index.json` file placed in the `out` directory."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {
                "id": "lz3Aem5WNQ2R"
            },
            "outputs": [],
            "source": [
                "from streaming.base.util import merge_index\n",
                "merge_index(out_root, keep_local=True)"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {
                "id": "NByc1ZGINcXe"
            },
            "source": [
                "Let's checkout the root directory, where you can see one `index.json` file along with subdirectories that contain shard files."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {
                "id": "hrD9yGgPNclL"
            },
            "outputs": [],
            "source": [
                "%ll $out_root"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {
                "id": "XZ1-fIIMNVO8"
            },
            "source": [
                "### Load MDS dataset\n",
                "\n",
                "Read the sample using `StreamingDataset`. Here, we just print the sample IDs."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {
                "id": "FlRod4cnNVlM"
            },
            "outputs": [],
            "source": [
                "# read the sample\n",
                "dataset = StreamingDataset(local=out_root,\n",
                "                           remote=None,\n",
                "                           shuffle=False)\n",
                "for ix, sample in enumerate(dataset):\n",
                "    print(sample['number'])"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {
                "id": "AKRdDahKWyki"
            },
            "source": [
                "### Cleanup"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {
                "id": "_esj_q2MNiny"
            },
            "outputs": [],
            "source": [
                "%rm -rf $out_root"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {
                "id": "lf00woabVLWG"
            },
            "source": [
                "\n",
                "## What next?\n",
                "\n",
                "You've now seen an in-depth tutorial on converting a dataset into MDS format using multiple process. If you are interested in some real-world examples, then, check out the [WebVid](https://github.com/mosaicml/streaming/blob/main/streaming/multimodal/convert/webvid/crawl_webvid.py) and [Pile](https://github.com/mosaicml/streaming/blob/main/streaming/text/convert/pile.py) dataset conversion scripts which convert datasets into MDS format via multiprocessing."
            ]
        }
    ],
    "metadata": {
        "colab": {
            "provenance": []
        },
        "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.10.12"
        }
    },
    "nbformat": 4,
    "nbformat_minor": 4
}