DVM hardening and logging

doc/user_guides/platform.rst

@@ -393,7 +393,8 @@ The platform configuration file contains platform specific information that the
   USE_ACCURATE_NODES = ON
 
 **HOST**
-        name of the platform.  Used by the portal.
+        name of the platform, which is used by the portal.
+        If left blank, the framework will detect the host name and use that.
 **MPIRUN**
         command to launch parallel applications.  Used by the task
 	manager to launch parallel tasks on compute nodes.  If you

examples-proposed/020-simple-ensemble/perlmutter.slurm

@@ -27,6 +27,6 @@ module load openmpi
 
 # The 2>&1 binds stderr to stdout so that both are captured in the tee log file.  The
 # `tee` command allows you to see the output on the terminal as well as save it.
-ips.py --simulation=ensemble.conf --platform=perlmutter.conf \
+ips.py -d --simulation=ensemble.conf --platform=perlmutter.conf \
     2>&1 | tee ${SLURM_JOBID}_ips.log
 

examples-proposed/021-ensembles-from-CSV/perlmutter.slurm

@@ -27,6 +27,6 @@ module load openmpi
 
 # The 2>&1 binds stderr to stdout so that both are captured in the tee log file.  The
 # `tee` command allows you to see the output on the terminal as well as save it.
-ips.py --simulation=ensemble.conf --platform=perlmutter.conf \
+ips.py -d --simulation=ensemble.conf --platform=perlmutter.conf \
     2>&1 | tee ${SLURM_JOBID}_ips.log
 

examples-proposed/022-tasks-and-ensembles/perlmutter.slurm

@@ -32,6 +32,6 @@ echo "Using MPI_STATS_EXEC: ${MPI_STATS_EXEC}"
 
 # The 2>&1 binds stderr to stdout so that both are captured in the tee log file.  The
 # `tee` command allows you to see the output on the terminal as well as save it.
-ips.py --simulation=ensemble.conf --platform=perlmutter.conf \
+ips.py -d --simulation=ensemble.conf --platform=perlmutter.conf \
     2>&1 | tee ${SLURM_JOBID}_ips.log
 

examples-proposed/023-simple-ensemble-with-portal/driver.py

@@ -71,7 +71,7 @@ def step(self, timestamp=0.0):
             template,
             variables,
             run_dir=Path('.').absolute(),
-            name='my_simple_ensemble',
+            name='INSTANCE_',
             num_nodes=1,
             cores_per_instance=1,
         )

examples-proposed/023-simple-ensemble-with-portal/perlmutter.slurm

@@ -33,6 +33,6 @@ source /global/common/software/atom/ips-portal/credentials/ips-portal-developmen
 
 # The 2>&1 binds stderr to stdout so that both are captured in the tee log file.  The
 # `tee` command allows you to see the output on the terminal as well as save it.
-ips.py --simulation=ensemble.conf --platform=perlmutter.conf \
+ips.py -d --simulation=ensemble.conf --platform=perlmutter.conf \
     2>&1 | tee ${SLURM_JOBID}_ips.log
 

examples-proposed/024-aggregated-compute-ensemble/perlmutter.slurm

@@ -30,6 +30,6 @@ source /global/common/software/atom/ips-portal/credentials/ips-portal-developmen
 
 # The 2>&1 binds stderr to stdout so that both are captured in the tee log file.  The
 # `tee` command allows you to see the output on the terminal as well as save it.
-ips.py --simulation=ensemble.conf --platform=perlmutter.conf \
+ips.py -d --simulation=ensemble.conf --platform=perlmutter.conf \
     2>&1 | tee ${SLURM_JOBID}_ips.log
 

ipsframework/resource_helper.py

@@ -357,7 +357,13 @@ def get_resource_list(services, host, partial_nodes=False):
         else:
             accurate_nodes = True
     elif node_detect_str == 'slurm_env':
-        num_nodes, ppn, mixed_nodes, list_of_nodes = get_slurm_info()
+        try:
+            num_nodes, ppn, mixed_nodes, list_of_nodes = get_slurm_info()
+        except Exception:
+            services.critical(
+                'MISCONFIGURED NODE_DETECTION parameter: you are not running in a Slurm environment'
+            )
+            raise
         accurate_nodes = True
     elif node_detect_str == 'manual':
         num_nodes, ppn, mixed_nodes, list_of_nodes = manual_detection(services)

ipsframework/services.py

@@ -162,6 +162,21 @@ def launch(executable: Any, task_name: str, working_dir: str | os.PathLike, *arg
                 log.error(f'DVM URI file {dvm_uri_file} does not exist')
             else:
                 log.debug(f'Using DVM URI file: {dvm_uri_file}')
+            if worker.dvm_proc.poll() is not None:
+                log.error(
+                    f'DVM (prte) exited with code {worker.dvm_proc.returncode}, see {worker.dvm_log_file}'
+                )
+            # Workers run inside an srun step, which may export its own PMIx
+            # identity/rendezvous variables; drop them so prun can only reach
+            # this worker's DVM.  PMIX_MCA_* are tuning knobs and are kept.
+            for key in [
+                k for k in new_env if k.startswith('PMIX_') and not k.startswith('PMIX_MCA_')
+            ]:
+                if key not in task_env:
+                    del new_env[key]
+            new_env['PMIX_SERVER_URI41'] = worker.dvm_uri
+            # read by TaskManager to pass `prun --dvm-uri file:...`
+            new_env['IPS_DVM_URI_FILE'] = str(dvm_uri_file)
 
         # PMIX_SERVER_URI41 is used by prun to figure out how to talk to the DVM
         # It can be defined in `task_env` or in `os.environ`, so we look in
@@ -2983,9 +2998,13 @@ def create_platform_config_file(prefix, working_dir, cores_per_instance, **kwarg
             platform_config = ConfigObj()
             platform_config.filename = str(platform_config_file_path)
 
-            # Though in a batch submission context this may not have much
-            # meaning.
-            platform_config['HOST'] = socket.gethostname()
+            # Note that in a batch submission context this may not have much
+            # meaning.  In any case, this is intentionally set to the empty
+            # string so that the instance will set to the actual hostname it
+            # is running on; if we set it here it will always be set to the
+            # head node regardless of the actual node on which the instance
+            # runs.
+            platform_config['HOST'] = ''
 
             # Regardless, faithfully duplicate the MPIRUN setting from the
             # top-level platform config, which is what the user has set. Same
@@ -3214,19 +3233,22 @@ def send_ensemble_instance_to_portal(ensemble_name: str, data_path: Path) -> Non
 
 
 class DVMPlugin(WorkerPlugin):
-    def __init__(self, logger, oversubscribe=False, hwthreads=False):
+    def __init__(self, logger, working_dir, oversubscribe=False, hwthreads=False):
         """
         Dask worker plugin to launch and manage an OpenMPI PRTE DVM on each
         worker node.
 
         :param logger: Logger object
+        :param working_dir: Directory for the DVM URI and log files; shared by
+            all workers, so file names are suffixed with host and PID.
         :param oversubscribe: Whether to allow oversubscription of nodes
             when launching the ensemble runs. Default is False.
         :param hwthreads: Whether to use hardware threads
         """
         super().__init__()
 
         self.logger = logger
+        self.working_dir = Path(working_dir)
         self.oversubscribe = oversubscribe
         self.hwthreads = hwthreads
 
@@ -3250,7 +3272,11 @@ def setup(self, worker: Worker):
         worker.logger = self.logger
 
         self.logger.info('Launching DVM')
-        self.worker.dvm_uri_file = f'/tmp/dvm.uri.{os.getpid()}'
+        # The working dir is shared by all workers (possibly across nodes), so
+        # each worker needs a unique suffix.
+        suffix = f'{socket.gethostname()}.{os.getpid()}'
+        self.worker.dvm_uri_file = str(self.working_dir / f'dvm.uri.{suffix}')
+        Path(self.worker.dvm_uri_file).unlink(missing_ok=True)  # never read a stale URI
         # The command to launch the DVM.  Note that we use `prte` instead of
         # `mpirun` or `prun`.  This is because we want to launch the DVM
         # itself, not a specific application.  The DVM will then manage the
@@ -3283,27 +3309,33 @@ def setup(self, worker: Worker):
         # This environment variable is specific to OpenMPI's PRTE
         os.environ['PRTE_MCA_rmaps_default_mapping_policy'] = mapping_policy
 
+        # prte output goes to a file, not a PIPE: an undrained pipe can fill
+        # up and block the DVM, and the file survives for post-mortems.
+        self.worker.dvm_log_file = self.working_dir / f'dvm.log.{suffix}'
+        self.worker.dvm_log = open(self.worker.dvm_log_file, 'w')
         try:
             self.logger.debug(f'Executing command: {command!s}')
             self.worker.dvm_proc = subprocess.Popen(
-                command, stdout=subprocess.PIPE, stderr=subprocess.STDOUT
+                command, stdout=self.worker.dvm_log, stderr=subprocess.STDOUT
             )
         except Exception as e:
             print(f'Exception during setting up DVM: {e}', file=sys.stderr)
             console.print(Traceback.from_exception(type(e), e, e.__traceback__))
+            raise
 
-            # If there was an exception, dump any stdout/stderr we have
-            if hasattr(self.worker.dvm_proc, 'stdout'):
-                print(self.worker.dvm_proc.stdout, file=sys.stderr, flush=True)
-                print(self.worker.dvm_proc.stderr, file=sys.stderr, flush=True)
-
-        # TODO What if there was a subprocess exception?
-
-        ready = self.worker.dvm_proc.stdout.readline()
-        self.logger.info(f'Ready Message : {ready}')
+        # Wait for prte to write its URI; fail loudly if it dies or stalls.
+        uri_path = Path(self.worker.dvm_uri_file)
+        deadline = time.time() + 30
+        while not (uri_path.exists() and uri_path.stat().st_size > 0):
+            if self.worker.dvm_proc.poll() is not None or time.time() > deadline:
+                raise RuntimeError(
+                    f'DVM failed to start (exit code {self.worker.dvm_proc.poll()}), '
+                    f'see {self.worker.dvm_log_file}:\n{self.worker.dvm_log_file.read_text()}'
+                )
+            time.sleep(0.1)
 
         with open(self.worker.dvm_uri_file, 'r') as f:
-            self.worker.dvm_uri = f.readline()
+            self.worker.dvm_uri = f.readline().strip()
             self.logger.debug(f'Read DVM URI: {self.worker.dvm_uri}')
 
         os.environ['PMIX_SERVER_URI41'] = self.worker.dvm_uri
@@ -3329,6 +3361,7 @@ def teardown(self, worker: Worker):
         # kill the process directly.
         self.worker.dvm_proc.terminate()
         self.worker.dvm_proc.kill()
+        self.worker.dvm_log.close()
 
         self.logger.info('DVM shutdown')
 
@@ -3648,7 +3681,7 @@ def _make_worker_args(
                 # TODO better error handling than just re-raising the exception
                 raise
 
-        dask_nodes = 1 if dask_nodes is None else dask_nodes
+        dask_nodes = dask_nodes or 1  # None or 0 (nodes unspecified) -> 1
         self.services.info(
             f'Launching {dask_nodes} Dask workers with {dask_ppw} processes per worker'
         )
@@ -3748,7 +3781,12 @@ def _make_worker_args(
         # the DVM for the workers so that OpenMPI can work properly.
         self.services.debug('Registering DVMPlugin')
         self.dask_client.register_plugin(
-            DVMPlugin(logger=services.logger, oversubscribe=oversubscribe, hwthreads=hwthreads)
+            DVMPlugin(
+                logger=services.logger,
+                working_dir=services.get_working_dir(),
+                oversubscribe=oversubscribe,
+                hwthreads=hwthreads,
+            )
         )
         self.services.debug('Registered DVMPlugin')
 

ipsframework/task_manager.py

@@ -436,6 +436,14 @@ def build_launch_cmd(
                 ppn_flag = '-npernode'
                 host_select = '-H'
                 if mpi_binary == 'prun':
+                    # Point prun at this node's DVM explicitly instead of
+                    # letting it discover one; set by services.launch().
+                    dvm_uri_file = os.environ.get('IPS_DVM_URI_FILE')
+                    if not dvm_uri_file:
+                        raise RuntimeError(
+                            'MPIRUN_VERSION is OPENMPI-DVM but IPS_DVM_URI_FILE is not set; '
+                            'this instance was not launched from a Dask worker with a DVMPlugin'
+                        )
                     # --display MAP-DEVEL is added to show the DVM state when
                     # invoking this prun. We do this so that we can verify the
                     # resources managed by DVM for this task as displayed in
@@ -444,6 +452,8 @@ def build_launch_cmd(
                     cmd = ' '.join(
                         [
                             mpicmd,
+                            '--dvm-uri',
+                            f'file:{dvm_uri_file}',
                             '--display',
                             'ALLOCATION,MAP-DEVEL,BINDINGS',
                             nproc_flag,

tests/new/test_task_manager.py

@@ -1027,3 +1027,32 @@ def init_final_task_pool(nproc=1, tppn=0, number_of_tasks=1, tcpp=0, msg=None):
     }
     with pytest.raises(ResourceRequestMismatchError):
         init_final_task_pool(msg=msg)
+
+
+def test_build_launch_cmd_prun_uses_explicit_dvm_uri(monkeypatch):
+    tm = TaskManager(mock.Mock())
+    tm.task_launch_cmd = 'mpirun'
+    tm.mpicmd = '/usr/bin/prun'
+    tm.resource_mgr = mock.Mock(nodes=['node1'])
+    tm.config_mgr = mock.Mock(get_platform_parameter=lambda *a, **k: 'OPENMPI-DVM')
+
+    kwargs = {
+        'nproc': 1,
+        'binary': 'executable',
+        'cmd_args': (),
+        'working_dir': None,
+        'ppn': 1,
+        'max_ppn': 1,
+        'nodes': 'dummynode0',
+        'accurate_nodes': False,
+        'partial_nodes': False,
+        'task_id': None,
+    }
+
+    monkeypatch.delenv('IPS_DVM_URI_FILE', raising=False)
+    with pytest.raises(RuntimeError):
+        tm.build_launch_cmd(**kwargs)
+
+    monkeypatch.setenv('IPS_DVM_URI_FILE', '/work/dvm.uri.nid001.42')
+    cmd, _ = tm.build_launch_cmd(**kwargs)
+    assert cmd.startswith('/usr/bin/prun --dvm-uri file:/work/dvm.uri.nid001.42 ')