Remote Execution#
The remote execution layer runs Pressio-based ROM workflows on HPC clusters
from your local environment. Everything goes through the Dispatcher class,
which manages the SSH connection, file transfers, job submission, monitoring,
and result collection on your behalf.
The Dispatcher#
The Dispatcher has both Local and Remote subclasses, allowing users to define models and workflows in such a way that they can be run from anywhere.
The RemoteDispatcher sends work to a remote cluster, while a
LocalDispatcher executes that work on the current host. Both
dispatchers expose the same public interface for run commands, directory
management, and file transfers.
Run Methods#
There are three core run functions on the dispatchers:
run(cmd, run_directory): Run the given shell commandcmdfrom therun_directoryon the execution host.submit_job(cmd, run_directory): Submitcmd(executed fromrun_directory) to SLURM, wait for it, and report the results. Ifcmdis omitted, submits the pre-configured SLURM script instead (see Configuring the dispatcher).call(target, *args, run_directory, **kwargs): Run the Python callable named bytarget(as"module:qualname") on the execution host and return its result.
Note
The target invoked by the call method must be importable on the
execution host. You can point the Dispatcher to Python virtual environments
or custom Python installations during configuration (see
Configuring the dispatcher).
Directory Management#
The dispatcher classes share the following interface for managing files and directories on the execution host.
path_exists(path): Whetherpathexists on the execution host.create_empty_dir(dir_name): Createdir_name, and any missing parents, on the execution host.list_dir(path): Names of the entries inpathon the execution host.remove(path): Delete the file atpathon the execution host.write_text(path, content): Writecontentto a text file on the execution host.np_savetxt(path, arr, fmt)andnp_savez(path, **arrays): Write numpy data to the execution host.
File Transfer#
These methods transfer files from one location to another. For the
RemoteDispatcher, those locations are on different computers. For the
LocalDispatcher, these equate to local cp commands.
put(local_path, remote_path): Copy a local file to the remote host.get(remote_path, local_path): Copy a remote file to the local host.
Configuration#
The last public method allows the user to access the configuration parameters of the dispatcher classes.
get_config(param=None): Return the value of the specified param, or the whole config dict if no argument is given.
Using the Dispatcher#
The Dispatcher can be used in any context to send work to/from
a remote host. In the context of the Pressio ecosystem and rom-tools, it
is commonly used to execute ROM workflows on HPC clusters and gather the
results back on the user’s local machine.
To use the Dispatcher with rom-tools, you must add it into both your
model and your workflow.
Updating your model#
Assume your model (Model) has the two
fundamental methods populate_run_directory() and run_model(). The same
ideas apply to any other methods your model contains.
Step 1: Add a constructor#
Construct your model with the dispatcher as a member variable:
from typing import Optional
from romtools.hpc.dispatchers import RemoteDispatcher
class MyModel:
def __init__(self, dispatcher: RemoteDispatcher):
self.dispatcher = dispatcher
Tip
You may see some sample models where the dispatcher argument is optional, defaulting to
None. In these cases, the model calls resolve_dispatcher(dispatcher) (imported
from romtools.hpc.dispatchers), which falls back to a basic
LocalDispatcher if none is provided.
Step 2: Set up the run directory#
Update populate_run_directory() so that, instead of copying files locally,
you use self.dispatcher.put(local_file, remote_path) to copy the necessary
files to the remote host.
Note
You can pass relative paths to put() and get(). The base directory
on the remote host where all commands are executed is called the
remote_root and is configured at runtime (see
Configuring the dispatcher).
To run shell commands, e.g. to validate input decks or load modules:
cmd = "load my_module && my_input_validator -i input_file.yaml"
result = self.dispatcher.run(cmd)
if not result.ok:
raise RuntimeError(result.stderr)
Note
run() executes the command directly, so it is the right choice for quick
work such as validation. Anything long or parallel belongs in
submit_job(), so that it lands on compute nodes rather than the login
node.
Note
A command that exits non-zero is reported through the returned Result,
not raised. Every dispatcher behaves this way, so a model can move between
them unchanged. Check result.ok or result.exit_code and decide what
a failure means for your model.
Step 3: Define run_model()#
There are two primary ways to run the model through the dispatcher.
SLURM script (RECOMMENDED). Create a SLURM script locally that executes
your model and configure the dispatcher with that script (see
Configuring the dispatcher). Then run_model() can
be as simple as:
def run_model(self, run_directory: str, parameter_sample: dict) -> int:
return self.dispatcher.submit_job().exit_code
This copies your local SLURM script onto the remote host, submits it, and polls it until it completes.
Note
submit_job() returns a Result carrying the job’s exit code and
captured output. The workflows in rom-tools treat a non-zero
run_model() return as a failed sample, so pass the exit code through
rather than returning 0.
Manual commands. Alternatively, have the dispatcher create the SLURM script for you. Define the command you wish to execute and let the dispatcher wrap it in a SLURM script and submit it:
def run_model(self, run_directory: str, parameter_sample: dict) -> int:
cmd = "srun --ntasks=$SLURM_NNODES --ntasks-per-node=1 my_app"
return self.dispatcher.submit_job(cmd, run_directory).exit_code
The dispatcher creates a SLURM script that executes this command, submits it, and polls the job until it completes. Details of the SLURM script, like the user account, number of nodes, etc., are configured as explained in Configuring the dispatcher.
Step 4: Update other methods#
Any other methods that execute commands or read files can be wired with the same core dispatcher functionality.
For a QoiModel that executes a functor, the
call() method may be useful to dispatch Python functions to the execution
host. Note that any Callable objects passed to call must be
importable on the execution host.
Updating your workflow#
After updating your model class, instantiate the dispatcher in your workflow. The safest approach is to scope it with a context manager:
from romtools.hpc.dispatchers import RemoteDispatcher
campaign_dir = "sample_00"
with RemoteDispatcher(campaign_dir) as dispatcher:
model = ExampleModel(dispatcher)
romtools.workflows.run_sampling(
model=model,
dispatcher=dispatcher,
)
Note that the dispatcher is passed to both the model and the sampling workflow.
The campaign_directory you pass to the dispatcher is created both locally
and remotely, as a subdirectory of your local current directory and your remote
remote_root. This mirrored directory structure simplifies sending files
back and forth between hosts, and gives each run its own subdirectory.
Because it names the same directory on both machines, a RemoteDispatcher
requires a relative path. A nested one such as runs/sample_00 is kept
whole; an absolute path, or one climbing out with .., is rejected when the
dispatcher is built.
Supported workflows#
Each of these workflows accepts a dispatcher argument:
run_sampling()run_eki(),run_mf_eki(),mf_eki_with_auto_rom()run_vi(),run_mf_vi(),mf_vi_with_auto_rom()
Every one of them falls back to a LocalDispatcher when you pass nothing, so
existing workflows and models keep running unchanged.
Inverse workflows#
Inverse workflows evaluate the model many times per iteration, under
<work_directory>/iteration_<k>/run_*. Pass the dispatcher to the model and
to the driver as explained above:
with RemoteDispatcher("eki_00") as dispatcher:
model = MyModel(dispatcher)
parameter_samples, qois = romtools.workflows.run_eki(
model=model,
parameter_space=my_parameter_space,
observations=obs,
observations_covariance=obs_cov,
absolute_work_dir="eki_00",
dispatcher=dispatcher,
)
The driver uses the dispatcher to create each run directory and to write its
restart, history, and per-iteration statistics files. Your model still decides
how a single evaluation runs, through the self.dispatcher calls you added
above.
Note
The working directory argument absolute_work_dir must be an absolute path for local runs. When you
pass a RemoteDispatcher, a relative path is also accepted and is
resolved against remote_root.
Note
In the multifidelity workflows (run_mf_eki(), run_mf_vi(), and their
*_with_auto_rom() wrappers), only the high-fidelity (FOM) evaluations
are dispatched. Surrogate (ROM) models are fit and evaluated in-process, so
their run directories always stay on the local machine.
Warning
evaluation_concurrency greater than 1 is not supported with a
RemoteDispatcher. Concurrent evaluation runs each sample in a separate
process, which a remote connection is not set up to share. Use
evaluation_concurrency=1 for remote runs or keep concurrency with a
LocalDispatcher executing on the cluster (bear in mind that
evaluation_concurrency=N then puts N jobs in the queue at once, which
your site’s submission limits may cap).
Warning
Restart files written through a RemoteDispatcher land on the remote host,
but the drivers read restart_file from the local filesystem. To restart a
remote run, retrieve the restart file first, either with the collect
configuration option or with dispatcher.get().
Configuring the dispatcher#
Once the model is updated and the workflow is defined, run the workflow and configure the dispatcher:
python my_workflow.py <config args>
There are three ways to configure:
YAML. Define a YAML file with all configurable params and pass it with
--hpc-config path/to/your/config.yaml(or-c, see below).CLI. Set params on the command line. Every argument is preceded by the
--hpc-prefix to avoid conflicts with nested applications: set theremote_rootby passing--hpc-remote-root /path/to/remote/root.Combination. CLI arguments override YAML parameters, so you can use a YAML file for the bulk of configuration and CLI args to vary settings from run to run:
python my_workflow.py -c path/to/config.yaml --hpc-collect '*.log'
Both dispatchers read configuration the same way, so a LocalDispatcher
takes the same YAML and the same switches as a RemoteDispatcher, although
it won’t use all the parameters.
Tip
Run python -m romtools.hpc, or refer to the SCHEMA in
hpc/configuration.py, to see the full schema of configuration arguments.
You do not need to specify every argument; the Configuration class
defines a default value for each parameter.
Core configuration arguments#
The dispatcher parses your workflow’s own command line, so every switch it
claims carries an --hpc- prefix and it ignores everything else. The
argument named remote_root in the YAML is --hpc-remote-root on the
command line, num_nodes is --hpc-num-nodes, and so on. Your workflow
keeps every switch of its own, -h included.
Short |
Long |
Meaning |
|---|---|---|
|
|
Path to the YAML configuration file |
Warning
If you nest a dispatcher inside a program that already
uses -c for something of its own, the two collide and the dispatcher
reads your program’s value as a YAML path. In this case, use
--hpc-config instead, or pass the YAML to the dispatcher in code:
with RemoteDispatcher(campaign_dir, config="my_hpc_config.yaml") as dispatcher:
...
config also accepts a Configuration object, for a workflow that
builds its settings itself rather than reading them from a file.
ssh — establish the connection with the remote host:
remote: Name of the remote host.user: Username for the connection.port: Port for the connection.
ssh:
remote: my_cluster
user: my_name
port: 22
workflow — file management and setup:
remote_root: Directory on the remote host where commands are executed, absolute or relative to the home directory.collect: Comma-separated list of files, directories, or glob patterns to retrieve from the remote run directory. If omitted, nothing is retrieved.upload: Comma-separated list of files, directories, or glob patterns to place in the run directory before work starts. ARemoteDispatchersends them to the remote host; aLocalDispatchercopies them from the current directory to the run directory. If omitted, nothing is uploaded. Bothcollectanduploadname what lives under the directory they are read from, so an absolute pattern or one containing..is rejected when the dispatcher is built.python_setup: Shell commands that set up the environment before invoking Python, such as loading modules or activating a virtual environment. Used bycall().python_command: Command that invokes the Python with the necessary libraries installed (python3when a remote call needs one and none is given). Used bycall().
A RemoteDispatcher always runs call() under a separate interpreter,
so it uses these whenever they are set. A LocalDispatcher imports the
target into the current process by default; setting either one makes it stage
the call out to a subprocess instead, which is what you want on a cluster node
whose Python comes from a module or a wrapper such as srun.
workflow:
remote_root: my_campaigns
collect: "*.log, passed.txt"
upload: "input.yaml, mesh/"
python_setup: "module load python/3.11"
python_command: python3
slurm — schedule jobs:
script: Path to a local SLURM script, uploaded to the remote host and submitted on calls tosubmit_job().account: Account WCID to charge for the job.job_namenum_nodestasks_per_nodewall_timepartitionpoll_interval: Seconds betweensqueuepolls.timeout: Seconds to keep retrying thesacctquery for a finished job’s exit code before giving up.
All arguments other than script and poll_interval are only used
when the dispatcher creates the SLURM script for you from a command.
If you provide a pre-made script to use, configuration can just be:
slurm:
script: /path/to/script
poll_interval: 10 # optional, defaults to 30
output — toggle output modes:
debug: Enable verbose debug logging.
output:
debug: True
user_defined — define your own parameters:
user_defined:
my_param: 4
Access them in your model class with:
self.my_param = self.dispatcher.get_config("my_param")
Example#
Run the example workflow with:
python romtools/hpc/example/workflow.py --hpc-remote <remote-host> --hpc-user <username> --hpc-account <account/wcid>
See all available arguments with:
python -m romtools.hpc