Hackathon/tpu pod

python/ray/autoscaler/_private/gcp/config.py

@@ -3,6 +3,7 @@
 import logging
 import os
 import time
+from typing import Any, Dict, Tuple
 from functools import partial
 
 from cryptography.hazmat.backends import default_backend
@@ -68,13 +69,6 @@ def get_node_type(node: dict) -> GCPNodeType:
         )
 
     if "machineType" not in node and "acceleratorType" in node:
-        # remove after TPU pod support is added!
-        if node["acceleratorType"] not in ("v2-8", "v3-8"):
-            raise ValueError(
-                "For now, only v2-8' and 'v3-8' accelerator types are "
-                "supported. Support for TPU pods will be added in the future."
-            )
-
         return GCPNodeType.TPU
     return GCPNodeType.COMPUTE
 
@@ -288,15 +282,16 @@ def bootstrap_gcp(config):
         config["provider"][HAS_TPU_PROVIDER_FIELD] = True
 
         # We can't run autoscaling through a serviceAccount on TPUs (atm)
-        if _is_head_node_a_tpu(config):
-            raise RuntimeError("TPUs are not supported as head nodes.")
+        # if _is_head_node_a_tpu(config):
+        #     raise RuntimeError("TPUs are not supported as head nodes.")
 
     crm, iam, compute, tpu = construct_clients_from_provider_config(config["provider"])
 
     config = _configure_project(config, crm)
     config = _configure_iam_role(config, crm, iam)
     config = _configure_key_pair(config, compute)
     config = _configure_subnet(config, compute)
+    config = _hack_in_tpu_chip_type(config)
 
     return config
 
@@ -547,6 +542,71 @@ def _configure_subnet(config, compute):
     return config
 
 
+def _hack_in_tpu_chip_type(config: Dict[str, Any]) -> Dict[str, Any]:
+    """Add "virtual" instance types for TPU chips.
+    Also update TPU resources for the real instance type.
+    """
+    config = copy.deepcopy(config)
+    available_node_types = config["available_node_types"]
+    for node_type_name, node_type in copy.deepcopy(available_node_types).items():
+        num_tpus = num_tpus_from_node_config(node_type.get("node_config"))
+        if num_tpus > 0:
+            available_node_types[node_type_name]["resources"]["TPU"] = num_tpus
+            # Hacks, don't mind me.
+            available_node_types[node_type_name]["min_workers"] = 0
+            available_node_types[node_type_name]["max_workers"] = 0
+        if num_tpus > 1:
+            tpu_chip_type_name, tpu_chip_type = _get_tpu_chip_type(
+                node_type_name, node_type
+            )
+            available_node_types[tpu_chip_type_name] = tpu_chip_type
+    # Idle downscaling not supported!
+    config["idle_timeout_minutes"] = 1000000000
+    return config
+
+
+def num_tpus_from_node_config(node_type: Dict[str, Any]) -> int:
+    """Get number of tpus for a tpu instance type.
+
+    Return 0 if it's not a TPU type or parsing the number of TPUs failed.
+    """
+    accelerator_type = node_type.get("acceleratorType", "")
+    return num_tpus_from_accelerator_type(accelerator_type)
+
+
+def num_tpus_from_accelerator_type(accelerator_type: str) -> int:
+    type_components = accelerator_type.split("-")
+    if not len(type_components) == 2:
+        return 0
+    version, suffix = accelerator_type.split("-")
+    if version.startswith("v") and suffix.isnumeric():
+        num_tpu, remainder = divmod(int(suffix), 8)
+        if remainder:
+            # Not divisible by 8, hmmmmm.
+            return 0
+        else:
+            return num_tpu
+    else:
+        return 0
+
+
+def _get_tpu_chip_type(
+    tpu_instance_type_name, tpu_node_type
+) -> Tuple[str, Dict[str, Any]]:
+    """Get a virtual type node type, used for autoscaler book-keeping."""
+    tpu_chip_type_name = f"{tpu_instance_type_name}-chip"
+    tpu_chip_type = {
+        # Pure hackery, don't mind me.
+        "min_workers": tpu_node_type["min_workers"],
+        # Prevent the autoscaler from attempting to directly terminate this node type.
+        "max_workers": 100000000000000,
+        "resources": {"TPU": 1},
+        # Copy parent type's config for this hack.
+        "node_config": tpu_node_type["node_config"],
+    }
+    return tpu_chip_type_name, tpu_chip_type
+
+
 def _list_subnets(config, compute):
     response = (
         compute.subnetworks()

python/ray/autoscaler/_private/gcp/node.py

@@ -96,7 +96,14 @@ def _generate_node_name(labels: dict, node_suffix: str) -> str:
     The suffix is expected to be one of 'compute' or 'tpu'
     (as in ``GCPNodeType``).
     """
-    name_label = labels[TAG_RAY_NODE_NAME]
+    node_name_key = None
+    # Get one of the copies of the node name key if we're dealing with a TPU pod.
+    for key in labels:
+        if TAG_RAY_NODE_NAME in key:
+            node_name_key = key
+            break
+    assert node_name_key, "Couldn't get a node name key."
+    name_label = labels[node_name_key]
     assert len(name_label) <= (INSTANCE_NAME_MAX_LEN - INSTANCE_NAME_UUID_LEN - 1), (
         name_label,
         len(name_label),
@@ -156,7 +163,7 @@ def get_external_ip(self) -> str:
         return
 
     @abc.abstractmethod
-    def get_internal_ip(self) -> str:
+    def get_internal_ip(self, index: int = 0) -> str:
         return
 
     def __repr__(self) -> str:
@@ -180,8 +187,14 @@ def get_external_ip(self) -> str:
             .get("natIP", None)
         )
 
-    def get_internal_ip(self) -> str:
-        return self.get("networkInterfaces", [{}])[0].get("networkIP")
+    def get_internal_ip(self, index=0) -> str:
+        network_interfaces = self.get("networkInterfaces", [{}])
+        if index < len(network_interfaces):
+            return self.get("networkInterfaces", [{}])[index].get("networkIP")
+        else:
+            return None
+
+    # netWorkerEndpoints[*]["ipAddress"]
 
 
 class GCPTPUNode(GCPNode):
@@ -203,8 +216,12 @@ def get_external_ip(self) -> str:
             .get("externalIp", None)
         )
 
-    def get_internal_ip(self) -> str:
-        return self.get("networkEndpoints", [{}])[0].get("ipAddress", None)
+    def get_internal_ip(self, index=0) -> str:
+        network_interfaces = self.get("networkEndpoints", [{}])
+        if index < len(network_interfaces):
+            return self.get("networkEndpoints", [{}])[index].get("ipAddress")
+        else:
+            return None
 
 
 class GCPResource(metaclass=abc.ABCMeta):

python/ray/autoscaler/_private/gcp/node_provider.py

@@ -2,14 +2,16 @@
 import time
 from functools import wraps
 from threading import RLock
-from typing import Dict, List, Tuple
+from typing import Dict, List, Tuple, Optional
 
 import googleapiclient
 
 from ray.autoscaler._private.gcp.config import (
     bootstrap_gcp,
     construct_clients_from_provider_config,
     get_node_type,
+    num_tpus_from_accelerator_type,
+    num_tpus_from_node_config,
 )
 
 # The logic has been abstracted away here to allow for different GCP resources
@@ -26,6 +28,8 @@
 
 logger = logging.getLogger(__name__)
 
+TPUCHIP = "tpuchip"
+
 
 def _retry(method, max_tries=5, backoff_s=1):
     """Retry decorator for methods of GCPNodeProvider.
@@ -110,55 +114,113 @@ def non_terminated_nodes(self, tag_filters: dict):
 
             # Note: All the operations use "name" as the unique instance id
             self.cached_nodes = {i["name"]: i for i in instances}
-            return [i["name"] for i in instances]
+            node_names = []
+            for i in instances:
+                instance_name = i["name"]
+                num_tpus = self._num_tpus_from_instance(i)
+                if num_tpus > 1:
+                    for tpu_index in range(num_tpus):
+                        tpu_node_name = self._add_tpu_chip_suffix(
+                            instance_name, tpu_index
+                        )
+                        node_names.append(tpu_node_name)
+                else:
+                    node_names.append(instance_name)
+
+            return node_names
+
+    def _add_tpu_chip_suffix(self, name, tpu_index) -> str:
+        return f"{name}-{TPUCHIP}-{tpu_index}"
+
+    def _name_and_tpu_index(self, suffixed_name) -> Tuple[str, int]:
+        components = suffixed_name.split("-")
+        if components[-2] == TPUCHIP and components[-1].isnumeric():
+            return "-".join(components[:-2]), int(components[-1])
+        else:
+            return "", -1
+
+    def _is_tpu_chip(self, node_name):
+        return TPUCHIP in node_name
+
+    def _num_tpus_from_instance(self, instance) -> int:
+        accelerator_type = instance.get("acceleratorType")
+        if accelerator_type:
+            return num_tpus_from_accelerator_type(accelerator_type)
+        else:
+            return 0
 
     def is_running(self, node_id: str):
         with self.lock:
+            if self._is_tpu_chip(node_id):
+                node_id, _ = self._name_and_tpu_index(node_id)
             node = self._get_cached_node(node_id)
             return node.is_running()
 
     def is_terminated(self, node_id: str):
         with self.lock:
+            if self._is_tpu_chip(node_id):
+                node_id, _ = self._name_and_tpu_index(node_id)
             node = self._get_cached_node(node_id)
             return node.is_terminated()
 
     def node_tags(self, node_id: str):
         with self.lock:
-            node = self._get_cached_node(node_id)
-            return node.get_labels()
+            if self._is_tpu_chip(node_id):
+                node_id, tpu_index = self._name_and_tpu_index(node_id)
+                node = self._get_cached_node(node_id)
+                labels = node.get_labels()
+                return self._get_tags_for_tpu_chip(labels, tpu_index)
+            else:
+                node = self._get_cached_node(node_id)
+                return node.get_labels()
+
+    def _get_tags_for_tpu_chip(
+        self, labels: Dict[str, str], tpu_index
+    ) -> Dict[str, str]:
+        tags: Dict[str, str] = {}
+        for key in labels:
+            tag_key, index = self._name_and_tpu_index(key)
+            if tag_key == "":
+                # Mis-formatted key.
+                continue
+            if index == tpu_index:
+                tags[tag_key] = labels[key]
+        return tags
 
     @_retry
     def set_node_tags(self, node_id: str, tags: dict):
         with self.lock:
-            labels = tags
-            node = self._get_node(node_id)
+            if self._is_tpu_chip(node_id):
+                node_id, tpu_index = self._name_and_tpu_index(node_id)
+                labels = {
+                    f"{self._add_tpu_chip_suffix(key, tpu_index)}": value
+                    for key, value in tags.items()
+                }
+            else:
+                labels = tags
 
+            node = self._get_node(node_id)
             resource = self._get_resource_depending_on_node_name(node_id)
-
             result = resource.set_labels(node=node, labels=labels)
 
             return result
 
-    def external_ip(self, node_id: str):
+    def external_ip(self, node_id: str) -> Optional[str]:
+        # (Not necessary to add TPU host logic here.)
         with self.lock:
             node = self._get_cached_node(node_id)
 
             ip = node.get_external_ip()
-            if ip is None:
-                node = self._get_node(node_id)
-                ip = node.get_external_ip()
 
             return ip
 
-    def internal_ip(self, node_id: str):
+    def internal_ip(self, node_id: str) -> Optional[str]:
         with self.lock:
+            ip_index = 0
+            if self._is_tpu_chip(node_id):
+                node_id, ip_index = self._name_and_tpu_index(node_id)
             node = self._get_cached_node(node_id)
-
-            ip = node.get_internal_ip()
-            if ip is None:
-                node = self._get_node(node_id)
-                ip = node.get_internal_ip()
-
+            ip = node.get_internal_ip(ip_index)
             return ip
 
     @_retry
@@ -170,6 +232,11 @@ def create_node(self, base_config: dict, tags: dict, count: int) -> Dict[str, di
         """
         with self.lock:
             labels = tags  # gcp uses "labels" instead of aws "tags"
+            num_tpus = num_tpus_from_node_config(base_config)
+            if num_tpus > 1:
+                labels = self._format_tpu_chip_labels(tags, num_tpus)
+            else:
+                labels = tags
 
             node_type = get_node_type(base_config)
             resource = self.resources[node_type]
@@ -179,9 +246,22 @@ def create_node(self, base_config: dict, tags: dict, count: int) -> Dict[str, di
             )  # type: List[Tuple[dict, str]]
             return {instance_id: result for result, instance_id in results}
 
+    def _format_tpu_chip_labels(self, tags, num_tpus):
+        """Represent each label key for tpu node in a pod as the usual
+        tag key dash a numeric suffix.
+        """
+        return {
+            f"{self._add_tpu_chip_suffix(key, tpu_index)}": value
+            for key, value in tags.items()
+            for tpu_index in range(num_tpus)
+        }
+
     @_retry
     def terminate_node(self, node_id: str):
         with self.lock:
+            if self._is_tpu_chip(node_id):
+                # Kill the parent instance instead.
+                node_id, _ = self._name_and_tpu_index(node_id)
             resource = self._get_resource_depending_on_node_name(node_id)
             try:
                 result = resource.delete_instance(