<commit_before>from __future__ import absolute_import, division, print_function

import errno
import json
import logging
import signal
import socket
import sys
import threading
import warnings
from collections import OrderedDict

import six
from six.moves import queue
import torch
import torch.multiprocessing as mp

import pyro
from pyro.infer import TracePosterior
from pyro.infer.mcmc.logger import initialize_logger, initialize_progbar, DIAGNOSTIC_MSG, TqdmHandler
from pyro.util import optional


def logger_thread(log_queue, warmup_steps, num_samples, num_chains):
    """
    Logging thread that asynchronously consumes logging events from `log_queue`,
    and handles them appropriately.
    """
    progress_bars = [initialize_progbar(warmup_steps, s, pos=i)
                     for i, s in enumerate(num_samples)]
    logger = logging.getLogger(__name__)
    logger.propagate = False
    logger.addHandler(TqdmHandler())
    num_samples = [0] * len(num_samples)
    try:
        while True:
            try:
                record = log_queue.get_nowait()
            except queue.Empty:
                continue
            if record is None:
                break
            metadata, msg = record.getMessage().split("]", 1)
            _, msg_type, logger_id = metadata[1:].split()
            if msg_type == DIAGNOSTIC_MSG:
                pbar_pos = int(logger_id.split(":")[-1]) - 1
                num_samples[pbar_pos] += 1
                if num_samples[pbar_pos] == warmup_steps:
                    progress_bars[pbar_pos].set_description("Sample [{}]".format(pbar_pos + 1))
                diagnostics = json.loads(msg, object_pairs_hook=OrderedDict)
                progress_bars[pbar_pos].set_postfix(diagnostics)
                progress_bars[pbar_pos].update()
            else:
                logger.handle(record)
    finally:
        for pbar in progress_bars:
            pbar.close()
        # Required to not overwrite multiple progress bars on exit.
        sys.stderr.write("\n" * num_chains)


class _Worker(object):
    def __init__(self, chain_id, result_queue, log_queue,
                 kernel, num_samples, warmup_steps=0,
                 args=None, kwargs=None):
        self.chain_id = chain_id
        self.trace_gen = _SingleSampler(kernel, num_samples=num_samples, warmup_steps=warmup_steps)
        self.args = args if args is not None else []
        self.kwargs = kwargs if kwargs is not None else {}
        self.rng_seed = torch.initial_seed()
        self.log_queue = log_queue
        self.result_queue = result_queue
        self.default_tensor_type = torch.Tensor().type()

    def run(self, *args, **kwargs):
        pyro.set_rng_seed(self.chain_id + self.rng_seed)
        torch.set_default_tensor_type(self.default_tensor_type)
        kwargs["logger_id"] = "CHAIN:{}".format(self.chain_id)
        kwargs["log_queue"] = self.log_queue
        try:
            for sample in self.trace_gen._traces(*args, **kwargs):
                self.result_queue.put_nowait((self.chain_id, sample))
            self.result_queue.put_nowait((self.chain_id, None))
        except Exception as e:
            self.trace_gen.logger.exception(e)
            self.result_queue.put_nowait((self.chain_id, e))


class _ParallelSampler(TracePosterior):
    """
    Parallel runner class for running MCMC chains in parallel. This uses the
    `torch.multiprocessing` module (itself a light wrapper over the python
    `multiprocessing` module) to spin up parallel workers.
    """
    def __init__(self, kernel, num_samples, warmup_steps, num_chains, mp_context):
        super(_ParallelSampler, self).__init__()
        self.kernel = kernel
        self.warmup_steps = warmup_steps
        self.num_chains = num_chains
        self.workers = []
        self.ctx = mp
        if mp_context:
            if six.PY2:
                raise ValueError("multiprocessing.get_context() is "
                                 "not supported in Python 2.")
            self.ctx = mp.get_context(mp_context)
        self.result_queue = self.ctx.Manager().Queue()
        self.log_queue = self.ctx.Manager().Queue()
        self.logger = initialize_logger(logging.getLogger("pyro.infer.mcmc"),
                                        "MAIN", log_queue=self.log_queue)
        # initialize number of samples per chain
        samples_per_chain = num_samples // num_chains
        self.num_samples = [samples_per_chain] * num_chains
        for i in range(num_samples % num_chains):
            self.num_samples[i] += 1
        self.log_thread = threading.Thread(target=logger_thread,
                                           args=(self.log_queue, self.warmup_steps,
                                                 self.num_samples, self.num_chains))
        self.log_thread.daemon = True
        self.log_thread.start()

    def init_workers(self, *args, **kwargs):
        self.workers = []
        for i in range(self.num_chains):
            worker = _Worker(i + 1, self.result_queue, self.log_queue, self.kernel,
                             self.num_samples[i], self.warmup_steps)
            worker.daemon = True
            self.workers.append(self.ctx.Process(name=str(i), target=worker.run,
                                                 args=args, kwargs=kwargs))

    def terminate(self):
        if self.log_thread.is_alive():
            self.log_queue.put_nowait(None)
            self.log_thread.join(timeout=1)
        for w in self.workers:
            if w.is_alive():
                w.terminate()

    def _traces(self, *args, **kwargs):
        # Ignore sigint in worker processes; they will be shut down
        # when the main process terminates.
        sigint_handler = signal.signal(signal.SIGINT, signal.SIG_IGN)
        self.init_workers(*args, **kwargs)
        # restore original handler
        signal.signal(signal.SIGINT, sigint_handler)
        active_workers = self.num_chains
        try:
            for w in self.workers:
                w.start()
            while active_workers:
                try:
                    chain_id, val = self.result_queue.get_nowait()
                # This can happen when the worker process has terminated.
                # See https://github.com/pytorch/pytorch/pull/5380 for motivation.
                except socket.error as e:
                    if getattr(e, "errno", None) == errno.ENOENT:
                        pass
                    else:
                        raise e
                except queue.Empty:
                    continue
                if isinstance(val, Exception):
                    # Exception trace is already logged by worker.
                    raise val
                elif val is not None:
                    yield val
                else:
                    active_workers -= 1<commit_msg>Return a deterministic trace generator when using HMC with parallel chains (#1506)<commit_after>from __future__ import absolute_import, division, print_function

import errno
import json
import logging
import signal
import socket
import sys
import threading
import warnings
from collections import OrderedDict, deque

import six
from six.moves import queue
import torch
import torch.multiprocessing as mp

import pyro
from pyro.infer import TracePosterior
from pyro.infer.mcmc.logger import initialize_logger, initialize_progbar, DIAGNOSTIC_MSG, TqdmHandler
from pyro.util import optional


def logger_thread(log_queue, warmup_steps, num_samples, num_chains):
    """
    Logging thread that asynchronously consumes logging events from `log_queue`,
    and handles them appropriately.
    """
    progress_bars = [initialize_progbar(warmup_steps, s, pos=i)
                     for i, s in enumerate(num_samples)]
    logger = logging.getLogger(__name__)
    logger.propagate = False
    logger.addHandler(TqdmHandler())
    num_samples = [0] * len(num_samples)
    try:
        while True:
            try:
                record = log_queue.get_nowait()
            except queue.Empty:
                continue
            if record is None:
                break
            metadata, msg = record.getMessage().split("]", 1)
            _, msg_type, logger_id = metadata[1:].split()
            if msg_type == DIAGNOSTIC_MSG:
                pbar_pos = int(logger_id.split(":")[-1]) - 1
                num_samples[pbar_pos] += 1
                if num_samples[pbar_pos] == warmup_steps:
                    progress_bars[pbar_pos].set_description("Sample [{}]".format(pbar_pos + 1))
                diagnostics = json.loads(msg, object_pairs_hook=OrderedDict)
                progress_bars[pbar_pos].set_postfix(diagnostics)
                progress_bars[pbar_pos].update()
            else:
                logger.handle(record)
    finally:
        for pbar in progress_bars:
            pbar.close()
        # Required to not overwrite multiple progress bars on exit.
        sys.stderr.write("\n" * num_chains)


class _Worker(object):
    def __init__(self, chain_id, result_queue, log_queue,
                 kernel, num_samples, warmup_steps=0,
                 args=None, kwargs=None):
        self.chain_id = chain_id
        self.trace_gen = _SingleSampler(kernel, num_samples=num_samples, warmup_steps=warmup_steps)
        self.args = args if args is not None else []
        self.kwargs = kwargs if kwargs is not None else {}
        self.rng_seed = torch.initial_seed()
        self.log_queue = log_queue
        self.result_queue = result_queue
        self.default_tensor_type = torch.Tensor().type()

    def run(self, *args, **kwargs):
        pyro.set_rng_seed(self.chain_id + self.rng_seed)
        torch.set_default_tensor_type(self.default_tensor_type)
        kwargs["logger_id"] = "CHAIN:{}".format(self.chain_id)
        kwargs["log_queue"] = self.log_queue
        try:
            for sample in self.trace_gen._traces(*args, **kwargs):
                self.result_queue.put_nowait((self.chain_id, sample))
            self.result_queue.put_nowait((self.chain_id, None))
        except Exception as e:
            self.trace_gen.logger.exception(e)
            self.result_queue.put_nowait((self.chain_id, e))


class _ParallelSampler(TracePosterior):
    """
    Parallel runner class for running MCMC chains in parallel. This uses the
    `torch.multiprocessing` module (itself a light wrapper over the python
    `multiprocessing` module) to spin up parallel workers.
    """
    def __init__(self, kernel, num_samples, warmup_steps, num_chains, mp_context):
        super(_ParallelSampler, self).__init__()
        self.kernel = kernel
        self.warmup_steps = warmup_steps
        self.num_chains = num_chains
        self.workers = []
        self.ctx = mp
        if mp_context:
            if six.PY2:
                raise ValueError("multiprocessing.get_context() is "
                                 "not supported in Python 2.")
            self.ctx = mp.get_context(mp_context)
        self.result_queue = self.ctx.Manager().Queue()
        self.log_queue = self.ctx.Manager().Queue()
        self.logger = initialize_logger(logging.getLogger("pyro.infer.mcmc"),
                                        "MAIN", log_queue=self.log_queue)
        # initialize number of samples per chain
        samples_per_chain = num_samples // num_chains
        self.num_samples = [samples_per_chain] * num_chains
        for i in range(num_samples % num_chains):
            self.num_samples[i] += 1
        self.log_thread = threading.Thread(target=logger_thread,
                                           args=(self.log_queue, self.warmup_steps,
                                                 self.num_samples, self.num_chains))
        self.log_thread.daemon = True
        self.log_thread.start()

    def init_workers(self, *args, **kwargs):
        self.workers = []
        for i in range(self.num_chains):
            worker = _Worker(i + 1, self.result_queue, self.log_queue, self.kernel,
                             self.num_samples[i], self.warmup_steps)
            worker.daemon = True
            self.workers.append(self.ctx.Process(name=str(i), target=worker.run,
                                                 args=args, kwargs=kwargs))

    def terminate(self):
        if self.log_thread.is_alive():
            self.log_queue.put_nowait(None)
            self.log_thread.join(timeout=1)
        for w in self.workers:
            if w.is_alive():
                w.terminate()

    def _traces(self, *args, **kwargs):
        # Ignore sigint in worker processes; they will be shut down
        # when the main process terminates.
        sigint_handler = signal.signal(signal.SIGINT, signal.SIG_IGN)
        self.init_workers(*args, **kwargs)
        # restore original handler
        signal.signal(signal.SIGINT, sigint_handler)
        active_workers = self.num_chains
        # To yield a deterministic ordering, we hold intermediate traces
        # from each of the workers in its own queue in `results_buffer`
        # and yield these in a round robin fashion.
        buffer_idx = 0
        results_buffer = [deque() for _ in range(self.num_chains)]
        try:
            for w in self.workers:
                w.start()
            while active_workers:
                try:
                    chain_id, val = self.result_queue.get_nowait()
                # This can happen when the worker process has terminated.
                # See https://github.com/pytorch/pytorch/pull/5380 for motivation.
                except socket.error as e:
                    if getattr(e, "errno", None) == errno.ENOENT:
                        pass
                    else:
                        raise e
                except queue.Empty:
                    continue
                if isinstance(val, Exception):
                    # Exception trace is already logged by worker.
                    raise val
                if val is None:
                    active_workers -= 1
                else:
                    results_buffer[chain_id - 1].append(val)
                while results_buffer[buffer_idx]:
                    yield results_buffer[buffer_idx].popleft()
                    buffer_idx = (buffer_idx + 1) % self.num_chains
            # empty out the results buffer
            non_empty_buffers = set(range(self.num_chains))
            while non_empty_buffers:
                if results_buffer[buffer_idx]:
                    yield results_buffer[buffer_idx].popleft()
                else:
                    if buffer_idx in non_empty_buffers:
                        non_empty_buffers.remove(buffer_idx)
                buffer_idx = (buffer_idx + 1) % self.num_chains
{'fraction_non_alphanumeric': 0.08574731402148783, 'fraction_numerical': 0.0029919760641914863, 'mean_word_length': 1.9729128764908026, 'pattern_counts': {'":': 2, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 16, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3914124', 'n_tokens_mistral': 4168, 'n_tokens_neox': 4004, 'n_words': 970}