Enable TPUs to use Shared VPC

products/tpu/api.yaml

@@ -50,7 +50,7 @@ objects:
       guides:
         'Official Documentation':
           'https://cloud.google.com/tpu/docs/'
-      api: 'https://cloud.google.com/tpu/docs/reference/rest/'
+      api: 'https://cloud.google.com/tpu/docs/reference/rest/v1/projects.locations.nodes'
     parameters:
       - !ruby/object:Api::Type::String # TODO: resourceref?
         name: 'zone'
@@ -94,8 +94,9 @@ objects:
           used.
       - !ruby/object:Api::Type::String
         name: 'cidrBlock'
-        required: true
         input: true
+        conflicts:
+          - use_service_networking
         description: |
           The CIDR block that the TPU node will use when selecting an IP
           address. This CIDR block must be a /29 block; the Compute Engine
@@ -114,6 +115,17 @@ objects:
           node. To share resources, including Google Cloud Storage data, with
           the Tensorflow job running in the Node, this account must have
           permissions to that data.
+      - !ruby/object:Api::Type::Boolean
+        name: 'useServiceNetworking'
+        description: |
+          Whether the VPC peering for the node is set up through Service Networking API.
+          The VPC Peering should be set up before provisioning the node. If this field is set,
+          cidr_block field should not be specified. If the network that you want to peer the
+          TPU Node to is a Shared VPC network, the node must be created with this this field enabled.
+        input: true
+        default_value: false
+        conflicts:
+          - cidr_block
       - !ruby/object:Api::Type::NestedObject
         name: 'schedulingConfig'
         input: true

products/tpu/terraform.yaml

@@ -33,13 +33,15 @@ overrides: !ruby/object:Overrides::ResourceOverrides
     properties:
       name: !ruby/object:Overrides::Terraform::PropertyOverride
         custom_flatten: 'templates/terraform/custom_flatten/name_from_self_link.erb'
-      network: !ruby/object:Overrides::Terraform::PropertyOverride
+      cidrBlock: !ruby/object:Overrides::Terraform::PropertyOverride
         default_from_api: true
-        diff_suppress_func: 'compareSelfLinkOrResourceName'
       schedulingConfig: !ruby/object:Overrides::Terraform::PropertyOverride
         diff_suppress_func: 'compareTpuNodeSchedulingConfig'
       schedulingConfig.preemptible: !ruby/object:Overrides::Terraform::PropertyOverride
         diff_suppress_func: 'compareTpuNodeSchedulingConfig'
+      network: !ruby/object:Overrides::Terraform::PropertyOverride
+        default_from_api: true
+        diff_suppress_func: 'compareTpuNodeNetwork'
       zone: !ruby/object:Overrides::Terraform::PropertyOverride
         ignore_read: true
     custom_code: !ruby/object:Provider::Terraform::CustomCode

templates/terraform/constants/tpu_node.erb

@@ -12,6 +12,39 @@ func compareTpuNodeSchedulingConfig(k, old, new string, d *schema.ResourceData)
 	return false
 }
 
+
+// The API takes the project_id in the network field and returns it as a number which causes permadiff.
+func compareTpuNodeNetwork(k, old, new string, d *schema.ResourceData) bool {
+	// 1st part of diffsuppress
+
+	retval := false
+	network := GetResourceNameFromSelfLink(old)
+
+	serviceNetworkingNetworkName, err := retrieveServiceNetworkingNetworkName(d, config, network)
+	if err != nil {
+		return err
+	}
+
+	if new == serviceNetworkingNetworkName {
+		retval = true
+	}
+
+	// 2nd part of diffsuppress
+	newParts := strings.Split(new, "/")
+
+	if len(newParts) == 1 {
+		// `new` is a name
+		// `old` is always a self_link
+		if GetResourceNameFromSelfLink(old) == newParts[0] {
+			retval = true
+		}
+	}
+
+	// The `new` string is a self_link
+	retval = retval && compareSelfLinkRelativePaths("", old, new, nil)
+	return retval
+}
+
 func validateHttpHeaders() schema.SchemaValidateFunc {
 	return func(i interface{}, k string) (s []string, es []error) {
 		headers := i.(map[string]interface{})

templates/terraform/examples/tpu_node_full.tf.erb

@@ -13,17 +13,17 @@ resource "google_tpu_node" "<%= ctx[:primary_resource_id] %>" {
 
   accelerator_type = "v3-8"
 
-  cidr_block         = "10.3.0.0/29"
   tensorflow_version = data.google_tpu_tensorflow_versions.available.versions[0]
 
   description = "Terraform Google Provider test TPU"
+  use_service_networking = true
 <%#-
   We previously used a separate network resource here, but TPUs only allow using 50
   different network names, ever. This caused our tests to start failing, so just
   use the default network in order to still demonstrate using as many fields as
   possible on the resource.
 -%>
-  network = "default"
+  network = google_service_networking_connection.private_service_connection.network
 
   labels = {
     foo = "bar"
@@ -33,3 +33,22 @@ resource "google_tpu_node" "<%= ctx[:primary_resource_id] %>" {
     preemptible = true
   }
 }
+
+data "google_compute_network" "network" {
+  name = "default"
+}
+
+resource "google_compute_global_address" "service_range" {
+  provider = google-beta
+  name          = "tf-test%{random_suffix}"
+  purpose       = "VPC_PEERING"
+  address_type  = "INTERNAL"
+  prefix_length = 16
+  network       = data.google_compute_network.network.id
+}
+
+resource "google_service_networking_connection" "private_service_connection" {
+  network                 = data.google_compute_network.network.id
+  service                 = "servicenetworking.googleapis.com"
+  reserved_peering_ranges = [google_compute_global_address.service_range.name]
+}