{
  "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 the ``return_generator=True`` option allows to progressively consumes\nthe outputs as they arrive and keeps the memory at an acceptable level.\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_generator=False`\n(default), we should expect a usage of more than 2GB in RAM. Indeed, all the\nresults are 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_generator=False...\")\nres = Parallel(n_jobs=2, return_generator=False)(\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_generator=True``, ``res`` is simply a generator with 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_generator=True...\")\nres = Parallel(n_jobs=2, return_generator=True)(\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.semilogy(\n    np.maximum.accumulate(monitor.memory_buffer),\n    label='return_generator=False'\n)\nplt.semilogy(\n    np.maximum.accumulate(monitor_gen.memory_buffer),\n    label='return_generator=True'\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_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"
      ]
    }
  ],
  "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
}