def kill_remote_ipcluster(hostname="hpc05",
username=None,
password=None,
env_path=None):
"""Kill your remote ipcluster and cleanup the files.
This should do the same as the following bash function (recommended:
add this in your `.bash_profile` / `.bashrc`):
```bash
del() {
qselect -u $USER | xargs qdel
rm -f *.hpc05.hpc* ipengine* ipcontroller* pbs_*
pkill -f hpc05_culler 2> /dev/null
pkill -f ipcluster 2> /dev/null
pkill -f ipengine 2> /dev/null
pkill -f ipyparallel.controller 2> /dev/null
pkill -f ipyparallel.engines 2> /dev/null
}
```
"""
if env_path is None:
env_path = ""
python_exec = "python"
else:
python_exec = os.path.join(env_path, "bin", "python")
with setup_ssh(hostname, username, password) as ssh:
cmd = f"import hpc05; hpc05.connect.kill_ipcluster()"
cmd = f'{python_exec} -c "{cmd}"'
stdin, stdout, stderr = ssh.exec_command(cmd, get_pty=True)
with suppress(Exception):
lines = stdout.readlines()
for line in lines:
print(line)
def get_remote_env(env=None):
# XXX: improve this function and pass all argmuments!
with setup_ssh() as ssh:
cmd = "conda list --export"
if env:
cmd += f" -n {env}"
stdin, stdout, stderr = ssh.exec_command(cmd, get_pty=True)
remote_env = [
l[:-1].rstrip("\r") for l in stdout.readlines()
if not l.startswith("# ")
]
return remote_env
def start_remote_ipcluster(
n,
profile="pbs",
hostname="hpc05",
username=None,
password=None,
env_path=None,
timeout=300,
):
"""Starts an `ipcluster` over ssh on `hostname` and wait untill it's
successfully started.
Parameters
----------
n : int
Number of engines to be started.
profile : str, default 'pbs'
Profile name of IPython profile.
hostname : str
Hostname of machine where the ipcluster runs.
username : str
Username to log into `hostname`. If not provided, it tries to look it up in
your `.ssh/config`.
password : str
Password for `ssh username@hostname`.
env_path : str, default=None
Path of the Python environment, '/path/to/ENV/' if Python is in `/path/to/ENV/bin/python`.
Examples '~/miniconda3/envs/dev/', 'miniconda3/envs/dev', '~/miniconda3'.
Defaults to the environment that is sourced in `.bashrc` or `.bash_profile`.
timeout : int
Time for which we try to connect to get all the engines.
Returns
-------
None
"""
if env_path is None:
env_path = ""
python_exec = "python"
else:
python_exec = os.path.join(env_path, "bin", "python")
with setup_ssh(hostname, username, password) as ssh:
cmd = f"import hpc05; hpc05.start_ipcluster({n}, '{profile}', '{env_path}', {timeout})"
cmd = f'{python_exec} -c "{cmd}"'
stdin, stdout, stderr = ssh.exec_command(cmd, get_pty=True)
wait_for_succesful_start(stdout, timeout=timeout)
def _create_remote_profile(
hostname="hpc05",
username=None,
password=None,
profile="pbs",
local_controller=False,
custom_template=None,
batch_type="pbs",
):
assert batch_type in ("pbs", "slurm")
if custom_template is not None:
raise NotImplementedError(
f"Use `create_local_{batch_type}_profile` on the"
" cluster locally or implement this function.")
with setup_ssh(hostname, username, password) as ssh:
cmd = f'import hpc05; hpc05.create_local_{batch_type}_profile("{profile}", {local_controller})'
cmd = f"python -c '{cmd}'"
stdin, stdout, stderr = ssh.exec_command(cmd, get_pty=True)
out, err = stdout.readlines(), stderr.readlines()
for lines in [out, err]:
for line in lines:
print(line.rstrip("\n"))
def __init__(
self,
profile="pbs",
hostname="hpc05",
username=None,
password=None,
culler=True,
culler_args=None,
env_path=None,
local=False,
*args,
**kwargs,
):
culler_cmd = get_culler_cmd(profile, env_path, culler_args=culler_args)
if local or on_hostname(hostname):
# Don't connect over ssh if this is run on hostname.
if culler:
subprocess.Popen(
culler_cmd,
shell=True,
stdout=open("/dev/null", "w"),
stderr=open("logfile.log", "a"),
)
super().__init__(profile=profile, *args, **kwargs)
else:
import pexpect
json_file, self.json_filename = tempfile.mkstemp(
) # Create temporary file
os.close(json_file)
# Try to get the json 10 times.
remote_json = (
fr".ipython/profile_{profile}/security/ipcontroller-client.json"
)
print_same_line(f"Trying to copy over {remote_json}.")
for i in range(10):
print_same_line(
f"Trying to copy over {remote_json}. Attempt: {i}/10.")
with setup_ssh(hostname, username, password) as ssh:
with suppress(FileNotFoundError), ssh.open_sftp() as sftp:
sftp.chdir(os.path.dirname(remote_json))
sftp.get(os.path.basename(remote_json),
self.json_filename)
break
if i == 9:
raise FileNotFoundError(
f'Could not copy the json file: "{remote_json}" of the pbs '
"cluster. This could have several reasons, most likely it is "
"because the `ipcluster` probably is not running or "
"you have no `profile_pbs`, create with "
"`hpc05.profile.create_remote_pbs_profile()`.")
time.sleep(1)
print_same_line(f"Copied over {remote_json} in {i+1} attempt.",
new_line_end=True)
# Read the json file
with open(self.json_filename) as json_file:
json_data = json.load(json_file)
keys = ("control", "iopub", "mux", "notification", "registration",
"task")
local_json_data = json_data.copy()
local_json_data["location"] = "localhost"
# Select six random ports for the local machine
local_ports = zmq.ssh.tunnel.select_random_ports(6)
for port, key in zip(local_ports, keys):
local_json_data[key] = port
# Replace remote ports by local ports
with open(self.json_filename, "w") as json_file:
json.dump(local_json_data, json_file)
ips = [
"{}:{}:{} ".format(local_json_data[key], json_data["location"],
json_data[key]) for key in keys
]
ssh_forward_cmd = ("ssh -o ConnectTimeout=10 -N -L " +
"-L ".join(ips) + hostname)
self.tunnel = pexpect.spawn(ssh_forward_cmd)
self.tunnel.expect(
[pexpect.TIMEOUT, pexpect.EOF, "[Pp]assword", "passphrase"],
timeout=6)
if culler:
# Using `with setup_ssh` here results in the culler not being started.
ssh = setup_ssh(hostname, username, password)
ssh.exec_command(culler_cmd, get_pty=True)
super().__init__(self.json_filename, *args, **kwargs)
if not on_hostname(hostname):
def __del__(self):
if self.tunnel:
self.tunnel.close()