Use variants to produce separate builds with and without cufile support

ci/build_cpp.sh

@@ -25,10 +25,15 @@ export RAPIDS_ARTIFACTS_DIR
 # populates `RATTLER_CHANNELS` array and `RATTLER_ARGS` array
 source rapids-rattler-channel-string
 
+# Use sed to replace the RAPIDS_CUDA_VERSION string in the variants file with
+# the value of that variable in this shell.
+sed -i "s/RAPIDS_CUDA_VERSION/${RAPIDS_CUDA_VERSION}/g" "conda/recipes/libkvikio/variants_$(arch).yaml"
+
 # --no-build-id allows for caching with `sccache`
 # more info is available at
 # https://rattler.build/latest/tips_and_tricks/#using-sccache-or-ccache-with-rattler-build
 rattler-build build --recipe conda/recipes/libkvikio \
+                    --variant-config "conda/recipes/libkvikio/variants_$(arch).yaml" \
                     "${RATTLER_ARGS[@]}" \
                     "${RATTLER_CHANNELS[@]}"
 

conda/recipes/libkvikio/conda_build_config.yaml

@@ -1,9 +1,3 @@
-c_compiler_version:
-  - 13
-
-cxx_compiler_version:
-  - 13
-
 cmake_version:
   - ">=3.30.4"
 

conda/recipes/libkvikio/recipe.yaml

@@ -4,11 +4,18 @@ schema_version: 1
 context:
   version: ${{ env.get("RAPIDS_PACKAGE_VERSION") }}
   minor_version: ${{ (version | split("."))[:2] | join(".") }}
-  cuda_version: ${{ (env.get("RAPIDS_CUDA_VERSION") | split("."))[:2] | join(".") }}
-  cuda_major: '${{ (env.get("RAPIDS_CUDA_VERSION") | split("."))[0] }}'
+  # We need to support three cases:
+  # 1. Linux x86_64, which always uses libcufile
+  # 2. Linux aarch64 with CUDA >= 12.2, which uses libcufile
+  # 3. Linux aarch64 with CUDA < 12.2, which does not use libcufile
+  # Each case has different cuda-version constraints as expressed below
+  should_use_cufile: ${{ x86_64 or (aarch64 and cuda_version >= "12.2") }}
+  # When reverting, instances of cuda_key_string can be replaced with cuda_major
+  cuda_key_string: ${{ cuda_version | replace(".", "_") }}
+  #cuda_version: ${{ (env.get("RAPIDS_CUDA_VERSION") | split("."))[:2] | join(".") }}
+  #cuda_major: '${{ (env.get("RAPIDS_CUDA_VERSION") | split("."))[0] }}'
   date_string: '${{ env.get("RAPIDS_DATE_STRING") }}'
   head_rev: '${{ git.head_rev(".")[:8] }}'
-  linux64: ${{ linux and x86_64 }}
 
 recipe:
   name: libkvikio-split
@@ -46,7 +53,7 @@ cache:
         SCCACHE_REGION: ${{ env.get("SCCACHE_REGION") }}
         SCCACHE_S3_USE_SSL: ${{ env.get("SCCACHE_S3_USE_SSL") }}
         SCCACHE_S3_NO_CREDENTIALS: ${{ env.get("SCCACHE_S3_NO_CREDENTIALS") }}
-        SCCACHE_S3_KEY_PREFIX: libkvikio/${{ env.get("RAPIDS_CONDA_ARCH") }}/cuda${{ cuda_major }}
+        SCCACHE_S3_KEY_PREFIX: libkvikio/${{ env.get("RAPIDS_CONDA_ARCH") }}/cuda${{ cuda_key_string }}
   requirements:
     build:
       - ${{ compiler("c") }}
@@ -59,7 +66,7 @@ cache:
     host:
       - cuda-version =${{ cuda_version }}
       - libcurl ${{ libcurl_version }}
-      - if: (linux and x86_64) or (linux and aarch64 and cuda_version >= "12.2")
+      - if: should_use_cufile
         then:
         - libcufile-dev
       - libnuma
@@ -72,7 +79,7 @@ outputs:
       script:
         content: |
           cmake --install cpp/build
-      string: cuda${{ cuda_major }}_${{ date_string }}_${{ head_rev }}
+      string: cuda${{ cuda_key_string }}_${{ date_string }}_${{ head_rev }}
       dynamic_linking:
         overlinking_behavior: "error"
       prefix_detection:
@@ -86,15 +93,25 @@ outputs:
         - cuda-version =${{ cuda_version }}
         - libcurl ${{ libcurl_version }}
       run:
-        - ${{ pin_compatible("cuda-version", upper_bound="x", lower_bound="x") }}
-        - if: linux and x86_64
+        - if: x86_64
+          then:
+          - ${{ pin_compatible("cuda-version", upper_bound="x", lower_bound="x") }}
+          else:
+          - if: aarch64 and cuda_version >= "12.2"
+            then:
+            - ${{ pin_compatible("cuda-version", upper_bound="x", lower_bound="12.2") }}
+            else:
+            - ${{ pin_compatible("cuda-version", upper_bound="12.2.0a0", lower_bound="12.0") }}
+        - if: should_use_cufile
           then:
           - libcufile-dev
       ignore_run_exports:
         by_name:
           - cuda-version
-          - libcufile
           - libcurl
+          - if: should_use_cufile
+            then:
+            - libcufile
     tests:
       - script:
           - test -f $PREFIX/include/kvikio/file_handle.hpp
@@ -107,7 +124,7 @@ outputs:
       name: libkvikio-tests
       version: ${{ version }}
     build:
-      string: cuda${{ cuda_major }}_${{ date_string }}_${{ head_rev }}
+      string: cuda${{ cuda_key_string }}_${{ date_string }}_${{ head_rev }}
       dynamic_linking:
         overlinking_behavior: "error"
       script:
@@ -121,20 +138,29 @@ outputs:
         - ${{ pin_subpackage("libkvikio", exact=True) }}
         - cuda-version =${{ cuda_version }}
         - cuda-cudart-dev
-        - if: linux and x86_64
+        - if: should_use_cufile
           then:
           - libcufile-dev
       run:
-        - ${{ pin_compatible("cuda-version", upper_bound="x", lower_bound="x") }}
+        - if: x86_64
+          then:
+          - ${{ pin_compatible("cuda-version", upper_bound="x", lower_bound="x") }}
+          else:
+          - if: aarch64 and cuda_version >= "12.2"
+            then:
+            - ${{ pin_compatible("cuda-version", upper_bound="x", lower_bound="12.2") }}
+            else:
+            - ${{ pin_compatible("cuda-version", upper_bound="12.2.0a0", lower_bound="12.0") }}
         - cuda-cudart
-        - libcufile
       ignore_run_exports:
         by_name:
           - cuda-cudart
           - cuda-version
-          - libcufile
           - libcurl
           - libnuma
+          - if: should_use_cufile
+            then:
+            - libcufile
     about:
       homepage: ${{ load_from_file("python/libkvikio/pyproject.toml").project.urls.Homepage }}
       license: ${{ load_from_file("python/libkvikio/pyproject.toml").project.license.text }}

conda/recipes/libkvikio/variants_aarch64.yaml

@@ -0,0 +1,18 @@
+zip_keys:
+  - [c_compiler_version, cxx_compiler_version, cuda_version, should_use_cufile]
+
+c_compiler_version:
+  - 12
+  - 13
+
+cxx_compiler_version:
+  - 12
+  - 13
+
+cuda_version:
+  - 12.1 # The last version to not support cufile
+  - RAPIDS_CUDA_VERSION
+
+should_use_cufile:
+  - false
+  - true

conda/recipes/libkvikio/variants_x86_64.yaml

@@ -0,0 +1,11 @@
+c_compiler_version:
+  - 13
+
+cxx_compiler_version:
+  - 13
+
+cuda_version:
+  - RAPIDS_CUDA_VERSION
+
+should_use_cufile:
+  - true

cpp/include/kvikio/shim/cufile.hpp

@@ -128,11 +128,7 @@ bool is_cufile_available() noexcept;
  *
  * @return The version (1000*major + 10*minor) or zero if older than 1080.
  */
-#ifdef KVIKIO_CUFILE_FOUND
 int cufile_version() noexcept;
-#else
-constexpr int cufile_version() noexcept { return 0; }
-#endif
 
 /**
  * @brief Check if cuFile's batch API is available.

cpp/src/shim/cufile.cpp

@@ -147,6 +147,8 @@ int cufile_version() noexcept
     return 0;
   }
 }
+#else
+int cufile_version() noexcept { return 0; }
 #endif
 
 bool is_batch_api_available() noexcept { return cufile_version() >= 1060; }

python/kvikio/tests/test_cufile_driver.py

@@ -13,47 +13,59 @@ def test_version():
 
 
 @pytest.mark.cufile
-def test_open_and_close():
-    kvikio.cufile_driver.driver_open()
-    kvikio.cufile_driver.driver_close()
+def test_open_and_close(request):
+    try:
+        kvikio.cufile_driver.driver_open()
+        kvikio.cufile_driver.driver_close()
+    except RuntimeError as e:
+        if "KvikIO not compiled with cuFile.h" in str(e):
+            pytest.skip("KvikIO not compiled with cuFile.h, skipping cuFile tests")
 
 
 @pytest.mark.cufile
 def test_property_accessor():
     """Test the method `get` and `set`"""
 
-    # Attempt to set a nonexistent property
-    with pytest.raises(KeyError):
-        kvikio.cufile_driver.set("nonexistent_property", 123)
-
-    # Attempt to get a nonexistent property
-    with pytest.raises(KeyError):
-        kvikio.cufile_driver.get("nonexistent_property")
-
-    # Attempt to set a read-only property
-    with pytest.raises(KeyError, match="read-only"):
-        kvikio.cufile_driver.set("major_version", 2077)
-
-    # Nested context managers
-    poll_thresh_size_default = kvikio.cufile_driver.get("poll_thresh_size")
-    with kvikio.cufile_driver.set("poll_thresh_size", 1024):
-        assert kvikio.cufile_driver.get("poll_thresh_size") == 1024
-        with kvikio.cufile_driver.set("poll_thresh_size", 2048):
-            assert kvikio.cufile_driver.get("poll_thresh_size") == 2048
-            with kvikio.cufile_driver.set("poll_thresh_size", 4096):
-                assert kvikio.cufile_driver.get("poll_thresh_size") == 4096
-            assert kvikio.cufile_driver.get("poll_thresh_size") == 2048
-        assert kvikio.cufile_driver.get("poll_thresh_size") == 1024
-    assert kvikio.cufile_driver.get("poll_thresh_size") == poll_thresh_size_default
-
-    # Multiple context managers
-    poll_mode_default = kvikio.cufile_driver.get("poll_mode")
-    max_device_cache_size_default = kvikio.cufile_driver.get("max_device_cache_size")
-    with kvikio.cufile_driver.set({"poll_mode": True, "max_device_cache_size": 2048}):
-        assert kvikio.cufile_driver.get("poll_mode") and (
-            kvikio.cufile_driver.get("max_device_cache_size") == 2048
+    try:
+        # Attempt to set a nonexistent property
+        with pytest.raises(KeyError):
+            kvikio.cufile_driver.set("nonexistent_property", 123)
+
+        # Attempt to get a nonexistent property
+        with pytest.raises(KeyError):
+            kvikio.cufile_driver.get("nonexistent_property")
+
+        # Attempt to set a read-only property
+        with pytest.raises(KeyError, match="read-only"):
+            kvikio.cufile_driver.set("major_version", 2077)
+
+        # Nested context managers
+        poll_thresh_size_default = kvikio.cufile_driver.get("poll_thresh_size")
+        with kvikio.cufile_driver.set("poll_thresh_size", 1024):
+            assert kvikio.cufile_driver.get("poll_thresh_size") == 1024
+            with kvikio.cufile_driver.set("poll_thresh_size", 2048):
+                assert kvikio.cufile_driver.get("poll_thresh_size") == 2048
+                with kvikio.cufile_driver.set("poll_thresh_size", 4096):
+                    assert kvikio.cufile_driver.get("poll_thresh_size") == 4096
+                assert kvikio.cufile_driver.get("poll_thresh_size") == 2048
+            assert kvikio.cufile_driver.get("poll_thresh_size") == 1024
+        assert kvikio.cufile_driver.get("poll_thresh_size") == poll_thresh_size_default
+
+        # Multiple context managers
+        poll_mode_default = kvikio.cufile_driver.get("poll_mode")
+        max_device_cache_size_default = kvikio.cufile_driver.get(
+            "max_device_cache_size"
+        )
+        with kvikio.cufile_driver.set(
+            {"poll_mode": True, "max_device_cache_size": 2048}
+        ):
+            assert kvikio.cufile_driver.get("poll_mode") and (
+                kvikio.cufile_driver.get("max_device_cache_size") == 2048
+            )
+        assert (kvikio.cufile_driver.get("poll_mode") == poll_mode_default) and (
+            kvikio.cufile_driver.get("max_device_cache_size")
+            == max_device_cache_size_default
         )
-    assert (kvikio.cufile_driver.get("poll_mode") == poll_mode_default) and (
-        kvikio.cufile_driver.get("max_device_cache_size")
-        == max_device_cache_size_default
-    )
+    except RuntimeError as e:
+        if "KvikIO not compiled with cuFile.h" in str(e):
+            pytest.skip("KvikIO not compiled with cuFile.h, skipping cuFile tests")