Skip to content
This repository has been archived by the owner on Jul 6, 2023. It is now read-only.

feat: Add MultisliceNode #229

Merged
merged 2 commits into from Jun 1, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
19 changes: 19 additions & 0 deletions google/cloud/tpu_v2alpha1/types/cloud_tpu.py
Expand Up @@ -396,6 +396,9 @@ class Node(proto.Message):
The AccleratorConfig for the TPU Node.
shielded_instance_config (google.cloud.tpu_v2alpha1.types.ShieldedInstanceConfig):
Shielded Instance options.
multislice_node (bool):
Output only. Whether the Node belongs to a
Multislice group.
"""

class State(proto.Enum):
Expand Down Expand Up @@ -603,6 +606,10 @@ class ApiVersion(proto.Enum):
number=45,
message="ShieldedInstanceConfig",
)
multislice_node: bool = proto.Field(
proto.BOOL,
number=47,
)


class QueuedResource(proto.Message):
Expand Down Expand Up @@ -1317,6 +1324,14 @@ class DeleteQueuedResourceRequest(proto.Message):
Required. The resource name.
request_id (str):
Idempotent request UUID.
force (bool):
If set to true, all running nodes belonging
to this queued resource will be deleted first
and then the queued resource will be deleted.
Otherwise (i.e. force=false), the queued
resource will only be deleted if its nodes have
already been deleted or the queued resource is
in the ACCEPTED, FAILED, or SUSPENDED state.
"""

name: str = proto.Field(
Expand All @@ -1327,6 +1342,10 @@ class DeleteQueuedResourceRequest(proto.Message):
proto.STRING,
number=2,
)
force: bool = proto.Field(
proto.BOOL,
number=3,
)


class ServiceIdentity(proto.Message):
Expand Down
Expand Up @@ -8,7 +8,7 @@
],
"language": "PYTHON",
"name": "google-cloud-tpu",
"version": "1.10.1"
"version": "0.1.0"
},
"snippets": [
{
Expand Down
Expand Up @@ -8,7 +8,7 @@
],
"language": "PYTHON",
"name": "google-cloud-tpu",
"version": "1.10.1"
"version": "0.1.0"
},
"snippets": [
{
Expand Down
Expand Up @@ -8,7 +8,7 @@
],
"language": "PYTHON",
"name": "google-cloud-tpu",
"version": "1.10.1"
"version": "0.1.0"
},
"snippets": [
{
Expand Down
2 changes: 1 addition & 1 deletion scripts/fixup_tpu_v2alpha1_keywords.py
Expand Up @@ -42,7 +42,7 @@ class tpuCallTransformer(cst.CSTTransformer):
'create_node': ('parent', 'node', 'node_id', 'request_id', ),
'create_queued_resource': ('parent', 'queued_resource', 'queued_resource_id', 'request_id', ),
'delete_node': ('name', 'request_id', ),
'delete_queued_resource': ('name', 'request_id', ),
'delete_queued_resource': ('name', 'request_id', 'force', ),
'generate_service_identity': ('parent', ),
'get_accelerator_type': ('name', ),
'get_guest_attributes': ('name', 'query_path', 'worker_ids', ),
Expand Down
4 changes: 4 additions & 0 deletions tests/unit/gapic/tpu_v2alpha1/test_tpu.py
Expand Up @@ -1092,6 +1092,7 @@ def test_get_node(request_type, transport: str = "grpc"):
id=205,
api_version=cloud_tpu.Node.ApiVersion.V1_ALPHA1,
queued_resource="queued_resource_value",
multislice_node=True,
)
response = client.get_node(request)

Expand All @@ -1114,6 +1115,7 @@ def test_get_node(request_type, transport: str = "grpc"):
assert response.id == 205
assert response.api_version == cloud_tpu.Node.ApiVersion.V1_ALPHA1
assert response.queued_resource == "queued_resource_value"
assert response.multislice_node is True


def test_get_node_empty_call():
Expand Down Expand Up @@ -1162,6 +1164,7 @@ async def test_get_node_async(
id=205,
api_version=cloud_tpu.Node.ApiVersion.V1_ALPHA1,
queued_resource="queued_resource_value",
multislice_node=True,
)
)
response = await client.get_node(request)
Expand All @@ -1185,6 +1188,7 @@ async def test_get_node_async(
assert response.id == 205
assert response.api_version == cloud_tpu.Node.ApiVersion.V1_ALPHA1
assert response.queued_resource == "queued_resource_value"
assert response.multislice_node is True


@pytest.mark.asyncio
Expand Down