Merge Pull Request #3052 from E3SM-Project/scream/bartgol/eamxx/flush-output-with-rhist

Automatically Merged using E3SM Pull Request AutoTester
PR Title: A couple of fixes and an addition for IO
PR Author: bartgol
PR LABELS: enhancement, I/O, AT: AUTOMERGE, bugfix

Author: E3SM-Bot

components/eamxx/src/share/io/scream_io_file_specs.hpp

@@ -14,17 +14,21 @@ namespace scream
 {
 
 // How the file capacity is specified
+// NOTE: Keep Yearly=0, Monthly=1, Daily=2, so you can use them
+//       to access a TimeStamp date at the correct index in StorageSpecs methods
 enum StorageType {
-  NumSnaps,   // Fixed number of snaps per file
-  Monthly,    // Each file contains output for one month
-  Yearly      // Each file contains output for one year
+  Yearly  = 0,  // Each file contains output for one year
+  Monthly = 1,  // Each file contains output for one month
+  Daily   = 2,  // Each file contains output for one day
+  NumSnaps      // Fixed number of snaps per file
 };
 
 inline std::string e2str (const StorageType st) {
   switch (st) {
     case NumSnaps: return "num_snapshots";
     case Yearly:   return "one_year";
     case Monthly:  return "one_month";
+    case Daily:    return "one_day";
     default:       return "unknown";
   }
 }
@@ -33,21 +37,20 @@ struct StorageSpecs {
 
   StorageType type = NumSnaps;
 
-  // Current index ***in terms of this->type***
-  // If type==NumSnaps, curr_idx=num_snapshots_in_file,
-  // otherwise it is the month/year index stored in this file
+  // Current index for type!=NumSnaps. It stores the year/month/day
+  // index associated with this file
   int  curr_idx = -1;
 
   // A snapshot fits if
   //  - type=NumSnaps: the number of stored snaps is less than the max allowed per file.
-  //  - otherwise: the snapshot month/year index match the one currently stored in the file
+  //  - otherwise: the snapshot year/month/day index match the one currently stored in the file
   //               or the file has no snapshot stored yet
   bool snapshot_fits (const util::TimeStamp& t) const {
-    const auto& idx = type==Monthly ? t.get_month() : t.get_year();
     switch (type) {
-      case Yearly:
-      case Monthly:
-        return curr_idx==-1 or curr_idx==idx;
+      case Yearly:  [[fallthrough]];
+      case Monthly: [[fallthrough]];
+      case Daily:
+        return curr_idx==-1 or curr_idx==t.get_date()[static_cast<int>(type)];
       case NumSnaps:
         return num_snapshots_in_file<max_snapshots_in_file;
       default:
@@ -56,10 +59,15 @@ struct StorageSpecs {
   }
 
   void update_storage (const util::TimeStamp& t) {
+    // We always update the snap counter (regardless of storage type),
+    // so that FileSpecs can correctly detect if it needs flushing
     switch (type) {
-      case Yearly:    curr_idx = t.get_year();  break;
-      case Monthly:   curr_idx = t.get_month(); break;
-      case NumSnaps:  ++num_snapshots_in_file;  break;
+      case Yearly:  [[fallthrough]];
+      case Monthly: [[fallthrough]];
+      case Daily:
+          curr_idx = t.get_date()[static_cast<int>(type)]; [[fallthrough]];
+      case NumSnaps:
+        ++num_snapshots_in_file;  break;
       default:
         EKAT_ERROR_MSG ("Error! Unrecognized/unsupported file storage type.\n");
     }

components/eamxx/src/share/io/scream_output_manager.cpp

@@ -228,8 +228,7 @@ setup (const std::map<std::string,std::shared_ptr<fm_type>>& field_mgrs,
       const auto& last_output_filename = get_attribute<std::string>(rhist_file,"GLOBAL","last_output_filename");
       m_resume_output_file = last_output_filename!="" and not restart_pl.get("force_new_file",false);
       if (m_resume_output_file) {
-        int num_snaps = scorpio::get_attribute<int>(rhist_file,"GLOBAL","last_output_file_num_snaps");
-        m_output_file_specs.storage.num_snapshots_in_file = num_snaps;
+        m_output_file_specs.storage.num_snapshots_in_file = scorpio::get_attribute<int>(rhist_file,"GLOBAL","last_output_file_num_snaps");
 
         if (m_output_file_specs.storage.snapshot_fits(m_output_control.next_write_ts)) {
           // The setup_file call will not register any new variable (the file is in Append mode,
@@ -500,10 +499,7 @@ void OutputManager::run(const util::TimeStamp& timestamp)
           write_timestamp (filespecs.filename,"last_write",m_output_control.last_write_ts,true);
           scorpio::set_attribute (filespecs.filename,"GLOBAL","last_output_filename",m_output_file_specs.filename);
           scorpio::set_attribute (filespecs.filename,"GLOBAL","num_snapshots_since_last_write",m_output_control.nsamples_since_last_write);
-
-          int nsnaps = m_output_file_specs.is_open
-                     ? scorpio::get_dimlen(m_output_file_specs.filename,"time") : 0;
-          scorpio::set_attribute (filespecs.filename,"GLOBAL","last_output_file_num_snaps",nsnaps);
+          scorpio::set_attribute (filespecs.filename,"GLOBAL","last_output_file_num_snaps",m_output_file_specs.storage.num_snapshots_in_file);
         }
         // Write these in both output and rhist file. The former, b/c we need these info when we postprocess
         // output, and the latter b/c we want to make sure these params don't change across restarts
@@ -563,6 +559,11 @@ void OutputManager::run(const util::TimeStamp& timestamp)
     }
     if (is_checkpoint_step) {
       write_global_data(m_checkpoint_control,m_checkpoint_file_specs);
+
+      // Always flush output during checkpoints (assuming we opened it already)
+      if (m_output_file_specs.is_open) {
+        scorpio::flush_file (m_output_file_specs.filename);
+      }
     }
     stop_timer(timer_root+"::update_snapshot_tally");
     if (is_output_step && m_time_bnds.size()>0) {
@@ -636,16 +637,16 @@ compute_filename (const IOFileSpecs& file_specs,
   auto ts = (m_avg_type==OutputAvgType::Instant || file_specs.ftype==FileType::HistoryRestart)
           ? timestamp : control.last_write_ts;
 
+  int ts_string_len = 0;
   switch (file_specs.storage.type) {
-    case NumSnaps:
-      filename += "." + ts.to_string(); break;
-    case Yearly:
-      filename += "." + std::to_string(ts.get_year()); break;
-    case Monthly:
-      filename += "." + std::to_string(ts.get_year()) + "-" + std::to_string(ts.get_month()); break;
+    case Yearly:   ts_string_len = 4;  break; // YYYY
+    case Monthly:  ts_string_len = 7;  break; // YYYY-MM
+    case Daily:    ts_string_len = 10; break; // YYYY-MM-DD
+    case NumSnaps: ts_string_len = 16; break; // YYYY-MM-DD-XXXXX
     default:
       EKAT_ERROR_MSG ("Error! Unrecognized/unsupported file storage type.\n");
   }
+  filename += "." + ts.to_string().substr(0,ts_string_len);
 
   return filename + ".nc";
 }
@@ -701,6 +702,8 @@ setup_internals (const std::map<std::string,std::shared_ptr<fm_type>>& field_mgr
       storage.type = Yearly;
     } else if (storage_type=="one_month") {
       storage.type = Monthly;
+    } else if (storage_type=="one_day") {
+      storage.type = Daily;
     } else {
       EKAT_ERROR_MSG ("Error! Unrecognized/unsupported file storage type.\n");
     }

components/eamxx/src/share/io/tests/io_monthly.cpp

@@ -178,11 +178,11 @@ void read (const int seed, const ekat::Comm& comm)
   // Get filename from timestamp
   std::string casename = "io_monthly";
   auto get_filename = [&](const util::TimeStamp& t) {
+    auto t_str = t.to_string().substr(0,7);
     std::string fname = casename
                       + ".INSTANT.nsteps_x1"
                       + ".np" + std::to_string(comm.size())
-                      + "." + std::to_string(t.get_year())
-                      + "-" + std::to_string(t.get_month())
+                      + "." + t_str
                       + ".nc";
     return fname;
   };