S3 datapipes

.flake8

@@ -4,6 +4,7 @@ ignore = E203,E402,E501,F821,W503,W504,
 per-file-ignores =
     __init__.py: F401, F403, F405
     test/*: F401
+    _extension.py: F401
 exclude =
     ./.git,
     ./third_party,

.github/workflows/ci.yml

@@ -29,27 +29,77 @@ jobs:
           - 3.7
           - 3.8
           - 3.9
+        with-s3:
+          - 1
+          - 0
     steps:
       - name: Setup additional system libraries
         if: startsWith( matrix.os, 'ubuntu' )
         run: |
           sudo add-apt-repository multiverse
           sudo apt update
-          sudo apt install rar unrar
+          sudo apt install rar unrar libssl-dev libcurl4-openssl-dev zlib1g-dev
       - name: Setup Python ${{ matrix.python-version }}
         uses: actions/setup-python@v2
         with:
           python-version: ${{ matrix.python-version }}
+      - name: Setup msbuild on Windows
+        if: matrix.with-s3 == 1 && matrix.os == 'windows-latest'
+        uses: microsoft/[email protected]
+      - name: Set up Visual Studio shell
+        if: matrix.with-s3 == 1 && matrix.os == 'windows-latest'
+        uses: egor-tensin/vs-shell@v2
+        with:
+          arch: x64
       - name: Check out source repository
         uses: actions/checkout@v2
       - name: Install dependencies
         run: |
           pip3 install -r requirements.txt
           pip3 install --pre torch -f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html
+          pip3 install cmake ninja pybind11
+          echo "/home/runner/.local/bin" >> $GITHUB_PATH
+      - name: Export AWS-SDK-CPP & PYBIND11
+        if: matrix.with-s3 == 1
+        shell: bash
+        run: |
+          if [[ ${{ matrix.os }} == 'windows-latest' ]]; then
+            AWSSDK_PATH="$GITHUB_WORKSPACE\\aws-sdk-cpp\\sdk-lib"
+          else
+            AWSSDK_PATH="$GITHUB_WORKSPACE/aws-sdk-cpp/sdk-lib"
+          fi
+          PYBIND11_PATH=`pybind11-config --cmakedir`
+          echo "::set-output name=awssdk::$AWSSDK_PATH"
+          echo "::set-output name=pybind11::$PYBIND11_PATH"
+        id: export_path
+      - name: Install AWS-SDK-CPP on Windows for S3 IO datapipes
+        if: matrix.with-s3 == 1 && matrix.os == 'windows-latest'
+        run: |
+          git clone --recurse-submodules https://github.com/aws/aws-sdk-cpp
+          cd aws-sdk-cpp
+          mkdir sdk-lib
+          cmake -S . -B build -GNinja -DBUILD_ONLY="s3;transfer" -DBUILD_SHARED_LIBS=OFF -DENABLE_TESTING=OFF -DCMAKE_BUILD_TYPE=Release -DCMAKE_INSTALL_PREFIX=sdk-lib
+          cmake --build build --config Release
+          cmake --install build --config Release
+      - name: Install AWS-SDK-CPP on Non-Windows for S3 IO datapipes
+        if: matrix.with-s3 == 1 && matrix.os != 'windows-latest'
+        run: |
+          git clone --recurse-submodules https://github.com/aws/aws-sdk-cpp
+          cd aws-sdk-cpp/
+          mkdir sdk-build sdk-lib
+          cd sdk-build
+          cmake .. -DCMAKE_BUILD_TYPE=Release -DBUILD_ONLY="s3;transfer" -DENABLE_TESTING=OFF -DCMAKE_INSTALL_PREFIX=../sdk-lib
+          make
+          sudo make install
+      - name: Build TorchData
+        run: |
+          python setup.py develop
+        env:
+          BUILD_S3: ${{ matrix.with-s3 }}
+          pybind11_DIR: ${{ steps.export_path.outputs.pybind11 }}
+          AWSSDK_DIR: ${{ steps.export_path.outputs.awssdk }}
       - name: Install test requirements
         run: pip3 install expecttest fsspec iopath==0.1.9 numpy pytest rarfile
-      - name: Build TorchData
-        run: python setup.py develop
       - name: Run DataPipes tests with pytest
         if: ${{ ! contains(github.event.pull_request.labels.*.name, 'ciflow/slow') }}
         run:

.gitignore

@@ -1,6 +1,6 @@
 build/*
 dist/*
-torchdata.egg-info/*
+*.egg-info/*
 
 torchdata/version.py
 torchdata/datapipes/iter/__init__.pyi
@@ -17,6 +17,24 @@ torchdata/datapipes/iter/__init__.pyi
 # macOS dir files
 .DS_Store
 
+## General
+
+*/*.so*
+*/**/*.so*
+torchdata/*.so*
+
+# Compiled Object files
+*.slo
+*.lo
+*.o
+*.cuo
+*.obj
+
+# Compiled Dynamic libraries
+*.so
+*.dylib
+*.dll
+
 # Compiled python
 *.pyc
 *.pyd

CMakeLists.txt

@@ -0,0 +1,59 @@
+cmake_minimum_required(VERSION 3.13 FATAL_ERROR)
+
+# Most of the configurations are taken from PyTorch
+# https://github.com/pytorch/pytorch/blob/0c9fb4aff0d60eaadb04e4d5d099fb1e1d5701a9/CMakeLists.txt
+
+# Use compiler ID "AppleClang" instead of "Clang" for XCode.
+# Not setting this sometimes makes XCode C compiler gets detected as "Clang",
+# even when the C++ one is detected as "AppleClang".
+cmake_policy(SET CMP0010 NEW)
+cmake_policy(SET CMP0025 NEW)
+
+# Suppress warning flags in default MSVC configuration.  It's not
+# mandatory that we do this (and we don't if cmake is old), but it's
+# nice when it's possible, and it's possible on our Windows configs.
+if(NOT CMAKE_VERSION VERSION_LESS 3.15.0)
+  cmake_policy(SET CMP0092 NEW)
+endif()
+
+project(torchdata)
+
+# check and set CMAKE_CXX_STANDARD
+string(FIND "${CMAKE_CXX_FLAGS}" "-std=c++" env_cxx_standard)
+if(env_cxx_standard GREATER -1)
+  message(
+      WARNING "C++ standard version definition detected in environment variable."
+      "PyTorch requires -std=c++14. Please remove -std=c++ settings in your environment.")
+endif()
+
+set(CMAKE_CXX_STANDARD 14)
+set(CMAKE_C_STANDARD 11)
+
+# https://developercommunity.visualstudio.com/t/VS-16100-isnt-compatible-with-CUDA-11/1433342
+if(MSVC)
+  if(USE_CUDA)
+    set(CMAKE_CXX_STANDARD 17)
+  endif()
+endif()
+
+
+set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
+set(CMAKE_POSITION_INDEPENDENT_CODE ON)
+
+# Apple specific
+if(APPLE)
+  # Get clang version on macOS
+  execute_process( COMMAND ${CMAKE_CXX_COMPILER} --version OUTPUT_VARIABLE clang_full_version_string )
+  string(REGEX REPLACE "Apple LLVM version ([0-9]+\\.[0-9]+).*" "\\1" CLANG_VERSION_STRING ${clang_full_version_string})
+  message( STATUS "CLANG_VERSION_STRING:         " ${CLANG_VERSION_STRING} )
+
+  # RPATH stuff
+  set(CMAKE_MACOSX_RPATH ON)
+
+  set(CMAKE_SHARED_LIBRARY_SUFFIX ".so")
+endif()
+
+# Options
+option(BUILD_S3 "Build s3 io functionality" OFF)
+
+add_subdirectory(torchdata/csrc)

setup.py

@@ -1,13 +1,17 @@
 #!/usr/bin/env python
 # Copyright (c) Facebook, Inc. and its affiliates.
+import distutils.command.clean
 import os
+import shutil
 import subprocess
+import sys
 
 from pathlib import Path
 
 from setuptools import find_packages, setup
-from torchdata.datapipes.gen_pyi import gen_pyi
 
+from tools import setup_helpers
+from tools.gen_pyi import gen_pyi
 
 ROOT_DIR = Path(__file__).parent.resolve()
 
@@ -52,12 +56,41 @@ def _export_version(version, sha):
 ]
 
 
+class clean(distutils.command.clean.clean):
+    def run(self):
+        # Run default behavior first
+        distutils.command.clean.clean.run(self)
+
+        # Remove torchdata extension
+        def remove_extension(pattern):
+            for path in (ROOT_DIR / "torchdata").glob(pattern):
+                print(f"removing extension '{path}'")
+                path.unlink()
+
+        for ext in ["so", "dylib", "pyd"]:
+            remove_extension("**/*." + ext)
+
+        # Remove build directory
+        build_dirs = [
+            ROOT_DIR / "build",
+        ]
+        for path in build_dirs:
+            if path.exists():
+                print(f"removing '{path}' (and everything under it)")
+                shutil.rmtree(str(path), ignore_errors=True)
+
+
 if __name__ == "__main__":
     VERSION, SHA = _get_version()
     _export_version(VERSION, SHA)
 
     print("-- Building version " + VERSION)
 
+    if sys.argv[1] != "clean":
+        gen_pyi()
+        # TODO: Fix #343
+        os.chdir(ROOT_DIR)
+
     setup(
         # Metadata
         name="torchdata",
@@ -82,8 +115,18 @@ def _export_version(version, sha):
             "Programming Language :: Python :: Implementation :: CPython",
             "Topic :: Scientific/Engineering :: Artificial Intelligence",
         ],
+        package_data={
+            "torchdata": [
+                "datapipes/iter/*.pyi",
+            ],
+        },
         # Package Info
-        packages=find_packages(exclude=["test*", "examples*"]),
+        packages=find_packages(exclude=["test*", "examples*", "tools*", "torchdata.csrc*", "build*"]),
         zip_safe=False,
+        # C++ Extension Modules
+        ext_modules=setup_helpers.get_ext_modules(),
+        cmdclass={
+            "build_ext": setup_helpers.CMakeBuild,
+            "clean": clean,
+        },
     )
-    gen_pyi()

test/test_remote_io.py

@@ -6,9 +6,19 @@
 
 import expecttest
 
+import torchdata
+
 from _utils._common_utils_for_test import check_hash_fn, create_temp_dir
 
-from torchdata.datapipes.iter import EndOnDiskCacheHolder, FileOpener, HttpReader, IterableWrapper, OnDiskCacheHolder
+from torchdata.datapipes.iter import (
+    EndOnDiskCacheHolder,
+    FileOpener,
+    HttpReader,
+    IterableWrapper,
+    OnDiskCacheHolder,
+    S3FileLister,
+    S3FileLoader,
+)
 
 
 class TestDataPipeRemoteIO(expecttest.TestCase):
@@ -161,6 +171,110 @@ def _read_and_decode(x):
             self.assertTrue(os.path.exists(expected_csv_path))
             self.assertEqual(expected_csv_path, csv_path)
 
+    def test_s3_io_iterdatapipe(self):
+        # sanity test
+        file_urls = ["s3://ai2-public-datasets"]
+        try:
+            s3_lister_dp = S3FileLister(IterableWrapper(file_urls))
+            s3_loader_dp = S3FileLoader(IterableWrapper(file_urls))
+        except ModuleNotFoundError:
+            warnings.warn(
+                "S3 IO datapipes or C++ extension '_torchdata' isn't built in the current 'torchdata' package"
+            )
+            return
+
+        # S3FileLister: different inputs
+        input_list = [
+            [["s3://ai2-public-datasets"], 71],  # bucket without '/'
+            [["s3://ai2-public-datasets/"], 71],  # bucket with '/'
+            [["s3://ai2-public-datasets/charades"], 18],  # folder without '/'
+            [["s3://ai2-public-datasets/charades/"], 18],  # folder without '/'
+            [["s3://ai2-public-datasets/charad"], 18],  # prefix
+            [
+                [
+                    "s3://ai2-public-datasets/charades/Charades_v1",
+                    "s3://ai2-public-datasets/charades/Charades_vu17",
+                ],
+                12,
+            ],  # prefixes
+            [["s3://ai2-public-datasets/charades/Charades_v1.zip"], 1],  # single file
+            [
+                [
+                    "s3://ai2-public-datasets/charades/Charades_v1.zip",
+                    "s3://ai2-public-datasets/charades/Charades_v1_flow.tar",
+                    "s3://ai2-public-datasets/charades/Charades_v1_rgb.tar",
+                    "s3://ai2-public-datasets/charades/Charades_v1_480.zip",
+                ],
+                4,
+            ],  # multiple files
+            [
+                [
+                    "s3://ai2-public-datasets/charades/Charades_v1.zip",
+                    "s3://ai2-public-datasets/charades/Charades_v1_flow.tar",
+                    "s3://ai2-public-datasets/charades/Charades_v1_rgb.tar",
+                    "s3://ai2-public-datasets/charades/Charades_v1_480.zip",
+                    "s3://ai2-public-datasets/charades/Charades_vu17",
+                ],
+                10,
+            ],  # files + prefixes
+        ]
+        for input in input_list:
+            s3_lister_dp = S3FileLister(IterableWrapper(input[0]), region="us-west-2")
+            self.assertEqual(sum(1 for _ in s3_lister_dp), input[1], f"{input[0]} failed")
+
+        # S3FileLister: prefixes + different region
+        file_urls = [
+            "s3://aft-vbi-pds/bin-images/111",
+            "s3://aft-vbi-pds/bin-images/222",
+        ]
+        s3_lister_dp = S3FileLister(IterableWrapper(file_urls), region="us-east-1")
+        self.assertEqual(sum(1 for _ in s3_lister_dp), 2212, f"{input} failed")
+
+        # S3FileLister: incorrect inputs
+        input_list = [
+            [""],
+            ["ai2-public-datasets"],
+            ["s3://"],
+            ["s3:///bin-images"],
+        ]
+        for input in input_list:
+            with self.assertRaises(ValueError, msg=f"{input} should raise ValueError."):
+                s3_lister_dp = S3FileLister(IterableWrapper(input), region="us-east-1")
+                for _ in s3_lister_dp:
+                    pass
+
+        # S3FileLoader: loader
+        input = [
+            "s3://charades-tar-shards/charades-video-0.tar",
+            "s3://charades-tar-shards/charades-video-1.tar",
+        ]  # multiple files
+        s3_loader_dp = S3FileLoader(input, region="us-west-2")
+        self.assertEqual(sum(1 for _ in s3_loader_dp), 2, f"{input} failed")
+
+        input = [["s3://aft-vbi-pds/bin-images/100730.jpg"], 1]
+        s3_loader_dp = S3FileLoader(input[0], region="us-east-1")
+        self.assertEqual(sum(1 for _ in s3_loader_dp), input[1], f"{input[0]} failed")
+
+        # S3FileLoader: incorrect inputs
+        input_list = [
+            [""],
+            ["ai2-public-datasets"],
+            ["s3://"],
+            ["s3:///bin-images"],
+            ["s3://ai2-public-datasets/bin-image"],
+        ]
+        for input in input_list:
+            with self.assertRaises(ValueError, msg=f"{input} should raise ValueError."):
+                s3_loader_dp = S3FileLoader(input, region="us-east-1")
+                for _ in s3_loader_dp:
+                    pass
+
+        # integration test
+        input = [["s3://charades-tar-shards/"], 10]
+        s3_lister_dp = S3FileLister(IterableWrapper(input[0]), region="us-west-2")
+        s3_loader_dp = S3FileLoader(s3_lister_dp, region="us-west-2")
+        self.assertEqual(sum(1 for _ in s3_loader_dp), input[1], f"{input[0]} failed")
+
 
 if __name__ == "__main__":
     unittest.main()