adding aistore datapipe (#545)

Summary:
Fixes #{[517](#517)}

### Changes
- Added `aisio.py` (Iterable Datapipe for AIStore backends)
- Added unit tests in ~~`test/test_local_io.py`~~ `test/test_aistore.py`
- Added GitHub action for running AIStore in ~~`CI.yml`~~ `.github/aistore_ci.yml` workflow

### Questions to maintainers
- We are unsure about the documentation generated on PyTorch to refer it in `README.md`, so I have tentatively added a URL similar to s3io functions. (see ```torchdata/datapipes/iter/load/README.md```)

Signed-off-by: Abhishek Gaikwad <[email protected]>

Pull Request resolved: #545

Reviewed By: VitalyFedyunin

Differential Revision: D37620194

Pulled By: msaroufim

fbshipit-source-id: 9df099586dd39d47f8fdf2b760b17503f8a9822d

Author: gaikwadabhishek

.github/workflows/aistore_ci.yml

@@ -0,0 +1,59 @@
+name: Run AIStore Datapipe Test
+on:
+  push:
+    branches:
+      - main
+      - release/*
+    tags:
+  pull_request:
+    types: [opened, synchronize, reopened, labeled]
+    branches:
+      - main
+      # For PR created by ghstack
+      - gh/*/*/base
+      - release/*
+
+jobs:
+  test:
+    if: ${{ github.repository_owner == 'pytorch' }}
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        os:
+          - macos-latest
+          - ubuntu-latest
+        python-version:
+          - 3.7
+          - 3.8
+          - 3.9
+    steps:
+      - name: Get PyTorch Channel
+        shell: bash
+        run: |
+          if [[ "${{ github.base_ref }}" == release/* ]] || [[ "${{ github.ref }}" == refs/heads/release/* ]] || [[ "${{ github.ref }}" == refs/tags/v* ]]; then
+            PT_CHANNEL="https://download.pytorch.org/whl/test/cpu/torch_test.html"
+          else
+            PT_CHANNEL="https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html"
+          fi
+          echo "::set-output name=value::$PT_CHANNEL"
+        id: pytorch_channel
+      - name: Setup Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v2
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Check out source repository
+        uses: actions/checkout@v2
+      - name: Install dependencies
+        run: |
+          pip3 install -r requirements.txt
+          pip3 install --pre torch -f "${{ steps.pytorch_channel.outputs.value }}"
+      - name: Run AIStore local deployment
+        uses: NVIDIA/aistore@master
+      - name: Build TorchData
+        run: |
+          python setup.py install
+      - name: Install test requirements
+        run: pip3 install -r test/requirements_aistore.txt
+      - name: Run AIStore DataPipe tests with pytest
+        run: pytest --no-header -v test/test_aistore.py

.github/workflows/ci.yml

@@ -80,18 +80,18 @@ jobs:
         if: ${{ ! contains(github.event.pull_request.labels.*.name, 'ciflow/slow') }}
         run:
           pytest --no-header -v test --ignore=test/test_period.py --ignore=test/test_text_examples.py
-          --ignore=test/test_audio_examples.py
+          --ignore=test/test_audio_examples.py --ignore=test/test_aistore.py
       - name: Run DataPipes tests with pytest (including slow tests)
         if: ${{ contains(github.event.pull_request.labels.*.name, 'ciflow/slow') }}
         run:
           pytest --no-header -v test --ignore=test/test_period.py --ignore=test/test_text_examples.py
-          --ignore=test/test_audio_examples.py
+          --ignore=test/test_audio_examples.py --ignore=test/test_aistore.py
         env:
           PYTORCH_TEST_WITH_SLOW: 1
       - name: Run DataPipes period tests with pytest
         if: ${{ contains(github.event.pull_request.labels.*.name, 'ciflow/period') }}
         run:
           pytest --no-header -v test/test_period.py --ignore=test/test_text_examples.py
-          --ignore=test/test_audio_examples.py
+          --ignore=test/test_audio_examples.py --ignore=test/test_aistore.py
         env:
           PYTORCH_TEST_WITH_SLOW: 1

docs/source/torchdata.datapipes.iter.rst

@@ -130,6 +130,8 @@ saving files, and listing the files in directories).
     :toctree: generated/
     :template: datapipe.rst
 
+    AISFileLister
+    AISFileLoader
     FSSpecFileLister
     FSSpecFileOpener
     FSSpecSaver

mypy.ini

@@ -47,3 +47,6 @@ ignore_missing_imports = True
 
 [mypy-graphviz.*]
 ignore_missing_imports = True
+
+[mypy-aistore.*]
+ignore_missing_imports = True

test/requirements_aistore.txt

@@ -0,0 +1,2 @@
+pytest
+aistore > 0.9.1

test/test_aistore.py

@@ -0,0 +1,101 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import random
+import string
+import tempfile
+import unittest
+
+from torchdata.datapipes.iter import AISFileLister, AISFileLoader
+
+try:
+    from aistore.client.api import Client
+    from aistore.client.errors import AISError, ErrBckNotFound
+
+    AIS_CLUSTER_ENDPT = "http://localhost:8080"
+
+    HAS_AIS = Client(AIS_CLUSTER_ENDPT).is_aistore_running()
+except (ImportError, ConnectionError):
+    HAS_AIS = False
+skipIfNoAIS = unittest.skipIf(not HAS_AIS, "AIS not running or library not installed")
+
+
+@skipIfNoAIS
+class TestDataPipeLocalIO(unittest.TestCase):
+    def setUp(self):
+        # initialize client and create new bucket
+        self.client = Client(AIS_CLUSTER_ENDPT)
+        letters = string.ascii_lowercase
+        self.bck_name = "".join(random.choice(letters) for _ in range(10))
+        self.client.create_bucket(self.bck_name)
+        # create temp files
+        num_objs = 10
+
+        # create 10 objects in the `/temp` dir
+        for i in range(num_objs):
+            object_body = "test string" * random.randrange(1, 10)
+            content = object_body.encode("utf-8")
+            obj_name = f"temp/obj{ i }"
+            with tempfile.NamedTemporaryFile() as file:
+                file.write(content)
+                file.flush()
+                self.client.put_object(self.bck_name, obj_name, file.name)
+
+        # create 10 objects in the `/`dir
+        for i in range(num_objs):
+            object_body = "test string" * random.randrange(1, 10)
+            content = object_body.encode("utf-8")
+            obj_name = f"obj{ i }"
+            with tempfile.NamedTemporaryFile() as file:
+                file.write(content)
+                file.flush()
+                self.client.put_object(self.bck_name, obj_name, file.name)
+
+    def tearDown(self):
+        # Try to destroy bucket and its items
+        try:
+            self.client.destroy_bucket(self.bck_name)
+        except ErrBckNotFound:
+            pass
+
+    def test_ais_io_iterdatapipe(self):
+
+        prefixes = [
+            ["ais://" + self.bck_name],
+            ["ais://" + self.bck_name + "/"],
+            ["ais://" + self.bck_name + "/temp/", "ais://" + self.bck_name + "/obj"],
+        ]
+
+        # check if the created files exist
+        for prefix in prefixes:
+            urls = AISFileLister(url=AIS_CLUSTER_ENDPT, source_datapipe=prefix)
+            ais_loader = AISFileLoader(url=AIS_CLUSTER_ENDPT, source_datapipe=urls)
+            with self.assertRaises(TypeError):
+                len(urls)
+            self.assertEqual(len(list(urls)), 20)
+            self.assertEqual(sum(1 for _ in ais_loader), 20)
+
+        # check for incorrect prefixes
+        prefixes = ["ais://asdasd"]
+
+        # AISFileLister: Bucket not found
+        try:
+            list(AISFileLister(url=AIS_CLUSTER_ENDPT, source_datapipe=prefixes))
+        except ErrBckNotFound as err:
+            self.assertEqual(err.status_code, 404)
+
+        # AISFileLoader: incorrect inputs
+        url_list = [[""], ["ais:"], ["ais://"], ["s3:///unkown-bucket"]]
+
+        for url in url_list:
+            with self.assertRaises(AISError):
+                file_loader = AISFileLoader(url=AIS_CLUSTER_ENDPT, source_datapipe=url)
+                for _ in file_loader:
+                    pass
+
+
+if __name__ == "__main__":
+    unittest.main()

torchdata/datapipes/iter/__init__.py

@@ -29,6 +29,10 @@
     UnBatcher,
     Zipper,
 )
+from torchdata.datapipes.iter.load.aisio import (
+    AISFileListerIterDataPipe as AISFileLister,
+    AISFileLoaderIterDataPipe as AISFileLoader,
+)
 
 ###############################################################################
 # TorchData
@@ -125,6 +129,8 @@
 from torchdata.datapipes.map.util.converter import MapToIterConverterIterDataPipe as MapToIterConverter
 
 __all__ = [
+    "AISFileLister",
+    "AISFileLoader",
     "BatchMapper",
     "Batcher",
     "BucketBatcher",

torchdata/datapipes/iter/load/README.md

@@ -1,6 +1,8 @@
-# S3 IO Datapipe Documentation
+# Iterable Datapipes
 
-## Build from Source
+## S3 IO Datapipe Documentation
+
+### Build from Source
 
 `ninja` is required to link PyThon implementation to C++ source code.
 
@@ -44,8 +46,37 @@ It's recommended to set up a detailed configuration file with the `AWS_CONFIG_FI
 environment variables are also parsed: `HOME`, `S3_USE_HTTPS`, `S3_VERIFY_SSL`, `S3_ENDPOINT_URL`, `AWS_REGION` (would
 be overwritten by the `region` variable).
 
-## Troubleshooting
+### Troubleshooting
 
 If you get `Access Denied`, it's very possibly a
 [wrong region configuration](https://github.com/aws/aws-sdk-cpp/issues/1211) or an
 [accessing issue with `aws-sdk-cpp`](https://aws.amazon.com/premiumsupport/knowledge-center/s3-access-denied-aws-sdk/).
+
+## AIStore IO Datapipe
+
+[AIStore](https://github.com/NVIDIA/aistore) (AIS for short) is a highly available lightweight object storage system
+that specifically focuses on petascale deep learning. As a reliable redundant storage, AIS supports n-way mirroring and
+erasure coding. But it is not purely – or not only – a storage system: it’ll shuffle user datasets and run custom
+extract-transform-load workloads.
+
+AIS is an elastic cluster that can grow and shrink at runtime and can be ad-hoc deployed, with or without Kubernetes,
+anywhere from a single Linux machine to a bare-metal cluster of any size.
+
+AIS fully supports Amazon S3, Google Cloud, and Microsoft Azure backends, providing a unified namespace across multiple
+connected backends and/or other AIS clusters, and [more](https://github.com/NVIDIA/aistore#features). Getting started
+with AIS will take only a few minutes (prerequisites boil down to having a Linux with a disk) and can be done either by
+running a prebuilt all-in-one docker image or directly from the open-source.
+
+### Dependency
+
+The `AISFileLister` and `AISFileLoader` under [`aisio.py`](/torchdata/datapipes/iter/load/aisio.py) internally use the
+[Python SDK](https://github.com/NVIDIA/aistore/tree/master/sdk/python) for AIStore.
+
+Run `pip install aistore` or `conda install aistore` to install the [python package](https://pypi.org/project/aistore/).
+
+### Example
+
+Please refer to the documentation:
+
+- [`AISFileLister`](https://pytorch.org/data/main/generated/torchdata.datapipes.iter.AISFileLister.html#aisfilelister)
+- [`AISFileLoader`](https://pytorch.org/data/main/generated/torchdata.datapipes.iter.AISFileLoader.html#aisfileloader)

torchdata/datapipes/iter/load/aisio.py

@@ -0,0 +1,126 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from io import BytesIO
+from typing import Iterator, Tuple
+
+from torchdata.datapipes import functional_datapipe
+
+from torchdata.datapipes.iter import IterDataPipe
+from torchdata.datapipes.utils import StreamWrapper
+
+try:
+    from aistore.client import Client
+    from aistore.pytorch.utils import parse_url, unparse_url
+
+    HAS_AIS = True
+except ImportError:
+    HAS_AIS = False
+
+
+def _assert_aistore() -> None:
+    if not HAS_AIS:
+        raise ModuleNotFoundError(
+            "Package `aistore` is required to be installed to use this datapipe."
+            "Please run `pip install aistore` or `conda install aistore` to install the package"
+            "For more info visit: https://github.com/NVIDIA/aistore/blob/master/sdk/python/"
+        )
+
+
+@functional_datapipe("list_files_by_ais")
+class AISFileListerIterDataPipe(IterDataPipe[str]):
+    """
+    Iterable Datapipe that lists files from the AIStore backends with the given URL prefixes. (functional name: ``list_files_by_ais``).
+    Acceptable prefixes include but not limited to - `ais://bucket-name`, `ais://bucket-name/`
+
+    Note:
+    -   This function also supports files from multiple backends (`aws://..`, `gcp://..`, `hdfs://..`, etc)
+    -   Input must be a list and direct URLs are not supported.
+    -   length is -1 by default, all calls to len() are invalid as
+        not all items are iterated at the start.
+    -   This internally uses AIStore Python SDK.
+
+    Args:
+        source_datapipe(IterDataPipe[str]): a DataPipe that contains URLs/URL
+                                            prefixes to objects on AIS
+        url(str): AIStore endpoint
+        length(int): length of the datapipe
+
+    Example:
+        >>> from torchdata.datapipes.iter import IterableWrapper, AISFileLister
+        >>> ais_prefixes = IterableWrapper(['ais://bucket-name/folder/', 'aws:bucket-name/folder/', ...])
+        >>> dp_ais_urls = AISFileLister(url='localhost:8080', source_datapipe=prefix)
+        >>> for d in dp_ais_urls:
+        ...     pass
+        >>> # Functional API
+        >>> dp_ais_urls = dp_ais_urls.list_files_by_ais(url='localhost:8080')
+        >>> for d in dp_ais_urls:
+        ...     pass
+    """
+
+    def __init__(self, source_datapipe: IterDataPipe[str], url: str, length: int = -1) -> None:
+        _assert_aistore()
+        self.source_datapipe: IterDataPipe[str] = source_datapipe
+        self.length: int = length
+        self.client = Client(url)
+
+    def __iter__(self) -> Iterator[str]:
+        for prefix in self.source_datapipe:
+            provider, bck_name, prefix = parse_url(prefix)
+            obj_iter = self.client.list_objects_iter(bck_name=bck_name, provider=provider, prefix=prefix)
+            for entry in obj_iter:
+                yield unparse_url(provider=provider, bck_name=bck_name, obj_name=entry.name)
+
+    def __len__(self) -> int:
+        if self.length == -1:
+            raise TypeError(f"{type(self).__name__} instance doesn't have valid length")
+        return self.length
+
+
+@functional_datapipe("load_files_by_ais")
+class AISFileLoaderIterDataPipe(IterDataPipe[Tuple[str, StreamWrapper]]):
+    """
+    Iterable DataPipe that loads files from AIStore with the given URLs (functional name: ``load_files_by_ais``).
+    Iterates all files in BytesIO format and returns a tuple (url, BytesIO).
+
+    Note:
+    -   This function also supports files from multiple backends (`aws://..`, `gcp://..`, etc)
+    -   Input must be a list and direct URLs are not supported.
+    -   This internally uses AIStore Python SDK.
+
+    Args:
+        source_datapipe(IterDataPipe[str]): a DataPipe that contains URLs/URL prefixes to objects
+        url(str): AIStore endpoint
+        length(int): length of the datapipe
+
+    Example:
+        >>> from torchdata.datapipes.iter import IterableWrapper, AISFileLister,AISFileLoader
+        >>> ais_prefixes = IterableWrapper(['ais://bucket-name/folder/', 'aws:bucket-name/folder/', ...])
+        >>> dp_ais_urls = AISFileLister(url='localhost:8080', source_datapipe=prefix)
+        >>> dp_s3_files = AISFileLoader(url='localhost:8080', source_datapipe=dp_ais_urls)
+        >>> for url, file in dp_ais_urls:
+        ...     pass
+        >>> # Functional API
+        >>> dp_ais_urls = dp_ais_urls.load_files_by_ais(url='localhost:8080')
+        >>> for url, file in dp_ais_urls:
+        ...     pass
+    """
+
+    def __init__(self, source_datapipe: IterDataPipe[str], url: str, length: int = -1) -> None:
+        _assert_aistore()
+        self.source_datapipe: IterDataPipe[str] = source_datapipe
+        self.length = length
+        self.client = Client(url)
+
+    def __iter__(self) -> Iterator[Tuple[str, StreamWrapper]]:
+        for url in self.source_datapipe:
+            provider, bck_name, obj_name = parse_url(url)
+            yield url, StreamWrapper(
+                BytesIO(self.client.get_object(bck_name=bck_name, provider=provider, obj_name=obj_name).read_all())
+            )
+
+    def __len__(self) -> int:
+        return len(self.source_datapipe)