diff --git a/doc/user_guides/platform.rst b/doc/user_guides/platform.rst index 44750791..94042eb8 100644 --- a/doc/user_guides/platform.rst +++ b/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 diff --git a/examples-proposed/020-simple-ensemble/perlmutter.slurm b/examples-proposed/020-simple-ensemble/perlmutter.slurm index 111337dd..ac49e989 100644 --- a/examples-proposed/020-simple-ensemble/perlmutter.slurm +++ b/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 diff --git a/examples-proposed/021-ensembles-from-CSV/perlmutter.slurm b/examples-proposed/021-ensembles-from-CSV/perlmutter.slurm index 7d08fb33..b4687834 100644 --- a/examples-proposed/021-ensembles-from-CSV/perlmutter.slurm +++ b/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 diff --git a/examples-proposed/022-tasks-and-ensembles/perlmutter.slurm b/examples-proposed/022-tasks-and-ensembles/perlmutter.slurm index 8ad6e0e1..1281565b 100644 --- a/examples-proposed/022-tasks-and-ensembles/perlmutter.slurm +++ b/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 diff --git a/examples-proposed/023-simple-ensemble-with-portal/driver.py b/examples-proposed/023-simple-ensemble-with-portal/driver.py index 632fe1ea..f4fbe825 100644 --- a/examples-proposed/023-simple-ensemble-with-portal/driver.py +++ b/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, ) diff --git a/examples-proposed/023-simple-ensemble-with-portal/perlmutter.slurm b/examples-proposed/023-simple-ensemble-with-portal/perlmutter.slurm index 0b5a9215..eea6f4c6 100644 --- a/examples-proposed/023-simple-ensemble-with-portal/perlmutter.slurm +++ b/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 diff --git a/examples-proposed/024-aggregated-compute-ensemble/perlmutter.slurm b/examples-proposed/024-aggregated-compute-ensemble/perlmutter.slurm index ad8cb157..97ce365e 100644 --- a/examples-proposed/024-aggregated-compute-ensemble/perlmutter.slurm +++ b/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 diff --git a/ipsframework/resource_helper.py b/ipsframework/resource_helper.py index d5311c23..3d06c57b 100644 --- a/ipsframework/resource_helper.py +++ b/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) diff --git a/ipsframework/services.py b/ipsframework/services.py index 0117d6a3..da4b8c2f 100644 --- a/ipsframework/services.py +++ b/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,12 +3233,14 @@ 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 @@ -3227,6 +3248,7 @@ def __init__(self, logger, oversubscribe=False, hwthreads=False): 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') diff --git a/ipsframework/task_manager.py b/ipsframework/task_manager.py index bf972d22..e3c19736 100644 --- a/ipsframework/task_manager.py +++ b/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, diff --git a/tests/new/test_task_manager.py b/tests/new/test_task_manager.py index da729830..de563c96 100644 --- a/tests/new/test_task_manager.py +++ b/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 ')