# Path: taskloaf/shmem.py
# @contextlib.contextmanager
# def alloc_shmem(size, filepath):
#     assert not os.path.exists(filepath)
#     try:
#         init_shmem_file(filepath, size)
#         logger.info(f"creating shm {filepath}")
#         yield filepath
#     finally:
#         try:
#             logger.info(f"deleting shm {filepath}")
#             os.remove(filepath)
#         except FileNotFoundError:
#             logger.info(
#                 f"failed to delete shm {filepath} because it no longer exists"
#             )
#             pass
# 
# class Shmem:
#     def __init__(self, filepath, track_refs=False):
#         self.filepath = filepath
#         self.track_refs = track_refs
# 
#     def __enter__(self):
#         self.file = open(self.filepath, "r+b")
#         self.size = get_size_from_fd(self.file.fileno())
#         self.mmap = mmap_full_file(self.file.fileno())
# 
#         if not self.track_refs:
#             # mmap tracks references internally. This is annoying when trying
#             # to delete the mmap.  This code is some crooked trickery using
#             # ctypes and grabbing the pointer from a numpy buffer to create a
#             # memoryview from the mmap without mmap knowing about it so that
#             # the mmap can be closed without tracking its references (this WILL
#             # cause seg faults if there are existing references to the mmap
#             # segment when it's deleted)
#             # When using Shmem blocks through the allocator system, this is
#             # fine since then taskloaf performs its own memory tracking
#             temp_np = np.frombuffer(self.mmap, dtype=np.uint8)
#             ptr = temp_np.ctypes.data
#             del temp_np
#             ptrc = ctypes.cast(ptr, ctypes.POINTER(ctypes.c_byte))
#             new_array = np.ctypeslib.as_array(ptrc, shape=(self.size,))
#             self.mem = memoryview(new_array.data.cast("B"))
#         else:
#             self.mem = memoryview(self.mmap)
# 
#         return self
# 
#     def __exit__(self, exc_type, exc_value, traceback):
#         self.file.close()
#         self.mmap.close()
# 
# Path: taskloaf/timer.py
# class Timer(object):
#     def __init__(self, output_fnc=print):
#         if output_fnc is None:
#             output_fnc = ignore_args
#         self.output_fnc = output_fnc
#         self.restart()
# 
#     def restart(self):
#         self.start = get_time()
# 
#     def report(self, name, should_restart=True):
#         text = name + " took " + str(get_time() - self.start)
#         self.output_fnc(text)
#         if should_restart:
#             self.restart()


import os
import time
import numpy as np
import multiprocessing
from taskloaf.shmem import alloc_shmem, Shmem
from taskloaf.timer import Timer
def shmem_zeros(filepath):
    with Shmem(filepath) as sm:
        np.frombuffer(sm.mem)[:] = 0

def test_shmem_edit():
    A = np.random.rand(100)
    with alloc_shmem(A.nbytes, 'test') as filepath:
        with Shmem(filepath) as sm:
            sm.mem[:] = A.data.cast('B')
            out = multiprocessing.Pool(1).map(shmem_zeros, [sm.filepath])[0]
            np.testing.assert_almost_equal(np.frombuffer(sm.mem), 0.0)

"""
There's some complexity in this benchmark. The performance of allocating mmap
is pretty much identical to a alloc+memcpy if you can write the relevant data
to the mmap-ed region before mmap-ing it. However, if you fill the memory
through the pointer returned by mmap, mmap will zero the memory region first.

An approach that can avoid this problem is to write directly to the file
associated with the mmap-ed region. Because the memory is not actually
allocated until it's used, this will avoid the time cost of mmap zeroing the
memory

Also, note that the lazy nature of python garbage collection means that munmap
is called at unexpected times and cause some problems with benchmarking.

LATER: i THINK THAT ALL THESE WOES WERE CAUSED ALMOST ENTIRELY BY THE LACK OF
USING HUGE PAGES
"""
def benchmark_shmem():

# Path: taskloaf/shmem.py
# @contextlib.contextmanager
# def alloc_shmem(size, filepath):
#     assert not os.path.exists(filepath)
#     try:
#         init_shmem_file(filepath, size)
#         logger.info(f"creating shm {filepath}")
#         yield filepath
#     finally:
#         try:
#             logger.info(f"deleting shm {filepath}")
#             os.remove(filepath)
#         except FileNotFoundError:
#             logger.info(
#                 f"failed to delete shm {filepath} because it no longer exists"
#             )
#             pass
# 
# class Shmem:
#     def __init__(self, filepath, track_refs=False):
#         self.filepath = filepath
#         self.track_refs = track_refs
# 
#     def __enter__(self):
#         self.file = open(self.filepath, "r+b")
#         self.size = get_size_from_fd(self.file.fileno())
#         self.mmap = mmap_full_file(self.file.fileno())
# 
#         if not self.track_refs:
#             # mmap tracks references internally. This is annoying when trying
#             # to delete the mmap.  This code is some crooked trickery using
#             # ctypes and grabbing the pointer from a numpy buffer to create a
#             # memoryview from the mmap without mmap knowing about it so that
#             # the mmap can be closed without tracking its references (this WILL
#             # cause seg faults if there are existing references to the mmap
#             # segment when it's deleted)
#             # When using Shmem blocks through the allocator system, this is
#             # fine since then taskloaf performs its own memory tracking
#             temp_np = np.frombuffer(self.mmap, dtype=np.uint8)
#             ptr = temp_np.ctypes.data
#             del temp_np
#             ptrc = ctypes.cast(ptr, ctypes.POINTER(ctypes.c_byte))
#             new_array = np.ctypeslib.as_array(ptrc, shape=(self.size,))
#             self.mem = memoryview(new_array.data.cast("B"))
#         else:
#             self.mem = memoryview(self.mmap)
# 
#         return self
# 
#     def __exit__(self, exc_type, exc_value, traceback):
#         self.file.close()
#         self.mmap.close()
# 
# Path: taskloaf/timer.py
# class Timer(object):
#     def __init__(self, output_fnc=print):
#         if output_fnc is None:
#             output_fnc = ignore_args
#         self.output_fnc = output_fnc
#         self.restart()
# 
#     def restart(self):
#         self.start = get_time()
# 
#     def report(self, name, should_restart=True):
#         text = name + " took " + str(get_time() - self.start)
#         self.output_fnc(text)
#         if should_restart:
#             self.restart()


# Path: tests/test_shmem.py
import os
import time
import numpy as np
import multiprocessing
from taskloaf.shmem import alloc_shmem, Shmem
from taskloaf.timer import Timer

def shmem_zeros(filepath):
    with Shmem(filepath) as sm:
        np.frombuffer(sm.mem)[:] = 0

def test_shmem_edit():
    A = np.random.rand(100)
    with alloc_shmem(A.nbytes, 'test') as filepath:
        with Shmem(filepath) as sm:
            sm.mem[:] = A.data.cast('B')
            out = multiprocessing.Pool(1).map(shmem_zeros, [sm.filepath])[0]
            np.testing.assert_almost_equal(np.frombuffer(sm.mem), 0.0)

"""
There's some complexity in this benchmark. The performance of allocating mmap
is pretty much identical to a alloc+memcpy if you can write the relevant data
to the mmap-ed region before mmap-ing it. However, if you fill the memory
through the pointer returned by mmap, mmap will zero the memory region first.

An approach that can avoid this problem is to write directly to the file
associated with the mmap-ed region. Because the memory is not actually
allocated until it's used, this will avoid the time cost of mmap zeroing the
memory

Also, note that the lazy nature of python garbage collection means that munmap
is called at unexpected times and cause some problems with benchmarking.

LATER: i THINK THAT ALL THESE WOES WERE CAUSED ALMOST ENTIRELY BY THE LACK OF
USING HUGE PAGES
"""
def benchmark_shmem():

    t = Timer()