Skip to content

Commit ff29a19

Browse files
marius311marius311
committed
tweak parallel utils
1 parent 33f4d60 commit ff29a19

File tree

3 files changed

+106
-105
lines changed

3 files changed

+106
-105
lines changed

‎src/CMBLensing.jl‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -99,6 +99,7 @@ export gibbs_initialize_f!, gibbs_initialize_ϕ!, gibbs_initialize_θ!,
9999
# generic stuff
100100
include("util.jl")
101101
include("util_fft.jl")
102+
include("util_parallel.jl")
102103
include("numerical_algorithms.jl")
103104
include("generic.jl")
104105
include("cls.jl")

‎src/util.jl‎

Lines changed: 0 additions & 105 deletions
Original file line numberDiff line numberDiff line change
@@ -276,58 +276,6 @@ macro ismain()
276276
end
277277

278278

279-
@init @require MPIClusterManagers="e7922434-ae4b-11e9-05c5-9780451d2c66" begin
280-
281-
using .MPIClusterManagers: MPI, start_main_loop, TCP_TRANSPORT_ALL, MPI_TRANSPORT_ALL
282-
283-
"""
284-
init_MPI_workers()
285-
286-
Initialize MPI processes as Julia workers. Should be called from all MPI
287-
processes, and will only return on the master process.
288-
289-
`transport` should be `"MPI"` or `"TCP"`, which is by default read from the
290-
environment variable `JULIA_MPI_TRANSPORT`, and otherwise defaults to `"TCP"`.
291-
292-
If CUDA is loaded and functional in the Main module, additionally calls
293-
[`assign_GPU_workers()`](@ref)
294-
"""
295-
function init_MPI_workers(;
296-
stdout_to_master = false,
297-
stderr_to_master = false,
298-
transport = get(ENV,"JULIA_MPI_TRANSPORT","TCP")
299-
)
300-
301-
if !MPI.Initialized()
302-
MPI.Init()
303-
end
304-
size = MPI.Comm_size(MPI.COMM_WORLD)
305-
rank = MPI.Comm_rank(MPI.COMM_WORLD)
306-
307-
if size>1
308-
# workers don't return from this call:
309-
start_main_loop(
310-
Dict("TCP"=>TCP_TRANSPORT_ALL,"MPI"=>MPI_TRANSPORT_ALL)[transport],
311-
stdout_to_master=stdout_to_master,
312-
stderr_to_master=stderr_to_master
313-
)
314-
315-
if @isdefined(CUDA) && CUDA.functional()
316-
assign_GPU_workers()
317-
end
318-
@everywhere begin
319-
typ = (myid()==1) ? "(master)" : "(worker)"
320-
dev = (@isdefined(CUDA) && CUDA.functional()) ? device() : "CPU"
321-
@info "MPI process $(myid()) $typ is running on $(gethostname())::$dev"
322-
end
323-
end
324-
325-
end
326-
327-
328-
end
329-
330-
331279
firsthalf(x) = x[1:end÷2]
332280
lasthalf(x) = x[end÷2:end]
333281

@@ -463,59 +411,6 @@ macro auto_adjoint(funcdef)
463411
esc(Expr(:block, defs...))
464412
end
465413

466-
467-
468-
"""
469-
assign_GPU_workers()
470-
471-
Assign each Julia worker process a unique GPU using `CUDA.device!`.
472-
Workers may be distributed across different hosts, and each host can have
473-
multiple GPUs.
474-
"""
475-
function assign_GPU_workers()
476-
@everywhere @eval Main using Distributed, CMBLensing
477-
master_uuid = @isdefined(CUDA) ? CUDA.uuid(device()) : nothing
478-
accessible_gpus = Dict(map(workers()) do id
479-
@eval Main @fetchfrom $id begin
480-
ds = CUDA.devices()
481-
# put master's GPU last so we don't double up on it unless we need to
482-
$id => sort((CUDA.deviceid.(ds) .=> CUDA.uuid.(ds)), by=(((k,v),)->v==$master_uuid ? Inf : k))
483-
end
484-
end)
485-
claimed = Set()
486-
assignments = Dict(map(workers()) do myid
487-
for (gpu_id, gpu_uuid) in accessible_gpus[myid]
488-
if !(gpu_uuid in claimed)
489-
push!(claimed, gpu_uuid)
490-
return myid => gpu_id
491-
end
492-
end
493-
error("Can't assign a unique GPU to every worker, process $myid has no free GPUs left.")
494-
end)
495-
@everywhere workers() device!($assignments[myid()])
496-
println(GPU_worker_info())
497-
end
498-
499-
"""
500-
GPU_worker_info()
501-
502-
Returns string showing info about assigned GPU workers.
503-
"""
504-
function GPU_worker_info()
505-
lines = @eval Main map(procs()) do id
506-
@fetchfrom id begin
507-
if @isdefined(CUDA)
508-
device = "device = $(sprint(io->show(io, MIME("text/plain"), CUDA.device()))) $(split(string(CUDA.uuid(CUDA.device())),'-')[1]))"
509-
else
510-
device = ""
511-
end
512-
join(["($(id==1 ? "master" : "worker") = $id", "host = $(gethostname())", device], ", ")
513-
end
514-
end
515-
join(["GPU_worker_info:"; lines], "\n")
516-
end
517-
518-
519414
string_trunc(x) = Base._truncate_at_width_or_chars(string(x), displaysize(stdout)[2]-14)
520415

521416
import NamedTupleTools

‎src/util_parallel.jl‎

Lines changed: 105 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,105 @@
1+
2+
_mpi_rank() = nothing
3+
4+
@init @require MPIClusterManagers="e7922434-ae4b-11e9-05c5-9780451d2c66" begin
5+
6+
using .MPIClusterManagers: MPI, start_main_loop, TCP_TRANSPORT_ALL, MPI_TRANSPORT_ALL
7+
8+
"""
9+
init_MPI_workers()
10+
11+
Initialize MPI processes as Julia workers. Should be called from all MPI
12+
processes, and will only return on the master process.
13+
14+
`transport` should be `"MPI"` or `"TCP"`, which is by default read from the
15+
environment variable `JULIA_MPI_TRANSPORT`, and otherwise defaults to `"TCP"`.
16+
17+
If CUDA is loaded and functional in the Main module, additionally calls
18+
[`assign_GPU_workers()`](@ref)
19+
"""
20+
function init_MPI_workers(;
21+
stdout_to_master = false,
22+
stderr_to_master = false,
23+
transport = get(ENV,"JULIA_MPI_TRANSPORT","TCP")
24+
)
25+
26+
if !MPI.Initialized()
27+
MPI.Init()
28+
end
29+
size = MPI.Comm_size(MPI.COMM_WORLD)
30+
rank = MPI.Comm_rank(MPI.COMM_WORLD)
31+
32+
if size>1
33+
# workers don't return from this call:
34+
start_main_loop(
35+
Dict("TCP"=>TCP_TRANSPORT_ALL,"MPI"=>MPI_TRANSPORT_ALL)[transport],
36+
stdout_to_master=stdout_to_master,
37+
stderr_to_master=stderr_to_master
38+
)
39+
40+
if @isdefined(CUDA) && CUDA.functional()
41+
assign_GPU_workers()
42+
end
43+
@everywhere begin
44+
typ = (myid()==1) ? "(master)" : "(worker)"
45+
dev = (@isdefined(CUDA) && CUDA.functional()) ? device() : "CPU"
46+
@info "MPI process $(myid()) $typ is running on $(gethostname())::$dev"
47+
end
48+
end
49+
50+
end
51+
52+
_mpi_rank() = MPI.Comm_rank(MPI.COMM_WORLD)
53+
54+
end
55+
56+
57+
"""
58+
assign_GPU_workers()
59+
60+
Assign each Julia worker process a unique GPU using `CUDA.device!`.
61+
Workers may be distributed across different hosts, and each host can have
62+
multiple GPUs.
63+
"""
64+
function assign_GPU_workers()
65+
@everywhere @eval Main using Distributed, CMBLensing
66+
master_uuid = @isdefined(CUDA) ? CUDA.uuid(device()) : nothing
67+
accessible_gpus = Dict(map(workers()) do id
68+
@eval Main @fetchfrom $id begin
69+
ds = CUDA.devices()
70+
# put master's GPU last so we don't double up on it unless we need to
71+
$id => sort((CUDA.deviceid.(ds) .=> CUDA.uuid.(ds)), by=(((k,v),)->v==$master_uuid ? Inf : k))
72+
end
73+
end)
74+
claimed = Set()
75+
assignments = Dict(map(workers()) do myid
76+
for (gpu_id, gpu_uuid) in accessible_gpus[myid]
77+
if !(gpu_uuid in claimed)
78+
push!(claimed, gpu_uuid)
79+
return myid => gpu_id
80+
end
81+
end
82+
error("Can't assign a unique GPU to every worker, process $myid has no free GPUs left.")
83+
end)
84+
@everywhere workers() device!($assignments[myid()])
85+
println(worker_info())
86+
end
87+
88+
89+
"""
90+
worker_info()
91+
92+
Returns string showing info about assigned workers.
93+
"""
94+
function worker_info()
95+
lines = @eval Main map(procs()) do id
96+
@fetchfrom id begin
97+
info = ["myid = $id"]
98+
!isnothing(CMBLensing._mpi_rank()) && push!(info, "mpi-rank = $(CMBLensing._mpi_rank())")
99+
push!(info, "host = $(gethostname())")
100+
@isdefined(CUDA) && push!(info, "device = $(sprint(io->show(io, MIME("text/plain"), CUDA.device()))) $(split(string(CUDA.uuid(CUDA.device())),'-')[1]))")
101+
" ("*join(info, ", ")*")"
102+
end
103+
end
104+
join(["Worker info:"; lines], "\n")
105+
end

0 commit comments

Comments
 (0)