diff --git a/python/rmm/__init__.py b/python/rmm/__init__.py index 8c753828e..bdb7f7c56 100644 --- a/python/rmm/__init__.py +++ b/python/rmm/__init__.py @@ -15,7 +15,8 @@ from rmm import mr from rmm._lib.device_buffer import DeviceBuffer -from rmm.mr import disable_logging, enable_logging +from rmm._version import get_versions +from rmm.mr import disable_logging, enable_logging, get_log_filenames from rmm.rmm import ( RMMError, RMMNumbaManager, @@ -25,6 +26,4 @@ rmm_cupy_allocator, ) -from rmm._version import get_versions - __version__ = get_versions()["version"] diff --git a/python/rmm/_lib/memory_resource.pyx b/python/rmm/_lib/memory_resource.pyx index 0c8764ce6..fb06b53f5 100644 --- a/python/rmm/_lib/memory_resource.pyx +++ b/python/rmm/_lib/memory_resource.pyx @@ -382,17 +382,18 @@ cdef class LoggingResourceAdaptor(UpstreamResourceAdaptor): if log_file_name is None: log_file_name = os.getenv("RMM_LOG_FILE") if not log_file_name: - raise TypeError( + raise ValueError( "RMM log file must be specified either using " "log_file_name= argument or RMM_LOG_FILE " "environment variable" ) + # Append the device ID before the file extension log_file_name = _append_id( log_file_name, getDevice() ) - - _log_file_name = log_file_name + log_file_name = os.path.abspath(log_file_name) + self._log_file_name = log_file_name self.c_obj.reset( new logging_resource_adaptor[device_memory_resource]( @@ -627,7 +628,27 @@ cpdef _flush_logs(): def enable_logging(log_file_name=None): """ - Enable logging of run-time events. + Enable logging of run-time events for all devices. + + Parameters + ---------- + log_file_name: str, optional + Name of the log file. If not specified, the environment variable + RMM_LOG_FILE is used. A ValueError is thrown if neither is available. + A separate log file is produced for each device, + and the suffix `".dev{id}"` is automatically added to the log file + name. + + Notes + ----- + Note that if you use the environment variable CUDA_VISIBLE_DEVICES + with logging enabled, the suffix may not be what you expect. For + example, if you set CUDA_VISIBLE_DEVICES=1, the log file produced + will still have suffix `0`. Similarly, if you set + CUDA_VISIBLE_DEVICES=1,0 and use devices 0 and 1, the log file + with suffix `0` will correspond to the GPU with device ID `1`. + Use `rmm.get_log_filenames()` to get the log file names + corresponding to each device. """ global _per_device_mrs @@ -651,3 +672,26 @@ def disable_logging(): for i, each_mr in _per_device_mrs.items(): if isinstance(each_mr, LoggingResourceAdaptor): set_per_device_resource(i, each_mr.get_upstream()) + + +def get_log_filenames(): + """ + Returns the log filename (or `None` if not writing logs) + for each device in use. + + Examples + -------- + >>> import rmm + >>> rmm.reinitialize(devices=[0, 1], logging=True, log_file_name="rmm.log") + >>> rmm.get_log_filenames() + {0: '/home/user/workspace/rapids/rmm/python/rmm.dev0.log', + 1: '/home/user/workspace/rapids/rmm/python/rmm.dev1.log'} + """ + global _per_device_mrs + + return { + i: each_mr.get_file_name() + if isinstance(each_mr, LoggingResourceAdaptor) + else None + for i, each_mr in _per_device_mrs.items() + } diff --git a/python/rmm/mr.py b/python/rmm/mr.py index 44be50318..15cec90b0 100644 --- a/python/rmm/mr.py +++ b/python/rmm/mr.py @@ -14,6 +14,7 @@ enable_logging, get_current_device_resource, get_current_device_resource_type, + get_log_filenames, get_per_device_resource, get_per_device_resource_type, is_initialized, @@ -40,5 +41,6 @@ "get_current_device_resource", "get_per_device_resource_type", "get_current_device_resource_type", + "get_log_filenames", "is_initialized", ] diff --git a/python/rmm/rmm.py b/python/rmm/rmm.py index 3f79debf4..9c57c6ac6 100644 --- a/python/rmm/rmm.py +++ b/python/rmm/rmm.py @@ -65,7 +65,21 @@ def reinitialize( This has significant performance impact. log_file_name : str Name of the log file. If not specified, the environment variable - RMM_LOG_FILE is used. A TypeError is thrown if neither is available. + RMM_LOG_FILE is used. A ValueError is thrown if neither is available. + A separate log file is produced for each device, + and the suffix `".dev{id}"` is automatically added to the log file + name. + + Notes + ----- + Note that if you use the environment variable CUDA_VISIBLE_DEVICES + with logging enabled, the suffix may not be what you expect. For + example, if you set CUDA_VISIBLE_DEVICES=1, the log file produced + will still have suffix `0`. Similarly, if you set + CUDA_VISIBLE_DEVICES=1,0 and use devices 0 and 1, the log file + with suffix `0` will correspond to the GPU with device ID `1`. + Use `rmm.get_log_filenames()` to get the log file names + corresponding to each device. """ rmm.mr._initialize( pool_allocator=pool_allocator,