def test_openssh_config(self):
"""Test reading and using OpenSSH config file"""
config_file = tempfile.NamedTemporaryFile()
_host = "127.0.0.2"
_user = "******"
_listen_socket = make_socket(_host)
_server = start_server(_listen_socket)
_port = _listen_socket.getsockname()[1]
_key = USER_KEY_PATH
content = [
("""Host %s\n""" % (_host, )),
(""" User %s\n""" % (_user, )),
(""" Port %s\n""" % (_port, )),
(""" IdentityFile %s\n""" % (_key, )),
]
config_file.writelines([s.encode('utf-8') for s in content])
config_file.flush()
host, user, port, pkey = utils.read_openssh_config(
_host, config_file=config_file.name)
client = SSHClient(_host, _openssh_config_file=config_file.name)
config_file.close()
self.assertEqual(host, _host)
self.assertEqual(user, _user)
self.assertEqual(port, _port)
self.assertTrue(pkey)
self.assertEqual(client.host, _host)
self.assertEqual(client.user, _user)
self.assertEqual(client.port, _port)
self.assertTrue(client.pkey)
del _server, _listen_socket
def test_openssh_config(self):
"""Test reading and using OpenSSH config file"""
config_file = tempfile.NamedTemporaryFile()
_host = "127.0.0.2"
_user = "******"
_listen_socket = make_socket(_host)
_server = start_server(_listen_socket)
_port = _listen_socket.getsockname()[1]
_key = USER_KEY_PATH
content = [("""Host %s\n""" % (_host,)),
(""" User %s\n""" % (_user,)),
(""" Port %s\n""" % (_port,)),
(""" IdentityFile %s\n""" % (_key,)),
]
config_file.writelines([s.encode('utf-8') for s in content])
config_file.flush()
host, user, port, pkey = utils.read_openssh_config(
_host, config_file=config_file.name)
client = SSHClient(_host, _openssh_config_file=config_file.name)
config_file.close()
self.assertEqual(host, _host)
self.assertEqual(user, _user)
self.assertEqual(port, _port)
self.assertTrue(pkey)
self.assertEqual(client.host, _host)
self.assertEqual(client.user, _user)
self.assertEqual(client.port, _port)
self.assertTrue(client.pkey)
del _server, _listen_socket
def test_openssh_config_missing(self):
self.assertFalse(utils.read_openssh_config('test', config_file=str(uuid4())))
def test_openssh_config_missing(self):
self.assertFalse(utils.read_openssh_config('test', config_file=str(uuid4())))
def init(config_path: str = "~/config.json") -> None:
""" Public-facing API for SSHMPI initialization. """
with open(os.path.expanduser(config_path), "r") as config_file:
config = json.load(config_file)
server_ip = config["server_ip"]
port = config["port"]
hosts = config.get("hostnames", [])
# Per-host config dictionaries.
host_config = {}
# Get hostnames of remote nodes.
if not hosts:
hosts = get_available_hostnames_from_sshconfig()
for hostname in hosts:
_, _, remote_port, _ = read_openssh_config(hostname)
host_config[hostname] = {"port": remote_port}
else:
for hostname in hosts:
host_config[hostname] = {}
cellar.HOSTNAMES = hosts
print("Hosts:", hosts)
# Get private key.
pkey = os.path.expanduser("~/.ssh/id_rsa")
# Start the ssh client.
sshclient = ParallelSSHClient(hosts, host_config=host_config, pkey=pkey)
# Reserve local UDP ports for each remote node.
# TODO: Address possibility that ports are already in-use by another program.
# HARDCODE
i = 50000
head_ip = ""
ports: Mapping[str, int] = {}
for name in hosts:
_, external_ip, external_port = stun.get_ip_info(source_port=i)
head_ip = external_ip
ports[name] = external_port
i += 1
# NOTE: We're not using the rendezvous server anymore.
# Reset the channel map of the rendezvous server.
# reset(server_ip, port)
# Command string format arguments are in ``host_args``.
host_args = [(head_ip, ports[name], name) for name in hosts]
sshclient.run_command(
"meadclient %s %s %s > mead_global.log 2>&1",
host_args=host_args,
shell="bash -ic",
)
# Create and start the head node client (one for each remote node).
head_processes: Dict[str, mp.Process] = {}
for hostname in hosts:
# The ``in_spout`` receives data coming from the remote node.
in_funnel, in_spout = mp.Pipe()
cellar.HEAD_SPOUTS[hostname] = in_spout
# The ``out_queue`` sends data going to the remote node.
out_queue: mp.Queue = mp.Queue()
cellar.HEAD_QUEUES[hostname] = out_queue
# We use the hostname as the channel.
leader = Client(server_ip, port, hostname, in_funnel, out_queue)
p_client = mp.Process(target=leader.main)
p_client.start()
head_processes[hostname] = p_client
# Store references to the head processes and SSH client.
cellar.HEAD_PROCESSES = head_processes
cellar.SSHCLIENT = sshclient