skypilot-nightly 1.0.0.dev20241227__py3-none-any.whl → 1.0.0.dev20250124__py3-none-any.whl

sky/provision/oci/query_utils.py

@@ -7,6 +7,8 @@ History:
    find_compartment: allow search subtree when find a compartment.
  - Hysun He (hysun.he@oracle.com) @ Nov.12, 2024: Add methods to
    Add/remove security rules: create_nsg_rules & remove_nsg
+ - Hysun He (hysun.he@oracle.com) @ Jan.01, 2025: Support reuse existing
+   VCN for SkyServe.
 """
 from datetime import datetime
 import functools
@@ -17,7 +19,6 @@ import traceback
 import typing
 from typing import List, Optional, Tuple
 
-from sky import exceptions
 from sky import sky_logging
 from sky.adaptors import common as adaptors_common
 from sky.adaptors import oci as oci_adaptor
@@ -496,23 +497,25 @@ class QueryHelper:
 
         compartment = cls.find_compartment(region)
 
-        list_vcns_resp = net_client.list_vcns(
-            compartment_id=compartment,
-            display_name=oci_utils.oci_config.VCN_NAME,
-            lifecycle_state='AVAILABLE',
-        )
+        vcn_id = oci_utils.oci_config.get_vcn_ocid(region)
+        if vcn_id is None:
+            list_vcns_resp = net_client.list_vcns(
+                compartment_id=compartment,
+                display_name=oci_utils.oci_config.VCN_NAME,
+                lifecycle_state='AVAILABLE',
+            )
 
-        if not list_vcns_resp:
-            raise exceptions.ResourcesUnavailableError(
-                'The VCN is not available')
+            # Get the primary vnic. The vnic might be an empty list for the
+            # corner case when the cluster was exited during provision.
+            if not list_vcns_resp.data:
+                return None
 
-        # Get the primary vnic.
-        assert len(list_vcns_resp.data) > 0
-        vcn = list_vcns_resp.data[0]
+            vcn = list_vcns_resp.data[0]
+            vcn_id = vcn.id
 
         list_nsg_resp = net_client.list_network_security_groups(
             compartment_id=compartment,
-            vcn_id=vcn.id,
+            vcn_id=vcn_id,
             limit=1,
             display_name=nsg_name,
         )
@@ -529,7 +532,7 @@ class QueryHelper:
             create_network_security_group_details=oci_adaptor.oci.core.models.
             CreateNetworkSecurityGroupDetails(
                 compartment_id=compartment,
-                vcn_id=vcn.id,
+                vcn_id=vcn_id,
                 display_name=nsg_name,
             ))
         get_nsg_resp = net_client.get_network_security_group(

sky/provision/provisioner.py

@@ -415,7 +415,6 @@ def _post_provision_setup(
         f'{json.dumps(dataclasses.asdict(provision_record), indent=2)}\n'
         'Cluster info:\n'
         f'{json.dumps(dataclasses.asdict(cluster_info), indent=2)}')
-
     head_instance = cluster_info.get_head_instance()
     if head_instance is None:
         e = RuntimeError(f'Provision failed for cluster {cluster_name!r}. '

sky/provision/runpod/instance.py

@@ -83,7 +83,8 @@ def run_instances(region: str, cluster_name_on_cloud: str,
         node_type = 'head' if head_instance_id is None else 'worker'
         try:
             instance_id = utils.launch(
-                name=f'{cluster_name_on_cloud}-{node_type}',
+                cluster_name=cluster_name_on_cloud,
+                node_type=node_type,
                 instance_type=config.node_config['InstanceType'],
                 region=region,
                 disk_size=config.node_config['DiskSize'],
@@ -92,6 +93,8 @@ def run_instances(region: str, cluster_name_on_cloud: str,
                 public_key=config.node_config['PublicKey'],
                 preemptible=config.node_config['Preemptible'],
                 bid_per_gpu=config.node_config['BidPerGPU'],
+                docker_login_config=config.provider_config.get(
+                    'docker_login_config'),
             )
         except Exception as e:  # pylint: disable=broad-except
             logger.warning(f'run_instances error: {e}')
@@ -145,6 +148,8 @@ def terminate_instances(
     """See sky/provision/__init__.py"""
     del provider_config  # unused
     instances = _filter_instances(cluster_name_on_cloud, None)
+    template_name, registry_auth_id = utils.get_registry_auth_resources(
+        cluster_name_on_cloud)
     for inst_id, inst in instances.items():
         logger.debug(f'Terminating instance {inst_id}: {inst}')
         if worker_only and inst['name'].endswith('-head'):
@@ -157,6 +162,10 @@ def terminate_instances(
                     f'Failed to terminate instance {inst_id}: '
                     f'{common_utils.format_exception(e, use_bracket=False)}'
                 ) from e
+    if template_name is not None:
+        utils.delete_pod_template(template_name)
+    if registry_auth_id is not None:
+        utils.delete_register_auth(registry_auth_id)
 
 
 def get_cluster_info(

sky/provision/runpod/utils.py

@@ -2,10 +2,11 @@
 
 import base64
 import time
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Tuple
 
 from sky import sky_logging
 from sky.adaptors import runpod
+from sky.provision import docker_utils
 import sky.provision.runpod.api.commands as runpod_commands
 from sky.skylet import constants
 from sky.utils import common_utils
@@ -47,6 +48,11 @@ GPU_NAME_MAP = {
 }
 
 
+def _construct_docker_login_template_name(cluster_name: str) -> str:
+    """Constructs the registry auth template name."""
+    return f'{cluster_name}-docker-login-template'
+
+
 def retry(func):
     """Decorator to retry a function."""
 
@@ -66,9 +72,83 @@ def retry(func):
     return wrapper
 
 
+# Adapted from runpod.api.queries.pods.py::QUERY_POD.
+# Adding containerRegistryAuthId to the query.
+_QUERY_POD = """
+query myPods {
+    myself {
+        pods {
+            id
+            containerDiskInGb
+            containerRegistryAuthId
+            costPerHr
+            desiredStatus
+            dockerArgs
+            dockerId
+            env
+            gpuCount
+            imageName
+            lastStatusChange
+            machineId
+            memoryInGb
+            name
+            podType
+            port
+            ports
+            uptimeSeconds
+            vcpuCount
+            volumeInGb
+            volumeMountPath
+            runtime {
+                ports{
+                    ip
+                    isIpPublic
+                    privatePort
+                    publicPort
+                    type
+                }
+            }
+            machine {
+                gpuDisplayName
+            }
+        }
+    }
+}
+"""
+
+
+def _sky_get_pods() -> dict:
+    """List all pods with extra registry auth information.
+
+    Adapted from runpod.get_pods() to include containerRegistryAuthId.
+    """
+    raw_return = runpod.runpod.api.graphql.run_graphql_query(_QUERY_POD)
+    cleaned_return = raw_return['data']['myself']['pods']
+    return cleaned_return
+
+
+_QUERY_POD_TEMPLATE_WITH_REGISTRY_AUTH = """
+query myself {
+    myself {
+        podTemplates {
+            name
+            containerRegistryAuthId
+        }
+    }
+}
+"""
+
+
+def _list_pod_templates_with_container_registry() -> dict:
+    """List all pod templates."""
+    raw_return = runpod.runpod.api.graphql.run_graphql_query(
+        _QUERY_POD_TEMPLATE_WITH_REGISTRY_AUTH)
+    return raw_return['data']['myself']['podTemplates']
+
+
 def list_instances() -> Dict[str, Dict[str, Any]]:
     """Lists instances associated with API key."""
-    instances = runpod.runpod.get_pods()
+    instances = _sky_get_pods()
 
     instance_dict: Dict[str, Dict[str, Any]] = {}
     for instance in instances:
@@ -100,14 +180,75 @@ def list_instances() -> Dict[str, Dict[str, Any]]:
     return instance_dict
 
 
-def launch(name: str, instance_type: str, region: str, disk_size: int,
-           image_name: str, ports: Optional[List[int]], public_key: str,
-           preemptible: Optional[bool], bid_per_gpu: float) -> str:
+def delete_pod_template(template_name: str) -> None:
+    """Deletes a pod template."""
+    try:
+        runpod.runpod.api.graphql.run_graphql_query(
+            f'mutation {{deleteTemplate(templateName: "{template_name}")}}')
+    except runpod.runpod.error.QueryError as e:
+        logger.warning(f'Failed to delete template {template_name}: {e}'
+                       'Please delete it manually.')
+
+
+def delete_register_auth(registry_auth_id: str) -> None:
+    """Deletes a registry auth."""
+    try:
+        runpod.runpod.delete_container_registry_auth(registry_auth_id)
+    except runpod.runpod.error.QueryError as e:
+        logger.warning(f'Failed to delete registry auth {registry_auth_id}: {e}'
+                       'Please delete it manually.')
+
+
+def _create_template_for_docker_login(
+    cluster_name: str,
+    image_name: str,
+    docker_login_config: Optional[Dict[str, str]],
+) -> Tuple[str, Optional[str]]:
+    """Creates a template for the given image with the docker login config.
+
+    Returns:
+        formatted_image_name: The formatted image name.
+        template_id: The template ID. None for no docker login config.
+    """
+    if docker_login_config is None:
+        return image_name, None
+    login_config = docker_utils.DockerLoginConfig(**docker_login_config)
+    container_registry_auth_name = f'{cluster_name}-registry-auth'
+    container_template_name = _construct_docker_login_template_name(
+        cluster_name)
+    # The `name` argument is only for display purpose and the registry server
+    # will be splitted from the docker image name (Tested with AWS ECR).
+    # Here we only need the username and password to create the registry auth.
+    # TODO(tian): Now we create a template and a registry auth for each cluster.
+    # Consider create one for each server and reuse them. Challenges including
+    # calculate the reference count and delete them when no longer needed.
+    create_auth_resp = runpod.runpod.create_container_registry_auth(
+        name=container_registry_auth_name,
+        username=login_config.username,
+        password=login_config.password,
+    )
+    registry_auth_id = create_auth_resp['id']
+    create_template_resp = runpod.runpod.create_template(
+        name=container_template_name,
+        image_name=None,
+        registry_auth_id=registry_auth_id,
+    )
+    return login_config.format_image(image_name), create_template_resp['id']
+
+
+def launch(cluster_name: str, node_type: str, instance_type: str, region: str,
+           disk_size: int, image_name: str, ports: Optional[List[int]],
+           public_key: str, preemptible: Optional[bool], bid_per_gpu: float,
+           docker_login_config: Optional[Dict[str, str]]) -> str:
     """Launches an instance with the given parameters.
 
     Converts the instance_type to the RunPod GPU name, finds the specs for the
     GPU, and launches the instance.
+
+    Returns:
+        instance_id: The instance ID.
     """
+    name = f'{cluster_name}-{node_type}'
     gpu_type = GPU_NAME_MAP[instance_type.split('_')[1]]
     gpu_quantity = int(instance_type.split('_')[0].replace('x', ''))
     cloud_type = instance_type.split('_')[2]
@@ -139,21 +280,24 @@ def launch(name: str, instance_type: str, region: str, disk_size: int,
     # Use base64 to deal with the tricky quoting issues caused by runpod API.
     encoded = base64.b64encode(setup_cmd.encode('utf-8')).decode('utf-8')
 
+    docker_args = (f'bash -c \'echo {encoded} | base64 --decode > init.sh; '
+                   f'bash init.sh\'')
+
     # Port 8081 is occupied for nginx in the base image.
     custom_ports_str = ''
     if ports is not None:
         custom_ports_str = ''.join([f'{p}/tcp,' for p in ports])
+    ports_str = (f'22/tcp,'
+                 f'{custom_ports_str}'
+                 f'{constants.SKY_REMOTE_RAY_DASHBOARD_PORT}/http,'
+                 f'{constants.SKY_REMOTE_RAY_PORT}/http')
 
-    docker_args = (f'bash -c \'echo {encoded} | base64 --decode > init.sh; '
-                   f'bash init.sh\'')
-    ports = (f'22/tcp,'
-             f'{custom_ports_str}'
-             f'{constants.SKY_REMOTE_RAY_DASHBOARD_PORT}/http,'
-             f'{constants.SKY_REMOTE_RAY_PORT}/http')
+    image_name_formatted, template_id = _create_template_for_docker_login(
+        cluster_name, image_name, docker_login_config)
 
     params = {
         'name': name,
-        'image_name': image_name,
+        'image_name': image_name_formatted,
         'gpu_type_id': gpu_type,
         'cloud_type': cloud_type,
         'container_disk_in_gb': disk_size,
@@ -161,9 +305,10 @@ def launch(name: str, instance_type: str, region: str, disk_size: int,
         'min_memory_in_gb': gpu_specs['memoryInGb'] * gpu_quantity,
         'gpu_count': gpu_quantity,
         'country_code': region,
-        'ports': ports,
+        'ports': ports_str,
         'support_public_ip': True,
         'docker_args': docker_args,
+        'template_id': template_id,
     }
 
     if preemptible is None or not preemptible:
@@ -177,6 +322,18 @@ def launch(name: str, instance_type: str, region: str, disk_size: int,
     return new_instance['id']
 
 
+def get_registry_auth_resources(
+        cluster_name: str) -> Tuple[Optional[str], Optional[str]]:
+    """Gets the registry auth resources."""
+    container_registry_auth_name = _construct_docker_login_template_name(
+        cluster_name)
+    for template in _list_pod_templates_with_container_registry():
+        if template['name'] == container_registry_auth_name:
+            return container_registry_auth_name, template[
+                'containerRegistryAuthId']
+    return None, None
+
+
 def remove(instance_id: str) -> None:
     """Terminates the given instance."""
     runpod.runpod.terminate_pod(instance_id)

sky/resources.py

@@ -540,7 +540,7 @@ class Resources:
         if memory_gb <= 0:
             with ux_utils.print_exception_no_traceback():
                 raise ValueError(
-                    f'The "cpus" field should be positive. Found: {memory!r}')
+                    f'The "memory" field should be positive. Found: {memory!r}')
 
     def _set_accelerators(
         self,