We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
HMC gives non deterministic results with parallel chaining
Even though the `ParallelSampler` yields the same set of traces (each worker gives deterministic results), the order in which the traces are consumed might be different, which gives different results when we use `TracePredictive`. 

One way to get deterministic results from HMC/NUTS with `num_chains > 1` would be to store intermediate results, and yield results from each of the workers in a round-robin fashion. 

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `pyro/infer/mcmc/mcmc.py`
Content:
```
1 from __future__ import absolute_import, division, print_function
2 
3 import errno
4 import json
5 import logging
6 import signal
7 import socket
8 import sys
9 import threading
10 import warnings
11 from collections import OrderedDict
12 
13 import six
14 from six.moves import queue
15 import torch
16 import torch.multiprocessing as mp
17 
18 import pyro
19 from pyro.infer import TracePosterior
20 from pyro.infer.mcmc.logger import initialize_logger, initialize_progbar, DIAGNOSTIC_MSG, TqdmHandler
21 from pyro.util import optional
22 
23 
24 def logger_thread(log_queue, warmup_steps, num_samples, num_chains):
25     """
26     Logging thread that asynchronously consumes logging events from `log_queue`,
27     and handles them appropriately.
28     """
29     progress_bars = [initialize_progbar(warmup_steps, s, pos=i)
30                      for i, s in enumerate(num_samples)]
31     logger = logging.getLogger(__name__)
32     logger.propagate = False
33     logger.addHandler(TqdmHandler())
34     num_samples = [0] * len(num_samples)
35     try:
36         while True:
37             try:
38                 record = log_queue.get_nowait()
39             except queue.Empty:
40                 continue
41             if record is None:
42                 break
43             metadata, msg = record.getMessage().split("]", 1)
44             _, msg_type, logger_id = metadata[1:].split()
45             if msg_type == DIAGNOSTIC_MSG:
46                 pbar_pos = int(logger_id.split(":")[-1]) - 1
47                 num_samples[pbar_pos] += 1
48                 if num_samples[pbar_pos] == warmup_steps:
49                     progress_bars[pbar_pos].set_description("Sample [{}]".format(pbar_pos + 1))
50                 diagnostics = json.loads(msg, object_pairs_hook=OrderedDict)
51                 progress_bars[pbar_pos].set_postfix(diagnostics)
52                 progress_bars[pbar_pos].update()
53             else:
54                 logger.handle(record)
55     finally:
56         for pbar in progress_bars:
57             pbar.close()
58         # Required to not overwrite multiple progress bars on exit.
59         sys.stderr.write("\n" * num_chains)
60 
61 
62 class _Worker(object):
63     def __init__(self, chain_id, result_queue, log_queue,
64                  kernel, num_samples, warmup_steps=0,
65                  args=None, kwargs=None):
66         self.chain_id = chain_id
67         self.trace_gen = _SingleSampler(kernel, num_samples=num_samples, warmup_steps=warmup_steps)
68         self.args = args if args is not None else []
69         self.kwargs = kwargs if kwargs is not None else {}
70         self.rng_seed = torch.initial_seed()
71         self.log_queue = log_queue
72         self.result_queue = result_queue
73         self.default_tensor_type = torch.Tensor().type()
74 
75     def run(self, *args, **kwargs):
76         pyro.set_rng_seed(self.chain_id + self.rng_seed)
77         torch.set_default_tensor_type(self.default_tensor_type)
78         kwargs["logger_id"] = "CHAIN:{}".format(self.chain_id)
79         kwargs["log_queue"] = self.log_queue
80         try:
81             for sample in self.trace_gen._traces(*args, **kwargs):
82                 self.result_queue.put_nowait((self.chain_id, sample))
83             self.result_queue.put_nowait((self.chain_id, None))
84         except Exception as e:
85             self.trace_gen.logger.exception(e)
86             self.result_queue.put_nowait((self.chain_id, e))
87 
88 
89 class _ParallelSampler(TracePosterior):
90     """
91     Parallel runner class for running MCMC chains in parallel. This uses the
92     `torch.multiprocessing` module (itself a light wrapper over the python
93     `multiprocessing` module) to spin up parallel workers.
94     """
95     def __init__(self, kernel, num_samples, warmup_steps, num_chains, mp_context):
96         super(_ParallelSampler, self).__init__()
97         self.kernel = kernel
98         self.warmup_steps = warmup_steps
99         self.num_chains = num_chains
100         self.workers = []
101         self.ctx = mp
102         if mp_context:
103             if six.PY2:
104                 raise ValueError("multiprocessing.get_context() is "
105                                  "not supported in Python 2.")
106             self.ctx = mp.get_context(mp_context)
107         self.result_queue = self.ctx.Manager().Queue()
108         self.log_queue = self.ctx.Manager().Queue()
109         self.logger = initialize_logger(logging.getLogger("pyro.infer.mcmc"),
110                                         "MAIN", log_queue=self.log_queue)
111         # initialize number of samples per chain
112         samples_per_chain = num_samples // num_chains
113         self.num_samples = [samples_per_chain] * num_chains
114         for i in range(num_samples % num_chains):
115             self.num_samples[i] += 1
116         self.log_thread = threading.Thread(target=logger_thread,
117                                            args=(self.log_queue, self.warmup_steps,
118                                                  self.num_samples, self.num_chains))
119         self.log_thread.daemon = True
120         self.log_thread.start()
121 
122     def init_workers(self, *args, **kwargs):
123         self.workers = []
124         for i in range(self.num_chains):
125             worker = _Worker(i + 1, self.result_queue, self.log_queue, self.kernel,
126                              self.num_samples[i], self.warmup_steps)
127             worker.daemon = True
128             self.workers.append(self.ctx.Process(name=str(i), target=worker.run,
129                                                  args=args, kwargs=kwargs))
130 
131     def terminate(self):
132         if self.log_thread.is_alive():
133             self.log_queue.put_nowait(None)
134             self.log_thread.join(timeout=1)
135         for w in self.workers:
136             if w.is_alive():
137                 w.terminate()
138 
139     def _traces(self, *args, **kwargs):
140         # Ignore sigint in worker processes; they will be shut down
141         # when the main process terminates.
142         sigint_handler = signal.signal(signal.SIGINT, signal.SIG_IGN)
143         self.init_workers(*args, **kwargs)
144         # restore original handler
145         signal.signal(signal.SIGINT, sigint_handler)
146         active_workers = self.num_chains
147         try:
148             for w in self.workers:
149                 w.start()
150             while active_workers:
151                 try:
152                     chain_id, val = self.result_queue.get_nowait()
153                 # This can happen when the worker process has terminated.
154                 # See https://github.com/pytorch/pytorch/pull/5380 for motivation.
155                 except socket.error as e:
156                     if getattr(e, "errno", None) == errno.ENOENT:
157                         pass
158                     else:
159                         raise e
160                 except queue.Empty:
161                     continue
162                 if isinstance(val, Exception):
163                     # Exception trace is already logged by worker.
164                     raise val
165                 elif val is not None:
166                     yield val
167                 else:
168                     active_workers -= 1
169         finally:
170             self.terminate()
171 
172 
173 class _SingleSampler(TracePosterior):
174     """
175     Single process runner class optimized for the case `num_chains=1`.
176     """
177     def __init__(self, kernel, num_samples, warmup_steps):
178         self.kernel = kernel
179         self.warmup_steps = warmup_steps
180         self.num_samples = num_samples
181         self.logger = None
182         super(_SingleSampler, self).__init__()
183 
184     def _gen_samples(self, num_samples, init_trace):
185         trace = init_trace
186         for _ in range(num_samples):
187             trace = self.kernel.sample(trace)
188             diagnostics = json.dumps(self.kernel.diagnostics())
189             self.logger.info(diagnostics, extra={"msg_type": DIAGNOSTIC_MSG})
190             yield trace
191 
192     def _traces(self, *args, **kwargs):
193         logger_id = kwargs.pop("logger_id", "")
194         log_queue = kwargs.pop("log_queue", None)
195         self.logger = logging.getLogger("pyro.infer.mcmc")
196         is_multiprocessing = log_queue is not None
197         progress_bar = initialize_progbar(self.warmup_steps, self.num_samples) \
198             if not is_multiprocessing else None
199         self.logger = initialize_logger(self.logger, logger_id, progress_bar, log_queue)
200         self.kernel.setup(self.warmup_steps, *args, **kwargs)
201         trace = self.kernel.initial_trace()
202         with optional(progress_bar, not is_multiprocessing):
203             for trace in self._gen_samples(self.warmup_steps, trace):
204                 continue
205             if progress_bar:
206                 progress_bar.set_description("Sample")
207             for trace in self._gen_samples(self.num_samples, trace):
208                 yield (trace, 1.0)
209         self.kernel.cleanup()
210 
211 
212 class MCMC(TracePosterior):
213     """
214     Wrapper class for Markov Chain Monte Carlo algorithms. Specific MCMC algorithms
215     are TraceKernel instances and need to be supplied as a ``kernel`` argument
216     to the constructor.
217 
218     .. note:: The case of `num_chains > 1` uses python multiprocessing to
219         run parallel chains in multiple processes. This goes with the usual
220         caveats around multiprocessing in python, e.g. the model used to
221         initialize the ``kernel`` must be serializable via `pickle`, and the
222         performance / constraints will be platform dependent (e.g. only
223         the "spawn" context is available in Windows). This has also not
224         been extensively tested on the Windows platform.
225 
226     :param kernel: An instance of the ``TraceKernel`` class, which when
227         given an execution trace returns another sample trace from the target
228         (posterior) distribution.
229     :param int num_samples: The number of samples that need to be generated,
230         excluding the samples discarded during the warmup phase.
231     :param int warmup_steps: Number of warmup iterations. The samples generated
232         during the warmup phase are discarded. If not provided, default is
233         half of `num_samples`.
234     :param int num_chains: Number of MCMC chains to run in parallel. Depending on
235         whether `num_chains` is 1 or more than 1, this class internally dispatches
236         to either `_SingleSampler` or `_ParallelSampler`.
237     :param str mp_context: Multiprocessing context to use when `num_chains > 1`.
238         Only applicable for Python 3.5 and above. Use `mp_context="spawn"` for
239         CUDA.
240     """
241     def __init__(self, kernel, num_samples, warmup_steps=0,
242                  num_chains=1, mp_context=None):
243         self.warmup_steps = warmup_steps if warmup_steps is not None else num_samples // 2  # Stan
244         self.num_samples = num_samples
245         if num_chains > 1:
246             cpu_count = mp.cpu_count()
247             if num_chains > cpu_count:
248                 warnings.warn("`num_chains` is more than CPU count - {}. "
249                               "Resetting num_chains to CPU count.").format(cpu_count)
250             self.sampler = _ParallelSampler(kernel, num_samples, warmup_steps,
251                                             num_chains, mp_context)
252         else:
253             self.sampler = _SingleSampler(kernel, num_samples, warmup_steps)
254         super(MCMC, self).__init__()
255 
256     def _traces(self, *args, **kwargs):
257         for sample in self.sampler._traces(*args, **kwargs):
258             yield sample
259
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/pyro/infer/mcmc/mcmc.py b/pyro/infer/mcmc/mcmc.py
--- a/pyro/infer/mcmc/mcmc.py
+++ b/pyro/infer/mcmc/mcmc.py
@@ -8,7 +8,7 @@
 import sys
 import threading
 import warnings
-from collections import OrderedDict
+from collections import OrderedDict, deque
 
 import six
 from six.moves import queue
@@ -144,6 +144,11 @@
         # 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()
@@ -162,10 +167,22 @@
                 if isinstance(val, Exception):
                     # Exception trace is already logged by worker.
                     raise val
-                elif val is not None:
-                    yield val
-                else:
+                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
         finally:
             self.terminate()

{"golden_diff": "diff --git a/pyro/infer/mcmc/mcmc.py b/pyro/infer/mcmc/mcmc.py\n--- a/pyro/infer/mcmc/mcmc.py\n+++ b/pyro/infer/mcmc/mcmc.py\n@@ -8,7 +8,7 @@\n import sys\n import threading\n import warnings\n-from collections import OrderedDict\n+from collections import OrderedDict, deque\n \n import six\n from six.moves import queue\n@@ -144,6 +144,11 @@\n         # restore original handler\n         signal.signal(signal.SIGINT, sigint_handler)\n         active_workers = self.num_chains\n+        # To yield a deterministic ordering, we hold intermediate traces\n+        # from each of the workers in its own queue in `results_buffer`\n+        # and yield these in a round robin fashion.\n+        buffer_idx = 0\n+        results_buffer = [deque() for _ in range(self.num_chains)]\n         try:\n             for w in self.workers:\n                 w.start()\n@@ -162,10 +167,22 @@\n                 if isinstance(val, Exception):\n                     # Exception trace is already logged by worker.\n                     raise val\n-                elif val is not None:\n-                    yield val\n-                else:\n+                if val is None:\n                     active_workers -= 1\n+                else:\n+                    results_buffer[chain_id - 1].append(val)\n+                while results_buffer[buffer_idx]:\n+                    yield results_buffer[buffer_idx].popleft()\n+                    buffer_idx = (buffer_idx + 1) % self.num_chains\n+            # empty out the results buffer\n+            non_empty_buffers = set(range(self.num_chains))\n+            while non_empty_buffers:\n+                if results_buffer[buffer_idx]:\n+                    yield results_buffer[buffer_idx].popleft()\n+                else:\n+                    if buffer_idx in non_empty_buffers:\n+                        non_empty_buffers.remove(buffer_idx)\n+                buffer_idx = (buffer_idx + 1) % self.num_chains\n         finally:\n             self.terminate()\n", "issue": "HMC gives non deterministic results with parallel chaining\nEven though the `ParallelSampler` yields the same set of traces (each worker gives deterministic results), the order in which the traces are consumed might be different, which gives different results when we use `TracePredictive`. \r\n\r\nOne way to get deterministic results from HMC/NUTS with `num_chains > 1` would be to store intermediate results, and yield results from each of the workers in a round-robin fashion. \n", "before_files": [{"content": "from __future__ import absolute_import, division, print_function\n\nimport errno\nimport json\nimport logging\nimport signal\nimport socket\nimport sys\nimport threading\nimport warnings\nfrom collections import OrderedDict\n\nimport six\nfrom six.moves import queue\nimport torch\nimport torch.multiprocessing as mp\n\nimport pyro\nfrom pyro.infer import TracePosterior\nfrom pyro.infer.mcmc.logger import initialize_logger, initialize_progbar, DIAGNOSTIC_MSG, TqdmHandler\nfrom pyro.util import optional\n\n\ndef logger_thread(log_queue, warmup_steps, num_samples, num_chains):\n    \"\"\"\n    Logging thread that asynchronously consumes logging events from `log_queue`,\n    and handles them appropriately.\n    \"\"\"\n    progress_bars = [initialize_progbar(warmup_steps, s, pos=i)\n                     for i, s in enumerate(num_samples)]\n    logger = logging.getLogger(__name__)\n    logger.propagate = False\n    logger.addHandler(TqdmHandler())\n    num_samples = [0] * len(num_samples)\n    try:\n        while True:\n            try:\n                record = log_queue.get_nowait()\n            except queue.Empty:\n                continue\n            if record is None:\n                break\n            metadata, msg = record.getMessage().split(\"]\", 1)\n            _, msg_type, logger_id = metadata[1:].split()\n            if msg_type == DIAGNOSTIC_MSG:\n                pbar_pos = int(logger_id.split(\":\")[-1]) - 1\n                num_samples[pbar_pos] += 1\n                if num_samples[pbar_pos] == warmup_steps:\n                    progress_bars[pbar_pos].set_description(\"Sample [{}]\".format(pbar_pos + 1))\n                diagnostics = json.loads(msg, object_pairs_hook=OrderedDict)\n                progress_bars[pbar_pos].set_postfix(diagnostics)\n                progress_bars[pbar_pos].update()\n            else:\n                logger.handle(record)\n    finally:\n        for pbar in progress_bars:\n            pbar.close()\n        # Required to not overwrite multiple progress bars on exit.\n        sys.stderr.write(\"\\n\" * num_chains)\n\n\nclass _Worker(object):\n    def __init__(self, chain_id, result_queue, log_queue,\n                 kernel, num_samples, warmup_steps=0,\n                 args=None, kwargs=None):\n        self.chain_id = chain_id\n        self.trace_gen = _SingleSampler(kernel, num_samples=num_samples, warmup_steps=warmup_steps)\n        self.args = args if args is not None else []\n        self.kwargs = kwargs if kwargs is not None else {}\n        self.rng_seed = torch.initial_seed()\n        self.log_queue = log_queue\n        self.result_queue = result_queue\n        self.default_tensor_type = torch.Tensor().type()\n\n    def run(self, *args, **kwargs):\n        pyro.set_rng_seed(self.chain_id + self.rng_seed)\n        torch.set_default_tensor_type(self.default_tensor_type)\n        kwargs[\"logger_id\"] = \"CHAIN:{}\".format(self.chain_id)\n        kwargs[\"log_queue\"] = self.log_queue\n        try:\n            for sample in self.trace_gen._traces(*args, **kwargs):\n                self.result_queue.put_nowait((self.chain_id, sample))\n            self.result_queue.put_nowait((self.chain_id, None))\n        except Exception as e:\n            self.trace_gen.logger.exception(e)\n            self.result_queue.put_nowait((self.chain_id, e))\n\n\nclass _ParallelSampler(TracePosterior):\n    \"\"\"\n    Parallel runner class for running MCMC chains in parallel. This uses the\n    `torch.multiprocessing` module (itself a light wrapper over the python\n    `multiprocessing` module) to spin up parallel workers.\n    \"\"\"\n    def __init__(self, kernel, num_samples, warmup_steps, num_chains, mp_context):\n        super(_ParallelSampler, self).__init__()\n        self.kernel = kernel\n        self.warmup_steps = warmup_steps\n        self.num_chains = num_chains\n        self.workers = []\n        self.ctx = mp\n        if mp_context:\n            if six.PY2:\n                raise ValueError(\"multiprocessing.get_context() is \"\n                                 \"not supported in Python 2.\")\n            self.ctx = mp.get_context(mp_context)\n        self.result_queue = self.ctx.Manager().Queue()\n        self.log_queue = self.ctx.Manager().Queue()\n        self.logger = initialize_logger(logging.getLogger(\"pyro.infer.mcmc\"),\n                                        \"MAIN\", log_queue=self.log_queue)\n        # initialize number of samples per chain\n        samples_per_chain = num_samples // num_chains\n        self.num_samples = [samples_per_chain] * num_chains\n        for i in range(num_samples % num_chains):\n            self.num_samples[i] += 1\n        self.log_thread = threading.Thread(target=logger_thread,\n                                           args=(self.log_queue, self.warmup_steps,\n                                                 self.num_samples, self.num_chains))\n        self.log_thread.daemon = True\n        self.log_thread.start()\n\n    def init_workers(self, *args, **kwargs):\n        self.workers = []\n        for i in range(self.num_chains):\n            worker = _Worker(i + 1, self.result_queue, self.log_queue, self.kernel,\n                             self.num_samples[i], self.warmup_steps)\n            worker.daemon = True\n            self.workers.append(self.ctx.Process(name=str(i), target=worker.run,\n                                                 args=args, kwargs=kwargs))\n\n    def terminate(self):\n        if self.log_thread.is_alive():\n            self.log_queue.put_nowait(None)\n            self.log_thread.join(timeout=1)\n        for w in self.workers:\n            if w.is_alive():\n                w.terminate()\n\n    def _traces(self, *args, **kwargs):\n        # Ignore sigint in worker processes; they will be shut down\n        # when the main process terminates.\n        sigint_handler = signal.signal(signal.SIGINT, signal.SIG_IGN)\n        self.init_workers(*args, **kwargs)\n        # restore original handler\n        signal.signal(signal.SIGINT, sigint_handler)\n        active_workers = self.num_chains\n        try:\n            for w in self.workers:\n                w.start()\n            while active_workers:\n                try:\n                    chain_id, val = self.result_queue.get_nowait()\n                # This can happen when the worker process has terminated.\n                # See https://github.com/pytorch/pytorch/pull/5380 for motivation.\n                except socket.error as e:\n                    if getattr(e, \"errno\", None) == errno.ENOENT:\n                        pass\n                    else:\n                        raise e\n                except queue.Empty:\n                    continue\n                if isinstance(val, Exception):\n                    # Exception trace is already logged by worker.\n                    raise val\n                elif val is not None:\n                    yield val\n                else:\n                    active_workers -= 1\n        finally:\n            self.terminate()\n\n\nclass _SingleSampler(TracePosterior):\n    \"\"\"\n    Single process runner class optimized for the case `num_chains=1`.\n    \"\"\"\n    def __init__(self, kernel, num_samples, warmup_steps):\n        self.kernel = kernel\n        self.warmup_steps = warmup_steps\n        self.num_samples = num_samples\n        self.logger = None\n        super(_SingleSampler, self).__init__()\n\n    def _gen_samples(self, num_samples, init_trace):\n        trace = init_trace\n        for _ in range(num_samples):\n            trace = self.kernel.sample(trace)\n            diagnostics = json.dumps(self.kernel.diagnostics())\n            self.logger.info(diagnostics, extra={\"msg_type\": DIAGNOSTIC_MSG})\n            yield trace\n\n    def _traces(self, *args, **kwargs):\n        logger_id = kwargs.pop(\"logger_id\", \"\")\n        log_queue = kwargs.pop(\"log_queue\", None)\n        self.logger = logging.getLogger(\"pyro.infer.mcmc\")\n        is_multiprocessing = log_queue is not None\n        progress_bar = initialize_progbar(self.warmup_steps, self.num_samples) \\\n            if not is_multiprocessing else None\n        self.logger = initialize_logger(self.logger, logger_id, progress_bar, log_queue)\n        self.kernel.setup(self.warmup_steps, *args, **kwargs)\n        trace = self.kernel.initial_trace()\n        with optional(progress_bar, not is_multiprocessing):\n            for trace in self._gen_samples(self.warmup_steps, trace):\n                continue\n            if progress_bar:\n                progress_bar.set_description(\"Sample\")\n            for trace in self._gen_samples(self.num_samples, trace):\n                yield (trace, 1.0)\n        self.kernel.cleanup()\n\n\nclass MCMC(TracePosterior):\n    \"\"\"\n    Wrapper class for Markov Chain Monte Carlo algorithms. Specific MCMC algorithms\n    are TraceKernel instances and need to be supplied as a ``kernel`` argument\n    to the constructor.\n\n    .. note:: The case of `num_chains > 1` uses python multiprocessing to\n        run parallel chains in multiple processes. This goes with the usual\n        caveats around multiprocessing in python, e.g. the model used to\n        initialize the ``kernel`` must be serializable via `pickle`, and the\n        performance / constraints will be platform dependent (e.g. only\n        the \"spawn\" context is available in Windows). This has also not\n        been extensively tested on the Windows platform.\n\n    :param kernel: An instance of the ``TraceKernel`` class, which when\n        given an execution trace returns another sample trace from the target\n        (posterior) distribution.\n    :param int num_samples: The number of samples that need to be generated,\n        excluding the samples discarded during the warmup phase.\n    :param int warmup_steps: Number of warmup iterations. The samples generated\n        during the warmup phase are discarded. If not provided, default is\n        half of `num_samples`.\n    :param int num_chains: Number of MCMC chains to run in parallel. Depending on\n        whether `num_chains` is 1 or more than 1, this class internally dispatches\n        to either `_SingleSampler` or `_ParallelSampler`.\n    :param str mp_context: Multiprocessing context to use when `num_chains > 1`.\n        Only applicable for Python 3.5 and above. Use `mp_context=\"spawn\"` for\n        CUDA.\n    \"\"\"\n    def __init__(self, kernel, num_samples, warmup_steps=0,\n                 num_chains=1, mp_context=None):\n        self.warmup_steps = warmup_steps if warmup_steps is not None else num_samples // 2  # Stan\n        self.num_samples = num_samples\n        if num_chains > 1:\n            cpu_count = mp.cpu_count()\n            if num_chains > cpu_count:\n                warnings.warn(\"`num_chains` is more than CPU count - {}. \"\n                              \"Resetting num_chains to CPU count.\").format(cpu_count)\n            self.sampler = _ParallelSampler(kernel, num_samples, warmup_steps,\n                                            num_chains, mp_context)\n        else:\n            self.sampler = _SingleSampler(kernel, num_samples, warmup_steps)\n        super(MCMC, self).__init__()\n\n    def _traces(self, *args, **kwargs):\n        for sample in self.sampler._traces(*args, **kwargs):\n            yield sample\n", "path": "pyro/infer/mcmc/mcmc.py"}], "after_files": [{"content": "from __future__ import absolute_import, division, print_function\n\nimport errno\nimport json\nimport logging\nimport signal\nimport socket\nimport sys\nimport threading\nimport warnings\nfrom collections import OrderedDict, deque\n\nimport six\nfrom six.moves import queue\nimport torch\nimport torch.multiprocessing as mp\n\nimport pyro\nfrom pyro.infer import TracePosterior\nfrom pyro.infer.mcmc.logger import initialize_logger, initialize_progbar, DIAGNOSTIC_MSG, TqdmHandler\nfrom pyro.util import optional\n\n\ndef logger_thread(log_queue, warmup_steps, num_samples, num_chains):\n    \"\"\"\n    Logging thread that asynchronously consumes logging events from `log_queue`,\n    and handles them appropriately.\n    \"\"\"\n    progress_bars = [initialize_progbar(warmup_steps, s, pos=i)\n                     for i, s in enumerate(num_samples)]\n    logger = logging.getLogger(__name__)\n    logger.propagate = False\n    logger.addHandler(TqdmHandler())\n    num_samples = [0] * len(num_samples)\n    try:\n        while True:\n            try:\n                record = log_queue.get_nowait()\n            except queue.Empty:\n                continue\n            if record is None:\n                break\n            metadata, msg = record.getMessage().split(\"]\", 1)\n            _, msg_type, logger_id = metadata[1:].split()\n            if msg_type == DIAGNOSTIC_MSG:\n                pbar_pos = int(logger_id.split(\":\")[-1]) - 1\n                num_samples[pbar_pos] += 1\n                if num_samples[pbar_pos] == warmup_steps:\n                    progress_bars[pbar_pos].set_description(\"Sample [{}]\".format(pbar_pos + 1))\n                diagnostics = json.loads(msg, object_pairs_hook=OrderedDict)\n                progress_bars[pbar_pos].set_postfix(diagnostics)\n                progress_bars[pbar_pos].update()\n            else:\n                logger.handle(record)\n    finally:\n        for pbar in progress_bars:\n            pbar.close()\n        # Required to not overwrite multiple progress bars on exit.\n        sys.stderr.write(\"\\n\" * num_chains)\n\n\nclass _Worker(object):\n    def __init__(self, chain_id, result_queue, log_queue,\n                 kernel, num_samples, warmup_steps=0,\n                 args=None, kwargs=None):\n        self.chain_id = chain_id\n        self.trace_gen = _SingleSampler(kernel, num_samples=num_samples, warmup_steps=warmup_steps)\n        self.args = args if args is not None else []\n        self.kwargs = kwargs if kwargs is not None else {}\n        self.rng_seed = torch.initial_seed()\n        self.log_queue = log_queue\n        self.result_queue = result_queue\n        self.default_tensor_type = torch.Tensor().type()\n\n    def run(self, *args, **kwargs):\n        pyro.set_rng_seed(self.chain_id + self.rng_seed)\n        torch.set_default_tensor_type(self.default_tensor_type)\n        kwargs[\"logger_id\"] = \"CHAIN:{}\".format(self.chain_id)\n        kwargs[\"log_queue\"] = self.log_queue\n        try:\n            for sample in self.trace_gen._traces(*args, **kwargs):\n                self.result_queue.put_nowait((self.chain_id, sample))\n            self.result_queue.put_nowait((self.chain_id, None))\n        except Exception as e:\n            self.trace_gen.logger.exception(e)\n            self.result_queue.put_nowait((self.chain_id, e))\n\n\nclass _ParallelSampler(TracePosterior):\n    \"\"\"\n    Parallel runner class for running MCMC chains in parallel. This uses the\n    `torch.multiprocessing` module (itself a light wrapper over the python\n    `multiprocessing` module) to spin up parallel workers.\n    \"\"\"\n    def __init__(self, kernel, num_samples, warmup_steps, num_chains, mp_context):\n        super(_ParallelSampler, self).__init__()\n        self.kernel = kernel\n        self.warmup_steps = warmup_steps\n        self.num_chains = num_chains\n        self.workers = []\n        self.ctx = mp\n        if mp_context:\n            if six.PY2:\n                raise ValueError(\"multiprocessing.get_context() is \"\n                                 \"not supported in Python 2.\")\n            self.ctx = mp.get_context(mp_context)\n        self.result_queue = self.ctx.Manager().Queue()\n        self.log_queue = self.ctx.Manager().Queue()\n        self.logger = initialize_logger(logging.getLogger(\"pyro.infer.mcmc\"),\n                                        \"MAIN\", log_queue=self.log_queue)\n        # initialize number of samples per chain\n        samples_per_chain = num_samples // num_chains\n        self.num_samples = [samples_per_chain] * num_chains\n        for i in range(num_samples % num_chains):\n            self.num_samples[i] += 1\n        self.log_thread = threading.Thread(target=logger_thread,\n                                           args=(self.log_queue, self.warmup_steps,\n                                                 self.num_samples, self.num_chains))\n        self.log_thread.daemon = True\n        self.log_thread.start()\n\n    def init_workers(self, *args, **kwargs):\n        self.workers = []\n        for i in range(self.num_chains):\n            worker = _Worker(i + 1, self.result_queue, self.log_queue, self.kernel,\n                             self.num_samples[i], self.warmup_steps)\n            worker.daemon = True\n            self.workers.append(self.ctx.Process(name=str(i), target=worker.run,\n                                                 args=args, kwargs=kwargs))\n\n    def terminate(self):\n        if self.log_thread.is_alive():\n            self.log_queue.put_nowait(None)\n            self.log_thread.join(timeout=1)\n        for w in self.workers:\n            if w.is_alive():\n                w.terminate()\n\n    def _traces(self, *args, **kwargs):\n        # Ignore sigint in worker processes; they will be shut down\n        # when the main process terminates.\n        sigint_handler = signal.signal(signal.SIGINT, signal.SIG_IGN)\n        self.init_workers(*args, **kwargs)\n        # restore original handler\n        signal.signal(signal.SIGINT, sigint_handler)\n        active_workers = self.num_chains\n        # To yield a deterministic ordering, we hold intermediate traces\n        # from each of the workers in its own queue in `results_buffer`\n        # and yield these in a round robin fashion.\n        buffer_idx = 0\n        results_buffer = [deque() for _ in range(self.num_chains)]\n        try:\n            for w in self.workers:\n                w.start()\n            while active_workers:\n                try:\n                    chain_id, val = self.result_queue.get_nowait()\n                # This can happen when the worker process has terminated.\n                # See https://github.com/pytorch/pytorch/pull/5380 for motivation.\n                except socket.error as e:\n                    if getattr(e, \"errno\", None) == errno.ENOENT:\n                        pass\n                    else:\n                        raise e\n                except queue.Empty:\n                    continue\n                if isinstance(val, Exception):\n                    # Exception trace is already logged by worker.\n                    raise val\n                if val is None:\n                    active_workers -= 1\n                else:\n                    results_buffer[chain_id - 1].append(val)\n                while results_buffer[buffer_idx]:\n                    yield results_buffer[buffer_idx].popleft()\n                    buffer_idx = (buffer_idx + 1) % self.num_chains\n            # empty out the results buffer\n            non_empty_buffers = set(range(self.num_chains))\n            while non_empty_buffers:\n                if results_buffer[buffer_idx]:\n                    yield results_buffer[buffer_idx].popleft()\n                else:\n                    if buffer_idx in non_empty_buffers:\n                        non_empty_buffers.remove(buffer_idx)\n                buffer_idx = (buffer_idx + 1) % self.num_chains\n        finally:\n            self.terminate()\n\n\nclass _SingleSampler(TracePosterior):\n    \"\"\"\n    Single process runner class optimized for the case `num_chains=1`.\n    \"\"\"\n    def __init__(self, kernel, num_samples, warmup_steps):\n        self.kernel = kernel\n        self.warmup_steps = warmup_steps\n        self.num_samples = num_samples\n        self.logger = None\n        super(_SingleSampler, self).__init__()\n\n    def _gen_samples(self, num_samples, init_trace):\n        trace = init_trace\n        for _ in range(num_samples):\n            trace = self.kernel.sample(trace)\n            diagnostics = json.dumps(self.kernel.diagnostics())\n            self.logger.info(diagnostics, extra={\"msg_type\": DIAGNOSTIC_MSG})\n            yield trace\n\n    def _traces(self, *args, **kwargs):\n        logger_id = kwargs.pop(\"logger_id\", \"\")\n        log_queue = kwargs.pop(\"log_queue\", None)\n        self.logger = logging.getLogger(\"pyro.infer.mcmc\")\n        is_multiprocessing = log_queue is not None\n        progress_bar = initialize_progbar(self.warmup_steps, self.num_samples) \\\n            if not is_multiprocessing else None\n        self.logger = initialize_logger(self.logger, logger_id, progress_bar, log_queue)\n        self.kernel.setup(self.warmup_steps, *args, **kwargs)\n        trace = self.kernel.initial_trace()\n        with optional(progress_bar, not is_multiprocessing):\n            for trace in self._gen_samples(self.warmup_steps, trace):\n                continue\n            if progress_bar:\n                progress_bar.set_description(\"Sample\")\n            for trace in self._gen_samples(self.num_samples, trace):\n                yield (trace, 1.0)\n        self.kernel.cleanup()\n\n\nclass MCMC(TracePosterior):\n    \"\"\"\n    Wrapper class for Markov Chain Monte Carlo algorithms. Specific MCMC algorithms\n    are TraceKernel instances and need to be supplied as a ``kernel`` argument\n    to the constructor.\n\n    .. note:: The case of `num_chains > 1` uses python multiprocessing to\n        run parallel chains in multiple processes. This goes with the usual\n        caveats around multiprocessing in python, e.g. the model used to\n        initialize the ``kernel`` must be serializable via `pickle`, and the\n        performance / constraints will be platform dependent (e.g. only\n        the \"spawn\" context is available in Windows). This has also not\n        been extensively tested on the Windows platform.\n\n    :param kernel: An instance of the ``TraceKernel`` class, which when\n        given an execution trace returns another sample trace from the target\n        (posterior) distribution.\n    :param int num_samples: The number of samples that need to be generated,\n        excluding the samples discarded during the warmup phase.\n    :param int warmup_steps: Number of warmup iterations. The samples generated\n        during the warmup phase are discarded. If not provided, default is\n        half of `num_samples`.\n    :param int num_chains: Number of MCMC chains to run in parallel. Depending on\n        whether `num_chains` is 1 or more than 1, this class internally dispatches\n        to either `_SingleSampler` or `_ParallelSampler`.\n    :param str mp_context: Multiprocessing context to use when `num_chains > 1`.\n        Only applicable for Python 3.5 and above. Use `mp_context=\"spawn\"` for\n        CUDA.\n    \"\"\"\n    def __init__(self, kernel, num_samples, warmup_steps=0,\n                 num_chains=1, mp_context=None):\n        self.warmup_steps = warmup_steps if warmup_steps is not None else num_samples // 2  # Stan\n        self.num_samples = num_samples\n        if num_chains > 1:\n            cpu_count = mp.cpu_count()\n            if num_chains > cpu_count:\n                warnings.warn(\"`num_chains` is more than CPU count - {}. \"\n                              \"Resetting num_chains to CPU count.\").format(cpu_count)\n            self.sampler = _ParallelSampler(kernel, num_samples, warmup_steps,\n                                            num_chains, mp_context)\n        else:\n            self.sampler = _SingleSampler(kernel, num_samples, warmup_steps)\n        super(MCMC, self).__init__()\n\n    def _traces(self, *args, **kwargs):\n        for sample in self.sampler._traces(*args, **kwargs):\n            yield sample\n", "path": "pyro/infer/mcmc/mcmc.py"}]}