Is the memory usage lower for NetCDFOutputWriter compared to JLD2OutputWriter on GPU?

[liuchihl]

My simulation, which is (150x300x100) in size, was running on a GPU V100. After 532 iterations, which took about 1 hour of wall time, it terminated due to an out-of-memory (OOM) error when I was using the JLD2OutputWriter to write the output. However, when I switched to using the NetCDFOutputWriter for the exact same simulation and diagnostics, the simulation completed in about 10 minutes of wall time (2219 iterations) without any memory issues. I'm unsure what caused this significant difference in memory usage between the two output writers. Does this issue specifically occur for large (or not-small) simulations on GPUs? I'm pleased with how much more efficient the NetCDF output writer is, but I'd like to understand the reason behind this difference. Any insight would be much appreciated! Thank you!

(I am using Oceananigans v.90.9 and Julia 1.9.3)

[glwagner]

Do you mind sharing your script? Both the outputs and the model configuration can affect how the output writers behave.

There are not any major differences in memory consumption between the two output writers that I know of. There may be some small differences because they are initialized differently (with JLD2, we serialize some model properties for convenience). So, your experience could be indicative of a bug. I also don't think there is any reason why the simulations would involve different numbers of iterations (if the schedule is the same, the number of iterations should be similar modulo any randomness in the initial condition).

[liuchihl]

Sure, no problem, here is the script.

Yes, the schedule is identical for both writers. The difference lies in the fact that with JLD2OutputWriter, the model stopped at iteration 532, whereas with the NetCDF writer, it completed 2219 iterations, which is the number of iterations to finish this example.

[liuchihl]

Sure I can open an issue.

Could this be the cause of the discrepancy?

Ah, you're absolutely right, I should have saved the same diagnostics. That is indeed the discrepancy. After using all_diags, the initialization time for NetCDFOutputWriter increased to pretty much the same as JLD2OutputWriter (which is ~36 minutes).
This long initialization issue is consistent with what people have found in #3381.

Interesting that "all diags" takes way longer to compile than just "custom diags". It doesn't really make sense.

This is very surprising, but when upgrading to Julia 1.10, this long-initialization issue no longer persists when using both writers.

[glwagner]

Ok, I should have read more closely the first time!! Glad we caught the issue though.

To summarize:

• Memory usage is indeed fairly similar between NetCDF and JLD2 output writers
• It'd be nice to save aspects of the immersed boundary grid when using NetCDF
• There's a long initialization issue on 1.9 but it's solved by 1.10, as far as we can tell...

[liuchihl]

Thanks for spotting that!!

I think I actually found the main issue of my case:
I noticed the way I define my diagnostics is to merge different types of data arrays. For example, model.velocities and model.tracers are different from custom_diags = (B=B, uhat=û, what=ŵ):
https://github.com/liuchihl/internal-tide-mixing/blob/45bc947888d77db9926d330ef1e686431ce91413/tilted_internal_tide_experiment.jl#L177

The long initialization issue has been resolved, and now only takes a few seconds once I separate the two output writing processes. My guess is that the output writer operates most efficiently when the diagnostics are consistent in types, though merging them into just one file would be ideal.

[glwagner]

That does seem like a culprit --- mixing Field with AbstractOperation.

That said, users should be able to mix these two. It's inconvenient to have to use two different output writers! I bet we can fix things by using dictionaries rather than NamedTuple for outputs --- or some silly thing like that.

Do you still have long initialization time on julia 1.10?

[liuchihl]

No, the good news is that there is no long initialization time on Julia 1.10 even when mixing different types of diagnostics.

[simone-silvestri]

Could you also share the error? The GPU should not be a problem as the output is transferred to the CPU before writing. If it's an OOM error it means we are consuming too much RAM, but, as far as I know, there should be no heavy memory allocation in the output writers so it might mean that there is a bug somewhere.
Knowing the line where the failure happens might indicate where the memory allocation happens

[liuchihl]

Yes, of course. Below are the final lines of the output message along with the error when using JLD2outputwriter:

[ Info: Fetching JLD2 output (:u, :v, :w, :b, :B, :uhat, :what)...
[ Info: Fetching time: 116.340 ms
[ Info: Writing JLD2 output (:u, :v, :w, :b, :B, :uhat, :what) to ./output/internal_tide_3days-theta=0.0036_realtopo3D_Nx150_fields.jld2...
[ Info: Writing done: time=1.047 seconds, size=10.1 GiB, Δsize=241.0 MiB
[ Info: Fetching JLD2 output (:B, :uhat, :what)...
[ Info: Fetching time: 502.184 μs
[ Info: Writing JLD2 output (:B, :uhat, :what) to ./output/internal_tide_3days-theta=0.0036_realtopo3D_Nx150_slices.jld2...
[ Info: Writing done: time=49.503 ms, size=15.6 MiB, Δsize=353.1 KiB
/export/spool/slurm/slurmd.spool/job30694661/slurm_script: line 21: 1106691 Killed                  julia --project $rundir > $rundir_out.txt
slurmstepd: error: Detected 1 oom_kill event in StepId=30694661.batch. Some of the step tasks have been OOM Killed.

Knowing the line where the failure happens might indicate where the memory allocation happens

I agree, but it's not clear to me where the failure might have occurred.