feat: TPU *Create / *get / *create_with_script / *delete Samples (GoogleCloudPlatform#12690)

* Added create /get / delete / create_startup_script TPU Samples with tests and NOX configuration * Update CODEOWNERS and blunderbuss.yml

Author: Thoughtseize1

.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}"