Use driver-owned PIO settings for coupled Omega initialization

components/omega/cime_config/omega_buildnml/tests/test_validate_user_overrides.py

@@ -13,6 +13,13 @@ def defaults():
             "StartTime": "0001-01-01_00:00:00",
         },
         "Tendencies": {"SurfaceTracerRestoringEnable": False},
+        "IO": {
+            "IOTasks": 1,
+            "IOStride": 1,
+            "IOBaseTask": 0,
+            "IORearranger": "box",
+            "IODefaultFormat": "pnetcdf",
+        },
         "IOStreams": {
             "InitialState": {"Filename": "ocean.nc"},
             "History": {"Freq": 1, "FreqUnits": "months"},
@@ -85,3 +92,24 @@ def test_overriding_an_existing_stream_is_allowed(defaults):
     validated = validate_user_overrides(user_overrides, defaults)
 
     assert validated == user_overrides
+
+
+@pytest.mark.parametrize("option", ["IOBaseTask", "IORearranger"])
+def test_driver_owned_io_options_are_rejected(option, defaults):
+    """
+    The base IO task and rearranger are owned by the driver (CIME/shr_pio),
+    so a user may not override them in ``user_nl_omega``.
+    """
+    user_overrides = {"IO": {option: 4 if option == "IOBaseTask" else "subset"}}
+
+    with pytest.raises(ValueError, match="cannot be overridden"):
+        validate_user_overrides(user_overrides, defaults)
+
+
+def test_component_configurable_io_options_are_allowed(defaults):
+    """``IOTasks`` and ``IOStride`` remain component-configurable."""
+    user_overrides = {"IO": {"IOTasks": 8, "IOStride": 2}}
+
+    validated = validate_user_overrides(user_overrides, defaults)
+
+    assert validated == user_overrides

components/omega/cime_config/omega_buildnml/validate.py

@@ -47,6 +47,11 @@
         "TimeIntegration.RunDuration",
         # calendar must agree with the CIME ``CALENDAR`` setting
         "TimeIntegration.CalendarType",
+        # base IO task and rearranger are owned by the driver (CIME/shr_pio)
+        # so they stay consistent with the rest of the case. IOTasks and
+        # IOStride remain component-configurable.
+        "IO.IOBaseTask",
+        "IO.IORearranger",
     }
 )
 

components/omega/src/base/IO.cpp

@@ -51,6 +51,24 @@ RearrFromString(const std::string &Rearr // [in] choice of IO rearranger
 
 } // End RearrFromString
 
+//------------------------------------------------------------------------------
+// Converts the PIO rearranger enum back to the name used in the config file.
+// Only used for log messages. RearrMap is searched rather than adding a second
+// table so the two cannot drift apart. "default" is skipped because it aliases
+// "box" and the recognisable name is wanted here. Falls back to the numeric
+// value for anything not in the map.
+static std::string RearrToString(const Rearranger Rearr // [in] rearranger enum
+) {
+
+   for (const auto &Entry : RearrMap) {
+      if (Entry.second == Rearr && Entry.first != "default")
+         return Entry.first;
+   }
+
+   return std::to_string(static_cast<int>(Rearr));
+
+} // End RearrToString
+
 //------------------------------------------------------------------------------
 // Converts string choice for File Format to an enum
 FileFmt
@@ -125,11 +143,54 @@ IfExists IfExistsFromString(
 // Methods
 //------------------------------------------------------------------------------
 // Initializes the IO system based on configuration inputs and
-// default MPI communicator
+// default MPI communicator. The base task and rearranger are read from the
+// Omega configuration and forwarded to the driver-owned overload below.
 void init(const MPI_Comm &InComm // [in] MPI communicator to use
 ) {
 
-   // Retrieve parallel IO parameters from the Omega configuration
+   // Retrieve the driver-owned parallel IO parameters from the Omega
+   // configuration. These are read here for standalone runs and unit tests
+   // where there is no coupler to supply them.
+
+   Error Err;
+   Config *OmegaConfig = Config::getOmegaConfig();
+   OMEGA_REQUIRE(OmegaConfig, "Null OmegaConfig pointer in IO::init");
+
+   // Read IO subconfiguration
+   Config IOConfig("IO");
+   Err = OmegaConfig->get(IOConfig);
+   CHECK_ERROR_ABORT(Err, "IO: IO group not found in input Config");
+
+   // Read parallel IO settings - default to single-task if config
+   // values do not exist
+   int IOBaseTask           = 0;
+   std::string InRearranger = "box";
+
+   Err = IOConfig.get("IOBaseTask", IOBaseTask);
+   CHECK_ERROR_WARN(Err, "IO: IOBaseTask not found in Config - using {}",
+                    IOBaseTask);
+
+   Err = IOConfig.get("IORearranger", InRearranger);
+   CHECK_ERROR_WARN(Err, "IO: Rearranger not found in Config - using {}",
+                    InRearranger);
+
+   IOInitParams IOParams{IOBaseTask, RearrFromString(InRearranger)};
+   init(InComm, IOParams);
+
+   return;
+
+} // end init
+
+//------------------------------------------------------------------------------
+// Initializes the IO system using driver-owned IO parameters for the base
+// task and rearranger. The number of IO tasks and the IO stride are still
+// read from the Omega configuration.
+void init(const MPI_Comm &InComm,      // [in] MPI communicator to use
+          const IOInitParams &IOParams // [in] driver-owned IO parameters
+) {
+
+   // Retrieve the remaining parallel IO parameters from the Omega
+   // configuration
 
    Error Err;
    Config *OmegaConfig = Config::getOmegaConfig();
@@ -147,12 +208,10 @@ void init(const MPI_Comm &InComm // [in] MPI communicator to use
                     InFileFmt);
    DefaultFileFmt = FileFmtFromString(InFileFmt);
 
-   // Read parallel IO settings - default to single-task if config
-   // values do not exist
-   int NumIOTasks           = 1;
-   int IOStride             = 1;
-   int IOBaseTask           = 0;
-   std::string InRearranger = "box";
+   // Read component-configurable parallel IO settings - default to
+   // single-task if config values do not exist
+   int NumIOTasks = 1;
+   int IOStride   = 1;
 
    Err = IOConfig.get("IOTasks", NumIOTasks);
    CHECK_ERROR_WARN(Err, "IO: NumIOTasks not found in Config - using {}",
@@ -162,22 +221,26 @@ void init(const MPI_Comm &InComm // [in] MPI communicator to use
    CHECK_ERROR_WARN(Err, "IO: IOStride not found in Config - using {}",
                     IOStride);
 
-   Err = IOConfig.get("IOBaseTask", IOBaseTask);
-   CHECK_ERROR_WARN(Err, "IO: IOBaseTask not found in Config - using {}",
-                    IOBaseTask);
-
-   Err = IOConfig.get("IORearranger", InRearranger);
-   CHECK_ERROR_WARN(Err, "IO: Rearranger not found in Config - using {}",
-                    InRearranger);
-   Rearranger Rearrange = RearrFromString(InRearranger);
-
-   // Call PIO routine to initialize
-   DefaultRearr = Rearrange;
-   int PIOErr   = PIOc_Init_Intracomm(InComm, NumIOTasks, IOStride, IOBaseTask,
-                                      Rearrange, &SysID);
+   // Base task and rearranger are supplied by the caller (driver-owned in a
+   // coupled run) rather than read from the component config.
+   DefaultRearr = IOParams.IORearranger;
+   int PIOErr =
+       PIOc_Init_Intracomm(InComm, NumIOTasks, IOStride, IOParams.IOBaseTask,
+                           IOParams.IORearranger, &SysID);
    if (PIOErr != 0)
       ABORT_ERROR("IO::init: Error initializing SCORPIO");
 
+   // Report the settings SCORPIO was actually initialized with. IOBaseTask and
+   // IORearranger are the interesting ones: in a coupled run they come from the
+   // driver (shr_pio_getioroot/shr_pio_getrearranger) rather than from Omega's
+   // config, and without this message there is no way to observe which values
+   // were used, since they are passed straight into PIOc_Init_Intracomm.
+   LOG_INFO(
+       "IO::init: IOTasks={} IOStride={} IOBaseTask={} IORearranger={} ({})",
+       NumIOTasks, IOStride, IOParams.IOBaseTask,
+       RearrToString(IOParams.IORearranger),
+       static_cast<int>(IOParams.IORearranger));
+
    return;
 
 } // end init

components/omega/src/base/IO.h

@@ -163,13 +163,32 @@ IfExists IfExistsFromString(
     const std::string &IfExists ///< [in] choice of behavior on file existence
 );
 
+/// Parameters for initializing the IO subsystem that are owned by the
+/// driver/coupler rather than by the Omega component config. In a coupled
+/// run the coupler (via CIME/shr_pio) chooses the base IO task and the
+/// rearranger so that Omega's parallel IO layout is consistent with the
+/// rest of the E3SM case.
+struct IOInitParams {
+   int IOBaseTask;          ///< base (root) MPI task for IO
+   Rearranger IORearranger; ///< parallel IO rearranger algorithm
+};
+
 // Methods
 
 /// Initializes the IO system based on configuration inputs and
 /// default MPI communicator
 void init(const MPI_Comm &InComm ///< [in] MPI communicator to use
 );
 
+/// Initializes the IO system using driver-owned IO parameters for the base
+/// task and rearranger. The number of IO tasks and the IO stride are still
+/// read from the Omega config; only the base task and rearranger are taken
+/// from the supplied parameters. Used in coupled runs where the coupler owns
+/// these settings.
+void init(const MPI_Comm &InComm,       ///< [in] MPI communicator to use
+          const IOInitParams &IOParams  ///< [in] driver-owned IO parameters
+);
+
 /// This routine opens a file for reading. The filename with full path must be
 /// supplied and a FileID is returned to be used by other IO functions.
 /// The format of the file is assumed to be the default defined on init

components/omega/src/drivers/coupled/ocn_comp_mct.F90

@@ -69,6 +69,7 @@ subroutine ocn_init_mct(EClock, cdata, x2o, o2x, NLFilename)
       use shr_sys_mod, only: shr_sys_flush
       use shr_cal_mod, only: shr_cal_noleap, shr_cal_gregorian
       use shr_file_mod, only: shr_file_getunit, shr_file_setIO
+      use shr_pio_mod, only: shr_pio_getioroot, shr_pio_getrearranger
 
       ! !INPUT/OUTPUT PARAMETERS:
       type(ESMF_Clock), intent(inout) :: EClock
@@ -97,6 +98,8 @@ subroutine ocn_init_mct(EClock, cdata, x2o, o2x, NLFilename)
          coupling_time_step, case_start_tod, case_start_ymd, cur_tod, cur_ymd
       integer(kind=c_int) :: start_type_c
       integer(kind=c_int) :: layout
+      integer(kind=c_int) :: io_base_task  ! driver-owned base (root) IO task
+      integer(kind=c_int) :: io_rearranger ! driver-owned PIO rearranger
       character(kind=c_char, len=CL), target :: calendar_c
       character(kind=c_char, len=CL), target :: ocn_log_fname_c
 
@@ -186,6 +189,12 @@ subroutine ocn_init_mct(EClock, cdata, x2o, o2x, NLFilename)
       ! populate the import/export field name and index arrays
       call omega_set_cpl_indices()
 
+      ! The base IO task and rearranger are owned by the driver/coupler (set
+      ! by CIME via shr_pio) so they stay consistent with the rest of the
+      ! case; they are passed to Omega rather than read from omega.yml.
+      io_base_task = shr_pio_getioroot(OCN_ID)
+      io_rearranger = shr_pio_getrearranger(OCN_ID)
+
 #ifdef HAVE_MOAB
       layout = omega_get_layout_moab()
 #else
@@ -209,7 +218,9 @@ subroutine ocn_init_mct(EClock, cdata, x2o, o2x, NLFilename)
          c_loc(import_field_names), &
          c_loc(export_field_names), &
          c_loc(import_field_indices), &
-         c_loc(export_field_indices) &
+         c_loc(export_field_indices), &
+         io_base_task, &
+         io_rearranger &
          )
 
       !-------------------------------------------------------------------------

components/omega/src/drivers/coupled/omega_cxx2f_interface.cpp

@@ -3,6 +3,7 @@
 //
 //===----------------------------------------------------------------------===//
 #include "DataTypes.h"
+#include "IO.h"
 #include "Logging.h"
 #include "MachEnv.h"
 #include "OceanDriver.h"
@@ -60,7 +61,9 @@ void omega_ocn_init1(
     const char *ImportFieldNames,  // [in] array of import field names
     const char *ExportFieldNames,  // [in] array of export field names
     const int *ImportFieldIndices, // [in] array of import field indices
-    const int *ExportFieldIndices  // [in] array of export field indices
+    const int *ExportFieldIndices, // [in] array of export field indices
+    const int IOBaseTask,          // [in] driver-owned base (root) IO task
+    const int IORearranger         // [in] driver-owned PIO rearranger (int)
 ) {
 
    // Create the C MPI_Comm from the Fortran one
@@ -103,8 +106,14 @@ void omega_ocn_init1(
        NCouplerImports, NCouplerExports,  ImportIdxMap,
        ExportIdxMap,    CouplingInterval, OMEGA::CouplingLayout::MCT};
 
+   // The base IO task and rearranger are owned by the driver/coupler (via
+   // CIME/shr_pio). The rearranger int uses the same PIO_REARR_* values as
+   // Omega's IO::Rearranger enum (box = 1, subset = 2).
+   OMEGA::IO::IOInitParams IOParams{
+       IOBaseTask, static_cast<OMEGA::IO::Rearranger>(IORearranger)};
+
    OMEGA::ocnInit1(Comm, OcnID, YamlConfigFile, OcnLogFile, StartTypeEnum,
-                   TimeParams, CouplingParams);
+                   TimeParams, CouplingParams, IOParams);
 
    Pacer::stop("Init1", 0);
 

components/omega/src/drivers/coupled/omega_f2cxx_interface.F90

@@ -27,7 +27,9 @@ subroutine omega_ocn_init1( &
          import_field_names, &
          export_field_names, &
          import_field_indices, &
-         export_field_indices) bind(c)
+         export_field_indices, &
+         io_base_task, &
+         io_rearranger) bind(c)
 
          use, intrinsic :: iso_c_binding, only: c_int, c_char, c_ptr
 
@@ -43,7 +45,9 @@ subroutine omega_ocn_init1( &
             n_coupler_imports, &
             n_coupler_exports, &
             n_omega_imports, &
-            n_omega_exports
+            n_omega_exports, &
+            io_base_task, &
+            io_rearranger
 
          character(kind=c_char), target, intent(in) :: &
             yaml_config_name, ocn_log_name, calendar_name

components/omega/src/ocn/OceanDriver.h

@@ -12,6 +12,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "Config.h"
+#include "IO.h"
 #include "SfcCoupling.h"
 #include "TimeMgr.h"
 #include "TimeStepper.h"
@@ -40,9 +41,10 @@ int ocnInit1(
     const int OcnId,                         ///< [in] mct comp id for ocean
     const std::string &ConfigFile,           ///< [in] path to yaml config file
     const std::string &LogFile,              ///< [in] path to log file
-    const StartType StartType,               ///< [in] simulation start type
-    const TimeInitParams &TimeParams,        ///< [in] time parameters
-    const CouplingInitParams &CouplingParams ///< [in] coupling parameters
+    const StartType StartType,                ///< [in] simulation start type
+    const TimeInitParams &TimeParams,         ///< [in] time parameters
+    const CouplingInitParams &CouplingParams, ///< [in] coupling parameters
+    const IO::IOInitParams &IOParams          ///< [in] driver-owned IO params
 );
 
 /// Coupled init phase 2: runs once the coupler has allocated its MCT buffers;
@@ -63,9 +65,10 @@ int ocnFinalize(const TimeInstant &CurrTime);
 /// Initialize Omega modules needed to run ocean model
 int initOmegaModules(MPI_Comm Comm);
 
-/// Initialize Omega modules with coupler-provided time parameters
+/// Initialize Omega modules with coupler-provided time and IO parameters
 int initOmegaModules(MPI_Comm Comm, const TimeInitParams &TParams,
-                     const CouplingInitParams &CParams);
+                     const CouplingInitParams &CParams,
+                     const IO::IOInitParams &IOParams);
 
 /// Update Halo/Host arrays with new state, auxiliary state, and tracer fields
 int initUpdateHaloAndHostArrays();