Merge pull request #114 from seung-lab/heartbeat

Shutdown workers if manager node died

Extend memory monitor script to check manager node heartbeat regularly, if the check fails for half an hour, shutdown the google cloud workers.

Author: ranlu

cloud/google/common.py

@@ -79,7 +79,7 @@ def GenerateAirflowVar(context, hostname_manager):
 
     sqlalchemy_conn = f'''postgresql+psycopg2://{postgres_user}:{postgres_password}@{hostname_manager}/{postgres_db}'''
     airflow_variable = {
-        'AIRFLOW__CORE__HOSTNAME_CALLABLE': 'google_metadata.gce_hostname',
+        'AIRFLOW__CORE__HOSTNAME_CALLABLE': 'common.google_api.gce_hostname',
         'AIRFLOW__DATABASE__SQL_ALCHEMY_CONN': sqlalchemy_conn,
         'AIRFLOW__CORE__FERNET_KEY': context.properties['airflow'].get('fernetKey', fernet_key),
         'AIRFLOW__CELERY__BROKER_URL': f'amqp://{hostname_manager}',

common/google_api.py

@@ -0,0 +1,130 @@
+from googleapiclient import discovery
+import requests
+import json
+
+def get_project_data(key):
+    apiurl = f"http://metadata/computeMetadata/v1/project/{key}"
+    response = requests.get(apiurl, headers={"Metadata-Flavor": "Google"})
+    response.raise_for_status()
+    return response.text
+
+def get_instance_data(key):
+    apiurl = f"http://metadata/computeMetadata/v1/instance/{key}"
+    response = requests.get(apiurl, headers={"Metadata-Flavor": "Google"})
+    response.raise_for_status()
+    return response.text
+
+def get_instance_metadata(project, zone, instance):
+    service = discovery.build('compute', 'v1')
+    request = service.instances().get(project=project, zone=zone, instance=instance)
+    info = request.execute()
+    return info['metadata']
+
+def set_instance_metadata(project, zone, instance, data):
+    service = discovery.build('compute', 'v1')
+    request = service.instances().setMetadata(body=data, project=project, zone=zone, instance=instance)
+    return request.execute()
+
+def gce_external_ip():
+    return get_instance_data("network-interfaces/0/access-configs/0/external-ip")
+
+def gce_internal_ip():
+    return get_instance_data("network-interfaces/0/ip")
+
+def gce_hostname():
+    return get_instance_data("hostname")
+
+def get_project_id():
+    apiurl = "http://metadata.google.internal/computeMetadata/v1/project/project-id"
+    response = requests.get(apiurl, headers={"Metadata-Flavor": "Google"})
+    response.raise_for_status()
+    return response.text
+
+def get_zone():
+    apiurl = "http://metadata.google.internal/computeMetadata/v1/instance/zone"
+    response = requests.get(apiurl, headers={"Metadata-Flavor": "Google"})
+    response.raise_for_status()
+    return response.text.split('/')[-1]
+
+def delete_instances(ig, instances):
+    project_id = get_project_id()
+    request_body = {
+        "instances": instances,
+        "skipInstancesOnValidationError": True,
+    }
+    service = discovery.build('compute', 'v1')
+    request = service.instanceGroupManagers().deleteInstances(project=project_id, zone=ig["zone"], instanceGroupManager=ig["name"], body=request_body)
+    return request.execute()
+
+def get_created_by():
+    apiurl = "http://metadata.google.internal/computeMetadata/v1/instance/attributes/created-by"
+    response = requests.get(apiurl, headers={"Metadata-Flavor": "Google"})
+    response.raise_for_status()
+    return response.text
+
+def instance_group_manager_info(project_id, instance_group):
+    service = discovery.build('compute', 'v1')
+    request = service.instanceGroupManagers().get(project=project_id, zone=instance_group['zone'], instanceGroupManager=instance_group['name'])
+    return request.execute()
+
+def instance_group_manager_error(project_id, instance_group):
+    service = discovery.build('compute', 'v1')
+    request = service.instanceGroupManagers().listErrors(project=project_id, zone=instance_group['zone'], instanceGroupManager=instance_group['name'], orderBy="creationTimestamp desc")
+    return request.execute()
+
+def instance_group_info(project_id, instance_group):
+    service = discovery.build('compute', 'v1')
+    request = service.instanceGroups().get(project=project_id, zone=instance_group['zone'], instanceGroup=instance_group['name'])
+    return request.execute()
+
+def list_managed_instances(instance_group):
+    project_id = get_project_id()
+    service = discovery.build("compute", "v1")
+    page_token = None
+    instances = []
+    while True:
+        request = service.instanceGroupManagers().listManagedInstances(project=project_id, zone=instance_group["zone"], instanceGroupManager=instance_group["name"], pageToken=page_token, maxResults=20)
+        ret = request.execute()
+        if not ret:
+            return instances
+        instances += [r["instance"] for r in ret['managedInstances']]
+        page_token = ret.get("nextPageToken", None)
+        if not page_token:
+            break
+
+    return instances
+
+def get_instance_property(instance_zone, instance, key):
+    project_id = get_project_id()
+    service = discovery.build("compute", "v1")
+    request = service.instances().get(project=project_id, zone=instance_zone, instance=instance)
+    ret = request.execute()
+    return ret[key]
+
+def resize_instance_group(ig, size):
+    project_id = get_project_id()
+    service = discovery.build('compute', 'v1')
+    request = service.instanceGroupManagers().resize(project=project_id, zone=ig['zone'], instanceGroupManager=ig['name'], size=size)
+    return request.execute()
+
+def start_instance(instance_name, zone):
+    service = discovery.build('compute', 'v1')
+    request = service.instances().start(
+        project=get_project_id(),
+        zone=zone,
+        instance=instance_name
+    )
+    response = request.execute()
+    return response
+
+
+def stop_instance(instance_name, zone):
+    service = discovery.build('compute', 'v1')
+    request = service.instances().stop(
+        discardLocalSsd=True,
+        project=get_project_id(),
+        zone=zone,
+        instance=instance_name
+    )
+    response = request.execute()
+    return response