Automatically add extension when doing distributed output writing (#4179

)
CliMA · Mar 9, 2025 · d2e9e27 · d2e9e27
1 parent 8e8745d
commit d2e9e27
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 0 deletions.
diff --git a/src/OutputWriters/jld2_output_writer.jl b/src/OutputWriters/jld2_output_writer.jl
@@ -157,6 +157,7 @@ function JLD2OutputWriter(model, outputs; filename, schedule,
 
     mkpath(dir)
     filename = auto_extension(filename, ".jld2")
+    filename = with_architecture_suffix(architecture(model), filename, ".jld2")
     filepath = abspath(joinpath(dir, filename))
 
     initialize!(file_splitting, model)

diff --git a/src/OutputWriters/output_writer_utils.jl b/src/OutputWriters/output_writer_utils.jl
@@ -211,3 +211,17 @@ output_averaging_schedule(output) = nothing # fallback
 
 show_array_type(a::Type{Array{T}}) where T = "Array{$T}"
 
+#####
+##### Architecture suffix
+#####
+
+with_architecture_suffix(arch, filename, ext) = filename
+
+function with_architecture_suffix(arch::Distributed, filename, ext)
+    Ne = length(ext)
+    prefix = filename[1:end-Ne]
+    rank = arch.local_rank
+    prefix *= "_rank$rank"
+    return prefix * ext
+end
+