Clusters & Compute Management

import requests
import json

DATABRICKS_HOST = "https://adb-1234567890.12.azuredatabricks.net"
TOKEN = "dapi_your_token"
headers = {"Authorization": f"Bearer {TOKEN}"}

# Create an All-Purpose Cluster
all_purpose_config = {
    "cluster_name": "dev-interactive",
    "spark_version": "14.3.x-scala2.12",
    "node_type_id": "Standard_DS3_v2",
    "num_workers": 2,
    "autotermination_minutes": 30,
    "spark_conf": {
        "spark.databricks.cluster.profile": "serverless"
    }
}

response = requests.post(
    f"{DATABRICKS_HOST}/api/2.0/clusters/create",
    headers=headers,
    json=all_purpose_config
)
cluster_id = response.json()["cluster_id"]
print(f"Created cluster: {cluster_id}")

# Job clusters are defined within job configurations
job_cluster_config = {
    "new_cluster": {
        "spark_version": "14.3.x-scala2.12",
        "node_type_id": "Standard_DS3_v2",
        "num_workers": 4,
    }
}

# List available node types for your cloud
response = requests.get(
    f"{DATABRICKS_HOST}/api/2.0/clusters/list-node-types",
    headers=headers
)

node_types = response.json()["node_types"]
for nt in node_types[:5]:
    print(f"{nt['node_type_id']}: {nt['memory_mb']}MB, {nt['num_cores']} cores")

# Common sizing guidelines:
# Memory-optimized: Large shuffles, caching, ML training
# Compute-optimized: CPU-intensive transformations
# Storage-optimized: Large Delta Lake reads/writes
# GPU-accelerated: Deep learning, NLP models

# Key Spark configurations for cluster setup
spark_conf = {
    # Adaptive Query Execution (enabled by default in DBR 14+)
    "spark.sql.adaptive.enabled": "true",

    # Dynamic partition pruning
    "spark.sql.optimizer.dynamicPartitionPruning.enabled": "true",

    # Shuffle partitions -- tune based on data size
    "spark.sql.shuffle.partitions": "auto",

    # Delta Lake optimizations
    "spark.databricks.delta.optimizeWrite.enabled": "true",
    "spark.databricks.delta.autoCompact.enabled": "true",
}

# Set configs at runtime in a notebook
spark.conf.set("spark.sql.shuffle.partitions", "200")

# Check current configuration
print(spark.conf.get("spark.sql.shuffle.partitions"))

# Create cluster with autoscaling enabled
autoscaling_cluster = {
    "cluster_name": "autoscaling-etl",
    "spark_version": "14.3.x-scala2.12",
    "node_type_id": "Standard_DS4_v2",
    "autoscale": {
        "min_workers": 2,
        "max_workers": 10
    },
    "autotermination_minutes": 30
}

# Monitor autoscaling activity
cluster_info = requests.get(
    f"{DATABRICKS_HOST}/api/2.0/clusters/get",
    headers=headers,
    params={"cluster_id": cluster_id}
).json()

print(f"Current workers: {cluster_info.get('num_workers', 'N/A')}")
print(f"State: {cluster_info['state']}")

# Cluster with spot instances for workers
spot_cluster = {
    "cluster_name": "cost-optimized-etl",
    "spark_version": "14.3.x-scala2.12",
    "node_type_id": "Standard_DS4_v2",
    "num_workers": 4,

    # AWS: use spot instances
    "aws_attributes": {
        "first_on_demand": 1,  # Driver on-demand
        "availability": "SPOT_WITH_FALLBACK",
        "spot_bid_price_percent": 100,
        "zone_id": "auto"
    },

    # Azure: use spot VMs
    "azure_attributes": {
        "first_on_demand": 1,
        "availability": "SPOT_WITH_FALLBACK_AZURE",
        "spot_bid_max_price": -1
    }
}

# Availability options:
# SPOT             - Only spot (cheapest, risk of interruption)
# SPOT_WITH_FALLBACK - Try spot, fall back to on-demand
# ON_DEMAND        - Only on-demand (most reliable)

# Start a terminated cluster
requests.post(
    f"{DATABRICKS_HOST}/api/2.0/clusters/start",
    headers=headers,
    json={"cluster_id": cluster_id}
)

# Terminate a running cluster (graceful shutdown)
requests.post(
    f"{DATABRICKS_HOST}/api/2.0/clusters/delete",
    headers=headers,
    json={"cluster_id": cluster_id}
)

# Restart a cluster (useful after config changes)
requests.post(
    f"{DATABRICKS_HOST}/api/2.0/clusters/restart",
    headers=headers,
    json={"cluster_id": cluster_id}
)

# Poll cluster state until running
import time

def wait_for_cluster(cluster_id, target="RUNNING", timeout=600):
    start = time.time()
    while time.time() - start < timeout:
        info = requests.get(
            f"{DATABRICKS_HOST}/api/2.0/clusters/get",
            headers=headers,
            params={"cluster_id": cluster_id}
        ).json()
        state = info["state"]
        print(f"State: {state}")
        if state == target:
            return info
        if state in ["TERMINATED", "ERROR"]:
            raise Exception(f"Cluster entered {state}")
        time.sleep(15)

{
    "name": "Cost-Controlled Development",
    "definition": {
        "autotermination_minutes": {
            "type": "range",
            "maxValue": 60,
            "defaultValue": 30
        },
        "num_workers": {
            "type": "range",
            "maxValue": 8,
            "defaultValue": 2
        },
        "node_type_id": {
            "type": "allowlist",
            "values": ["Standard_DS3_v2", "Standard_DS4_v2"]
        }
    }
}