Enable TPU in GKE (kubeflow#1766) (kubeflow#1774)

* Add enable_tpu setting * Merge change about Stackdriver GKE enabled only when using v1beta1 * Add enable_tpu setting
StatCan · Oct 24, 2018 · f98dbaa · f98dbaa
1 parent 9640323
commit f98dbaa
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 1 deletion.
diff --git a/deployment/gke/deployment_manager_configs/cluster-kubeflow.yaml b/deployment/gke/deployment_manager_configs/cluster-kubeflow.yaml
@@ -49,6 +49,8 @@ resources:
     gpu-pool-enable-autoscaling: true
     gpu-pool-min-nodes: 0
     gpu-pool-max-nodes: 10
+    # Whether to enable TPUs
+    enable_tpu: false
     securityConfig:
       # Whether to use a cluster with private IPs
       # Use v1beta1 api

diff --git a/deployment/gke/deployment_manager_configs/cluster.jinja b/deployment/gke/deployment_manager_configs/cluster.jinja
@@ -74,9 +74,14 @@ resources:
       name: {{ CLUSTER_NAME }}
       initialClusterVersion: "{{ properties['cluster-version'] }}"
       {% if properties['gkeApiVersion'] == 'v1beta1' %}
-      # We need 1.10.2 to support Stackdrivier GKE.
+      # We need 1.10.2 to support Stackdriver GKE.
       loggingService: logging.googleapis.com/kubernetes
       monitoringService: monitoring.googleapis.com/kubernetes
+      {% if properties['enable_tpu'] %}
+      enable_tpu: {{ properties['enable_tpu'] }}
+      ipAllocationPolicy:
+        useIpAliases: {{ properties['enable_tpu'] }}
+      {% endif %}
       podSecurityPolicyConfig:
         enabled: {{ properties['securityConfig']['podSecurityPolicy'] }}
       {% endif %}