Enable RMM Debug Logging via Python (#1339)

Enables RMM debug logging in Python and exposes controls for the level of logging. Fixes #1336.

Authors:
  - Mark Harris (https://github.com/harrism)
  - Lawrence Mitchell (https://github.com/wence-)

Approvers:
  - Lawrence Mitchell (https://github.com/wence-)
  - Vyas Ramasubramani (https://github.com/vyasr)

URL: #1339

README.md

@@ -581,9 +581,12 @@ of more detailed logging. The default is `INFO`. Available levels are `TRACE`, `
 
 The log relies on the [spdlog](https://github.com/gabime/spdlog.git) library.
 
-Note that to see logging below the `INFO` level, the C++ application must also call
-`rmm::logger().set_level()`, e.g. to enable all levels of logging down to `TRACE`, call
-`rmm::logger().set_level(spdlog::level::trace)` (and compile with `-DRMM_LOGGING_LEVEL=TRACE`).
+Note that to see logging below the `INFO` level, the application must also set the logging level at
+run time. C++ applications must must call `rmm::logger().set_level()`, for example to enable all
+levels of logging down to `TRACE`, call `rmm::logger().set_level(spdlog::level::trace)` (and compile
+librmm with `-DRMM_LOGGING_LEVEL=TRACE`). Python applications must call `rmm.set_logging_level()`,
+for example to enable all levels of logging down to `TRACE`, call `rmm.set_logging_level("trace")`
+(and compile the RMM Python module with `-DRMM_LOGGING_LEVEL=TRACE`).
 
 Note that debug logging is different from the CSV memory allocation logging provided by
 `rmm::mr::logging_resource_adapter`. The latter is for logging a history of allocation /

python/CMakeLists.txt

@@ -56,5 +56,7 @@ endif()
 include(rapids-cython)
 rapids_cython_init()
 
+add_compile_definitions("SPDLOG_ACTIVE_LEVEL=SPDLOG_LEVEL_${RMM_LOGGING_LEVEL}")
+
 add_subdirectory(rmm/_cuda)
 add_subdirectory(rmm/_lib)

python/rmm/__init__.py

@@ -14,6 +14,15 @@
 
 from rmm import mr
 from rmm._lib.device_buffer import DeviceBuffer
+from rmm._lib.logger import (
+    flush_logger,
+    get_flush_level,
+    get_logging_level,
+    logging_level,
+    set_flush_level,
+    set_logging_level,
+    should_log,
+)
 from rmm.mr import disable_logging, enable_logging, get_log_filenames
 from rmm.rmm import (
     RMMError,
@@ -25,14 +34,21 @@
 
 __all__ = [
     "DeviceBuffer",
-    "RMMError",
     "disable_logging",
+    "RMMError",
     "enable_logging",
+    "flush_logger",
+    "get_flush_level",
     "get_log_filenames",
+    "get_logging_level",
     "is_initialized",
+    "logging_level",
     "mr",
     "register_reinitialize_hook",
     "reinitialize",
+    "set_flush_level",
+    "set_logging_level",
+    "should_log",
     "unregister_reinitialize_hook",
 ]
 

python/rmm/_lib/CMakeLists.txt

@@ -12,7 +12,7 @@
 # the License.
 # =============================================================================
 
-set(cython_sources device_buffer.pyx lib.pyx memory_resource.pyx cuda_stream.pyx
+set(cython_sources device_buffer.pyx lib.pyx logger.pyx memory_resource.pyx cuda_stream.pyx
                    torch_allocator.pyx)
 set(linked_libraries rmm::rmm)
 

python/rmm/_lib/logger.pyx

@@ -0,0 +1,260 @@
+# Copyright (c) 2023, NVIDIA CORPORATION.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import warnings
+
+from libcpp cimport bool
+
+
+cdef extern from "spdlog/common.h" namespace "spdlog::level" nogil:
+    cpdef enum logging_level "spdlog::level::level_enum":
+        """
+        The debug logging level for RMM.
+
+        Debug logging prints messages to a log file. See
+        `Debug Logging <https://github.com/rapidsai/rmm#debug-logging>`_
+        for more information.
+
+        Valid levels, in decreasing order of verbosity, are TRACE, DEBUG,
+        INFO, WARN, ERR, CRITICAL, and OFF. Default is INFO.
+
+        Examples
+        --------
+        >>> import rmm
+        >>> rmm.logging_level.DEBUG
+        <logging_level.DEBUG: 1>
+        >>> rmm.logging_level.DEBUG.value
+        1
+        >>> rmm.logging_level.DEBUG.name
+        'DEBUG'
+
+        See Also
+        --------
+        set_logging_level : Set the debug logging level
+        get_logging_level : Get the current debug logging level
+        """
+        TRACE "spdlog::level::trace"
+        DEBUG "spdlog::level::debug"
+        INFO "spdlog::level::info"
+        WARN "spdlog::level::warn"
+        ERR "spdlog::level::err"
+        CRITICAL "spdlog::level::critical"
+        OFF "spdlog::level::off"
+
+
+cdef extern from "spdlog/spdlog.h" namespace "spdlog" nogil:
+    cdef cppclass spdlog_logger "spdlog::logger":
+        spdlog_logger() except +
+        void set_level(logging_level level)
+        logging_level level()
+        void flush() except +
+        void flush_on(logging_level level)
+        logging_level flush_level()
+        bool should_log(logging_level msg_level)
+
+
+cdef extern from "rmm/logger.hpp" namespace "rmm" nogil:
+    cdef spdlog_logger& logger() except +
+
+
+def _validate_level_type(level):
+    if not isinstance(level, logging_level):
+        raise TypeError("level must be an instance of the logging_level enum")
+
+
+def should_log(level):
+    """
+    Check if a message at the given level would be logged.
+
+    A message at the given level would be logged if the current debug logging
+    level is set to a level that is at least as verbose than the given level,
+    *and* the RMM module is compiled for a logging level at least as verbose.
+    If these conditions are not both met, this function will return false.
+
+    Debug logging prints messages to a log file. See
+    `Debug Logging <https://github.com/rapidsai/rmm#debug-logging>`_
+    for more information.
+
+    Parameters
+    ----------
+    level : logging_level
+        The debug logging level. Valid values are instances of the
+        ``logging_level`` enum.
+
+    Returns
+    -------
+    should_log : bool
+        True if a message at the given level would be logged, False otherwise.
+
+    Raises
+    ------
+    TypeError
+        If the logging level is not an instance of the ``logging_level`` enum.
+    """
+    _validate_level_type(level)
+    return logger().should_log(level)
+
+
+def set_logging_level(level):
+    """
+    Set the debug logging level.
+
+    Debug logging prints messages to a log file. See
+    `Debug Logging <https://github.com/rapidsai/rmm#debug-logging>`_
+    for more information.
+
+    Parameters
+    ----------
+    level : logging_level
+        The debug logging level. Valid values are instances of the
+        ``logging_level`` enum.
+
+    Raises
+    ------
+    TypeError
+        If the logging level is not an instance of the ``logging_level`` enum.
+
+    See Also
+    --------
+    get_logging_level : Get the current debug logging level.
+
+    Examples
+    --------
+    >>> import rmm
+    >>> rmm.set_logging_level(rmm.logging_level.WARN) # set logging level to warn
+    """
+    _validate_level_type(level)
+    logger().set_level(level)
+
+    if not should_log(level):
+        warnings.warn(f"RMM will not log logging_level.{level.name}. This "
+                      "may be because the C++ library is compiled for a "
+                      "less-verbose logging level.")
+
+
+def get_logging_level():
+    """
+    Get the current debug logging level.
+
+    Debug logging prints messages to a log file. See
+    `Debug Logging <https://github.com/rapidsai/rmm#debug-logging>`_
+    for more information.
+
+    Returns
+    -------
+    level : logging_level
+        The current debug logging level, an instance of the ``logging_level``
+        enum.
+
+    See Also
+    --------
+    set_logging_level : Set the debug logging level.
+
+    Examples
+    --------
+    >>> import rmm
+    >>> rmm.get_logging_level() # get current logging level
+    <logging_level.INFO: 2>
+    """
+    return logging_level(logger().level())
+
+
+def flush_logger():
+    """
+    Flush the debug logger. This will cause any buffered log messages to
+    be written to the log file.
+
+    Debug logging prints messages to a log file. See
+    `Debug Logging <https://github.com/rapidsai/rmm#debug-logging>`_
+    for more information.
+
+    See Also
+    --------
+    set_flush_level : Set the flush level for the debug logger.
+    get_flush_level : Get the current debug logging flush level.
+
+    Examples
+    --------
+    >>> import rmm
+    >>> rmm.flush_logger() # flush the logger
+    """
+    logger().flush()
+
+
+def set_flush_level(level):
+    """
+    Set the flush level for the debug logger. Messages of this level or higher
+    will automatically flush to the file.
+
+    Debug logging prints messages to a log file. See
+    `Debug Logging <https://github.com/rapidsai/rmm#debug-logging>`_
+    for more information.
+
+    Parameters
+    ----------
+    level : logging_level
+        The debug logging level. Valid values are instances of the
+        ``logging_level`` enum.
+
+    Raises
+    ------
+    TypeError
+        If the logging level is not an instance of the ``logging_level`` enum.
+
+    See Also
+    --------
+    get_flush_level : Get the current debug logging flush level.
+    flush_logger : Flush the logger.
+
+    Examples
+    --------
+    >>> import rmm
+    >>> rmm.flush_on(rmm.logging_level.WARN) # set flush level to warn
+    """
+    _validate_level_type(level)
+    logger().flush_on(level)
+
+    if not should_log(level):
+        warnings.warn(f"RMM will not log logging_level.{level.name}. This "
+                      "may be because the C++ library is compiled for a "
+                      "less-verbose logging level.")
+
+
+def get_flush_level():
+    """
+    Get the current debug logging flush level for the RMM logger. Messages of
+    this level or higher will automatically flush to the file.
+
+    Debug logging prints messages to a log file. See
+    `Debug Logging <https://github.com/rapidsai/rmm#debug-logging>`_
+    for more information.
+
+    Returns
+    -------
+    logging_level
+        The current flush level, an instance of the ``logging_level``
+        enum.
+
+    See Also
+    --------
+    set_flush_level : Set the flush level for the logger.
+    flush_logger : Flush the logger.
+
+    Examples
+    --------
+    >>> import rmm
+    >>> rmm.flush_level() # get current flush level
+    <logging_level.INFO: 2>
+    """
+    return logging_level(logger().flush_level())

python/rmm/tests/test_rmm.py

@@ -16,6 +16,7 @@
 import gc
 import os
 import pickle
+import warnings
 from itertools import product
 
 import numpy as np
@@ -942,3 +943,35 @@ def test_rmm_device_buffer_copy(cuda_ary, make_copy):
     result = db_copy.copy_to_host()
 
     np.testing.assert_equal(expected, result)
+
+
+@pytest.mark.parametrize("level", rmm.logging_level)
+def test_valid_logging_level(level):
+    with warnings.catch_warnings():
+        warnings.filterwarnings(
+            "ignore", message="RMM will not log logging_level.TRACE."
+        )
+        warnings.filterwarnings(
+            "ignore", message="RMM will not log logging_level.DEBUG."
+        )
+        rmm.set_logging_level(level)
+        assert rmm.get_logging_level() == level
+        rmm.set_logging_level(rmm.logging_level.INFO)  # reset to default
+
+        rmm.set_flush_level(level)
+        assert rmm.get_flush_level() == level
+        rmm.set_flush_level(rmm.logging_level.INFO)  # reset to default
+
+        rmm.should_log(level)
+
+
+@pytest.mark.parametrize(
+    "level", ["INFO", 3, "invalid", 100, None, 1.2345, [1, 2, 3]]
+)
+def test_invalid_logging_level(level):
+    with pytest.raises(TypeError):
+        rmm.set_logging_level(level)
+    with pytest.raises(TypeError):
+        rmm.set_flush_level(level)
+    with pytest.raises(TypeError):
+        rmm.should_log(level)