Allow remote components on NAS compute nodes; check for either k or NAS in examples

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):

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

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 <port number>
     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
         )