skypilot-nightly 1.0.0.dev2024053101__py3-none-any.whl → 1.0.0.dev2025022801__py3-none-any.whl

sky/utils/dag_utils.py

@@ -3,16 +3,16 @@ import copy
 from typing import Any, Dict, List, Optional, Tuple
 
 from sky import dag as dag_lib
-from sky import jobs
 from sky import sky_logging
 from sky import task as task_lib
-from sky.backends import backend_utils
+from sky.utils import cluster_utils
 from sky.utils import common_utils
+from sky.utils import registry
 from sky.utils import ux_utils
 
 logger = sky_logging.init_logger(__name__)
 
-# Message thrown when APIs sky.{exec,launch,spot.launch}() received a string
+# Message thrown when APIs sky.{exec,launch,jobs.launch}() received a string
 # instead of a Dag.  CLI (cli.py) is implemented by us so should not trigger
 # this.
 _ENTRYPOINT_STRING_AS_DAG_MESSAGE = """\
@@ -31,54 +31,43 @@ The command can then be run as:
 
   sky.launch(task, ...)
 
-  sky.spot.launch(task, ...)
+  sky.jobs.launch(task, ...)
 """.strip()
 
 
 def convert_entrypoint_to_dag(entrypoint: Any) -> 'dag_lib.Dag':
-    """Convert the entrypoint to a sky.Dag.
+    """Converts the entrypoint to a sky.Dag and applies the policy.
 
     Raises TypeError if 'entrypoint' is not a 'sky.Task' or 'sky.Dag'.
     """
     # Not suppressing stacktrace: when calling this via API user may want to
     # see their own program in the stacktrace. Our CLI impl would not trigger
     # these errors.
+    converted_dag: 'dag_lib.Dag'
     if isinstance(entrypoint, str):
         with ux_utils.print_exception_no_traceback():
             raise TypeError(_ENTRYPOINT_STRING_AS_DAG_MESSAGE)
     elif isinstance(entrypoint, dag_lib.Dag):
-        return copy.deepcopy(entrypoint)
+        converted_dag = copy.deepcopy(entrypoint)
     elif isinstance(entrypoint, task_lib.Task):
         entrypoint = copy.deepcopy(entrypoint)
         with dag_lib.Dag() as dag:
             dag.add(entrypoint)
             dag.name = entrypoint.name
-        return dag
+        converted_dag = dag
     else:
         with ux_utils.print_exception_no_traceback():
             raise TypeError(
                 'Expected a sky.Task or sky.Dag but received argument of type: '
                 f'{type(entrypoint)}')
 
+    return converted_dag
 
-def load_chain_dag_from_yaml(
-    path: str,
-    env_overrides: Optional[List[Tuple[str, str]]] = None,
-) -> dag_lib.Dag:
-    """Loads a chain DAG from a YAML file.
-
-    Has special handling for an initial section in YAML that contains only the
-    'name' field, which is the DAG name.
 
-    'env_overrides' is a list of (key, value) pairs that will be used to update
-    the task's 'envs' section. If it is a chain dag, the envs will be updated
-    for all tasks in the chain.
-
-    Returns:
-      A chain Dag with 1 or more tasks (an empty entrypoint would create a
-      trivial task).
-    """
-    configs = common_utils.read_yaml_all(path)
+def _load_chain_dag(
+        configs: List[Dict[str, Any]],
+        env_overrides: Optional[List[Tuple[str, str]]] = None) -> dag_lib.Dag:
+    """Loads a chain DAG from a list of YAML configs."""
     dag_name = None
     if set(configs[0].keys()) == {'name'}:
         dag_name = configs[0]['name']
@@ -86,7 +75,7 @@ def load_chain_dag_from_yaml(
     elif len(configs) == 1:
         dag_name = configs[0].get('name')
 
-    if len(configs) == 0:
+    if not configs:
         # YAML has only `name: xxx`. Still instantiate a task.
         configs = [{'name': dag_name}]
 
@@ -103,12 +92,74 @@ def load_chain_dag_from_yaml(
     return dag
 
 
-def dump_chain_dag_to_yaml(dag: dag_lib.Dag, path: str) -> None:
+def load_chain_dag_from_yaml(
+    path: str,
+    env_overrides: Optional[List[Tuple[str, str]]] = None,
+) -> dag_lib.Dag:
+    """Loads a chain DAG from a YAML file.
+
+    Has special handling for an initial section in YAML that contains only the
+    'name' field, which is the DAG name.
+
+    'env_overrides' is a list of (key, value) pairs that will be used to update
+    the task's 'envs' section. If it is a chain dag, the envs will be updated
+    for all tasks in the chain.
+
+    Returns:
+      A chain Dag with 1 or more tasks (an empty entrypoint would create a
+      trivial task).
+    """
+    configs = common_utils.read_yaml_all(path)
+    return _load_chain_dag(configs, env_overrides)
+
+
+def load_chain_dag_from_yaml_str(
+    yaml_str: str,
+    env_overrides: Optional[List[Tuple[str, str]]] = None,
+) -> dag_lib.Dag:
+    """Loads a chain DAG from a YAML string.
+
+    Has special handling for an initial section in YAML that contains only the
+    'name' field, which is the DAG name.
+
+    'env_overrides' is a list of (key, value) pairs that will be used to update
+    the task's 'envs' section. If it is a chain dag, the envs will be updated
+    for all tasks in the chain.
+
+    Returns:
+      A chain Dag with 1 or more tasks (an empty entrypoint would create a
+      trivial task).
+    """
+    configs = common_utils.read_yaml_all_str(yaml_str)
+    return _load_chain_dag(configs, env_overrides)
+
+
+def dump_chain_dag_to_yaml_str(dag: dag_lib.Dag) -> str:
+    """Dumps a chain DAG to a YAML string.
+
+    Args:
+        dag: the DAG to dump.
+
+    Returns:
+        The YAML string.
+    """
     assert dag.is_chain(), dag
     configs = [{'name': dag.name}]
     for task in dag.tasks:
         configs.append(task.to_yaml_config())
-    common_utils.dump_yaml(path, configs)
+    return common_utils.dump_yaml_str(configs)
+
+
+def dump_chain_dag_to_yaml(dag: dag_lib.Dag, path: str) -> None:
+    """Dumps a chain DAG to a YAML file.
+
+    Args:
+        dag: the DAG to dump.
+        path: the path to the YAML file.
+    """
+    dag_str = dump_chain_dag_to_yaml_str(dag)
+    with open(path, 'w', encoding='utf-8') as f:
+        f.write(dag_str)
 
 
 def maybe_infer_and_fill_dag_and_task_names(dag: dag_lib.Dag) -> None:
@@ -125,7 +176,7 @@ def maybe_infer_and_fill_dag_and_task_names(dag: dag_lib.Dag) -> None:
             dag.name = first_task.name
 
     if dag.name is None:
-        dag.name = backend_utils.generate_cluster_name()
+        dag.name = cluster_utils.generate_cluster_name()
 
     if len(dag.tasks) == 1:
         if first_task.name is None:
@@ -140,11 +191,21 @@ def fill_default_config_in_dag_for_job_launch(dag: dag_lib.Dag) -> None:
     for task_ in dag.tasks:
 
         new_resources_list = []
+        default_strategy = registry.JOBS_RECOVERY_STRATEGY_REGISTRY.default
+        assert default_strategy is not None
         for resources in list(task_.resources):
-            change_default_value: Dict[str, Any] = {}
-            if resources.job_recovery is None:
-                change_default_value[
-                    'job_recovery'] = jobs.DEFAULT_RECOVERY_STRATEGY
+            original_job_recovery = resources.job_recovery
+            job_recovery = {'strategy': default_strategy}
+            if isinstance(original_job_recovery, str):
+                job_recovery['strategy'] = original_job_recovery
+            elif isinstance(original_job_recovery, dict):
+                job_recovery.update(original_job_recovery)
+                strategy = job_recovery.get('strategy')
+                if strategy is None:
+                    job_recovery['strategy'] = default_strategy
+            change_default_value: Dict[str, Any] = {
+                'job_recovery': job_recovery
+            }
 
             new_resources = resources.copy(**change_default_value)
             new_resources_list.append(new_resources)

sky/utils/db_utils.py

@@ -4,11 +4,27 @@ import sqlite3
 import threading
 from typing import Any, Callable, Optional
 
+# This parameter (passed to sqlite3.connect) controls how long we will wait to
+# obtains a database lock (not necessarily during connection, but whenever it is
+# needed). It is not a connection timeout.
+# Even in WAL mode, only a single writer is allowed at a time. Other writers
+# will block until the write lock can be obtained. This behavior is described in
+# the SQLite documentation for WAL: https://www.sqlite.org/wal.html
+# Python's default timeout is 5s. In normal usage, lock contention is very low,
+# and this is more than sufficient. However, in some highly concurrent cases,
+# such as a jobs controller suddenly recovering thousands of jobs at once, we
+# can see a small number of processes that take much longer to obtain the lock.
+# In contrived highly contentious cases, around 0.1% of transactions will take
+# >30s to take the lock. We have not seen cases that take >60s. For cases up to
+# 1000x parallelism, this is thus thought to be a conservative setting.
+# For more info, see the PR description for #4552.
+_DB_TIMEOUT_S = 60
+
 
 @contextlib.contextmanager
 def safe_cursor(db_path: str):
     """A newly created, auto-committing, auto-closing cursor."""
-    conn = sqlite3.connect(db_path)
+    conn = sqlite3.connect(db_path, timeout=_DB_TIMEOUT_S)
     cursor = conn.cursor()
     try:
         yield cursor
@@ -79,8 +95,6 @@ class SQLiteConn(threading.local):
     def __init__(self, db_path: str, create_table: Callable):
         super().__init__()
         self.db_path = db_path
-        # NOTE: We use a timeout of 10 seconds to avoid database locked
-        # errors. This is a hack, but it works.
-        self.conn = sqlite3.connect(db_path, timeout=10)
+        self.conn = sqlite3.connect(db_path, timeout=_DB_TIMEOUT_S)
         self.cursor = self.conn.cursor()
         create_table(self.cursor, self.conn)

sky/utils/env_options.py

@@ -1,21 +1,43 @@
 """Global environment options for sky."""
 import enum
 import os
+from typing import Dict
 
 
 class Options(enum.Enum):
     """Environment variables for SkyPilot."""
-    IS_DEVELOPER = 'SKYPILOT_DEV'
-    SHOW_DEBUG_INFO = 'SKYPILOT_DEBUG'
-    DISABLE_LOGGING = 'SKYPILOT_DISABLE_USAGE_COLLECTION'
-    MINIMIZE_LOGGING = 'SKYPILOT_MINIMIZE_LOGGING'
+
+    # (env var name, default value)
+    IS_DEVELOPER = ('SKYPILOT_DEV', False)
+    SHOW_DEBUG_INFO = ('SKYPILOT_DEBUG', False)
+    DISABLE_LOGGING = ('SKYPILOT_DISABLE_USAGE_COLLECTION', False)
+    MINIMIZE_LOGGING = ('SKYPILOT_MINIMIZE_LOGGING', True)
+    SUPPRESS_SENSITIVE_LOG = ('SKYPILOT_SUPPRESS_SENSITIVE_LOG', False)
     # Internal: this is used to skip the cloud user identity check, which is
     # used to protect cluster operations in a multi-identity scenario.
     # Currently, this is only used in the job and serve controller, as there
     # will not be multiple identities, and skipping the check can increase
     # robustness.
-    SKIP_CLOUD_IDENTITY_CHECK = 'SKYPILOT_SKIP_CLOUD_IDENTITY_CHECK'
+    SKIP_CLOUD_IDENTITY_CHECK = ('SKYPILOT_SKIP_CLOUD_IDENTITY_CHECK', False)
+
+    def __init__(self, env_var: str, default: bool) -> None:
+        self.env_var = env_var
+        self.default = default
+
+    def __repr__(self) -> str:
+        return self.env_var
 
-    def get(self):
+    def get(self) -> bool:
         """Check if an environment variable is set to True."""
-        return os.getenv(self.value, 'False').lower() in ('true', '1')
+        return os.getenv(self.env_var,
+                         str(self.default)).lower() in ('true', '1')
+
+    @property
+    def env_key(self) -> str:
+        """The environment variable key name."""
+        return self.value[0]
+
+    @classmethod
+    def all_options(cls) -> Dict[str, bool]:
+        """Returns all options as a dictionary."""
+        return {option.env_key: option.get() for option in list(Options)}

sky/utils/kubernetes/create_cluster.sh

@@ -12,9 +12,11 @@ IMAGE_GPU="us-central1-docker.pkg.dev/skypilot-375900/skypilotk8s/skypilot-gpu:l
 PORT_RANGE_START=30000
 PORT_RANGE_END=30100
 
+USER_HASH=$1
+
 # Check for GPU flag
 ENABLE_GPUS=false
-if [[ "$1" == "--gpus" ]]; then
+if [[ "$2" == "--gpus" ]]; then
     ENABLE_GPUS=true
 fi
 
@@ -88,45 +90,20 @@ if kind get clusters | grep -q skypilot; then
 fi
 
 # Generate cluster YAML
-echo "Generating /tmp/skypilot-kind.yaml"
+YAML_PATH="/tmp/skypilot-kind-$USER_HASH.yaml"
+echo "Generating $YAML_PATH"
 
 # Add GPUs flag to the generate_kind_config.py command if GPUs are enabled
 if $ENABLE_GPUS; then
-    python -m sky.utils.kubernetes.generate_kind_config --path /tmp/skypilot-kind.yaml --port-start ${PORT_RANGE_START} --port-end ${PORT_RANGE_END} --gpus
+    python -m sky.utils.kubernetes.generate_kind_config --path $YAML_PATH --port-start ${PORT_RANGE_START} --port-end ${PORT_RANGE_END} --gpus
 else
-  python -m sky.utils.kubernetes.generate_kind_config --path /tmp/skypilot-kind.yaml --port-start ${PORT_RANGE_START} --port-end ${PORT_RANGE_END}
+  python -m sky.utils.kubernetes.generate_kind_config --path $YAML_PATH --port-start ${PORT_RANGE_START} --port-end ${PORT_RANGE_END}
 fi
 
-kind create cluster --config /tmp/skypilot-kind.yaml --name skypilot
+kind create cluster --config $YAML_PATH --name skypilot
 
 echo "Kind cluster created."
 
-# Function to wait for SkyPilot GPU labeling jobs to complete
-wait_for_gpu_labeling_jobs() {
-    echo "Starting wait for SkyPilot GPU labeling jobs to complete..."
-
-    SECONDS=0
-    TIMEOUT=600  # 10 minutes in seconds
-
-    while true; do
-        TOTAL_JOBS=$(kubectl get jobs -n kube-system -l job=sky-gpu-labeler --no-headers | wc -l)
-        COMPLETED_JOBS=$(kubectl get jobs -n kube-system -l job=sky-gpu-labeler --no-headers | grep "1/1" | wc -l)
-
-        if [[ $COMPLETED_JOBS -eq $TOTAL_JOBS ]]; then
-            echo "All SkyPilot GPU labeling jobs completed ($TOTAL_JOBS)."
-            break
-        elif [ $SECONDS -ge $TIMEOUT ]; then
-            echo "Timeout reached while waiting for GPU labeling jobs."
-            exit 1
-        else
-            echo "Waiting for GPU labeling jobs to complete... ($COMPLETED_JOBS/$TOTAL_JOBS completed)"
-            echo "To check status, see GPU labeling pods:"
-            echo "kubectl get jobs -n kube-system -l job=sky-gpu-labeler"
-            sleep 5
-        fi
-    done
-}
-
 # Function to wait for GPU operator to be correctly installed
 wait_for_gpu_operator_installation() {
     echo "Starting wait for GPU operator installation..."
@@ -150,22 +127,6 @@ wait_for_gpu_operator_installation() {
     done
 }
 
-wait_for_skypilot_gpu_image_pull() {
-    echo "Pulling SkyPilot GPU image..."
-    docker pull ${IMAGE_GPU}
-    echo "Loading SkyPilot GPU image into kind cluster..."
-    kind load docker-image --name skypilot ${IMAGE_GPU}
-    echo "SkyPilot GPU image loaded into kind cluster."
-}
-
-wait_for_skypilot_cpu_image_pull() {
-    echo "Pulling SkyPilot CPU image..."
-    docker pull ${IMAGE}
-    echo "Loading SkyPilot CPU image into kind cluster..."
-    kind load docker-image --name skypilot ${IMAGE}
-    echo "SkyPilot CPU image loaded into kind cluster."
-}
-
 wait_for_nginx_ingress_controller_install() {
     echo "Starting installation of Nginx Ingress Controller..."
 
@@ -206,21 +167,8 @@ if $ENABLE_GPUS; then
          nvidia/gpu-operator --set driver.enabled=false
     # Wait for GPU operator installation to succeed
     wait_for_gpu_operator_installation
-
-    # Load the SkyPilot GPU image into the cluster for faster labelling
-    wait_for_skypilot_gpu_image_pull
-
-    # Label nodes with GPUs
-    echo "Labelling nodes with GPUs..."
-    python -m sky.utils.kubernetes.gpu_labeler
-
-    # Wait for all the GPU labeling jobs to complete
-    wait_for_gpu_labeling_jobs
 fi
 
-# Load local skypilot image on to the cluster for faster startup
-wait_for_skypilot_cpu_image_pull
-
 # Install the Nginx Ingress Controller
 wait_for_nginx_ingress_controller_install
 

sky/utils/kubernetes/deploy_remote_cluster.sh

@@ -0,0 +1,243 @@
+#!/bin/bash
+# Refer to https://docs.skypilot.co/en/latest/reservations/existing-machines.html for details on how to use this script.
+set -e
+
+# Colors for nicer UX
+RED='\033[0;31m'
+GREEN='\033[0;32m'
+YELLOW='\033[1;33m'
+NC='\033[0m' # No color
+
+# Variables
+IPS_FILE=$1
+USER=$2
+SSH_KEY=$3
+K3S_TOKEN=mytoken  # Any string can be used as the token
+CLEANUP=false
+INSTALL_GPU=false
+
+if [[ "$4" == "--cleanup" ]]; then
+    CLEANUP=true
+fi
+
+# Basic argument checks
+if [ -z "$IPS_FILE" ] || [ -z "$USER" ] || [ -z "$SSH_KEY" ]; then
+    >&2 echo -e "${RED}Error: Missing required arguments.${NC}"
+    >&2 echo "Usage: ./deploy_remote_cluster.sh ips.txt username path/to/ssh/key [--cleanup]"
+    exit 1
+fi
+
+# Check if SSH key exists
+if [ ! -f "$SSH_KEY" ]; then
+    >&2 echo -e "${RED}Error: SSH key not found: $SSH_KEY${NC}"
+    exit 1
+fi
+
+# Check if IPs file exists
+if [ ! -f "$IPS_FILE" ]; then
+    >&2 echo -e "${RED}Error: IPs file not found: $IPS_FILE${NC}"
+    exit 1
+fi
+
+# Get head node and worker nodes from the IPs file
+HEAD_NODE=$(head -n 1 "$IPS_FILE")
+WORKER_NODES=$(tail -n +2 "$IPS_FILE")
+
+# Check if the IPs file is empty or not formatted correctly
+if [ -z "$HEAD_NODE" ]; then
+    >&2 echo -e "${RED}Error: IPs file is empty or not formatted correctly.${NC}"
+    exit 1
+fi
+
+# Function to show a progress message
+progress_message() {
+    echo -e "${YELLOW}➜ $1${NC}"
+}
+
+# Step to display success
+success_message() {
+    echo -e "${GREEN}✔ $1${NC}"
+}
+
+# Function to run a command on a remote machine via SSH
+run_remote() {
+    local NODE_IP=$1
+    local CMD=$2
+    # echo -e "${YELLOW}Running command on $NODE_IP...${NC}"
+    ssh -o StrictHostKeyChecking=no -o IdentitiesOnly=yes -i "$SSH_KEY" "$USER@$NODE_IP" "$CMD"
+}
+
+# Function to uninstall k3s and clean up the state on a remote machine
+cleanup_server_node() {
+    local NODE_IP=$1
+    echo -e "${YELLOW}Cleaning up head node $NODE_IP...${NC}"
+    run_remote "$NODE_IP" "
+        echo 'Uninstalling k3s...' &&
+        /usr/local/bin/k3s-uninstall.sh || true &&
+        sudo rm -rf /etc/rancher /var/lib/rancher /var/lib/kubelet /etc/kubernetes ~/.kube
+    "
+    echo -e "${GREEN}Node $NODE_IP cleaned up successfully.${NC}"
+}
+
+# Function to uninstall k3s and clean up the state on a remote machine
+cleanup_agent_node() {
+    local NODE_IP=$1
+    echo -e "${YELLOW}Cleaning up node $NODE_IP...${NC}"
+    run_remote "$NODE_IP" "
+        echo 'Uninstalling k3s...' &&
+        /usr/local/bin/k3s-agent-uninstall.sh || true &&
+        sudo rm -rf /etc/rancher /var/lib/rancher /var/lib/kubelet /etc/kubernetes ~/.kube
+    "
+    echo -e "${GREEN}Node $NODE_IP cleaned up successfully.${NC}"
+}
+
+check_gpu() {
+    local NODE_IP=$1
+    if run_remote "$NODE_IP" "command -v nvidia-smi &> /dev/null && nvidia-smi --query-gpu=gpu_name --format=csv,noheader &> /dev/null"; then
+        return 0  # GPU detected
+    else
+        return 1  # No GPU detected
+    fi
+}
+
+# Pre-flight checks
+run_remote "$HEAD_NODE" "echo 'SSH connection successful'"
+# TODO: Add more pre-flight checks here, including checking if port 6443 is accessible
+
+# If --cleanup flag is set, uninstall k3s and exit
+if [ "$CLEANUP" == "true" ]; then
+    echo -e "${YELLOW}Starting cleanup...${NC}"
+
+    # Clean up head node
+    cleanup_server_node "$HEAD_NODE"
+
+    # Clean up worker nodes
+    for NODE in $WORKER_NODES; do
+        cleanup_agent_node "$NODE"
+    done
+
+    echo -e "${GREEN}Cleanup completed successfully.${NC}"
+    exit 0
+fi
+
+# Step 1: Install k3s on the head node
+progress_message "Deploying Kubernetes on head node ($HEAD_NODE)..."
+run_remote "$HEAD_NODE" "
+    curl -sfL https://get.k3s.io | K3S_TOKEN=$K3S_TOKEN sh - &&
+    mkdir -p ~/.kube &&
+    sudo cp /etc/rancher/k3s/k3s.yaml ~/.kube/config &&
+    sudo chown \$(id -u):\$(id -g) ~/.kube/config &&
+    for i in {1..3}; do
+        if kubectl wait --for=condition=ready node --all --timeout=2m --kubeconfig ~/.kube/config; then
+            break
+        else
+            echo 'Waiting for nodes to be ready...'
+            sleep 5
+        fi
+    done
+    if [ $i -eq 3 ]; then
+        echo 'Failed to wait for nodes to be ready after 3 attempts'
+        exit 1
+    fi"
+success_message "K3s deployed on head node."
+
+# Check if head node has a GPU
+if check_gpu "$HEAD_NODE"; then
+    echo -e "${YELLOW}GPU detected on head node ($HEAD_NODE).${NC}"
+    INSTALL_GPU=true
+fi
+
+# Fetch the head node's internal IP (this will be passed to worker nodes)
+MASTER_ADDR=$(run_remote "$HEAD_NODE" "hostname -I | awk '{print \$1}'")
+
+echo -e "${GREEN}Master node internal IP: $MASTER_ADDR${NC}"
+
+# Step 2: Install k3s on worker nodes and join them to the master node
+for NODE in $WORKER_NODES; do
+    progress_message "Deploying Kubernetes on worker node ($NODE)..."
+    run_remote "$NODE" "
+        curl -sfL https://get.k3s.io | K3S_URL=https://$MASTER_ADDR:6443 K3S_TOKEN=$K3S_TOKEN sh -"
+    success_message "Kubernetes deployed on worker node ($NODE)."
+
+    # Check if worker node has a GPU
+    if check_gpu "$NODE"; then
+        echo -e "${YELLOW}GPU detected on worker node ($NODE).${NC}"
+        INSTALL_GPU=true
+    fi
+done
+# Step 3: Configure local kubectl to connect to the cluster
+progress_message "Configuring local kubectl to connect to the cluster..."
+scp -o StrictHostKeyChecking=no -o IdentitiesOnly=yes -i "$SSH_KEY" "$USER@$HEAD_NODE":~/.kube/config ~/.kube/config
+
+# Back up the original kubeconfig file if it exists
+KUBECONFIG_FILE="$HOME/.kube/config"
+if [[ -f "$KUBECONFIG_FILE" ]]; then
+    echo "Backing up existing kubeconfig to $KUBECONFIG_FILE.bak"
+    cp "$KUBECONFIG_FILE" "$KUBECONFIG_FILE.bak"
+fi
+
+# Update kubeconfig for the local machine to use the master node's IP
+# Temporary file to hold the modified kubeconfig
+TEMP_FILE=$(mktemp)
+
+# Remove the certificate-authority-data, and replace the server with the master address
+awk '
+  BEGIN { in_cluster = 0 }
+  /^clusters:/ { in_cluster = 1 }
+  /^users:/ { in_cluster = 0 }
+  in_cluster && /^ *certificate-authority-data:/ { next }
+  in_cluster && /^ *server:/ {
+    print "    server: https://'${HEAD_NODE}:6443'"
+    print "    insecure-skip-tls-verify: true"
+    next
+  }
+  { print }
+' "$KUBECONFIG_FILE" > "$TEMP_FILE"
+
+# Replace the original kubeconfig with the modified one
+mv "$TEMP_FILE" "$KUBECONFIG_FILE"
+
+success_message "kubectl configured to connect to the cluster."
+
+echo "Cluster deployment completed. You can now run 'kubectl get nodes' to verify the setup."
+
+# Install GPU operator if a GPU was detected on any node
+if [ "$INSTALL_GPU" == "true" ]; then
+    echo -e "${YELLOW}GPU detected in the cluster. Installing Nvidia GPU Operator...${NC}"
+    run_remote "$HEAD_NODE" "
+        curl -fsSL -o get_helm.sh https://raw.githubusercontent.com/helm/helm/master/scripts/get-helm-3 &&
+        chmod 700 get_helm.sh &&
+        ./get_helm.sh &&
+        helm repo add nvidia https://helm.ngc.nvidia.com/nvidia && helm repo update &&
+        kubectl create namespace gpu-operator --kubeconfig ~/.kube/config || true &&
+        sudo ln -s /sbin/ldconfig /sbin/ldconfig.real || true &&
+        helm install gpu-operator -n gpu-operator --create-namespace nvidia/gpu-operator \
+        --set 'toolkit.env[0].name=CONTAINERD_CONFIG' \
+        --set 'toolkit.env[0].value=/var/lib/rancher/k3s/agent/etc/containerd/config.toml' \
+        --set 'toolkit.env[1].name=CONTAINERD_SOCKET' \
+        --set 'toolkit.env[1].value=/run/k3s/containerd/containerd.sock' \
+        --set 'toolkit.env[2].name=CONTAINERD_RUNTIME_CLASS' \
+        --set 'toolkit.env[2].value=nvidia' &&
+        echo 'Waiting for GPU operator installation...' &&
+        while ! kubectl describe nodes --kubeconfig ~/.kube/config | grep -q 'nvidia.com/gpu:'; do
+            echo 'Waiting for GPU operator...'
+            sleep 5
+        done
+        echo 'GPU operator installed successfully.'"
+    success_message "GPU Operator installed."
+else
+    echo -e "${YELLOW}No GPUs detected. Skipping GPU Operator installation.${NC}"
+fi
+
+# Configure SkyPilot
+progress_message "Configuring SkyPilot..."
+sky check kubernetes
+success_message "SkyPilot configured successfully."
+
+# Display final success message
+echo -e "${GREEN}==== 🎉 Kubernetes cluster deployment completed successfully 🎉 ====${NC}"
+echo "You can now interact with your Kubernetes cluster through SkyPilot: "
+echo "  • List available GPUs: sky show-gpus --cloud kubernetes"
+echo "  • Launch a GPU development pod: sky launch -c devbox --cloud kubernetes --gpus A100:1"
+echo "  • Connect to pod with SSH: ssh devbox"
+echo "  • Connect to pod with VSCode: code --remote ssh-remote+devbox '/'"