Add support for DistributedNext

Project.toml

@@ -6,12 +6,14 @@ version = "0.18.13"
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
+DistributedNext = "fab6aee4-877b-4bac-a744-3eca44acbb6f"
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 MemPool = "f9f48841-c794-520a-933b-121f7ba6ed94"
 OnlineStats = "a15396b6-48d5-5d58-9928-6d29437db91e"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
+Preferences = "21216c6a-2e73-6563-6e65-726566657250"
 Profile = "9abbd945-dff8-562f-b5e8-e1ebf5ef1b79"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
@@ -47,15 +49,17 @@ Adapt = "4.0.4"
 Colors = "0.12, 0.13"
 DataFrames = "1"
 DataStructures = "0.18"
+DistributedNext = "1.0.0"
 Distributions = "0.25"
 GraphViz = "0.2"
 Graphs = "1"
 JSON3 = "1"
 MacroTools = "0.5"
-MemPool = "0.4.6"
+MemPool = "0.4.11"
 OnlineStats = "1"
 Plots = "1"
 PrecompileTools = "1.2"
+Preferences = "1.4.3"
 PythonCall = "0.9"
 Requires = "1"
 ScopedValues = "1.1"
@@ -64,12 +68,3 @@ StatsBase = "0.28, 0.29, 0.30, 0.31, 0.32, 0.33, 0.34"
 TaskLocalValues = "0.1"
 TimespanLogging = "0.1"
 julia = "1.9"
-
-[extras]
-Colors = "5ae59095-9a9b-59fe-a467-6f913c188581"
-DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
-Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
-GraphViz = "f526b714-d49f-11e8-06ff-31ed36ee7ee0"
-JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1"
-Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
-PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d"

lib/TimespanLogging/Project.toml

@@ -5,7 +5,11 @@ version = "0.1.0"
 
 [deps]
 Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
+DistributedNext = "fab6aee4-877b-4bac-a744-3eca44acbb6f"
+Preferences = "21216c6a-2e73-6563-6e65-726566657250"
 Profile = "9abbd945-dff8-562f-b5e8-e1ebf5ef1b79"
 
 [compat]
+DistributedNext = "1.0.0"
+Preferences = "1.4.3"
 julia = "1"

lib/TimespanLogging/src/core.jl

@@ -1,4 +1,10 @@
-using Distributed
+import Preferences: @load_preference, @set_preferences!
+if @load_preference("distributed-package") == "DistributedNext"
+    using DistributedNext
+else
+    using Distributed
+end
+
 import Profile
 import Base.gc_num
 
@@ -16,6 +22,18 @@ ProfilerResult(samples, lineinfo, tasks::Vector{Task}) =
 ProfilerResult(samples, lineinfo, tasks::Nothing) =
     ProfilerResult(samples, lineinfo, map(Base.pointer_from_objref, UInt[]))
 
+"""
+    set_distributed_package!(value[="Distributed|DistributedNext"])
+
+Set a [preference](https://github.com/JuliaPackaging/Preferences.jl) for using
+either the Distributed.jl stdlib or DistributedNext.jl. You will need to restart
+Julia after setting a new preference.
+"""
+function set_distributed_package!(value)
+    @set_preferences!("distributed-package" => value)
+    @info "TimespanLogging.jl preference has been set, restart your Julia session for this change to take effect!"
+end
+
 """
     Timespan
 

src/Dagger.jl

@@ -8,8 +8,6 @@ import MemPool
 import MemPool: DRef, FileRef, poolget, poolset
 
 import Base: collect, reduce
-import Distributed
-import Distributed: Future, RemoteChannel, myid, workers, nworkers, procs, remotecall, remotecall_wait, remotecall_fetch
 
 import LinearAlgebra
 import LinearAlgebra: Adjoint, BLAS, Diagonal, Bidiagonal, Tridiagonal, LAPACK, LowerTriangular, PosDefException, Transpose, UpperTriangular, UnitLowerTriangular, UnitUpperTriangular, diagind, ishermitian, issymmetric
@@ -33,6 +31,17 @@ import TimespanLogging: timespan_start, timespan_finish
 
 import Adapt
 
+# Preferences
+import Preferences: @load_preference, @set_preferences!
+
+if @load_preference("distributed-package") == "DistributedNext"
+    import DistributedNext
+    import DistributedNext: Future, RemoteChannel, myid, workers, nworkers, procs, remotecall, remotecall_wait, remotecall_fetch, check_same_host
+else
+    import Distributed
+    import Distributed: Future, RemoteChannel, myid, workers, nworkers, procs, remotecall, remotecall_wait, remotecall_fetch, check_same_host
+end
+
 include("lib/util.jl")
 include("utils/dagdebug.jl")
 
@@ -96,6 +105,21 @@ include("utils/logging-events.jl")
 include("utils/logging.jl")
 include("utils/viz.jl")
 
+"""
+    set_distributed_package!(value[="Distributed|DistributedNext"])
+
+Set a [preference](https://github.com/JuliaPackaging/Preferences.jl) for using
+either the Distributed.jl stdlib or DistributedNext.jl. You will need to restart
+Julia after setting a new preference.
+"""
+function set_distributed_package!(value)
+    MemPool.set_distributed_package!(value)
+    TimespanLogging.set_distributed_package!(value)
+
+    @set_preferences!("distributed-package" => value)
+    @info "Dagger.jl preference has been set, restart your Julia session for this change to take effect!"
+end
+
 # Precompilation
 import PrecompileTools: @compile_workload
 include("precompile.jl")

src/array/darray.jl

@@ -434,7 +434,7 @@ end
 
 function Distribute(data::AbstractArray{T,N}) where {T,N}
     nprocs = sum(w->length(Dagger.get_processors(OSProc(w))),
-                 Distributed.procs())
+                 procs())
     p = Blocks(ntuple(i->max(cld(size(data, i), nprocs), 1), N))
     return Distribute(partition(p, domain(data)), p, data)
 end

src/array/sort.jl

@@ -144,8 +144,8 @@ const use_shared_array = Ref(true)
 function _promote_array(x::AbstractArray{T}, y::AbstractArray{S}) where {T,S}
     Q = promote_type(T,S)
     ok = (isa(x, Array) || isa(x, SharedArray)) && (isa(y, Array) || isa(y, SharedArray))
-    if ok && isbitstype(Q) && use_shared_array[] && Distributed.check_same_host([workers()..., 1])
-        return SharedArray{Q}(length(x)+length(y), pids=Distributed.procs())
+    if ok && isbitstype(Q) && use_shared_array[] && check_same_host([workers()..., 1])
+        return SharedArray{Q}(length(x)+length(y), pids=procs())
     else
         return similar(x, Q, length(x)+length(y))
     end

src/sch/Sch.jl

@@ -1,6 +1,12 @@
 module Sch
 
-import Distributed: Future, ProcessExitedException, RemoteChannel, RemoteException, myid, remote_do, remotecall_fetch, remotecall_wait, workers
+import Preferences: @load_preference
+if @load_preference("distributed-package") == "DistributedNext"
+    import DistributedNext: Future, ProcessExitedException, RemoteChannel, RemoteException, myid, remote_do, remotecall_fetch, remotecall_wait, workers
+else
+    import Distributed: Future, ProcessExitedException, RemoteChannel, RemoteException, myid, remote_do, remotecall_fetch, remotecall_wait, workers
+end
+
 import MemPool
 import MemPool: DRef, StorageResource
 import MemPool: poolset, storage_capacity, storage_utilized

test/Project.toml

@@ -13,6 +13,7 @@ MemPool = "f9f48841-c794-520a-933b-121f7ba6ed94"
 OnlineStats = "a15396b6-48d5-5d58-9928-6d29437db91e"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
+Preferences = "21216c6a-2e73-6563-6e65-726566657250"
 PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Serialization = "9e88b42a-f829-5b0c-bbe9-9e923198166b"

test/preferences.jl

@@ -0,0 +1,21 @@
+import Preferences: load_preference, set_preferences!
+
+@testset "Preferences" begin
+    cmd = `$(Base.julia_cmd()) --startup-file=no --project -E 'using Dagger; parentmodule(Dagger.myid)'`
+
+    try
+        # Disabling the precompilation workload shaves off over half the time
+        # this test takes.
+        set_preferences!(Dagger, "precompile_workload" => false; force=true)
+
+        cd(dirname(Base.active_project())) do
+            Dagger.set_distributed_package!("Distributed")
+            @test readchomp(cmd) == "Distributed"
+
+            Dagger.set_distributed_package!("DistributedNext")
+            @test readchomp(cmd) == "DistributedNext"
+        end
+    finally
+        set_preferences!(Dagger, "precompile_workload" => true; force=true)
+    end
+end

test/runtests.jl

@@ -24,6 +24,7 @@ tests = [
     ("Disk Caching", "diskcaching.jl"),
     ("File IO", "file-io.jl"),
     ("External Languages - Python", "extlang/python.jl"),
+    ("Preferences", "preferences.jl")
     #("Fault Tolerance", "fault-tolerance.jl"),
 ]
 all_test_names = map(test -> replace(last(test), ".jl"=>""), tests)