[New Provisioner] New provisioner for k8s (#3019)

* Adopt new provisioner for k8s * format * rename names * move file template * resolve circular import * fixes * fix port * Fix IP to use * fix * format * align ray j2 * rename * format * Merge branch 'master' of github.com:skypilot-org/skypilot into new-provisioner-k8s * Support custom image * format * remove ray dependency for k8s support * chown * typo * make sure cpus are respected by ray cluster * fix tests for k8s * add sleep time * more robust initialization * Add pod_config_override * rename to pod_config * add schema * Fix waiting and initialize * backward compat for ssh_user field in handle * format * Fix svc * format * wait previous terminating pod * add more logging * format * format * Address comments * omit region name in message str for k8s * Fix ssh user * Fix the job owner * remove unecessary output * docs update * Add custom image tests * fix * lint * Add note for pvt repos * address comments * Update docs/source/reference/kubernetes/index.rst Co-authored-by: Zhanghao Wu <[email protected]> * Add pod_config docs * update docs * remove svc related code * fix * Fix stop * remove python dependency and uncomment ssh test * add kubernetes handler for surfacing provisioning errors * lint * Fix cyclic imports * update docker images * newline * Update sky/backends/cloud_vm_ray_backend.py * format * Move output to instance.py to keep align the UX with other clouds. * swap, no effect * Allowing image_id without `docker:` * format --------- Co-authored-by: Romil Bhardwaj <[email protected]> Co-authored-by: Romil Bhardwaj <[email protected]>
skypilot-org · Jan 30, 2024 · 9078de2 · 9078de2
1 parent 8bbd030
commit 9078de2
Show file tree

Hide file tree

Showing 42 changed files with 1,553 additions and 334 deletions.
diff --git a/Dockerfile_k8s b/Dockerfile_k8s
@@ -1,4 +1,4 @@
-FROM continuumio/miniconda3:22.11.1
+FROM continuumio/miniconda3:23.3.1-0
 
 # TODO(romilb): Investigate if this image can be consolidated with the skypilot
 #  client image (`Dockerfile`)
@@ -25,26 +25,21 @@ RUN useradd -m -s /bin/bash sky && \
 # Switch to sky user
 USER sky
 
-# Install SkyPilot pip dependencies
+# Install SkyPilot pip dependencies preemptively to speed up provisioning time
 RUN pip install wheel Click colorama cryptography jinja2 jsonschema && \
     pip install networkx oauth2client pandas pendulum PrettyTable && \
     pip install ray[default]==2.4.0 rich tabulate filelock && \
     pip install packaging 'protobuf<4.0.0' pulp && \
-    pip install awscli boto3 pycryptodome==3.12.0 && \
-    pip install docker kubernetes
+    pip install pycryptodome==3.12.0 && \
+    pip install docker kubernetes==28.1.0
 
 # Add /home/sky/.local/bin/ to PATH
 RUN echo 'export PATH="$PATH:$HOME/.local/bin"' >> ~/.bashrc
 
-# Install SkyPilot. This is purposely separate from installing SkyPilot
-# dependencies to optimize rebuild time
+# Copy SkyPilot code base. This is required for the ssh jump pod to find the
+# lifecycle management scripts
 COPY --chown=sky . /skypilot/sky/
 
-# TODO(romilb): Installing SkyPilot may not be necessary since ray up will do it
-RUN cd /skypilot/ && \
-    sudo mv -v sky/setup_files/* . && \
-    pip install ".[aws]"
-
 # Set PYTHONUNBUFFERED=1 to have Python print to stdout/stderr immediately
 ENV PYTHONUNBUFFERED=1
 

diff --git a/Dockerfile_k8s_gpu b/Dockerfile_k8s_gpu
@@ -38,21 +38,16 @@ RUN pip install wheel Click colorama cryptography jinja2 jsonschema && \
     pip install networkx oauth2client pandas pendulum PrettyTable && \
     pip install rich tabulate filelock && \
     pip install packaging 'protobuf<4.0.0' pulp && \
-    pip install awscli boto3 pycryptodome==3.12.0 && \
-    pip install docker kubernetes
+    pip install pycryptodome==3.12.0 && \
+    pip install docker kubernetes==28.1.0
 
 # Add /home/sky/.local/bin/ to PATH
 RUN echo 'export PATH="$PATH:$HOME/.local/bin"' >> ~/.bashrc
 
-# Install SkyPilot. This is purposely separate from installing SkyPilot
-# dependencies to optimize rebuild time
+# Copy SkyPilot code base. This is required for the ssh jump pod to find the
+# lifecycle management scripts
 COPY --chown=sky . /skypilot/sky/
 
-# TODO(romilb): Installing SkyPilot may not be necessary since ray up will do it
-RUN cd /skypilot/ && \
-    sudo mv -v sky/setup_files/* . && \
-    pip install ".[aws]"
-
 # Set PYTHONUNBUFFERED=1 to have Python print to stdout/stderr immediately
 ENV PYTHONUNBUFFERED=1
 

diff --git a/docs/source/reference/config.rst b/docs/source/reference/config.rst
@@ -173,6 +173,25 @@ Available fields and semantics:
     # is used as default if 'networking' is not specified.
     networking: portforward
 
+    # Additional fields to override the pod fields used by SkyPilot (optional)
+    #
+    # Any key:value pairs added here would get added to the pod spec used to
+    # create SkyPilot pods. The schema follows the same schema for a Pod object
+    # in the Kubernetes API:
+    # https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.26/#pod-v1-core
+    #
+    # Example use cases: adding custom labels to SkyPilot pods, specifying
+    # imagePullSecrets for pulling images from private registries, overriding
+    # the default runtimeClassName etc.
+    pod_config:
+      metadata:
+        labels:
+          my-label: my-value
+      spec:
+        runtimeClassName: nvidia
+        imagePullSecrets:
+          - name: my-secret
+
   # Advanced OCI configurations (optional).
   oci:
     # A dict mapping region names to region-specific configurations, or

diff --git a/docs/source/reference/kubernetes/index.rst b/docs/source/reference/kubernetes/index.rst
@@ -22,7 +22,7 @@ tasks can be submitted to your Kubernetes cluster just like any other cloud prov
 **Supported Kubernetes deployments:**
 
 * Hosted Kubernetes services (EKS, GKE)
-* On-prem clusters (Kubeadm, K3s, Rancher)
+* On-prem clusters (Kubeadm, Rancher)
 * Local development clusters (KinD, minikube)
 
 
@@ -123,17 +123,52 @@ Once your cluster administrator has :ref:`setup a Kubernetes cluster <kubernetes
     $ # Set a specific namespace to be used in the current-context
     $ kubectl config set-context --current --namespace=mynamespace
 
+
+Using Custom Images
+-------------------
+By default, we use and maintain a SkyPilot container image that has conda and a few other basic tools installed.
+
+To use your own image, add :code:`image_id: docker:<your image tag>` to the :code:`resources` section of your task YAML.
+
+.. code-block:: yaml
+
+    resources:
+      image_id: docker:myrepo/myimage:latest
+    ...
+
+Your image must satisfy the following requirements:
+
+* Image must be **debian-based** and must have the apt package manager installed.
+* The default user in the image must have root privileges or passwordless sudo access.
+
+.. note::
+
+    If your cluster runs on non-x86_64 architecture (e.g., Apple Silicon), your image must be built natively for that architecture. Otherwise, your job may get stuck at :code:`Start streaming logs ...`. See `GitHub issue <https://github.com/skypilot-org/skypilot/issues/3035>`_ for more.
+
+Using Images from Private Repositories
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+To use images from private repositories (e.g., Private DockerHub, Amazon ECR, Google Container Registry), create a `secret <https://kubernetes.io/docs/tasks/configure-pod-container/pull-image-private-registry/#create-a-secret-by-providing-credentials-on-the-command-line>`_ in your Kubernetes cluster and edit your :code:`~/.sky/config` to specify the secret like so:
+
+.. code-block:: yaml
+
+    kubernetes:
+      pod_config:
+        spec:
+          imagePullSecrets:
+            - name: your-secret-here
+
+.. tip::
+
+    If you use Amazon ECR, your secret credentials may expire every 12 hours. Consider using `k8s-ecr-login-renew <https://github.com/nabsul/k8s-ecr-login-renew>`_ to automatically refresh your secrets.
+
+
 FAQs
 ----
 
 * **Are autoscaling Kubernetes clusters supported?**
 
   To run on an autoscaling cluster, you may need to adjust the resource provisioning timeout (:code:`Kubernetes.TIMEOUT` in `clouds/kubernetes.py`) to a large value to give enough time for the cluster to autoscale. We are working on a better interface to adjust this timeout - stay tuned!
 
-* **What container image is used for tasks? Can I specify my own image?**
-
-  We use and maintain a SkyPilot container image that has conda and a few other basic tools installed. You can specify a custom image to use in `clouds/kubernetes.py`, but it must have rsync, conda and OpenSSH server installed. We are working on a interface to allow specifying custom images through the :code:`image_id` field in the task YAML - stay tuned!
-
 * **Can SkyPilot provision a Kubernetes cluster for me? Will SkyPilot add more nodes to my Kubernetes clusters?**
 
   The goal of Kubernetes support is to run SkyPilot tasks on an existing Kubernetes cluster. It does not provision any new Kubernetes clusters or add new nodes to an existing Kubernetes cluster.
@@ -151,8 +186,8 @@ Kubernetes support is under active development. Some features are in progress an
 * Auto-down - ✅ Available
 * Storage mounting - ✅ Available on x86_64 clusters
 * Multi-node tasks - ✅ Available
+* Custom images - ✅ Available
 * Opening ports and exposing services - 🚧 In progress
-* Custom images - 🚧 In progress
 * Multiple Kubernetes Clusters - 🚧 In progress
 
 

diff --git a/sky/authentication.py b/sky/authentication.py
@@ -43,9 +43,9 @@
 from sky.adaptors import ibm
 from sky.adaptors import runpod
 from sky.clouds.utils import lambda_utils
+from sky.provision.kubernetes import utils as kubernetes_utils
 from sky.utils import common_utils
 from sky.utils import kubernetes_enums
-from sky.utils import kubernetes_utils
 from sky.utils import subprocess_utils
 from sky.utils import ux_utils
 

diff --git a/sky/backends/backend_utils.py b/sky/backends/backend_utils.py
@@ -16,7 +16,6 @@
 
 import colorama
 import filelock
-import jinja2
 from packaging import version
 import requests
 from requests import adapters
@@ -33,14 +32,14 @@
 from sky import exceptions
 from sky import global_user_state
 from sky import provision as provision_lib
-from sky import serve as serve_lib
 from sky import sky_logging
 from sky import skypilot_config
 from sky import status_lib
 from sky.backends import onprem_utils
 from sky.clouds import cloud_registry
 from sky.clouds.utils import gcp_utils
 from sky.provision import instance_setup
+from sky.provision.kubernetes import utils as kubernetes_utils
 from sky.skylet import constants
 from sky.usage import usage_lib
 from sky.utils import cluster_yaml_utils
@@ -167,25 +166,6 @@ def _get_yaml_path_from_cluster_name(cluster_name: str,
     return str(output_path)
 
 
-def fill_template(template_name: str, variables: Dict,
-                  output_path: str) -> None:
-    """Create a file from a Jinja template and return the filename."""
-    assert template_name.endswith('.j2'), template_name
-    template_path = os.path.join(sky.__root_dir__, 'templates', template_name)
-    if not os.path.exists(template_path):
-        raise FileNotFoundError(f'Template "{template_name}" does not exist.')
-    with open(template_path) as fin:
-        template = fin.read()
-    output_path = os.path.abspath(os.path.expanduser(output_path))
-    os.makedirs(os.path.dirname(output_path), exist_ok=True)
-
-    # Write out yaml config.
-    j2_template = jinja2.Template(template)
-    content = j2_template.render(**variables)
-    with open(output_path, 'w') as fout:
-        fout.write(content)
-
-
 def _optimize_file_mounts(yaml_path: str) -> None:
     """Optimize file mounts in the given ray yaml file.
 
@@ -458,6 +438,7 @@ def add_cluster(
         auth_config: Dict[str, str],
         ports: List[int],
         docker_user: Optional[str] = None,
+        ssh_user: Optional[str] = None,
     ):
         """Add authentication information for cluster to local SSH config file.
 
@@ -476,8 +457,12 @@ def add_cluster(
             auth_config: read_yaml(handle.cluster_yaml)['auth']
             ports: List of port numbers for SSH corresponding to ips
             docker_user: If not None, use this user to ssh into the docker
+            ssh_user: Override the ssh_user in auth_config
         """
-        username = auth_config['ssh_user']
+        if ssh_user is None:
+            username = auth_config['ssh_user']
+        else:
+            username = ssh_user
         if docker_user is not None:
             username = docker_user
         key_path = os.path.expanduser(auth_config['ssh_private_key'])
@@ -875,7 +860,7 @@ def write_cluster_config(
     # Use a tmp file path to avoid incomplete YAML file being re-used in the
     # future.
     tmp_yaml_path = yaml_path + '.tmp'
-    fill_template(
+    common_utils.fill_template(
         cluster_config_template,
         dict(
             resources_vars,
@@ -949,6 +934,10 @@ def write_cluster_config(
         return config_dict
     _add_auth_to_cluster_config(cloud, tmp_yaml_path)
 
+    # Add kubernetes config fields from ~/.sky/config
+    if isinstance(cloud, clouds.Kubernetes):
+        kubernetes_utils.combine_pod_config_fields(tmp_yaml_path)
+
     # Restore the old yaml content for backward compatibility.
     if os.path.exists(yaml_path) and keep_launch_fields_in_existing_config:
         with open(yaml_path, 'r') as f:
@@ -1198,13 +1187,23 @@ def wait_until_ray_cluster_ready(
     return True, docker_user  # success
 
 
-def ssh_credential_from_yaml(cluster_yaml: str,
-                             docker_user: Optional[str] = None
-                            ) -> Dict[str, Any]:
-    """Returns ssh_user, ssh_private_key and ssh_control name."""
+def ssh_credential_from_yaml(
+    cluster_yaml: str,
+    docker_user: Optional[str] = None,
+    ssh_user: Optional[str] = None,
+) -> Dict[str, Any]:
+    """Returns ssh_user, ssh_private_key and ssh_control name.
+
+    Args:
+        cluster_yaml: path to the cluster yaml.
+        docker_user: when using custom docker image, use this user to ssh into
+            the docker container.
+        ssh_user: override the ssh_user in the cluster yaml.
+    """
     config = common_utils.read_yaml(cluster_yaml)
     auth_section = config['auth']
-    ssh_user = auth_section['ssh_user'].strip()
+    if ssh_user is None:
+        ssh_user = auth_section['ssh_user'].strip()
     ssh_private_key = auth_section.get('ssh_private_key')
     ssh_control_name = config.get('cluster_name', '__default__')
     ssh_proxy_command = auth_section.get('ssh_proxy_command')
@@ -1396,7 +1395,7 @@ def get_node_ips(
 
     if cloud.PROVISIONER_VERSION >= clouds.ProvisionerVersion.SKYPILOT:
         metadata = provision_lib.get_cluster_info(
-            provider_name, ray_config['provider']['region'],
+            provider_name, ray_config['provider'].get('region'),
             ray_config['cluster_name'], ray_config['provider'])
         if len(metadata.instances) < expected_num_nodes:
             # Simulate the exception when Ray head node is not up.
@@ -1786,7 +1785,8 @@ def run_ray_status_to_check_ray_cluster_healthy() -> bool:
 
             # Check if ray cluster status is healthy.
             ssh_credentials = ssh_credential_from_yaml(handle.cluster_yaml,
-                                                       handle.docker_user)
+                                                       handle.docker_user,
+                                                       handle.ssh_user)
 
             runner = command_runner.SSHCommandRunner(external_ips[0],
                                                      **ssh_credentials,
@@ -2515,7 +2515,7 @@ def get_task_demands_dict(task: 'task_lib.Task') -> Dict[str, float]:
     # For sky serve controller task, we set the CPU resource to a smaller
     # value to support a larger number of services.
     resources_dict = {
-        'CPU': (serve_lib.SERVICES_TASK_CPU_DEMAND
+        'CPU': (constants.SERVICES_TASK_CPU_DEMAND
                 if task.service_name is not None else DEFAULT_TASK_CPU_DEMAND)
     }
     if task.best_resources is not None:
@@ -2536,7 +2536,7 @@ def get_task_resources_str(task: 'task_lib.Task') -> str:
     The resources string is only used as a display purpose, so we only show
     the accelerator demands (if any). Otherwise, the CPU demand is shown.
     """
-    task_cpu_demand = (serve_lib.SERVICES_TASK_CPU_DEMAND if task.service_name
+    task_cpu_demand = (constants.SERVICES_TASK_CPU_DEMAND if task.service_name
                        is not None else DEFAULT_TASK_CPU_DEMAND)
     if task.best_resources is not None:
         accelerator_dict = task.best_resources.accelerators