def enable_logging(self):
"""
Enable RMM logging. RMM creates a CSV output file derived from
provided file name that looks like: log_file_prefix + ".devX", where
X is the GPU number.
"""
rmm.enable_logging(log_file_name=self._log_file_prefix)
def setup(self, worker=None):
if self.async_alloc:
import rmm
rmm.mr.set_current_device_resource(rmm.mr.CudaAsyncMemoryResource())
if self.logging:
rmm.enable_logging(
log_file_name=get_rmm_log_file_name(
worker, self.logging, self.log_directory
)
)
elif self.initial_pool_size is not None or self.managed_memory:
import rmm
pool_allocator = False if self.initial_pool_size is None else True
rmm.reinitialize(
pool_allocator=pool_allocator,
managed_memory=self.managed_memory,
initial_pool_size=self.initial_pool_size,
maximum_pool_size=self.maximum_pool_size,
logging=self.logging,
log_file_name=get_rmm_log_file_name(
worker, self.logging, self.log_directory
),
)
def test_rmm_enable_disable_logging(dtype, nelem, alloc, tmpdir):
suffix = ".csv"
base_name = str(tmpdir.join("rmm_log.csv"))
rmm.enable_logging(log_file_name=base_name)
print(rmm.mr.get_per_device_resource(0))
array_tester(dtype, nelem, alloc)
rmm.mr._flush_logs()
# Need to open separately because the device ID is appended to filename
fname = base_name[:-len(suffix)] + ".dev0" + suffix
try:
with open(fname, "rb") as f:
csv = f.read()
assert csv.find(b"Time,Action,Pointer,Size,Stream") >= 0
finally:
os.remove(fname)
rmm.disable_logging()