{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Returning a generator in joblib.Parallel\n\nThis example illustrates memory optimization enabled by using\n:class:`joblib.Parallel` to get a generator on the outputs of parallel jobs.\nWe first create tasks that return results with large memory footprints.\nIf we call :class:`~joblib.Parallel` for several of these tasks directly, we\nobserve a high memory usage, as all the results are held in RAM before being\nprocessed\n\nUsing ``return_as='generator'`` allows to progressively consume the outputs\nas they arrive and keeps the memory at an acceptable level.\n\nIn this case, the output of the `Parallel` call is a generator that yields the\nresults in the order the tasks have been submitted with. If the order of the\ntasks does not matter (for instance if they are consumed by a commutative\naggregation function), then using ``return_as='generator_unordered'`` can be\neven more efficient.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## ``MemoryMonitor`` helper\n\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The following class is an helper to monitor the memory of the process and its\nchildren in another thread, so we can display it afterward.\n\nWe will use ``psutil`` to monitor the memory usage in the code. Make sure it\nis installed with ``pip install psutil`` for this example.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import time\nfrom psutil import Process\nfrom threading import Thread\n\n\nclass MemoryMonitor(Thread):\n    \"\"\"Monitor the memory usage in MB in a separate thread.\n\n    Note that this class is good enough to highlight the memory profile of\n    Parallel in this example, but is not a general purpose profiler fit for\n    all cases.\n    \"\"\"\n    def __init__(self):\n        super().__init__()\n        self.stop = False\n        self.memory_buffer = []\n        self.start()\n\n    def get_memory(self):\n        \"Get memory of a process and its children.\"\n        p = Process()\n        memory = p.memory_info().rss\n        for c in p.children():\n            memory += c.memory_info().rss\n        return memory\n\n    def run(self):\n        memory_start = self.get_memory()\n        while not self.stop:\n            self.memory_buffer.append(self.get_memory() - memory_start)\n            time.sleep(0.2)\n\n    def join(self):\n        self.stop = True\n        super().join()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Save memory by consuming the outputs of the tasks as fast as possible\n\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We create a task whose output takes about 15MB of RAM.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import numpy as np\n\n\ndef return_big_object(i):\n    time.sleep(.1)\n    return i * np.ones((10000, 200), dtype=np.float64)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We create a reduce step. The input will be a generator on big objects\ngenerated in parallel by several instances of ``return_big_object``.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "def accumulator_sum(generator):\n    result = 0\n    for value in generator:\n        result += value\n        print(\".\", end=\"\", flush=True)\n    print(\"\")\n    return result"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We process many of the tasks in parallel. If ``return_as=\"list\"`` (default),\nwe should expect a usage of more than 2GB in RAM. Indeed, all the results\nare computed and stored in ``res`` before being processed by\n`accumulator_sum` and collected by the gc.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from joblib import Parallel, delayed\n\nmonitor = MemoryMonitor()\nprint(\"Running tasks with return_as='list'...\")\nres = Parallel(n_jobs=2, return_as=\"list\")(\n    delayed(return_big_object)(i) for i in range(150)\n)\nprint(\"Accumulate results:\", end='')\nres = accumulator_sum(res)\nprint('All tasks completed and reduced successfully.')\n\n# Report memory usage\ndel res  # we clean the result to avoid memory border effects\nmonitor.join()\npeak = max(monitor.memory_buffer) / 1e9\nprint(f\"Peak memory usage: {peak:.2f}GB\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "If we use ``return_as=\"generator\"``, ``res`` is simply a generator on the\nresults that are ready. Here we consume the results as soon as they arrive\nwith the ``accumulator_sum`` and once they have been used, they are collected\nby the gc. The memory footprint is thus reduced, typically around 300MB.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "monitor_gen = MemoryMonitor()\nprint(\"Create result generator with return_as='generator'...\")\nres = Parallel(n_jobs=2, return_as=\"generator\")(\n    delayed(return_big_object)(i) for i in range(150)\n)\nprint(\"Accumulate results:\", end='')\nres = accumulator_sum(res)\nprint('All tasks completed and reduced successfully.')\n\n# Report memory usage\ndel res  # we clean the result to avoid memory border effects\nmonitor_gen.join()\npeak = max(monitor_gen.memory_buffer) / 1e6\nprint(f\"Peak memory usage: {peak:.2f}MB\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We can then report the memory usage accross time of the two runs using the\nMemoryMonitor.\n\nIn the first case, as the results accumulate in ``res``, the memory grows\nlinearly and it is freed once the ``accumulator_sum`` function finishes.\n\nIn the second case, the results are processed by the accumulator as soon as\nthey arrive, and the memory does not need to be able to contain all\nthe results.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import matplotlib.pyplot as plt\nplt.figure(0)\nplt.semilogy(\n    np.maximum.accumulate(monitor.memory_buffer),\n    label='return_as=\"list\"'\n)\nplt.semilogy(\n    np.maximum.accumulate(monitor_gen.memory_buffer),\n    label='return_as=\"generator\"'\n)\nplt.xlabel(\"Time\")\nplt.xticks([], [])\nplt.ylabel(\"Memory usage\")\nplt.yticks([1e7, 1e8, 1e9], ['10MB', '100MB', '1GB'])\nplt.legend()\nplt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "It is important to note that with ``return_as=\"generator\"``, the results are\nstill accumulated in RAM after computation. But as we asynchronously process\nthem, they can be freed sooner. However, if the generator is not consumed\nthe memory still grows linearly.\n\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Further memory efficiency for commutative aggregation\n\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "There is still room for improving the relief on memory allocation we get\nusing ``return_as=\"generator\"``. Indeed, notice how the generator of the\nprevious example respects the order the tasks have been submitted with. This\nbehavior can cause a build up in memory of results waiting to be consumed,\nin case some tasks finished before other tasks despite being submitted\nlater. The corresponding results will be kept in memory until the slower\ntasks submitted earlier are done and have been iterated over.\n\nIn case the downstream consumer of the results is reliant on the assumption\nthat the results are yielded in the same order that the tasks were submitted,\nit can't be helped. But in our example, since the `+` operator is\ncommutative, the function ``accumulator_sum`` does not need the generator to\nreturn the results with any particular order. In this case it's safe to use\nthe option ``return_as=\"generator_unordered\"``, so that the results are\nreturned as soon as a task is completed, ignoring the order of task\nsubmission.\n\nBeware that the downstream consumer of the results must not expect them be\nreturned with any deterministic or predictable order at all, since the\nprogress of the tasks can depend on the availability of the workers, which\ncan be affected by external events, such as system load, implementation\ndetails in the backend, etc.\n\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "To better highlight improvements in memory usage when using the parameter\n``return_as=\"generator_unordered\"``, let's explcitly add delay in some of\nthe submitted tasks.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "def return_big_object_delayed(i):\n    if (i + 20) % 60:\n        time.sleep(0.1)\n    else:\n        time.sleep(5)\n    return i * np.ones((10000, 200), dtype=np.float64)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's check memory usage when using ``return_as=\"generator\"``...\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "monitor_delayed_gen = MemoryMonitor()\nprint(\"Create result generator on delayed tasks with return_as='generator'...\")\nres = Parallel(n_jobs=2, return_as=\"generator\")(\n    delayed(return_big_object_delayed)(i) for i in range(150)\n)\nprint(\"Accumulate results:\", end='')\nres = accumulator_sum(res)\nprint('All tasks completed and reduced successfully.')\n\n# Report memory usage\ndel res  # we clean the result to avoid memory border effects\nmonitor_delayed_gen.join()\npeak = max(monitor_delayed_gen.memory_buffer) / 1e6\nprint(f\"Peak memory usage: {peak:.2f}MB\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "If we use ``return_as=\"generator_unordered\"``, ``res`` will not enforce any\norder when returning the results, and will simply enable iterating on the\nresults as soon as it's available. The peak memory usage is now controlled\nto an even lower level, since that results can be consumed immediately\nrather than being delayed by the compute of slower tasks that have been\nsubmitted earlier.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "monitor_delayed_gen_unordered = MemoryMonitor()\nprint(\n  \"Create result generator on delayed tasks with \"\n  \"return_as='generator_unordered'...\"\n)\nres = Parallel(n_jobs=2, return_as=\"generator_unordered\")(\n    delayed(return_big_object_delayed)(i) for i in range(150)\n)\nprint(\"Accumulate results:\", end='')\nres = accumulator_sum(res)\nprint('All tasks completed and reduced successfully.')\n\n# Report memory usage\ndel res  # we clean the result to avoid memory border effects\nmonitor_delayed_gen_unordered.join()\npeak = max(monitor_delayed_gen_unordered.memory_buffer) / 1e6\nprint(f\"Peak memory usage: {peak:.2f}MB\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Notice how the plot for ``'return_as=\"generator'`` now shows a high memory\nusage plateau when slow jobs cause a congestion of intermediate results\nwaiting in RAM before in-order aggregation. This high memory usage is never\nobserved when using ``'return_as=\"generator_unordered\"``.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "plt.figure(1)\nplt.semilogy(\n    np.maximum.accumulate(monitor_delayed_gen.memory_buffer),\n    label='return_as=\"generator\"'\n)\nplt.semilogy(\n    np.maximum.accumulate(monitor_delayed_gen_unordered.memory_buffer),\n    label='return_as=\"generator_unordered\"'\n)\nplt.xlabel(\"Time\")\nplt.xticks([], [])\nplt.ylabel(\"Memory usage\")\nplt.yticks([1e7, 1e8, 1e9], ['10MB', '100MB', '1GB'])\nplt.legend()\nplt.show()"
      ]
    }
  ],
  "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.10.9"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}