Add memory_limit functions

Author: oschulz

docs/src/index.md

@@ -14,7 +14,9 @@ The elastic cluster manager automatically adds new workers to an automatically c
 Since workers can appear and disappear dynamically, initializing them (loading packages, etc.) via the standard `Distributed.@everywhere` macro is problematic, as workers added afterwards won't be initialized. Parallel processing tools provides the macro [`@always_everywhere`](@ref) to run code globally on all current processes, but also store the code so it can be run again on future new worker processes. Workers that are part of a [`FlexWorkerPool`](@ref) will be updated automatically on `take!` and `onworker`. You can also use [`ensure_procinit`](@ref) to manually update all workers
 to all `@always_everywhere` used so far.
 
-The function [`pinthreads_auto`](@ref)  (used inside of `@always_everywhere`) provides a convenient way to perform some automatic thread pinning on all processes. Note that it needs to follow an [`import ThreadPinning`](https://github.com/carstenbauer/ThreadPinning.jl/), and that more complex use cased may require customized thread pinning for best performance.
+The function [`pinthreads_auto`](@ref)  (used inside of `@always_everywhere`) provides a convenient way to perform some automatic thread pinning on all processes. Note that it needs to follow an [`import ThreadPinning`](https://github.com/carstenbauer/ThreadPinning.jl/), and that more complex use cases may require customized thread pinning for best performance.
+
+Some batch system configurations can result in whole Julia processes, or even a whole batch job, being terminated if a process exceeds its memory limit. In such cases, you can try to gain a softer failure mode by setting a custom (slightly smaller) memory limit using [`memory_limit!`](@ref).
 
 For example:
 
@@ -30,6 +32,9 @@ using ParallelProcessingTools, Distributed
 
     import ThreadPinning
     pinthreads_auto()
+
+    # Optional: Set a custom memory limit for worker processes:
+    # myid() != 1 && memory_limit!(8 * 1000^3) # 8 GB
 end
 
 runmode = OnLocalhost(n = 4)

src/ParallelProcessingTools.jl

@@ -25,6 +25,7 @@ using Parameters: @with_kw
 include("custom_cluster_managers.jl")
 using .CustomClusterManagers: ElasticManager
 
+include("memory.jl")
 include("display.jl")
 include("waiting.jl")
 include("exceptions.jl")

src/memory.jl

@@ -0,0 +1,79 @@
+# This file is a part of ParallelProcessingTools.jl, licensed under the MIT License (MIT).
+
+const _RLIMIT_AS = 9 # maximum size of the process's virtual memory
+
+const _EINVAL = 22 # libc errno for "Invalid argument"
+
+# Struct for libc getrlimit and setrlimit:
+mutable struct _RLimit
+    cur::Clong
+    max::Clong
+end
+
+
+"""
+    memory_limit()
+
+Gets the virtual memory limit for the current Julia process.
+
+Returns a tuple `(soft_limit::Int64, hard_limit::Int64)` (in units of bytes).
+Values of `-1` mean unlimited. 
+
+!!! note
+    Currently only works on Linux, simply returns `(Int64(-1), Int64(-1))` on
+    other operationg systems.
+"""
+function memory_limit end
+export memory_limit
+
+@static if Sys.islinux()
+    function memory_limit()
+        rlim = Ref(_RLimit(0, 0))
+        rc = ccall(:getrlimit, Cint, (Cint, Ref{_RLimit}), _RLIMIT_AS, rlim)
+        if rc != 0
+            error("Failed to get memory limits: ", Base.Libc.strerror(Base.Libc.errno()))
+        end
+        return rlim[].cur, rlim[].max
+    end
+else
+    memory_limit() = Int64(-1), Int64(-1)
+end
+
+
+"""
+    memory_limit!(soft_limit::Integer, hard_limit::Integer = -1)
+
+Sets the virtual memory limit for the current Julia process.
+
+`soft_limit` and `hard_limit` are in units of bytes. Values of `-1` mean
+unlimited. `hard_limit` must not be stricter than `soft_limit`, and should
+typically be set to `-1`.
+
+Returns `(soft_limit::Int64, hard_limit::Int64)`.
+
+!!! note
+    Currently only has an effect on Linux, does nothing and simply returns
+    `(Int64(-1), Int64(-1))` on other operating systems.
+"""
+function memory_limit! end
+export memory_limit!
+
+@static if Sys.islinux()
+    function memory_limit!(soft_limit::Integer, hard_limit::Integer = Int64(-1))
+        GC.gc()
+
+        rlim = Ref(_RLimit(soft_limit, hard_limit))
+        rc = ccall(:setrlimit, Cint, (Cint, Ref{_RLimit}), _RLIMIT_AS, rlim)
+        if rc != 0
+            errno = Base.Libc.errno()
+            if errno == _EINVAL
+                throw(ArgumentError("Invalid soft/hard memory limit $soft_limit/$hard_limit."))
+            else
+                error("Failed to set memory limit: ", errno, Base.Libc.strerror(errno))
+            end
+        end
+        return rlim[].cur, rlim[].max
+    end
+else
+    memory_limit!(soft_limit::Integer, hard_limit::Integer) = Int64(-1), Int64(-1)
+end

test/runtests.jl

@@ -9,6 +9,7 @@ Test.@testset "Package ParallelProcessingTools" begin
     @info "Testing with $(Base.Threads.nthreads()) Julia threads."
 
     include("test_aqua.jl")
+    include("test_memory.jl")
     include("test_waiting.jl")
     include("test_states.jl")
     include("test_fileio.jl")

test/test_memory.jl

@@ -0,0 +1,21 @@
+# This file is a part of ParallelProcessingTools.jl, licensed under the MIT License (MIT).
+
+using Test
+using ParallelProcessingTools
+
+@testset "memory" begin
+    @testset "memory_limit" begin
+        @test @inferred(memory_limit()) isa Tuple{<:Integer,<:Integer}
+        limit = Int(min(typemax(Int), 8 * Int64(1024)^3))
+        new_limits = (limit, -1)
+        if Sys.islinux()
+            @test @inferred(memory_limit!(new_limits...)) == new_limits
+            @test @inferred(memory_limit()) == new_limits
+            stricter_limit = round(Int, 0.9*limit)
+            @test_throws ArgumentError @inferred(memory_limit!(limit, stricter_limit))
+        else
+            @test @inferred(memory_limit!(new_limits...)) == (-1, -1)
+            @test @inferred(memory_limit()) == (-1, -1)
+        end
+    end
+end