diff --git a/examples/aerostructural/supersonic_panel/as_opt_remote_parallel.py b/examples/aerostructural/supersonic_panel/as_opt_remote_parallel.py index 6f4208e7..92d16991 100644 --- a/examples/aerostructural/supersonic_panel/as_opt_remote_parallel.py +++ b/examples/aerostructural/supersonic_panel/as_opt_remote_parallel.py @@ -1,3 +1,6 @@ +import os +import socket + import openmdao.api as om from as_opt_parallel import run_check_totals, run_optimization from pbs4py import PBS @@ -14,11 +17,7 @@ def setup(self): # NOTE: make sure setup isn't called multiple times, otherwise the first jobs/port forwarding will go unused and you'll have to stop them manually for i in range(self.options["num_scenarios"]): - pbs_launcher = PBS.k4( - profile_filename="~/.bashrc", requested_number_of_nodes=1, time=1 - ) - pbs_launcher.mpiexec = "mpirun" - pbs_launcher.requested_number_of_nodes = 1 + pbs_launcher, hpc = self._get_pbs_launcher() # output functions of interest, which aren't already added as objective/constraints on server side if i == 0: @@ -43,6 +42,7 @@ def setup(self): pbs=pbs_launcher, port=start_port, acceptable_port_range=[start_port, end_port], + forward_through_frontend=True if hpc == "nas" else False, dump_separate_json=True, additional_remote_inputs=["mach", "qdyn", "aoa"], additional_remote_outputs=additional_remote_outputs, @@ -55,6 +55,41 @@ def setup(self): promotes_outputs=["*"], ) + def _get_pbs_launcher(self): + + # get hostname + if os.environ.get("PBS_O_HOST") is not None: # running from HPC job + host = os.environ.get("PBS_O_HOST") + else: # running from login node + host = socket.gethostname() + + # check if using nas or k + if host.startswith("k4-li"): + hpc = "k" + elif host.startswith("pfe"): + hpc = "nas" + else: + raise ValueError( + f"Unable to determine if running from NAS or K based on hostname '{host}'" + ) + + if hpc == "k": + pbs_launcher = PBS.k4( + profile_filename="~/.bashrc", + requested_number_of_nodes=1, + time=1, + ) + elif hpc == "nas": + pbs_launcher = PBS.nas( + profile_filename="~/.bashrc", + requested_number_of_nodes=1, + time=1, + # group_list=None, # add group list here + proc_type="bro", + ) + + return pbs_launcher, hpc + class TopLevelGroup(om.Group): def setup(self): diff --git a/examples/aerostructural/supersonic_panel/as_opt_remote_serial.py b/examples/aerostructural/supersonic_panel/as_opt_remote_serial.py index 09a147e2..2ce51d3a 100644 --- a/examples/aerostructural/supersonic_panel/as_opt_remote_serial.py +++ b/examples/aerostructural/supersonic_panel/as_opt_remote_serial.py @@ -1,3 +1,6 @@ +import os +import socket + import openmdao.api as om from as_opt_parallel import write_out_optimization_data from pbs4py import PBS @@ -60,7 +63,22 @@ def run_optimization(prob: om.Problem): def main(): check_totals = False - hpc = "k" # nas or k + + # get hostname + if os.environ.get("PBS_O_HOST") is not None: # running from HPC job + host = os.environ.get("PBS_O_HOST") + else: # running from login node + host = socket.gethostname() + + # check if using nas or k + if host.startswith("k4-li"): + hpc = "k" + elif host.startswith("pfe"): + hpc = "nas" + else: + raise ValueError( + f"Unable to determine if running from NAS or K based on hostname '{host}'" + ) if hpc == "nas": @@ -86,6 +104,7 @@ def main(): RemoteZeroMQComp( run_server_filename="mphys_server.py", # default server filename pbs=pbs, + forward_through_frontend=True if hpc == "nas" else False, additional_server_args="--model_filename as_opt_parallel " + "--scenario_name cruise pullup", ), # customizable options for server file diff --git a/mphys/network/zmq_pbs.py b/mphys/network/zmq_pbs.py index c52b72bf..53a114b4 100644 --- a/mphys/network/zmq_pbs.py +++ b/mphys/network/zmq_pbs.py @@ -1,5 +1,6 @@ import argparse import json +import os import socket import subprocess import time @@ -27,6 +28,11 @@ def initialize(self): default=[5081, 6000], desc="port range to look through if 'port' is currently busy", ) + self.options.declare( + "forward_through_frontend", + default=False, + desc="whether to have ssh port forwarding jump through frontend node, in case compute nodes cannot communicate", + ) self.options.declare( "additional_server_args", default="", @@ -67,6 +73,7 @@ def _setup_server_manager(self): component_name=self.name, port=self.options["port"], acceptable_port_range=self.options["acceptable_port_range"], + forward_through_frontend=self.options["forward_through_frontend"], additional_server_args=self.options["additional_server_args"], job_expiration_max_restarts=self.options["job_expiration_max_restarts"], ) @@ -89,6 +96,8 @@ class MPhysZeroMQServerManager(ServerManager): Desired port number for ssh port forwarding acceptable_port_range : list Range of alternative port numbers if specified port is already in use + forward_through_frontend: bool + Setup ssh forwarding to jump through frontend node ($PBS_O_HOST). For cases where compute nodes cannot communicate additional_server_args : str Optional arguments to give server, in addition to --port job_expiration_max_restarts : int @@ -102,6 +111,7 @@ def __init__( component_name: str, port=5081, acceptable_port_range=[5081, 6000], + forward_through_frontend=False, additional_server_args="", job_expiration_max_restarts=None, ): @@ -110,6 +120,7 @@ def __init__( self.component_name = component_name self.port = port self.acceptable_port_range = acceptable_port_range + self.forward_through_frontend = forward_through_frontend self.additional_server_args = additional_server_args self.job_expiration_max_restarts = job_expiration_max_restarts self.queue_time_delay = ( @@ -225,7 +236,11 @@ def _wait_for_job_to_start(self): ) def _setup_ssh(self): - ssh_command = f"ssh -4 -o ServerAliveCountMax=40 -o ServerAliveInterval=15 -N -L {self.port}:localhost:{self.port} {self.job.hostname} &" + front_end_host = os.environ.get("PBS_O_HOST") + if front_end_host is not None and self.forward_through_frontend: + ssh_command = f"ssh -4 -o ServerAliveCountMax=40 -o ServerAliveInterval=15 -N -L {self.port}:localhost:{self.port} -J {front_end_host} {self.job.hostname} &" + else: + ssh_command = f"ssh -4 -o ServerAliveCountMax=40 -o ServerAliveInterval=15 -N -L {self.port}:localhost:{self.port} {self.job.hostname} &" self.ssh_proc = subprocess.Popen( ssh_command.split(), stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL )