Repository: voxel51/fiftyone

Issue:
[BUG] Large RAM usage when downloading Open Images dataset
### System information

-   **OS Platform and Distribution (windows 10):
-   **FiftyOne installed from pip :
-   **FiftyOne version : v0.13.2;
-   **Python version: 3.9:

### Commands to reproduce

# shell commands  
_________________________________________
set HTTP_PROXY=http://127.0.0.1:1081
set HTTPS_PROXY=http://127.0.0.1:1081
python .\PycharmProjects\pythonProject\main.py

#code in main.py
_______________________________________________
import fiftyone.zoo as foz
import os

#dataset = foz.load_zoo_dataset("quickstart")
dataset1 = foz.load_zoo_dataset(
              "open-images-v6",
              split="train",
              label_types=["detections"],
              classes=["Bathtub"],
              max_samples=1000,
              seed=51,
              shuffle=True,
              dataset_name="open-images-Bathtub"
          )

_________________________________________
### Describe the problem
###This is happend after the CSV was parsed and it started to download the images. Crazy increase of memory,as  shown in image below:

![image](https://user-images.githubusercontent.com/15262893/132808119-c4c9ed7d-e997-4a63-9999-e3a9836e05db.png)










Discussion:
Hi @HelloCentaurus 👋 

By default, downloading the images for the Open Images dataset will use `multiprocessing.cpu_count()` processes to do the job ([here's the code](https://github.com/voxel51/fiftyone/blob/9e492a276edd48f0c0bfed7344dfaa2124973bea/fiftyone/utils/openimages.py#L1580-L1646)). Looks like you're working on a pretty beefy machine with lots of resources. Maybe such a large number of processes is becoming a bit too resource greedy for your liking.

You can pass an optional `num_workers` flag to `load_zoo_dataset()` to specify how many processes can be used. Can you try setting, say `num_workers=4` to see if that resolves it?

Note that you'll have to delete your existing download in order for the images to be actually downloaded again:

```py
import fiftyone.zoo as foz

foz.delete_zoo_dataset("open-images-v6", split="train")
```

One comment is that train labels CSV file is huge (~3 GB if I remember correctly), and we load it into memory order to find some suitable images to download. I wonder if the multiprocessing could be causing copies of that to be shipped to each process? It is not explicitly done, might be a case of lazy garbage collection?

We really should have use multithreading as opposed to multiprocessing for the downloads. It would be interesting to see if making that change in the code I linked above would prevent the memory usage.
I see `C:` and `D:` drives in your screenshot, so I believe you are running Windows. That reminds me that other users have had issues with multiprocessing code on Windows (eg https://github.com/voxel51/fiftyone/issues/1119), for which the recommendation was to use `num_workers=1`, which uses no multiprocessing at all.

Basically I think multiprocessing goes awry on Windows (in your case RAM hungry, in other cases flat out errors) and the best option seems to be to not use it.

Resolving patch:
diff --git a/fiftyone/utils/aws.py b/fiftyone/utils/aws.py
new file mode 100644
--- /dev/null
+++ b/fiftyone/utils/aws.py
@@ -0,0 +1,132 @@
+"""
+Utilities for working with
+`Amazon Web Services <https://aws.amazon.com/>`
+S3 storage buckets.
+
+| Copyright 2017-2021, Voxel51, Inc.
+| `voxel51.com <https://voxel51.com/>`_
+|
+"""
+import logging
+import multiprocessing
+import multiprocessing.dummy
+import os
+from urllib.parse import urlparse
+
+import boto3
+import botocore
+
+import fiftyone.core.utils as fou
+
+import eta.core.utils as etau
+
+
+logger = logging.getLogger(__name__)
+
+
+def download_public_s3_files(
+    urls, download_dir=None, num_workers=None, overwrite=True,
+):
+    """Download files from a public AWS S3 bucket using unsigned urls.
+
+    The `url` argument either accepts:
+
+        * A list of paths to objects in the s3 bucket::
+
+            urls = ["s3://bucket_name/dir1/file1.ext", ...]
+
+          When `urls` is a list, then the `download_dir` argument is required
+          and all objects will be downloaded into that directory
+
+        * A dictionary mapping the paths of objects to files on disk
+          to store each object::
+
+            urls = {
+                "s3://bucket_name/dir1/file1.ext": "/path/to/local/file1.ext",
+                ...
+            }
+
+    Args:
+        urls: either a list of urls to objects in an s3 bucket, or a dict
+            mapping these urls to locations on disk. If `urls` is a list, then
+            the `download_dir` argument is required
+        download_dir (None): the directory to store all downloaded objects.
+            This is only used if `urls` is a list
+        num_workers (None): the number of processes to use when downloading
+            files. By default, ``multiprocessing.cpu_count()`` is used
+        overwrite (True): whether to overwrite existing files
+    """
+    if isinstance(urls, list):
+        if download_dir is None:
+            raise ValueError(
+                "When `urls` is a list, `download_dir` is required but was found to be `None`."
+            )
+        urls = {url: None for url in urls}
+
+    if download_dir:
+        etau.ensure_dir(download_dir)
+
+    if num_workers is None or num_workers < 1:
+        num_workers = multiprocessing.cpu_count()
+
+    s3_client = boto3.client(
+        "s3",
+        config=botocore.config.Config(
+            signature_version=botocore.UNSIGNED,
+            max_pool_connections=max(10, num_workers),
+        ),
+    )
+
+    inputs = _build_inputs(
+        urls, s3_client, download_dir=download_dir, overwrite=overwrite
+    )
+
+    if not inputs:
+        return
+
+    if num_workers == 1:
+        _single_thread_download(inputs)
+    else:
+        _multi_thread_download(inputs, num_workers)
+
+
+def _build_inputs(urls, s3_client, download_dir=None, overwrite=True):
+    inputs = []
+    for url, filepath in urls.items():
+        bucket_name, object_path = _parse_url(url)
+        if filepath is None:
+            filepath = os.path.join(download_dir, object_path)
+        if not os.path.isfile(filepath):
+            inputs.append((bucket_name, object_path, filepath, s3_client))
+        else:
+            if overwrite:
+                os.remove(filepath)
+            else:
+                logger.warning(
+                    "File `%s` already exists, skipping..." % filepath
+                )
+
+    return inputs
+
+
+def _parse_url(url):
+    result = urlparse(url, allow_fragments=False)
+    return result.netloc, result.path.lstrip("/")
+
+
+def _single_thread_download(inputs):
+    with fou.ProgressBar(total=len(inputs), iters_str="files") as pb:
+        for bucket_name, obj, filepath, s3_client in pb(inputs):
+            s3_client.download_file(bucket_name, obj, filepath)
+
+
+def _multi_thread_download(inputs, num_workers):
+    with fou.ProgressBar(total=len(inputs), iters_str="files") as pb:
+        with multiprocessing.dummy.Pool(num_workers) as pool:
+            for _ in pool.imap_unordered(_do_s3_download, inputs):
+                pb.update()
+
+
+def _do_s3_download(args):
+    bucket_name, obj, filepath, s3_client = args
+    s3_client.download_file(bucket_name, obj, filepath)
diff --git a/fiftyone/utils/coco.py b/fiftyone/utils/coco.py
--- a/fiftyone/utils/coco.py
+++ b/fiftyone/utils/coco.py
@@ -12,6 +12,7 @@
 from itertools import groupby
 import logging
 import multiprocessing
+import multiprocessing.dummy
 import os
 import random
 import shutil
@@ -1491,7 +1492,7 @@ def download_coco_dataset_split(
     if split not in _IMAGE_DOWNLOAD_LINKS[year]:
         raise ValueError(
             "Unsupported split '%s'; supported values are %s"
-            % (year, tuple(_IMAGE_DOWNLOAD_LINKS[year].keys()))
+            % (split, tuple(_IMAGE_DOWNLOAD_LINKS[year].keys()))
         )
 
     if classes is not None and split == "test":
@@ -1841,7 +1842,7 @@ def _download_images(images_dir, image_ids, images, num_workers):
                 _do_download(task)
     else:
         with fou.ProgressBar(total=len(tasks), iters_str="images") as pb:
-            with multiprocessing.Pool(num_workers) as pool:
+            with multiprocessing.dummy.Pool(num_workers) as pool:
                 for _ in pool.imap_unordered(_do_download, tasks):
                     pb.update()
 
diff --git a/fiftyone/utils/openimages.py b/fiftyone/utils/openimages.py
--- a/fiftyone/utils/openimages.py
+++ b/fiftyone/utils/openimages.py
@@ -10,13 +10,10 @@
 from collections import defaultdict
 import csv
 import logging
-import multiprocessing
 import os
 import random
 import warnings
 
-import boto3
-import botocore
 import pandas as pd
 
 import eta.core.image as etai
@@ -27,7 +24,7 @@
 import fiftyone as fo
 import fiftyone.core.fields as fof
 import fiftyone.core.labels as fol
-import fiftyone.core.utils as fou
+import fiftyone.utils.aws as foua
 import fiftyone.utils.data as foud
 
 
@@ -1581,23 +1578,21 @@ def _download_masks_if_necessary(image_ids, dataset_dir, split, download=True):
 def _download_images_if_necessary(
     image_ids, split, dataset_dir, num_workers=None, download=True
 ):
-    if num_workers is None or num_workers < 1:
-        num_workers = multiprocessing.cpu_count()
-
     data_dir = os.path.join(dataset_dir, "data")
     etau.ensure_dir(data_dir)
 
-    inputs = []
+    urls = {}
     num_existing = 0
     for image_id in image_ids:
         filepath = os.path.join(data_dir, image_id + ".jpg")
         obj = split + "/" + image_id + ".jpg"  # AWS path, always use "/"
         if not os.path.isfile(filepath):
-            inputs.append((filepath, obj))
+            url = "s3://%s/%s" % (_BUCKET_NAME, obj)
+            urls[url] = filepath
         else:
             num_existing += 1
 
-    num_images = len(inputs)
+    num_images = len(urls)
 
     if num_images == 0:
         if download:
@@ -1617,36 +1612,11 @@ def _download_images_if_necessary(
     else:
         logger.info("Downloading %d images", num_images)
 
-    if num_workers == 1:
-        s3_client = boto3.client(
-            "s3",
-            config=botocore.config.Config(signature_version=botocore.UNSIGNED),
-        )
-        with fou.ProgressBar(iters_str="images") as pb:
-            for filepath, obj in pb(inputs):
-                s3_client.download_file(_BUCKET_NAME, obj, filepath)
-    else:
-        with fou.ProgressBar(total=num_images, iters_str="images") as pb:
-            with multiprocessing.Pool(num_workers, _initialize_worker) as pool:
-                for _ in pool.imap_unordered(_do_s3_download, inputs):
-                    pb.update()
+    foua.download_public_s3_files(urls, num_workers=num_workers)
 
     return num_images
 
 
-def _initialize_worker():
-    global s3_client
-    s3_client = boto3.client(
-        "s3",
-        config=botocore.config.Config(signature_version=botocore.UNSIGNED),
-    )
-
-
-def _do_s3_download(args):
-    filepath, obj = args
-    s3_client.download_file(_BUCKET_NAME, obj, filepath)
-
-
 def _verify_version(version):
     if version not in _SUPPORTED_VERSIONS:
         raise ValueError(


Test patch: