<commit_before>import logging
import multiprocessing

import numpy

from smqtk.utils import SmqtkObject


__author__ = '<EMAIL>'


__all__ = [
    'elements_to_matrix',
]


def elements_to_matrix(descr_elements, mat=None, procs=None):
    """
    Add to or create a numpy matrix, adding to it the vector data contained in
    a sequence of DescriptorElement instances using asynchronous processing.

    If ``mat`` is provided, its shape must equal:
        ( len(descr_elements) , descr_elements[0].size )

    :param descr_elements:
    :param mat:
    :param procs:

    :return: Create or input matrix


    """
    log = logging.getLogger(__name__)

    # Special case for in-memory storage of descriptors
    from smqtk.representation.descriptor_element.local_elements \
        import DescriptorMemoryElement

    # Create/check matrix
    if mat is None:
        mat = numpy.ndarray((len(descr_elements),
                             descr_elements[0].vector().size),
                            float)
    else:
        assert mat.shape[0] == len(descr_elements)
        assert mat.shape[1] == descr_elements[0].vector().size

    if procs is None:
        procs = multiprocessing.cpu_count()

    in_q = multiprocessing.Queue()
    out_q = multiprocessing.Queue(procs * 2)

    # Workers for async extraction
    workers = [_ElemVectorExtractor(in_q, out_q) for _ in xrange(procs)]

    async_packets_sent = 0
    for r, d in enumerate(descr_elements):
        if isinstance(d, DescriptorMemoryElement):
            # No loading required, already in memory
            mat[r] = d.vector()
        else:
            in_q.put((r, d))
            async_packets_sent += 1

    # Collect work from async
    for _ in xrange(async_packets_sent):
        r, v = out_q.get()
        mat[r] = v
    out_q.close()

    # All work should be exhausted at this point
    assert in_q.qsize() == 0
    assert out_q.qsize() == 0

    # Shutdown workers
    # - Workers should exit upon getting a None packet
    for _ in workers:
        # One for each worker
        in_q.put(None)
    in_q.close()
    for w in workers:
        w.join()

    return mat


class _ElemVectorExtractor (SmqtkObject, multiprocessing.Process):
    """
    Helper process for extracting DescriptorElement vectors on a separate
    process. This terminates with a None packet fed to in_q. Otherwise, in_q
    values are expected to be (row, element) pairs. Tuples of the form
    (row, vector) are published to the out_q.

    Terminal value: None

    """

    def __init__(self, in_q, out_q):
        super(_ElemVectorExtractor, self)\
            .__init__()
        self.in_q = in_q
        self.out_q = out_q

    def run(self):
        packet = self.in_q.get()
        while packet is not None:
            row, elem = packet
            self.out_q.put((row, elem.vector()))
            packet = self.in_q.get()<commit_msg>Fixes to async element-to-matrix func<commit_after>import logging
import multiprocessing
import time

import numpy

from smqtk.utils import SmqtkObject


__author__ = '<EMAIL>'


__all__ = [
    'elements_to_matrix',
]


def elements_to_matrix(descr_elements, mat=None, procs=None):
    """
    Add to or create a numpy matrix, adding to it the vector data contained in
    a sequence of DescriptorElement instances using asynchronous processing.

    If ``mat`` is provided, its shape must equal:
        ( len(descr_elements) , descr_elements[0].size )

    :param descr_elements:
    :param mat:
    :param procs:

    :return: Create or input matrix


    """
    log = logging.getLogger(__name__)

    # Special case for in-memory storage of descriptors
    from smqtk.representation.descriptor_element.local_elements \
        import DescriptorMemoryElement

    # Create/check matrix
    if mat is None:
        shp = (len(descr_elements),
               descr_elements[0].vector().size)
        log.debug("Creating new matrix with shape: %s", shp)
        mat = numpy.ndarray(shp, float)
    else:
        assert mat.shape[0] == len(descr_elements)
        assert mat.shape[1] == descr_elements[0].vector().size

    if procs is None:
        procs = multiprocessing.cpu_count()

    in_q = multiprocessing.Queue()
    out_q = multiprocessing.Queue(procs * 2)
    log.debug("Output queue size: %d", out_q._maxsize)

    # Workers for async extraction
    log.debug("constructing worker processes")
    workers = [_ElemVectorExtractor(i, in_q, out_q) for i in xrange(procs)]
    for w in workers:
        w.daemon = True

    try:
        # Start worker processes
        log.debug("starting worker processes")
        for w in workers:
            w.start()

        log.debug("Sending work packets")
        async_packets_sent = 0
        for r, d in enumerate(descr_elements):
            if isinstance(d, DescriptorMemoryElement):
                # No loading required, already in memory
                mat[r] = d.vector()
            else:
                in_q.put((r, d))
                async_packets_sent += 1

        # Collect work from async
        log.debug("Aggregating async results")
        lf = f = 0
        t = time.time()
        for _ in xrange(async_packets_sent):
            r, v = out_q.get()
            mat[r] = v

            f += 1
            if time.time() - t > 1:
                log.debug("Rows per second: %f, Total: %d",  (f-lf) / (time.time()/t), f)
                lf = f
                t = time.time()

        log.debug("Closing output queue")
        out_q.close()

        # All work should be exhausted at this point
        assert in_q.qsize() == 0
        assert out_q.qsize() == 0

        # Shutdown workers
        # - Workers should exit upon getting a None packet
        log.debug("Sending worker terminal signals")
        for _ in workers:
            # One for each worker
            in_q.put(None)
        log.debug("Closing input queue")
        in_q.close()

        log.debug("Done")
        return mat
    finally:
        # Forcibly terminate worker processes if still alive
        log.debug("Joining/Terminating workers")
        for w in workers:
            if w.is_alive():
                w.terminate()
            w.join()
        for q in (in_q, out_q):
            q.close()
            q.join_thread()


class _ElemVectorExtractor (SmqtkObject, multiprocessing.Process):
    """
    Helper process for extracting DescriptorElement vectors on a separate
    process. This terminates with a None packet fed to in_q. Otherwise, in_q
    values are expected to be (row, element) pairs. Tuples of the form
    (row, vector) are published to the out_q.

    Terminal value: None

    """

    def __init__(self, i, in_q, out_q):
        super(_ElemVectorExtractor, self)\
            .__init__(name='[w%d]' % i)
        self._log.debug("Making worker (%d, %s, %s)", i, in_q, out_q)
        self.i = i
        self.in_q = in_q
        self.out_q = out_q

    def run(self):
        packet = self.in_q.get()
        #self._log.debug("[%d] Received packet: %s", self.i, packet)
        while packet is not None:
            row, elem = packet
            self.out_q.put((row, elem.vector()))
            packet = self.in_q.get()
            #self._log.debug("[%d] Received packet: %s", self.i, packet)
{'fraction_non_alphanumeric': 0.0965081110805609, 'fraction_numerical': 0.0031619466593346165, 'mean_word_length': 2.1852014010507883, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 6, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5826787', 'n_tokens_mistral': 2264, 'n_tokens_neox': 2164, 'n_words': 653}