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

sky/provision/aws/config.py

@@ -8,7 +8,6 @@ _default_ec2_resource() to avoid version mismatch issues.
 # https://github.com/ray-project/ray/tree/ray-2.0.1/python/ray/autoscaler/_private/aws/config.py
 # Git commit of the release 2.0.1: 03b6bc7b5a305877501110ec04710a9c57011479
 import copy
-import functools
 import json
 import logging
 import time
@@ -16,10 +15,13 @@ from typing import Any, Dict, List, Optional, Set, Tuple
 
 import colorama
 
+from sky import exceptions
 from sky import sky_logging
 from sky.adaptors import aws
 from sky.provision import common
 from sky.provision.aws import utils
+from sky.utils import annotations
+from sky.utils import common_utils
 
 logger = sky_logging.init_logger(__name__)
 
@@ -40,8 +42,9 @@ def _skypilot_log_error_and_exit_for_failover(error: str) -> None:
     Mainly used for handling VPC/subnet errors before nodes are launched.
     """
     # NOTE: keep. The backend looks for this to know no nodes are launched.
-    prefix = 'SKYPILOT_ERROR_NO_NODES_LAUNCHED: '
-    raise RuntimeError(prefix + error)
+    full_error = f'SKYPILOT_ERROR_NO_NODES_LAUNCHED: {error}'
+    logger.error(full_error)
+    raise RuntimeError(full_error)
 
 
 def bootstrap_instances(
@@ -191,16 +194,56 @@ def _configure_iam_role(iam) -> Dict[str, Any]:
             for policy_arn in attach_policy_arns:
                 role.attach_policy(PolicyArn=policy_arn)
 
+            # SkyPilot: 'PassRole' is required by the controllers (jobs and
+            # services) created with `aws.remote_identity: SERVICE_ACCOUNT` to
+            # create instances with the IAM role.
+            skypilot_pass_role_policy_doc = {
+                'Statement': [
+                    {
+                        'Effect': 'Allow',
+                        'Action': [
+                            'iam:GetRole',
+                            'iam:PassRole',
+                        ],
+                        'Resource': role.arn,
+                    },
+                    {
+                        'Effect': 'Allow',
+                        'Action': 'iam:GetInstanceProfile',
+                        'Resource': profile.arn,
+                    },
+                ]
+            }
+            role.Policy('SkyPilotPassRolePolicy').put(
+                PolicyDocument=json.dumps(skypilot_pass_role_policy_doc))
+
         profile.add_role(RoleName=role.name)
         time.sleep(15)  # wait for propagation
     return {'Arn': profile.arn}
 
 
-@functools.lru_cache(maxsize=128)  # Keep bounded.
-def _get_route_tables(ec2, vpc_id: Optional[str], main: bool) -> List[Any]:
+@annotations.lru_cache(scope='request', maxsize=128)  # Keep bounded.
+def _get_route_tables(ec2, vpc_id: Optional[str], region: str,
+                      main: bool) -> List[Any]:
+    """Get route tables associated with a VPC and region
+
+    Args:
+        ec2: ec2 resource object
+        vpc_id: vpc_id is optional, if not provided, all route tables in the
+                region will be returned
+        region: region is mandatory to allow the lru cache
+                   to return the corect results
+        main: if True, only main route tables will be returned otherwise
+                only non-main route tables will be returned
+
+    Returns:
+        A list of route tables associated with the options VPC and region
+    """
     filters = [{'Name': 'association.main', 'Values': [str(main).lower()]}]
     if vpc_id is not None:
         filters.append({'Name': 'vpc-id', 'Values': [vpc_id]})
+    logger.debug(
+        f'Getting route tables with filters: {filters} in region: {region}')
     return ec2.meta.client.describe_route_tables(Filters=filters).get(
         'RouteTables', [])
 
@@ -213,7 +256,8 @@ def _is_subnet_public(ec2, subnet_id, vpc_id: Optional[str]) -> bool:
     https://docs.aws.amazon.com/vpc/latest/userguide/VPC_Internet_Gateway.html
     """
     # Get the route tables associated with the subnet
-    all_route_tables = _get_route_tables(ec2, vpc_id, main=False)
+    region = ec2.meta.client.meta.region_name
+    all_route_tables = _get_route_tables(ec2, vpc_id, region, main=False)
     route_tables = [
         rt for rt in all_route_tables
         # An RT can be associated with multiple subnets, i.e.,
@@ -235,14 +279,15 @@ def _is_subnet_public(ec2, subnet_id, vpc_id: Optional[str]) -> bool:
     logger.debug(f'subnet {subnet_id} route tables: {route_tables}')
     if _has_igw_route(route_tables):
         return True
-    if len(route_tables) > 0:
+    if route_tables:
         return False
 
     # Handle the case that a "main" route table is implicitly associated with
     # subnets. Since the associations are implicit, the filter above won't find
     # any. Check there exists a main route table with routes pointing to an IGW.
     logger.debug('Checking main route table')
-    main_route_tables = _get_route_tables(ec2, vpc_id, main=True)
+    region = ec2.meta.client.meta.region_name
+    main_route_tables = _get_route_tables(ec2, vpc_id, region, main=True)
     return _has_igw_route(main_route_tables)
 
 
@@ -338,10 +383,13 @@ def _usable_subnets(
         raise exc
 
     if not subnets:
+        vpc_msg = (f'Does a default VPC exist in region '
+                   f'{ec2.meta.client.meta.region_name}? ') if (
+                       vpc_id_of_sg is None) else ''
         _skypilot_log_error_and_exit_for_failover(
-            'No usable subnets found, try '
-            'manually creating an instance in your specified region to '
-            'populate the list of subnets and trying this again. '
+            f'No usable subnets found. {vpc_msg}'
+            'Try manually creating an instance in your specified region to '
+            'populate the list of subnets and try again. '
             'Note that the subnet must map public IPs '
             'on instance launch unless you set `use_internal_ips: true` in '
             'the `provider` config.')
@@ -409,7 +457,7 @@ def _vpc_id_from_security_group_ids(ec2, sg_ids: List[str]) -> Any:
 
     no_sg_msg = ('Failed to detect a security group with id equal to any of '
                  'the configured SecurityGroupIds.')
-    assert len(vpc_ids) > 0, no_sg_msg
+    assert vpc_ids, no_sg_msg
 
     return vpc_ids[0]
 
@@ -450,6 +498,11 @@ def _get_subnet_and_vpc_id(ec2, security_group_ids: Optional[List[str]],
         vpc_id_of_sg = None
 
     all_subnets = list(ec2.subnets.all())
+    # If no VPC is specified, use the default VPC.
+    # We filter only for default VPCs to avoid using subnets that users may
+    # not want SkyPilot to use.
+    if vpc_id_of_sg is None:
+        all_subnets = [s for s in all_subnets if s.vpc.is_default]
     subnets, vpc_id = _usable_subnets(
         ec2,
         user_specified_subnets=None,
@@ -500,47 +553,76 @@ def _configure_security_group(ec2, vpc_id: str, expected_sg_name: str,
 
 def _get_or_create_vpc_security_group(ec2, vpc_id: str,
                                       expected_sg_name: str) -> Any:
-    # Figure out which security groups with this name exist for each VPC...
-    vpc_to_existing_sg = {
-        sg.vpc_id: sg for sg in _get_security_groups_from_vpc_ids(
-            ec2,
-            [vpc_id],
-            [expected_sg_name],
-        )
-    }
+    """Find or create a security group in the specified VPC.
 
-    if vpc_id in vpc_to_existing_sg:
-        return vpc_to_existing_sg[vpc_id]
+    Args:
+        ec2: The initialized EC2 client object.
+        vpc_id: The ID of the VPC where the security group should be queried
+            or created.
+        expected_sg_name: The expected name of the security group.
 
-    # create a new security group
-    ec2.meta.client.create_security_group(
-        Description='Auto-created security group for Ray workers',
-        GroupName=expected_sg_name,
-        VpcId=vpc_id,
-    )
-    security_group = _get_security_groups_from_vpc_ids(ec2, [vpc_id],
-                                                       [expected_sg_name])
+    Returns:
+        The security group object containing the details of the security group.
 
-    assert security_group, 'Failed to create security group'
-    security_group = security_group[0]
+    Raises:
+        exceptions.NoClusterLaunchedError: If the security group creation fails
+            and is not due to an existing duplicate.
+        botocore.exceptions.ClientError: If the security group creation fails
+            due to AWS service issues.
+    """
+    # Figure out which security groups with this name exist for each VPC...
+    security_group = _get_security_group_from_vpc_id(ec2, vpc_id,
+                                                     expected_sg_name)
+    if security_group is not None:
+        return security_group
 
+    try:
+        # create a new security group
+        ec2.meta.client.create_security_group(
+            Description='Auto-created security group for Ray workers',
+            GroupName=expected_sg_name,
+            VpcId=vpc_id,
+        )
+    except ec2.meta.client.exceptions.ClientError as e:
+        if e.response['Error']['Code'] == 'InvalidGroup.Duplicate':
+            # The security group already exists, but we didn't see it
+            # because of eventual consistency.
+            logger.warning(f'{expected_sg_name} already exists when creating.')
+            security_group = _get_security_group_from_vpc_id(
+                ec2, vpc_id, expected_sg_name)
+            assert (security_group is not None and
+                    security_group.group_name == expected_sg_name), (
+                        f'Expected {expected_sg_name} but got {security_group}')
+            logger.info(
+                f'Found existing security group {colorama.Style.BRIGHT}'
+                f'{security_group.group_name}{colorama.Style.RESET_ALL} '
+                f'[id={security_group.id}]')
+            return security_group
+        message = ('Failed to create security group. Error: '
+                   f'{common_utils.format_exception(e)}')
+        logger.warning(message)
+        raise exceptions.NoClusterLaunchedError(message) from e
+
+    security_group = _get_security_group_from_vpc_id(ec2, vpc_id,
+                                                     expected_sg_name)
+    assert security_group is not None, 'Failed to create security group'
     logger.info(f'Created new security group {colorama.Style.BRIGHT}'
                 f'{security_group.group_name}{colorama.Style.RESET_ALL} '
                 f'[id={security_group.id}]')
     return security_group
 
 
-def _get_security_groups_from_vpc_ids(ec2, vpc_ids: List[str],
-                                      group_names: List[str]) -> List[Any]:
-    unique_vpc_ids = list(set(vpc_ids))
-    unique_group_names = set(group_names)
-
+def _get_security_group_from_vpc_id(ec2, vpc_id: str,
+                                    group_name: str) -> Optional[Any]:
+    """Get security group by VPC ID and group name."""
     existing_groups = list(
         ec2.security_groups.filter(Filters=[{
             'Name': 'vpc-id',
-            'Values': unique_vpc_ids
+            'Values': [vpc_id]
         }]))
-    filtered_groups = [
-        sg for sg in existing_groups if sg.group_name in unique_group_names
-    ]
-    return filtered_groups
+
+    for sg in existing_groups:
+        if sg.group_name == group_name:
+            return sg
+
+    return None

sky/provision/aws/instance.py

@@ -12,24 +12,21 @@ import time
 from typing import Any, Callable, Dict, List, Optional, Set, TypeVar
 
 from sky import sky_logging
-from sky import status_lib
 from sky.adaptors import aws
 from sky.clouds import aws as aws_cloud
+from sky.clouds.utils import aws_utils
 from sky.provision import common
+from sky.provision import constants
 from sky.provision.aws import utils
 from sky.utils import common_utils
 from sky.utils import resources_utils
+from sky.utils import status_lib
 from sky.utils import ux_utils
 
 logger = sky_logging.init_logger(__name__)
 
 _T = TypeVar('_T')
 
-# Tag uniquely identifying all nodes of a cluster
-TAG_RAY_CLUSTER_NAME = 'ray-cluster-name'
-TAG_SKYPILOT_CLUSTER_NAME = 'skypilot-cluster-name'
-TAG_RAY_NODE_KIND = 'ray-node-type'  # legacy tag for backward compatibility
-TAG_SKYPILOT_HEAD_NODE = 'skypilot-head-node'
 # Max retries for general AWS API calls.
 BOTO_MAX_RETRIES = 12
 # Max retries for creating an instance.
@@ -58,7 +55,7 @@ _RESUME_PER_INSTANCE_TIMEOUT = 120  # 2 minutes
 # https://aws.amazon.com/ec2/pricing/on-demand/#Data_Transfer_within_the_same_AWS_Region
 
 
-def _default_ec2_resource(region: str) -> Any:
+def _default_ec2_resource(region: str, check_credentials: bool = True) -> Any:
     if not hasattr(aws, 'version'):
         # For backward compatibility, reload the module if the aws module was
         # imported before and stale. Used for, e.g., a live jobs controller
@@ -98,12 +95,13 @@ def _default_ec2_resource(region: str) -> Any:
         importlib.reload(aws)
     return aws.resource('ec2',
                         region_name=region,
-                        max_attempts=BOTO_MAX_RETRIES)
+                        max_attempts=BOTO_MAX_RETRIES,
+                        check_credentials=check_credentials)
 
 
 def _cluster_name_filter(cluster_name_on_cloud: str) -> List[Dict[str, Any]]:
     return [{
-        'Name': f'tag:{TAG_RAY_CLUSTER_NAME}',
+        'Name': f'tag:{constants.TAG_RAY_CLUSTER_NAME}',
         'Values': [cluster_name_on_cloud],
     }]
 
@@ -181,8 +179,8 @@ def _create_instances(ec2_fail_fast, cluster_name: str,
                       count: int, associate_public_ip_address: bool) -> List:
     tags = {
         'Name': cluster_name,
-        TAG_RAY_CLUSTER_NAME: cluster_name,
-        TAG_SKYPILOT_CLUSTER_NAME: cluster_name,
+        constants.TAG_RAY_CLUSTER_NAME: cluster_name,
+        constants.TAG_SKYPILOT_CLUSTER_NAME: cluster_name,
         **tags
     }
     conf = node_config.copy()
@@ -212,6 +210,8 @@ def _create_instances(ec2_fail_fast, cluster_name: str,
     assert 'NetworkInterfaces' not in conf, conf
     assert security_group_ids is not None, conf
 
+    logger.debug(f'Creating {count} instances with config: \n{conf}')
+
     # NOTE: This ensures that we try ALL availability zones before
     # throwing an error.
     num_subnets = len(subnet_ids)
@@ -250,10 +250,8 @@ def _create_instances(ec2_fail_fast, cluster_name: str,
 
 def _get_head_instance_id(instances: List) -> Optional[str]:
     head_instance_id = None
-    head_node_markers = (
-        (TAG_SKYPILOT_HEAD_NODE, '1'),
-        (TAG_RAY_NODE_KIND, 'head'),  # backward compat with Ray
-    )
+    head_node_markers = tuple(constants.HEAD_NODE_TAGS.items())
+
     for inst in instances:
         for t in inst.tags:
             if (t['Key'], t['Value']) in head_node_markers:
@@ -288,7 +286,7 @@ def run_instances(region: str, cluster_name_on_cloud: str,
         'Name': 'instance-state-name',
         'Values': ['pending', 'running', 'stopping', 'stopped'],
     }, {
-        'Name': f'tag:{TAG_RAY_CLUSTER_NAME}',
+        'Name': f'tag:{constants.TAG_RAY_CLUSTER_NAME}',
         'Values': [cluster_name_on_cloud],
     }]
     exist_instances = list(ec2.instances.filter(Filters=filters))
@@ -314,31 +312,27 @@ def run_instances(region: str, cluster_name_on_cloud: str,
             raise RuntimeError(f'Impossible state "{state}".')
 
     def _create_node_tag(target_instance, is_head: bool = True) -> str:
+        node_type_tags = (constants.HEAD_NODE_TAGS
+                          if is_head else constants.WORKER_NODE_TAGS)
+        node_tag = [{'Key': k, 'Value': v} for k, v in node_type_tags.items()]
         if is_head:
-            node_tag = [{
-                'Key': TAG_SKYPILOT_HEAD_NODE,
-                'Value': '1'
-            }, {
-                'Key': TAG_RAY_NODE_KIND,
-                'Value': 'head'
-            }, {
+            node_tag.append({
                 'Key': 'Name',
                 'Value': f'sky-{cluster_name_on_cloud}-head'
-            }]
+            })
         else:
-            node_tag = [{
-                'Key': TAG_SKYPILOT_HEAD_NODE,
-                'Value': '0'
-            }, {
-                'Key': TAG_RAY_NODE_KIND,
-                'Value': 'worker'
-            }, {
+            node_tag.append({
                 'Key': 'Name',
                 'Value': f'sky-{cluster_name_on_cloud}-worker'
-            }]
+            })
+        # Remove AWS internal tags, as they are not allowed to be set by users.
+        target_instance_tags = [
+            tag for tag in target_instance.tags
+            if not tag['Key'].startswith('aws:')
+        ]
         ec2.meta.client.create_tags(
             Resources=[target_instance.id],
-            Tags=target_instance.tags + node_tag,
+            Tags=target_instance_tags + node_tag,
         )
         return target_instance.id
 
@@ -444,19 +438,87 @@ def run_instances(region: str, cluster_name_on_cloud: str,
             head_instance_id = _create_node_tag(resumed_instances[0])
 
     if to_start_count > 0:
+        target_reservation_names = (config.node_config.get(
+            'CapacityReservationSpecification',
+            {}).get('CapacityReservationTarget',
+                    {}).get('CapacityReservationId', []))
+        created_instances = []
+        if target_reservation_names:
+            node_config = copy.deepcopy(config.node_config)
+            # Clear the capacity reservation specification settings in the
+            # original node config, as we will create instances with
+            # reservations with specific settings for each reservation.
+            node_config['CapacityReservationSpecification'] = {
+                'CapacityReservationTarget': {}
+            }
+
+            reservations = aws_utils.list_reservations_for_instance_type(
+                node_config['InstanceType'], region=region)
+            # Filter the reservations by the user-specified ones, because
+            # reservations contain 'open' reservations as well, which do not
+            # need to explicitly specify in the config for creating instances.
+            target_reservations = []
+            for r in reservations:
+                if (r.targeted and r.name in target_reservation_names):
+                    target_reservations.append(r)
+            logger.debug(f'Reservations: {reservations}')
+            logger.debug(f'Target reservations: {target_reservations}')
+
+            target_reservations_list = sorted(
+                target_reservations,
+                key=lambda x: x.available_resources,
+                reverse=True)
+            for r in target_reservations_list:
+                if r.available_resources <= 0:
+                    # We have sorted the reservations by the available
+                    # resources, so if the reservation is not available, the
+                    # following reservations are not available either.
+                    break
+                reservation_count = min(r.available_resources, to_start_count)
+                logger.debug(f'Creating {reservation_count} instances '
+                             f'with reservation {r.name}')
+                node_config['CapacityReservationSpecification'][
+                    'CapacityReservationTarget'] = {
+                        'CapacityReservationId': r.name
+                    }
+                if r.type == aws_utils.ReservationType.BLOCK:
+                    # Capacity block reservations needs to specify the market
+                    # type during instance creation.
+                    node_config['InstanceMarketOptions'] = {
+                        'MarketType': aws_utils.ReservationType.BLOCK.value
+                    }
+                created_reserved_instances = _create_instances(
+                    ec2_fail_fast,
+                    cluster_name_on_cloud,
+                    node_config,
+                    tags,
+                    reservation_count,
+                    associate_public_ip_address=(
+                        not config.provider_config['use_internal_ips']))
+                created_instances.extend(created_reserved_instances)
+                to_start_count -= reservation_count
+                if to_start_count <= 0:
+                    break
+
         # TODO(suquark): If there are existing instances (already running or
         #  resumed), then we cannot guarantee that they will be in the same
         #  availability zone (when there are multiple zones specified).
         #  This is a known issue before.
 
-        created_instances = _create_instances(
-            ec2_fail_fast,
-            cluster_name_on_cloud,
-            config.node_config,
-            tags,
-            to_start_count,
-            associate_public_ip_address=(
-                not config.provider_config['use_internal_ips']))
+        if to_start_count > 0:
+            # Remove the capacity reservation specification from the node config
+            # as we have already created the instances with the reservations.
+            config.node_config.get('CapacityReservationSpecification',
+                                   {}).pop('CapacityReservationTarget', None)
+            created_remaining_instances = _create_instances(
+                ec2_fail_fast,
+                cluster_name_on_cloud,
+                config.node_config,
+                tags,
+                to_start_count,
+                associate_public_ip_address=(
+                    not config.provider_config['use_internal_ips']))
+            created_instances.extend(created_remaining_instances)
         created_instances.sort(key=lambda x: x.id)
 
         created_instance_ids = [n.id for n in created_instances]
@@ -563,7 +625,7 @@ def stop_instances(
     ]
     if worker_only:
         filters.append({
-            'Name': f'tag:{TAG_RAY_NODE_KIND}',
+            'Name': f'tag:{constants.TAG_RAY_NODE_KIND}',
             'Values': ['worker'],
         })
     instances = _filter_instances(ec2,
@@ -601,7 +663,7 @@ def terminate_instances(
     ]
     if worker_only:
         filters.append({
-            'Name': f'tag:{TAG_RAY_NODE_KIND}',
+            'Name': f'tag:{constants.TAG_RAY_NODE_KIND}',
             'Values': ['worker'],
         })
     # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/ec2.html#EC2.Instance
@@ -717,16 +779,31 @@ def open_ports(
 
     existing_ports: Set[int] = set()
     for existing_rule in sg.ip_permissions:
-        # Skip any non-tcp rules.
-        if existing_rule['IpProtocol'] != 'tcp':
+        # Skip any non-tcp rules or if all traffic (-1) is specified.
+        if existing_rule['IpProtocol'] not in ['tcp', '-1']:
             continue
         # Skip any rules that don't have a FromPort or ToPort.
-        if 'FromPort' not in existing_rule or 'ToPort' not in existing_rule:
-            continue
-        existing_ports.update(
-            range(existing_rule['FromPort'], existing_rule['ToPort'] + 1))
-    ports_to_open = resources_utils.port_set_to_ranges(
-        resources_utils.port_ranges_to_set(ports) - existing_ports)
+        if 'FromPort' in existing_rule and 'ToPort' in existing_rule:
+            existing_ports.update(
+                range(existing_rule['FromPort'], existing_rule['ToPort'] + 1))
+        elif existing_rule['IpProtocol'] == '-1':
+            # For AWS, IpProtocol = -1 means all traffic
+            for group_pairs in existing_rule['UserIdGroupPairs']:
+                if group_pairs['GroupId'] != sg.id:
+                    # We skip the port opening when the rule allows access from
+                    # other security groups, as that is likely added by a user
+                    # manually and satisfy their requirement.
+                    # The security group created by SkyPilot allows all traffic
+                    # from the same security group, which should not be skipped.
+                    existing_ports.add(-1)
+                    break
+            break
+
+    ports_to_open = []
+    # Do not need to open any ports when all traffic is already allowed.
+    if -1 not in existing_ports:
+        ports_to_open = resources_utils.port_set_to_ranges(
+            resources_utils.port_ranges_to_set(ports) - existing_ports)
 
     ip_permissions = []
     for port in ports_to_open:
@@ -799,7 +876,7 @@ def wait_instances(region: str, cluster_name_on_cloud: str,
 
     filters = [
         {
-            'Name': f'tag:{TAG_RAY_CLUSTER_NAME}',
+            'Name': f'tag:{constants.TAG_RAY_CLUSTER_NAME}',
             'Values': [cluster_name_on_cloud],
         },
     ]
@@ -850,7 +927,7 @@ def get_cluster_info(
             'Values': ['running'],
         },
         {
-            'Name': f'tag:{TAG_RAY_CLUSTER_NAME}',
+            'Name': f'tag:{constants.TAG_RAY_CLUSTER_NAME}',
             'Values': [cluster_name_on_cloud],
         },
     ]

sky/provision/azure/__init__.py

@@ -1,4 +1,11 @@
 """Azure provisioner for SkyPilot."""
 
+from sky.provision.azure.config import bootstrap_instances
 from sky.provision.azure.instance import cleanup_ports
+from sky.provision.azure.instance import get_cluster_info
 from sky.provision.azure.instance import open_ports
+from sky.provision.azure.instance import query_instances
+from sky.provision.azure.instance import run_instances
+from sky.provision.azure.instance import stop_instances
+from sky.provision.azure.instance import terminate_instances
+from sky.provision.azure.instance import wait_instances

sky/{skylet/providers → provision}/azure/azure-config-template.json

@@ -5,7 +5,7 @@
         "clusterId": {
             "type": "string",
             "metadata": {
-                "description": "Unique string appended to resource names to isolate resources from different ray clusters."
+                "description": "Unique string appended to resource names to isolate resources from different SkyPilot clusters."
             }
         },
         "subnet": {
@@ -13,17 +13,29 @@
             "metadata": {
                 "description": "Subnet parameters."
             }
+        },
+        "location": {
+            "type": "string",
+            "metadata": {
+                "description": "Location of where the resources are allocated."
+            }
+        },
+        "nsgName": {
+            "type": "string",
+            "metadata": {
+                "description": "Name of the Network Security Group associated with the SkyPilot cluster."
+            }
         }
     },
     "variables": {
         "contributor": "[subscriptionResourceId('Microsoft.Authorization/roleDefinitions', 'b24988ac-6180-42a0-ab88-20f7382dd24c')]",
-        "location": "[resourceGroup().location]",
-        "msiName": "[concat('ray-', parameters('clusterId'), '-msi')]",
-        "roleAssignmentName": "[concat('ray-', parameters('clusterId'), '-ra')]",
-        "nsgName": "[concat('ray-', parameters('clusterId'), '-nsg')]",
+        "location": "[parameters('location')]",
+        "msiName": "[concat('sky-', parameters('clusterId'), '-msi')]",
+        "roleAssignmentName": "[concat('sky-', parameters('clusterId'), '-ra')]",
+        "nsgName": "[parameters('nsgName')]",
         "nsg": "[resourceId('Microsoft.Network/networkSecurityGroups', variables('nsgName'))]",
-        "vnetName": "[concat('ray-', parameters('clusterId'), '-vnet')]",
-        "subnetName": "[concat('ray-', parameters('clusterId'), '-subnet')]"
+        "vnetName": "[concat('sky-', parameters('clusterId'), '-vnet')]",
+        "subnetName": "[concat('sky-', parameters('clusterId'), '-subnet')]"
     },
     "resources": [
        {