S3 support

.github/workflows/pr.yaml

@@ -18,7 +18,7 @@ jobs:
       - conda-python-build
       - conda-python-tests
       - docs-build
-      - devcontainer
+      # - devcontainer
       - wheel-cpp-build
       - wheel-python-build
       - wheel-python-tests
@@ -61,16 +61,17 @@ jobs:
       arch: "amd64"
       container_image: "rapidsai/ci-conda:latest"
       run_script: "ci/build_docs.sh"
-  devcontainer:
-    secrets: inherit
-    uses: rapidsai/shared-workflows/.github/workflows/[email protected]
-    with:
-      arch: '["amd64"]'
-      cuda: '["12.5"]'
-      build_command: |
-        sccache -z;
-        build-all --verbose;
-        sccache -s;
+  # TODO: uncomment when https://github.com/rapidsai/kvikio/issues/446 has been fixed.
+  # devcontainer:
+  #   secrets: inherit
+  #   uses: rapidsai/shared-workflows/.github/workflows/[email protected]
+  #   with:
+  #     arch: '["amd64"]'
+  #     cuda: '["12.5"]'
+  #     build_command: |
+  #       sccache -z;
+  #       build-all --verbose;
+  #       sccache -s;
   wheel-cpp-build:
     secrets: inherit
     uses: rapidsai/shared-workflows/.github/workflows/[email protected]

README.md

@@ -15,9 +15,89 @@ The C++ library is header-only making it easy to include in [existing projects](
 * A Python [Zarr](https://zarr.readthedocs.io/en/stable/) backend for reading and writing GPU data to file seamlessly.
 * Concurrent reads and writes using an internal thread pool.
 * Non-blocking API.
-* Handle both host and device IO seamlessly.
+* Read/write to both host and device memory seamlessly.
+* Provides remote read from AWS S3 storage seamlessly.
 * Provides Python bindings to [nvCOMP](https://github.com/NVIDIA/nvcomp).
 
+
 ### Documentation
  * Python: <https://docs.rapids.ai/api/kvikio/nightly/>
  * C++: <https://docs.rapids.ai/api/libkvikio/nightly/>
+
+
+### Examples
+
+#### Python
+```python
+import cupy
+import kvikio
+
+def main(path):
+    a = cupy.arange(100)
+    f = kvikio.CuFile(path, "w")
+    # Write whole array to file
+    f.write(a)
+    f.close()
+
+    b = cupy.empty_like(a)
+    f = kvikio.CuFile(path, "r")
+    # Read whole array from file
+    f.read(b)
+    assert all(a == b)
+
+    # Use contexmanager
+    c = cupy.empty_like(a)
+    with kvikio.CuFile(path, "r") as f:
+        f.read(c)
+    assert all(a == c)
+
+    # Non-blocking read
+    d = cupy.empty_like(a)
+    with kvikio.CuFile(path, "r") as f:
+        future1 = f.pread(d[:50])
+        future2 = f.pread(d[50:], file_offset=d[:50].nbytes)
+        future1.get()  # Wait for first read
+        future2.get()  # Wait for second read
+    assert all(a == d)
+
+
+if __name__ == "__main__":
+    main("/tmp/kvikio-hello-world-file")
+```
+
+#### C++
+```c++
+#include <cstddef>
+#include <cuda_runtime.h>
+#include <kvikio/file_handle.hpp>
+using namespace std;
+
+int main()
+{
+  // Create two arrays `a` and `b`
+  constexpr std::size_t size = 100;
+  void *a = nullptr;
+  void *b = nullptr;
+  cudaMalloc(&a, size);
+  cudaMalloc(&b, size);
+
+  // Write `a` to file
+  kvikio::FileHandle fw("test-file", "w");
+  size_t written = fw.write(a, size);
+  fw.close();
+
+  // Read file into `b`
+  kvikio::FileHandle fr("test-file", "r");
+  size_t read = fr.read(b, size);
+  fr.close();
+
+  // Read file into `b` in parallel using 16 threads
+  kvikio::default_thread_pool::reset(16);
+  {
+    kvikio::FileHandle f("test-file", "r");
+    future<size_t> future = f.pread(b_dev, sizeof(a), 0);  // Non-blocking
+    size_t read = future.get(); // Blocking
+    // Notice, `f` closes automatically on destruction.
+  }
+}
+```

build.sh

@@ -18,15 +18,16 @@ ARGS=$*
 # script, and that this script resides in the repo dir!
 REPODIR=$(cd $(dirname $0); pwd)
 
-VALIDARGS="clean libkvikio kvikio -v -g -n --pydevelop -h"
-HELP="$0 [clean] [libkvikio] [kvikio] [-v] [-g] [-n] [--cmake-args=\"<args>\"] [-h]
+VALIDARGS="clean libkvikio kvikio -v -g -n --pydevelop --no-s3 -h"
+HELP="$0 [clean] [libkvikio] [kvikio] [--no-s3] [-v] [-g] [-n] [--pydevelop] [--cmake-args=\"<args>\"] [-h]
    clean                       - remove all existing build artifacts and configuration (start over)
    libkvikio                   - build and install the libkvikio C++ code
    kvikio                      - build and install the kvikio Python package
+   --no-s3                     - build with no AWS S3 support
    -v                          - verbose build mode
    -g                          - build for debug
    -n                          - no install step
-   --pydevelop                 - Install Python packages in editable mode
+   --pydevelop                 - install Python packages in editable mode
    --cmake-args=\\\"<args>\\\" - pass arbitrary list of CMake configuration options (escape all quotes in argument)
    -h                          - print this text
    default action (no args) is to build and install 'libkvikio' and 'kvikio' targets
@@ -36,6 +37,7 @@ KVIKIO_BUILD_DIR="${REPODIR}/python/build ${REPODIR}/python/_skbuild"
 BUILD_DIRS="${LIBKVIKIO_BUILD_DIR} ${KVIKIO_BUILD_DIR}"
 
 # Set defaults for vars modified by flags to this script
+ENABLE_S3_SUPPORT="-DKvikIO_AWSSDK_SUPPORT=ON"
 VERBOSE_FLAG=""
 BUILD_TYPE=Release
 INSTALL_TARGET=install
@@ -86,6 +88,7 @@ function ensureCMakeRan {
         cmake -B "${LIBKVIKIO_BUILD_DIR}" -S . \
               -DCMAKE_INSTALL_PREFIX="${INSTALL_PREFIX}" \
               -DCMAKE_BUILD_TYPE=${BUILD_TYPE} \
+              ${ENABLE_S3_SUPPORT} \
               ${EXTRA_CMAKE_ARGS}
         RAN_CMAKE=1
     fi
@@ -109,6 +112,9 @@ if (( ${NUMARGS} != 0 )); then
 fi
 
 # Process flags
+if hasArg --no-s3; then
+    ENABLE_S3_SUPPORT="-DKvikIO_AWSSDK_SUPPORT=OFF"
+fi
 if hasArg -v; then
     VERBOSE_FLAG=-v
     set -x
@@ -150,7 +156,7 @@ if (( NUMARGS == 0 )) || hasArg libkvikio; then
     cmake --build "${LIBKVIKIO_BUILD_DIR}" -j${PARALLEL_LEVEL} ${VERBOSE_FLAG}
     if [[ ${INSTALL_TARGET} != "" ]]; then
         echo "installing libkvikio..."
-        cmake --build "${LIBKVIKIO_BUILD_DIR}" --target install -v ${VERBOSE_FLAG}
+        cmake --build "${LIBKVIKIO_BUILD_DIR}" --target install ${VERBOSE_FLAG}
     fi
 fi
 

conda/environments/all_cuda-118_arch-aarch64.yaml

@@ -6,6 +6,8 @@ channels:
 - conda-forge
 - nvidia
 dependencies:
+- aws-sdk-cpp>=1.11.267
+- boto3>=1.21.21
 - c-compiler
 - cmake>=3.26.4,!=3.30.0
 - cuda-python>=11.7.1,<12.0a0
@@ -17,6 +19,7 @@ dependencies:
 - dask>=2022.05.2
 - doxygen=1.9.1
 - gcc_linux-aarch64=11.*
+- moto>=4.0.8
 - ninja
 - numcodecs <0.12.0
 - numpy>=1.23,<3.0a0

conda/environments/all_cuda-118_arch-x86_64.yaml

@@ -6,6 +6,8 @@ channels:
 - conda-forge
 - nvidia
 dependencies:
+- aws-sdk-cpp>=1.11.267
+- boto3>=1.21.21
 - c-compiler
 - cmake>=3.26.4,!=3.30.0
 - cuda-python>=11.7.1,<12.0a0
@@ -19,6 +21,7 @@ dependencies:
 - gcc_linux-64=11.*
 - libcufile-dev=1.4.0.31
 - libcufile=1.4.0.31
+- moto>=4.0.8
 - ninja
 - numcodecs <0.12.0
 - numpy>=1.23,<3.0a0

conda/environments/all_cuda-125_arch-aarch64.yaml

@@ -6,6 +6,8 @@ channels:
 - conda-forge
 - nvidia
 dependencies:
+- aws-sdk-cpp>=1.11.267
+- boto3>=1.21.21
 - c-compiler
 - cmake>=3.26.4,!=3.30.0
 - cuda-nvcc
@@ -18,6 +20,7 @@ dependencies:
 - doxygen=1.9.1
 - gcc_linux-aarch64=11.*
 - libcufile-dev
+- moto>=4.0.8
 - ninja
 - numcodecs <0.12.0
 - numpy>=1.23,<3.0a0

conda/environments/all_cuda-125_arch-x86_64.yaml

@@ -6,6 +6,8 @@ channels:
 - conda-forge
 - nvidia
 dependencies:
+- aws-sdk-cpp>=1.11.267
+- boto3>=1.21.21
 - c-compiler
 - cmake>=3.26.4,!=3.30.0
 - cuda-nvcc
@@ -18,6 +20,7 @@ dependencies:
 - doxygen=1.9.1
 - gcc_linux-64=11.*
 - libcufile-dev
+- moto>=4.0.8
 - ninja
 - numcodecs <0.12.0
 - numpy>=1.23,<3.0a0

conda/recipes/kvikio/meta.yaml

@@ -64,11 +64,13 @@ requirements:
     - rapids-build-backend >=0.3.0,<0.4.0.dev0
     - scikit-build-core >=0.10.0
     - libkvikio ={{ version }}
+    - aws-sdk-cpp>=1.11.267
   run:
     - python
     - numpy >=1.23,<3.0a0
     - cupy >=12.0.0
     - zarr
+    - aws-sdk-cpp>=1.11.267
     # See https://github.com/zarr-developers/numcodecs/pull/475
     - numcodecs <0.12.0
     - packaging

conda/recipes/libkvikio/meta.yaml

@@ -52,6 +52,7 @@ requirements:
     {% else %}
     - libcufile-dev  # [linux]
     {% endif %}
+    - aws-sdk-cpp>=1.11.267
 
 outputs:
   - name: libkvikio
@@ -83,6 +84,7 @@ outputs:
         {% else %}
         - libcufile-dev  # [linux]
         {% endif %}
+        - aws-sdk-cpp>=1.11.267
     test:
         commands:
           - test -f $PREFIX/include/kvikio/file_handle.hpp
@@ -106,6 +108,7 @@ outputs:
         - cuda-cudart-dev
         - libcufile-dev  # [linux]
         {% endif %}
+        - aws-sdk-cpp>=1.11.267
     requirements:
       build:
         - cmake {{ cmake_version }}
@@ -118,6 +121,7 @@ outputs:
         - cuda-cudart-dev
         - libcufile-dev  # [linux]
         {% endif %}
+        - aws-sdk-cpp>=1.11.267
       run:
         - {{ pin_compatible('cuda-version', max_pin='x', min_pin='x') }}
         {% if cuda_major == "11" %}
@@ -127,6 +131,7 @@ outputs:
         - cuda-cudart
         - libcufile  # [linux]
         {% endif %}
+        - aws-sdk-cpp>=1.11.267
     about:
       home: https://rapids.ai
       license: Apache-2.0

cpp/CMakeLists.txt

@@ -37,6 +37,7 @@ rapids_cmake_build_type(Release)
 
 # build options
 option(KvikIO_BUILD_EXAMPLES "Configure CMake to build examples" ON)
+option(KvikIO_AWSSDK_SUPPORT "Configure CMake to build with AWS S3 support" ON)
 
 rapids_cmake_support_conda_env(conda_env MODIFY_PREFIX_PATH)
 
@@ -55,6 +56,31 @@ rapids_find_package(
   INSTALL_EXPORT_SET kvikio-exports
 )
 
+# If AWSSDK isn't found, the Cython module remote_handle.pyx isn't built and C++ users shouldn't
+# include <kvikio/remote_handle.hpp>
+if(KvikIO_AWSSDK_SUPPORT)
+  rapids_find_package(
+    AWSSDK REQUIRED COMPONENTS s3
+    BUILD_EXPORT_SET kvikio-exports
+    INSTALL_EXPORT_SET kvikio-exports
+  )
+endif()
+
+if(AWSSDK_FOUND)
+  get_property(
+    _lib_type
+    TARGET aws-cpp-sdk-s3
+    PROPERTY TYPE
+  )
+  if(_lib_type STREQUAL "STATIC_LIBRARY")
+    rapids_find_package(
+      ZLIB
+      BUILD_EXPORT_SET kvikio-exports
+      INSTALL_EXPORT_SET kvikio-exports
+    )
+  endif()
+endif()
+
 rapids_find_package(
   cuFile
   BUILD_EXPORT_SET kvikio-exports
@@ -131,6 +157,10 @@ target_include_directories(
 target_link_libraries(
   kvikio INTERFACE Threads::Threads ${CMAKE_DL_LIBS} nvtx3::nvtx3-cpp BS::thread_pool
 )
+if(AWSSDK_FOUND)
+  target_link_libraries(kvikio INTERFACE aws-cpp-sdk-s3)
+  target_compile_definitions(kvikio INTERFACE $<BUILD_LOCAL_INTERFACE:KVIKIO_AWS_SDK_FOUND>)
+endif()
 target_compile_features(kvikio INTERFACE cxx_std_17)
 
 # optionally build examples
@@ -217,6 +247,13 @@ if(NOT already_set_kvikio)
       target_compile_definitions(kvikio::kvikio INTERFACE KVIKIO_CUFILE_STREAM_API_FOUND)
     endif()
   endif()
+
+  if(KvikIO_AWSSDK_SUPPORT)
+    find_package(AWSSDK COMPONENTS s3 QUIET)
+  endif()
+  if(AWSSDK_FOUND)
+    target_compile_definitions(kvikio::kvikio INTERFACE $<BUILD_LOCAL_INTERFACE:KVIKIO_AWS_SDK_FOUND>)
+  endif()
 endif()
 ]=]
 )

cpp/examples/CMakeLists.txt

@@ -14,6 +14,8 @@
 
 set(TEST_INSTALL_PATH bin/tests/libkvikio)
 
+# Example: basic_io
+
 if(CUDAToolkit_FOUND)
   add_executable(BASIC_IO_TEST basic_io.cpp)
   set_target_properties(BASIC_IO_TEST PROPERTIES INSTALL_RPATH "\$ORIGIN/../../lib")
@@ -35,6 +37,8 @@ else()
   message(STATUS "Cannot build the basic_io example when CUDA is not found")
 endif()
 
+# Example: basic_no_cuda
+
 add_executable(BASIC_NO_CUDA_TEST basic_no_cuda.cpp)
 set_target_properties(BASIC_NO_CUDA_TEST PROPERTIES INSTALL_RPATH "\$ORIGIN/../../lib")
 target_include_directories(BASIC_NO_CUDA_TEST PRIVATE ../include)