feat: TPU *Create / *get / *create_with_script / *delete Samples

.github/CODEOWNERS

@@ -31,6 +31,7 @@
 /secretmanager/**/*                    @GoogleCloudPlatform/dee-infra @GoogleCloudPlatform/python-samples-reviewers @GoogleCloudPlatform/cloud-samples-reviewers @GoogleCloudPlatform/cloud-secrets-team
 /securitycenter/**/*                   @GoogleCloudPlatform/dee-infra @GoogleCloudPlatform/python-samples-reviewers @GoogleCloudPlatform/cloud-samples-reviewers @GoogleCloudPlatform/gcp-security-command-center
 /service_extensions/**/*               @GoogleCloudPlatform/service-extensions-samples-reviewers @GoogleCloudPlatform/dee-infra @GoogleCloudPlatform/python-samples-reviewers @GoogleCloudPlatform/cloud-samples-reviewers
+/tpu/**/*                              @GoogleCloudPlatform/dee-infra @GoogleCloudPlatform/python-samples-reviewers @GoogleCloudPlatform/cloud-samples-reviewers
 /vmwareengine/**/*                     @GoogleCloudPlatform/dee-infra @GoogleCloudPlatform/python-samples-reviewers @GoogleCloudPlatform/cloud-samples-reviewers
 /webrisk/**/*                          @GoogleCloudPlatform/dee-infra @GoogleCloudPlatform/python-samples-reviewers @GoogleCloudPlatform/cloud-samples-reviewers
 

.github/blunderbuss.yml

@@ -27,6 +27,7 @@ assign_issues_by:
       - "api: recaptchaenterprise"
       - "api: secretmanager"
       - "api: securitycenter"
+      - "api: tpu"
       - "api: vmwareengine"
     to:
       - GoogleCloudPlatform/dee-infra
@@ -154,6 +155,7 @@ assign_prs_by:
       - "api: privateca"
       - "api: recaptchaenterprise"
       - "api: secretmanager"
+      - "api: tpu"
       - "api: securitycenter"
     to:
       - GoogleCloudPlatform/dee-infra

tpu/create_tpu.py

@@ -0,0 +1,79 @@
+# 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 create_cloud_tpu(
+    project_id: str,
+    zone: str,
+    tpu_name: str,
+    tpu_type: str = "v2-8",
+    runtime_version: str = "tpu-vm-tf-2.17.0-pjrt",
+) -> Node:
+    """Creates a Cloud TPU node.
+    Args:
+        project_id (str): The ID of the Google Cloud project.
+        zone (str): The zone where the TPU node will be created.
+        tpu_name (str): The name of the TPU node.
+        tpu_type (str, optional): The type of TPU to create.
+        runtime_version (str, optional): The runtime version for the TPU.
+    Returns:
+        Node: The created TPU node.
+    """
+    # [START tpu_vm_create]
+    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"
+    # tpu_type = "v2-8"
+    # runtime_version = "tpu-vm-tf-2.17.0-pjrt"
+
+    # Create a TPU node
+    node = tpu_v2.Node()
+    node.accelerator_type = tpu_type
+    # To see available runtime version use command:
+    # gcloud compute tpus versions list --zone={ZONE}
+    node.runtime_version = runtime_version
+
+    request = tpu_v2.CreateNodeRequest(
+        parent=f"projects/{project_id}/locations/{zone}",
+        node_id=tpu_name,
+        node=node,
+    )
+
+    # Create a TPU client
+    client = tpu_v2.TpuClient()
+    operation = client.create_node(request=request)
+    print("Waiting for operation to complete...")
+
+    response = operation.result()
+    print(response)
+    # Example response:
+    # name: "projects/[project_id]/locations/[zone]/nodes/my-tpu"
+    # accelerator_type: "v2-8"
+    # state: READY
+    # ...
+
+    # [END tpu_vm_create]
+    return response
+
+
+if __name__ == "__main__":
+    PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
+    ZONE = "us-central1-b"
+    create_cloud_tpu(PROJECT_ID, ZONE, "tpu-name")

tpu/create_tpu_with_script.py

@@ -0,0 +1,76 @@
+# 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 create_cloud_tpu_with_script(
+    project_id: str,
+    zone: str,
+    tpu_name: str,
+    tpu_type: str = "v2-8",
+    runtime_version: str = "tpu-vm-tf-2.17.0-pjrt",
+) -> Node:
+    # [START tpu_vm_create_startup_script]
+    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"
+    # tpu_type = "v2-8"
+    # runtime_version = "tpu-vm-tf-2.17.0-pjrt"
+
+    node = tpu_v2.Node()
+    node.accelerator_type = tpu_type
+    node.runtime_version = runtime_version
+
+    # This startup script updates numpy to the latest version and logs the output to a file.
+    metadata = {
+        "startup-script": """#!/bin/bash
+    echo "Hello World" > /var/log/hello.log
+    sudo pip3 install --upgrade numpy >> /var/log/hello.log 2>&1
+    """
+    }
+
+    # Adding metadata with startup script to the TPU node.
+    node.metadata = metadata
+    # Enabling external IPs for internet access from the TPU node.
+    node.network_config = tpu_v2.NetworkConfig(enable_external_ips=True)
+
+    request = tpu_v2.CreateNodeRequest(
+        parent=f"projects/{project_id}/locations/{zone}",
+        node_id=tpu_name,
+        node=node,
+    )
+
+    client = tpu_v2.TpuClient()
+    operation = client.create_node(request=request)
+    print("Waiting for operation to complete...")
+
+    response = operation.result()
+    print(response.metadata)
+    # Example response:
+    # {'startup-script': '#!/bin/bash\n    echo "Hello World" > /var/log/hello.log\n
+    # ...
+
+    # [END tpu_vm_create_startup_script]
+    return response
+
+
+if __name__ == "__main__":
+    PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
+    ZONE = "us-central1-b"
+    create_cloud_tpu_with_script(PROJECT_ID, ZONE, "tpu-name")

tpu/delete_tpu.py

@@ -0,0 +1,48 @@
+# 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
+
+
+def delete_cloud_tpu(project_id: str, zone: str, tpu_name: str = "tpu-name") -> None:
+    """Deletes a Cloud TPU node.
+    Args:
+        project_id (str): The ID of the Google Cloud project.
+        zone (str): The zone where the TPU node is located.
+        tpu_name (str, optional): The name of the TPU node.
+    Returns: None
+    """
+    # [START tpu_vm_delete]
+    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()
+    try:
+        client.delete_node(
+            name=f"projects/{project_id}/locations/{zone}/nodes/{tpu_name}"
+        )
+        print("The TPU node was deleted.")
+    except Exception as e:
+        print(e)
+
+    # [END tpu_vm_delete]
+
+
+if __name__ == "__main__":
+    PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
+    ZONE = "us-central1-b"
+    delete_cloud_tpu(PROJECT_ID, ZONE, "tpu-name12")

tpu/get_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 import Node
+
+
+def get_cloud_tpu(project_id: str, zone: str, tpu_name: str = "tpu-name") -> Node:
+    """Retrieves a Cloud TPU node.
+    Args:
+        project_id (str): The ID of the Google Cloud project.
+        zone (str): The zone where the TPU node is located.
+        tpu_name (str, optional): The name of the TPU node.
+    Returns:
+        Node: The retrieved TPU node.
+    """
+    # [START tpu_vm_get]
+    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()
+    node = client.get_node(
+        name=f"projects/{project_id}/locations/{zone}/nodes/{tpu_name}"
+    )
+
+    print(node)
+    # Example response:
+    # name: "projects/[project_id]/locations/[zone]/nodes/tpu-name"
+    # state: "READY"
+    # runtime_version: ...
+
+    # [END tpu_vm_get]
+    return node
+
+
+if __name__ == "__main__":
+    PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
+    ZONE = "us-central1-b"
+    get_cloud_tpu(PROJECT_ID, ZONE, "tpu-name")

tpu/noxfile_config.py

@@ -0,0 +1,42 @@
+# Copyright 2021 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
+#
+#      http://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.
+
+# Default TEST_CONFIG_OVERRIDE for python repos.
+
+# You can copy this file into your directory, then it will be imported from
+# the noxfile.py.
+
+# The source of truth:
+# https://github.com/GoogleCloudPlatform/python-docs-samples/blob/main/noxfile_config.py
+
+TEST_CONFIG_OVERRIDE = {
+    # You can opt out from the test for specific Python versions.
+    "ignored_versions": ["2.7", "3.7", "3.9", "3.10", "3.11"],
+    # Old samples are opted out of enforcing Python type hints
+    # All new samples should feature them
+    "enforce_type_hints": True,
+    # An envvar key for determining the project id to use. Change it
+    # to 'BUILD_SPECIFIC_GCLOUD_PROJECT' if you want to opt in using a
+    # build specific Cloud project. You can also use your own string
+    # to use your own Cloud project.
+    "gcloud_project_env": "GOOGLE_CLOUD_PROJECT",
+    # 'gcloud_project_env': 'BUILD_SPECIFIC_GCLOUD_PROJECT',
+    # If you need to use a specific version of pip,
+    # change pip_version_override to the string representation
+    # of the version number, for example, "20.2.4"
+    "pip_version_override": None,
+    # A dictionary you want to inject into your test. Don't put any
+    # secrets here. These values will override predefined values.
+    "envs": {},
+}

tpu/requirements-test.txt

@@ -0,0 +1 @@
+pytest==8.2.0

tpu/requirements.txt

@@ -0,0 +1,3 @@
+google-cloud-tpu==1.18.5
+
+

tpu/test_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
+import uuid
+
+from google.cloud.tpu_v2.types import Node
+
+import pytest
+
+import create_tpu
+import create_tpu_with_script
+import delete_tpu
+import get_tpu
+
+TPU_NAME = "test-tpu-" + uuid.uuid4().hex[:10]
+PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
+ZONE = "us-south1-a"
+TPU_TYPE = "v5litepod-1"
+TPU_VERSION = "tpu-vm-tf-2.17.0-pjrt"
+
+
+# Instance of TPU
+@pytest.fixture(scope="session")
+def tpu_instance() -> Node:
+    yield create_tpu.create_cloud_tpu(PROJECT_ID, ZONE, TPU_NAME, TPU_TYPE, TPU_VERSION)
+    try:
+        delete_tpu.delete_cloud_tpu(PROJECT_ID, ZONE, TPU_NAME)
+    except Exception as e:
+        print(f"Error during cleanup: {e}")
+
+
+def test_creating_tpu(tpu_instance: Node) -> None:
+    assert tpu_instance.state == Node.State.READY
+
+
+def test_creating_with_startup_script() -> None:
+    tpu_name_with_script = "script-tpu-" + uuid.uuid4().hex[:5]
+    try:
+        tpu_with_script = create_tpu_with_script.create_cloud_tpu_with_script(
+            PROJECT_ID, ZONE, tpu_name_with_script, TPU_TYPE, TPU_VERSION
+        )
+        assert "--upgrade numpy" in tpu_with_script.metadata["startup-script"]
+    finally:
+        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}"