def add_known_host(address, port, known_hosts="~/.ssh/known_hosts"):
result = execute(["ssh-keygen", "-R", "[%s]:%s" % (address, port)])
result = execute(["ssh-keyscan", "-p", str(port), address])
if result["returncode"] == 0:
fingerprint = result["stdout"]
known_hosts_file = os.path.expanduser(known_hosts)
with open(known_hosts_file, "a") as fd:
fd.write("\n%s" % fingerprint)
print(" => Known hosts fingerprint added for %s:%s\n" % (address, port))
else:
print_warning(" => Could not add know_hosts fingerprint for %s:%s\n" % (address, port))
def install_libs(cluster_id, host, token):
"""Install ipywidgets, databrickslabs_jupyterlab libraries on the driver
Args:
host (str): host from databricks cli config for given profile string
python_path (str): Remote python path to be used for kernel
"""
python_path = get_python_path(cluster_id)
packages = get_local_libs()
deps = {
p["name"]: p["version"]
for p in packages if p["name"] in ["ipywidgets", "ipykernel"]
}
libs = [
"ipywidgets==%s" % deps["ipywidgets"],
"ipykernel==%s" % deps["ipykernel"],
"databrickslabs-jupyterlab==%s" % __version__,
]
print(" => Installing ", " ".join(libs))
ssh = os.environ.get("SSH") or "ssh"
cmd = [ssh, cluster_id, "sudo", python_path + "/pip", "install"] + libs
result = execute(cmd)
if result["returncode"] == 0:
if result["stdout"].startswith("Requirement already satisfied"):
return (0, "Requirement already satisfied")
else:
return (0, "Installed")
return (result["returncode"], result["stdout"] + "\n" + result["stderr"])
def conda_version():
"""Check conda version"""
result = execute(["conda", "--version"])
if result["returncode"] != 0:
print_error(result["stderr"])
sys.exit(1)
return result["stdout"].strip().split(" ")[1]
def get_local_libs():
"""Get installed libraries of the currtent conda environment"""
result = execute(["conda", "list", "--json"])
if result["returncode"] != 0:
print_error(result["stderr"])
return False
try:
return json.loads(result["stdout"])
except Exception as ex: # pylint: disable=broad-except
print_error(ex)
return False
def configure_ssh(profile, cluster_id):
"""Configure SSH for the remote cluster
Args:
profile (str): Databricks CLI profile string
host (str): host from databricks cli config for given profile string
token (str): token from databricks cli config for given profile string
cluster_id (str): cluster ID
"""
sshkey_file = os.path.expanduser("~/.ssh/id_%s" % profile)
if not os.path.exists(sshkey_file):
print("\n => ssh key '%s' does not exist" % sshkey_file)
answer = input(" => Shall it be created (y/n)? (default = n): ")
if answer.lower() == "y":
print(" => Creating ssh key %s" % sshkey_file)
result = execute(
["ssh-keygen", "-b", "2048", "-N", "", "-f", sshkey_file])
if result["returncode"] == 0:
print_ok(" => OK")
else:
print_error(result["stderr"])
bye()
else:
bye()
else:
print_ok("\n => ssh key '%s' already exists" % sshkey_file)
with open(sshkey_file + ".pub", "r") as fd:
sshkey = fd.read().strip()
try:
apiclient = connect(profile)
client = ClusterApi(apiclient)
except Exception as ex: # pylint: disable=broad-except
print_error(ex)
return None
try:
response = client.get_cluster(cluster_id)
except Exception as ex: # pylint: disable=broad-except
print_error(ex)
return None
ssh_public_keys = response.get("ssh_public_keys", [])
if sshkey in [key.strip() for key in ssh_public_keys]:
print_ok(" => public ssh key already configured for cluster %s" %
cluster_id)
bye()
request = {}
for key, val in response.items():
if key not in [
"driver",
"executors",
"spark_context_id",
"state",
"state_message",
"start_time",
"terminated_time",
"last_state_loss_time",
"last_activity_time",
"disk_spec",
]: # omit runtime attributes
request[key] = val
request["ssh_public_keys"] = ssh_public_keys + [sshkey]
print_warning(" => The ssh key will be added to the cluster." +
"\n Note: The cluster will be restarted immediately!")
answer = input(
" => Shall the ssh key be added and the cluster be restarted (y/n)? (default = n): "
)
if answer.lower() == "y":
try:
response = client.edit_cluster(request)
except DatabricksApiException as ex:
print_error(str(ex))
return None
print_ok(" => OK")
else:
print_error(" => Cancelled")
def configure_ssh(profile, cluster_id):
"""Configure SSH for the remote cluster
Args:
profile (str): Databricks CLI profile string
host (str): host from databricks cli config for given profile string
token (str): token from databricks cli config for given profile string
cluster_id (str): cluster ID
"""
sshkey_file = os.path.expanduser("~/.ssh/id_%s" % profile)
if not os.path.exists(sshkey_file):
print("\n => ssh key '%s' does not exist" % sshkey_file)
answer = input(" => Shall it be created (y/n)? (default = n): ")
if answer.lower() == "y":
print(" => Creating ssh key %s" % sshkey_file)
result = execute(
["ssh-keygen", "-b", "2048", "-N", "", "-f", sshkey_file])
if result["returncode"] == 0:
print_ok(" => OK")
else:
print_error(result["stderr"])
bye()
else:
bye()
else:
print_ok("\n => ssh key '%s' already exists" % sshkey_file)
with open(sshkey_file + ".pub", "r") as fd:
sshkey = fd.read().strip()
try:
apiclient = connect(profile)
client = ClusterApi(apiclient)
except Exception as ex: # pylint: disable=broad-except
print_error(ex)
return None
try:
response = client.get_cluster(cluster_id)
except Exception as ex: # pylint: disable=broad-except
print_error(ex)
return None
ssh_public_keys = response.get("ssh_public_keys", [])
if sshkey in [key.strip() for key in ssh_public_keys]:
print_ok(" => public ssh key already configured for cluster %s" %
cluster_id)
bye()
request = {}
for key in [
"autotermination_minutes",
"cluster_id",
"cluster_name",
"cluster_source",
"creator_user_name",
"default_tags",
"driver_node_type_id",
"enable_elastic_disk",
"init_scripts_safe_mode",
"node_type_id",
"spark_version",
]:
request[key] = response[key]
if response.get("spark_env_vars", None) is not None:
request["spark_env_vars"] = response["spark_env_vars"]
if response.get("aws_attributes", None) is not None:
request["aws_attributes"] = response["aws_attributes"]
if response.get("num_workers", None) is not None:
request["num_workers"] = response["num_workers"]
if response.get("autoscale", None) is not None:
request["autoscale"] = response["autoscale"]
request["ssh_public_keys"] = ssh_public_keys + [sshkey]
print_warning(" => The ssh key will be added to the cluster." +
"\n Note: The cluster will be restarted immediately!")
answer = input(
" => Shall the ssh key be added and the cluster be restarted (y/n)? (default = n): "
)
if answer.lower() == "y":
try:
response = client.edit_cluster(request)
except DatabricksApiException as ex:
print_error(str(ex))
return None
print_ok(" => OK")
else:
print_error(" => Cancelled")