diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index f1d26b9d6e5..490692272cd 100644 --- a/.github/CODEOWNERS +++ b/.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 diff --git a/.github/blunderbuss.yml b/.github/blunderbuss.yml index 9bc43049794..fa02379d2af 100644 --- a/.github/blunderbuss.yml +++ b/.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 diff --git a/tpu/create_tpu.py b/tpu/create_tpu.py new file mode 100644 index 00000000000..da2cdeaaf5b --- /dev/null +++ b/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") diff --git a/tpu/create_tpu_with_script.py b/tpu/create_tpu_with_script.py new file mode 100644 index 00000000000..53a2ada4fb1 --- /dev/null +++ b/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") diff --git a/tpu/delete_tpu.py b/tpu/delete_tpu.py new file mode 100644 index 00000000000..b185aed3ac2 --- /dev/null +++ b/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") diff --git a/tpu/get_tpu.py b/tpu/get_tpu.py new file mode 100644 index 00000000000..9af19039af5 --- /dev/null +++ b/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") diff --git a/tpu/noxfile_config.py b/tpu/noxfile_config.py new file mode 100644 index 00000000000..457e86f5413 --- /dev/null +++ b/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": {}, +} diff --git a/tpu/requirements-test.txt b/tpu/requirements-test.txt new file mode 100644 index 00000000000..15d066af319 --- /dev/null +++ b/tpu/requirements-test.txt @@ -0,0 +1 @@ +pytest==8.2.0 diff --git a/tpu/requirements.txt b/tpu/requirements.txt new file mode 100644 index 00000000000..341dbbfd3a0 --- /dev/null +++ b/tpu/requirements.txt @@ -0,0 +1,3 @@ +google-cloud-tpu==1.18.5 + + diff --git a/tpu/test_tpu.py b/tpu/test_tpu.py new file mode 100644 index 00000000000..6a69c98795a --- /dev/null +++ b/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}"