<commit_before>import os
from multiprocessing.managers import BaseProxy, MakeProxyType, AutoProxy, \
                                     DictProxy, SyncManager, Server
import multiprocessing as mp
import numpy as np
from shared_ndarray.shared_ndarray import SharedNDArray
import posix_ipc


class SharedMemoryTracker:
    "Manages one or more shared memory segments."

    def __init__(self, name, segment_names=[]):
        self.shared_memory_context_name = name
        self.segment_names = segment_names

    def register_segment(self, segment):
        print(f"Registering segment {segment.name} in pid {os.getpid()}")
        self.segment_names.append(segment.name)

    def destroy_segment(self, segment_name):
        print(f"Destroying segment {segment_name} in pid {os.getpid()}")
        self.segment_names.remove(segment_name)
        segment = posix_ipc.SharedMemory(segment_name)
        segment.close_fd()
        segment.unlink()

    def unlink(self):
        for segment_name in self.segment_names:
            print(f"Unlinking segment {segment_name} in pid {os.getpid()}")
            segment = posix_ipc.SharedMemory(segment_name)
            segment.close_fd()
            segment.unlink()
        self.segment_names[:] = []

    def __del__(self):
        print(f"somebody called {self.__class__.__name__}.__del__: {os.getpid()}")
        self.unlink()

    def __getstate__(self):
        return (self.shared_memory_context_name, self.segment_names)

    def __setstate__(self, state):
        self.__init__(*state)

    def ndarray(self, arr, dtype=None):
        "To be removed -- needs to be register-able to intro numpy support."
        try:
            shape = arr.shape
            dtype = arr.dtype
            shared_array = SharedNDArray.copy(arr)
        except AttributeError:
            shape = arr    # assume is a tuple
            dtype = dtype  # assume contains numpy.dtype
            shared_array = SharedNDArray(shape, dtype)
        self.register_segment(shared_array._shm)
        return shared_array


class AugmentedServer(Server):
    def __init__(self, *args, **kwargs):
        Server.__init__(self, *args, **kwargs)
        self.shared_memory_context = \
            SharedMemoryTracker(f"shmm_{self.address}_{os.getpid()}")
        print(f"AugmentedServer started by pid {os.getpid()}")

    def create(self, c, typeid, *args, **kwargs):
        # Unless set up as a shared proxy, don't make shared_memory_context
        # a standard part of kwargs.  This makes things easier for supplying
        # simple functions.
        if hasattr(self.registry[typeid][-1], "_shared_memory_proxy"):
            kwargs['shared_memory_context'] = self.shared_memory_context
        return Server.create(self, c, typeid, *args, **kwargs)

    def shutdown(self, c):
        self.shared_memory_context.unlink()
        return Server.shutdown(self, c)


class SharedMemoryManager(SyncManager):
    """Like SyncManager but uses AugmentedServer instead of Server.

    TODO: relocate/merge into managers submodule."""

    _Server = AugmentedServer

    def __init__(self, *args, **kwargs):
        # TODO: Remove after debugging satisfied
        SyncManager.__init__(self, *args, **kwargs)
        print(f"{self.__class__.__name__} created by pid {os.getpid()}")

    def __del__(self):
        # TODO: Remove after debugging satisfied
        print(f"{self.__class__.__name__} told to die by pid {os.getpid()}")
        pass

    def get_server(self):
        'Better than monkeypatching for now; merge into Server ultimately'
        if self._state.value != State.INITIAL:
            if self._state.value == State.STARTED:
                raise ProcessError("Already started server")
            elif self._state.value == State.SHUTDOWN:
                raise ProcessError("Manager has shut down")
            else:
                raise ProcessError(
                    "Unknown state {!r}".format(self._state.value))
        return AugmentedServer(self._registry, self._address,
                               self._authkey, self._serializer)


class SharedList(SharedNDArray):

    def __init__(self, iterable, *, shared_memory_context):
        self.shared_memory_context = shared_memory_context
        if isinstance(iterable, (list, tuple)):
            # Skip creation of short-lived duplicate list.
            shared_array = SharedNDArray.copy(np.array(iterable))
        else:
            shared_array = SharedNDArray.copy(np.array(list(iterable)))
        self.allocated_size = -1  # TODO: overallocate more than needed
        self.replace_held_shared_array(shared_array, destroy_old=False)

    def replace_held_shared_array(self, shared_array, destroy_old=True):
        if destroy_old:
            self.shared_memory_context.destroy_segment(self._shm.name)
        self.array = shared_array.array
        self._shm = shared_array._shm
        self._buf = shared_array._buf
        self.shared_memory_context.register_segment(self._shm)
        # Preserve it or the SharedNDArray kills it and all these mutation
        # operations will fail catastrophically.
        # But if it's preserved, then its refcount won't drop to 0 even
        # after the last Proxy reference is killed.
        # Ultimately, replace SharedNDArray's functionality.
        self._shared_array = shared_array

    def __del__(self):
        self._buf.close()

    def __contains__(self, value):
        return self.array.__contains__(value)

    def __getitem__(self, position):
        return self.array.__getitem__(position)

    def __setitem__(self, position, value):
        return self.array.__setitem__(position, value)

    def _getstate(self):
        return self.array.shape, self.array.dtype, self._shm.name

    def append(self, value):
        arr_of_one = np.array([value], dtype=self.array.dtype)
        shared_array = SharedNDArray.copy(
            np.concatenate([self.array, arr_of_one]))
        self.allocated_size = -1  # TODO: overallocate more than needed
        self.replace_held_shared_array(shared_array)

    def clear(self):
        shared_array = SharedNDArray.copy(np.array([], dtype=self.array.dtype))
        self.allocated_size = -1  # TODO: overallocate more than needed
        self.replace_held_shared_array(shared_array)
<commit_msg>Convert to use shared_memory instead of posix_ipc.<commit_after>import os
from multiprocessing.managers import BaseProxy, MakeProxyType, AutoProxy, \
                                     DictProxy, SyncManager, Server
import multiprocessing as mp
import numpy as np
from shared_ndarray.shared_ndarray import SharedNDArray
import shared_memory


class SharedMemoryTracker:
    "Manages one or more shared memory segments."

    def __init__(self, name, segment_names=[]):
        self.shared_memory_context_name = name
        self.segment_names = segment_names

    def register_segment(self, segment):
        print(f"Registering segment {segment.name} in pid {os.getpid()}")
        self.segment_names.append(segment.name)

    def destroy_segment(self, segment_name):
        print(f"Destroying segment {segment_name} in pid {os.getpid()}")
        self.segment_names.remove(segment_name)
        segment = shared_memory.SharedMemory(segment_name)
        segment.close_fd()
        segment.unlink()

    def unlink(self):
        for segment_name in self.segment_names:
            print(f"Unlinking segment {segment_name} in pid {os.getpid()}")
            segment = shared_memory.SharedMemory(segment_name)
            segment.close_fd()
            segment.unlink()
        self.segment_names[:] = []

    def __del__(self):
        print(f"somebody called {self.__class__.__name__}.__del__: {os.getpid()}")
        self.unlink()

    def __getstate__(self):
        return (self.shared_memory_context_name, self.segment_names)

    def __setstate__(self, state):
        self.__init__(*state)

    def ndarray(self, arr, dtype=None):
        "To be removed -- needs to be register-able to intro numpy support."
        try:
            shape = arr.shape
            dtype = arr.dtype
            shared_array = SharedNDArray.copy(arr)
        except AttributeError:
            shape = arr    # assume is a tuple
            dtype = dtype  # assume contains numpy.dtype
            shared_array = SharedNDArray(shape, dtype)
        self.register_segment(shared_array._shm)
        return shared_array


class AugmentedServer(Server):
    def __init__(self, *args, **kwargs):
        Server.__init__(self, *args, **kwargs)
        self.shared_memory_context = \
            SharedMemoryTracker(f"shmm_{self.address}_{os.getpid()}")
        print(f"AugmentedServer started by pid {os.getpid()}")

    def create(self, c, typeid, *args, **kwargs):
        # Unless set up as a shared proxy, don't make shared_memory_context
        # a standard part of kwargs.  This makes things easier for supplying
        # simple functions.
        if hasattr(self.registry[typeid][-1], "_shared_memory_proxy"):
            kwargs['shared_memory_context'] = self.shared_memory_context
        return Server.create(self, c, typeid, *args, **kwargs)

    def shutdown(self, c):
        self.shared_memory_context.unlink()
        return Server.shutdown(self, c)


class SharedMemoryManager(SyncManager):
    """Like SyncManager but uses AugmentedServer instead of Server.

    TODO: relocate/merge into managers submodule."""

    _Server = AugmentedServer

    def __init__(self, *args, **kwargs):
        # TODO: Remove after debugging satisfied
        SyncManager.__init__(self, *args, **kwargs)
        print(f"{self.__class__.__name__} created by pid {os.getpid()}")

    def __del__(self):
        # TODO: Remove after debugging satisfied
        print(f"{self.__class__.__name__} told to die by pid {os.getpid()}")
        pass

    def get_server(self):
        'Better than monkeypatching for now; merge into Server ultimately'
        if self._state.value != State.INITIAL:
            if self._state.value == State.STARTED:
                raise ProcessError("Already started server")
            elif self._state.value == State.SHUTDOWN:
                raise ProcessError("Manager has shut down")
            else:
                raise ProcessError(
                    "Unknown state {!r}".format(self._state.value))
        return AugmentedServer(self._registry, self._address,
                               self._authkey, self._serializer)


class SharedList(SharedNDArray):

    def __init__(self, iterable, *, shared_memory_context):
        self.shared_memory_context = shared_memory_context
        if isinstance(iterable, (list, tuple)):
            # Skip creation of short-lived duplicate list.
            shared_array = SharedNDArray.copy(np.array(iterable))
        else:
            shared_array = SharedNDArray.copy(np.array(list(iterable)))
        self.allocated_size = -1  # TODO: overallocate more than needed
        self.replace_held_shared_array(shared_array, destroy_old=False)

    def replace_held_shared_array(self, shared_array, destroy_old=True):
        if destroy_old:
            self.shared_memory_context.destroy_segment(self._shm.name)
        self.array = shared_array.array
        self._shm = shared_array._shm
        self.shared_memory_context.register_segment(self._shm)
        # Preserve it or the SharedNDArray kills it and all these mutation
        # operations will fail catastrophically.
        # But if it's preserved, then its refcount won't drop to 0 even
        # after the last Proxy reference is killed.
        # Ultimately, replace SharedNDArray's functionality.
        self._shared_array = shared_array

    def __del__(self):
        self._shm.close()

    def __contains__(self, value):
        return self.array.__contains__(value)

    def __getitem__(self, position):
        return self.array.__getitem__(position)

    def __setitem__(self, position, value):
        return self.array.__setitem__(position, value)

    def _getstate(self):
        return self.array.shape, self.array.dtype, self._shm.name

    def append(self, value):
        arr_of_one = np.array([value], dtype=self.array.dtype)
        shared_array = SharedNDArray.copy(
            np.concatenate([self.array, arr_of_one]))
        self.allocated_size = -1  # TODO: overallocate more than needed
        self.replace_held_shared_array(shared_array)

    def clear(self):
        shared_array = SharedNDArray.copy(np.array([], dtype=self.array.dtype))
        self.allocated_size = -1  # TODO: overallocate more than needed
        self.replace_held_shared_array(shared_array)

{'fraction_non_alphanumeric': 0.09340615847383836, 'fraction_numerical': 0.0007915776141850709, 'mean_word_length': 2.7545319465081723, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3253269', 'n_tokens_mistral': 3629, 'n_tokens_neox': 3516, 'n_words': 874}