Merge pull request #361 from JuliaParallel/kr/enable_caching

`enable_disk_caching` function for handling MemPool storage setup

Author: jpsamaroo

.buildkite/pipeline.yml

@@ -76,7 +76,7 @@ steps:
           version: "1.7"
       - JuliaCI/julia-coverage#v1:
           codecov: true
-    command: "julia --project -e 'using Pkg; Pkg.develop(;path=\"lib/TimespanLogging\"); Pkg.test(\"TimespanLogging\")'"
+    command: "julia --project -e 'using Pkg; Pkg.instantiate(); Pkg.develop(;path=\"lib/TimespanLogging\"); Pkg.test(\"TimespanLogging\")'"
   - label: Julia 1.7 - DaggerWebDash
     timeout_in_minutes: 60
     <<: *test

Manifest.toml

@@ -17,10 +17,10 @@ uuid = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 version = "1.16.0"
 
 [[deps.ChangesOfVariables]]
-deps = ["LinearAlgebra", "Test"]
-git-tree-sha1 = "f84967c4497e0e1955f9a582c232b02847c5f589"
+deps = ["InverseFunctions", "LinearAlgebra", "Test"]
+git-tree-sha1 = "2fba81a302a7be671aefe194f0525ef231104e7f"
 uuid = "9e997f8a-9a97-42d5-a9f1-ce6bfc15e2c0"
-version = "0.1.7"
+version = "0.1.8"
 
 [[deps.ColorTypes]]
 deps = ["FixedPointNumbers", "Random"]
@@ -36,9 +36,9 @@ version = "0.12.10"
 
 [[deps.Compat]]
 deps = ["Dates", "LinearAlgebra", "UUIDs"]
-git-tree-sha1 = "7a60c856b9fa189eb34f5f8a6f6b5529b7942957"
+git-tree-sha1 = "4e88377ae7ebeaf29a047aa1ee40826e0b708a5d"
 uuid = "34da2185-b29b-5c13-b0c7-acf172513d20"
-version = "4.6.1"
+version = "4.7.0"
 
 [[deps.CompilerSupportLibraries_jll]]
 deps = ["Artifacts", "Libdl"]
@@ -57,9 +57,9 @@ version = "1.15.0"
 
 [[deps.DataStructures]]
 deps = ["Compat", "InteractiveUtils", "OrderedCollections"]
-git-tree-sha1 = "d1fff3a548102f48987a52a2e0d114fa97d730f0"
+git-tree-sha1 = "cf25ccb972fec4e4817764d01c82386ae94f77b4"
 uuid = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
-version = "0.18.13"
+version = "0.18.14"
 
 [[deps.Dates]]
 deps = ["Printf"]
@@ -87,9 +87,9 @@ uuid = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 
 [[deps.InverseFunctions]]
 deps = ["Test"]
-git-tree-sha1 = "6667aadd1cdee2c6cd068128b3d226ebc4fb0c67"
+git-tree-sha1 = "eabe3125edba5c9c10b60a160b1779a000dc8b29"
 uuid = "3587e190-3f89-42d0-90ee-14403ec27112"
-version = "0.1.9"
+version = "0.1.11"
 
 [[deps.IrrationalConstants]]
 git-tree-sha1 = "630b497eafcc20001bba38a4651b327dcfc491d2"
@@ -128,7 +128,9 @@ uuid = "d6f4376e-aef5-505a-96c1-9c027394607a"
 
 [[deps.MemPool]]
 deps = ["DataStructures", "Distributed", "Mmap", "Random", "Serialization", "Sockets"]
-git-tree-sha1 = "303f55cede4947d0acb325ac0bd7595fbcad6b99"
+git-tree-sha1 = "8862fe03a7edf2faa9cf4d0584d453a05377bf9c"
+repo-rev = "master"
+repo-url = "https://github.com/JuliaData/MemPool.jl.git"
 uuid = "f9f48841-c794-520a-933b-121f7ba6ed94"
 version = "0.4.2"
 
@@ -149,9 +151,9 @@ deps = ["Artifacts", "CompilerSupportLibraries_jll", "Libdl"]
 uuid = "4536629a-c528-5b80-bd46-f80d51c5b363"
 
 [[deps.OrderedCollections]]
-git-tree-sha1 = "d321bf2de576bf25ec4d3e4360faca399afca282"
+git-tree-sha1 = "2e73fe17cac3c62ad1aebe70d44c963c3cfdc3e3"
 uuid = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
-version = "1.6.0"
+version = "1.6.2"
 
 [[deps.Printf]]
 deps = ["Unicode"]
@@ -191,9 +193,9 @@ uuid = "6462fe0b-24de-5631-8697-dd941f90decc"
 
 [[deps.SortingAlgorithms]]
 deps = ["DataStructures"]
-git-tree-sha1 = "a4ada03f999bd01b3a25dcaa30b2d929fe537e00"
+git-tree-sha1 = "c60ec5c62180f27efea3ba2908480f8055e17cee"
 uuid = "a2af1166-a08f-5f64-846c-94a0d3cef48c"
-version = "1.1.0"
+version = "1.1.1"
 
 [[deps.SparseArrays]]
 deps = ["LinearAlgebra", "Random"]

Project.toml

@@ -34,6 +34,7 @@ julia = "1.7"
 
 [extras]
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 
 [targets]
-test = ["Test"]
+test = ["Test", "Pkg"]

src/Dagger.jl

@@ -36,6 +36,7 @@ include("compute.jl")
 include("utils/clock.jl")
 include("utils/system_uuid.jl")
 include("utils/locked-object.jl")
+include("utils/caching.jl")
 include("sch/Sch.jl"); using .Sch
 
 # Array computations

src/utils/caching.jl

@@ -0,0 +1,94 @@
+"""
+    mem_limit(total::UInt, percentage_limit::Int, nprocs::Int)
+
+Returns the per process mem limit in MiB based on the provided `total` memory,
+`percentage_limit` (1-100) and the number of processes in the current setup.
+"""
+function mem_limit(total::UInt, percentage_limit::Int, nprocs::Int)
+    return total * percentage_limit / 100 ÷ nprocs
+end
+
+"""
+    enable_disk_caching!(ram_percentage_limit::Int=30, disk_limit_mb::Int=16*2^10, processes::Vector{Int}=procs())
+
+Sets up disk caching on all processes available in the environment according to the provided
+limits. The user should provide the percentage, which will decide what's the memory limit on
+each participating machine (differentiated by hostname). The disk limit is set strictly per
+process and doesn't include any hostname related logic.
+"""
+function enable_disk_caching!(
+    ram_percentage_limit::Int=30,
+    disk_limit_mb::Int=16*2^10,
+    processes::Vector{Int}=procs(),
+)
+    if !(0 < ram_percentage_limit <= 100)
+        throw(ArgumentError("Ram limit values must be in (1, 100] range"))
+    end
+
+    total_mem = @static VERSION >= v"1.8-" ? Sys.total_physical_memory : Sys.total_memory
+    process_info = [
+        id => remotecall(id, total_mem) do total_mem
+            return (; total_memory=total_mem(), hostname=gethostname())
+        end for id in processes
+    ]
+
+    machines = Dict{NamedTuple,Vector{Int}}()
+    for (id, info) in process_info
+        key = fetch(info)
+        if key isa NamedTuple
+            machines[key] = push!(get(machines, key, Int[]), id)
+        else
+            @error("Error querying hardware information on process id = $id", ex = key)
+        end
+    end
+
+    mem_limits = Dict{Int,Int}()
+    for (info, ids) in machines
+        for id in ids
+            mem_limits[id] = mem_limit(info.total_memory, ram_percentage_limit, length(ids))
+        end
+    end
+
+    return enable_disk_caching!(mem_limits, disk_limit_mb, processes)
+end
+
+
+"""
+    enable_disk_caching!(mem_limits::Dict{Int,Int}, disk_limit_mb::Int=16*2^10, processes::Vector{Int}=procs())
+
+Sets up disk caching on participating processes.
+This is a low level method for applying the `mem_limits` directly onto the processes.
+This skips the process discovery stage and the limit calculation.
+"""
+function enable_disk_caching!(
+    mem_limits::Dict{Int,Int},
+    disk_limit_mb::Int=16*2^10,
+    processes::Vector{Int}=procs(),
+)
+    results = [
+        remotecall(id) do
+            !isdefined(Main, :Dagger) && Main.eval(:(using Dagger))
+            Dagger.MemPool.setup_global_device!(
+                Dagger.MemPool.DiskCacheConfig(;
+                    toggle=true, membound=mem_limits[id], diskbound=disk_limit_mb * 2^20
+                ),
+            )
+            nothing
+        end for id in processes
+    ]
+    any_error = false
+    for (i, id) in enumerate(processes)
+        r = fetch(results[i])
+        any_error |= r !== nothing
+        if r !== nothing
+            @error("Error setting up disk caching on process id = $id", ex = r)
+        end
+    end
+
+    return if any_error
+        @error("Disk cache setup failed")
+        false
+    else
+        true
+    end
+end

test/cache_setup_test.jl

@@ -0,0 +1,9 @@
+using Distributed
+@everywhere using Pkg
+@everywhere Pkg.activate("..")
+using Dagger
+@assert Dagger.enable_disk_caching!()
+total_mem = @static VERSION >= v"1.8-" ? Sys.total_physical_memory : Sys.total_memory
+expected_mem_limit = Dagger.mem_limit(total_mem(), 30, length(procs()))
+applied_mem_limits = fetch.(remotecall.(() -> Dagger.MemPool.GLOBAL_DEVICE[].mem_limit, procs()))
+@assert all(applied_mem_limits .== expected_mem_limit)

test/diskcaching.jl

@@ -0,0 +1,19 @@
+@testset "Disk caching setup on multiple processes (single machine)" begin
+    runs = []
+    for p in 0:3
+        j = if p == 0
+            Cmd(`julia --startup-file=no`)
+        else
+            Cmd(`julia --startup-file=no -p $p`)
+        end
+        withenv("JULIA_MEMPOOL_EXPERIMENTAL_FANCY_ALLOCATOR"=>nothing,
+                "JULIA_MEMPOOL_EXPERIMENTAL_MEMORY_BOUND"=>nothing,
+                "JULIA_MEMPOOL_EXPERIMENTAL_DISK_CACHE"=>nothing,
+                "JULIA_MEMPOOL_EXPERIMENTAL_DISK_BOUND"=>nothing,
+                "JULIA_MEMPOOL_EXPERIMENTAL_ALLOCATOR_KIND"=>nothing) do
+            push!(runs, run(`$j cache_setup_test.jl`; wait=true))
+        end
+    end
+    wait.(runs)
+    @test all(getproperty.(runs, :exitcode) .== 0)
+end

test/runtests.jl

@@ -37,6 +37,8 @@ include("task-queues.jl")
 include("domain.jl")
 include("array.jl")
 include("cache.jl")
+include("diskcaching.jl")
+
 try # TODO: Fault tolerance is sometimes unreliable
 #include("fault-tolerance.jl")
 catch