Added support for explicit SSH port.

SvenskaSpel · Sep 23, 2024 · fa389c9 · fa389c9
1 parent 6eb56a8
commit fa389c9
Showing 1 changed file with 21 additions and 10 deletions.
diff --git a/locust_swarm/swarm.py b/locust_swarm/swarm.py
@@ -27,6 +27,9 @@
     import tomli as tomllib
 
 
+DEFAULT_SSH_PORT = 22
+
+
 class LocustTomlConfigParser(configargparse.TomlConfigParser):
     def parse(self, stream):
         try:
@@ -193,6 +196,12 @@ def parse(self, stream):
     default=False,
     help="Upload locust-plugins to load gens (useful if you are developing locust-plugins)",
 )
+parser.add_argument(
+    "--ssh-port",
+    type=int,
+    default=DEFAULT_SSH_PORT,
+    help=f"Port to use with SSH, default ({DEFAULT_SSH_PORT})"
+)
 
 parser.add_argument(
     "--version",
@@ -252,7 +261,7 @@ def check_and_lock_server(server):
     # a server is considered busy if it is either running a locust process or
     # is "locked" by a sleep command with somewhat unique syntax.
     # the regex uses a character class ([.]) to avoid matching with the pgrep command itself
-    check_command = f"ssh -o LogLevel=error {server} \"pgrep -f '^sleep 1 19|[l]ocust --worker' && echo busy || (echo available && sleep 1 19)\""
+    check_command = f"ssh -p {args.ssh_port} -o LogLevel=error {server} \"pgrep -f '^sleep 1 19|[l]ocust --worker' && echo busy || (echo available && sleep 1 19)\""
 
     logging.debug(check_command)
     p = subprocess.Popen(check_command, stdout=subprocess.PIPE, shell=True)
@@ -285,7 +294,7 @@ def cleanup(server_list):
             pass
     psutil.wait_procs(procs, timeout=3)
     check_output_multiple(
-        f"ssh -q {server} 'pkill -9 -u $USER -f \"locust --worker\"' 2>&1 | grep -v 'No such process' || true"
+        f"ssh -p {args.ssh_port} -q {server} 'pkill -9 -u $USER -f \"locust --worker\"' 2>&1 | grep -v 'No such process' || true"
         for server in server_list
     )
     logging.debug("cleanup complete")
@@ -325,8 +334,8 @@ def start_worker_process(server, port):
     upload(server)
 
     if args.selenium:
-        check_output(f"ssh -q {server} 'rm -rf /tmp/.com.google.Chrome.*' || true")
-        selenium_cmd = f"ssh -q {server} 'pkill -f \"^java -jar selenium-server-4.\"; java -jar selenium-server-4.0.0.jar standalone > selenium.log 2>&1' &"
+        check_output(f"ssh -p {args.ssh_port} -q {server} 'rm -rf /tmp/.com.google.Chrome.*' || true")
+        selenium_cmd = f"ssh -p {args.ssh_port} -q {server} 'pkill -f \"^java -jar selenium-server-4.\"; java -jar selenium-server-4.0.0.jar standalone > selenium.log 2>&1' &"
         logging.info(selenium_cmd)
         subprocess.Popen(
             selenium_cmd,
@@ -337,12 +346,12 @@ def start_worker_process(server, port):
         )
         time.sleep(0.2)
         check_output(
-            f"ssh -q {server} 'pgrep -f \"^java -jar selenium-server-4.\"' # check to see selenium actually launched"
+            f"ssh -p {args.ssh_port} -q {server} 'pgrep -f \"^java -jar selenium-server-4.\"' # check to see selenium actually launched"
         )
         time.sleep(1)
 
     if args.playwright:
-        check_output(f"ssh -q {server} 'rm -rf tmp/* && pkill playwright.sh || true'")
+        check_output(f"ssh -p {args.ssh_port} -q {server} 'rm -rf tmp/* && pkill playwright.sh || true'")
 
     if args.remote_master:
         port_forwarding_parameters = []
@@ -375,6 +384,8 @@ def start_worker_process(server, port):
 
     cmd = " ".join([
         "ssh",
+        "-p",
+        str(args.ssh_port),
         "-q",
         *port_forwarding_parameters,
         server,
@@ -438,7 +449,7 @@ def main():
 
     try:
         subprocess.check_output(
-            f"ssh -o LogLevel=error -o BatchMode=yes {loadgen_list[0]} true 2>&1",
+            f"ssh -p {args.ssh_port} -o LogLevel=error -o BatchMode=yes {loadgen_list[0]} true 2>&1",
             shell=True,
             timeout=10,
         )
@@ -447,7 +458,7 @@ def main():
             # add all loadgens to known hosts
             for loadgen in loadgen_list:
                 subprocess.check_output(
-                    f"ssh -o LogLevel=error -o BatchMode=yes -o StrictHostKeyChecking=accept-new {loadgen} true",
+                    f"ssh -p {args.ssh_port} -o LogLevel=error -o BatchMode=yes -o StrictHostKeyChecking=accept-new {loadgen} true",
                     shell=True,
                 )
         else:
@@ -492,10 +503,10 @@ def get_available_servers_and_lock_them():
         env_vars = ["PYTHONUNBUFFERED=1"]
         if args.test_env:
             env_vars.append("LOCUST_TEST_ENV=" + args.test_env)
-        ssh_command = ["ssh", "-q", args.remote_master, "'", *env_vars, "sudo", "-E", "nohup"]
+        ssh_command = ["ssh", "-p", str(args.ssh_port), "-q", args.remote_master, "'", *env_vars, "sudo", "-E", "nohup"]
         bind_only_localhost = []
         ssh_command_end = ["'"]
-        check_output(f"ssh -q {args.remote_master} 'pkill -9 -u $USER locust' || true")
+        check_output(f"ssh -p {args.ssh_port} -q {args.remote_master} 'pkill -9 -u $USER locust' || true")
         upload(args.remote_master)
     else:
         # avoid firewall popups by only binding localhost if running local master (ssh port forwarding):