def __init__(self, logger, **kwargs):
"""
:param **kwargs:
reserved python word for unlimited parameters
keys should only include: token, host
:type **kwargs: dict
"""
self.api_client = ApiClient(**kwargs)
self.cluster_client = ClusterApi(self.api_client)
self.libraries_client = LibrariesApi(self.api_client)
self.logger = logger
def get_cli(api_client, cluster_id, cluster_name):
"""
Retrieves metadata about a cluster.
"""
if cluster_id:
cluster = ClusterApi(api_client).get_cluster(cluster_id)
elif cluster_name:
cluster = ClusterApi(api_client).get_cluster_by_name(cluster_name)
else:
raise RuntimeError('cluster_name and cluster_id were empty?')
click.echo(pretty_format(cluster))
class SdkClient():
def __init__(self, profile=None):
client = utils.get_api_client(profile)
self.cluster_client = ClusterApi(client)
self.jobs_client = JobsApi(client)
def list_clusters(self):
return self.cluster_client.list_clusters()
def get_cluster(self, cluster_id):
return self.cluster_client.get_cluster(cluster_id)
def list_jobs(self):
return self.jobs_client.list_jobs()
def restart_cli(api_client, cluster_id):
"""
Restarts a Databricks cluster given its ID.
If the cluster is not currently in a RUNNING state, nothing will happen
"""
ClusterApi(api_client).restart_cluster(cluster_id)
def list_zones_cli(api_client):
"""
Lists zones where clusters can be created.
The output format is specified in
https://docs.databricks.com/api/latest/clusters.html#list-zones
"""
click.echo(pretty_format(ClusterApi(api_client).list_zones()))
def create_cli(api_client, json_file, json):
"""
Creates a Databricks cluster.
The specification for the request json can be found at
https://docs.databricks.com/api/latest/clusters.html#create
"""
json_cli_base(json_file, json, lambda json: ClusterApi(api_client).create_cluster(json))
def list_node_types_cli(api_client):
"""
Lists possible node types for a cluster.
The output format is specified in
https://docs.databricks.com/api/latest/clusters.html#list-node-types
"""
click.echo(pretty_format(ClusterApi(api_client).list_node_types()))
def spark_versions_cli(api_client):
"""
Lists possible Databricks Runtime versions for a cluster.
The output format is specified in
https://docs.databricks.com/api/latest/clusters.html#spark-versions
"""
click.echo(pretty_format(ClusterApi(api_client).spark_versions()))
def permanent_delete_cli(api_client, cluster_id):
"""
Permanently deletes a Spark cluster.
If the cluster is running, it is terminated and its resources are asynchronously removed.
If the cluster is terminated, then it is immediately removed.
"""
ClusterApi(api_client).permanent_delete(cluster_id)
def start_cli(api_client, cluster_id):
"""
Starts a terminated Databricks cluster given its ID.
If the cluster is not currently in a TERMINATED state, nothing will happen.
"""
ClusterApi(api_client).start_cluster(cluster_id)
def resize_cli(api_client, cluster_id, num_workers):
"""Resizes a Databricks cluster given its ID.
Provide a `--num-workers` parameter to indicate the new cluster size.
If the cluster is not currently in a RUNNING state, this will cause an
error to occur.
"""
ClusterApi(api_client).resize_cluster(cluster_id, num_workers)
def delete_cli(api_client, cluster_id):
"""
Removes a Databricks cluster given its ID.
The cluster is removed asynchronously. Once the deletion has completed,
the cluster will be in a TERMINATED state. If the cluster is already in
a TERMINATING or TERMINATED state, nothing will happen.
Use ``databricks clusters get --cluster-id CLUSTER_ID`` to check termination states.
"""
ClusterApi(api_client).delete_cluster(cluster_id)
def cluster_events_cli(api_client, cluster_id, start_time, end_time, order, event_type, offset,
limit, output):
"""
Gets events for a Spark cluster.
"""
events_json = ClusterApi(api_client).get_events(
cluster_id=cluster_id, start_time=start_time, end_time=end_time, order=order,
event_types=event_type, offset=offset, limit=limit)
if OutputClickType.is_json(output):
click.echo(pretty_format(events_json))
else:
click.echo(tabulate(_cluster_events_to_table(events_json), tablefmt='plain'))
def edit_cli(api_client, json_file, json):
"""
Edits a Databricks cluster.
The specification for the request json can be found at
https://docs.databricks.com/api/latest/clusters.html#edit
"""
if not bool(json_file) ^ bool(json):
raise RuntimeError('Either --json-file or --json should be provided')
if json_file:
with open(json_file, 'r') as f:
json = f.read()
deser_json = json_loads(json)
ClusterApi(api_client).edit_cluster(deser_json)
def get_cluster_state(profile, cluster_id):
"""Get cluster state
Args:
profile (str): Databricks CLI profile string
cluster_id (str): Cluster ID
Returns:
dict: ClusterStatus
"""
apiclient = connect(profile)
try:
client = ClusterApi(apiclient)
cluster = client.get_cluster(cluster_id)
state = cluster.get("state", None)
except:
_logger.warning("DbStatusHandler cluster not reachable")
state = "TERMINATED"
# if state == "TERMINATED":
# _logger.warning("DbStatusHandler cluster state = %s" % state)
# else:
# _logger.info("DbStatusHandler cluster state = %s" % state)
return state
def list_cli(api_client, output):
"""
Lists active and recently terminated clusters.
Returns information about all currently active clusters, and up
to 100 most recently terminated clusters in the past 7 days.
By default the output format will be a human readable table with the following fields
- Cluster ID
- Cluster name
- Cluster state
"""
clusters_json = ClusterApi(api_client).list_clusters()
if OutputClickType.is_json(output):
click.echo(pretty_format(clusters_json))
else:
click.echo(tabulate(_clusters_to_table(clusters_json), tablefmt='plain'))
def tgt_cluster_api(tgt_api_client:ApiClient):
return ClusterApi(tgt_api_client)
class ClusterManagement:
def __init__(self, logger, **kwargs):
"""
:param **kwargs:
reserved python word for unlimited parameters
keys should only include: token, host
:type **kwargs: dict
"""
self.api_client = ApiClient(**kwargs)
self.cluster_client = ClusterApi(self.api_client)
self.libraries_client = LibrariesApi(self.api_client)
self.logger = logger
def create_cluster(self, cluster_specs):
"""function to build/edit cluster and start
:param cluster_specs: cluster specs in clusterconf.yaml
:type cluster_specs: dict
"""
# self.cluster_client.get_cluster_by_name("unknown")
try:
cluster = self.cluster_client.get_cluster_by_name(
cluster_specs["cluster_name"])
self.logger.info(f"cluster {cluster['cluster_name']} exists "
f"with id {cluster['cluster_id']}")
self.logger.debug(cluster_specs)
self.logger.debug(cluster)
if not cluster_specs.items() <= cluster.items():
self.logger.warning(
"cluster spec doesn't match existing cluster")
cluster_specs['cluster_id'] = cluster['cluster_id']
self.cluster_client.edit_cluster(cluster_specs)
else:
self.logger.info("cluster spec matches")
except Exception:
cluster = self.cluster_client.create_cluster(cluster_specs)
self.logger.info(f"the cluster {cluster} is being created")
time.sleep(30)
cluster_id = cluster['cluster_id']
status = self._cluster_status(cluster_id)
while status['state'] in ["RESTARTING", "RESIZING", "TERMINATING"]:
self.logger.info(
f"waiting for the cluster. status {status['state']}")
time.sleep(10)
status = self._cluster_status(cluster_id)
while status['state'] in ["TERMINATED", "PENDING"]:
self.logger.info(f"cluster status {status['state']}")
if status['state'] == "TERMINATED":
self.logger.info(f"starting cluster, status {status['state']}")
self.cluster_client.start_cluster(cluster_id)
time.sleep(10)
status = self._cluster_status(cluster_id)
self.logger.info(f"cluster is up. final status: {status['state']}")
return cluster_id
def install_cluster_library(self, cluster_id, cluster_libraries):
"""function to install libraries on cluster
:param cluster_id: id of cluster in Databricks to install libs on
:type cluster_id: str
:param cluster_libraries: clusterlib.yaml
:type cluster_libraries: list(dict)
"""
try:
if not isinstance(cluster_libraries, list):
raise ValueError(
f"cluster_libraries is not a list: {cluster_libraries}")
current_libs = self.libraries_client.cluster_status(cluster_id)
# parse the libs to match the yaml
parsed_currentlibs = []
if current_libs.get("library_statuses"):
for lib in current_libs["library_statuses"]:
parsed_currentlibs.append(lib["library"])
install_libs = [
x for x in cluster_libraries if x not in parsed_currentlibs
]
self.logger.info(f"install libraries: {install_libs}")
self.libraries_client.install_libraries(cluster_id, install_libs)
uninstall_libs = [
x for x in parsed_currentlibs if x not in cluster_libraries
]
self.logger.warning(f"uninstall libraries: {uninstall_libs}")
self.libraries_client.uninstall_libraries(cluster_id,
uninstall_libs)
except Exception as error:
self.logger.error(f"install_cluster_library error: {repr(error)}")
def _cluster_status(self, cluster_id):
"""internal method to get cluster status
:param cluster_id: id of databricks cluster
:type cluster_id: str
"""
try:
status = self.cluster_client.get_cluster(cluster_id)
return status
except Exception as error:
self.logger.error(f"cluster status error: {error}")
def delete_unmanaged_clusters(self, cluster_config):
"""function to delete clusters that are not in clusterconf.yaml
:param cluster_config: clusterconf.yaml
:type cluster_config: list(dict)
"""
existing_clusters = self.cluster_client.list_clusters()
if existing_clusters.get("clusters"):
existing_clusters = [
c for c in existing_clusters.get("clusters")
if c["cluster_source"].upper() != "JOB"
]
self.logger.debug(existing_clusters)
cluster_list = [c["cluster_name"] for c in cluster_config]
remove_cluster = [(c["cluster_name"], c["cluster_id"])
for c in existing_clusters
if c["cluster_name"] not in cluster_list]
self.logger.warning("removing unmanaged clusters:")
self.logger.warning(remove_cluster)
for c in remove_cluster:
self.logger.debug(f"deleting {c[1]}")
self.cluster_client.permanent_delete(c[1])
return
def main(self, cluster_specs, cluster_libraries):
"""main method to build/edit clusters and install libs
:cluster_spec: cluster spec in clusterconf.yaml
:type cluster_spec: dict
:param cluster_libraries: clusterlib.yaml
:type cluster_libraries: list(dict)
"""
# self.logger.info("=======================================================")
self.logger.info(
f"create/update cluster: {cluster_specs['cluster_name']}")
cluster_id = self.create_cluster(cluster_specs)
self.logger.info("installing libraries")
self.install_cluster_library(cluster_id, cluster_libraries)
overwrite=args.overwrite)
print("Created: {}".format(myegg))
############################
# Interact with Databricks:
############################
# first make sure you are using the correct profile and connecting to the intended workspace
my_api_client = _get_api_client(
ProfileConfigProvider(args.profile).get_config())
# Create a cluster if flagged
if args.create_cluster:
# treat args.cluster_id as the name, because if you create a cluster, you do not know its id yet.
DEFAULT_CLUSTER_CONFIG["cluster_name"] = args.cluster_id
cluster_info = ClusterApi(my_api_client).create_cluster(
DEFAULT_CLUSTER_CONFIG)
args.cluster_id = cluster_info["cluster_id"]
print("Creating a new cluster with name {}. New cluster_id={}".format(
DEFAULT_CLUSTER_CONFIG["cluster_name"], args.cluster_id))
# Upload the egg:
upload_path = Path(args.dbfs_path, args.eggname).as_posix()
# Check if file exists to alert user.
print("Uploading {} to databricks at {}".format(args.eggname, upload_path))
if dbfs_file_exists(my_api_client, upload_path):
if args.overwrite:
print("Overwriting file at {}".format(upload_path))
else:
raise IOError("""
{} already exists on databricks cluster.
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")
getContext().apiUrl().get())
# COMMAND ----------
apiclient.perform_query("GET", "/clusters/list?state=RUNNING")
# COMMAND ----------
from databricks_cli.clusters.api import ClusterApi
import json
from pyspark.sql.types import *
from pyspark.sql.functions import *
# COMMAND ----------
clusters_api = ClusterApi(apiclient)
# COMMAND ----------
data = clusters_api.list_clusters()
rdd = sc.parallelize(data['clusters']).map(lambda x: json.dumps(x))
#
# COMMAND ----------
raw_clusters_df = spark.read.json(rdd)
# COMMAND ----------
parsed_clusters = raw_clusters_df.filter(col('init_scripts').isNotNull())
display(parsed_clusters)
from databrickslabs_jupyterlab.remote import connect
from helpers import get_profile, get_running_clusters, remove_running_clusters
def delete_cluster(client, cluster_id):
try:
response = client.permanent_delete(cluster_id)
return response
except Exception as ex: # pylint: disable=broad-except
print(ex)
return None
clusters = get_running_clusters()
profile = get_profile()
try:
apiclient = connect(profile)
client = ClusterApi(apiclient)
except Exception as ex: # pylint: disable=broad-except
print(ex)
sys.exit(1)
for name, cluster_id in clusters.items():
print(name)
delete_cluster(client, cluster_id)
remove_running_clusters()
def __init__(self, profile=None):
client = utils.get_api_client(profile)
self.cluster_client = ClusterApi(client)
self.jobs_client = JobsApi(client)
def src_cluster_api(src_api_client:ApiClient):
return ClusterApi(src_api_client)
def get_cli(api_client, cluster_id):
"""
Retrieves metadata about a cluster.
"""
click.echo(pretty_format(ClusterApi(api_client).get_cluster(cluster_id)))
def configure_ssh(profile, host, token, 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:
print_error(ex)
return None
try:
response = client.get_cluster(cluster_id)
except Exception as ex:
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")
def get_cluster(profile, cluster_id=None, status=None):
"""Get the cluster configuration from remote
Args:
profile (str): Databricks CLI profile string
cluster_id (str, optional): If cluster_id is given, the user will not be asked to select one.
Defaults to None.
status (Status, optional): A Status class providing set_status. Defaults to None.
Returns:
tuple: Cluster configs: cluster_id, public_ip, cluster_name and a flag whether it was started
Returns:
tuple: (cluster_id, public_ip, cluster_name, started)
"""
with open("%s/.ssh/id_%s.pub" % (expanduser("~"), profile)) as fd:
try:
ssh_pub = fd.read().strip()
except:
print_error(
" Error: ssh key for profile 'id_%s.pub' does not exist in %s/.ssh"
% (profile, expanduser("~")))
bye()
try:
apiclient = connect(profile)
client = ClusterApi(apiclient)
clusters = client.list_clusters()
except Exception as ex:
print_error(ex)
return (None, None, None, None)
clusters = clusters["clusters"]
if cluster_id is not None:
cluster = None
for c in clusters:
if c["cluster_id"] == cluster_id:
cluster = c
break
if cluster is None:
print_error(
" Error: A cluster with id '%s' does not exist in the workspace of profile '%s'"
% (cluster_id, profile))
return (None, None, None, None)
if ssh_pub not in [
c.strip() for c in cluster.get("ssh_public_keys", [])
]:
print_error(
" Error: Cluster with id '%s' does not have ssh key '~/.ssh/id_%s' configured"
% (cluster_id, profile))
return (None, None, None, None)
else:
my_clusters = [
cluster for cluster in clusters if ssh_pub in
[c.strip() for c in cluster.get("ssh_public_keys", [])]
]
if not my_clusters:
print_error(
" Error: There is no cluster in the workspace for profile '%s' configured with ssh key '~/.ssh/id_%s':"
% (profile, profile))
print(
" Use 'databrickslabs_jupyterlab %s -s' to configure ssh for clusters in this workspace\n"
% profile)
return (None, None, None, None)
current_conda_env = os.environ.get("CONDA_DEFAULT_ENV", None)
found = None
for i, c in enumerate(my_clusters):
if c["cluster_name"].replace(" ", "_") == current_conda_env:
found = c["cluster_name"]
break
if found is not None:
print_warning(
"\n => The current conda environment is '%s'.\n You might want to select cluster %d with the name '%s'?\n"
% (current_conda_env, i, found))
cluster = select_cluster(my_clusters)
cluster_id = cluster["cluster_id"]
cluster_name = cluster["cluster_name"]
try:
response = client.get_cluster(cluster_id)
except Exception as ex:
print_error(ex)
return (None, None, None, None)
state = response["state"]
if not state in ["RUNNING", "RESIZING"]:
if state == "TERMINATED":
print(" => Starting cluster %s" % cluster_id)
if status is not None:
status.set_status(profile, cluster_id, "Cluster Starting")
try:
response = client.start_cluster(cluster_id)
except Exception as ex:
print_error(ex)
return (None, None, None, None)
print(
" => Waiting for cluster %s being started (this can take up to 5 min)"
% cluster_id)
print(" ", end="", flush=True)
while not state in ("RUNNING", "RESIZING"):
if status is not None:
status.set_status(profile, cluster_id, "Cluster Starting",
False)
print(".", end="", flush=True)
time.sleep(5)
try:
response = client.get_cluster(cluster_id)
except Exception as ex:
print_error(ex)
return (None, None, None, None)
if response.get("error", None) is not None:
print_error(response["error"])
return (None, None, None, None)
else:
state = response["state"]
print_ok("\n => OK")
if status is not None:
status.set_status(profile, cluster_id, "Cluster started")
print(
"\n => Waiting for libraries on cluster %s being installed (this can take some time)"
% cluster_id)
print(" ", end="", flush=True)
done = False
while not done:
try:
states = get_library_state(profile, cluster_id)
except DatabricksApiException as ex:
print_error(ex)
return (None, None, None, None)
installing = any(
[s in ["PENDING", "RESOLVING", "INSTALLING"] for s in states])
if installing:
if status is not None:
status.set_status(profile, cluster_id,
"Installing cluster libraries", False)
print(".", end="", flush=True)
time.sleep(5)
else:
done = True
print_ok("\n => OK\n")
if status is not None:
status.set_status(profile, cluster_id,
"Cluster libraries installed", False)
public_ip = response["driver"].get("public_dns", None)
if public_ip is None:
print_error(" Error: Cluster does not have public DNS name")
return (None, None, None, None)
print_ok(" => Selected cluster: %s (%s)" % (cluster_name, public_ip))
return (cluster_id, public_ip, cluster_name, None)
def test_src_clusters(src_cluster_api: ClusterApi):
print(src_cluster_api.list_clusters())
assert src_cluster_api.list_clusters() is not None