Loading/using data from GPU model run?

[sdbrenner]

This question relates to #2360, but I don't think it's a repeat. Apologies if this is a naive question: I am new to Oceananigans, Julia, and GPU computing.

Summary

I am running a simulation where my surface boundary condition is defined in terms of an array rather than a function. I am running the simulation on a GPU, so the array is a CuArray. When I try to read in the output data from the simulation after, using FieldTimeSeries, I receive an error that seems to stem from the definition of the boundary condition as a CuArray:

ERROR: ArgumentError: top FluxBoundaryCondition: 100×100 CUDA.CuArray{Float64, 2, CUDA.DeviceMemory} must use `Array` rather than `CuArray` on CPU architectures!

I am using Oceananigans v0.93.1, CUDA v5.5.2

Minimal Working Example

The following minimal simulation can be used to reproduce the issue—the simulation runs without error, but loading the resulting output data produces the error

using Oceananigans
using CUDA

# Grid
grid = RectilinearGrid(GPU(),
                       size = (100, 100,10),
                       x = (0, 10e3),
                       y = (0, 10e3),
                       z = (0, 100),
                       topology = (Periodic, Periodic, Bounded))

# Boundary conditions
const sbf = 1e-7
Qb = sbf * ones(100,100) |> CuArray
buoyancy_bcs = FieldBoundaryConditions(top = FluxBoundaryCondition(Qb))

# Model
model = NonhydrostaticModel(; grid, 
                              buoyancy = BuoyancyTracer(),
                              tracers = (:b),
                              boundary_conditions = (; b=buoyancy_bcs))

# Initial conditions
ϵ(x, y, z) = 1e-4*(2rand() - 1) # noise
set!(model, b=ϵ )

# Instansiate simulation
simulation = Simulation(model; Δt=30, stop_iteration=100)

# Output
simulation.output_writers[:b] =
    JLD2OutputWriter(model, model.tracers,
                     filename = "test_Qb_array.jld2",
                     schedule=IterationInterval(10),
                     overwrite_existing = true)

# Run
run!(simulation)

After running the simulation, trying to load the output data with FieldTimeSeries produces the error:

julia> FieldTimeSeries("test_Qb_array.jld2", "b")
ERROR: ArgumentError: top FluxBoundaryCondition: 100×100 CUDA.CuArray{Float64, 2, CUDA.DeviceMemory} must use `Array` rather than `CuArray` on CPU architectures!
Stacktrace:
  [1] validate_boundary_condition_architecture(::CUDA.CuArray{…}, ::CPU, bc::BoundaryCondition{…}, side::Symbol)
    @ Oceananigans.BoundaryConditions ~/.julia/packages/Oceananigans/s1DfC/src/BoundaryConditions/boundary_condition.jl:148
  [2] validate_boundary_condition_architecture(bc::BoundaryCondition{Oceananigans.BoundaryConditions.Flux, CUDA.CuArray{…}}, arch::CPU, side::Symbol)
    @ Oceananigans.BoundaryConditions ~/.julia/packages/Oceananigans/s1DfC/src/BoundaryConditions/boundary_condition.jl:141
  [3] validate_boundary_conditions(loc::Tuple{…}, grid::RectilinearGrid{…}, bcs::FieldBoundaryConditions{…})
    @ Oceananigans.Fields ~/.julia/packages/Oceananigans/s1DfC/src/Fields/field.jl:77
  [4] #apply_regionally!#56
    @ ~/.julia/packages/Oceananigans/s1DfC/src/Utils/multi_region_transformation.jl:121 [inlined]
  [5] apply_regionally!
    @ ~/.julia/packages/Oceananigans/s1DfC/src/Utils/multi_region_transformation.jl:118 [inlined]
  [6] macro expansion
    @ ~/.julia/packages/Oceananigans/s1DfC/src/Utils/multi_region_transformation.jl:206 [inlined]
  [7] Field(loc::Tuple{…}, grid::RectilinearGrid{…}, data::OffsetArrays.OffsetArray{…}, bcs::FieldBoundaryConditions{…}, indices::Tuple{…}, op::Nothing, status::Nothing)
    @ Oceananigans.Fields ~/.julia/packages/Oceananigans/s1DfC/src/Fields/field.jl:91
  [8] #Field#13
    @ ~/.julia/packages/Oceananigans/s1DfC/src/Fields/field.jl:180 [inlined]
  [9] Field
    @ ~/.julia/packages/Oceananigans/s1DfC/src/Fields/field.jl:171 [inlined]
 [10] Field(location::Tuple{…}, path::String, name::String, iter::Int64; grid::Nothing, architecture::CPU, indices::Tuple{…}, boundary_conditions::FieldBoundaryConditions{…})
    @ Oceananigans.OutputReaders ~/.julia/packages/Oceananigans/s1DfC/src/OutputReaders/field_time_series.jl:569
 [11] set!(fts::FieldTimeSeries{…}, path::String, name::String)
    @ Oceananigans.OutputReaders ~/.julia/packages/Oceananigans/s1DfC/src/OutputReaders/set_field_time_series.jl:38
 [12] FieldTimeSeries(path::String, name::String; backend::InMemory{…}, architecture::Nothing, grid::Nothing, location::Nothing, boundary_conditions::Oceananigans.OutputReaders.UnspecifiedBoundaryConditions, time_indexing::Oceananigans.OutputReaders.Linear, iterations::Nothing, times::Nothing)
    @ Oceananigans.OutputReaders ~/.julia/packages/Oceananigans/s1DfC/src/OutputReaders/field_time_series.jl:526
 [13] FieldTimeSeries(path::String, name::String)
    @ Oceananigans.OutputReaders ~/.julia/packages/Oceananigans/s1DfC/src/OutputReaders/field_time_series.jl:431
 [14] top-level scope
    @ REPL[4]:1
Some type information was truncated. Use `show(err)` to see complete types.

Note that jldopen("test_Qb_array.jld2") does not produce errors.

Comparison to #2360

#2360 also related to issues using FieldTimeSeries to open output from a GPU-simulation; however, the error that occurred was different and the issue was resolved by updating Oceananigans to v0.72.4. Morover, the author of that question noted that:

But all initial and boundary conditions are defined using functions.

If, in my MWE, I replace Qb = sbf * ones(100,100) |> CuArray with Qb(x,y,t) = sbf (i.e., defining the boundary condition using a function), then I no longer receive an error when loading the data with FieldTimeSeries.

More context

I am trying to run coupled simulations with Oceananigans and a sea ice discrete-element model. The ice model responds to the ocean currents and produces surface stress fields that affect the ocean, so those fields must be passed back and forth. The ice model produces stress fields as arrays.
Coupling is performed by stepping each model forward in a loop, something like:

for n = 1:N
    global Qu, Qv 
    
    # Pass ocean currents to sea ice model
    iceSimulation.model.ocean.u .= Array(interior(ocnModel.velocities.u))[:, :, end]
    iceSimulation.model.ocean.v .= Array(interior(ocnModel.velocities.v))[:, :, end]  

    # Pass surface stresses to ocean model
    Qu .= -iceSimulation.model.ocean.τx/ρₒ |> CuArray
    Qv .= -iceSimulation.model.ocean.τy/ρₒ |> CuArray

    # Run sea ice model
    <run sea ice model>

    # Run Ocean Model
    Oceananigans.time_step!(ocnSimulation)
end

While I can run the model like this, I don't know how to load the model output.

Follow-up question

In the various Oceananigans examples in the documentation, timeseries output is loaded to create animations of the results. Are there any simple examples available or suggested practices for further analysis/calculations using those outputs (e.g., computing various budgets, spectra, etc.)? I'm still wrapping my head around the Oceananigans/Julia variable structures and how to use them effectively.

[glwagner]

Not a naive question! Though naive questions are fine.

I would actually regard this as a bug. I am surprised we haven't run into it before...

I think it's because we serialize the CuArray boundary condition directly (which we should not do). But when it's a function we do nothing, eg

Oceananigans.jl/src/OutputWriters/output_writer_utils.jl

Line 147 in 6c40d7e

serializeproperty!(file, address * "/boundary_conditions", boundary_conditions(f))

Oceananigans.jl/src/OutputWriters/output_writer_utils.jl

Lines 133 to 141 in 6c40d7e

	function serializeproperty!(file, address, p::FieldBoundaryConditions)
	# TODO: it'd be better to "filter" `FieldBoundaryCondition` and then serialize
	# rather than punting with `missing` instead.
	if has_reference(Function, p)
	file[address] = missing
	else
	file[address] = p
	end
	end

we need an on_architecture(CPU(), bcs::FieldBoundaryCondition)!

As a matter of best practice though I would try to use

Qb = Field{Center, Center, Nothing}(grid)

As your boundary condition, which comes with a bunch of benefits including participating in abstract operations, etc. I'm not sure that will fix your issue though...

Coupling is performed by stepping each model forward in a loop, something like:

You should be able to use Callback for this, rather than manually executing a loop. But that's up to you, the loop works too. It would also be interesting to support this in ClimaOcean, which implements an OceanSeaIceModel abstraction for coupled modeling. ClimaOcean is still something of a WIP, but also this functionality is probably not too difficult to implement.

While I can run the model like this, I don't know how to load the model output.

I'm not sure what you mean. I think the output writers should still work right?

Are there any simple examples available or suggested practices for further analysis/calculations using those outputs (e.g., computing various budgets, spectra, etc.)?

The same operations that produce diagnostics for output can be used for post-processing. The horizontal convection example probably has the most post-processing. We don't have any example of computing spectra. Oceanostics might have something though.

[sdbrenner]

Thanks for the help and feedback. I think swapping to Qb = Field{Center, Center, Nothing}(grid) might fix the problem. Adapting your MWE:

using Oceananigans
using CUDA

x = y = z = (0, 1)
grid = RectilinearGrid(GPU(); size=(1, 1, 1), x, y, z)

τx = Field{Center, Center, Nothing}(grid)
u_bcs = FieldBoundaryConditions(top = FluxBoundaryCondition(τx))
model = NonhydrostaticModel(; grid, boundary_conditions = (; u=u_bcs))
simulation = Simulation(model; Δt=1, stop_iteration=1)

simulation.output_writers[:jld2] = JLD2OutputWriter(model, model.velocities,
                                                    filename = "test_cuarray_bc.jld2",
                                                    schedule=IterationInterval(1),
                                                    overwrite_existing = true)

τx .= CuArray(zeros(size(grid)...))
run!(simulation)

Then, after running this, I am able to load the output without the error appearing.

julia> FieldTimeSeries("test_cuarray_bc.jld2", "u")
1×1×1×1 FieldTimeSeries{InMemory} located at (Face, Center, Center) of u at test_cuarray_bc.jld2
├── grid: 1×1×1 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 1×1×1 halo
├── indices: (:, :, :)
├── time_indexing: Linear()
├── backend: InMemory()
├── path: test_cuarray_bc.jld2
├── name: u
└── data: 3×3×3×1 OffsetArray(::Array{Float64, 4}, 0:2, 0:2, 0:2, 1:1) with eltype Float64 with indices 0:2×0:2×0:2×1:1
    └── max=0.0, min=0.0, mean=0.0

I will try this approach in my actual script and report back if I run into issues.

To address your other comments:

Coupling is performed by stepping each model forward in a loop, something like:

You should be able to use Callback for this, rather than manually executing a loop. But that's up to you, the loop works too. It would also be interesting to support this in ClimaOcean, which implements an OceanSeaIceModel abstraction for coupled modeling. ClimaOcean is still something of a WIP, but also this functionality is probably not too difficult to implement.

I did notice #2693 talking about callbacks for setting up coupling. I may experiment with that in the future.

While I can run the model like this, I don't know how to load the model output.

I'm not sure what you mean. I think the output writers should still work right?

Sorry, this was a poorly-worded statement. I was referring specifically to this issue; i.e., I can run the model and it writes output data, but if I try to load that model output with FieldTimeSeries, I encounter this error... so because of the error, I don't know how to load/use the output. Hopefully the swap to Qb = Field{Center, Center, Nothing}(grid) remedies this.

The same operations that produce diagnostics for output can be used for post-processing. The horizontal convection example probably has the most post-processing.

I hadn't noticed some of the additional operations at the end of that example. That is actually really useful. I'll take a look at Oceanostics too.

Thanks again!

[glwagner]

Let us know how we can make things better and easier to find if you think of something!

[glwagner]

Okay, PR #3893 may help.

I think the example in the original post has a few elements that aren't needed. Is this more minimal?

using Oceananigans
using CUDA

x = y = z = (0, 1)
grid = RectilinearGrid(GPU(); size=(1, 1, 1), x, y, z)

τx = CuArray(zeros(size(grid)...))
u_bcs = FieldBoundaryConditions(top = FluxBoundaryCondition(τx))
model = NonhydrostaticModel(; grid, boundary_conditions = (; u=u_bcs))
simulation = Simulation(model; Δt=1, stop_iteration=1)

simulation.output_writers[:jld2] = JLD2OutputWriter(model, model.velocities,
                                                    filename = "test_cuarray_bc.jld2",
                                                    schedule=IterationInterval(1),
                                                    overwrite_existing = true)

run!(simulation)

# Later, without a CPU...
using Oceananigans
ut = FieldTimeSeries("test_cuarray_bc.jld2", "u")

Have to think about how to run this because we have to hide the GPU from the process.

Actually no, this error works even when a GPU is available. Which makes sense, because the problem is that the grid is converted to CPU but the boundary conditions are not.

[sdbrenner]

Thanks! I guess I am also still learning which pieces are really necessary... this minimal example is much better.

[glwagner]

It takes time to learn how to build a minimal example, but its a very useful thing to learn because it is how we develop a precise understanding of what the code means, allowing us to be very intentional when designing numerical experiments