Add debug scripts for ESMF/MPI concurrency error on E3SM Unified 1.11.0rc6 (#934)

Author: tomvothecoder

.vscode/e3sm_diags.code-workspace

@@ -64,6 +64,22 @@
         "env": {
           "PYTHONPATH": "${workspaceFolder}"
         }
+      },
+      {
+        "name": "Python: Remote Attach",
+        "type": "debugpy",
+        "request": "attach",
+        "justMyCode": false,
+        "connect": {
+          "host": "localhost",
+          "port": 5678
+        },
+        "pathMappings": [
+          {
+            "localRoot": "${workspaceFolder}",
+            "remoteRoot": "/path/to/your/code/on/remote"
+          }
+        ]
       }
     ]
   }

auxiliary_tools/cdat_regression_testing/933-esmf-mpi/esmpy_manager.py

@@ -0,0 +1,21 @@
+# srun -N 1 -t 01:00:00 --pty bash
+# source /lcrc/soft/climate/e3sm-unified/test_e3sm_unified_1.11.0rc7_chrysalis.sh
+# ssh -4 -L 5678:localhost:5678 user@compute-node
+# Run and Debug View -> Python: Remote Attach -> Attach to Remote
+#
+# Compute env (Conda): source /lcrc/soft/climate/e3sm-unified/base/etc/profile.d/conda.sh && conda activate e3sm_unified_1.11.0rc7_chrysalis
+# Compute env (Spack): source /lcrc/soft/climate/e3sm-unified/test_e3sm_unified_1.11.0rc7_chrysalis.sh
+
+import os
+
+# ESMF_MPIRUN: Disables the use of mpirun by setting it to "no".
+os.environ["ESMF_MPIRUN"] = "no"
+# Set the communication method to mpiuni to disable MPI
+os.environ["ESMF_COMM"] = "mpiuni"
+import esmpy as ESMF
+
+# Initialize ESMF Manager with MPI. -- still crashes
+ESMF.Manager(debug=True)
+
+# Finalize ESMF with MPI. -- still crashes
+ESMF.ESMP_Finalize()

auxiliary_tools/cdat_regression_testing/933-esmf-mpi/mvce.py

@@ -0,0 +1,41 @@
+# srun -N 1 -t 01:00:00 --pty bash
+# source /lcrc/soft/climate/e3sm-unified/test_e3sm_unified_1.11.0rc7_chrysalis.sh
+# ssh -4 -L 5678:localhost:5678 user@compute-node
+# Run and Debug View -> Python: Remote Attach -> Attach to Remote
+#
+# Compute env (Conda): source /lcrc/soft/climate/e3sm-unified/base/etc/profile.d/conda.sh && conda activate e3sm_unified_1.11.0rc7_chrysalis
+# Compute env (Spack): source /lcrc/soft/climate/e3sm-unified/test_e3sm_unified_1.11.0rc7_chrysalis.sh
+
+
+import dask
+import esmpy as ESMF
+import numpy as np
+from dask import bag
+from xesmf.backend import Grid
+
+# Initialize ESMF with MPI only once
+ESMF.Manager(debug=True)
+
+def xesmf_code(_):
+    Grid(
+        np.array((576, 361)),
+        staggerloc=ESMF.StaggerLoc.CENTER,
+        coord_sys=ESMF.CoordSys.SPH_DEG,
+        num_peri_dims=None
+    )
+
+config = {"scheduler": "processes", "multiprocessing.context": "fork"}
+
+with dask.config.set(config):
+    print(f"Running with 2 workers and 'fork'")
+    print("*"*50)
+    try:
+        b = bag.from_sequence(range(2))
+        results = b.map(xesmf_code).compute(num_workers=2)
+    except Exception as e:
+        print(f"{type(e).__name__}: Message: {str(e)}")
+    else:
+        print("Done")
+
+# Finalize ESMF with MPI.
+ESMF.ESMP_Finalize()

auxiliary_tools/cdat_regression_testing/933-esmf-mpi/mvce_dask_dist.py

@@ -0,0 +1,46 @@
+# 1. srun -N 1 -t 01:00:00 --pty bash
+# 2. source /lcrc/soft/climate/e3sm-unified/test_e3sm_unified_1.11.0rc6_chrysalis.sh
+# 3. ssh -4 -L 5678:localhost:5678 user@compute-node
+# 4. Run and Debug View -> Python: Remote Attach -> Attach to Remote
+
+import esmpy as ESMF
+import numpy as np
+from dask.distributed import Client
+from xesmf.backend import Grid
+
+ESMF.Manager(debug=True, mpi=False)
+
+
+import debugpy
+debugpy.listen(("0.0.0.0", 5678))
+print("Waiting for debugger attach...")
+debugpy.wait_for_client()
+
+def xesmf_code(_):
+    Grid(
+        np.array((576, 361)),
+        staggerloc=ESMF.StaggerLoc.CENTER,
+        coord_sys=ESMF.CoordSys.SPH_DEG,
+        num_peri_dims=None
+    )
+
+# Open MPI has detected that this process has attempted to initialize
+# MPI (via MPI_INIT or MPI_INIT_THREAD) more than once.  This is
+# # erroneous.
+if __name__ == "__main__":
+    run_workers = [2]  # Example worker counts
+    client = Client(processes=False, n_workers=2)
+
+    for num_workers in run_workers:
+        print(f"\nRunning with {num_workers} workers using Dask distributed")
+        print("*" * 50)
+        try:
+            futures = client.map(xesmf_code, range(num_workers))
+            results = client.gather(futures)
+        except Exception as e:
+            print(f"{type(e).__name__}: Message: {str(e)}")
+        else:
+            print("Done")
+
+
+    client.close()

auxiliary_tools/cdat_regression_testing/933-esmf-mpi/mvce_grid_only.py

@@ -0,0 +1,32 @@
+import dask
+import esmpy as ESMF
+import numpy as np
+from dask import bag
+from xesmf.backend import Grid
+import multiprocessing
+
+# Explicitly initialize ESMF in the main process
+ESMF.Manager(debug=True)
+
+def xesmf_code(_):
+    Grid(
+        np.array((576, 361)),
+        staggerloc=ESMF.StaggerLoc.CENTER,
+        coord_sys=ESMF.CoordSys.SPH_DEG,
+        num_peri_dims=None
+    )
+
+if __name__ == '__main__':
+    multiprocessing.set_start_method('forkserver', force=True)  # Avoid 'fork'
+
+    config = {"scheduler": "processes", "multiprocessing.context": "forkserver"}
+    with dask.config.set(config):
+        print(f"Running with 2 workers and 'forkserver'")
+        print("*"*50)
+        try:
+            b = bag.from_sequence(range(2))
+            results = b.map(xesmf_code).compute(num_workers=2)
+        except Exception as e:
+            print(f"{type(e).__name__}: Message: {str(e)}")
+        else:
+            print("Done")

auxiliary_tools/cdat_regression_testing/933-esmf-mpi/mvce_mpi.py

@@ -0,0 +1,22 @@
+from mpi4py import MPI
+
+# Initialize MPI
+if not MPI.Is_initialized():
+    print("MPI is not initialized. Initializing.")
+    MPI.Init()
+else:
+    print("MPI is already initialized.")
+
+# Perform your MPI operations here
+comm = MPI.COMM_WORLD
+rank = comm.Get_rank()
+size = comm.Get_size()
+
+print(f"MPI process with rank {rank} out of {size} is running.")
+
+# Finalize MPI
+if not MPI.Is_finalized():
+    print("MPI is not finalized. Finalizing.")
+    MPI.Finalize()
+else:
+    print("MPI is already finalized.")

auxiliary_tools/cdat_regression_testing/933-esmf-mpi/mvce_mpi_import.py

@@ -0,0 +1,14 @@
+"""
+# Need to do this before importing MPI.
+import mpi4py
+mpi4py.rc.initialize = False  # do not initialize MPI automatically
+mpi4py.rc.finalize = False    # do not finalize MPI automatically
+"""
+
+# import mpi4py
+# mpi4py.rc.initialize = False  # do not initialize MPI automatically
+# mpi4py.rc.finalize = False    # do not finalize MPI automatically
+
+# PMI2_Init failed to intialize.  Return code: 14
+from mpi4py import MPI
+

auxiliary_tools/cdat_regression_testing/933-esmf-mpi/run_script.py

@@ -0,0 +1,54 @@
+# srun -N 1 -t 01:00:00 --pty bash
+# source /lcrc/soft/climate/e3sm-unified/test_e3sm_unified_1.11.0rc8_chrysalis.sh
+# ssh -4 -L 5678:localhost:5678 user@compute-node
+# Run and Debug View -> Python: Remote Attach -> Attach to Remote
+#
+# Compute env (Conda): source /lcrc/soft/climate/e3sm-unified/base/etc/profile.d/conda.sh && conda activate e3sm_unified_1.11.0rc8_chrysalis
+# Compute env (Spack): source /lcrc/soft/climate/e3sm-unified/test_e3sm_unified_1.11.0rc8_chrysalis.sh
+
+import os
+import sys
+
+import numpy
+
+from e3sm_diags.parameter.core_parameter import CoreParameter
+from e3sm_diags.run import runner
+
+# import debugpy
+# debugpy.listen(("0.0.0.0", 5678))
+# print("Waiting for debugger attach...")
+# debugpy.wait_for_client()
+
+short_name = 'v2.LR.historical_0201'
+test_ts = '/lcrc/group/e3sm/ac.forsyth2/zppy_weekly_comprehensive_v2_output/test_pr651_both_commits_20250117/v2.LR.historical_0201/post/atm/180x360_aave/ts/monthly/2yr'
+
+param = CoreParameter()
+
+# Model
+param.test_data_path = '/lcrc/group/e3sm/ac.forsyth2/zppy_weekly_comprehensive_v2_output/test_pr651_both_commits_20250117/v2.LR.historical_0201/post/atm/180x360_aave/clim/2yr'
+param.test_name = 'v2.LR.historical_0201'
+param.short_test_name = short_name
+
+# Ref
+
+# Obs
+param.reference_data_path = '/lcrc/group/e3sm/diagnostics/observations/Atm/climatology/'
+
+
+# Output dir
+param.results_dir = 'model_vs_obs_1982-1983'
+
+# Additional settings
+param.run_type = 'model_vs_obs'
+param.diff_title = 'Model - Observations'
+param.multiprocessing = True
+param.num_workers = 8
+#param.fail_on_incomplete = True
+params = [param]
+
+# Run
+cfg_path = "auxiliary_tools/cdat_regression_testing/933-esmf-mpi/v2_run.cfg"
+sys.argv.extend(["--diags", cfg_path])
+
+runner.sets_to_run = ['lat_lon', ]
+runner.run_diags(params)

auxiliary_tools/cdat_regression_testing/933-esmf-mpi/v2_run.cfg

@@ -0,0 +1,11 @@
+[#]
+sets = ["lat_lon"]
+case_id = "MERRA2"
+variables = ["TREFMXAV"]
+regions = ["global"]
+ref_name = "MERRA2"
+reference_name = "MERRA2 Reanalysis"
+regions = ["global"]
+seasons=["ANN"]
+contour_levels = [-35, -30, -25, -20, -15, -10, -5, 0, 5, 10, 15, 20, 25, 30, 35, 40]
+diff_levels = [-12, -8, -4, -2, -1, -0.5, 0.5, 1, 2, 4, 8, 12]