feat: New Samples for TPU: Start, Stop, List (#12702)

* Added Samples for TPU: List, Start, Stop
* Updated auto-label.yaml

Author: Thoughtseize1

.github/auto-label.yaml

@@ -82,6 +82,7 @@ path:
         storagetransfer: "storagetransfer"
         talent: "jobs"
         texttospeech: "texttospeech"
+        tpu: "tpu"
         trace: "cloudtrace"
         translate: "translate"
         vision: "vision"

tpu/list_tpu.py

@@ -0,0 +1,54 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+
+from google.cloud.tpu_v2.services.tpu.pagers import ListNodesPager
+
+
+def list_cloud_tpu(project_id: str, zone: str) -> ListNodesPager:
+    """List all TPU nodes in the project and zone.
+    Args:
+        project_id (str): The ID of Google Cloud project.
+        zone (str): The zone of the TPU nodes.
+    Returns:
+        ListNodesPager: The list of TPU nodes.
+    """
+    # [START tpu_vm_list]
+    from google.cloud import tpu_v2
+
+    # TODO(developer): Update and un-comment below lines
+    # project_id = "your-project-id"
+    # zone = "us-central1-b"
+
+    client = tpu_v2.TpuClient()
+
+    nodes = client.list_nodes(parent=f"projects/{project_id}/locations/{zone}")
+    for node in nodes:
+        print(node.name)
+        print(node.state)
+        print(node.accelerator_type)
+    # Example response:
+    # projects/[project_id]/locations/[zone]/nodes/node-name
+    # State.READY
+    # v2-8
+
+    # [END tpu_vm_list]
+    return nodes
+
+
+if __name__ == "__main__":
+    PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
+    ZONE = "us-central1-b"
+    list_cloud_tpu(PROJECT_ID, ZONE)

tpu/start_tpu.py

@@ -0,0 +1,61 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+
+from google.cloud.tpu_v2 import Node
+
+
+def start_cloud_tpu(project_id: str, zone: str, tpu_name: str = "tpu-name") -> Node:
+    """Start a TPU node.
+    Args:
+        project_id (str): The ID of Google Cloud project.
+        zone (str): The zone of the TPU node.
+        tpu_name (str): The name of the TPU node to start.
+    Returns:
+        Node: The started TPU node.
+    """
+    # [START tpu_vm_start]
+    from google.cloud import tpu_v2
+
+    # TODO(developer): Update and un-comment below lines
+    # project_id = "your-project-id"
+    # zone = "us-central1-b"
+    # tpu_name = "tpu-name"
+
+    client = tpu_v2.TpuClient()
+
+    request = tpu_v2.StartNodeRequest(
+        name=f"projects/{project_id}/locations/{zone}/nodes/{tpu_name}",
+    )
+    try:
+        operation = client.start_node(request=request)
+        print("Waiting for start operation to complete...")
+        response = operation.result()
+        print(f"TPU {tpu_name} has been started")
+        print(response.state)
+        # Example response:
+        # State.READY
+
+        return response
+    except Exception as e:
+        print(e)
+
+    # [END tpu_vm_start]
+
+
+if __name__ == "__main__":
+    PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
+    ZONE = "us-central1-b"
+    start_cloud_tpu(PROJECT_ID, ZONE, "tpu-name")

tpu/stop_tpu.py

@@ -0,0 +1,62 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+
+from google.cloud.tpu_v2 import Node
+
+
+def stop_cloud_tpu(project_id: str, zone: str, tpu_name: str = "tpu-name") -> Node:
+    """Stop a TPU node.
+    Args:
+        project_id (str): The ID of Google Cloud project.
+        zone (str): The zone of the TPU node.
+        tpu_name (str): The name of the TPU node to stop.
+    Returns:
+        Node: The stopped TPU node.
+    """
+
+    # [START tpu_vm_stop]
+    from google.cloud import tpu_v2
+
+    # TODO(developer): Update and un-comment below lines
+    # project_id = "your-project-id"
+    # zone = "us-central1-b"
+    # tpu_name = "tpu-name"
+
+    client = tpu_v2.TpuClient()
+
+    request = tpu_v2.StopNodeRequest(
+        name=f"projects/{project_id}/locations/{zone}/nodes/{tpu_name}",
+    )
+    try:
+        operation = client.stop_node(request=request)
+        print("Waiting for stop operation to complete...")
+        response = operation.result()
+        print(f"This TPU {tpu_name} has been stopped")
+        print(response.state)
+        # Example response:
+        # State.STOPPED
+
+        return response
+    except Exception as e:
+        print(e)
+
+    # [END tpu_vm_stop]
+
+
+if __name__ == "__main__":
+    PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
+    ZONE = "us-central1-b"
+    stop_cloud_tpu(PROJECT_ID, ZONE, "tpu-name")

tpu/test_tpu.py

@@ -22,6 +22,9 @@
 import create_tpu_with_script
 import delete_tpu
 import get_tpu
+import list_tpu
+import start_tpu
+import stop_tpu
 
 TPU_NAME = "test-tpu-" + uuid.uuid4().hex[:10]
 PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
@@ -52,10 +55,26 @@ def test_creating_with_startup_script() -> None:
         )
         assert "--upgrade numpy" in tpu_with_script.metadata["startup-script"]
     finally:
+        print(f"\n\n ------------ Deleting TPU {TPU_NAME}\n ------------")
         delete_tpu.delete_cloud_tpu(PROJECT_ID, ZONE, tpu_name_with_script)
 
 
 def test_get_tpu() -> None:
     tpu = get_tpu.get_cloud_tpu(PROJECT_ID, ZONE, TPU_NAME)
     assert tpu.state == Node.State.READY
     assert tpu.name == f"projects/{PROJECT_ID}/locations/{ZONE}/nodes/{TPU_NAME}"
+
+
+def test_list_tpu() -> None:
+    nodes = list_tpu.list_cloud_tpu(PROJECT_ID, ZONE)
+    assert len(list(nodes)) > 0
+
+
+def test_stop_tpu() -> None:
+    node = stop_tpu.stop_cloud_tpu(PROJECT_ID, ZONE, TPU_NAME)
+    assert node.state == Node.State.STOPPED
+
+
+def test_start_tpu() -> None:
+    node = start_tpu.start_cloud_tpu(PROJECT_ID, ZONE, TPU_NAME)
+    assert node.state == Node.State.READY