removed stream and device_id

rapidsai · Feb 27, 2025 · 0337847 · 0337847
1 parent 2a8f203
commit 0337847
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 78 deletions.
diff --git a/python/kvikio/kvikio/_lib/libnvcomp.pyx b/python/kvikio/kvikio/_lib/libnvcomp.pyx
@@ -150,8 +150,6 @@ cdef class _ANSManager(_nvcompManager):
     def __cinit__(
         self,
         size_t uncomp_chunk_size,
-        user_stream,
-        const int device_id,
     ):
         self._impl = <nvcompManagerBase*>new ANSManager(
             uncomp_chunk_size,
@@ -165,8 +163,6 @@ cdef class _BitcompManager(_nvcompManager):
         size_t uncomp_chunk_size,
         nvcompType_t data_type,
         int bitcomp_algo,
-        user_stream,
-        const int device_id
     ):
         cdef opts = nvcompBatchedBitcompFormatOpts(bitcomp_algo, data_type)
         self._impl = <nvcompManagerBase*>new BitcompManager(
@@ -179,8 +175,6 @@ cdef class _CascadedManager(_nvcompManager):
     def __cinit__(
         self,
         _options,
-        user_stream,
-        const int device_id,
     ):
         self._impl = <nvcompManagerBase*>new CascadedManager(
             _options["chunk_size"],
@@ -193,8 +187,6 @@ cdef class _GdeflateManager(_nvcompManager):
         self,
         size_t uncomp_chunk_size,
         int algo,
-        user_stream,
-        const int device_id
     ):
         cdef opts = nvcompBatchedGdeflateOpts_t(algo)
         self._impl = <nvcompManagerBase*>new GdeflateManager(
@@ -208,8 +200,6 @@ cdef class _LZ4Manager(_nvcompManager):
         self,
         size_t uncomp_chunk_size,
         nvcompType_t data_type,
-        user_stream,
-        const int device_id,
     ):
         # TODO: Doesn't work with user specified streams passed down
         # from anywhere up. I'm not going to rabbit hole on it until
@@ -226,8 +216,6 @@ cdef class _SnappyManager(_nvcompManager):
     def __cinit__(
         self,
         size_t uncomp_chunk_size,
-        user_stream,
-        const int device_id,
     ):
         # TODO: Doesn't work with user specified streams passed down
         # from anywhere up. I'm not going to rabbit hole on it until

diff --git a/python/kvikio/kvikio/nvcomp.py b/python/kvikio/kvikio/nvcomp.py
@@ -1,4 +1,4 @@
-# Copyright (c) 2021-2023, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2021-2025, NVIDIA CORPORATION. All rights reserved.
 # See file LICENSE for terms.
 
 from enum import Enum
@@ -62,12 +62,10 @@ class nvCompManager:
     # Default options exist for every option type for every class that inherits
     # from nvCompManager, which takes advantage of the below property-setting
     # code.
-    stream: cp.cuda.Stream = cp.cuda.Stream()
     chunk_size: int = 1 << 16
     data_type: _lib.pyNvcompType_t = _lib.pyNvcompType_t.pyNVCOMP_TYPE_UCHAR
     # Some classes have this defined as type, some as data_type.
     type: _lib.pyNvcompType_t = _lib.pyNvcompType_t.pyNVCOMP_TYPE_UCHAR
-    device_id: int = 0
 
     # Bitcomp Defaults
     bitcomp_algo: int = 0
@@ -84,12 +82,6 @@ def __init__(self, kwargs):
 
         Special case: Convert data_type to a _lib.pyNvcompType_t
         """
-        # Special case: Throw error if stream or device_id are specified
-        if kwargs.get("stream") is not None:
-            raise NotImplementedError(
-                "stream argument not yet supported: " "Use the default argument"
-            )
-
         # data_type will be passed in as a python object. Convert it to
         # a C++ nvcompType_t here.
         if kwargs.get("data_type"):
@@ -221,13 +213,10 @@ def __init__(self, **kwargs):
         ----------
         chunk_size: int (optional)
             Defaults to 4096.
-        device_id: int (optional)
-            Specify which device_id on the node to use for allocation and compression.
-            Defaults to 0.
         """
         super().__init__(kwargs)
 
-        self._manager = _lib._ANSManager(self.chunk_size, self.stream, self.device_id)
+        self._manager = _lib._ANSManager(self.chunk_size)
 
 
 class BitcompManager(nvCompManager):
@@ -241,18 +230,13 @@ def __init__(self, **kwargs):
         ----------
         chunk_size: int (optional)
             Defaults to 4096.
-        device_id: int (optional)
-            Specify which device_id on the node to use
-            Defaults to 0.
         """
         super().__init__(kwargs)
 
         self._manager = _lib._BitcompManager(
             self.chunk_size,
             self.data_type.value,
             self.bitcomp_algo,
-            self.stream,
-            self.device_id,
         )
 
 
@@ -278,9 +262,6 @@ def __init__(self, **kwargs):
         use_bp: bool (optional)
             Enable Bitpacking, see [algorithms overview.md](
                 https://github.com/NVIDIA/nvcomp/blob/main/doc/algorithms_overview.md#bitpacking)  # noqa: E501
-        device_id: int (optional)
-            Specify which device_id on the node to use
-            Defaults to 0.
         """
         super().__init__(kwargs)
         default_options = {
@@ -304,9 +285,7 @@ def __init__(self, **kwargs):
             "num_deltas": self.num_deltas,
             "use_bp": self.use_bp,
         }
-        self._manager = _lib._CascadedManager(
-            default_options, self.stream, self.device_id
-        )
+        self._manager = _lib._CascadedManager(default_options)
 
 
 class GdeflateManager(nvCompManager):
@@ -322,18 +301,10 @@ def __init__(self, **kwargs):
         algo: int (optional)
             Integer in the range [0, 1, 2]. Only algorithm #0 is currently
             supported.
-        stream: cudaStream_t (optional)
-            Which CUDA stream to perform the operation on. Not currently
-            supported.
-        device_id: int (optional)
-            Specify which device_id on the node to use
-            Defaults to 0.
         """
         super().__init__(kwargs)
 
-        self._manager = _lib._GdeflateManager(
-            self.chunk_size, self.algo, self.stream, self.device_id
-        )
+        self._manager = _lib._GdeflateManager(self.chunk_size, self.algo)
 
 
 class LZ4Manager(nvCompManager):
@@ -354,17 +325,9 @@ def __init__(self, **kwargs):
         data_type: pyNVCOMP_TYPE (optional)
             The data type returned for decompression.
             Defaults to pyNVCOMP_TYPE.UCHAR
-        stream: cudaStream_t (optional)
-            Which CUDA stream to perform the operation on. Not currently
-            supported.
-        device_id: int (optional)
-            Specify which device_id on the node to use
-            Defaults to 0.
         """
         super().__init__(kwargs)
-        self._manager = _lib._LZ4Manager(
-            self.chunk_size, self.data_type.value, self.stream, self.device_id
-        )
+        self._manager = _lib._LZ4Manager(self.chunk_size, self.data_type.value)
 
 
 class SnappyManager(nvCompManager):
@@ -377,17 +340,9 @@ def __init__(self, **kwargs):
         Parameters
         ----------
         chunk_size: int (optional)
-        stream: cudaStream_t (optional)
-            Which CUDA stream to perform the operation on. Not currently
-            supported.
-        device_id: int (optional)
-            Specify which device_id on the node to use
-            Defaults to 0.
         """
         super().__init__(kwargs)
-        self._manager = _lib._SnappyManager(
-            self.chunk_size, self.stream, self.device_id
-        )
+        self._manager = _lib._SnappyManager(self.chunk_size)
 
 
 class ManagedDecompressionManager(nvCompManager):

diff --git a/python/kvikio/tests/test_nvcomp.py b/python/kvikio/tests/test_nvcomp.py
@@ -1,4 +1,4 @@
-# Copyright (c) 2021-2023, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2021-2025, NVIDIA CORPORATION. All rights reserved.
 # See file LICENSE for terms.
 
 import pytest
@@ -77,12 +77,11 @@ def test_round_trip_dtypes(manager, dtype):
     "inputs",
     [
         {},
-        {"chunk_size": 1 << 16, "device_id": 0},
         {
             "chunk_size": 1 << 16,
         },
         {
-            "device_id": 0,
+            "chunk_size": 1 << 16,
         },
     ],
 )
@@ -99,13 +98,13 @@ def test_ans_inputs(inputs):
     "inputs",
     [
         {},
-        {"data_type": np.uint8, "algo": 0, "device_id": 0},
-        {"data_type": np.uint8},
         {
+            "data_type": np.uint8,
             "algo": 0,
         },
+        {"data_type": np.uint8},
         {
-            "device_id": 0,
+            "algo": 0,
         },
     ],
 )
@@ -179,7 +178,6 @@ def test_bitcomp_algorithms(inputs, expected):
                 "num_deltas": 1,
                 "use_bp": True,
             },
-            "device_id": 0,
         },
     ],
 )
@@ -196,15 +194,15 @@ def test_cascaded_inputs(inputs):
     "inputs",
     [
         {},
-        {"chunk_size": 1 << 16, "algo": 0, "device_id": 0},
         {
             "chunk_size": 1 << 16,
+            "algo": 0,
         },
         {
-            "algo": 0,
+            "chunk_size": 1 << 16,
         },
         {
-            "device_id": 0,
+            "algo": 0,
         },
     ],
 )
@@ -253,15 +251,15 @@ def test_gdeflate_algorithms_not_implemented(inputs, expected):
     "inputs",
     [
         {},
-        {"chunk_size": 1 << 16, "data_type": np.uint8, "device_id": 0},
         {
             "chunk_size": 1 << 16,
+            "data_type": np.uint8,
         },
         {
-            "data_type": np.uint8,
+            "chunk_size": 1 << 16,
         },
         {
-            "device_id": 0,
+            "data_type": np.uint8,
         },
     ],
 )
@@ -278,11 +276,13 @@ def test_lz4_inputs(inputs):
     "inputs",
     [
         {},
-        {"chunk_size": 1 << 16, "device_id": 0},
         {
             "chunk_size": 1 << 16,
         },
-        {"device_id": 0},
+        {
+            "chunk_size": 1 << 16,
+        },
+        {},
     ],
 )
 def test_snappy_inputs(inputs):