fix: clear stale PMI metadata before YAC remaps

docs/ucar.md

@@ -190,22 +190,34 @@ synthesized from the prefix, so there is nothing to edit in the script.
 **`doctor` passes but YAC tools fail** — run `yac_smoke_test.py`, which checks
 importability on the compute worker specifically rather than on the login node.
 
-**A conservative remap kills the worker (`WorkerLost`)** — almost certainly
-MPI, not memory, and the two look identical from the outside. Importing YAC
-calls `MPI_Init`: Open MPI starts fine as a singleton, MPICH aborts the
-process, and an abort inside a Globus worker arrives as a lost worker with
-nothing mentioning MPI.
-
-Tell them apart with `check_remote_yac(endpoint=...)`. It runs the same
-conservative remap in a *subprocess*, where an abort is survivable:
-
-- `conservative_ok: true` — YAC is fine. The worker needs to start inside an
-  MPI context (`GlobusMPIEngine`, or a launcher prefix on the engine). An
-  export in `worker_init` cannot do this; that runs before the worker exists.
-- `conservative_ok: false` — the build is the problem, not the launcher.
-
-`./scripts/ucar_endpoint.sh check-mpi` runs the bare and launched imports side
-by side and tells you which case you are in.
+**A conservative remap kills the worker (`WorkerLost`)** — MPI, not memory,
+and the two look identical from the outside. Importing YAC calls `MPI_Init`,
+which works out whether this process is one rank of many by reading the
+environment. A worker started inside a batch job inherits `PMI_FD` from that
+job, pointing at a descriptor whose owner has already exited. MPI writes to
+it, gets `Bad file descriptor`, and aborts — which reaches you as a lost
+worker with nothing mentioning MPI.
+
+The server unsets `PMI_FD`, `PMI_RANK` and `PMI_SIZE` before every YAC-backed
+remap, so a current endpoint should not hit this. If you still see it, the
+endpoint is probably running an older `uxarray-mcp`; check the version in
+`check_remote_yac(endpoint=...)` and redeploy. To fix it at the endpoint
+instead, add to `worker_init`:
+
+```bash
+unset PMI_FD PMI_RANK PMI_SIZE
+```
+
+`conservative_ok: false` from `check_remote_yac` means something else — the
+YAC build itself. `./scripts/ucar_endpoint.sh check-mpi` distinguishes the two.
+
+**Do not fix this with an MPI launcher.** `mpiexec -n 1` also makes the import
+work, and earlier versions of this page recommended it. At more than one rank
+it silently corrupts results: YAC splits the work across ranks, nothing here
+gathers it back, and each rank returns its own partition as if it were the
+whole field. Measured against a known-correct serial answer of 73344.0, four
+ranks gave 56635, 58331, 53862 and 60373 on four consecutive runs — full-length
+arrays, no NaNs, no error.
 
 Meanwhile `inverse_distance_weighted` and `bilinear` work — they use uxarray's
 own engine. Neither conserves the field integral, which is usually the reason

scripts/chrysalis_endpoint.sh

@@ -170,6 +170,10 @@ engine:
       export PATH="$YAC_WORKER_PATH"
       export PYTHONPATH="$YAC_PYTHONPATH"
       export LD_LIBRARY_PATH="$YAC_LD_LIBRARY_PATH:\${LD_LIBRARY_PATH:-}"
+      # Stale PMI vars from the batch job make MPI_Init abort when YAC is
+      # imported (the fd PMI_FD names is already closed). The server also
+      # clears these per-remap; doing it here covers older clients too.
+      unset PMI_FD PMI_RANK PMI_SIZE
 
 idle_heartbeats_soft: 10
 idle_heartbeats_hard: 5760
@@ -199,6 +203,10 @@ engine:
       export PATH="$YAC_WORKER_PATH"
       export PYTHONPATH="$YAC_PYTHONPATH"
       export LD_LIBRARY_PATH="$YAC_LD_LIBRARY_PATH:\${LD_LIBRARY_PATH:-}"
+      # Stale PMI vars from the batch job make MPI_Init abort when YAC is
+      # imported (the fd PMI_FD names is already closed). The server also
+      # clears these per-remap; doing it here covers older clients too.
+      unset PMI_FD PMI_RANK PMI_SIZE
 
     launcher:
       type: SrunLauncher

scripts/ucar_endpoint.sh

@@ -23,18 +23,17 @@ WITH_YAC="${WITH_YAC:-0}"
 # HEALPix remap ran. worker_init sources activate-yac.sh from this prefix, so
 # `restart` is what picks the change up.
 YAC_VERSION="${YAC_VERSION:-3.20.2}"
-# Whether this site needs YAC imported under an MPI launcher.
+# Importing YAC calls MPI_Init, which reads the environment to decide whether
+# it is one rank of many. A worker inside a batch allocation inherits PMI_FD
+# from the job; the descriptor it names is already closed, so MPI writes to
+# it, gets EBADF and aborts the worker. `check-mpi` detects this; worker_init
+# below unsets the variables so it cannot happen.
 #
-# `check-mpi` answers that question; acting on the answer is still manual.
-# Open MPI initialises as a singleton and needs nothing. MPICH aborts, and
-# the endpoint then needs its *workers* started inside an MPI context --
-# which means `GlobusMPIEngine`, or a `launcher_prefix` on the engine, not an
-# export in worker_init. Casper was observed failing this way and has not yet
-# been reconfigured, so this is deliberately not automated: a wrong engine
-# block takes the endpoint down for every tool, not just the remap.
-#
-# Tracked rather than guessed at. See the WorkerLost message in
-# remote/agent.py for what a caller is told in the meantime.
+# Not fixed with GlobusMPIEngine or a launcher prefix, which earlier versions
+# of this script recommended. `mpiexec -n 1` works, but at more than one rank
+# YAC splits the work and nothing gathers it back, so each rank returns its
+# own partition as the whole field -- silently, with no NaNs. Unsetting one
+# variable has no such failure mode.
 YAC_PREFIX="$HOME/opt/yac-$YAC_VERSION"
 YAC_ACTIVATE="$YAC_PREFIX/activate-yac.sh"
 TMUX_SESSION="uxarray-endpoint"
@@ -51,7 +50,7 @@ Usage (run on a Casper login node):
   ucar_endpoint.sh start        Load env + start endpoint in tmux
   ucar_endpoint.sh restart      Stop running endpoint, then start fresh
   ucar_endpoint.sh status       Show endpoint list
-  ucar_endpoint.sh check-mpi    Does importing YAC need an MPI launcher?
+  ucar_endpoint.sh check-mpi    Can YAC import without a stale PMI_FD?
 
 Environment overrides:
   ENDPOINT_NAME   Globus Compute endpoint profile name (default: ucar-uxarray-yac)
@@ -122,60 +121,58 @@ EOF
   echo "  wrote $YAC_ACTIVATE (YAC $YAC_VERSION, python $pyver)"
 }
 
-# Does importing YAC survive outside an MPI launcher?
+# Can YAC import here, and is a stale PMI_FD the reason it cannot?
 #
-# Importing YAC calls MPI_Init. Open MPI initialises as a singleton and is
-# fine; MPICH aborts the process instead. Inside a Globus Compute worker that
-# abort takes the worker down, and the caller sees `WorkerLost` with nothing
-# mentioning MPI -- which reads as an out-of-memory kill and sends people to
-# measure their mesh. Observed on Casper: `method="conservative"` killed the
-# worker while the identical remap finished in 0.09s under
-# `check_remote_yac`, which runs it in a subprocess.
+# Importing YAC calls MPI_Init, which reads the environment to decide whether
+# it is one rank of many. A worker inside a batch allocation inherits PMI_FD
+# from the job, but that descriptor's owner has exited, so MPI writes to a
+# closed fd, gets EBADF and aborts. Inside a Globus worker the abort takes
+# the worker down and the caller sees `WorkerLost` with nothing mentioning
+# MPI -- which reads as an out-of-memory kill and sends people to measure
+# their mesh.
 #
-# This answers the question directly, before anyone submits a remap. If the
-# bare import dies and the launched one lives, the endpoint's worker_init
-# needs a launcher.
+# Confirmed on Chrysalis by dropping one variable at a time: PMI_FD alone is
+# the trigger. PMI_RANK, PMI_SIZE, PMI_JOBID, SLURM_* and HYDRA_* are all
+# harmless on their own.
 _check_mpi() {
   _load_modules
   _activate_env
   local probe='import yac.core, sys; print("import ok:", sys.executable)'
 
-  echo "==> Bare import (what a Globus worker does today)"
-  if python -c "$probe" 2>&1 | sed 's/^/    /'; then
-    local bare_ok=1
-  else
-    local bare_ok=0
-    echo "    FAILED -- this is what takes the worker down."
-  fi
+  echo "==> Inherited PMI variables in this shell"
+  local pmi_seen
+  pmi_seen="$(env | grep -E '^PMI' | cut -d= -f1 | tr '\n' ' ')"
+  echo "    ${pmi_seen:-(none)}"
 
   echo
-  echo "==> Under a launcher"
-  local launcher=""
-  for cand in "mpiexec -n 1" "srun --ntasks 1"; do
-    if command -v "${cand%% *}" >/dev/null 2>&1; then launcher="$cand"; break; fi
-  done
-  if [[ -z "$launcher" ]]; then
-    echo "    No mpiexec or srun on PATH; cannot compare."
-    return 1
-  fi
-  echo "    using: $launcher"
-  local launched_ok=1
-  $launcher python -c "$probe" 2>&1 | sed 's/^/    /' || launched_ok=0
+  echo "==> Bare import (what a worker does without the fix)"
+  local bare_ok=1
+  python -c "$probe" 2>&1 | sed 's/^/    /' || bare_ok=0
+  [[ "$bare_ok" == "0" ]] && echo "    FAILED -- this is what takes the worker down."
+
+  echo
+  echo "==> With PMI_FD/PMI_RANK/PMI_SIZE unset (what the server now does)"
+  local unset_ok=1
+  env -u PMI_FD -u PMI_RANK -u PMI_SIZE python -c "$probe" 2>&1 \
+    | sed 's/^/    /' || unset_ok=0
 
   echo
   if [[ "$bare_ok" == "1" ]]; then
     echo "VERDICT: bare import works. MPI is not what is killing the worker;"
     echo "  if conservative remaps still fail, look at memory next."
-  elif [[ "$launched_ok" == "1" ]]; then
-    echo "VERDICT: bare import dies, launched import works -- this is the MPI"
-    echo "  singleton problem."
-    echo "  The workers must start inside an MPI context. In the endpoint's"
-    echo "  user_config_template.yaml.j2, either switch the engine to"
-    echo "  GlobusMPIEngine, or give the current engine a launcher prefix so"
-    echo "  each worker runs under: $launcher"
+  elif [[ "$unset_ok" == "1" ]]; then
+    echo "VERDICT: a stale PMI_FD is the problem, and unsetting it fixes it."
+    echo "  uxarray-mcp >= 2026.10.1 does this itself before every YAC remap,"
+    echo "  so first check the endpoint is not running an older server."
+    echo "  To fix it at the endpoint, add to worker_init:"
+    echo "      unset PMI_FD PMI_RANK PMI_SIZE"
     echo "  Then: $(basename "$0") restart"
+    echo
+    echo "  Do NOT switch the engine to GlobusMPIEngine or add a launcher"
+    echo "  prefix. That also makes the import work, but above one rank each"
+    echo "  rank returns only its own partition of the field, with no error."
   else
-    echo "VERDICT: YAC does not import either way. This is not MPI --"
+    echo "VERDICT: YAC does not import either way. This is not PMI --"
     echo "  check the build at $YAC_PREFIX."
   fi
 }
@@ -205,12 +202,6 @@ _configure() {
   if [[ "$WITH_YAC" == "1" ]]; then
     _ensure_yac_activate
     yac_line=$'\n      source '"$YAC_ACTIVATE"
-    # NOTE: if `check-mpi` says this site needs a launcher, the change goes
-    # in the engine block below rather than here -- worker_init runs *before*
-    # the worker process starts, so exporting something here cannot put that
-    # process inside an MPI context. Left unwired on purpose rather than
-    # shipped as a setting that looks applied and does nothing. See the
-    # YAC_LAUNCHER note at the top of this file.
   elif [[ "$ENDPOINT_NAME" == *yac* ]]; then
     echo "ERROR: '$ENDPOINT_NAME' is named for YAC but WITH_YAC is 0." >&2
     echo "  Rebuild it with YAC:    WITH_YAC=1 $(basename "$0") configure" >&2
@@ -236,6 +227,10 @@ engine:
     init_blocks: 1
     worker_init: |
       unset PYTHONPATH
+      # Stale PMI vars from the batch job make MPI_Init abort when YAC is
+      # imported (the fd PMI_FD names is already closed). The server also
+      # clears these per-remap; doing it here covers older clients too.
+      unset PMI_FD PMI_RANK PMI_SIZE
       source "\$(conda info --base)/etc/profile.d/conda.sh"
       conda activate $CONDA_ENV$yac_line
 

src/uxarray_mcp/remote/agent.py

@@ -136,27 +136,31 @@ def _normalize_remote_error(
         )
         if yac_backed:
             # Memory is the usual cause of a lost worker, but not for this
-            # call. Importing YAC runs ``MPI_Init``; Open MPI initialises
-            # fine as a singleton, MPICH aborts the process instead, and an
-            # abort inside a Globus worker arrives here as WorkerLost with
-            # nothing mentioning MPI. Observed on Casper, where the same
-            # conservative remap that kills the worker completes in 0.09s
-            # under ``check_remote_yac`` -- which runs it in a worker-side
-            # *subprocess*. Same node, same YAC, same method: the difference
-            # is only whether an abort can take the worker with it.
+            # call. Importing YAC runs ``MPI_Init``, which decides whether
+            # this process is one rank of many by reading the environment.
+            # A worker inside a batch allocation inherits ``PMI_FD`` from
+            # the job, pointing at a descriptor that is already closed, so
+            # MPI writes to it, gets EBADF and *aborts* -- which arrives
+            # here as WorkerLost with nothing mentioning MPI.
             #
-            # Sending the caller to look at mesh size here costs them the
-            # afternoon, because subsetting cannot fix it.
+            # The server now unsets PMI_FD/PMI_RANK/PMI_SIZE before every
+            # YAC-backed remap, so reaching this message means something
+            # else: an older server on the worker, or a genuinely broken
+            # YAC build. Sending the caller to look at mesh size would cost
+            # them the afternoon, because subsetting cannot fix either.
             return RuntimeError(
                 f"The worker on {endpoint} died running a YAC-backed remap"
                 f"{where}. This is usually MPI, not memory: importing YAC "
-                "calls MPI_Init, and on an MPICH site that aborts the "
-                "process rather than initialising as a singleton, which "
-                "arrives here as a lost worker. Confirm with "
-                "`check_remote_yac(endpoint=...)`: if its `conservative_ok` "
-                "is true, YAC itself is fine and the endpoint needs its "
-                "workers launched under an MPI launcher (`mpiexec -n 1` or "
-                "`srun --ntasks 1`) in worker_init. Until then, "
+                "calls MPI_Init, and a worker that inherited a stale PMI_FD "
+                "from its batch job aborts instead of initialising as a "
+                "singleton. This server unsets those variables before "
+                "remapping, so the likely cause is an older uxarray-mcp on "
+                "the endpoint -- compare versions with "
+                "`check_remote_yac(endpoint=...)` and redeploy, or unset "
+                "PMI_FD in the endpoint's worker_init. Do not reach for an "
+                "MPI launcher: at more than one rank each rank remaps only "
+                "its own partition and returns it as the whole field, with "
+                "no error and no NaNs. Meanwhile "
                 "method='inverse_distance_weighted' or 'bilinear' run on "
                 "uxarray's own engine and do not need YAC -- but note that "
                 "neither conserves the field integral, which is the reason "

src/uxarray_mcp/remote/compute_functions.py

@@ -3342,6 +3342,10 @@ def _surface(mod):
 """
 
     env = os.environ.copy()
+    # The probe must match the real in-process path; otherwise it can report
+    # conservative_ok while the worker itself still aborts in MPI_Init.
+    for _pmi_var in ("PMI_FD", "PMI_RANK", "PMI_SIZE"):
+        env.pop(_pmi_var, None)
     prefix_applied = None
     if yac_prefix:
         prefix = os.path.abspath(os.path.expanduser(yac_prefix))
@@ -4619,6 +4623,14 @@ def remote_remap_variable(
     """
     import os
 
+    # Clear stale batch-job PMI metadata before YAC calls MPI_Init. The
+    # worker may inherit PMI_FD/PMI_RANK/PMI_SIZE even though PMI_FD is a
+    # closed descriptor; MPICH then aborts instead of choosing singleton mode.
+    # A launcher is not a safe substitute: multiple ranks need explicit
+    # gather logic, which these worker functions do not have.
+    for _pmi_var in ("PMI_FD", "PMI_RANK", "PMI_SIZE"):
+        os.environ.pop(_pmi_var, None)
+
     import numpy as np
     import uxarray as ux
 
@@ -4885,6 +4897,14 @@ def remote_regrid_dataset(
     """
     import os
 
+    # Clear stale batch-job PMI metadata before YAC calls MPI_Init. The
+    # worker may inherit PMI_FD/PMI_RANK/PMI_SIZE even though PMI_FD is a
+    # closed descriptor; MPICH then aborts instead of choosing singleton mode.
+    # A launcher is not a safe substitute: multiple ranks need explicit
+    # gather logic, which these worker functions do not have.
+    for _pmi_var in ("PMI_FD", "PMI_RANK", "PMI_SIZE"):
+        os.environ.pop(_pmi_var, None)
+
     import numpy as np
     import uxarray as ux
 
@@ -5149,6 +5169,14 @@ def remote_remap_to_rectilinear(
     """
     import os
 
+    # Clear stale batch-job PMI metadata before YAC calls MPI_Init. The
+    # worker may inherit PMI_FD/PMI_RANK/PMI_SIZE even though PMI_FD is a
+    # closed descriptor; MPICH then aborts instead of choosing singleton mode.
+    # A launcher is not a safe substitute: multiple ranks need explicit
+    # gather logic, which these worker functions do not have.
+    for _pmi_var in ("PMI_FD", "PMI_RANK", "PMI_SIZE"):
+        os.environ.pop(_pmi_var, None)
+
     import numpy as np
     import uxarray as ux
     import xarray as xr