def node_pools(options: Dict[str, Any]) -> Iterator[List[str]]:
credentials, project = google.auth.default()
if options["project"] is None:
options["project"] = project
gke = ClusterManagerClient(credentials=credentials)
# build node pool configurations
pools = {}
if options["generator"]:
if options["load_balancer"] is None or options[
"server_uri"] is not None:
pools["generator"] = NodePool(initial_node_count=1)
pools["generator"].config.machine_type = "n1-highcpu-2"
if options["load_balancer"] is None:
pools["server"] = NodePool(initial_node_count=4)
pools["server"].config.machine_type = "n1-highcpu-2"
if options["emulator"]:
pools["emulator"] = NodePool(initial_node_count=1)
else:
# need pubsub permissions
pools["server"].config.oauth_scopes.append(
"https://www.googleapis.com/auth/pubsub")
# add labels
for name, pool in pools.items():
pool.name = name
pool.config.preemptible = options["preemptible"]
pool.config.labels["name"] = name
if options["location"][-2] == "-":
# triple node count for single zone cluster
pool.initial_node_count *= 3
# create cluster
if not pools:
yield [] # nothing to create
else:
kwargs = {
"cluster":
Cluster(
name=options["cluster"],
logging_service=None,
monitoring_service=None,
node_pools=list(pools.values()),
),
"parent":
f"projects/{options['project']}/locations/{options['location']}",
}
name = f"{kwargs['parent']}/clusters/{options['cluster']}"
try:
operation = gke.create_cluster(**kwargs)
except AlreadyExists:
pass
else:
# wait for operation to complete
request = GetOperationRequest(
name=operation.self_link.split("projects").pop())
while gke.get_operation(
request).status <= Operation.Status.RUNNING:
time.sleep(15)
# set kube credentials
cluster = gke.get_cluster(name=name)
config = kube.Configuration()
config.host = f"https://{cluster.endpoint}:443"
config.verify_ssl = False
config.api_key = {"authorization": f"Bearer {credentials.token}"}
kube.Configuration.set_default(config)
# delete cluster after test completes
try:
yield list(pools)
finally:
gke.delete_cluster(name=name)
def cluster(options: Dict[str, Any]) -> Iterator[Optional[str]]:
if options["cluster"] is None:
load_kube_config()
if options["project"] is None:
options["project"] = "test"
yield None
else:
credentials, project = google.auth.default()
if options["project"] is None:
options["project"] = project
gke = ClusterManagerClient(credentials=credentials)
# create cluster
parent = f"projects/{options['project']}/locations/{options['location']}"
name = f"{parent}/clusters/{options['cluster']}"
try:
operation = gke.create_cluster(
cluster=Cluster(
name=options["cluster"],
logging_service=None,
monitoring_service=None,
node_pools=[
NodePool(
initial_node_count=1,
name="test",
config={
"preemptible": options["preemptible"],
"machine_type": "n1-highcpu-2",
},
)
],
),
parent=parent,
)
except AlreadyExists:
pass
else:
# wait for operation to complete
request = GetOperationRequest(
name=operation.self_link.split("projects").pop())
while gke.get_operation(
request).status <= Operation.Status.RUNNING:
time.sleep(15)
# set kube credentials
cluster = gke.get_cluster(name=name)
config = kube.Configuration()
config.host = f"https://{cluster.endpoint}:443"
config.verify_ssl = False
config.api_key = {"authorization": f"Bearer {credentials.token}"}
kube.Configuration.set_default(config)
# delete cluster after test completes
try:
yield options["cluster"]
finally:
try:
# delete persistent volumes because gke cluster delete won't do it
# https://cloud.google.com/kubernetes-engine/docs/how-to/deleting-a-cluster#overview
api = kube.CoreV1Api()
for pv in api.list_persistent_volume().items:
try:
pv.spec.persistent_volume_reclaim_policy = "Delete"
api.patch_persistent_volume(
name=pv.metadata.name,
body=pv,
)
api.delete_persistent_volume(
name=pv.metadata.name,
grace_period_seconds=0,
propagation_policy="Foreground",
)
except ApiException:
print_exc()
# wait for pv deletes to complete
for _ in range(60):
if not api.list_persistent_volume().items:
break
time.sleep(1)
else:
print("FAILED TO CLEANUP PERSISTENT VOLUMES")
finally:
gke.delete_cluster(name=name)