Asynchronous NetCDF output.

Manifest.toml

@@ -146,9 +146,9 @@ version = "1.0.5"
 
 [[FillArrays]]
 deps = ["LinearAlgebra", "Random", "SparseArrays", "Test"]
-git-tree-sha1 = "471b7e33dc9c9c5b9170045dd57c8ba0927b2918"
+git-tree-sha1 = "2def0123a4f3572234405b0e3d80bfe5d3e1a2a4"
 uuid = "1a297f60-69ca-5386-bcde-b61e274b549b"
-version = "0.4.0"
+version = "0.5.0"
 
 [[Formatting]]
 deps = ["Compat"]

Project.toml

@@ -8,6 +8,7 @@ CUDAapi = "3895d2a7-ec45-59b8-82bb-cfc6a382f9b3"
 CUDAdrv = "c5f51814-7f29-56b8-a69c-e4d8f6be1fde"
 CUDAnative = "be33ccc6-a3ff-5ff2-a52e-74243cff1e17"
 CuArrays = "3a865a2d-5b23-5a0f-bc46-62713ec82fae"
+Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
 GPUifyLoops = "ba82f77b-6841-5d2e-bd9f-4daf811aec27"
 JLD = "4138dd39-2aa7-5051-a626-17a0bb65d9c8"

examples/free_convection.jl

@@ -1,42 +1,52 @@
-using Oceananigans
-using CuArrays
-
-# physical constants
-p0 = 1027
-cp = 4181.3
-
-# set simulation parameters
-Nx, Ny, Nz = 128, 128, 128
-Lx, Ly, Lz = 100, 100, 100
-Nt, Δt = 10000, 6
-ν, κ = 1e-4, 1e-4
-Tz = 0.01
-bottom_gradient = Tz
-top_flux = 75 / (p0 * cp) #for flux bc this should be positive for cooling
-
-# create the model
-model = Model(N=(Nx, Ny, Nz), L=(Lx, Ly, Lz), ν=ν, κ=κ, arch=:GPU, float_type=Float32)
-
-# set boundary conditions
-model.boundary_conditions.T.z.left = BoundaryCondition(Flux, top_flux)
-model.boundary_conditions.T.z.right = BoundaryCondition(Gradient, bottom_gradient)
-# the default boundary conditions for velocity is free-slip
-
-# set initial condition
-T_prof = 273.15 .+ 20 .+ Tz .* model.grid.zC
-# make initial condition into 3D array
-T_3d = repeat(reshape(T_prof, 1, 1, Nz), Nx, Ny, 1)
-@. T_3d[:, :, 1:round(Int, Nz/2)] += 0.01*randn() #add noise to the top half of the domain for convection
-
-model.tracers.T.data .= CuArray(T_3d)
-
-# Write temperature field to disk every 10 time steps.
-output_writer = NetCDFOutputWriter(dir=".", prefix="convection", frequency=2500)
-push!(model.output_writers, output_writer)
-
-# Time stepping
-Ni = 10
-for i = 1:Ni
-    time_step!(model, ceil(Nt/Ni), Δt)
-    println("Time step: $(model.clock.iteration)")
-end
+using Distributed
+addprocs(1)  # For asynchronous output writing.
+
+# Apparently I have to do this even when starting julia with "-p 2 --project"
+# otherwise we get "ERROR: LoadError: On worker 2: ArgumentError: Package Oceananigans
+# not found in current path:"
+@everywhere using Pkg
+@everywhere Pkg.activate(".");
+
+@everywhere using Oceananigans
+@everywhere using CuArrays
+
+# physical constants
+p0 = 1027
+cp = 4181.3
+
+# set simulation parameters
+Nx, Ny, Nz = 128, 128, 128
+Lx, Ly, Lz = 100, 100, 100
+Nt, Δt = 1000, 1
+ν, κ = 1e-4, 1e-4
+Tz = 0.01
+bottom_gradient = Tz
+top_flux = 75 / (p0 * cp) #for flux bc this should be positive for cooling
+
+# create the model
+model = Model(N=(Nx, Ny, Nz), L=(Lx, Ly, Lz), ν=ν, κ=κ, arch=:GPU, float_type=Float32)
+
+# set boundary conditions
+model.boundary_conditions.T.z.left = BoundaryCondition(Flux, top_flux)
+model.boundary_conditions.T.z.right = BoundaryCondition(Gradient, bottom_gradient)
+# the default boundary conditions for velocity is free-slip
+
+# set initial condition
+T_prof = 273.15 .+ 20 .+ Tz .* model.grid.zC
+# make initial condition into 3D array
+T_3d = repeat(reshape(T_prof, 1, 1, Nz), Nx, Ny, 1)
+@. T_3d[:, :, 1:round(Int, Nz/2)] += 0.01*randn() #add noise to the top half of the domain for convection
+
+model.tracers.T.data .= CuArray(T_3d)
+
+# Write temperature field to disk every 10 time steps.
+output_writer = NetCDFOutputWriter(dir=".", prefix="convection", frequency=200, async=true)
+push!(model.output_writers, output_writer)
+
+# Time stepping
+for i = 1:Nt
+    tic = time_ns()
+    print("Time: $(model.clock.time) ")
+    time_step!(model, 1, Δt)
+    println(prettytime(time_ns() - tic))
+end

src/fields.jl

@@ -3,6 +3,8 @@ import Base:
     getindex, lastindex, setindex!,
     iterate, similar, *, +, -
 
+@hascuda using CuArrays
+
 """
     CellField{T,G<:Grid{T}} <: Field
 

src/output_writers.jl

@@ -1,4 +1,7 @@
 import JLD
+using Distributed
+
+using NetCDF
 
 "A type for writing checkpoints."
 struct Checkpointer <: OutputWriter
@@ -15,6 +18,7 @@ mutable struct NetCDFOutputWriter <: OutputWriter
     output_frequency::Int
     padding::Int
     compression::Int
+    async::Bool
 end
 
 "A type for writing Binary output."
@@ -25,8 +29,8 @@ mutable struct BinaryOutputWriter <: OutputWriter
     padding::Int
 end
 
-function NetCDFOutputWriter(; dir=".", prefix="", frequency=1, padding=9, compression=5)
-    NetCDFOutputWriter(dir, prefix, frequency, padding, compression)
+function NetCDFOutputWriter(; dir=".", prefix="", frequency=1, padding=9, compression=5, async=false)
+    NetCDFOutputWriter(dir, prefix, frequency, padding, compression, async)
 end
 
 "Return the filename extension for the `OutputWriter` filetype."
@@ -122,15 +126,40 @@ end
 # etc; so this API needs to be designed. For now, we simply save u, v, w, and T.
 
 function write_output(model::Model, fw::NetCDFOutputWriter)
+    fields = Dict(
+        "xC" => collect(model.grid.xC),
+        "yC" => collect(model.grid.yC),
+        "zC" => collect(model.grid.zC),
+        "xF" => collect(model.grid.xF),
+        "yF" => collect(model.grid.yF),
+        "zF" => collect(model.grid.zF),
+        "u" => Array(model.velocities.u.data),
+        "v" => Array(model.velocities.v.data),
+        "w" => Array(model.velocities.w.data),
+        "T" => Array(model.tracers.T.data),
+        "S" => Array(model.tracers.S.data)
+    )
+
+    if fw.async
+        # Execute asynchronously on worker 2.
+        println("Using @async...")
+        println("nprocs()=$(nprocs())")
+        @async remotecall(write_output_netcdf, 2, fw, fields, model.clock.iteration)
+    else
+        println("Regular call...")
+        write_output_netcdf(fw, fields, model.clock.iteration)
+    end
 
-    xC = collect(model.grid.xC)
-    yC = collect(model.grid.yC)
-    zC = collect(model.grid.zC)
-
-    xF = collect(model.grid.xF)
-    yF = collect(model.grid.yF)
-    zF = collect(model.grid.zF)
+    return nothing
+end
 
+function write_output_netcdf(fw::NetCDFOutputWriter, fields, iteration)
+    xC, yC, zC = fields["xC"], fields["yC"], fields["zC"]
+    xF, yF, zF = fields["xF"], fields["yF"], fields["zF"]
+
+    u, v, w = fields["u"], fields["v"], fields["w"]
+    T, S    = fields["T"], fields["S"]
+
     xC_attr = Dict("longname" => "Locations of the cell centers in the x-direction.", "units" => "m")
     yC_attr = Dict("longname" => "Locations of the cell centers in the y-direction.", "units" => "m")
     zC_attr = Dict("longname" => "Locations of the cell centers in the z-direction.", "units" => "m")
@@ -145,9 +174,13 @@ function write_output(model::Model, fw::NetCDFOutputWriter)
     T_attr = Dict("longname" => "Temperature", "units" => "K")
     S_attr = Dict("longname" => "Salinity", "units" => "g/kg")
 
-    filepath = joinpath(fw.dir, filename(fw, "", model.clock.iteration))
+    filepath = joinpath(fw.dir, filename(fw, "", iteration))
 
-    println("[NetCDFOutputWriter] Writing fields to disk: $filepath")
+    if fw.async
+        println("[Worker $(Distributed.myid()): NetCDFOutputWriter] Writing fields to disk: $filepath")
+    else
+        println("[NetCDFOutputWriter] Writing fields to disk: $filepath")
+    end
 
     isfile(filepath) && rm(filepath)
 
@@ -176,11 +209,11 @@ function write_output(model::Model, fw::NetCDFOutputWriter)
                             "zC", zC, zC_attr,
                             atts=S_attr, compress=fw.compression)
 
-    ncwrite(Array(model.velocities.u.data), filepath, "u")
-    ncwrite(Array(model.velocities.v.data), filepath, "v")
-    ncwrite(Array(model.velocities.w.data), filepath, "w")
-    ncwrite(Array(model.tracers.T.data), filepath, "T")
-    ncwrite(Array(model.tracers.S.data), filepath, "S")
+    ncwrite(u, filepath, "u")
+    ncwrite(v, filepath, "v")
+    ncwrite(w, filepath, "w")
+    ncwrite(T, filepath, "T")
+    ncwrite(S, filepath, "S")
 
     ncclose(filepath)