Cython None guards (rapidsai#497)

Closes rapidsai#492

Authors:
  - Mads R. B. Kristensen (https://github.com/madsbk)

Approvers:
  - Tom Augspurger (https://github.com/TomAugspurger)
  - Kyle Edwards (https://github.com/KyleFromNVIDIA)

URL: rapidsai#497

Author: madsbk

python/rapidsmpf/rapidsmpf/CMakeLists.txt

@@ -3,8 +3,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # =================================================================================
 
-set(cython_modules config.pyx progress_thread.pyx rmm_resource_adaptor.pyx shuffler.pyx
-                   statistics.pyx
+set(cython_modules config.pyx cuda_stream.pyx progress_thread.pyx rmm_resource_adaptor.pyx
+                   shuffler.pyx statistics.pyx
 )
 
 rapids_cython_create_modules(

python/rapidsmpf/rapidsmpf/buffer/packed_data.pyx

@@ -46,7 +46,10 @@ cdef class PackedData:
 
     @classmethod
     def from_cudf_packed_columns(
-        cls, PackedColumns packed_columns, Stream stream, BufferResource br
+        cls,
+        PackedColumns packed_columns not None,
+        Stream stream not None,
+        BufferResource br not None,
     ):
         """
         Constructs a PackedData from CudfPackedColumns by taking the ownership of the

python/rapidsmpf/rapidsmpf/buffer/resource.pyx

@@ -61,7 +61,7 @@ cdef class BufferResource:
     """
     def __cinit__(
         self,
-        DeviceMemoryResource device_mr,
+        DeviceMemoryResource device_mr not None,
         memory_available = None,
         periodic_spill_check = 1e-3
     ):
@@ -177,7 +177,7 @@ cdef class LimitAvailableMemory:
     >>> mr = RmmResourceAdaptor(...)
     >>> memory_limiter = LimitAvailableMemory(mr, limit=1_000_000)
     """
-    def __init__(self, RmmResourceAdaptor mr, int64_t limit):
+    def __init__(self, RmmResourceAdaptor mr not None, int64_t limit):
         self._mr = mr  # Keep the mr alive.
         cdef cpp_RmmResourceAdaptor* handle = mr.get_handle()
         with nogil:

python/rapidsmpf/rapidsmpf/buffer/spill_manager.pyx

@@ -88,7 +88,7 @@ cdef class SpillManager:
             self._handle = NULL
 
     @classmethod
-    def _create(cls, BufferResource br):
+    def _create(cls, BufferResource br not None):
         """Construct a SpillManager associated the specified buffer resource.
 
         This shouldn't be used directly instead use `BufferResource.spill_manage)`.

python/rapidsmpf/rapidsmpf/communicator/communicator.pyx

@@ -56,7 +56,7 @@ cdef class Logger:
         """
         return cpp_verbosity_level(self._comm._handle)
 
-    def print(self, msg: str):
+    def print(self, str msg not None):
         """
         Logs a print message.
 
@@ -68,7 +68,7 @@ cdef class Logger:
         cdef string _msg = msg.encode()
         cpp_log(LOG_LEVEL.PRINT, self._comm._handle, move(_msg))
 
-    def warn(self, msg: str):
+    def warn(self, str msg not None):
         """
         Logs a warning message.
 
@@ -80,7 +80,7 @@ cdef class Logger:
         cdef string _msg = msg.encode()
         cpp_log(LOG_LEVEL.WARN, self._comm._handle, move(_msg))
 
-    def info(self, msg: str):
+    def info(self, str msg not None):
         """
         Logs an informational message.
 
@@ -92,7 +92,7 @@ cdef class Logger:
         cdef string _msg = msg.encode()
         cpp_log(LOG_LEVEL.INFO, self._comm._handle, move(_msg))
 
-    def debug(self, msg: str):
+    def debug(self, str msg not None):
         """
         Logs a debug message.
 
@@ -104,7 +104,7 @@ cdef class Logger:
         cdef string _msg = msg.encode()
         cpp_log(LOG_LEVEL.DEBUG, self._comm._handle, move(_msg))
 
-    def trace(self, msg: str):
+    def trace(self, str msg not None):
         """
         Logs a trace message.
 

python/rapidsmpf/rapidsmpf/communicator/mpi.pyx

@@ -15,7 +15,7 @@ cdef extern from "<rapidsmpf/communicator/mpi.hpp>" nogil:
         cpp_MPI_Communicator(libmpi.MPI_Comm comm, cpp_Options options) except +
 
 
-def new_communicator(Intracomm comm, Options options):
+def new_communicator(Intracomm comm not None, Options options not None):
     """
     Create a new RapidsMPF-MPI communicator based on an existing mpi4py communicator.
 

python/rapidsmpf/rapidsmpf/communicator/single.pyx

@@ -13,7 +13,7 @@ cdef extern from "<rapidsmpf/communicator/single.hpp>" nogil:
         cpp_Single_Communicator(cpp_Options options) except +
 
 
-def new_communicator(Options options):
+def new_communicator(Options options not None):
     """
     Create a new RapidsMPF single-process communicator.
 

python/rapidsmpf/rapidsmpf/communicator/ucxx.pyi

@@ -7,7 +7,10 @@ from rapidsmpf.communicator.communicator import Communicator
 from rapidsmpf.config import Options
 
 def new_communicator(
-    nranks: int, ucx_worker: UCXWorker, root_ucxx_address: UCXAddress, options: Options
+    nranks: int,
+    ucx_worker: UCXWorker | None,
+    root_ucxx_address: UCXAddress | None,
+    options: Options,
 ) -> Communicator: ...
 def get_root_ucxx_address(comm: Communicator) -> bytes: ...
 def barrier(comm: Communicator) -> None: ...

python/rapidsmpf/rapidsmpf/communicator/ucxx.pyx

@@ -83,7 +83,7 @@ def new_communicator(
     Rank nranks,
     UCXWorker ucx_worker,
     UCXAddress root_ucxx_address,
-    Options options,
+    Options options not None,
 ):
     """
     Create a new UCXX communicator with the given number of ranks.
@@ -121,7 +121,7 @@ def new_communicator(
     )
 
 
-def get_root_ucxx_address(Communicator comm):
+def get_root_ucxx_address(Communicator comm not None):
     """
     Get the address of the communicator's UCXX worker.
 
@@ -164,7 +164,7 @@ def get_root_ucxx_address(Communicator comm):
         assert host_port_pair  # Prevent "defined but unused" error
 
 
-def barrier(Communicator comm):
+def barrier(Communicator comm not None):
     """
     Execute a barrier on the UCXX communicator.
 

python/rapidsmpf/rapidsmpf/config.pyx

@@ -64,7 +64,7 @@ cdef class Options:
             ret = self._handle.insert_if_absent(move(opts))
         return ret
 
-    def get(self, str key, *, return_type, factory):
+    def get(self, str key not None, *, return_type, factory):
         """
         Retrieves a configuration option by key.
 
@@ -119,7 +119,7 @@ cdef class Options:
             return config_options_get.get_str(self, key, factory)
         return config_options_get.get_py_obj(self, key, factory)
 
-    def get_or_default(self, str key, *, default_value):
+    def get_or_default(self, str key not None, *, default_value):
         """
         Retrieve a configuration option by key, using a default value if not present.
 
@@ -235,7 +235,7 @@ cdef class Options:
         return <bytes>PyBytes_FromStringAndSize(<const char*>&vec[0], vec.size())
 
     @staticmethod
-    def deserialize(bytes serialized_buffer):
+    def deserialize(bytes serialized_buffer not None):
         """
         Deserialize a binary buffer into an Options object.