Extend tests for restarts

CliMA · Sep 20, 2024 · 2e14bd6 · 2e14bd6
1 parent 448d485
commit 2e14bd6
Show file tree

Hide file tree

Showing 3 changed files with 310 additions and 204 deletions.
diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml
@@ -447,17 +447,23 @@ steps:
         #retry:
         #  automatic: true
 
-      - label: ":computer: MPI GPU test restart"
+      - label: ":computer: test restart"
         command: >
-          srun julia --color=yes --project=examples test/restart.jl
-        env:
-          CLIMACOMMS_CONTEXT: "MPI"
-          CLIMACOMMS_DEVICE: "CUDA"
+          julia --color=yes --project=examples test/restart.jl
         agents:
-          slurm_gpus_per_task: 1
-          slurm_ntasks: 2
           slurm_mem: 16GB
 
+      # - label: ":computer: MPI GPU test restart"
+      #   command: >
+      #     srun julia --color=yes --project=examples test/restart.jl
+      #   env:
+      #     CLIMACOMMS_CONTEXT: "MPI"
+      #     CLIMACOMMS_DEVICE: "CUDA"
+      #   agents:
+      #     slurm_gpus_per_task: 1
+      #     slurm_ntasks: 2
+      #     slurm_mem: 16GB
+
       - label: ":computer: MPI no lim aquaplanet (ρe) equilmoist clearsky radiation"
         command: >
           srun julia --color=yes --project=examples examples/hybrid/driver.jl

diff --git a/src/solver/type_getters.jl b/src/solver/type_getters.jl
@@ -301,9 +301,7 @@ function get_state_restart(config::AtmosConfig, restart_file, atmos_model_hash)
     atmos_model_hash_in_restart =
         InputOutput.HDF5.read_attribute(reader.file, "atmos_model_hash")
     if atmos_model_hash_in_restart != atmos_model_hash
-        error(
-            "Restart file $(restart_file) was constructed with a different AtmosModel",
-        )
+        @warn "Restart file $(restart_file) was constructed with a different AtmosModel, no consistency check was performed"
     end
     return (Y, t_start)
 end
@@ -484,6 +482,42 @@ thermo_state_type(::EquilMoistModel, ::Type{FT}) where {FT} = TD.PhaseEquil{FT}
 thermo_state_type(::NonEquilMoistModel, ::Type{FT}) where {FT} =
     TD.PhaseNonEquil{FT}
 
+function get_restart_file(config::AtmosConfig, base_output_dir)
+    restart_file = nothing
+    (; parsed_args) = config
+
+    if parsed_args["detect_restart_file"] && isdir(base_output_dir)
+        lowercase(parsed_args["output_dir_style"]) == "activelink" ||
+            error("detect_restart_file works only with ActiveLink")
+        # output_dir will be something like ABC/DEF/output_1234
+        name_rx = r"output_(\d\d\d\d)"
+        restart_file_rx = r"day\d+\.\w+\.hdf5"
+
+        existing_outputs =
+            filter(x -> !isnothing(match(name_rx, x)), readdir(base_output_dir))
+        if !isempty(existing_outputs)
+            latest_output = first(sort(existing_outputs, rev = true))
+            previous_folder = joinpath(base_output_dir, latest_output)
+            possible_restart_files = filter(
+                f -> occursin(restart_file_rx, f),
+                readdir(previous_folder),
+            )
+            if !isempty(possible_restart_files)
+                restart_file_name =
+                    last(CA.sort_files_by_time(possible_restart_files))
+                restart_file = joinpath(previous_folder, restart_file_name)
+            end
+        end
+    end
+
+    # If a restart file was passed, override what we detected automatically
+    if !isnothing(parsed_args["restart_file"])
+        restart_file = parsed_args["restart_file"]
+    end
+
+    return restart_file
+end
+
 function get_sim_info(config::AtmosConfig)
     (; parsed_args) = config
     FT = eltype(config)
@@ -503,53 +537,16 @@ function get_sim_info(config::AtmosConfig)
     haskey(allowed_dir_styles, lowercase(requested_style)) ||
         error("output_dir_style $(requested_style) not available")
 
+    # We look for a restart before creating a new output dir because we want to
+    # look for previous folders
+    restart_file = get_restart_file(config, base_output_dir)
+
     output_dir = OutputPathGenerator.generate_output_path(
         base_output_dir;
         context = config.comms_ctx,
         style = allowed_dir_styles[lowercase(requested_style)],
     )
 
-    restart_file = nothing
-
-    if parsed_args["detect_restart_file"]
-        lowercase(parsed_args["output_dir_style"]) == "activelink" ||
-            error("detect_restart_file works only with ActiveLink")
-        # output_dir will be something like ......./output_1234
-        name_rx = r"output_(\d\d\d\d)"
-        restart_file_rx = r"day\d+\.\w+\.hdf5"
-        counter_match = match(name_rx, output_dir)
-        if !isnothing(counter_match)
-            counter = parse(Int, counter_match[1])
-            if counter != 0
-                # As implemented in ClimaUtilities.OutputPathGenerator
-                previous_counter = counter - 1
-                previous_counter_str = lpad(previous_counter, 4, "0")
-                counter_str = lpad(counter, 4, "0")
-
-                previous_folder =
-                    replace(output_dir, counter_str => previous_counter_str)
-
-                isdir(previous_folder) ||
-                    error("Could not find a folder in $(previous_folder)")
-
-                possible_restart_files = filter(
-                    f -> occursin(restart_file_rx, f),
-                    readdir(previous_folder),
-                )
-                if !isempty(possible_restart_files)
-                    restart_file_name =
-                        last(CA.sort_files_by_time(possible_restart_files))
-                    restart_file = joinpath(previous_folder, restart_file_name)
-                end
-            end
-        end
-    end
-
-    # If a restart file was passed, override what we detected automatically
-    if !isnothing(parsed_args["restart_file"])
-        restart_file = parsed_args["restart_file"]
-    end
-
     isnothing(restart_file) ||
         @info "Restarting simulation from file $restart_file"