Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Added support for explicit SSH port. #23

Merged
merged 1 commit into from
Sep 24, 2024
Merged

Added support for explicit SSH port. #23

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 19 additions & 10 deletions locust_swarm/swarm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -193,6 +193,7 @@ 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, help="Port to use with SSH.")

parser.add_argument(
"--version",
Expand All @@ -204,6 +205,12 @@ def parse(self, stream):

args, unrecognized_args = parser.parse_known_args()
master_proc = None
ssh_port_args = []
if args.ssh_port:
ssh_port_args = [
"-p",
str(args.ssh_port),
]


def is_port_in_use(portno: int):
Expand Down Expand Up @@ -252,7 +259,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 {' '.join(ssh_port_args)} -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)
Expand Down Expand Up @@ -285,7 +292,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 {' '.join(ssh_port_args)} -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")
Expand Down Expand Up @@ -325,8 +332,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 {' '.join(ssh_port_args)} -q {server} 'rm -rf /tmp/.com.google.Chrome.*' || true")
selenium_cmd = f"ssh {' '.join(ssh_port_args)} -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,
Expand All @@ -337,12 +344,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 {' '.join(ssh_port_args)} -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 {' '.join(ssh_port_args)} -q {server} 'rm -rf tmp/* && pkill playwright.sh || true'")

if args.remote_master:
port_forwarding_parameters = []
Expand Down Expand Up @@ -375,6 +382,7 @@ def start_worker_process(server, port):

cmd = " ".join([
"ssh",
*ssh_port_args,
"-q",
*port_forwarding_parameters,
server,
Expand Down Expand Up @@ -437,8 +445,9 @@ def main():
worker_process_count = args.processes * args.loadgens

try:
print(f"ssh {' '.join(ssh_port_args)} -o LogLevel=error -o BatchMode=yes {loadgen_list[0]} true 2>&1")
subprocess.check_output(
f"ssh -o LogLevel=error -o BatchMode=yes {loadgen_list[0]} true 2>&1",
f"ssh {' '.join(ssh_port_args)} -o LogLevel=error -o BatchMode=yes {loadgen_list[0]} true 2>&1",
shell=True,
timeout=10,
)
Expand All @@ -447,7 +456,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 {' '.join(ssh_port_args)} -o LogLevel=error -o BatchMode=yes -o StrictHostKeyChecking=accept-new {loadgen} true",
shell=True,
)
else:
Expand Down Expand Up @@ -492,10 +501,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", *ssh_port_args, "-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 {' '.join(ssh_port_args)} -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):
Expand Down
Loading