From 3e393bf362613c543ad58aea899a24e6141f01be Mon Sep 17 00:00:00 2001 From: Andras Schmelczer Date: Wed, 29 Jun 2022 09:14:35 +0200 Subject: [PATCH] Improve parallel map --- docs/map.ipynb | 41 ++++++++++++------- .../utilities/parallel_map/parallel_map.py | 6 +-- .../parallel_map_configuration.py | 2 +- 3 files changed, 30 insertions(+), 19 deletions(-) diff --git a/docs/map.ipynb b/docs/map.ipynb index 7b493b2..9f07b8c 100644 --- a/docs/map.ipynb +++ b/docs/map.ipynb @@ -4,7 +4,16 @@ "cell_type": "code", "execution_count": 1, "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/data/projects/great_ai/.env/lib/python3.10/site-packages/tqdm/auto.py:22: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n", + " from .autonotebook import tqdm as notebook_tqdm\n" + ] + } + ], "source": [ "from great_ai.utilities import parallel_map" ] @@ -18,17 +27,19 @@ "name": "stderr", "output_type": "stream", "text": [ - "\u001b[38;5;39m2022-06-28 18:10:31,043 | INFO | Parallel map: configured ✅\u001b[0m\n", - "\u001b[38;5;39m2022-06-28 18:10:31,044 | INFO | ⚙️ concurrency: 12\u001b[0m\n", - "\u001b[38;5;39m2022-06-28 18:10:31,045 | INFO | ⚙️ chunk length: 8334\u001b[0m\n", - "\u001b[38;5;39m2022-06-28 18:10:31,045 | INFO | ⚙️ chunk count: 120\u001b[0m\n", - "\u001b[38;5;39m2022-06-28 18:10:31,046 | INFO | ⚙️ function size: 0 kB\u001b[0m\n", - "Parallel map: 100%|██████████| 1000000/1000000 [00:05<00:00, 181294.86it/s]\n" + "\u001b[38;5;39m2022-06-29 09:11:18,070 | INFO | Parallel map: configured ✅\u001b[0m\n", + "\u001b[38;5;39m2022-06-29 09:11:18,071 | INFO | ⚙️ concurrency: 200\u001b[0m\n", + "\u001b[38;5;39m2022-06-29 09:11:18,071 | INFO | ⚙️ chunk length: 1\u001b[0m\n", + "\u001b[38;5;39m2022-06-29 09:11:18,071 | INFO | ⚙️ chunk count: 1000\u001b[0m\n", + "\u001b[38;5;39m2022-06-29 09:11:18,072 | INFO | ⚙️ function size: 0 kB\u001b[0m\n", + "Parallel map: 100%|██████████| 1000/1000 [00:00<00:00, 1345.90it/s]\n" ] } ], "source": [ - "parallel_map(lambda x: x**3, range(int(1e6)))\n", + "from time import sleep\n", + "\n", + "parallel_map(lambda x: sleep(0.1), range(int(1000)), concurrency=200)\n", "None" ] }, @@ -41,12 +52,12 @@ "name": "stderr", "output_type": "stream", "text": [ - "\u001b[38;5;39m2022-06-28 18:10:36,757 | INFO | Parallel map: configured ✅\u001b[0m\n", - "\u001b[38;5;39m2022-06-28 18:10:36,758 | INFO | ⚙️ concurrency: 12\u001b[0m\n", - "\u001b[38;5;39m2022-06-28 18:10:36,759 | INFO | ⚙️ chunk length: 1\u001b[0m\n", - "\u001b[38;5;39m2022-06-28 18:10:36,759 | INFO | ⚙️ chunk count: unknown\u001b[0m\n", - "\u001b[38;5;39m2022-06-28 18:10:36,760 | INFO | ⚙️ function size: 0 kB\u001b[0m\n", - "Parallel map: 10it [00:21, 2.11s/it]\n" + "\u001b[38;5;39m2022-06-29 09:11:20,399 | INFO | Parallel map: configured ✅\u001b[0m\n", + "\u001b[38;5;39m2022-06-29 09:11:20,401 | INFO | ⚙️ concurrency: 12\u001b[0m\n", + "\u001b[38;5;39m2022-06-29 09:11:20,401 | INFO | ⚙️ chunk length: 1\u001b[0m\n", + "\u001b[38;5;39m2022-06-29 09:11:20,402 | INFO | ⚙️ chunk count: unknown\u001b[0m\n", + "\u001b[38;5;39m2022-06-29 09:11:20,403 | INFO | ⚙️ function size: 0 kB\u001b[0m\n", + "Parallel map: 10it [00:02, 4.93it/s]\n" ] }, { @@ -67,7 +78,7 @@ "def my_generator():\n", " for i in range(10):\n", " yield i\n", - " sleep(2)\n", + " sleep(0.2)\n", "\n", "\n", "parallel_map(lambda x: x**3, my_generator(), chunk_length=1)" diff --git a/src/great_ai/utilities/parallel_map/parallel_map.py b/src/great_ai/utilities/parallel_map/parallel_map.py index 967050f..a00a176 100644 --- a/src/great_ai/utilities/parallel_map/parallel_map.py +++ b/src/great_ai/utilities/parallel_map/parallel_map.py @@ -68,12 +68,13 @@ def parallel_map( desc="Parallel map", disable=disable_progress_bar, total=config.input_length, - miniters=1, ) ] - ctx = mp.get_context("spawn") + start_methods = mp.get_all_start_methods() + ctx = mp.get_context("fork") if "fork" in start_methods else mp.get_context("spawn") ctx.freeze_support() + input_queue = ctx.Queue(0 if config.chunk_count is None else config.chunk_count) output_queue = ctx.Queue(0 if config.chunk_count is None else config.chunk_count) should_stop = ctx.Event() @@ -99,7 +100,6 @@ def parallel_map( desc="Parallel map", disable=disable_progress_bar, total=config.input_length, - miniters=1, ) chunks = iter(chunk(enumerate(input_values), chunk_length=config.chunk_length)) diff --git a/src/great_ai/utilities/parallel_map/parallel_map_configuration.py b/src/great_ai/utilities/parallel_map/parallel_map_configuration.py index 8890173..e8fd14c 100644 --- a/src/great_ai/utilities/parallel_map/parallel_map_configuration.py +++ b/src/great_ai/utilities/parallel_map/parallel_map_configuration.py @@ -14,7 +14,7 @@ class ParallelMapConfiguration(BaseModel): input_length: Optional[int] serialized_map_function: bytes - def pretty_print(self, prefix=" ⚙️"): + def pretty_print(self, prefix=" ⚙️ "): logger.info(f"{prefix} concurrency: {self.concurrency}") logger.info(f"{prefix} chunk length: {self.chunk_length}") logger.info(