wehub-resource-sync
68 lines
2.1 KiB
Python
68 lines
2.1 KiB
Python
import multiprocessing as mp
|
|
import os
|
|
import subprocess
|
|
from typing import Optional
|
|
|
|
|
|
def run(ssh_cmd):
|
|
subprocess.check_call(ssh_cmd, shell=True)
|
|
|
|
|
|
def execute_remote(
|
|
cmd: str, ip: str, port: Optional[int] = 22, username: Optional[str] = ""
|
|
) -> mp.Process:
|
|
"""Execute command line on remote machine via ssh.
|
|
|
|
Args:
|
|
cmd: User-defined command (udf) to execute on the remote host.
|
|
ip: The ip-address of the host to run the command on.
|
|
port: Port number that the host is listening on.
|
|
username: Optional. If given, this will specify a username to use when issuing commands over SSH.
|
|
Useful when your infra requires you to explicitly specify a username to avoid permission issues.
|
|
|
|
Returns:
|
|
Process: The Process whose run() is to run the `cmd` on the remote host. Returns when the cmd completes
|
|
on the remote host.
|
|
"""
|
|
ip_prefix = ""
|
|
if username:
|
|
ip_prefix += "{username}@".format(username=username)
|
|
|
|
custom_port = os.getenv("DIST_DGL_TEST_SSH_PORT", "")
|
|
if custom_port:
|
|
port = custom_port
|
|
|
|
custom_ssh_key = os.getenv("DIST_DGL_TEST_SSH_KEY", "")
|
|
if custom_ssh_key:
|
|
custom_ssh_key = os.path.expanduser(custom_ssh_key)
|
|
custom_ssh_key = "-i " + custom_ssh_key
|
|
|
|
ssh_setup = os.getenv("DIST_DGL_TEST_SSH_SETUP", "")
|
|
if ssh_setup:
|
|
cmd = ssh_setup + ";" + cmd
|
|
# Construct ssh command that executes `cmd` on the remote host
|
|
ssh_cmd = "ssh -o StrictHostKeyChecking=no {ssh_key} -p {port} {ip_prefix}{ip} '{cmd}'".format(
|
|
ssh_key=custom_ssh_key,
|
|
port=str(port),
|
|
ip_prefix=ip_prefix,
|
|
ip=ip,
|
|
cmd=cmd,
|
|
)
|
|
ctx = mp.get_context("spawn")
|
|
proc = ctx.Process(target=run, args=(ssh_cmd,))
|
|
proc.start()
|
|
return proc
|
|
|
|
|
|
def get_ips(ip_config):
|
|
ips = []
|
|
with open(ip_config) as f:
|
|
for line in f:
|
|
result = line.strip().split()
|
|
if len(result) != 1:
|
|
raise RuntimeError(
|
|
"Invalid format of ip_config:{}".format(ip_config)
|
|
)
|
|
ips.append(result[0])
|
|
return ips
|