Merge pull request #33 from TuringLang/distributed

Add multicore sampling

Author: cpfiffer

src/AbstractMCMC.jl

@@ -4,13 +4,19 @@ import ConsoleProgressMonitor
 import LoggingExtras
 import ProgressLogging
 import StatsBase
-using StatsBase: sample
 import TerminalLoggers
 
 import Distributed
 import Logging
 import Random
 
+# Reexport sample
+using StatsBase: sample
+export sample
+
+# Parallel sampling types
+export MCMCThreads, MCMCDistributed
+
 """
     AbstractChains
 
@@ -39,6 +45,30 @@ An `AbstractModel` represents a generic model type that can be used to perform i
 """
 abstract type AbstractModel end
 
+"""
+    AbstractMCMCParallel
+
+An `AbstractMCMCParallel` algorithm represents a specific algorithm for sampling MCMC chains
+in parallel.
+"""
+abstract type AbstractMCMCParallel end
+
+"""
+    MCMCThreads
+
+The `MCMCThreads` algorithm allows to sample MCMC chains in parallel using multiple
+threads.
+"""
+struct MCMCThreads <: AbstractMCMCParallel end
+
+"""
+    MCMCDistributed
+
+The `MCMCDistributed` algorithm allows to sample MCMC chains in parallel using multiple
+processes.
+"""
+struct MCMCDistributed <: AbstractMCMCParallel end
+
 include("logging.jl")
 include("interface.jl")
 include("sample.jl")

src/sample.jl

@@ -1,4 +1,4 @@
-# Default implementations of `sample` and `psample`.
+# Default implementations of `sample`.
 
 function StatsBase.sample(
     model::AbstractModel,
@@ -19,25 +19,28 @@ function StatsBase.sample(
     return mcmcsample(rng, model, sampler, arg; kwargs...)
 end
 
-function psample(
+function StatsBase.sample(
     model::AbstractModel,
     sampler::AbstractSampler,
+    parallel::AbstractMCMCParallel,
     N::Integer,
     nchains::Integer;
     kwargs...
 )
-    return psample(Random.GLOBAL_RNG, model, sampler, N, nchains; kwargs...)
+    return StatsBase.sample(Random.GLOBAL_RNG, model, sampler, parallel, N, nchains;
+                            kwargs...)
 end
 
-function psample(
+function StatsBase.sample(
     rng::Random.AbstractRNG,
     model::AbstractModel,
     sampler::AbstractSampler,
+    parallel::AbstractMCMCParallel,
     N::Integer,
     nchains::Integer;
     kwargs...
 )
-    return mcmcpsample(rng, model, sampler, N, nchains; kwargs...)
+    return mcmcsample(rng, model, sampler, parallel, N, nchains; kwargs...)
 end
 
 # Default implementations of regular and parallel sampling.
@@ -173,23 +176,27 @@ function mcmcsample(
 end
 
 """
-    mcmcpsample([rng, ]model, sampler, N, nchains; kwargs...)
+    mcmcsample([rng, ]model, sampler, parallel, N, nchains; kwargs...)
 
-Sample `nchains` chains using the available threads, and combine them into a single chain.
-
-By default, the random number generator, the model and the samplers are deep copied for each
-thread to prevent contamination between threads.
+Sample `nchains` chains in parallel using the `parallel` algorithm, and combine them into a
+single chain.
 """
-function mcmcpsample(
+function mcmcsample(
     rng::Random.AbstractRNG,
     model::AbstractModel,
     sampler::AbstractSampler,
+    ::MCMCThreads,
     N::Integer,
     nchains::Integer;
     progress = true,
-    progressname = "Parallel sampling",
+    progressname = "Sampling ($(Threads.nthreads()) threads)",
     kwargs...
 )
+    # Check if actually multiple threads are used.
+    if Threads.nthreads() == 1
+        @warn "Only a single thread available: MCMC chains are not sampled in parallel"
+    end
+
     # Copy the random number generator, model, and sample for each thread
     rngs = [deepcopy(rng) for _ in 1:Threads.nthreads()]
     models = [deepcopy(model) for _ in 1:Threads.nthreads()]
@@ -204,7 +211,7 @@ function mcmcpsample(
     @ifwithprogresslogger progress name=progressname begin
         # Create a channel for progress logging.
         if progress
-            channel = Distributed.RemoteChannel(() -> Channel{Bool}(nchains), 1)
+            channel = Distributed.RemoteChannel(() -> Channel{Bool}(nchains))
         end
 
         Distributed.@sync begin
@@ -245,3 +252,77 @@ function mcmcpsample(
     # Concatenate the chains together.
     return reduce(chainscat, chains)
 end
+
+function mcmcsample(
+    rng::Random.AbstractRNG,
+    model::AbstractModel,
+    sampler::AbstractSampler,
+    ::MCMCDistributed,
+    N::Integer,
+    nchains::Integer;
+    progress = true,
+    progressname = "Sampling ($(Distributed.nworkers()) processes)",
+    kwargs...
+)
+    # Check if actually multiple processes are used.
+    if Distributed.nworkers() == 1
+        @warn "Only a single process available: MCMC chains are not sampled in parallel"
+    end
+
+    # Create a seed for each chain using the provided random number generator.
+    seeds = rand(rng, UInt, nchains)
+
+    # Set up worker pool.
+    pool = Distributed.CachingPool(Distributed.workers())
+
+    # Create a channel for progress logging.
+    channel = progress ? Distributed.RemoteChannel(() -> Channel{Bool}(nchains)) : nothing
+
+    local chains
+    @ifwithprogresslogger progress name=progressname begin
+        Distributed.@sync begin
+            # Update the progress bar.
+            if progress
+                Distributed.@async begin
+                    progresschains = 0
+                    while take!(channel)
+                        progresschains += 1
+                        ProgressLogging.@logprogress progresschains/nchains
+                    end
+                end
+            end
+
+            Distributed.@async begin
+                chains = let rng=rng, model=model, sampler=sampler, N=N, channel=channel,
+                    kwargs=kwargs
+                    Distributed.pmap(pool, seeds) do seed
+                        # Seed a new random number generator with the pre-made seed.
+                        subrng = deepcopy(rng)
+                        Random.seed!(subrng, seed)
+
+                        # Sample a chain.
+                        chain = StatsBase.sample(subrng, model, sampler, N;
+                                                 progress = false, kwargs...)
+
+                        # Update the progress bar.
+                        channel === nothing || put!(channel, true)
+
+                        # Return the new chain.
+                        return chain
+                    end
+                end
+
+                # Stop updating the progress bar.
+                progress && put!(channel, false)
+            end
+        end
+    end
+
+    # Concatenate the chains together.
+    return reduce(chainscat, chains)
+end
+
+# Deprecations.
+Base.@deprecate psample(model, sampler, N, nchains; kwargs...) sample(model, sampler, MCMCThreads(), N, nchains; kwargs...) false
+Base.@deprecate psample(rng, model, sampler, N, nchains; kwargs...) sample(rng, model, sampler, MCMCThreads(), N, nchains; kwargs...) false
+Base.@deprecate mcmcpsample(rng, model, sampler, N, nchains; kwargs...) mcmcsample(rng, model, sampler, MCMCThreads(), N, nchains; kwargs...) false

test/runtests.jl

@@ -1,11 +1,12 @@
 using AbstractMCMC
-using AbstractMCMC: sample, psample, steps!
+using AbstractMCMC: steps!
 using Atom.Progress: JunoProgressLogger
 using ConsoleProgressMonitor: ProgressLogger
 using IJulia
 using LoggingExtras: TeeLogger, EarlyFilteredLogger
 using TerminalLoggers: TerminalLogger
 
+using Distributed
 import Logging
 using Random
 using Statistics
@@ -99,14 +100,26 @@ include("interface.jl")
             @test first(LOGGERS) === logger
             @test Logging.current_logger() === CURRENT_LOGGER
         end
+
+        @testset "Suppress output" begin
+            logs, _ = collect_test_logs(; min_level=Logging.LogLevel(-1)) do
+                sample(MyModel(), MySampler(), 100; progress = false, sleepy = true)
+            end
+            @test all(l.level > Logging.LogLevel(-1) for l in logs)
+        end
     end
 
     if VERSION ≥ v"1.3"
-        @testset "Parallel sampling" begin
-            println("testing parallel sampling with ", Threads.nthreads(), " thread(s)...")
+        @testset "Multithreaded sampling" begin
+            if Threads.nthreads() == 1
+                warnregex = r"^Only a single thread available"
+                @test_logs (:warn, warnregex) sample(MyModel(), MySampler(), MCMCThreads(),
+                                                     10, 10; chain_type = MyChain)
+            end
 
             Random.seed!(1234)
-            chains = psample(MyModel(), MySampler(), 10_000, 1000; chain_type = MyChain)
+            chains = sample(MyModel(), MySampler(), MCMCThreads(), 10_000, 1000;
+                            chain_type = MyChain)
 
             # test output type and size
             @test chains isa Vector{MyChain}
@@ -121,12 +134,69 @@ include("interface.jl")
 
             # test reproducibility
             Random.seed!(1234)
-            chains2 = psample(MyModel(), MySampler(), 10_000, 1000; chain_type = MyChain)
+            chains2 = sample(MyModel(), MySampler(), MCMCThreads(), 10_000, 1000;
+                             chain_type = MyChain)
 
             @test all(((x, y),) -> x.as == y.as && x.bs == y.bs, zip(chains, chains2))
+
+            # Suppress output.
+            logs, _ = collect_test_logs(; min_level=Logging.LogLevel(-1)) do
+                sample(MyModel(), MySampler(), MCMCThreads(), 10_000, 1000;
+                        progress = false, chain_type = MyChain)
+            end
+            @test all(l.level > Logging.LogLevel(-1) for l in logs)
         end
     end
 
+    @testset "Multicore sampling" begin
+        if nworkers() == 1
+            warnregex = r"^Only a single process available"
+            @test_logs (:warn, warnregex) sample(MyModel(), MySampler(), MCMCDistributed(),
+                                                 10, 10; chain_type = MyChain)
+        end
+
+        # Add worker processes.
+        addprocs()
+
+        # Load all required packages (`interface.jl` needs Random).
+        @everywhere begin
+            using AbstractMCMC
+            using AbstractMCMC: sample
+
+            using Random
+            include("interface.jl")
+        end
+
+        Random.seed!(1234)
+        chains = sample(MyModel(), MySampler(), MCMCDistributed(), 10_000, 1000;
+                        chain_type = MyChain)
+
+        # Test output type and size.
+        @test chains isa Vector{MyChain}
+        @test length(chains) == 1000
+        @test all(x -> length(x.as) == length(x.bs) == 10_000, chains)
+
+        # Test some statistical properties.
+        @test all(x -> isapprox(mean(x.as), 0.5; atol=1e-2), chains)
+        @test all(x -> isapprox(var(x.as), 1 / 12; atol=5e-3), chains)
+        @test all(x -> isapprox(mean(x.bs), 0; atol=5e-2), chains)
+        @test all(x -> isapprox(var(x.bs), 1; atol=5e-2), chains)
+
+        # Test reproducibility.
+        Random.seed!(1234)
+        chains2 = sample(MyModel(), MySampler(), MCMCDistributed(), 10_000, 1000;
+                         chain_type = MyChain)
+
+        @test all(((x, y),) -> x.as == y.as && x.bs == y.bs, zip(chains, chains2))
+
+        # Suppress output.
+        logs, _ = collect_test_logs(; min_level=Logging.LogLevel(-1)) do
+            sample(MyModel(), MySampler(), MCMCDistributed(), 10_000, 100;
+                   progress = false, chain_type = MyChain)
+        end
+        @test all(l.level > Logging.LogLevel(-1) for l in logs)
+    end
+
     @testset "Chain constructors" begin
         chain1 = sample(MyModel(), MySampler(), 100; sleepy = true)
         chain2 = sample(MyModel(), MySampler(), 100; sleepy = true, chain_type = MyChain)
@@ -135,21 +205,6 @@ include("interface.jl")
         @test chain2 isa MyChain
     end
 
-    @testset "Suppress output" begin
-        logs, _ = collect_test_logs(; min_level=Logging.LogLevel(-1)) do
-            sample(MyModel(), MySampler(), 100; progress = false, sleepy = true)
-        end
-        @test isempty(logs)
-
-        if VERSION ≥ v"1.3"
-            logs, _ = collect_test_logs(; min_level=Logging.LogLevel(-1)) do
-                psample(MyModel(), MySampler(), 10_000, 1000;
-                        progress = false, chain_type = MyChain)
-            end
-            @test isempty(logs)
-        end
-    end
-
     @testset "Iterator sampling" begin
         Random.seed!(1234)
         as = []
@@ -182,4 +237,15 @@ include("interface.jl")
         bmean = mean(x.b for x in chain)
         @test abs(bmean) <= 0.001 && length(chain) < 10_000
     end
+
+    @testset "Deprecations" begin
+        @test_deprecated AbstractMCMC.psample(MyModel(), MySampler(), 10, 10;
+                                              chain_type = MyChain)
+        @test_deprecated AbstractMCMC.psample(Random.GLOBAL_RNG, MyModel(), MySampler(),
+                                              10, 10;
+                                              chain_type = MyChain)
+        @test_deprecated AbstractMCMC.mcmcpsample(Random.GLOBAL_RNG, MyModel(),
+                                                  MySampler(), 10, 10;
+                                                  chain_type = MyChain)
+    end
 end