<commit_before># encoding: utf-8
# Copyright (c) 2008-2014, IPython Development Team and Enthought, Inc.

__docformat__ = "restructuredtext en"

import uuid
from distarray.externals import six
import collections

from IPython.parallel import Client
import numpy

from distarray.client import DistArray


class Context(object):

    def __init__(self, view=None, targets=None):
        if view is None:
            self.client = Client()
            self.view = self.client[:]
        else:
            self.view = view
            self.client = view.client

        all_targets = self.view.targets
        if targets is None:
            self.targets = all_targets
        else:
            self.targets = []
            for target in targets:
                if target not in all_targets:
                    raise ValueError("Engine with id %r not registered" % target)
                else:
                    self.targets.append(target)

        # FIXME: IPython bug #4296: This doesn't work under Python 3
        #with self.view.sync_imports():
        #    import distarray
        self.view.execute("import distarray.local; import distarray.mpiutils;"
                          " import numpy")

        self._setup_key_records()
        self._make_intracomm()
        self._set_engine_rank_mapping()

    def __del__(self):
        """ Clean up keys we have put on the engines. """
        self._cleanup_keys()
        self._clear_comm_key()

    def _clear_comm_key(self):
        """ Delete our internal _comm_key from the engines. """
        # This could perhaps be called in the __exit__ of a context manager
        # for the Context. This _comm_key is not otherwise cleaned up by
        # the methods in this class.
        if hasattr(self, '_comm_key'):
            self._delete_engine_key(self._comm_key)
            self._comm_key = None

    def _set_engine_rank_mapping(self):
        # The MPI intracomm referred to by self._comm_key may have a different
        # mapping between IPython engines and MPI ranks than COMM_PRIVATE.  We
        # reorder self.targets so self.targets[i] is the IPython engine ID that
        # corresponds to MPI intracomm rank i.
        rank = self._generate_key()
        self.view.execute(
                '%s = %s.Get_rank()' % (rank, self._comm_key),
                block=True, targets=self.targets)

        # mapping target -> rank, rank -> target.
        target_to_rank = self.view.pull(rank, targets=self.targets).get_dict()
        rank_to_target = {v: k for (k, v) in target_to_rank.items()}
        # Probably wont need this...
        self.delete_key(rank)

        # ensure consistency
        assert set(self.targets) == set(target_to_rank)
        assert set(range(len(self.targets))) == set(rank_to_target)

        # reorder self.targets so that the targets are in MPI rank order for
        # the intracomm.
        self.targets = [rank_to_target[i] for i in range(len(rank_to_target))]

    def _make_intracomm(self):
        def get_rank():
            from distarray.mpiutils import COMM_PRIVATE
            return COMM_PRIVATE.Get_rank()

        # get a mapping of IPython engine ID to MPI rank
        rank_map = self.view.apply_async(get_rank).get_dict()
        ranks = [ rank_map[engine] for engine in self.targets ]

        # self.view's engines must encompass all ranks in the MPI communicator,
        # i.e., everything in rank_map.values().
        def get_size():
            from distarray.mpiutils import COMM_PRIVATE
            return COMM_PRIVATE.Get_size()

        comm_size = self.view.apply_async(get_size).get()[0]
        if set(rank_map.values()) != set(range(comm_size)):
            raise ValueError('Engines in view must encompass all MPI ranks.')

        # create a new communicator with the subset of engines note that
        # MPI_Comm_create must be called on all engines, not just those
        # involved in the new communicator.
        # This key is not added to key_records so it is not wrongly removed.
        self._comm_key = self._generate_key_name()
        self.view.execute(
            '%s = distarray.mpiutils.create_comm_with_list(%s)' % (self._comm_key, ranks),
            block=True
        )

    # Keeping track of the keys that we create on each engine.
    # This allows us to properly clean up after ourselves.
    # We store these as a dict, with the key as the key.
    # The value is the list of engines where the key has been pushed.
    # We can also find keys on the engines that are *not* tracked.
    # These should be kept track of better, but we can also delete them.

    def _setup_key_records(self):
        """ Create a dictionary, keyed by key name, with values as the
        set of targets the key is sent to.
        This can be used to properly clean up.
        """
        self.key_records = {}

    def _generate_key_name(self):
        """ Generate a unique name for a key. """
        uid = uuid.uuid4()
        key = '__distarray_%s' % uid.hex
        return key

    def _generate_key(self, targets=None):
        """ Generate a key name, and record that it will be present
        on the specified target engines.
        """
        if targets is None:
            targets = self.targets
        key = self._generate_key_name()
        self.key_records[key] = targets
        return key

    def _generate_key0(self):
        """ Generate a key for only target 0. """
        return self._generate_key(targets=[self.targets[0]])

    def _key_and_push(self, *values):
        keys = [self._generate_key() for value in values]
        self._push(dict(zip(keys, values)))
        return tuple(keys)

    def delete_key(self, key):
        """ Delete the key from the engines, and our records. """
        if key in self.key_records:
            targets = self.key_records[key]
            self._delete_engine_key(key, targets)
            del self.key_records[key]
        else:
            raise KeyError('Key %s is not known by Context.' % (key))

    def _cleanup_keys(self):
        """ Delete all the keys we are tracking from the engines.
        
        Ideally this should leave no keys remaining on the engines.
        Any leftovers can be removed via purge_keys().
        """
        keys_to_remove = list(self.key_records.keys())
        for key in keys_to_remove:
            self.delete_key(key)

    # Methods to directly operate on keys on the engines.
    # These make no consideration for what keys we *expect* to be there.

    def _get_engine_keys(self, targets=None):
        """ Get the keys that exist on each of the engines, via globals().
        
        The return value is a dict, keyed by the key, with the value as
        a list of engines where the key is present.
        """
        if targets is None:
            targets = self.targets
        globals_key = self._generate_key()
        prefix = '__distarray'
        cmd = '%s = [k for k in globals().keys() if k.startswith("%s")]' % (
            globals_key, prefix)
        self._execute(cmd)
        keylists = self._pull(globals_key)
        self.delete_key(globals_key)
        # Convert nested list to dict with key=key, value=list of engines.
        engine_keys = {}
        for iengine, keylist in enumerate(keylists):
            for key in keylist:
                if key not in engine_keys:
                    engine_keys[key] = []
                engine_keys[key].append(self.targets[iengine])
        return engine_keys

    def _clean_engine_keys(self, engine_keys):
        """ Clean the keys of engine objects to remove things
        that this Context created for itself internally.
        (We do not want to delete those keys!)
        """
        internal_keys = [self._comm_key]
        for internal_key in internal_keys:
            if internal_key in engine_keys:
                del engine_keys[internal_key]
    
    def _delete_engine_key(self, key, targets=None):
        """ Delete the key on the specified engines.
        No consideration is made of how the Context is aware of the key.
        """
        if targets is None:
            targets = self.targets
        cmd = 'del %s' % key
        self._execute(cmd, targets=targets)

    # Cleanup of leftover keys. Normally there should not be any of these.

    def _get_leftover_keys(self):
        """ Get the keys that exist on each of the engines.
        
        They are discovered via globals().
        Keys that the Context created for itself (i.e. _comm_key) are
        stripped from the result for clarity.
        The result is a dictionary, keyed by the key, with the value
        as the list of engine targets where the key exists.
        """
        engine_keys = self._get_engine_keys()
        self._clean_engine_keys(engine_keys)
        return engine_keys

    def _purge_keys(self, verbose=False):
        """ Delete leftover keys from the engines.
        Return True if any keys were found, False if not. 

        If verbose is True, then the leftover keys are
        printed out as they are found.

        This is intended to clean up unexpected keys.
        """
        leftover_keys = self._get_leftover_keys()
        leftovers = (len(leftover_keys) > 0)
        for key in leftover_keys:
            if verbose:
                print('Leftover key: %s' % (key))
            targets = leftover_keys[key]
            self._delete_engine_key(key, targets)
        return leftovers

    def _cleanup_all_keys(self, verbose=False):
        """ Cleanup all the keys on the engines.

        This first cleans up the keys we expect to exist.
        It then notices any remaining keys, and cleans those as well.

        This returns True if any unexpected keys were cleaned
        up, which may indicate that something should be tracked better,
        or False if nothing unexpected was found.<commit_msg>Committing WIP.<commit_after># encoding: utf-8
# Copyright (c) 2008-2014, IPython Development Team and Enthought, Inc.

__docformat__ = "restructuredtext en"

import uuid
from distarray.externals import six
import collections

from IPython.parallel import Client
import numpy

from distarray.client import DistArray


class Context(object):

    def __init__(self, view=None, targets=None):
        if view is None:
            self.client = Client()
            self.view = self.client[:]
        else:
            self.view = view
            self.client = view.client

        all_targets = self.view.targets
        if targets is None:
            self.targets = all_targets
        else:
            self.targets = []
            for target in targets:
                if target not in all_targets:
                    raise ValueError("Engine with id %r not registered" % target)
                else:
                    self.targets.append(target)

        # FIXME: IPython bug #4296: This doesn't work under Python 3
        #with self.view.sync_imports():
        #    import distarray
        self.view.execute("import distarray.local; import distarray.mpiutils;"
                          " import numpy")

        self._setup_key_records()
        self._make_intracomm()
        self._set_engine_rank_mapping()

    def __del__(self):
        """ Clean up keys we have put on the engines. """
        # There is an order-of-ops problem here.
        # We sometimes call Client.close() explicitly, while leaving
        # this call to run when garbage is collected.
        # BUT, once we call close(), we cannot clean up the keys anymore.
        return    # leak!
        ##self._cleanup_keys()
        self.purge_keys()
        self._clear_comm_key()

    def _clear_comm_key(self):
        """ Delete our internal _comm_key from the engines. """
        # This could perhaps be called in the __exit__ of a context manager
        # for the Context. This _comm_key is not otherwise cleaned up by
        # the methods in this class.
        if hasattr(self, '_comm_key'):
            print 'deleting comm key', self._comm_key
            self._delete_engine_key(self._comm_key)
            self._comm_key = None

    def _set_engine_rank_mapping(self):
        # The MPI intracomm referred to by self._comm_key may have a different
        # mapping between IPython engines and MPI ranks than COMM_PRIVATE.  We
        # reorder self.targets so self.targets[i] is the IPython engine ID that
        # corresponds to MPI intracomm rank i.
        rank = self._generate_key()
        self.view.execute(
                '%s = %s.Get_rank()' % (rank, self._comm_key),
                block=True, targets=self.targets)

        # mapping target -> rank, rank -> target.
        target_to_rank = self.view.pull(rank, targets=self.targets).get_dict()
        rank_to_target = {v: k for (k, v) in target_to_rank.items()}
        # Probably wont need this...
        self.delete_key(rank)

        # ensure consistency
        assert set(self.targets) == set(target_to_rank)
        assert set(range(len(self.targets))) == set(rank_to_target)

        # reorder self.targets so that the targets are in MPI rank order for
        # the intracomm.
        self.targets = [rank_to_target[i] for i in range(len(rank_to_target))]

    def _make_intracomm(self):
        def get_rank():
            from distarray.mpiutils import COMM_PRIVATE
            return COMM_PRIVATE.Get_rank()

        # get a mapping of IPython engine ID to MPI rank
        rank_map = self.view.apply_async(get_rank).get_dict()
        ranks = [ rank_map[engine] for engine in self.targets ]

        # self.view's engines must encompass all ranks in the MPI communicator,
        # i.e., everything in rank_map.values().
        def get_size():
            from distarray.mpiutils import COMM_PRIVATE
            return COMM_PRIVATE.Get_size()

        comm_size = self.view.apply_async(get_size).get()[0]
        if set(rank_map.values()) != set(range(comm_size)):
            raise ValueError('Engines in view must encompass all MPI ranks.')

        # create a new communicator with the subset of engines note that
        # MPI_Comm_create must be called on all engines, not just those
        # involved in the new communicator.
        # This key is not added to key_records so it is not wrongly removed.
        self._comm_key = self._generate_key_name(subprefix='comm')
        self.view.execute(
            '%s = distarray.mpiutils.create_comm_with_list(%s)' % (self._comm_key, ranks),
            block=True
        )

    # Keeping track of the keys that we create on each engine.
    # This allows us to properly clean up after ourselves.
    # We store these as a dict, with the key as the key.
    # The value is the list of engines where the key has been pushed.
    # We can also find keys on the engines that are *not* tracked.
    # These should be kept track of better, but we can also delete them.

    def _setup_key_records(self):
        """ Create a dictionary, keyed by key name, with values as the
        set of targets the key is sent to.
        This can be used to properly clean up.
        """
        self.key_records = {}
        # Create a unique prefix for the keys for this context.
        # Full length seems excessively verbose so use 16 characters.
        uid = uuid.uuid4()
        self.key_prefix = uid.hex[:16]
        print 'key_prefix', self.key_prefix

    def _key_prefix(self, subprefix=''):
        """ Get the prefix, unique for this Context, used by all keys. """
        prefix = '__distarray_%s_%s_' % (subprefix, self.key_prefix)
        return prefix

    def _generate_key_name(self, subprefix=''):
        """ Generate a unique name for a key. """
        uid = uuid.uuid4()
        ##key = '__distarray_%s_%s' % (self.key_prefix, uid.hex)
        #if prefix is None:
        #    prefix = self._key_prefix() + '_'
        key = '%s_%s' % (self._key_prefix(subprefix), uid.hex)
        #key = '%s_%s' % (prefix, uid.hex)
        print 'generated key:', key
        return key

    def _generate_key(self, targets=None):
        """ Generate a key name, and record that it will be present
        on the specified target engines.
        """
        if targets is None:
            targets = self.targets
        key = self._generate_key_name()
        ##print 'generated key', key, 'for targets', targets
        self.key_records[key] = targets
        return key

    def _generate_key0(self):
        """ Generate a key for only target 0. """
        ##return self._generate_key(targets=[self.targets[0]])
        # temporary workaround- this leaks the key but wont ask for invalid delete.
        return self._generate_key_name()

    def _key_and_push(self, *values):
        keys = [self._generate_key() for value in values]
        self._push(dict(zip(keys, values)))
        return tuple(keys)

    def dump_keys(self, prefix=None):
        """ Print out a listing of all names on the engines. """
        engine_keys = self._get_engine_names(prefix)
        print 'engine keys:'
        print engine_keys
        for key in sorted(engine_keys.keys()):
            targets = engine_keys[key]
            print key, targets

    def purge_keys(self):
        """ Delete keys that this context created from all the engines. """
        cmd = """for k in globals().keys():
                     if k.startswith('__dist'):
                         print('%s IS DIST' % (k))
                         del globals()[k]
                     else:
                         print('%s is not' % (k))"""
        print 'cmd:', cmd
        self._execute(cmd)

    def delete_key(self, key):
        """ Delete the key from the engines, and our records. """
        self._delete_engine_key(key)
        return
        ##print 'deleting key', key
        if key in self.key_records:
            targets = self.key_records[key]
            self._delete_engine_key(key, targets)
            del self.key_records[key]
        else:
            raise KeyError('Key %s is not known by Context.' % (key))

    def XXX_cleanup_keys(self):
        """ Delete all the keys we are tracking from the engines.
        
        Ideally this should leave no keys remaining on the engines.
        Any leftovers can be removed via purge_keys().
        """
        keys_to_remove = list(self.key_records.keys())
        for key in keys_to_remove:
            self.delete_key(key)

    # Methods to directly operate on keys on the engines.
    # These make no consideration for what keys we *expect* to be there.

    def XXX_get_engine_keys(self, targets=None):
        """ Get the keys that exist on each of the engines, via globals().
        
        The return value is a dict, keyed by the key, with the value as
        a list of engines where the key is present.
        """
        if targets is None:
            targets = self.targets
        globals_key = self._generate_key()
        prefix = '__distarray'
        cmd = '%s = [k for k in globals().keys() if k.startswith("%s")]' % (
            globals_key, prefix)
        self._execute(cmd)
        keylists = self._pull(globals_key)
        self.delete_key(globals_key)
        # Convert nested list to dict with key=key, value=list of engines.
        engine_keys = {}
        for iengine, keylist in enumerate(keylists):
            for key in keylist:
                if key not in engine_keys:
                    engine_keys[key] = []
                engine_keys[key].append(self.targets[iengine])
        return engine_keys

    def _get_engine_names(self, prefix=None):
        """ Get the keys that exist on each of the engines, via globals().
        
        The return value is a dict, keyed by the key, with the value as
        a list of engines where the key is present.
        """
        globals_key = self._generate_key()
        prefix = '__distarray'
        ##prefix = ''
        ##prefix = self._key_prefix()
        if prefix is None:
            prefix = self._key_prefix()

        cmd = '%s = [k for k in globals().keys() if k.startswith("%s")]' % (
            globals_key, prefix)
        print 'cmd:', cmd
        self._execute(cmd)
        keylists = self._pull(globals_key)
        self.delete_key(globals_key)
        # Convert nested list to dict with key=key, value=list of engines.
        engine_keys = {}
        for iengine, keylist in enumerate(keylists):
            for key in keylist:
                if key not in engine_keys:
                    engine_keys[key] = []
                engine_keys[key].append(self.targets[iengine])
        return engine_keys

    def XXX_clean_engine_keys(self, engine_keys):
        """ Clean the keys of engine objects to remove things
        that this Context created for itself internally.
        (We do not want to delete those keys!)
        """
        internal_keys = [self._comm_key]
        for internal_key in internal_keys:
            if internal_key in engine_keys:
                del engine_keys[internal_key]
    
    def _delete_engine_key(self, key, targets=None):
        """ Delete the key on the specified engines.
        No consideration is made of how the Context is aware of the key.
        """
        if targets is None:
            targets = self.targets
        cmd = 'del %s' % key
        # No more targets...
        #self._execute(cmd, targets=targets)
        self._execute(cmd)

    # Cleanup of leftover keys. Normally there should not be any of these.

    def XXX_get_leftover_keys(self):
        """ Get the keys that exist on each of the engines.
        
        They are discovered via globals().
        Keys that the Context created for itself (i.e. _comm_key) are
        stripped from the result for clarity.
        The result is a dictionary, keyed by the key, with the value
        as the list of engine targets where the key exists.
        """
        engine_keys = self._get_engine_keys()
        self._clean_engine_keys(engine_keys)
        return engine_keys

    def XXX_purge_keys(self, verbose=False):
        """ Delete leftover keys from the engines.
        Return True if any keys were found, False if not. 

        If verbose is True, then the leftover keys are
        printed out as they are found.

        This is intended to clean up unexpected keys.
        """
        leftover_keys = self._get_leftover_keys()
        leftovers = (len(leftover_keys) > 0)
        for key in leftover_keys:
            if verbose:
                print('Leftover key: %s' % (key))
            targets = leftover_keys[key]
            # No more targets arg...
            #self._delete_engine_key(key, targets)
            self._delete_engine_key(key)
        return leftovers

    def XXX_cleanup_all_keys(self, verbose=False):
        """ Cleanup all the keys on the engines.

        This first cleans up the keys we expect to exist.
        It then notices any remaining keys, and cleans those as well.

        This returns True if any unexpected keys were cleaned
        up, which may indicate that something should be tracked better,
        or False if nothing unexpected was found.
{'fraction_non_alphanumeric': 0.0818926296633303, 'fraction_numerical': 0.0019498245157935786, 'mean_word_length': 2.123984840281538, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 9, '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': '311316', 'n_tokens_mistral': 6525, 'n_tokens_neox': 6240, 'n_words': 2279}