Add GCP VPC creation if no VPC available (skypilot-org#1594)

infwinston · web-flow · commit b0894c803ee0 · 2023-01-24T12:50:07.000-08:00
* add vpc creation

* shorten

* create vpcnet in ray up instead

* implement best effort

* fix

* add constants.py

* fix

* quote..

* comments

* comment

* fix

* backward compatibility

* comment

* list_instances

* fix

* fix
diff --git a/sky/backends/backend_utils.py b/sky/backends/backend_utils.py
@@ -911,13 +911,19 @@ def write_cluster_config(
             tpu_name = cluster_name
 
         user_file_dir = os.path.expanduser(f'{SKY_USER_FILE_PATH}/')
+
+        from sky.skylet.providers.gcp import config as gcp_config  # pylint: disable=import-outside-toplevel
+        config = common_utils.read_yaml(os.path.expanduser(config_dict['ray']))
+        vpc_name = gcp_config.get_usable_vpc(config)
+
         scripts = tuple(
             fill_template(
                 template_name,
                 dict(
                     resources_vars, **{
                         'tpu_name': tpu_name,
                         'gcp_project_id': gcp_project_id,
+                        'vpc_name': vpc_name,
                     }),
                 # Use new names for TPU scripts so that different runs can use
                 # different TPUs.  Put in SKY_USER_FILE_PATH to be consistent
diff --git a/sky/skylet/providers/gcp/config.py b/sky/skylet/providers/gcp/config.py
@@ -12,7 +12,8 @@
 from google.oauth2.credentials import Credentials as OAuthCredentials
 from googleapiclient import discovery, errors
 
-from sky.skylet.providers.gcp.node import MAX_POLLS, POLL_INTERVAL, GCPNodeType
+from sky.skylet.providers.gcp.node import MAX_POLLS, POLL_INTERVAL, GCPNodeType, GCPCompute
+from sky.skylet.providers.gcp.constants import SKYPILOT_VPC_NAME, VPC_TEMPLATE, FIREWALL_RULES_TEMPLATE
 from ray.autoscaler._private.util import check_legacy_fields
 
 logger = logging.getLogger(__name__)
@@ -46,7 +47,6 @@
 # NOTE: iam.serviceAccountUser allows the Head Node to create worker nodes
 # with ServiceAccounts.
 
-
 def get_node_type(node: dict) -> GCPNodeType:
     """Returns node type based on the keys in ``node``.
 
@@ -488,6 +488,96 @@ def _configure_key_pair(config, compute):
     return config
 
 
+def _check_firewall_rules(vpc_name, config, compute):
+    """Check if the firewall rules in the VPC are sufficient."""
+    required_rules = FIREWALL_RULES_TEMPLATE.copy()
+
+    operation = (
+        compute.networks().
+        getEffectiveFirewalls(
+            project=config["provider"]["project_id"],
+            network=vpc_name
+        )
+    )
+    response = operation.execute()
+    if len(response) == 0:
+        return False
+    effective_rules = response["firewalls"]
+
+    def _get_refined_rule(rule):
+        KEY_TO_COMPARE = {"sourceRanges", "allowed", "direction"}
+        refined_rule = {}
+        for k in KEY_TO_COMPARE:
+            if k not in rule:
+                continue
+            if k == "allowed":
+                refined_rule[k] = sorted(rule[k], key=lambda x: x["IPProtocol"])
+            else:
+                refined_rule[k] = rule[k]
+        return refined_rule
+
+    required_rules = list(map(_get_refined_rule, required_rules))
+    effective_rules = list(map(_get_refined_rule, effective_rules))
+
+    for rule in required_rules:
+        if rule not in effective_rules:
+            return False
+    return True
+
+
+def get_usable_vpc(config):
+    """Return a usable VPC.
+
+    If not found, create a new one with sufficient firewall rules.
+    """
+    _, _, compute, _ = construct_clients_from_provider_config(config["provider"])
+
+    # For backward compatibility, reuse the VPC if the VM is launched.
+    resource = GCPCompute(
+        compute,
+        config["provider"]["project_id"],
+        config["provider"]["availability_zone"],
+        config["cluster_name"],
+    )
+    node = resource._list_instances(label_filters=None, status_filter=None)
+    if len(node) > 0:
+        netInterfaces = node[0].get("networkInterfaces", [])
+        if len(netInterfaces) > 0:
+            vpc_name = netInterfaces[0]["network"].split("/")[-1]
+            return vpc_name
+
+    vpcnets_all = _list_vpcnets(config, compute)
+
+    usable_vpc_name = None
+    for vpc in vpcnets_all:
+        if _check_firewall_rules(vpc["name"], config, compute):
+            usable_vpc_name = vpc["name"]
+            break
+
+    if usable_vpc_name is None:
+        logger.info(f"Creating a default VPC network, {SKYPILOT_VPC_NAME}...")
+
+        # Create a default VPC network
+        proj_id = config["provider"]["project_id"]
+        body = VPC_TEMPLATE.copy()
+        body["name"] = body["name"].format(VPC_NAME=SKYPILOT_VPC_NAME)
+        body["selfLink"] = body["selfLink"].format(PROJ_ID=proj_id, VPC_NAME=SKYPILOT_VPC_NAME)
+        _create_vpcnet(config, compute, body)
+
+        # Create firewall rules
+        for rule in FIREWALL_RULES_TEMPLATE:
+            body = rule.copy()
+            body["name"] = body["name"].format(VPC_NAME=SKYPILOT_VPC_NAME)
+            body["network"] = body["network"].format(PROJ_ID=proj_id, VPC_NAME=SKYPILOT_VPC_NAME)
+            body["selfLink"] = body["selfLink"].format(PROJ_ID=proj_id, VPC_NAME=SKYPILOT_VPC_NAME)
+            _create_firewall_rule(config, compute, body)
+
+        usable_vpc_name = SKYPILOT_VPC_NAME
+        logger.info(f"A VPC network {SKYPILOT_VPC_NAME} created.")
+
+    return usable_vpc_name
+
+
 def _configure_subnet(config, compute):
     """Pick a reasonable subnet if not specified by the config."""
     config = copy.deepcopy(config)
@@ -506,14 +596,9 @@ def _configure_subnet(config, compute):
     ):
         return config
 
-    subnets = _list_subnets(config, compute)
-
-    if not subnets:
-        raise NotImplementedError("Should be able to create subnet.")
-
-    # TODO: make sure that we have usable subnet. Maybe call
-    # compute.subnetworks().listUsable? For some reason it didn't
-    # work out-of-the-box
+    # SkyPilot: make sure there's a usable VPC
+    usable_vpc_name = get_usable_vpc(config)
+    subnets = _list_subnets(config, compute, filter=f"(name=\"{usable_vpc_name}\")")
     default_subnet = subnets[0]
 
     default_interfaces = [
@@ -542,17 +627,54 @@ def _configure_subnet(config, compute):
     return config
 
 
-def _list_subnets(config, compute):
+def _create_firewall_rule(config, compute, body):
+    operation = (
+        compute.firewalls()
+        .insert(project=config["provider"]["project_id"], body=body)
+        .execute()
+    )
+    response = wait_for_compute_global_operation(
+        config["provider"]["project_id"], operation, compute
+    )
+    return response
+
+
+def _create_vpcnet(config, compute, body):
+    operation = (
+        compute.networks()
+        .insert(project=config["provider"]["project_id"], body=body)
+        .execute()
+    )
+    response = wait_for_compute_global_operation(
+        config["provider"]["project_id"], operation, compute
+    )
+    return response
+
+
+def _list_vpcnets(config, compute):
+    response = (
+        compute.networks()
+        .list(
+            project=config["provider"]["project_id"],
+        )
+        .execute()
+    )
+
+    return response["items"] if "items" in response else []
+
+
+def _list_subnets(config, compute, filter=None):
     response = (
         compute.subnetworks()
         .list(
             project=config["provider"]["project_id"],
             region=config["provider"]["region"],
+            filter=filter,
         )
         .execute()
     )
 
-    return response["items"]
+    return response["items"] if "items" in response else []
 
 
 def _get_subnet(config, subnet_id, compute):
diff --git a/sky/skylet/providers/gcp/constants.py b/sky/skylet/providers/gcp/constants.py
@@ -0,0 +1,52 @@
+SKYPILOT_VPC_NAME = "skypilot-vpc"
+
+# Below parameters are from the default VPC on GCP.
+# https://cloud.google.com/vpc/docs/firewalls#more_rules_default_vpc
+VPC_TEMPLATE = {
+    "name": "{VPC_NAME}",
+    "selfLink": "projects/{PROJ_ID}/global/networks/{VPC_NAME}",
+    "autoCreateSubnetworks": True,
+    "mtu": 1460,
+    "routingConfig": {"routingMode": "GLOBAL"},
+}
+FIREWALL_RULES_TEMPLATE = [
+    {
+        "name": "{VPC_NAME}-allow-custom",
+        "description": "Allows connection from any source to any instance on the network using custom protocols.",
+        "network": "projects/{PROJ_ID}/global/networks/{VPC_NAME}",
+        "selfLink": "projects/{PROJ_ID}/global/firewalls/{VPC_NAME}-allow-custom",
+        "direction": "INGRESS",
+        "priority": 65534,
+        "allowed": [
+            {'IPProtocol': 'tcp', 'ports': ['0-65535']},
+            {'IPProtocol': 'udp', 'ports': ['0-65535']},
+            {'IPProtocol': 'icmp'}
+        ],
+        "sourceRanges": ["10.128.0.0/9"],
+    },
+    {
+        "name": "{VPC_NAME}-allow-ssh",
+        "description": "Allows TCP connections from any source to any instance on the network using port 22.",
+        "network": "projects/{PROJ_ID}/global/networks/{VPC_NAME}",
+        "selfLink": "projects/{PROJ_ID}/global/firewalls/{VPC_NAME}-allow-ssh",
+        "direction": "INGRESS",
+        "priority": 65534,
+        "allowed": [{
+            "IPProtocol": "tcp",
+            "ports": ["22"],
+        }],
+        "sourceRanges": ["0.0.0.0/0"],
+    },
+    {
+        "name": "{VPC_NAME}-allow-icmp",
+        "description": "Allows ICMP connections from any source to any instance on the network.",
+        "network": "projects/{PROJ_ID}/global/networks/{VPC_NAME}",
+        "selfLink": "projects/{PROJ_ID}/global/firewalls/{VPC_NAME}-allow-icmp",
+        "direction": "INGRESS",
+        "priority": 65534,
+        "allowed": [{
+            "IPProtocol": "icmp",
+        }],
+        "sourceRanges": ["0.0.0.0/0"],
+    },
+]
diff --git a/sky/skylet/providers/gcp/node.py b/sky/skylet/providers/gcp/node.py
@@ -378,7 +378,8 @@ def list_instances(self, label_filters: Optional[dict] = None,
         return self._list_instances(label_filters, non_terminated_status)
 
     def _list_instances(self, label_filters: Optional[dict],
-                        status_filter: List[str]) -> List[GCPComputeNode]:
+                        status_filter: Optional[List[str]]
+                        ) -> List[GCPComputeNode]:
         label_filters = label_filters or {}
 
         if label_filters:
@@ -395,16 +396,19 @@ def _list_instances(self, label_filters: Optional[dict],
         else:
             label_filter_expr = ""
 
-        instance_state_filter_expr = (
-            "("
-            + " OR ".join(
-                [
-                    "(status = {status})".format(status=status)
-                    for status in status_filter
-                ]
+        if status_filter:
+            instance_state_filter_expr = (
+                "("
+                + " OR ".join(
+                    [
+                        "(status = {status})".format(status=status)
+                        for status in status_filter
+                    ]
+                )
+                + ")"
             )
-            + ")"
-        )
+        else:
+            instance_state_filter_expr = ""
 
         cluster_name_filter_expr = "(labels.{key} = {value})".format(
             key=TAG_RAY_CLUSTER_NAME, value=self.cluster_name
diff --git a/sky/templates/gcp-tpu-create.sh.j2 b/sky/templates/gcp-tpu-create.sh.j2
@@ -9,4 +9,5 @@ gcloud config set project $PROJECT_ID
 yes | gcloud compute tpus create {{tpu_name}} \
 --zone={{zones}} \
 --version={{runtime_version}} \
---accelerator-type={{tpu_type}}
+--accelerator-type={{tpu_type}} \
+--network={{vpc_name}}
