/
githubmirror
/
julia
Обзор
Документация
Войти
/
githubmirror
/
julia
Код
Запросы
0
Пакеты
0
Релизы
0
Аналитика
Безопасность
master
stdlib/SharedArrays/src/SharedArrays.jl
709 строк
23 KB
Nicholas Bauer
Refactor Mmap shared memory API and interaction with SharedArrays (#60955)
21 июл 2026, 22:34
Не верифицирован
21 июл 2026, 22:34
01827ea
Код
Авторство
О чём код?
# This file is a part of Julia. License is MIT: https://julialang.org/license """ Provide the [`SharedArray`](@ref) type. It represents an array, which is shared across multiple processes, on a single machine. """ module SharedArrays using Mmap, Distributed, Random import Base: length, size, elsize, ndims, IndexStyle, reshape, convert, deepcopy_internal, show, getindex, setindex!, fill!, similar, reduce, map!, copyto!, cconvert import Base: Array import Random using Serialization using Serialization: serialize_cycle_header, serialize_type, writetag, UNDEFREF_TAG, serialize, deserialize import Serialization: serialize, deserialize import Distributed: RRID, procs, remotecall_fetch import Base.Filesystem: JL_O_CREAT, JL_O_RDWR, S_IRUSR, S_IWUSR export SharedArray, SharedVector, SharedMatrix, sdata, indexpids, localindices, unshare! mutable struct SharedArray{T,N} <: DenseArray{T,N} id::RRID dims::NTuple{N,Int} pids::Vector{Int} refs::Vector # The segname is currently used only in the test scripts to ensure that # the shmem segment has been unlinked. segname::String # Fields below are not to be serialized # Local shmem map. s::Array{T,N} # idx of current worker's pid in the pids vector, 0 if this shared array is not mapped locally. pidx::Int # the local partition into the array when viewed as a single dimensional array. # this can be removed when @distributed or its equivalent supports looping on # a subset of workers. loc_subarr_1d::SubArray{T,1,Array{T,1},Tuple{UnitRange{Int}},true} function SharedArray{T,N}(d,p,r,sn,s) where {T,N} S = new(RRID(),d,p,r,sn,s,0,view(Array{T}(undef, ntuple(d->0,N)), 1:0)) sa_refs[S.id] = WeakRef(S) S end end const sa_refs = Dict{RRID, WeakRef}() """ SharedArray{T}(dims::NTuple; init=false, pids=Int[]) SharedArray{T,N}(...) Construct a `SharedArray` of a bits type `T` and size `dims` across the processes specified by `pids` - all of which have to be on the same host. If `N` is specified by calling `SharedArray{T,N}(dims)`, then `N` must match the length of `dims`. If `pids` is left unspecified, the shared array will be mapped across all processes on the current host, including the master. But, `localindices` and `indexpids` will only refer to worker processes. This facilitates work distribution code to use workers for actual computation with the master process acting as a driver. If an `init` function of the type `initfn(S::SharedArray)` is specified, it is called on all the participating workers. The shared array is valid as long as a reference to the `SharedArray` object exists on the node which created the mapping. SharedArray{T}(filename::AbstractString, dims::NTuple, [offset=0]; mode=nothing, init=false, pids=Int[]) SharedArray{T,N}(...) Construct a `SharedArray` backed by the file `filename`, with element type `T` (must be a bits type) and size `dims`, across the processes specified by `pids` - all of which have to be on the same host. This file is mmapped into the host memory, with the following consequences: - The array data must be represented in binary format (e.g., an ASCII format like CSV cannot be supported) - Any changes you make to the array values (e.g., `A[3] = 0`) will also change the values on disk If `pids` is left unspecified, the shared array will be mapped across all processes on the current host, including the master. But, `localindices` and `indexpids` will only refer to worker processes. This facilitates work distribution code to use workers for actual computation with the master process acting as a driver. `mode` must be one of `"r"`, `"r+"`, `"w+"`, or `"a+"`, and defaults to `"r+"` if the file specified by `filename` already exists, or `"w+"` if not. If an `init` function of the type `initfn(S::SharedArray)` is specified, it is called on all the participating workers. You cannot specify an `init` function if the file is not writable. `offset` allows you to skip the specified number of bytes at the beginning of the file. """ SharedArray function SharedArray{T,N}(dims::Dims{N}; init=false, pids=Int[]) where {T,N} isbitstype(T) || throw(ArgumentError("type of SharedArray elements must be bits types, got $(T)")) pids, onlocalhost = shared_pids(pids) local s = Array{T}(undef, ntuple(d->0,N)) local io = nothing local S local shmmem_create_pid try # On OSX, the shm_seg_name length must be <= 31 characters (including the terminating NULL character) # Name only needs to be unique for the time it takes to instantiate the arrays; most likely case for # collision is during unit testing when RNG state is re-used. seg_name = "/jl" * "$(lpad(string(getpid() % 10^6), 6, "0"))" * # Ensure uniqueness from other processes "$(randstring(12))" * # Ensure uniqueness within process "$(bytes2hex(reverse(reinterpret(NTuple{4, UInt8}, time_ns() % UInt32))))" # Guard against RNG reset over ~4 s if onlocalhost shmmem_create_pid = myid() s, io = shm_mmap_array(T, dims, seg_name, JL_O_CREAT | JL_O_RDWR, false) else # The shared array is created on a remote machine shmmem_create_pid = pids[1] io = remotecall(pids[1]) do last(shm_mmap_array(T, dims, seg_name, JL_O_CREAT | JL_O_RDWR, false)) end wait(io) end func_mapshmem = () -> first(shm_mmap_array(T, dims, seg_name, JL_O_RDWR, true)) refs = Vector{Future}(undef, length(pids)) for (i, p) in enumerate(pids) refs[i] = remotecall(func_mapshmem, p) end # Wait till all the workers have mapped the segment for ref in refs wait(ref) end # All good, immediately unlink the segment. if (prod(dims) > 0) && (sizeof(T) > 0) if onlocalhost close(io) else remotecall_fetch(shmmem_create_pid, io) do fut close(fetch(fut)) end end end S = SharedArray{T,N}(dims, pids, refs, seg_name, s) initialize_shared_array(S, onlocalhost, init, pids) io = nothing finally if io !== nothing && @isdefined shmmem_create_pid if onlocalhost close(io) else remotecall_fetch(shmmem_create_pid, io) do fut close(fetch(fut)) end end end end S end SharedArray{T,N}(I::Integer...; kwargs...) where {T,N} = SharedArray{T,N}(I; kwargs...) SharedArray{T}(d::NTuple; kwargs...) where {T} = SharedArray{T,length(d)}(d; kwargs...) SharedArray{T}(I::Integer...; kwargs...) where {T} = SharedArray{T,length(I)}(I; kwargs...) SharedArray{T}(m::Integer; kwargs...) where {T} = SharedArray{T,1}(m; kwargs...) SharedArray{T}(m::Integer, n::Integer; kwargs...) where {T} = SharedArray{T,2}(m, n; kwargs...) SharedArray{T}(m::Integer, n::Integer, o::Integer; kwargs...) where {T} = SharedArray{T,3}(m, n, o; kwargs...) function SharedArray{T,N}(filename::AbstractString, dims::NTuple{N,Int}, offset::Integer=0; mode=nothing, init=false, pids::Vector{Int}=Int[]) where {T,N} if !isabspath(filename) throw(ArgumentError("$filename is not an absolute path; try abspath(filename)?")) end if !isbitstype(T) throw(ArgumentError("type of SharedArray elements must be bits types, got $(T)")) end pids, onlocalhost = shared_pids(pids) # If not supplied, determine the appropriate mode have_file = onlocalhost ? isfile(filename) : remotecall_fetch(isfile, pids[1], filename) mode_val = mode === nothing ? (have_file ? "r+" : "w+") : mode workermode = mode_val == "w+" ? "r+" : mode_val # workers don't truncate! # Ensure the file will be readable if !(mode_val in ("r", "r+", "w+", "a+")) throw(ArgumentError("mode must be readable, but $mode_val is not")) end if init !== false typeassert(init, Function) if !(mode_val in ("r+", "w+", "a+")) throw(ArgumentError("cannot initialize unwritable array (mode = $mode_val)")) end end if mode_val == "r" && !isfile(filename) throw(ArgumentError("file $filename does not exist, but mode $mode_val cannot create it")) end # Create the file if it doesn't exist, map it if it does refs = Vector{Future}(undef, length(pids)) func_mmap = mode -> open(filename, mode) do io mmap(io, Array{T,N}, dims, offset; shared=true) end s = Array{T}(undef, ntuple(d->0,N)) if onlocalhost s = func_mmap(mode_val) refs[1] = remotecall(pids[1]) do func_mmap(workermode) end else refs[1] = remotecall_wait(pids[1]) do func_mmap(mode_val) end end # Populate the rest of the workers for i = 2:length(pids) refs[i] = remotecall(pids[i]) do func_mmap(workermode) end end # Wait till all the workers have mapped the segment for ref in refs wait(ref) end S = SharedArray{T,N}(dims, pids, refs, filename, s) initialize_shared_array(S, onlocalhost, init, pids) S end SharedArray{T}(filename::AbstractString, dims::NTuple{N,Int}, offset::Integer=0; mode=nothing, init=false, pids::Vector{Int}=Int[]) where {T,N} = SharedArray{T,N}(filename, dims, offset; mode=mode, init=init, pids=pids) function initialize_shared_array(S, onlocalhost, init, pids) if onlocalhost init_loc_flds(S) else S.pidx = 0 end # if present, init function is called on each of the parts if isa(init, Function) @sync begin for p in pids @async remotecall_wait(init, p, S) end end end finalizer(finalize_refs, S) S end function finalize_refs(S::SharedArray{T,N}) where T where N if length(S.pids) > 0 for r in S.refs finalize(r) end empty!(S.pids) empty!(S.refs) init_loc_flds(S) S.s = Array{T}(undef, ntuple(d->0,N)) delete!(sa_refs, S.id) end S end """ SharedVector A one-dimensional [`SharedArray`](@ref). """ const SharedVector{T} = SharedArray{T,1} """ SharedMatrix A two-dimensional [`SharedArray`](@ref). """ const SharedMatrix{T} = SharedArray{T,2} SharedVector(A::Vector) = SharedArray(A) SharedMatrix(A::Matrix) = SharedArray(A) size(S::SharedArray) = S.dims elsize(::Type{SharedArray{T,N}}) where {T,N} = elsize(Array{T,N}) # aka fieldtype(T, :s) IndexStyle(::Type{<:SharedArray}) = IndexLinear() function local_array_by_id(refid) if isa(refid, Future) refid = remoteref_id(refid) end fetch(channel_from_id(refid)) end function reshape(a::SharedArray{T}, dims::NTuple{N,Int}) where {T,N} if length(a) != prod(dims) throw(DimensionMismatch("dimensions must be consistent with array size")) end refs = Vector{Future}(undef, length(a.pids)) for (i, p) in enumerate(a.pids) refs[i] = remotecall(p, a.refs[i], dims) do r, d reshape(local_array_by_id(r), d) end end A = SharedArray{T,N}(dims, a.pids, refs, a.segname, reshape(a.s, dims)) init_loc_flds(A) A end """ procs(S::SharedArray) Get the vector of processes mapping the shared array. """ procs(S::SharedArray) = S.pids """ indexpids(S::SharedArray) Return the current worker's index in the list of workers mapping the `SharedArray` (i.e. in the same list returned by `procs(S)`), or 0 if the `SharedArray` is not mapped locally. """ indexpids(S::SharedArray) = S.pidx """ sdata(S::SharedArray) Return the actual `Array` object backing `S`. """ sdata(S::SharedArray) = S.s sdata(A::AbstractArray) = A """ localindices(S::SharedArray) Return a range describing the "default" indices to be handled by the current process. This range should be interpreted in the sense of linear indexing, i.e., as a sub-range of `1:length(S)`. In multi-process contexts, returns an empty range in the parent process (or any process for which [`indexpids`](@ref) returns 0). It's worth emphasizing that `localindices` exists purely as a convenience, and you can partition work on the array among workers any way you wish. For a `SharedArray`, all indices should be equally fast for each worker process. """ localindices(S::SharedArray) = S.pidx > 0 ? range_1dim(S, S.pidx) : 1:0 cconvert(::Type{Ptr{T}}, S::SharedArray{T}) where {T} = cconvert(Ptr{T}, sdata(S)) cconvert(::Type{Ptr{T}}, S::SharedArray ) where {T} = cconvert(Ptr{T}, sdata(S)) function SharedArray(A::Array) S = SharedArray{eltype(A),ndims(A)}(size(A)) copyto!(S, A) end function SharedArray{T}(A::Array) where T S = SharedArray{T,ndims(A)}(size(A)) copyto!(S, A) end function SharedArray{TS,N}(A::Array{TA,N}) where {TS,TA,N} S = SharedArray{TS,ndims(A)}(size(A)) copyto!(S, A) end convert(T::Type{<:SharedArray}, a::Array) = T(a)::T function deepcopy_internal(S::SharedArray, stackdict::IdDict) haskey(stackdict, S) && return stackdict[S] R = SharedArray{eltype(S),ndims(S)}(size(S); pids = S.pids) copyto!(sdata(R), sdata(S)) stackdict[S] = R return R end function shared_pids(pids) if isempty(pids) # only use workers on the current host pids = procs(myid()) if length(pids) > 1 pids = filter(!=(1), pids) end onlocalhost = true else if !check_same_host(pids) throw(ArgumentError("SharedArray requires all requested processes to be on the same machine.")) end onlocalhost = myid() in procs(pids[1]) end pids, onlocalhost end function range_1dim(S::SharedArray, pidx) l = length(S) nw = length(S.pids) partlen = div(l, nw) if l < nw if pidx <= l return pidx:pidx else return 1:0 end elseif pidx == nw return (((pidx-1) * partlen) + 1):l else return (((pidx-1) * partlen) + 1):(pidx*partlen) end end sub_1dim(S::SharedArray, pidx) = view(S.s, range_1dim(S, pidx)) function init_loc_flds(S::SharedArray{T,N}, empty_local=false) where T where N if myid() in S.pids S.pidx = findfirst(isequal(myid()), S.pids) S.s = local_array_by_id(S.refs[S.pidx]) S.loc_subarr_1d = sub_1dim(S, S.pidx) else S.pidx = 0 if empty_local S.s = Array{T}(undef, ntuple(d->0,N)) end S.loc_subarr_1d = view(Array{T}(undef, ntuple(d->0,N)), 1:0) end end # Don't serialize s (it is the complete array) and # pidx, which is relevant to the current process only function serialize(s::AbstractSerializer, S::SharedArray) serialize_cycle_header(s, S) && return destpid = worker_id_from_socket(s.io) if S.id.whence == destpid # The shared array was created from destpid, hence a reference to it # must be available at destpid. serialize(s, true) serialize(s, S.id.whence) serialize(s, S.id.id) return end serialize(s, false) for n in fieldnames(SharedArray) if n in [:s, :pidx, :loc_subarr_1d] writetag(s.io, UNDEFREF_TAG) elseif n === :refs v = getfield(S, n) if isa(v[1], Future) # convert to ids to avoid distributed GC overhead ids = [remoteref_id(x) for x in v] serialize(s, ids) else serialize(s, v) end else serialize(s, getfield(S, n)) end end end function deserialize(s::AbstractSerializer, t::Type{<:SharedArray}) ref_exists = deserialize(s) if ref_exists sref = sa_refs[RRID(deserialize(s), deserialize(s))] if sref.value !== nothing return sref.value end error("Expected reference to shared array instance not found") end S = invoke(deserialize, Tuple{AbstractSerializer,DataType}, s, t) init_loc_flds(S, true) return S end function show(io::IO, S::SharedArray) if length(S.s) > 0 invoke(show, Tuple{IO,DenseArray}, io, S) else show(io, remotecall_fetch(sharr->sharr.s, S.pids[1], S)) end end function show(io::IO, mime::MIME"text/plain", S::SharedArray) if length(S.s) > 0 invoke(show, Tuple{IO,MIME"text/plain",DenseArray}, io, MIME"text/plain"(), S) else # retrieve from the first worker mapping the array. summary(io, S); println(io, ":") Base.print_array(io, remotecall_fetch(sharr->sharr.s, S.pids[1], S)) end end Array(S::SharedArray) = S.s # pass through getindex and setindex! - unlike DArrays, these always work on the complete array Base.@propagate_inbounds getindex(S::SharedArray, i::Real) = getindex(S.s, i) Base.@propagate_inbounds setindex!(S::SharedArray, x, i::Real) = setindex!(S.s, x, i) function fill!(S::SharedArray, v) vT = convert(eltype(S), v) f = S->fill!(S.loc_subarr_1d, vT) @sync for p in procs(S) @async remotecall_wait(f, p, S) end return S end function Random.rand!(S::SharedArray{T}) where T f = S->map!(x -> rand(T), S.loc_subarr_1d, S.loc_subarr_1d) @sync for p in procs(S) @async remotecall_wait(f, p, S) end return S end function Random.randn!(S::SharedArray) f = S->map!(x -> randn(), S.loc_subarr_1d, S.loc_subarr_1d) @sync for p in procs(S) @async remotecall_wait(f, p, S) end return S end # convenience constructors function shmem_fill(v, dims; kwargs...) SharedArray{typeof(v),length(dims)}(dims; init = S->fill!(S.loc_subarr_1d, v), kwargs...) end shmem_fill(v, I::Int...; kwargs...) = shmem_fill(v, I; kwargs...) # rand variant with range function shmem_rand(TR::Union{DataType, UnitRange}, dims; kwargs...) if isa(TR, UnitRange) SharedArray{Int,length(dims)}(dims; init = S -> map!(x -> rand(TR), S.loc_subarr_1d, S.loc_subarr_1d), kwargs...) else SharedArray{TR,length(dims)}(dims; init = S -> map!(x -> rand(TR), S.loc_subarr_1d, S.loc_subarr_1d), kwargs...) end end shmem_rand(TR::Union{DataType, UnitRange}, i::Int; kwargs...) = shmem_rand(TR, (i,); kwargs...) shmem_rand(TR::Union{DataType, UnitRange}, I::Int...; kwargs...) = shmem_rand(TR, I; kwargs...) shmem_rand(dims; kwargs...) = shmem_rand(Float64, dims; kwargs...) shmem_rand(I::Int...; kwargs...) = shmem_rand(I; kwargs...) function shmem_randn(dims; kwargs...) SharedArray{Float64,length(dims)}(dims; init = S-> map!(x -> randn(), S.loc_subarr_1d, S.loc_subarr_1d), kwargs...) end shmem_randn(I::Int...; kwargs...) = shmem_randn(I; kwargs...) similar(S::SharedArray, T::Type, dims::Dims) = similar(S.s, T, dims) similar(S::SharedArray, T::Type) = similar(S.s, T, size(S)) similar(S::SharedArray, dims::Dims) = similar(S.s, eltype(S), dims) similar(S::SharedArray) = similar(S.s, eltype(S), size(S)) reduce(f, S::SharedArray) = mapreduce(fetch, f, Any[ @spawnat p reduce(f, S.loc_subarr_1d) for p in procs(S) ]) reduce(::typeof(vcat), S::SharedVector) = invoke(reduce, Tuple{Any,SharedArray}, vcat, S) reduce(::typeof(hcat), S::SharedVector) = invoke(reduce, Tuple{Any,SharedArray}, hcat, S) function map!(f, S::SharedArray, Q::SharedArray) if (S !== Q) && (procs(S) != procs(Q) || localindices(S) != localindices(Q)) throw(ArgumentError("incompatible source and destination arguments")) end @sync for p in procs(S) @spawnat p begin for idx in localindices(S) S.s[idx] = f(Q.s[idx]) end end end return S end copyto!(S::SharedArray, A::Array) = (copyto!(S.s, A); S) function copyto!(S::SharedArray, R::SharedArray) length(S) == length(R) || throw(BoundsError()) ps = intersect(procs(S), procs(R)) isempty(ps) && throw(ArgumentError("source and destination arrays don't share any process")) l = length(S) length(ps) > l && (ps = ps[1:l]) nw = length(ps) partlen = div(l, nw) @sync for i = 1:nw p = ps[i] idx = i < nw ? ((i-1)*partlen+1:i*partlen) : ((i-1)*partlen+1:l) @spawnat p begin S.s[idx] = R.s[idx] end end return S end """ unshare!(S::SharedArray) Release resources from workers and make the memory no longer available to them. The array is still usable on the host process. Must be called from the process that created `S`; calling it from any other process throws an `ArgumentError`. !!! note Relying on the finalizers to perform cleanup requires multiple GC rounds to release the underlying mmap. Call this function proactively to ensure a single GC round is sufficient. """ function unshare!(S::SharedArray) if !isempty(S.pids) myid() == S.id.whence || throw(ArgumentError("unshare! must be called from the process that created the SharedArray")) @sync begin for i in eachindex(S.pids) S.pidx == i && continue @async remotecall_wait(S.pids[i], S.refs[i]) do r Mmap.munmap!(local_array_by_id(r)) end end end empty!(S.pids) empty!(S.refs) init_loc_flds(S) return nothing end end function print_shmem_limits(slen) try if Sys.islinux() return # Not relevant to Linux, which uses a tmpfs-backed system elseif Sys.isapple() pfx = "kern.sysv" elseif Sys.KERNEL === :FreeBSD || Sys.KERNEL === :DragonFly pfx = "kern.ipc" elseif Sys.KERNEL === :OpenBSD pfx = "kern.shminfo" else # seems NetBSD does not have *.shmall return end shmmax_MB = div(parse(Int, split(read(`sysctl $(pfx).shmmax`, String))[end]), 1024*1024) page_size = parse(Int, split(read(`getconf PAGE_SIZE`, String))[end]) shmall_MB = div(parse(Int, split(read(`sysctl $(pfx).shmall`, String))[end]) * page_size, 1024*1024) println("System max size of single shmem segment(MB) : ", shmmax_MB, "\nSystem max size of all shmem segments(MB) : ", shmall_MB, "\nRequested size(MB) : ", div(slen, 1024*1024), "\nPlease ensure requested size is within system limits.", "\nIf not, increase system limits and try again." ) catch e @warn "Unable to print shared memory limits." end end # utilities function shm_mmap_array(T, dims, shm_seg_name, mode, closeio) if (prod(dims) == 0) || (sizeof(T) == 0) return Array{T}(undef, dims), nothing end try return _shm_mmap_array(T, dims, shm_seg_name, mode, closeio) catch print_shmem_limits(prod(dims)*sizeof(T)) rethrow() end end function _shm_mmap_array(T, dims, shm_seg_name, mode, closeio) readonly = !((mode & JL_O_RDWR) == JL_O_RDWR) create = (mode & JL_O_CREAT) == JL_O_CREAT io = open(Mmap.SharedMemory, shm_seg_name, prod(dims) * sizeof(T); readonly, create) A = mmap(io, Array{T,length(dims)}, dims, zero(Int64)) # Workers can immediately close the virtual file after mapping if closeio close(io) io = nothing end return A, io end end # module